(self, inputs)
| 475 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 476 | @torch.compiler.disable() |
| 477 | def forward(self, inputs): |
| 478 | vx, ax, timestep, context, attention_mask = inputs |
| 479 | assert attention_mask.dtype == torch.bool and attention_mask.ndim == 2 |
| 480 | |
| 481 | context, attention_mask = self.preprocess_text_embeds(context, attention_mask) |
| 482 | # convert attention_mask bool -> int so _prepare_attention_mask() works |
| 483 | attention_mask = attention_mask.int() |
| 484 | |
| 485 | x = [vx] |
| 486 | if ax.numel() > 0: |
| 487 | x.append(ax) |
| 488 | transformer_options = {} |
| 489 | |
| 490 | input_dtype = x[0].dtype |
| 491 | batch_size = x[0].shape[0] |
| 492 | |
| 493 | merged_args = {**transformer_options, 'a_timestep': timestep} |
| 494 | # TODO: denoise_mask? |
| 495 | x, pixel_coords, additional_args = self._process_input(x, keyframe_idxs=None, denoise_mask=None) |
| 496 | merged_args.update(additional_args) |
| 497 | |
| 498 | timestep, (v_embedded_timestep, a_embedded_timestep), prompt_timestep = self._prepare_timestep(timestep, batch_size, input_dtype, **merged_args) |
| 499 | merged_args["prompt_timestep"] = prompt_timestep |
| 500 | context, attention_mask = self._prepare_context(context, batch_size, x, attention_mask) |
| 501 | |
| 502 | attention_mask = self._prepare_attention_mask(attention_mask, input_dtype) |
| 503 | pe = self._prepare_positional_embeddings(pixel_coords, self.framerate, input_dtype) |
| 504 | |
| 505 | # Always None |
| 506 | # self_attention_mask = self._build_guide_self_attention_mask( |
| 507 | # x, transformer_options, merged_args |
| 508 | # ) |
| 509 | |
| 510 | vx = x[0] |
| 511 | ax = x[1] |
| 512 | v_context = context[0] |
| 513 | a_context = context[1] |
| 514 | v_timestep = timestep[0] |
| 515 | a_timestep = timestep[1] |
| 516 | v_pe, av_cross_video_freq_cis = pe[0] |
| 517 | a_pe, av_cross_audio_freq_cis = pe[1] |
| 518 | v_pe_cos, v_pe_sin = v_pe[:2] |
| 519 | av_cross_video_freq_cos, av_cross_video_freq_sin = av_cross_video_freq_cis[:2] |
| 520 | a_pe_cos, a_pe_sin = a_pe[:2] |
| 521 | av_cross_audio_freq_cos, av_cross_audio_freq_sin = av_cross_audio_freq_cis[:2] |
| 522 | |
| 523 | ( |
| 524 | av_ca_audio_scale_shift_timestep, |
| 525 | av_ca_video_scale_shift_timestep, |
| 526 | av_ca_a2v_gate_noise_timestep, |
| 527 | av_ca_v2a_gate_noise_timestep, |
| 528 | ) = timestep[2] |
| 529 | v_prompt_timestep = timestep[3] |
| 530 | a_prompt_timestep = timestep[4] |
| 531 | |
| 532 | # Can't pass objects between PP layers, expand back to tensors. The compression doesn't get enabled anyway (why does CompressedTimestep even exist then?) |
| 533 | tmp = [v_timestep, a_timestep, av_ca_audio_scale_shift_timestep, av_ca_video_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] |
| 534 | for i, t in enumerate(tmp): |
nothing calls this directly
no test coverage detected