(
self,
sigma_sampler_config,
type="l2",
offset_noise_level=0.0,
linear_offset_noise=False,
batch2model_keys: Optional[Union[str, List[str], ListConfig]] = None,
)
| 17 | |
| 18 | class StandardDiffusionLoss(nn.Module): |
| 19 | def __init__( |
| 20 | self, |
| 21 | sigma_sampler_config, |
| 22 | type="l2", |
| 23 | offset_noise_level=0.0, |
| 24 | linear_offset_noise=False, |
| 25 | batch2model_keys: Optional[Union[str, List[str], ListConfig]] = None, |
| 26 | ): |
| 27 | super().__init__() |
| 28 | |
| 29 | assert type in ["l2", "l1", "lpips"] |
| 30 | |
| 31 | self.sigma_sampler = instantiate_from_config(sigma_sampler_config) |
| 32 | |
| 33 | self.type = type |
| 34 | self.offset_noise_level = offset_noise_level |
| 35 | self.linear_offset_noise = linear_offset_noise |
| 36 | |
| 37 | if type == "lpips": |
| 38 | self.lpips = LPIPS().eval() |
| 39 | |
| 40 | if not batch2model_keys: |
| 41 | batch2model_keys = [] |
| 42 | |
| 43 | if isinstance(batch2model_keys, str): |
| 44 | batch2model_keys = [batch2model_keys] |
| 45 | |
| 46 | self.batch2model_keys = set(batch2model_keys) |
| 47 | |
| 48 | def __call__(self, network, denoiser, conditioner, input, batch): |
| 49 | cond = conditioner(batch) |
nothing calls this directly
no test coverage detected