| 17 | |
| 18 | |
| 19 | class RescaleDiffusersLatent: |
| 20 | def __init__(self, divisor, exclude_keys=None): |
| 21 | self.divisor = divisor |
| 22 | self.exclude_keys = exclude_keys or [] |
| 23 | |
| 24 | def __call__(self, sample): |
| 25 | for k, v in sample.items(): |
| 26 | if k.endswith("_latent") and k not in self.exclude_keys: |
| 27 | latent = sample[k] |
| 28 | sample[k] = latent / self.divisor |
| 29 | return sample |
| 30 | |
| 31 | |
| 32 | class AddNoiseLatent: |
nothing calls this directly
no outgoing calls
no test coverage detected