| 39 | """ |
| 40 | |
| 41 | def __init__( |
| 42 | self, |
| 43 | ema_decay: Union[None, float] = None, |
| 44 | monitor: Union[None, str] = None, |
| 45 | input_key: str = "jpg", |
| 46 | ): |
| 47 | super().__init__() |
| 48 | |
| 49 | self.input_key = input_key |
| 50 | self.use_ema = ema_decay is not None |
| 51 | if monitor is not None: |
| 52 | self.monitor = monitor |
| 53 | |
| 54 | if self.use_ema: |
| 55 | self.model_ema = LitEma(self, decay=ema_decay) |
| 56 | logpy.info(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") |
| 57 | |
| 58 | if version.parse(torch.__version__) >= version.parse("2.0.0"): |
| 59 | self.automatic_optimization = False |
| 60 | |
| 61 | def apply_ckpt(self, ckpt: Union[None, str, dict]): |
| 62 | if ckpt is None: |