(self, input, sigma, **kwargs)
| 214 | return (sq_error * f_weight).flatten(1).mean(1) * c_weight |
| 215 | |
| 216 | def forward(self, input, sigma, **kwargs): |
| 217 | c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] |
| 218 | # denoised, _, _ = self.inner_model(input * c_in, sigma, **kwargs) |
| 219 | denoised = self.inner_model(input * c_in, sigma, **kwargs)[0] |
| 220 | # return denoised.to(torch.float32) * c_out + input * c_skip |
| 221 | # return denoised.to(torch.float32) |
| 222 | if self.parametrization =="v": |
| 223 | return denoised.to(torch.float32) * c_out + input * c_skip |
| 224 | elif self.parametrization =="x0": |
| 225 | #directly predicts the clean image |
| 226 | return denoised.to(torch.float32) |
| 227 | |
| 228 | |
| 229 | class DenoiserWithVariance(Denoiser): |
nothing calls this directly
no test coverage detected