Class FlowModel

diff2flow/flow.py:394–632

class FlowModel(nn.Module):
    def __init__(
            self,
            net_cfg: Union[dict, nn.Module],
            schedule: str = "linear",
            sigma_min: float = 0.0,
            timestep_sampler: dict = None,
        ):
        """
        Flow Matching, Stochastic Interpolants, or Rectified Flow model. :)
        
        Args:
            net: Neural network that takes in x and t and outputs the vector
                field at that point in time and space with the same shape as x.
            schedule: str, specifies the schedule for the flow. Currently
                supports "linear" and "gvp" (Generalized Variance Path) [3].
            sigma_min: a float representing the standard deviation of the
                Gaussian distribution around the mean of the probability
                path N(t * x1 + (1 - t) * x0, sigma), as used in [1].
            timestep_sampler: dict, configuration for the training timestep sampler.
        
        References:
            [1] Lipman et al. (2023). Flow Matching for Generative Modeling.
            [2] Tong et al. (2023). Improving and generalizing flow-based
                generative models with minibatch optimal transport.
            [3] Ma et al. (2024). SiT: Exploring flow and diffusion-based
                generative models with scalable interpolant transformers.
        """
        super().__init__()
        if isinstance(net_cfg, nn.Module):
            self.net = net_cfg
        else:
            self.net = instantiate_from_config(net_cfg)
        self.sigma_min = sigma_min

        if schedule == "linear":
            self.schedule = LinearSchedule()
        elif schedule == "gvp":
            assert sigma_min == 0.0, "GVP schedule does not support sigma_min."
            self.schedule = GVPSchedule()
        else:
            raise NotImplementedError(f"Schedule {schedule} not implemented.")
        
        if timestep_sampler is not None:
            self.t_sampler = instantiate_from_config(timestep_sampler)
        else:
            self.t_sampler = torch.rand             # default: uniform U(0, 1)

        self.sde_sampler = FlowSDE(schedule=self.schedule)

    def forward(self, x: Tensor, t: Tensor, cfg_scale=1.0, uc_cond=None, cond_key="y", **kwargs):
        if t.numel() == 1:
            t = t.expand(x.size(0))
        # _pred = self.net(x=x, t=t, **kwargs)
        _pred = forward_with_cfg(x, t, self.net, cfg_scale=cfg_scale, uc_cond=uc_cond, cond_key=cond_key, **kwargs)
        return _pred

    def ode_fn(self, t, x, **kwargs):