(self, target_channels, z_channels, depth, width, num_sampling_steps, grad_checkpointing=False)
| 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) |
no test coverage detected