| 595 | |
| 596 | |
| 597 | class FinalLayer(nn.Module): |
| 598 | def __init__(self, model): |
| 599 | super().__init__() |
| 600 | #self.scale_shift_table = model.scale_shift_table |
| 601 | self.norm_out = model.norm_out |
| 602 | self.proj_out = model.proj_out |
| 603 | #self.audio_scale_shift_table = model.audio_scale_shift_table |
| 604 | self.audio_norm_out = model.audio_norm_out |
| 605 | self.audio_proj_out = model.audio_proj_out |
| 606 | self.model = [model] |
| 607 | |
| 608 | def __getattr__(self, name): |
| 609 | return getattr(self.model[0], name) |
| 610 | |
| 611 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 612 | @torch.compiler.disable() |
| 613 | def forward(self, inputs): |
| 614 | ( |
| 615 | vx, ax, v_context, a_context, attention_mask, v_timestep, a_timestep, v_pe_cos, v_pe_sin, a_pe_cos, a_pe_sin, av_cross_video_freq_cos, av_cross_video_freq_sin, av_cross_audio_freq_cos, av_cross_audio_freq_sin, |
| 616 | av_ca_video_scale_shift_timestep, av_ca_audio_scale_shift_timestep, av_ca_a2v_gate_noise_timestep, av_ca_v2a_gate_noise_timestep, v_prompt_timestep, a_prompt_timestep, v_embedded_timestep, a_embedded_timestep, orig_shape |
| 617 | ) = inputs |
| 618 | # TODO: this will return a list with 2 elements when audio is enabled (currently single tensor) |
| 619 | return self._process_output( |
| 620 | [vx, ax], |
| 621 | [v_embedded_timestep, a_embedded_timestep], |
| 622 | keyframe_idxs=None, |
| 623 | orig_shape=orig_shape.tolist(), |
| 624 | ) |