(self, x_start, cond, t, noise=None)
| 477 | self.cond_stage_model = self.cond_stage_model.cuda() |
| 478 | |
| 479 | def p_losses(self, x_start, cond, t, noise=None): |
| 480 | |
| 481 | cond_mask = {} |
| 482 | cond_mask["c_crossattn"] = [cond["c_crossattn"][0]] |
| 483 | cond_mask["c_concat"] = [cond["c_concat_mask"][0]] |
| 484 | |
| 485 | cond_image = {} |
| 486 | cond_image["c_crossattn"] = [cond["c_crossattn"][0]] |
| 487 | cond_image["c_concat"] = [cond["c_concat_mask"][0]] |
| 488 | cond_image["c_concat_image"] = [cond["c_concat_image"][0]] |
| 489 | |
| 490 | weights_ones = torch.ones_like(t).to(x_start.device) |
| 491 | weights_thre = torch.where(t <= 200, torch.tensor(1), torch.tensor(0)) |
| 492 | |
| 493 | weights_mask = 1.0 * weights_ones # Loss 0 |
| 494 | weights_image = 1.0 * weights_ones # Loss 1 |
| 495 | weights_mask_2_image = 1.0 * weights_ones # Loss 2 |
| 496 | weights_mask_regularization = 1.0 * weights_thre # Loss 3 |
| 497 | |
| 498 | noise = default(noise, lambda: torch.randn_like(x_start)) |
| 499 | x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) |
| 500 | model_output_mask = self.apply_model(x_noisy, t, cond_mask) |
| 501 | |
| 502 | loss_dict = {} |
| 503 | prefix = 'train' if self.training else 'val' |
| 504 | |
| 505 | if self.parameterization == "x0": |
| 506 | target = x_start |
| 507 | elif self.parameterization == "eps": |
| 508 | target = noise |
| 509 | elif self.parameterization == "v": |
| 510 | target = self.get_v(x_start, noise, t) |
| 511 | else: |
| 512 | raise NotImplementedError() |
| 513 | |
| 514 | loss_simple = weights_mask * self.get_loss(model_output_mask, target, mean=False).mean([1, 2, 3]) |
| 515 | print(f"loss_simple_mask: {loss_simple.mean()}") |
| 516 | |
| 517 | if weights_image.all(): |
| 518 | model_output_image = self.apply_model(x_noisy, t, cond_image) |
| 519 | loss_simple_image = self.get_loss(model_output_image, target, mean=False).mean([1, 2, 3]) |
| 520 | print(f"loss_simple_image: {loss_simple_image.mean()}") |
| 521 | loss_simple = loss_simple + weights_image * loss_simple_image |
| 522 | |
| 523 | if weights_mask_2_image.all(): |
| 524 | loss_simple_mask_2_image = self.get_loss(model_output_mask, model_output_image.detach(), mean=False).mean([1, 2, 3]) |
| 525 | print(f"loss_simple_mask_2_image: {loss_simple_mask_2_image.mean()}") |
| 526 | loss_simple = loss_simple + weights_mask_2_image * loss_simple_mask_2_image |
| 527 | |
| 528 | if (self.global_step > (self.trainer.max_steps * 1 / 3)) and weights_mask_regularization.any(): # Done! |
| 529 | recon_output_image = self.predict_start_from_noise(x_noisy, t=t, noise=model_output_image) |
| 530 | noise_image_2_mask = default(noise, lambda: torch.randn_like(recon_output_image)) |
| 531 | x_noisy_mask_recon = self.q_sample(x_start=recon_output_image, t=t, noise=noise_image_2_mask) |
| 532 | |
| 533 | model_output_mask_xt = self.apply_model(x_noisy_mask_recon.detach(), t, cond_mask) |
| 534 | loss_simple_mask_regularization = self.get_loss(model_output_mask_xt, noise_image_2_mask, mean=False).mean([1, 2, 3]) |
| 535 | print(f"loss_simple_mask_regularization: {loss_simple_mask_regularization.mean()}") |
| 536 | loss_simple = loss_simple + weights_mask_regularization * loss_simple_mask_regularization |
nothing calls this directly
no test coverage detected