class

MeanFlowForImageGeneration

extendsImageGenerationModel

MeanFlowForImageGeneration(config: MeanFlowConfig)

source edit

MeanFlow's objective and its one-step sampler.

Parameters

configMeanFlowConfig

The variant to build.

Attributes

mean_flowMeanFlowModel

The average-velocity network.

Notes

Reference: Geng, Deng, Bai, Kolter, and He, "Mean Flows for One-step Generative Modeling", arXiv:2505.13447, 2025. Training is Algorithm 1, sampling is Algorithm 2 and Eq. 12.

Examples

>>> import lucid
>>> from lucid.models.generative.mean_flow import (
...     MeanFlowConfig, MeanFlowForImageGeneration)
>>> config = MeanFlowConfig(sample_size=8, patch_size=2, hidden_size=32,
...                         depth=2, num_heads=4, num_classes=10)
>>> model = MeanFlowForImageGeneration(config).eval()
>>> model.generate(2).samples.shape
(2, 4, 8, 8)

Used by 2

Constructors

dunder

init

→None

__init__(config: MeanFlowConfig)

source edit

Initialise the model. See the class docstring for parameters.

Instance methods

forward

→MeanFlowOutput

forward(images: Tensor, labels: Tensor | None = None)

source edit

One training step of the MeanFlow objective.

Parameters

imagesTensor

(B, C, H, W) — data, or the VAE latent of it.

labelsTensor or None= None

(B,) class indices. Dropped to the unconditional token with probability MeanFlowConfig.class_dropout, which is what leaves the model able to produce both fields.

Returns

MeanFlowOutput

The weighted loss, the prediction and the target.

Notes

The Jacobian-vector product is taken along $(v, 0, 1)$ — the tangent that the chain rule gives for $\frac{d}{dt}u(z_t, r, t)$ once $dz_t/dt = v$ and $dr/dt = 0$ are substituted. Any other tangent trains a different quantity: the paper's destructive ablation reports FID 61.06 for this one and 137–329 for the alternatives.

generate

→GenerationOutput

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

source edit

Sample by walking the average velocity backwards from noise.

Parameters

n_samplesint= 1

How many to draw.

labelsTensor or None= None

(n_samples,) class indices. None samples the unconditional field.

stepsint= 1

Network evaluations. One is the point of the method:

z_0 = z_1 - u(z_1, 0, 1)

covers the whole path. More subdivide

[0, 1]

and apply Eq. 12 on each piece, which the paper notes is straightforward and reports at two.

noiseTensor or None= None

Starting point

z_1

. Drawn when absent.

devicestr or None= None

Where to draw. Defaults to the model's own device.

Returns

GenerationOutput

samples of shape (n_samples, out_channels, H, W).

Raises

ValueError

If steps is not positive.

Examples

>>> import lucid
>>> from lucid.models.generative.mean_flow import (
...     MeanFlowConfig, MeanFlowForImageGeneration)
>>> config = MeanFlowConfig(sample_size=8, patch_size=2,
...                         hidden_size=32, depth=2, num_heads=4,
...                         num_classes=10)
>>> model = MeanFlowForImageGeneration(config).eval()
>>> model.generate(2).samples.shape
(2, 4, 8, 8)

>>> import lucid >>> from lucid.models.generative.mean_flow import ( ... MeanFlowConfig, MeanFlowForImageGeneration) >>> config = MeanFlowConfig(sample_size=8, patch_size=2, hidden_size=32, ... depth=2, num_heads=4, num_classes=10) >>> model = MeanFlowForImageGeneration(config).eval() >>> model.generate(2).samples.shape (2, 4, 8, 8)

forward

→MeanFlowOutput

forward(images: Tensor, labels: Tensor | None = None)

source edit

One training step of the MeanFlow objective.

Parameters

imagesTensor

(B, C, H, W) — data, or the VAE latent of it.

labelsTensor or None= None

(B,) class indices. Dropped to the unconditional token with probability MeanFlowConfig.class_dropout, which is what leaves the model able to produce both fields.

Returns

MeanFlowOutput

The weighted loss, the prediction and the target.

Notes

>>> import lucid >>> from lucid.models.generative.mean_flow import ( ... MeanFlowConfig, MeanFlowForImageGeneration) >>> config = MeanFlowConfig(sample_size=8, patch_size=2, ... hidden_size=32, depth=2, num_heads=4, ... num_classes=10) >>> model = MeanFlowForImageGeneration(config).eval() >>> model.generate(2).samples.shape (2, 4, 8, 8)