models

models

Model components: velocity-field backbone wrappers, projector heads, EMA.

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

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),
    )