(self, network, denoiser, conditioner, input, batch)
| 46 | self.batch2model_keys = set(batch2model_keys) |
| 47 | |
| 48 | def __call__(self, network, denoiser, conditioner, input, batch): |
| 49 | cond = conditioner(batch) |
| 50 | additional_model_inputs = {key: batch[key] for key in self.batch2model_keys.intersection(batch)} |
| 51 | |
| 52 | sigmas = self.sigma_sampler(input.shape[0]).to(input.device) |
| 53 | noise = torch.randn_like(input) |
| 54 | if self.offset_noise_level > 0.0: |
| 55 | noise = ( |
| 56 | noise + append_dims(torch.randn(input.shape[0]).to(input.device), input.ndim) * self.offset_noise_level |
| 57 | ) |
| 58 | noise = noise.to(input.dtype) |
| 59 | noised_input = input.float() + noise * append_dims(sigmas, input.ndim) |
| 60 | model_output = denoiser(network, noised_input, sigmas, cond, **additional_model_inputs) |
| 61 | w = append_dims(denoiser.w(sigmas), input.ndim) |
| 62 | return self.get_loss(model_output, input, w) |
| 63 | |
| 64 | def get_loss(self, model_output, target, w): |
| 65 | if self.type == "l2": |
nothing calls this directly
no test coverage detected