Diffusion Loss
| 7 | |
| 8 | |
| 9 | class DiffLoss(nn.Module): |
| 10 | """Diffusion Loss""" |
| 11 | def __init__(self, target_channels, z_channels, depth, width, num_sampling_steps, grad_checkpointing=False): |
| 12 | super(DiffLoss, self).__init__() |
| 13 | self.in_channels = target_channels |
| 14 | self.net = SimpleMLPAdaLN( |
| 15 | in_channels=target_channels, |
| 16 | model_channels=width, |
| 17 | out_channels=target_channels * 2, # for vlb loss |
| 18 | z_channels=z_channels, |
| 19 | num_res_blocks=depth, |
| 20 | grad_checkpointing=grad_checkpointing |
| 21 | ) |
| 22 | |
| 23 | self.train_diffusion = create_diffusion(timestep_respacing="", noise_schedule="cosine") |
| 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 |
| 37 | if not cfg == 1.0: |
| 38 | noise = torch.randn(z.shape[0] // 2, self.in_channels).cuda() |
| 39 | noise = torch.cat([noise, noise], dim=0) |
| 40 | model_kwargs = dict(c=z, cfg_scale=cfg) |
| 41 | sample_fn = self.net.forward_with_cfg |
| 42 | else: |
| 43 | noise = torch.randn(z.shape[0], self.in_channels).cuda() |
| 44 | model_kwargs = dict(c=z) |
| 45 | sample_fn = self.net.forward |
| 46 | |
| 47 | sampled_token_latent = self.gen_diffusion.p_sample_loop( |
| 48 | sample_fn, noise.shape, noise, clip_denoised=False, model_kwargs=model_kwargs, progress=False, |
| 49 | temperature=temperature |
| 50 | ) |
| 51 | |
| 52 | return sampled_token_latent |
| 53 | |
| 54 | |
| 55 | def modulate(x, shift, scale): |