(self, input_ids, encoder_hidden_states, pooled_text_emb, micro_conds, cross_attention_kwargs=None)
| 152 | pass |
| 153 | |
| 154 | def forward(self, input_ids, encoder_hidden_states, pooled_text_emb, micro_conds, cross_attention_kwargs=None): |
| 155 | encoder_hidden_states = self.encoder_proj(encoder_hidden_states) |
| 156 | encoder_hidden_states = self.encoder_proj_layer_norm(encoder_hidden_states) |
| 157 | |
| 158 | micro_cond_embeds = get_timestep_embedding( |
| 159 | micro_conds.flatten(), self.config.micro_cond_encode_dim, flip_sin_to_cos=True, downscale_freq_shift=0 |
| 160 | ) |
| 161 | |
| 162 | micro_cond_embeds = micro_cond_embeds.reshape((input_ids.shape[0], -1)) |
| 163 | |
| 164 | pooled_text_emb = torch.cat([pooled_text_emb, micro_cond_embeds], dim=1) |
| 165 | pooled_text_emb = pooled_text_emb.to(dtype=self.dtype) |
| 166 | pooled_text_emb = self.cond_embed(pooled_text_emb).to(encoder_hidden_states.dtype) |
| 167 | |
| 168 | hidden_states = self.embed(input_ids) |
| 169 | |
| 170 | hidden_states = self.down_block( |
| 171 | hidden_states, |
| 172 | pooled_text_emb=pooled_text_emb, |
| 173 | encoder_hidden_states=encoder_hidden_states, |
| 174 | cross_attention_kwargs=cross_attention_kwargs, |
| 175 | ) |
| 176 | |
| 177 | batch_size, channels, height, width = hidden_states.shape |
| 178 | hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels) |
| 179 | |
| 180 | hidden_states = self.project_to_hidden_norm(hidden_states) |
| 181 | hidden_states = self.project_to_hidden(hidden_states) |
| 182 | |
| 183 | for layer in self.transformer_layers: |
| 184 | if self.training and self.gradient_checkpointing: |
| 185 | |
| 186 | def layer_(*args): |
| 187 | return checkpoint(layer, *args) |
| 188 | |
| 189 | else: |
| 190 | layer_ = layer |
| 191 | |
| 192 | hidden_states = layer_( |
| 193 | hidden_states, |
| 194 | encoder_hidden_states=encoder_hidden_states, |
| 195 | cross_attention_kwargs=cross_attention_kwargs, |
| 196 | added_cond_kwargs={"pooled_text_emb": pooled_text_emb}, |
| 197 | ) |
| 198 | |
| 199 | hidden_states = self.project_from_hidden_norm(hidden_states) |
| 200 | hidden_states = self.project_from_hidden(hidden_states) |
| 201 | |
| 202 | hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2) |
| 203 | |
| 204 | hidden_states = self.up_block( |
| 205 | hidden_states, |
| 206 | pooled_text_emb=pooled_text_emb, |
| 207 | encoder_hidden_states=encoder_hidden_states, |
| 208 | cross_attention_kwargs=cross_attention_kwargs, |
| 209 | ) |
| 210 | |
| 211 | logits = self.mlm_layer(hidden_states) |
nothing calls this directly
no test coverage detected