(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: Optional[torch.Tensor] = None,
encoder_hidden_states_2: Optional[torch.Tensor] = None,
timestep: Optional[torch.LongTensor] = None,
added_cond_kwargs: Dict[str, torch.Tensor] = None,
encoder_attention_mask: Optional[torch.Tensor] = None,
fps: torch.Tensor=None,
return_dict: bool = False,
)
| 846 | |
| 847 | @torch.inference_mode() |
| 848 | def forward( |
| 849 | self, |
| 850 | hidden_states: torch.Tensor, |
| 851 | encoder_hidden_states: Optional[torch.Tensor] = None, |
| 852 | encoder_hidden_states_2: Optional[torch.Tensor] = None, |
| 853 | timestep: Optional[torch.LongTensor] = None, |
| 854 | added_cond_kwargs: Dict[str, torch.Tensor] = None, |
| 855 | encoder_attention_mask: Optional[torch.Tensor] = None, |
| 856 | fps: torch.Tensor=None, |
| 857 | return_dict: bool = False, |
| 858 | ): |
| 859 | assert hidden_states.ndim==5; "hidden_states's shape should be (bsz, f, ch, h ,w)" |
| 860 | |
| 861 | bsz, frame, _, height, width = hidden_states.shape |
| 862 | height, width = height // self.patch_size, width // self.patch_size |
| 863 | |
| 864 | hidden_states = self.patchfy(hidden_states) |
| 865 | len_frame = hidden_states.shape[1] |
| 866 | |
| 867 | if self.use_additional_conditions: |
| 868 | added_cond_kwargs = { |
| 869 | "resolution": torch.tensor([(height, width)]*bsz, device=hidden_states.device, dtype=hidden_states.dtype), |
| 870 | "nframe": torch.tensor([frame]*bsz, device=hidden_states.device, dtype=hidden_states.dtype), |
| 871 | "fps": fps |
| 872 | } |
| 873 | else: |
| 874 | added_cond_kwargs = {} |
| 875 | |
| 876 | timestep, embedded_timestep = self.adaln_single( |
| 877 | timestep, added_cond_kwargs=added_cond_kwargs |
| 878 | ) |
| 879 | |
| 880 | encoder_hidden_states = self.caption_projection(self.caption_norm(encoder_hidden_states)) |
| 881 | |
| 882 | if encoder_hidden_states_2 is not None and hasattr(self, 'clip_projection'): |
| 883 | clip_embedding = self.clip_projection(encoder_hidden_states_2) |
| 884 | encoder_hidden_states = torch.cat([clip_embedding, encoder_hidden_states], dim=1) |
| 885 | |
| 886 | hidden_states = rearrange(hidden_states, '(b f) l d-> b (f l) d', b=bsz, f=frame, l=len_frame).contiguous() |
| 887 | encoder_hidden_states, attn_mask = self.prepare_attn_mask(encoder_attention_mask, encoder_hidden_states, q_seqlen=frame*len_frame) |
| 888 | |
| 889 | hidden_states = self.block_forward( |
| 890 | hidden_states, |
| 891 | encoder_hidden_states, |
| 892 | timestep=timestep, |
| 893 | rope_positions=[frame, height, width], |
| 894 | attn_mask=attn_mask, |
| 895 | parallel=self.parallel |
| 896 | ) |
| 897 | |
| 898 | hidden_states = rearrange(hidden_states, 'b (f l) d -> (b f) l d', b=bsz, f=frame, l=len_frame) |
| 899 | |
| 900 | embedded_timestep = repeat(embedded_timestep, 'b d -> (b f) d', f=frame).contiguous() |
| 901 | |
| 902 | shift, scale = (self.scale_shift_table[None].to(dtype=embedded_timestep.dtype, device=embedded_timestep.device) + embedded_timestep[:, None]).chunk(2, dim=1) |
| 903 | hidden_states = self.norm_out(hidden_states) |
| 904 | # Modulation |
| 905 | hidden_states = hidden_states * (1 + scale) + shift |
nothing calls this directly
no test coverage detected