(self, sample, timestep, encoder_hidden_states, add_time_id, use_gradient_checkpointing=False, **kwargs)
| 371 | |
| 372 | |
| 373 | def forward(self, sample, timestep, encoder_hidden_states, add_time_id, use_gradient_checkpointing=False, **kwargs): |
| 374 | # 1. time |
| 375 | timestep = torch.tensor((timestep,)).to(sample.device) |
| 376 | t_emb = self.time_proj(timestep).to(sample.dtype) |
| 377 | t_emb = self.time_embedding(t_emb) |
| 378 | |
| 379 | add_embeds = self.add_time_proj(add_time_id.flatten()).to(sample.dtype) |
| 380 | add_embeds = add_embeds.reshape((-1, 768)) |
| 381 | add_embeds = self.add_time_embedding(add_embeds) |
| 382 | |
| 383 | time_emb = t_emb + add_embeds |
| 384 | |
| 385 | # 2. pre-process |
| 386 | height, width = sample.shape[2], sample.shape[3] |
| 387 | hidden_states = self.conv_in(sample) |
| 388 | text_emb = encoder_hidden_states |
| 389 | res_stack = [hidden_states] |
| 390 | |
| 391 | # 3. blocks |
| 392 | def create_custom_forward(module): |
| 393 | def custom_forward(*inputs): |
| 394 | return module(*inputs) |
| 395 | return custom_forward |
| 396 | for i, block in enumerate(self.blocks): |
| 397 | if self.training and use_gradient_checkpointing and not (isinstance(block, PushBlock) or isinstance(block, PopBlock) or isinstance(block, PopMixBlock)): |
| 398 | hidden_states, time_emb, text_emb, res_stack = torch.utils.checkpoint.checkpoint( |
| 399 | create_custom_forward(block), |
| 400 | hidden_states, time_emb, text_emb, res_stack, |
| 401 | use_reentrant=False, |
| 402 | ) |
| 403 | else: |
| 404 | hidden_states, time_emb, text_emb, res_stack = block(hidden_states, time_emb, text_emb, res_stack) |
| 405 | |
| 406 | # 4. output |
| 407 | hidden_states = self.conv_norm_out(hidden_states) |
| 408 | hidden_states = self.conv_act(hidden_states) |
| 409 | hidden_states = self.conv_out(hidden_states) |
| 410 | |
| 411 | return hidden_states |
| 412 | |
| 413 | @staticmethod |
| 414 | def state_dict_converter(): |
no test coverage detected