| 364 | |
| 365 | |
| 366 | class AudioProjModel(ModelMixin, ConfigMixin): |
| 367 | def __init__( |
| 368 | self, |
| 369 | seq_len=5, |
| 370 | seq_len_vf=12, |
| 371 | blocks=12, |
| 372 | channels=768, |
| 373 | intermediate_dim=512, |
| 374 | output_dim=768, |
| 375 | context_tokens=32, |
| 376 | norm_output_audio=False, |
| 377 | ): |
| 378 | super().__init__() |
| 379 | |
| 380 | self.seq_len = seq_len |
| 381 | self.blocks = blocks |
| 382 | self.channels = channels |
| 383 | self.input_dim = seq_len * blocks * channels |
| 384 | self.input_dim_vf = seq_len_vf * blocks * channels |
| 385 | self.intermediate_dim = intermediate_dim |
| 386 | self.context_tokens = context_tokens |
| 387 | self.output_dim = output_dim |
| 388 | |
| 389 | # define multiple linear layers |
| 390 | self.proj1 = nn.Linear(self.input_dim, intermediate_dim) |
| 391 | self.proj1_vf = nn.Linear(self.input_dim_vf, intermediate_dim) |
| 392 | self.proj2 = nn.Linear(intermediate_dim, intermediate_dim) |
| 393 | self.proj3 = nn.Linear(intermediate_dim, context_tokens * output_dim) |
| 394 | self.norm = nn.LayerNorm(output_dim) if norm_output_audio else nn.Identity() |
| 395 | |
| 396 | def forward(self, audio_embeds, audio_embeds_vf): |
| 397 | video_length = audio_embeds.shape[1] + audio_embeds_vf.shape[1] |
| 398 | B, _, _, S, C = audio_embeds.shape |
| 399 | |
| 400 | # process audio of first frame |
| 401 | audio_embeds = rearrange(audio_embeds, "bz f w b c -> (bz f) w b c") |
| 402 | batch_size, window_size, blocks, channels = audio_embeds.shape |
| 403 | audio_embeds = audio_embeds.view(batch_size, window_size * blocks * channels) |
| 404 | |
| 405 | # process audio of latter frame |
| 406 | audio_embeds_vf = rearrange(audio_embeds_vf, "bz f w b c -> (bz f) w b c") |
| 407 | batch_size_vf, window_size_vf, blocks_vf, channels_vf = audio_embeds_vf.shape |
| 408 | audio_embeds_vf = audio_embeds_vf.view(batch_size_vf, window_size_vf * blocks_vf * channels_vf) |
| 409 | |
| 410 | # first projection |
| 411 | audio_embeds = torch.relu(self.proj1(audio_embeds)) |
| 412 | audio_embeds_vf = torch.relu(self.proj1_vf(audio_embeds_vf)) |
| 413 | audio_embeds = rearrange(audio_embeds, "(bz f) c -> bz f c", bz=B) |
| 414 | audio_embeds_vf = rearrange(audio_embeds_vf, "(bz f) c -> bz f c", bz=B) |
| 415 | audio_embeds_c = torch.concat([audio_embeds, audio_embeds_vf], dim=1) |
| 416 | batch_size_c, N_t, C_a = audio_embeds_c.shape |
| 417 | audio_embeds_c = audio_embeds_c.view(batch_size_c*N_t, C_a) |
| 418 | |
| 419 | # second projection |
| 420 | audio_embeds_c = torch.relu(self.proj2(audio_embeds_c)) |
| 421 | |
| 422 | context_tokens = self.proj3(audio_embeds_c).reshape(batch_size_c*N_t, self.context_tokens, self.output_dim) |
| 423 | |