class

RectifiedFlowModel

extendsPretrainedModel

RectifiedFlowModel(config: RectifiedFlowConfig)

source edit

Velocity field trained to transport along straight paths.

forward(sample, t) evaluates $v(x, t; \theta)$ . Everything that makes it the paper's method is separate, and each piece maps onto one part of it:

path_sample — the straight line $t x_1 + (1-t) x_0$ .
conditional_target — its velocity, $x_1 - x_0$ .
rectified_flow_loss — the objective, eq. (1). With a noise argument it is the reflow objective; with a pinned t_schedule it is distillation.
sample — integrate $t: \varepsilon \to 1$ .
reflow_pairs — generate the (z0, z1) couplings the next round trains on.
straightness — the measure $S(\mathbf{Z})$ that reflow is supposed to reduce.
one_step — $z_0 + v(z_0, 0)$ , exact once the flow is straight.
log_prob / bits_per_dim — integrate back, tracking the divergence.

Parameters

configRectifiedFlowConfig

Architecture, schedule and solver settings.

Attributes

fieldnn.Module

The velocity field itself.

Notes

Reference: Liu, Gong, and Liu, "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow", ICLR, 2023 (arXiv:2209.03003). Eq. (1) is the objective, eq. (3) the straightness measure, and Theorem 3.3 the marginal preservation that makes reflow legitimate.

On the coupling. rectified_flow_loss(x1) draws x0 independently, which is the 1-rectified flow of the paper. Passing noise= supplies a paired x0 instead and is what turns the same call into reflow — there is no separate objective.

On transfer. Nothing here requires the source to be Gaussian. The paper's other half — transporting between two image domains rather than from noise — is the same objective with x0 drawn from the first domain instead of a prior, so it goes through rectified_flow_loss, sample and reflow_pairs unchanged by passing noise=. Only log_prob assumes a standard normal at t = 0, and a likelihood is not defined for a domain-to-domain transport anyway.

Where to run it. Training is a plain convolutional forward and backward with no solve, so it belongs on "metal". Generating reflow pairs is a solve per batch and is the expensive part of the method; the paper generates on the order of a million.

Not implemented, and why. Two settings the reference exposes are absent rather than overlooked:

A perceptual reflow loss. The reference offers lpips and lpips+l2 alongside l2 for the distillation stage. Both score images through a pre-trained external network, which Lucid's compute paths may not depend on, so only the squared error is available here. It is the reference's own default.
The stochastic sampler. The reference can add noise of magnitude (1 - t) * sigma_variance to turn the probability-flow ODE into an SDE. Its default is 0, and the paper is stated and evaluated as an ODE method, so sampling here is deterministic given the starting point.

Weight averaging (the reference's ema_rate) is a property of a training loop rather than of a model, and is not part of any family in this zoo.

Examples

>>> import lucid
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> loss, _, _ = model.rectified_flow_loss(lucid.randn((2, 3, 8, 8)))
>>> loss.shape
()
>>> model.sample(n_samples=2, steps=4).shape
(2, 3, 8, 8)

Used by 2

Constructors

dunder

init

→None

__init__(config: RectifiedFlowConfig)

source edit

Properties

prop

input_dim

→int

input_dim: int

source edit

int: Flattened width D of one sample.

prop

nfe

→int

nfe: int

source edit

int: Field evaluations spent by the most recent solve.

A training step solves nothing and leaves the count where the last solve put it. The number the paper is about: a straight flow needs one.

prop

t_schedule

→TimeSchedule | int

t_schedule: TimeSchedule | int

source edit

str or int: Which times the objective draws.

prop

trace_method

→str

trace_method: str

source edit

str: How the divergence is obtained when scoring likelihood.

Instance methods

bits_per_dim

→Tensor

bits_per_dim(x: Tensor)

source edit

Per-sample negative log-likelihood in bits per dimension.

conditional_target

→Tensor

conditional_target(x1: Tensor, x0: Tensor, t: Tensor)

source edit

The velocity of that line, $x_1 - x_0$ .

Independent of t — a straight line traversed at constant speed has one velocity — which is why t is accepted but unused. It stays in the signature because the quantity is a function of time in general, and a caller comparing paths should not have to special-case this one.

Parameters

x1Tensor

(B, C, H, W) endpoint samples.

x0Tensor

(B, C, H, W) endpoint samples.

tTensor

(B,) times, accepted for interface symmetry.

Returns

Tensor

(B, C, H, W) regression target.

Examples

>>> import lucid
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> x1, x0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8))
>>> target = model.conditional_target(x1, x0, lucid.tensor([0.2, 0.7]))
>>> target.shape
(2, 3, 8, 8)
>>> chord = (model.path_sample(x1, x0, lucid.ones(2))
...          - model.path_sample(x1, x0, lucid.zeros(2)))
>>> bool(lucid.allclose(target, chord))  # one velocity for the whole line
True

forward

→DiffusionModelOutput

forward(sample: Tensor, t: Tensor)

source edit

Evaluate the velocity field at (x, t).

log_prob

→Tensor

log_prob(x: Tensor)

source edit

Per-sample $\log p_1(x)$ in nats.

Integrates the field from t = 1 back to t = time_eps alongside the accumulated divergence, then scores the arrival against the standard normal the path starts from.

Parameters

xTensor

(B, C, H, W) data samples.

Returns

Tensor

(B,) log-likelihood.

Notes

With trace_method resolving to "hutchinson" — which it does at any image size — this is an unbiased estimate, and two calls on the same input will not agree.

The flow is only a normalizing flow while t_schedule is "uniform". A distilled model is trained to be stepped a fixed number of times, not integrated, so its density is not the one it generates from; this method does not refuse it, but the number means less than it appears to.

Examples

>>> import math
>>> import lucid
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> x = lucid.randn((2, 3, 8, 8))
>>> log_p = model.log_prob(x)
>>> log_p.shape
(2,)
>>> # A fresh field is effectively zero, so x is scored as N(0, I) itself.
>>> const = 0.5 * model.input_dim * math.log(2 * math.pi)
>>> gauss = -0.5 * (x.reshape(2, -1) ** 2).sum(dim=-1) - const
>>> bool(lucid.allclose(log_p, gauss, atol=1e-2))
True

one_step

→Tensor

one_step(noise: Tensor)

source edit

The one-step map $z_0 + v(z_0, \varepsilon)$ .

A single explicit Euler step across the whole interval. Exact if the flow is straight, which is what reflow and then t0 distillation are for; on a 1-rectified flow it is a crude approximation and is expected to look it.

Parameters

noiseTensor

(N, C, H, W) source samples.

Returns

Tensor

(N, C, H, W) one-step generations.

Examples

>>> import lucid
>>> import lucid.nn as nn
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> # The output layer starts near zero; give it weights so the check bites.
>>> _ = nn.init.normal_(model.field.conv_out.weight, std=0.05)
>>> noise = lucid.randn((2, 3, 8, 8))
>>> out = model.one_step(noise)
>>> out.shape, model.nfe
((2, 3, 8, 8), 1)
>>> bool(lucid.allclose(out, model.sample(noise=noise, steps=1)))
True

path_sample

→Tensor

path_sample(x1: Tensor, x0: Tensor, t: Tensor)

source edit

The point $x_t = t x_1 + (1 - t) x_0$ on the straight line.

Parameters

x1Tensor

(B, C, H, W) data samples.

x0Tensor

(B, C, H, W) samples from the source distribution.

tTensor

(B,) times in [0, 1].

Returns

Tensor

(B, C, H, W) interpolation. This is Flow Matching's optimal-transport path at $\sigma_{\min} = 0$ .

Examples

>>> import lucid
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> x1, x0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8))
>>> model.path_sample(x1, x0, lucid.tensor([0.25, 0.75])).shape
(2, 3, 8, 8)
>>> bool(lucid.allclose(model.path_sample(x1, x0, lucid.zeros(2)), x0))
True
>>> bool(lucid.allclose(model.path_sample(x1, x0, lucid.ones(2)), x1))
True
>>> half = model.path_sample(x1, x0, lucid.full((2,), 0.5))
>>> bool(lucid.allclose(half, (x0 + x1) / 2))  # the chord's midpoint
True

rectified_flow_loss

→tuple[Tensor, Tensor, Tensor]

rectified_flow_loss(x1: Tensor, noise: Tensor | None = None)

source edit

One optimisation step of eq. (1) — no ODE is solved.

Parameters

x1Tensor

(B, C, H, W) targets. Data for the first flow; the previous flow's samples during reflow.

noiseTensor= None

(B, C, H, W) sources. Omitted, they are drawn independently and this is the 1-rectified flow. Supplied as the noise that produced x1, the very same call becomes the reflow objective — the pairing is the whole difference.

Returns

tuple[Tensor, Tensor, Tensor]

(loss, prediction, target).

Raises

ValueError

If the model is configured to distil but no pairing is given. Distillation regresses at a fixed time, which is only meaningful against the coupling a previous flow produced; independent draws there would train the field to predict the mean of the data, and would do so silently.

Examples

>>> import dataclasses
>>> import lucid
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> x1, z0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8))
>>> loss, prediction, target = model.rectified_flow_loss(x1, noise=z0)
>>> loss.shape, prediction.shape
((), (2, 3, 8, 8))
>>> bool(lucid.allclose(target, x1 - z0))  # the pairing sets the target
True
>>> bool(lucid.allclose(loss, ((prediction - target) ** 2).mean()))
True
>>> distil = RectifiedFlowModel(dataclasses.replace(cfg, t_schedule="t0"))
>>> distil.rectified_flow_loss(x1)  # a pinned time needs paired noise
Traceback (most recent call last):
    ...
ValueError: t_schedule='t0' pins the time, ...

reflow_pairs

→tuple[Tensor, Tensor]

reflow_pairs(n_samples: int = 1, steps: int | None = None, device: str | None = None, noise: Tensor | None = None)

source edit

Couplings $(Z_0, Z_1)$ for the next rectified flow.

Draws from the source, solves the current flow forward, and returns both ends. Theorem 3.3 is what licenses feeding these back in: the solve preserves the marginals, so z1 is still distributed as the data while the pair is no longer independent — and it is that dependence the next round exploits.

Parameters

n_samplesint= 1

How many pairs to generate. Ignored when noise is given.

stepsint= None

Fixed Euler budget instead of an adaptive solve.

devicestr= None

Where to allocate the prior draw. Defaults to the device the model's parameters are on; an explicit value wins.

noiseTensor= None

Sources to use in place of a fresh draw.

Returns

tuple[Tensor, Tensor]

(z0, z1), both (N, C, H, W) — pass them straight to rectified_flow_loss(z1, noise=z0).

Notes

This is the expensive half of the method. The paper generates on the order of a million pairs before each reflow round, and it is the only place a solve appears in training at all.

Examples

>>> import lucid
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> noise = lucid.randn((2, 3, 8, 8))
>>> z0, z1 = model.reflow_pairs(noise=noise, steps=2)
>>> z0 is noise, z1.shape
(True, (2, 3, 8, 8))
>>> bool(lucid.allclose(z1, model.sample(noise=z0, steps=2)))  # z0's endpoint
True
>>> _, _, target = model.rectified_flow_loss(z1, noise=z0)
>>> bool(lucid.allclose(target, z1 - z0))  # the next round's target
True

sample

→Tensor

sample(n_samples: int = 1, steps: int | None = None, device: str | None = None, noise: Tensor | None = None)

source edit

Generate by integrating the field from noise to data.

Parameters

n_samplesint= 1

How many samples to draw. Ignored when noise is given.

stepsint= None

Force a fixed-budget Euler solve with this many uniform steps. This is the setting the method exists for: a straight flow is integrated exactly by Euler, so quality at steps=1 measures straightness directly. Left None, an adaptive solve runs instead and spends whatever the tolerance demands.

devicestr= None

Where to allocate the prior draw. Defaults to the device the model's parameters are on; an explicit value wins.

noiseTensor= None

(N, C, H, W) starting point in place of a fresh draw. Keeping it is what makes a reflow pair.

Returns

Tensor

(N, C, H, W) samples at t = 1.

Notes

Euler is used for the fixed-budget path rather than a higher-order method on purpose: the paper's claim is about the error a single first-order step makes, and a Runge–Kutta step would hide it by spending four evaluations per step.

Examples

>>> import lucid
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> model.sample(n_samples=2, steps=1).shape
(2, 3, 8, 8)
>>> model.nfe  # steps=1: a single field evaluation
1
>>> noise = lucid.randn((3, 3, 8, 8))
>>> model.sample(noise=noise, steps=4).shape  # the batch follows noise
(3, 3, 8, 8)
>>> model.nfe  # Euler, so the budget is exactly the step count
4

sample_times

→Tensor

sample_times(batch: int, device: str | None = None)

source edit

Draw the times the objective is evaluated at.

Four behaviours, one per t_schedule:

"uniform" — $t \sim U[\varepsilon, 1]$ , the objective itself and what reflow retrains under.
"t0" — every t at $\varepsilon$ . Regressing there makes one Euler step a generator: one-step distillation.
"t1" — every t at 1, distilling the reverse map.
k — the k-point grid a k-step Euler sampler visits, for k-step distillation.

Parameters

batchint

How many times to draw.

devicestr= None

Where to allocate them. Defaults to the device the model's parameters are on; an explicit value wins.

Returns

Tensor

(batch,) times in [time_eps, 1].

Examples

>>> import dataclasses
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> t = model.sample_times(6)
>>> t.shape
(6,)
>>> bool(((t >= cfg.time_eps) & (t <= 1.0)).all().item())
True
>>> # k = 2: only the two times a two-step Euler sampler visits.
>>> two_step = RectifiedFlowModel(dataclasses.replace(cfg, t_schedule=2))
>>> sorted({round(v, 4) for v in two_step.sample_times(64).tolist()})
[0.001, 0.5005]

straightness

→Tensor

straightness(noise: Tensor | None = None, n_samples: int = 1, steps: int = 32, device: str | None = None)

source edit

The measure $S(\mathbf{Z})$ of paper eq. (3).

S(\mathbf{Z}) = \int_0^1 \mathbb{E}\Bigl[ \bigl\| (Z_1 - Z_0) - \dot{Z}_t \bigr\|^2 \Bigr]\,\mathrm{d}t

Estimated by stepping the flow on a uniform grid, recording the field along the way, and comparing each velocity against the chord the trajectory actually spans. Zero exactly when every path is a straight line at constant speed — and therefore exactly when one Euler step is enough.

Parameters

noiseTensor= None

(N, C, H, W) starting points. Fixing them is the only way to compare two models fairly, since the measure is an expectation over the source.

n_samplesint= 1

How many paths to average over when noise is not given.

stepsint= 32

Grid resolution for both the solve and the quadrature.

devicestr= None

Where to draw the paths when noise is not given. Defaults to the device the model's parameters are on; an explicit value wins.

Returns

Tensor

Scalar estimate, in the sample's own squared units.

Notes

Reported per-dimension (a mean, not a sum), so the number is comparable across resolutions.

Examples

>>> import lucid
>>> import lucid.nn as nn
>>> from lucid.models.generative.rectified_flow import (
...     RectifiedFlowConfig, RectifiedFlowModel,
... )
>>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16,
...                           channel_mult=(1, 2), num_res_blocks=1,
...                           attention_resolutions=())
>>> model = RectifiedFlowModel(cfg).eval()
>>> noise = lucid.randn((2, 3, 8, 8))
>>> s = model.straightness(noise, steps=4)
>>> s.shape, model.nfe
((), 4)
>>> bool(s.item() < 1e-8)  # a fresh field barely moves, so nothing bends
True
>>> _ = nn.init.normal_(model.field.conv_out.weight, std=0.05)
>>> bool(model.straightness(noise, steps=4).item() > 0.0)
True

>>> import lucid >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> loss, _, _ = model.rectified_flow_loss(lucid.randn((2, 3, 8, 8))) >>> loss.shape () >>> model.sample(n_samples=2, steps=4).shape (2, 3, 8, 8)

>>> import lucid >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> x1, x0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8)) >>> target = model.conditional_target(x1, x0, lucid.tensor([0.2, 0.7])) >>> target.shape (2, 3, 8, 8) >>> chord = (model.path_sample(x1, x0, lucid.ones(2)) ... - model.path_sample(x1, x0, lucid.zeros(2))) >>> bool(lucid.allclose(target, chord)) # one velocity for the whole line True

>>> import math >>> import lucid >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> x = lucid.randn((2, 3, 8, 8)) >>> log_p = model.log_prob(x) >>> log_p.shape (2,) >>> # A fresh field is effectively zero, so x is scored as N(0, I) itself. >>> const = 0.5 * model.input_dim * math.log(2 * math.pi) >>> gauss = -0.5 * (x.reshape(2, -1) ** 2).sum(dim=-1) - const >>> bool(lucid.allclose(log_p, gauss, atol=1e-2)) True

>>> import lucid >>> import lucid.nn as nn >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> # The output layer starts near zero; give it weights so the check bites. >>> _ = nn.init.normal_(model.field.conv_out.weight, std=0.05) >>> noise = lucid.randn((2, 3, 8, 8)) >>> out = model.one_step(noise) >>> out.shape, model.nfe ((2, 3, 8, 8), 1) >>> bool(lucid.allclose(out, model.sample(noise=noise, steps=1))) True

>>> import lucid >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> x1, x0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8)) >>> model.path_sample(x1, x0, lucid.tensor([0.25, 0.75])).shape (2, 3, 8, 8) >>> bool(lucid.allclose(model.path_sample(x1, x0, lucid.zeros(2)), x0)) True >>> bool(lucid.allclose(model.path_sample(x1, x0, lucid.ones(2)), x1)) True >>> half = model.path_sample(x1, x0, lucid.full((2,), 0.5)) >>> bool(lucid.allclose(half, (x0 + x1) / 2)) # the chord's midpoint True

>>> import dataclasses >>> import lucid >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> x1, z0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8)) >>> loss, prediction, target = model.rectified_flow_loss(x1, noise=z0) >>> loss.shape, prediction.shape ((), (2, 3, 8, 8)) >>> bool(lucid.allclose(target, x1 - z0)) # the pairing sets the target True >>> bool(lucid.allclose(loss, ((prediction - target) ** 2).mean())) True >>> distil = RectifiedFlowModel(dataclasses.replace(cfg, t_schedule="t0")) >>> distil.rectified_flow_loss(x1) # a pinned time needs paired noise Traceback (most recent call last): ... ValueError: t_schedule='t0' pins the time, ...

>>> import lucid >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> noise = lucid.randn((2, 3, 8, 8)) >>> z0, z1 = model.reflow_pairs(noise=noise, steps=2) >>> z0 is noise, z1.shape (True, (2, 3, 8, 8)) >>> bool(lucid.allclose(z1, model.sample(noise=z0, steps=2))) # z0's endpoint True >>> _, _, target = model.rectified_flow_loss(z1, noise=z0) >>> bool(lucid.allclose(target, z1 - z0)) # the next round's target True

>>> import lucid >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> model.sample(n_samples=2, steps=1).shape (2, 3, 8, 8) >>> model.nfe # steps=1: a single field evaluation 1 >>> noise = lucid.randn((3, 3, 8, 8)) >>> model.sample(noise=noise, steps=4).shape # the batch follows noise (3, 3, 8, 8) >>> model.nfe # Euler, so the budget is exactly the step count 4

>>> import dataclasses >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> t = model.sample_times(6) >>> t.shape (6,) >>> bool(((t >= cfg.time_eps) & (t <= 1.0)).all().item()) True >>> # k = 2: only the two times a two-step Euler sampler visits. >>> two_step = RectifiedFlowModel(dataclasses.replace(cfg, t_schedule=2)) >>> sorted({round(v, 4) for v in two_step.sample_times(64).tolist()}) [0.001, 0.5005]

>>> import lucid >>> import lucid.nn as nn >>> from lucid.models.generative.rectified_flow import ( ... RectifiedFlowConfig, RectifiedFlowModel, ... ) >>> cfg = RectifiedFlowConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=()) >>> model = RectifiedFlowModel(cfg).eval() >>> noise = lucid.randn((2, 3, 8, 8)) >>> s = model.straightness(noise, steps=4) >>> s.shape, model.nfe ((), 4) >>> bool(s.item() < 1e-8) # a fresh field barely moves, so nothing bends True >>> _ = nn.init.normal_(model.field.conv_out.weight, std=0.05) >>> bool(model.straightness(noise, steps=4).item() > 0.0) True