(self, input, noise, sigma, **kwargs)
| 87 | return c_skip, c_out, c_in |
| 88 | |
| 89 | def loss(self, input, noise, sigma, **kwargs): |
| 90 | # print("Denoiser") |
| 91 | c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] |
| 92 | # print("sigma",sigma) |
| 93 | c_weight = self.weighting(sigma) |
| 94 | # print("c_weight is ", c_weight) |
| 95 | # edm2_weight = (sigma ** 2 + self.sigma_data ** 2) / (sigma * self.sigma_data) ** 2 |
| 96 | weight = c_weight |
| 97 | # weight = torch.ones_like(sigma) |
| 98 | |
| 99 | |
| 100 | noised_input = input + noise * utils.append_dims(sigma, input.ndim) |
| 101 | # model_output, multires_output, logvar = self.inner_model(noised_input * c_in, sigma, **kwargs) |
| 102 | if 'step' in kwargs: |
| 103 | step = kwargs['step'] |
| 104 | del kwargs['step'] |
| 105 | result = self.inner_model(noised_input * c_in, sigma, **kwargs) |
| 106 | if len(result)==2: |
| 107 | # model_output, multires_output, logvar = result |
| 108 | model_output, logvar = result |
| 109 | elif len(result)==4: |
| 110 | |
| 111 | model_output, multires_output, logvar,clip_feature_embedding = result |
| 112 | |
| 113 | #if you have channel weights they are usually computed from the script create_strand_latent_weights.py |
| 114 | if self.loss_weight_per_channel is None: |
| 115 | loss_weight_per_channel=torch.ones(model_output.shape[1], device=model_output.device) |
| 116 | else: |
| 117 | loss_weight_per_channel=self.loss_weight_per_channel |
| 118 | loss_weight_per_channel=loss_weight_per_channel.view(1,-1,1,1) |
| 119 | |
| 120 | |
| 121 | #loss for regions with zero density can be set to zero |
| 122 | density_map=input[:,-1:,:,:] |
| 123 | density_map=density_map*(0.5/self.sigma_data) + 0.5 #after training on lambda |
| 124 | density_map=density_map.clamp(0, 1) |
| 125 | #low density regions |
| 126 | low_density=density_map<0.001 |
| 127 | loss_weight_spatial=torch.ones_like(density_map, device=input.device) |
| 128 | loss_weight_spatial[low_density]=0.0 |
| 129 | loss_weight_spatial=loss_weight_spatial.repeat(1,model_output.shape[1],1,1) |
| 130 | loss_weight_spatial[:,-1:,:,:]=1.0 #density weight is all ones |
| 131 | # print("loss_weight_spatial",loss_weight_spatial.shape) |
| 132 | # print("model_output",model_output.shape) |
| 133 | |
| 134 | |
| 135 | |
| 136 | |
| 137 | |
| 138 | |
| 139 | if self.parametrization =="v": |
| 140 | #the original loss used in k-diffusion |
| 141 | target = (input - c_skip * noised_input) / c_out |
| 142 | mse_loss=(loss_weight_spatial*loss_weight_per_channel*(( (model_output.to(torch.float32)*c_out + noised_input*c_skip) - input) ** 2)).flatten(1).mean(1) |
| 143 | elif self.parametrization =="x0": |
| 144 | #directly predict target |
| 145 | target=input |
| 146 | mse_loss=(loss_weight_spatial*loss_weight_per_channel*(( model_output.to(torch.float32) - input) ** 2)).flatten(1).mean(1) |
no test coverage detected