| 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] |