wan22 ¶

Wan2.2-TI2V-5B dense MLX runtime — Track D.

The Wan2.2 TI2V-5B (FullAttn) differs from the ported Wan2.1-T2V only in:

Scale (24 heads x 128, hidden 3072, ffn 14336) — pure config, block math identical, so the dense loader mlx_dit_from_diffusers_safetensors loads the weights unchanged and we re-wrap the blocks here.
Per-token timestep conditioning (expand_timesteps=True): the timestep is [batch, seq_len] (a level per patch token — how TI2V keeps the conditioning image frame at t=0 while the video frames are noised). timestep_proj becomes [batch, seq_len, 6, dim] and the block/output modulation is per-token ([B, L, dim]), a direct broadcast — this module implements exactly that.

I2V rides on the same forward: encode the image, replace the first latent frame, and set that frame's timestep to 0 (handled by the caller / sampler). See docs/design/ti2v_5b_port_guide.md.

Classes¶

fastvideo.mlx_runtime.wan22.MLXWan22DiT ¶

MLXWan22DiT(weights: dict[str, array], blocks: list[MLXWan22TransformerBlock], config: dict, *, compile: bool = False)

Wan2.2-TI2V-5B dense DiT with per-token timestep conditioning.

Source code in fastvideo/mlx_runtime/wan22.py

def __init__(
    self,
    weights: dict[str, mx.array],
    blocks: list[MLXWan22TransformerBlock],
    config: dict,
    *,
    compile: bool = False,
) -> None:
    import os

    self.weights = weights
    self.blocks = blocks
    self.config = config
    self.num_heads = int(config["num_attention_heads"])
    self.head_dim = int(config["attention_head_dim"])
    self.hidden_size = self.num_heads * self.head_dim
    self.freq_dim = int(config["freq_dim"])
    self.patch_size = tuple(config["patch_size"])
    self.out_channels = int(config["out_channels"])
    self.eps = float(config.get("eps", 1e-6))
    self._enable_compile = compile or os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1"
    self._compiled_forward: Callable[..., Any] | None = None
    self._compiled_signature: tuple | None = None

fastvideo.mlx_runtime.wan22.MLXWan22TransformerBlock ¶

MLXWan22TransformerBlock(weights: dict[str, array], *, dim: int, ffn_dim: int, num_heads: int, eps: float = 1e-06)

Dense Wan block with per-token ([B, L, dim]) timestep modulation.

Source code in fastvideo/mlx_runtime/wan22.py

def __init__(
    self,
    weights: dict[str, mx.array],
    *,
    dim: int,
    ffn_dim: int,
    num_heads: int,
    eps: float = 1e-6,
):
    self.weights = weights
    self.dim = dim
    self.ffn_dim = ffn_dim
    self.num_heads = num_heads
    self.head_dim = dim // num_heads
    self.eps = eps
    self.attn2 = MLXWanT2VCrossAttention(weights, dim=dim, num_heads=num_heads, eps=eps)

Functions:¶

fastvideo.mlx_runtime.wan22.mlx_wan22_dit_from_diffusers_safetensors ¶

mlx_wan22_dit_from_diffusers_safetensors(checkpoint_path: str | Path, config_path: str | Path, *, dtype: str = 'fp16', num_blocks: int | None = None, quantization=None, compile: bool = False) -> MLXWan22DiT

Load Wan2.2-TI2V-5B (FullAttn) into MLXWan22DiT via the dense loader.

Source code in fastvideo/mlx_runtime/wan22.py

def mlx_wan22_dit_from_diffusers_safetensors(
    checkpoint_path: str | Path,
    config_path: str | Path,
    *,
    dtype: str = "fp16",
    num_blocks: int | None = None,
    quantization=None,
    compile: bool = False,
) -> MLXWan22DiT:
    """Load Wan2.2-TI2V-5B (FullAttn) into ``MLXWan22DiT`` via the dense loader."""
    dense = mlx_dit_from_diffusers_safetensors(checkpoint_path,
                                               config_path,
                                               dtype=dtype,
                                               num_blocks=num_blocks,
                                               quantization=quantization)
    inner_dim = int(dense.config["num_attention_heads"]) * int(dense.config["attention_head_dim"])
    blocks = [
        MLXWan22TransformerBlock(block.weights,
                                 dim=inner_dim,
                                 ffn_dim=int(dense.config["ffn_dim"]),
                                 num_heads=int(dense.config["num_attention_heads"]),
                                 eps=float(dense.config.get("eps", 1e-6))) for block in dense.blocks
    ]
    return MLXWan22DiT(dense.weights, blocks, dense.config, compile=compile)

fastvideo.mlx_runtime.wan22.mlx_wan22_dit_from_mlx_checkpoint ¶

mlx_wan22_dit_from_mlx_checkpoint(checkpoint_dir: str | Path, *, compile: bool = False) -> MLXWan22DiT

Rewrap a persisted MLX DiT checkpoint with Wan2.2 conditioning.

The generic checkpoint loader intentionally rebuilds MLXWanDiT because it is also used by the Wan2.1 runtime. Wan2.2 TI2V has the same weight layout but needs per-token timestep modulation, so callers must rewrap the loaded weights and blocks as :class:MLXWan22DiT before sampling.

Source code in fastvideo/mlx_runtime/wan22.py

def mlx_wan22_dit_from_mlx_checkpoint(
    checkpoint_dir: str | Path,
    *,
    compile: bool = False,
) -> MLXWan22DiT:
    """Rewrap a persisted MLX DiT checkpoint with Wan2.2 conditioning.

    The generic checkpoint loader intentionally rebuilds ``MLXWanDiT`` because
    it is also used by the Wan2.1 runtime.  Wan2.2 TI2V has the same weight
    layout but needs per-token timestep modulation, so callers must rewrap the
    loaded weights and blocks as :class:`MLXWan22DiT` before sampling.
    """
    from fastvideo.mlx_runtime.checkpoint import load_mlx_dit_checkpoint

    dense = load_mlx_dit_checkpoint(checkpoint_dir)
    inner_dim = int(dense.config["num_attention_heads"]) * int(dense.config["attention_head_dim"])
    blocks = [
        MLXWan22TransformerBlock(
            block.weights,
            dim=inner_dim,
            ffn_dim=int(dense.config["ffn_dim"]),
            num_heads=int(dense.config["num_attention_heads"]),
            eps=float(dense.config.get("eps", 1e-6)),
        ) for block in dense.blocks
    ]
    return MLXWan22DiT(dense.weights, blocks, dense.config, compile=compile)