(self, inputs)
| 180 | |
| 181 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 182 | def forward(self, inputs): |
| 183 | for item in inputs: |
| 184 | if torch.is_floating_point(item): |
| 185 | item.requires_grad_(True) |
| 186 | hidden_states, timestep, encoder_hidden_states, pooled_projections = inputs |
| 187 | |
| 188 | height, width = hidden_states.shape[-2:] |
| 189 | latent_size = torch.tensor([height, width]).to(hidden_states.device) |
| 190 | |
| 191 | hidden_states = self.pos_embed(hidden_states) # takes care of adding positional embeddings too. |
| 192 | temb = self.time_text_embed(timestep, pooled_projections) |
| 193 | encoder_hidden_states = self.context_embedder(encoder_hidden_states) |
| 194 | |
| 195 | return make_contiguous(hidden_states, temb, latent_size, encoder_hidden_states) |
| 196 | |
| 197 | |
| 198 | class TransformerLayer(nn.Module): |
nothing calls this directly
no test coverage detected