class
RectifiedFlowForImageGeneration
extends
ImageGenerationModelRectifiedFlowForImageGeneration(config: RectifiedFlowConfig)Rectified Flow with the training loss, reflow and .generate().
forward(x) returns a DiffusionModelOutput whose loss
is eq. (1). forward(x, noise=...) is the reflow objective over a
previous flow's couplings, which reflow_pairs produces.
Parameters
configRectifiedFlowConfigArchitecture, schedule and solver settings.
Attributes
rectified_flowRectifiedFlowModelUnderlying velocity field, exposing straightness and likelihood.
Notes
Reference: Liu, Gong, and Liu, "Flow Straight and Fast", ICLR, 2023 (arXiv:2209.03003). Reported: a one-step FID of 4.85 on CIFAR-10, state of the art among one-step diffusion and flow models at publication.
A full run is three stages, all through this one class:
- train with
forward(x)— the 1-rectified flow; reflow_pairs(), then train withforward(z1, noise=z0);- optionally rebuild with
t_schedule="t0"and train on the same pairs — one-step distillation.
Examples
>>> import lucid
>>> from lucid.models.generative.rectified_flow import (
... RectifiedFlowConfig, RectifiedFlowForImageGeneration,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
... channel_mult=(1, 2), num_res_blocks=1,
... attention_resolutions=())
>>> model = RectifiedFlowForImageGeneration(cfg).eval()
>>> model(lucid.randn((2, 3, 8, 8))).loss.shape
()
>>> model.generate(n_samples=2, steps=1).samples.shape
(2, 3, 8, 8)Used by 2
Constructors
1Properties
1Instance methods
3generate(n_samples: int = 1, steps: int | None = None, device: str | None = None, noise: Tensor | None = None)Sample by integrating the field from t = 0 to t = 1.
Parameters
n_samplesint= 1Number of samples to draw.
stepsint= NoneFixed Euler budget instead of an adaptive solve.
steps=1
is the regime the method exists for.devicestr= NoneWhere to allocate the prior draw. Defaults to the device the
model's parameters are on; an explicit value wins.
Starting point in place of a fresh draw.
Returns
GenerationOutputsamples of shape (n_samples, C, H, W).
Examples
>>> import lucid
>>> from lucid.models.generative.rectified_flow import (
... RectifiedFlowConfig, RectifiedFlowForImageGeneration,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
... channel_mult=(1, 2), num_res_blocks=1,
... attention_resolutions=())
>>> model = RectifiedFlowForImageGeneration(cfg).eval()
>>> model.generate(n_samples=3, steps=1).samples.shape
(3, 3, 8, 8)
>>> model.nfe # one step, one field evaluation: the regime it is for
1
>>> noise = lucid.randn((1, 3, 8, 8))
>>> model.generate(noise=noise, steps=2).samples.shape # batch follows noise
(1, 3, 8, 8)reflow_pairs(n_samples: int = 1, steps: int | None = None, device: str | None = None, noise: Tensor | None = None)Couplings for the next round, from the wrapped flow.
See RectifiedFlowModel.reflow_pairs for why these pairs are
what the next round trains on.
Parameters
n_samplesint= 1How many pairs to generate. Ignored when
noise is given.stepsint= NoneFixed Euler budget instead of an adaptive solve.
devicestr= NoneWhere to allocate the prior draw. Defaults to the device the
model's parameters are on; an explicit value wins.
Sources to use in place of a fresh draw.