(self, x, timesteps=None, context=None, y=None, **kwargs)
| 1019 | return |
| 1020 | |
| 1021 | def forward(self, x, timesteps=None, context=None, y=None, **kwargs): |
| 1022 | b, t, d, h, w = x.shape # * [1, 13, 16, 64, 112] |
| 1023 | if x.dtype != self.dtype: |
| 1024 | x = x.to(self.dtype) |
| 1025 | assert (y is not None) == ( |
| 1026 | self.num_classes is not None |
| 1027 | ), "must specify y if and only if the model is class-conditional" |
| 1028 | t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False, dtype=self.dtype) # * [1, 1920] |
| 1029 | |
| 1030 | aug_timesteps = kwargs.get("aug_t_chunk", None) |
| 1031 | if aug_timesteps is not None: |
| 1032 | aug_t_emb = timestep_embedding(aug_timesteps, self.model_channels, repeat_only=False, dtype=self.dtype) |
| 1033 | aug_emb = self.time_embed(aug_t_emb) |
| 1034 | |
| 1035 | emb = self.time_embed(t_emb) # * [1, 512] |
| 1036 | |
| 1037 | if self.num_classes is not None: |
| 1038 | # assert y.shape[0] == x.shape[0] |
| 1039 | assert x.shape[0] % y.shape[0] == 0 |
| 1040 | y = y.repeat_interleave(x.shape[0] // y.shape[0], dim=0) |
| 1041 | emb = emb + self.label_emb(y) |
| 1042 | |
| 1043 | kwargs["seq_length"] = t * h * w // (self.patch_size**2) |
| 1044 | kwargs["images"] = x |
| 1045 | kwargs["emb"] = emb # * Used in AdaLNMixin and FinalLayerMixin [1, 512] |
| 1046 | if aug_timesteps is not None: |
| 1047 | kwargs["aug_emb"] = aug_emb |
| 1048 | |
| 1049 | # Checked: Is it the reason for longer main3 not working? |
| 1050 | # Verified, No. For main3, split_cond_flag is False |
| 1051 | kwargs["split_cond_flag"] = self.allow_split_cond and \ |
| 1052 | 'cond_inds' in kwargs.keys() and \ |
| 1053 | kwargs['cond_inds'] != [] and \ |
| 1054 | len(kwargs['cond_inds']) > 0 and \ |
| 1055 | aug_timesteps is not None |
| 1056 | |
| 1057 | kwargs["encoder_outputs"] = context # torch.Size([2, 226, 4096]) |
| 1058 | kwargs["encoder_outputs"] = context # torch.Size([2, 226, 4096]) |
| 1059 | kwargs["text_length"] = context.shape[1] # 226 |
| 1060 | |
| 1061 | kwargs["input_ids"] = kwargs["position_ids"] = kwargs["attention_mask"] = torch.ones((1, 1)).to(x.dtype) |
| 1062 | output = super().forward(**kwargs)[0] |
| 1063 | |
| 1064 | return output |
nothing calls this directly
no test coverage detected