class

MeanFlowModel

extendsPretrainedModel

MeanFlowModel(config: MeanFlowConfig)

source edit

The average-velocity field $u_\theta(z, r, t)$ .

Parameters

configMeanFlowConfig

The variant to build.

Attributes

patch_embedlucid.nn.Conv2d

Patchifying projection — stride equals kernel, so patches do not overlap.

pos_embedTensor

Frozen two-dimensional sine-cosine table, added to the tokens.

blockslucid.nn.ModuleList

The DiT blocks.

final_FinalLayer

Modulated norm and projection back to patch space.

Notes

Reference: Geng, Deng, Bai, Kolter, and He, "Mean Flows for One-step Generative Modeling", arXiv:2505.13447, 2025. Backbone configurations are Table 4.

The forward takes both times. Which of them reach the embedding is MeanFlowConfig.time_conditioning; the signature does not change with it, because the Jacobian-vector product the objective needs is always taken with respect to $(z, r, t)$ no matter how the network chooses to encode them.

Examples

>>> import lucid
>>> from lucid.models.generative.mean_flow import (
...     MeanFlowConfig, MeanFlowModel)
>>> config = MeanFlowConfig(sample_size=8, patch_size=2, hidden_size=32,
...                         depth=2, num_heads=4, num_classes=10)
>>> model = MeanFlowModel(config).eval()
>>> z = lucid.randn((1, 4, 8, 8))
>>> t = lucid.tensor([1.0])
>>> r = lucid.tensor([0.0])
>>> model(z, r, t).shape
(1, 4, 8, 8)

Used by 2

Constructors

dunder

init

→None

__init__(config: MeanFlowConfig)

source edit

Initialise the network. See the class docstring for parameters.

Instance methods

forward

→Tensor

forward(z: Tensor, r: Tensor, t: Tensor, labels: Tensor | None = None)

source edit

Evaluate the average-velocity field.

Parameters

zTensor

(B, in_channels, H, W) — the point on the path.

rTensor

(B,) — the interval's start.

tTensor

(B,) — the interval's end, and where z lives.

labelsTensor or None= None

(B,) class indices. None uses the unconditional token, which is index num_classes.

Returns

Tensor

(B, out_channels, H, W) — $u_\theta(z, r, t)$ .

>>> import lucid >>> from lucid.models.generative.mean_flow import ( ... MeanFlowConfig, MeanFlowModel) >>> config = MeanFlowConfig(sample_size=8, patch_size=2, hidden_size=32, ... depth=2, num_heads=4, num_classes=10) >>> model = MeanFlowModel(config).eval() >>> z = lucid.randn((1, 4, 8, 8)) >>> t = lucid.tensor([1.0]) >>> r = lucid.tensor([0.0]) >>> model(z, r, t).shape (1, 4, 8, 8)