| 46 | """A Karras et al. preconditioner for denoising diffusion models.""" |
| 47 | |
| 48 | def __init__(self, inner_model, sigma_data=1., weighting='karras', scales=1, parametrization="v", loss_weight_per_channel=None): |
| 49 | super().__init__() |
| 50 | self.inner_model = inner_model |
| 51 | self.sigma_data = sigma_data |
| 52 | self.scales = scales |
| 53 | self.parametrization = parametrization |
| 54 | if callable(weighting): |
| 55 | self.weighting = weighting |
| 56 | if weighting == 'karras': |
| 57 | self.weighting = torch.ones_like |
| 58 | elif weighting == 'soft-min-snr': |
| 59 | self.weighting = self._weighting_soft_min_snr |
| 60 | elif weighting == 'snr': |
| 61 | self.weighting = self._weighting_snr |
| 62 | # elif weighting == 'edm2': |
| 63 | # self.weighting = self._weighting_edm2 |
| 64 | else: |
| 65 | raise ValueError(f'Unknown weighting type {weighting}') |
| 66 | |
| 67 | if loss_weight_per_channel is not None: |
| 68 | self.loss_weight_per_channel=torch.as_tensor(loss_weight_per_channel).cuda() |
| 69 | else: |
| 70 | self.loss_weight_per_channel=None |
| 71 | |
| 72 | |
| 73 | |