(
self,
enforce_zero_snr: bool = True,
diffusion_parameterization: str = 'v',
diffusion_schedule: str = 'linear',
*args,
**kwargs
)
| 18 | |
| 19 | class FlowModelObj(FlowModel): |
| 20 | def __init__( |
| 21 | self, |
| 22 | enforce_zero_snr: bool = True, |
| 23 | diffusion_parameterization: str = 'v', |
| 24 | diffusion_schedule: str = 'linear', |
| 25 | *args, |
| 26 | **kwargs |
| 27 | ): |
| 28 | super().__init__(*args, **kwargs) |
| 29 | self.register_sdv2_schedule(diffusion_schedule, enforce_zero_snr) |
| 30 | assert diffusion_parameterization in ['v', 'eps'], 'Diffusion parameterization has to be either v or eps' |
| 31 | self.diffusion_parameterization = diffusion_parameterization |
| 32 | self.diffusion_schedule = diffusion_schedule |
| 33 | |
| 34 | def ode_fn(self, t, x, **kwargs): |
| 35 | if t.numel() == 1: |
nothing calls this directly
no test coverage detected