(self, inputs)
| 410 | |
| 411 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 412 | def forward(self, inputs): |
| 413 | x, x_t, affline_emb_B_D, crossattn_emb, rope_emb_L_1_1_D, adaln_lora_B_3D, extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D, original_shape, sigma = inputs |
| 414 | original_shape = original_shape.tolist() |
| 415 | |
| 416 | x_B_T_H_W_D = rearrange(x, "T H W B D -> B T H W D") |
| 417 | output = self.decoder_head( |
| 418 | x_B_T_H_W_D=x_B_T_H_W_D, |
| 419 | emb_B_D=affline_emb_B_D, |
| 420 | crossattn_emb=None, |
| 421 | origin_shape=original_shape, |
| 422 | crossattn_mask=None, |
| 423 | adaln_lora_B_3D=adaln_lora_B_3D, |
| 424 | ) |
| 425 | |
| 426 | c_skip, c_out, c_in, c_noise = self.pipeline.model.scaling(sigma=sigma) |
| 427 | c_skip = c_skip.view(-1, 1, 1, 1, 1) |
| 428 | c_out = c_out.view(-1, 1, 1, 1, 1) |
| 429 | sigma = sigma.view(-1, 1, 1, 1, 1) |
| 430 | x0_pred = c_skip*x_t + c_out*output |
| 431 | weights_per_sigma = get_per_sigma_loss_weights(sigma) |
| 432 | return x0_pred, weights_per_sigma |
nothing calls this directly
no test coverage detected