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 sc
(
self,
net_cfg: Union[dict, nn.Module],
schedule: str = "linear",
sigma_min: float = 0.0,
timestep_sampler: dict = None,
)
| 393 | |
| 394 | class FlowModel(nn.Module): |
| 395 | def __init__( |
| 396 | self, |
| 397 | net_cfg: Union[dict, nn.Module], |
| 398 | schedule: str = "linear", |
| 399 | sigma_min: float = 0.0, |
| 400 | timestep_sampler: dict = None, |
| 401 | ): |
| 402 | """ |
| 403 | Flow Matching, Stochastic Interpolants, or Rectified Flow model. :) |
| 404 | |
| 405 | Args: |
| 406 | net: Neural network that takes in x and t and outputs the vector |
| 407 | field at that point in time and space with the same shape as x. |
| 408 | schedule: str, specifies the schedule for the flow. Currently |
| 409 | supports "linear" and "gvp" (Generalized Variance Path) [3]. |
| 410 | sigma_min: a float representing the standard deviation of the |
| 411 | Gaussian distribution around the mean of the probability |
| 412 | path N(t * x1 + (1 - t) * x0, sigma), as used in [1]. |
| 413 | timestep_sampler: dict, configuration for the training timestep sampler. |
| 414 | |
| 415 | References: |
| 416 | [1] Lipman et al. (2023). Flow Matching for Generative Modeling. |
| 417 | [2] Tong et al. (2023). Improving and generalizing flow-based |
| 418 | generative models with minibatch optimal transport. |
| 419 | [3] Ma et al. (2024). SiT: Exploring flow and diffusion-based |
| 420 | generative models with scalable interpolant transformers. |
| 421 | """ |
| 422 | super().__init__() |
| 423 | if isinstance(net_cfg, nn.Module): |
| 424 | self.net = net_cfg |
| 425 | else: |
| 426 | self.net = instantiate_from_config(net_cfg) |
| 427 | self.sigma_min = sigma_min |
| 428 | |
| 429 | if schedule == "linear": |
| 430 | self.schedule = LinearSchedule() |
| 431 | elif schedule == "gvp": |
| 432 | assert sigma_min == 0.0, "GVP schedule does not support sigma_min." |
| 433 | self.schedule = GVPSchedule() |
| 434 | else: |
| 435 | raise NotImplementedError(f"Schedule {schedule} not implemented.") |
| 436 | |
| 437 | if timestep_sampler is not None: |
| 438 | self.t_sampler = instantiate_from_config(timestep_sampler) |
| 439 | else: |
| 440 | self.t_sampler = torch.rand # default: uniform U(0, 1) |
| 441 | |
| 442 | self.sde_sampler = FlowSDE(schedule=self.schedule) |
| 443 | |
| 444 | def forward(self, x: Tensor, t: Tensor, cfg_scale=1.0, uc_cond=None, cond_key="y", **kwargs): |
| 445 | if t.numel() == 1: |
nothing calls this directly
no test coverage detected