Method ode_fn

diff2flow/flow_obj.py:34–38

(self, t, x, **kwargs)

        self.diffusion_schedule = diffusion_schedule

    def ode_fn(self, t, x, **kwargs):
        if t.numel() == 1:
            t = t.expand(x.shape[0])
        _pred = self.sample_vt(x, t, **kwargs)
        return _pred
    
    def register_sdv2_schedule(self, diffusion_schedule, enforce_zero_snr=True):
        # SDV2 schedule