(self, target, z, mask=None)
| 24 | self.gen_diffusion = create_diffusion(timestep_respacing=num_sampling_steps, noise_schedule="cosine") |
| 25 | |
| 26 | def forward(self, target, z, mask=None): |
| 27 | t = torch.randint(0, self.train_diffusion.num_timesteps, (target.shape[0],), device=target.device) |
| 28 | model_kwargs = dict(c=z) |
| 29 | loss_dict = self.train_diffusion.training_losses(self.net, target, t, model_kwargs) |
| 30 | loss = loss_dict["loss"] |
| 31 | if mask is not None: |
| 32 | loss = (loss * mask).sum() / mask.sum() |
| 33 | return loss.mean() |
| 34 | |
| 35 | def sample(self, z, temperature=1.0, cfg=1.0): |
| 36 | # diffusion loss sampling |
nothing calls this directly
no test coverage detected