x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images), t: (N,) tensor of diffusion timesteps y: (N,) tensor of class labels
(
self,
hidden_states,
t,
y=None,
)
| 909 | return ckpt_forward |
| 910 | |
| 911 | def forward( |
| 912 | self, |
| 913 | hidden_states, |
| 914 | t, |
| 915 | y=None, |
| 916 | ): |
| 917 | """ |
| 918 | x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images), |
| 919 | |
| 920 | t: (N,) tensor of diffusion timesteps |
| 921 | y: (N,) tensor of class labels |
| 922 | """ |
| 923 | hidden_states = self.x_embedder( |
| 924 | hidden_states |
| 925 | ) # (N, T, D), where T = H * W / patch_size ** 2, if video_frames>0, T = H * W * video_frames / patch_size ** 2 |
| 926 | _B, _T, _D = hidden_states.shape |
| 927 | |
| 928 | t = (t * 1000.0).to(hidden_states) |
| 929 | t = self.t_embedder(t) # (N, D) |
| 930 | if self.has_text: |
| 931 | # y = self.y_embedder(y, self.training) # (B, N, D) |
| 932 | y = self.y_embedder(y) # (B, N, D) |
| 933 | c = t + y.mean(dim=1) # (N, D) |
| 934 | elif self.num_classes > 0: |
| 935 | c = t + self.y_embedder(y, self.training) # (N, D) |
| 936 | else: |
| 937 | c = t |
| 938 | |
| 939 | if self.use_pe == 1 or self.use_pe == 2: |
| 940 | hidden_states = hidden_states + self.pos_embed |
| 941 | if self.video_frames > 0 and self.tpe: |
| 942 | # temporal pos |
| 943 | hidden_states = rearrange( |
| 944 | hidden_states, "b (t l) c -> (b l) t c", t=self.video_frames |
| 945 | ) |
| 946 | hidden_states = hidden_states + self.temporal_pos_embedding |
| 947 | hidden_states = rearrange(hidden_states, "(b l) t c -> b (t l) c", b=_B) |
| 948 | |
| 949 | residual = None |
| 950 | for layer_idx, block in enumerate(self.blocks): |
| 951 | if self.use_pe == 3: |
| 952 | hidden_states = hidden_states + self.pos_embed_list[layer_idx] |
| 953 | if self.use_checkpoint: |
| 954 | hidden_states, residual = torch.utils.checkpoint.checkpoint( |
| 955 | self.ckpt_wrapper(block), hidden_states, residual, c, y |
| 956 | ) |
| 957 | else: |
| 958 | hidden_states, residual = block( |
| 959 | hidden_states, residual=residual, c=c, text=y |
| 960 | ) # (N, T, D) |
| 961 | |
| 962 | ##### finished the Mamba blocks, here we apply the last Normalization layer |
| 963 | if not self.fused_add_norm: |
| 964 | if residual is None: |
| 965 | residual = hidden_states |
| 966 | else: |
| 967 | residual = residual + self.drop_path(hidden_states) |
| 968 | hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype)) |
no test coverage detected