models

models

Model components: velocity-field backbone wrappers, projector heads, EMA, and the DiT transformer with adaLN-Zero conditioning.

TinyVelocityField

TinyVelocityField(channels=1, hidden=32)

Bases: BaseVelocityField

A tiny convolutional velocity field for tests and toy runs.

Not a real backbone - just enough capacity to keep unit tests self-contained. For real training use WrappedBackbone around a proper UNet / DiT implementation.

Source code in deltaflow/models/backbone.py

def __init__(self, channels: int = 1, hidden: int = 32):
    super().__init__()
    self.channels = channels
    self.time_embed = nn.Sequential(
        nn.Linear(1, hidden),
        nn.SiLU(),
        nn.Linear(hidden, hidden),
    )
    self.net = nn.Sequential(
        nn.Conv2d(channels + hidden, hidden, 3, padding=1),
        nn.GroupNorm(8, hidden),
        nn.SiLU(),
        nn.Conv2d(hidden, hidden, 3, padding=1),
        nn.GroupNorm(8, hidden),
        nn.SiLU(),
        nn.Conv2d(hidden, channels, 3, padding=1),
    )

WrappedBackbone

WrappedBackbone(backbone, forward_fn=None, time_scale=1.0)

Bases: BaseVelocityField

Adapter that turns any nn.Module backbone into a velocity field.

Parameters:

Name	Type	Description	Default
backbone	Module	the underlying network (UNet, DiT, ...).	required
forward_fn	Optional[Callable]	callable (backbone, x, t, **cond) -> velocity. If None, the default is backbone(x, t, **cond).	None
time_scale	float	multiplier applied to t before it enters the backbone (useful when the backbone was trained on diffusion-style integer timesteps).	1.0

Source code in deltaflow/models/backbone.py

def __init__(
    self,
    backbone: nn.Module,
    forward_fn: Optional[Callable] = None,
    time_scale: float = 1.0,
):
    super().__init__()
    self.backbone = backbone
    self.forward_fn = forward_fn
    self.time_scale = time_scale

DiT

DiT(input_size=32, patch_size=4, in_channels=1, hidden_size=256, depth=4, num_heads=4, mlp_ratio=4.0, num_classes=10, class_dropout_prob=0.1, time_scale=1000.0)

Bases: BaseVelocityField

A class-conditional Diffusion Transformer velocity field.

The field patchifies an image x of shape (B, C, H, W) into a token sequence, adds a learned position embedding, runs it through depth DiTBlock layers conditioned on time + class embeddings, and unpatchifies the output back to a velocity of the same shape as x. It satisfies the BaseVelocityField contract, so it drops straight into DeltaFlow's losses and solvers.

Conditioning is passed through the opaque **cond channel as cond["y"] (integer class labels of shape (B,)). When y is omitted the field runs unconditionally via the learned null token.

Parameters:

Name	Type	Description	Default
input_size	int	spatial size of the (square) input image.	32
patch_size	int	side length of each square patch; must divide input_size.	4
in_channels	int	number of image channels.	1
hidden_size	int	transformer token width.	256
depth	int	number of DiTBlock layers.	4
num_heads	int	attention heads per block.	4
mlp_ratio	float	feed-forward expansion ratio.	4.0
num_classes	int	number of conditioning classes (a null token is added on top for classifier-free guidance).	10
class_dropout_prob	float	train-time label-dropout probability for CFG.	0.1
time_scale	float	multiplier applied to t inside TimestepEmbedding.	1000.0

Source code in deltaflow/models/dit.py

def __init__(
    self,
    input_size: int = 32,
    patch_size: int = 4,
    in_channels: int = 1,
    hidden_size: int = 256,
    depth: int = 4,
    num_heads: int = 4,
    mlp_ratio: float = 4.0,
    num_classes: int = 10,
    class_dropout_prob: float = 0.1,
    time_scale: float = 1000.0,
):
    super().__init__()
    if input_size % patch_size != 0:
        raise ValueError(
            f"input_size ({input_size}) must be divisible by patch_size ({patch_size})"
        )
    self.in_channels = in_channels
    self.out_channels = in_channels
    self.patch_size = patch_size
    self.input_size = input_size
    self.num_classes = num_classes
    self.num_patches_side = input_size // patch_size
    self.num_patches = self.num_patches_side**2

    self.patch_embed = nn.Conv2d(
        in_channels, hidden_size, kernel_size=patch_size, stride=patch_size
    )
    self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, hidden_size))
    self.t_embedder = TimestepEmbedding(hidden_size, time_scale=time_scale)
    self.y_embedder = LabelEmbedding(num_classes, hidden_size, class_dropout_prob)

    self.blocks = nn.ModuleList(
        [DiTBlock(hidden_size, num_heads, mlp_ratio) for _ in range(depth)]
    )
    self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)

    self.initialize_weights()

initialize_weights

initialize_weights()

Xavier-init linear layers, then zero-init every adaLN-Zero gate.

Source code in deltaflow/models/dit.py

def initialize_weights(self) -> None:
    """Xavier-init linear layers, then zero-init every adaLN-Zero gate."""

    def _basic_init(module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    self.apply(_basic_init)

    nn.init.normal_(self.pos_embed, std=0.02)

    w = self.patch_embed.weight.data
    nn.init.xavier_uniform_(w.view(w.shape[0], -1))
    nn.init.zeros_(self.patch_embed.bias)

    nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
    nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
    nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)

    # adaLN-Zero: zero the modulation heads so blocks start as the identity.
    for block in self.blocks:
        nn.init.zeros_(block.adaLN_modulation[-1].weight)
        nn.init.zeros_(block.adaLN_modulation[-1].bias)

    # Zero the final layer so the initial predicted velocity is exactly 0.
    nn.init.zeros_(self.final_layer.adaLN_modulation[-1].weight)
    nn.init.zeros_(self.final_layer.adaLN_modulation[-1].bias)
    nn.init.zeros_(self.final_layer.linear.weight)
    nn.init.zeros_(self.final_layer.linear.bias)

forward

forward(x, t, y=None, force_drop_ids=None, **cond)

Predict the velocity v_theta(x, t, y).

Parameters:

Name	Type	Description	Default
x	Tensor	input image batch (B, C, H, W).	required
t	Tensor	time of shape (B,) or a scalar (DeltaFlow convention t=0 noise, t=1 data).	required
y	Optional[Tensor]	optional integer class labels (B,) forwarded as cond["y"]; omit for unconditional inference.	None
force_drop_ids	Optional[Tensor]	optional mask selecting which labels to force to the null token (used by the unconditional branch of CFG).	None

Source code in deltaflow/models/dit.py

def forward(
    self,
    x: torch.Tensor,
    t: torch.Tensor,
    y: Optional[torch.Tensor] = None,
    force_drop_ids: Optional[torch.Tensor] = None,
    **cond,
) -> torch.Tensor:
    """Predict the velocity ``v_theta(x, t, y)``.

    Args:
        x: input image batch ``(B, C, H, W)``.
        t: time of shape ``(B,)`` or a scalar (DeltaFlow convention
            ``t=0`` noise, ``t=1`` data).
        y: optional integer class labels ``(B,)`` forwarded as ``cond["y"]``;
            omit for unconditional inference.
        force_drop_ids: optional mask selecting which labels to force to the
            null token (used by the unconditional branch of CFG).
    """
    if not isinstance(t, torch.Tensor):
        t = torch.tensor(t, device=x.device)
    if t.dim() == 0:
        t = t.expand(x.shape[0])

    h = self.patch_embed(x).flatten(2).transpose(1, 2)  # (B, N, D)
    h = h + self.pos_embed
    c = self._conditioning(x, t, y, force_drop_ids)
    for block in self.blocks:
        h = block(h, c)
    h = self.final_layer(h, c)
    return self._unpatchify(h)

forward_with_cfg

forward_with_cfg(x, t, y, cfg_scale=4.0)

Classifier-free-guided velocity for sampling.

Runs the field once with the real labels and once with the null token, then extrapolates

\[ v_{\mathrm{cfg}} = v_\text{uncond} + s\,\bigl(v_\text{cond} - v_\text{uncond}\bigr), \]

with guidance weight s = cfg_scale. The two passes are batched together, so this costs one forward over a doubled batch.

Source code in deltaflow/models/dit.py

@torch.no_grad()
def forward_with_cfg(
    self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor, cfg_scale: float = 4.0
) -> torch.Tensor:
    r"""Classifier-free-guided velocity for sampling.

    Runs the field once with the real labels and once with the null token,
    then extrapolates

    \[
    v_{\mathrm{cfg}} = v_\text{uncond}
        + s\,\bigl(v_\text{cond} - v_\text{uncond}\bigr),
    \]

    with guidance weight ``s = cfg_scale``. The two passes are batched
    together, so this costs one forward over a doubled batch.
    """
    half = x
    combined = torch.cat([half, half], dim=0)
    t_cat = torch.cat([t, t], dim=0) if t.dim() > 0 else t
    y_cat = torch.cat([y, y], dim=0)
    drop = torch.zeros(y_cat.shape[0], dtype=torch.bool, device=y.device)
    drop[y.shape[0] :] = True
    v = self.forward(combined, t_cat, y=y_cat, force_drop_ids=drop)
    v_cond, v_uncond = v.chunk(2, dim=0)
    return v_uncond + cfg_scale * (v_cond - v_uncond)

DiTBlock

DiTBlock(hidden_size, num_heads, mlp_ratio=4.0)

Bases: Module

A transformer block modulated by adaLN-Zero conditioning.

Both sub-layers (self-attention and MLP) are wrapped as

x = x + gate * sublayer(modulate(norm(x), shift, scale)),

where (shift, scale, gate) are produced per sub-layer from the shared conditioning embedding c. The six modulation vectors come from a single SiLU -> Linear head whose weights are zero-initialized (see DiT.initialize_weights), so at initialization every gate is zero and the block is the identity map.

Parameters:

Name	Type	Description	Default
hidden_size	int	token channel width.	required
num_heads	int	attention heads.	required
mlp_ratio	float	hidden expansion of the feed-forward MLP.	4.0

Source code in deltaflow/models/dit.py

def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0):
    super().__init__()
    self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
    self.attn = _Attention(hidden_size, num_heads)
    self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
    mlp_hidden = int(hidden_size * mlp_ratio)
    self.mlp = nn.Sequential(
        nn.Linear(hidden_size, mlp_hidden),
        nn.GELU(approximate="tanh"),
        nn.Linear(mlp_hidden, hidden_size),
    )
    self.adaLN_modulation = nn.Sequential(
        nn.SiLU(),
        nn.Linear(hidden_size, 6 * hidden_size),
    )

FinalLayer

FinalLayer(hidden_size, patch_size, out_channels)

Bases: Module

adaLN-Zero output head that maps tokens back to patch pixels.

The final normalization is modulated by (shift, scale) from the conditioning embedding, then a linear layer projects each token to patch_size**2 * out_channels values. Both the modulation head and the output projection are zero-initialized so the field predicts v = 0 at the start of training.

Source code in deltaflow/models/dit.py

def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
    super().__init__()
    self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
    self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels)
    self.adaLN_modulation = nn.Sequential(
        nn.SiLU(),
        nn.Linear(hidden_size, 2 * hidden_size),
    )

LabelEmbedding

LabelEmbedding(num_classes, hidden_size, dropout_prob=0.1)

Bases: Module

Class-label embedding with a learned null token for guidance.

The embedding table has num_classes + 1 rows; the extra row is the null (unconditional) token used by classifier-free guidance. During training a fraction dropout_prob of labels are replaced by the null token, teaching the field both the conditional and unconditional velocity with one set of weights.

Parameters:

Name	Type	Description	Default
num_classes	int	number of real classes.	required
hidden_size	int	embedding width.	required
dropout_prob	float	probability of dropping a label to the null token at train time (set 0 to disable CFG training).	0.1

Source code in deltaflow/models/dit.py

def __init__(self, num_classes: int, hidden_size: int, dropout_prob: float = 0.1):
    super().__init__()
    self.num_classes = num_classes
    self.dropout_prob = dropout_prob
    self.null_index = num_classes
    self.embedding_table = nn.Embedding(num_classes + 1, hidden_size)

token_drop

token_drop(labels, force_drop_ids=None)

Replace a random subset of labels with the null token.

If force_drop_ids is given (a boolean/0-1 mask), those positions are dropped deterministically instead, which is how the unconditional branch is requested at sampling time.

Source code in deltaflow/models/dit.py

def token_drop(
    self, labels: torch.Tensor, force_drop_ids: Optional[torch.Tensor] = None
) -> torch.Tensor:
    """Replace a random subset of ``labels`` with the null token.

    If ``force_drop_ids`` is given (a boolean/0-1 mask), those positions are
    dropped deterministically instead, which is how the unconditional branch
    is requested at sampling time.
    """
    if force_drop_ids is None:
        drop = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
    else:
        drop = force_drop_ids.to(torch.bool)
    return torch.where(drop, torch.full_like(labels, self.null_index), labels)

TimestepEmbedding

TimestepEmbedding(hidden_size, frequency_dim=256, time_scale=1000.0, max_period=10000)

Bases: Module

Reusable sinusoidal timestep embedding followed by a small MLP.

The continuous flow-matching time t in [0, 1] (DeltaFlow convention: t=0 is noise, t=1 is data) is first multiplied by time_scale to spread it across the sinusoidal frequency band, then embedded with the usual transformer sinusoids and projected by a two-layer MLP.

Parameters:

Name	Type	Description	Default
hidden_size	int	output (and MLP) width.	required
frequency_dim	int	width of the raw sinusoidal features before the MLP.	256
time_scale	float	multiplier applied to t before the sinusoids. The default of 1000 mirrors the diffusion-timestep range DiT was tuned on and gives continuous t in [0, 1] enough resolution.	1000.0
max_period	int	controls the lowest sinusoidal frequency.	10000

Source code in deltaflow/models/dit.py

def __init__(
    self,
    hidden_size: int,
    frequency_dim: int = 256,
    time_scale: float = 1000.0,
    max_period: int = 10000,
):
    super().__init__()
    self.frequency_dim = frequency_dim
    self.time_scale = time_scale
    self.max_period = max_period
    self.mlp = nn.Sequential(
        nn.Linear(frequency_dim, hidden_size),
        nn.SiLU(),
        nn.Linear(hidden_size, hidden_size),
    )

EMA

EMA(beta=0.995)

Exponential moving average helper for a shadow copy of a model.

Usage::

ema = EMA(beta=0.995)
ema_model = copy.deepcopy(model).eval().requires_grad_(False)
# after every optimizer step:
ema.update_model_average(ema_model, model)

Source code in deltaflow/models/ema.py

def __init__(self, beta: float = 0.995):
    self.beta = beta

MultiScaleProjector

MultiScaleProjector(feature_dims, hidden_dim=512, out_dim=256, layer_set=_DEFAULT_LAYER_SET)

Bases: Module

Collection of per-level projection heads.

Parameters:

Name	Type	Description	Default
feature_dims	Dict[str, int]	{layer_name: channel_count}, e.g. {"enc_1_4": 256, "bottleneck": 1024}. Only layers present in layer_set (default: four backbone hierarchy levels) get a projector.	required
hidden_dim	int	hidden width shared by every per-level projector.	512
out_dim	int	shared output embedding dimension.	256
layer_set	Sequence[str]	which layer names to project, defaults to ["enc_1_4", "enc_1_8", "bottleneck", "dec_1_8"].	_DEFAULT_LAYER_SET

Source code in deltaflow/models/projector.py

def __init__(
    self,
    feature_dims: Dict[str, int],
    hidden_dim: int = 512,
    out_dim: int = 256,
    layer_set: Sequence[str] = _DEFAULT_LAYER_SET,
):
    super().__init__()
    self.feature_dims = feature_dims
    self.hidden_dim = hidden_dim
    self.out_dim = out_dim
    self.layer_set = list(layer_set)

    self.projectors = nn.ModuleDict(
        {
            layer: ProjectorHead(in_dim=feature_dims[layer], hidden_dim=hidden_dim, out_dim=out_dim)
            for layer in self.layer_set
            if layer in feature_dims
        }
    )

pool_feature staticmethod

pool_feature(feature_map)

Global average pooling: [B, C, H, W] -> [B, C].

Source code in deltaflow/models/projector.py

@staticmethod
def pool_feature(feature_map: torch.Tensor) -> torch.Tensor:
    """Global average pooling: ``[B, C, H, W] -> [B, C]``."""
    return feature_map.mean(dim=(2, 3))

project_delta_h

project_delta_h(feats_cond, feats_uncond, normalize=True)

Compute delta_h = h_cond - h_uncond and project each level.

Returns a dict of (optionally L2-normalized) embeddings, one per layer present in both feature dicts and in self.projectors.

Source code in deltaflow/models/projector.py

def project_delta_h(
    self,
    feats_cond: Dict[str, torch.Tensor],
    feats_uncond: Dict[str, torch.Tensor],
    normalize: bool = True,
) -> Dict[str, torch.Tensor]:
    """Compute ``delta_h = h_cond - h_uncond`` and project each level.

    Returns a dict of (optionally L2-normalized) embeddings, one per
    layer present in both feature dicts and in ``self.projectors``.
    """
    projected = {}
    for layer in self.layer_set:
        if layer not in feats_cond or layer not in feats_uncond:
            continue
        if layer not in self.projectors:
            continue
        delta_h = feats_cond[layer] - feats_uncond[layer]
        pooled = self.pool_feature(delta_h)
        z = self.projectors[layer](pooled)
        if normalize:
            z = F.normalize(z, dim=-1)
        projected[layer] = z
    return projected

ProjectorHead

ProjectorHead(in_dim, hidden_dim=512, out_dim=256)

Bases: Module

Two-layer MLP projector (in_dim -> hidden_dim -> out_dim, GELU).

Source code in deltaflow/models/projector.py

def __init__(self, in_dim: int, hidden_dim: int = 512, out_dim: int = 256):
    super().__init__()
    self.mlp = nn.Sequential(
        nn.Linear(in_dim, hidden_dim),
        nn.GELU(),
        nn.Linear(hidden_dim, out_dim),
    )

modulate

modulate(x, shift, scale)

Apply adaLN modulation x * (1 + scale) + shift with broadcasting.

x has shape (B, N, D) (batch, tokens, channels), while shift and scale have shape (B, D). The 1 + keeps the transform centred on the identity, so a zero-initialized modulation MLP leaves x unchanged.

Source code in deltaflow/models/dit.py

def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """Apply adaLN modulation ``x * (1 + scale) + shift`` with broadcasting.

    ``x`` has shape ``(B, N, D)`` (batch, tokens, channels), while ``shift`` and
    ``scale`` have shape ``(B, D)``. The ``1 +`` keeps the transform centred on
    the identity, so a zero-initialized modulation MLP leaves ``x`` unchanged.
    """
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)