fastwan.py

# SPDX-License-Identifier: Apache-2.0
# mypy: disable-error-code=no-untyped-call
"""FastWan-oriented helpers for the experimental MLX runtime path."""

from __future__ import annotations

import json
import math
import os
import statistics
import time
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
from collections.abc import Callable

from fastvideo.logger import init_logger

if TYPE_CHECKING:
    import mlx.core as mx
    import torch

logger = init_logger(__name__)


@dataclass(frozen=True)
class FastWanShape:
    height: int
    width: int
    num_frames: int
    latent_frames: int
    latent_height: int
    latent_width: int
    patch_frames: int
    patch_height: int
    patch_width: int
    tokens: int
    hidden_size: int
    num_heads: int
    head_dim: int


class UnsupportedMLXQuantizationError(ValueError):
    """A quantization mode the installed MLX build cannot execute.

    Raised by :func:`ensure_quantization_supported` before any model weights
    are loaded, so callers (CLI flags, benchmark sweeps) can fail fast with an
    actionable message -- or skip the mode -- instead of crashing deep inside
    ``mx.quantize`` mid-load.
    """


@dataclass(frozen=True)
class MLXQuantizationSpec:
    """MLX quantized-matmul configuration for DiT linear weights."""

    mode: str
    bits: int | None = None
    group_size: int | None = None

    @classmethod
    def from_name(cls, name: str | None) -> MLXQuantizationSpec | None:
        if name is None or name in {"", "none", "fp16", "fp32"}:
            return None
        if name == "int8":
            return cls(mode="affine", bits=8, group_size=64)
        if name == "int6":
            return cls(mode="affine", bits=6, group_size=64)
        if name == "int4":
            return cls(mode="affine", bits=4, group_size=64)
        if name == "mxfp8":
            return cls(mode="mxfp8")
        if name == "mxfp4":
            return cls(mode="mxfp4")
        if name == "nvfp4":
            return cls(mode="nvfp4")
        raise ValueError(f"Unsupported MLX quantization mode: {name}")

    @property
    def label(self) -> str:
        if self.mode == "affine":
            return f"int{self.bits}"
        return self.mode


@dataclass(frozen=True)
class QuantizedMatrix:
    weight: mx.array
    scales: mx.array
    biases: mx.array | None
    spec: MLXQuantizationSpec
    dequantized_dtype: mx.Dtype


def fastwan_shape(
        *,
        height: int,
        width: int,
        num_frames: int,
        vae_temporal_compression: int = 4,
        vae_spatial_compression: int = 8,
        patch_size: tuple[int, int, int] = (1, 2, 2),
        num_heads: int = 12,
        head_dim: int = 128,
) -> FastWanShape:
    """Return the approximate DiT token shape for Wan/FastWan T2V inference."""
    latent_frames = (num_frames - 1) // vae_temporal_compression + 1
    latent_height = height // vae_spatial_compression
    latent_width = width // vae_spatial_compression
    patch_frames = latent_frames // patch_size[0]
    patch_height = latent_height // patch_size[1]
    patch_width = latent_width // patch_size[2]
    tokens = patch_frames * patch_height * patch_width
    return FastWanShape(
        height=height,
        width=width,
        num_frames=num_frames,
        latent_frames=latent_frames,
        latent_height=latent_height,
        latent_width=latent_width,
        patch_frames=patch_frames,
        patch_height=patch_height,
        patch_width=patch_width,
        tokens=tokens,
        hidden_size=num_heads * head_dim,
        num_heads=num_heads,
        head_dim=head_dim,
    )


def fastwan_shape_from_config(
    config_path: str | Path,
    *,
    height: int,
    width: int,
    num_frames: int,
) -> FastWanShape:
    config = json.loads(Path(config_path).read_text())
    return fastwan_shape(
        height=height,
        width=width,
        num_frames=num_frames,
        patch_size=tuple(config["patch_size"]),
        num_heads=int(config["num_attention_heads"]),
        head_dim=int(config["attention_head_dim"]),
    )


def replace_tokens(shape: FastWanShape, tokens: int) -> FastWanShape:
    return FastWanShape(**{**shape.__dict__, "tokens": tokens})


def median_ms(samples: list[float]) -> float:
    return statistics.median(samples) * 1000.0


def benchmark_mlx_attention(shape: FastWanShape, warmup: int, iters: int) -> float:
    import mlx.core as mx

    q = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16)
    k = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16)
    v = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16)
    scale = shape.head_dim**-0.5

    for _ in range(warmup):
        y = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
        mx.eval(y)

    samples = []
    for _ in range(iters):
        start = time.perf_counter()
        y = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
        mx.eval(y)
        samples.append(time.perf_counter() - start)
    return median_ms(samples)


def benchmark_mlx_linear(shape: FastWanShape, warmup: int, iters: int) -> float:
    import mlx.core as mx

    x = mx.random.normal((shape.tokens, shape.hidden_size), dtype=mx.float16)
    w = mx.random.normal((shape.hidden_size, shape.hidden_size), dtype=mx.float16)
    b = mx.zeros((shape.hidden_size, ), dtype=mx.float16)

    for _ in range(warmup):
        y = x @ w + b
        mx.eval(y)

    samples = []
    for _ in range(iters):
        start = time.perf_counter()
        y = x @ w + b
        mx.eval(y)
        samples.append(time.perf_counter() - start)
    return median_ms(samples)


def benchmark_torch_mps_attention(shape: FastWanShape, warmup: int, iters: int) -> float | None:
    try:
        import torch
        import torch.nn.functional as F
    except ImportError:
        return None

    if not torch.backends.mps.is_available():
        return None

    device = torch.device("mps")
    q = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16)
    k = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16)
    v = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16)

    for _ in range(warmup):
        y = F.scaled_dot_product_attention(q, k, v)
        torch.mps.synchronize()
        _ = y

    samples = []
    for _ in range(iters):
        start = time.perf_counter()
        y = F.scaled_dot_product_attention(q, k, v)
        torch.mps.synchronize()
        _ = y
        samples.append(time.perf_counter() - start)
    return median_ms(samples)


def torch_to_mx(tensor) -> mx.array:
    import mlx.core as mx

    return mx.array(tensor.detach().cpu().float().numpy())


def weight_dtype(weight):
    if isinstance(weight, QuantizedMatrix):
        return weight.dequantized_dtype
    return weight.dtype


_QUANT_SUPPORT_CACHE: dict[tuple[str, int | None, int | None], str | None] = {}


def quantization_support_error(spec: MLXQuantizationSpec) -> str | None:
    """Probe whether the installed MLX build supports ``spec``.

    Runs a tiny ``mx.quantize`` + ``mx.quantized_matmul`` with exactly the
    arguments :func:`quantize_matrix` / :func:`linear` use, so the result
    reflects the real runtime path. The affine (int8/int4) modes are stable
    across MLX releases, but the ``mxfp8``/``mxfp4``/``nvfp4`` mode strings
    require newer MLX builds and raise otherwise. Returns ``None`` when the
    mode works, else the underlying error message. Cached per spec.
    """
    key = (spec.mode, spec.bits, spec.group_size)
    if key not in _QUANT_SUPPORT_CACHE:
        import mlx.core as mx

        try:
            probe_dim = max(spec.group_size or 0, 64)
            weight = mx.zeros((probe_dim, probe_dim), dtype=mx.float16)
            quantized = quantize_matrix(weight, spec)
            y = linear(mx.zeros((1, probe_dim), dtype=mx.float16), quantized)
            mx.eval(y)
            _QUANT_SUPPORT_CACHE[key] = None
        except Exception as exc:  # noqa: BLE001 - MLX raises varied error types per backend/version.
            _QUANT_SUPPORT_CACHE[key] = f"{type(exc).__name__}: {exc}"
    return _QUANT_SUPPORT_CACHE[key]


def ensure_quantization_supported(spec: MLXQuantizationSpec | None) -> None:
    """Raise :class:`UnsupportedMLXQuantizationError` if ``spec`` cannot run here."""
    if spec is None:
        return
    error = quantization_support_error(spec)
    if error is None:
        return
    import mlx.core as mx

    mlx_version = getattr(mx, "__version__", "unknown")
    raise UnsupportedMLXQuantizationError(f"MLX quantization mode '{spec.label}' is not supported by the installed mlx "
                                          f"({mlx_version}): {error}. Upgrade mlx or pick a supported mode "
                                          f"(int8 is currently the most reliable quality/memory target).")


def quantize_matrix(weight, spec: MLXQuantizationSpec | None):
    if spec is None:
        return weight
    import mlx.core as mx

    if len(weight.shape) < 2:
        return weight
    q = mx.quantize(weight, group_size=spec.group_size, bits=spec.bits, mode=spec.mode)
    biases = q[2] if len(q) == 3 else None
    eval_args = [q[0], q[1]]
    if biases is not None:
        eval_args.append(biases)
    mx.eval(*eval_args)
    return QuantizedMatrix(
        weight=q[0],
        scales=q[1],
        biases=biases,
        spec=spec,
        dequantized_dtype=weight.dtype,
    )


# Affine quantized_matmul is slower than dequantize + steel GEMM at H3's packed
# token width. Measured on Apple M4 Max / MLX 0.32.2, INT6 group 64, BF16 acts,
# Q 5376→7168: M=256 qmm is faster; M=512 dequant+GEMM is +6.4%; M≥1024 ~10%.
# 832×480×124 packed M is ~14862–14994. H3 opts into this path explicitly;
# shared FastWan and Wan 2.2 linears stay on quantized_matmul. Do not cache
# dequantized weights. Override: FASTVIDEO_MLX_DQ_GEMM=0 off, =1 measured
# floor, =<int> explicit floor.
_AFFINE_DQ_GEMM_BITS = frozenset({2, 3, 4, 5, 6, 8})
MLX_AFFINE_DQ_GEMM_DEFAULT_MIN_M = 768
_dq_gemm_engaged = 0
_dq_gemm_logged = False


def reset_dq_gemm_telemetry() -> None:
    global _dq_gemm_engaged
    _dq_gemm_engaged = 0


def dq_gemm_engaged() -> int:
    return _dq_gemm_engaged


def affine_dq_gemm_min_m() -> int | None:
    raw = os.environ.get("FASTVIDEO_MLX_DQ_GEMM", "1").strip().lower()
    if raw in {"", "0", "off", "false", "no"}:
        return None
    if raw in {"1", "on", "true", "yes"}:
        return MLX_AFFINE_DQ_GEMM_DEFAULT_MIN_M
    try:
        value = int(raw)
    except ValueError:
        return MLX_AFFINE_DQ_GEMM_DEFAULT_MIN_M
    if value <= 0:
        return None
    return value


def _matmul_leading_rows(x) -> int:
    last = int(x.shape[-1]) if x.ndim else 0
    if last <= 0:
        return 0
    return int(x.size) // last


def _quantized_linear(x, weight: QuantizedMatrix, *, use_affine_dq_gemm: bool = False):
    import mlx.core as mx

    global _dq_gemm_engaged, _dq_gemm_logged
    spec = weight.spec
    min_m = affine_dq_gemm_min_m() if use_affine_dq_gemm else None
    rows = _matmul_leading_rows(x)
    if (min_m is not None and spec.mode == "affine" and spec.bits in _AFFINE_DQ_GEMM_BITS
            and spec.group_size is not None and rows >= min_m):
        dequantized = mx.dequantize(
            weight.weight,
            weight.scales,
            weight.biases,
            group_size=spec.group_size,
            bits=spec.bits,
            mode=spec.mode,
            dtype=x.dtype,
        )
        y = (x @ dequantized.T).astype(x.dtype)
        _dq_gemm_engaged += 1
        if not _dq_gemm_logged:
            _dq_gemm_logged = True
            logger.info("affine dequant+GEMM engaged (rows=%d, floor=%d, bits=%s)", rows, min_m, spec.bits)
        return y
    return mx.quantized_matmul(
        x,
        weight.weight,
        weight.scales,
        weight.biases,
        transpose=True,
        group_size=spec.group_size,
        bits=spec.bits,
        mode=spec.mode,
    ).astype(x.dtype)


def linear(x, weight, bias=None, *, use_affine_dq_gemm: bool = False):
    y = (_quantized_linear(x, weight, use_affine_dq_gemm=use_affine_dq_gemm)
         if isinstance(weight, QuantizedMatrix) else x @ weight.T)
    if bias is not None:
        y = y + bias
    return y


def _use_fast_norm() -> bool:
    """Opt-in to MLX's fused ``mx.fast`` normalization kernels for Wan.

    Off by default so the numerically-explicit reference path stays the
    baseline. Set ``FASTVIDEO_MLX_FAST_NORM=1`` to route LayerNorm/RMSNorm
    through single fused Metal kernels (fewer intermediates, less memory
    traffic) and benchmark the speedup.
    H3 uses its own fused RMSNorm default and does not consult this toggle.
    """
    import os

    return os.environ.get("FASTVIDEO_MLX_FAST_NORM", "0") == "1"


def layer_norm(x, weight=None, bias=None, eps: float = 1e-6):
    import mlx.core as mx

    if _use_fast_norm():
        # Compute in fp32 (matching the reference below) so downstream dtype
        # and precision are identical across call sites.
        w = weight.astype(mx.float32) if weight is not None else None
        b = bias.astype(mx.float32) if bias is not None else None
        return mx.fast.layer_norm(x.astype(mx.float32), w, b, eps)

    x_float = x.astype(mx.float32)
    mean = mx.mean(x_float, axis=-1, keepdims=True)
    var = mx.mean(mx.square(x_float - mean), axis=-1, keepdims=True)
    y = (x_float - mean) * mx.rsqrt(var + eps)
    if weight is not None:
        y = y * weight
    if bias is not None:
        y = y + bias
    return y


def rms_norm(x, weight, eps: float = 1e-6):
    import mlx.core as mx

    if _use_fast_norm():
        return mx.fast.rms_norm(x, weight, eps)

    orig_dtype = x.dtype
    x_float = x.astype(mx.float32)
    variance = mx.mean(mx.square(x_float), axis=-1, keepdims=True)
    y = x_float * mx.rsqrt(variance + eps)
    return y.astype(orig_dtype) * weight


def apply_rotary_emb(x, cos, sin, *, is_neox_style: bool = False):
    """Apply FastVideo's rotary convention to MLX tensors.

    Args:
        x: [batch, seq, heads, head_dim]
        cos/sin: [seq, head_dim] for Wan's full-dimension rotate-pair style,
          or [seq, head_dim // 2] for traditional RoPE.
    """
    import mlx.core as mx

    head_size = x.shape[-1]
    rope_dim = cos.shape[-1]
    cos = cos[None, :, None, :]
    sin = sin[None, :, None, :]
    x_float = x.astype(mx.float32)

    if rope_dim == head_size:
        x_pairs = x_float.reshape(*x.shape[:-1], -1, 2)
        x_real = x_pairs[..., 0]
        x_imag = x_pairs[..., 1]
        x_rotated = mx.stack([-x_imag, x_real], axis=-1).reshape(*x.shape)
        return (x_float * cos + x_rotated * sin).astype(x.dtype)

    if is_neox_style:
        x1, x2 = mx.split(x_float, 2, axis=-1)
        o1 = x1 * cos - x2 * sin
        o2 = x2 * cos + x1 * sin
        return mx.concatenate([o1, o2], axis=-1).astype(x.dtype)

    x1 = x_float[..., ::2]
    x2 = x_float[..., 1::2]
    o1 = x1 * cos - x2 * sin
    o2 = x2 * cos + x1 * sin
    return mx.stack([o1, o2], axis=-1).reshape(*x.shape).astype(x.dtype)


_WINDOWED_ATTENTION_WARNED = False


def _warn_windowed_attention_once(window: int) -> None:
    """Warn that FASTVIDEO_MLX_WINDOW degrades output on a dense-trained DiT.

    Sliding-window self-attention is fast (6.6x at a +-3-frame window on 1.3B)
    but these checkpoints were trained with dense attention, and restricting it
    at inference produces heavy colour-block noise: structural agreement with
    the dense baseline drops to 0.25 at +-3 frames and 0.03 at +-5. Sparsity of
    this kind is a training-time method. Kept as a research knob, but it should
    never be on by accident.
    """
    global _WINDOWED_ATTENTION_WARNED
    if _WINDOWED_ATTENTION_WARNED:
        return
    _WINDOWED_ATTENTION_WARNED = True
    logger.warning(
        "FASTVIDEO_MLX_WINDOW=%d enables sliding-window self-attention. These "
        "checkpoints are trained dense; expect severely degraded output. This is "
        "a research knob, not a speed setting — use --fast-spatial for real "
        "denoise savings.",
        window,
    )


def gelu_tanh(x):
    """tanh-approximate GELU, as used by Wan's FFN.

    ``mlx.nn.gelu_approx`` is the same tanh approximation behind a fused
    kernel. On the 1.3B FFN shape (32760x8960) it is bit-identical to the
    expanded expression below and 3.3x faster — 28.9ms -> 8.7ms per layer,
    which is 0.6s per denoise step across 30 layers.
    """
    import mlx.nn as nn

    return nn.gelu_approx(x)


def silu(x):
    import mlx.core as mx

    return x * mx.sigmoid(x)


def timestep_embedding(t, dim: int, max_period: int = 10000):
    import mlx.core as mx

    half = dim // 2
    freqs = mx.exp(-math.log(max_period) * mx.arange(0, half, dtype=mx.float32) / half)
    args = t[:, None].astype(mx.float32) * freqs[None]
    embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
    if dim % 2:
        embedding = mx.concatenate([embedding, mx.zeros_like(embedding[:, :1])], axis=-1)
    return embedding


def scale_residual(residual, x, gate):
    return residual + x * gate


def scale_residual_layer_norm_scale_shift(residual, x, gate, shift, scale, weight=None, bias=None, eps: float = 1e-6):
    if isinstance(gate, int):
        assert gate == 1
        residual_output = residual + x
    else:
        residual_output = residual + x * gate
    normalized = layer_norm(residual_output, weight=weight, bias=bias, eps=eps)
    modulated = normalized * (1.0 + scale) + shift
    return modulated, residual_output


class MLXWanT2VCrossAttention:

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

    def __call__(self, x, context):
        import mlx.core as mx

        batch = x.shape[0]
        q = linear(x, self.weights["attn2.to_q.weight"], self.weights.get("attn2.to_q.bias"))
        q = rms_norm(q, self.weights["attn2.norm_q.weight"], eps=self.eps).reshape(batch, -1, self.num_heads,
                                                                                   self.head_dim)

        if context.shape[1] == 0:
            attended = mx.zeros_like(q)
        else:
            k = linear(context, self.weights["attn2.to_k.weight"], self.weights.get("attn2.to_k.bias"))
            k = rms_norm(k, self.weights["attn2.norm_k.weight"],
                         eps=self.eps).reshape(batch, -1, self.num_heads, self.head_dim)
            v = linear(context, self.weights["attn2.to_v.weight"],
                       self.weights.get("attn2.to_v.bias")).reshape(batch, -1, self.num_heads, self.head_dim)
            attended = mx.fast.scaled_dot_product_attention(
                q.transpose(0, 2, 1, 3),
                k.transpose(0, 2, 1, 3),
                v.transpose(0, 2, 1, 3),
                scale=self.head_dim**-0.5,
            ).transpose(0, 2, 1, 3)

        attended = attended.reshape(batch, -1, self.dim)
        return linear(attended, self.weights["attn2.to_out.weight"], self.weights.get("attn2.to_out.bias"))


class MLXWanTransformerBlock:
    """Dense T2V Wan transformer block for the experimental MLX runtime.

    This mirrors the non-VSA PyTorch block for single-process dense attention.
    Rotary embeddings and sequence-parallel paths are intentionally left out of
    this first parity target.
    """

    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)

    def __call__(self, hidden_states, encoder_hidden_states, temb, freqs_cis=None):
        import mlx.core as mx

        orig_dtype = hidden_states.dtype
        e = self.weights["scale_shift_table"] + temb.astype(mx.float32)
        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = mx.split(e, 6, axis=1)

        norm_hidden_states = layer_norm(hidden_states.astype(mx.float32), eps=self.eps)
        norm_hidden_states = (norm_hidden_states * (1.0 + scale_msa) + shift_msa).astype(orig_dtype)

        query = linear(norm_hidden_states, self.weights["to_q.weight"], self.weights.get("to_q.bias"))
        key = linear(norm_hidden_states, self.weights["to_k.weight"], self.weights.get("to_k.bias"))
        value = linear(norm_hidden_states, self.weights["to_v.weight"], self.weights.get("to_v.bias"))

        query = rms_norm(query, self.weights["norm_q.weight"],
                         eps=self.eps).reshape(hidden_states.shape[0], -1, self.num_heads, self.head_dim)
        key = rms_norm(key, self.weights["norm_k.weight"], eps=self.eps).reshape(hidden_states.shape[0], -1,
                                                                                 self.num_heads, self.head_dim)
        value = value.reshape(hidden_states.shape[0], -1, self.num_heads, self.head_dim)

        if freqs_cis is not None:
            cos, sin = freqs_cis
            query = apply_rotary_emb(query, cos, sin, is_neox_style=False)
            key = apply_rotary_emb(key, cos, sin, is_neox_style=False)

        # Self-attention only. FASTVIDEO_MLX_WINDOW=0/unset → full SDPA (byte-identical
        # to the historical path). When >0, use chunked symmetric sliding-window
        # attention (see windowed_attention.py). Cross-attn (attn2) stays full.
        # Optional FASTVIDEO_MLX_WINDOW_SINK (default 0) adds global sink tokens.
        q_bh = query.transpose(0, 2, 1, 3)  # (B, H, S, D)
        k_bh = key.transpose(0, 2, 1, 3)
        v_bh = value.transpose(0, 2, 1, 3)
        scale = self.head_dim**-0.5
        window = int(os.environ.get("FASTVIDEO_MLX_WINDOW", "0") or "0")
        if window > 0:
            from fastvideo.mlx_runtime.windowed_attention import windowed_attention

            _warn_windowed_attention_once(window)
            sink = int(os.environ.get("FASTVIDEO_MLX_WINDOW_SINK", "0") or "0")
            attn_output = windowed_attention(q_bh, k_bh, v_bh, window=window, sink=sink, scale=scale)
        else:
            attn_output = mx.fast.scaled_dot_product_attention(q_bh, k_bh, v_bh, scale=scale)
        attn_output = attn_output.transpose(0, 2, 1, 3)
        attn_output = attn_output.reshape(hidden_states.shape[0], -1, self.dim)
        attn_output = linear(attn_output, self.weights["to_out.weight"], self.weights.get("to_out.bias"))

        norm_hidden_states, hidden_states = scale_residual_layer_norm_scale_shift(
            hidden_states,
            attn_output,
            gate_msa,
            0.0,
            0.0,
            weight=self.weights["self_attn_residual_norm.norm.weight"],
            bias=self.weights["self_attn_residual_norm.norm.bias"],
            eps=self.eps,
        )
        norm_hidden_states = norm_hidden_states.astype(orig_dtype)
        hidden_states = hidden_states.astype(orig_dtype)

        attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
        norm_hidden_states, hidden_states = scale_residual_layer_norm_scale_shift(
            hidden_states,
            attn_output,
            1,
            c_shift_msa,
            c_scale_msa,
            eps=self.eps,
        )
        norm_hidden_states = norm_hidden_states.astype(orig_dtype)
        hidden_states = hidden_states.astype(orig_dtype)

        ff_output = linear(norm_hidden_states, self.weights["ffn.fc_in.weight"], self.weights.get("ffn.fc_in.bias"))
        ff_output = gelu_tanh(ff_output)
        ff_output = linear(ff_output, self.weights["ffn.fc_out.weight"], self.weights.get("ffn.fc_out.bias"))
        hidden_states = scale_residual(hidden_states, ff_output, c_gate_msa)
        return hidden_states.astype(orig_dtype)


def mlx_block_weights_from_torch(torch_block) -> dict[str, mx.array]:
    return {name: torch_to_mx(value) for name, value in torch_block.state_dict().items()}


class MLXWanDiT:
    """Experimental FP16 Wan/FastWan DiT forward path in MLX."""

    def __init__(
        self,
        weights: dict[str, mx.array],
        blocks: list[MLXWanTransformerBlock],
        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.ffn_dim = int(config["ffn_dim"])
        self.in_channels = int(config["in_channels"])
        self.out_channels = int(config["out_channels"])
        self.patch_size = tuple(config["patch_size"])
        self.freq_dim = int(config["freq_dim"])
        # Opt-in graph fusion. With fixed weights and static shapes, the whole
        # denoise-step forward is a pure function of (latents, timestep) -- a
        # good mx.compile target. Off by default so the eager path stays the
        # baseline; enable via constructor or FASTVIDEO_MLX_COMPILE=1 and verify
        # with the benchmark's SSIM ~= 1.0 check.
        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

    def patch_embed(self, hidden_states):
        batch, channels, frames, height, width = hidden_states.shape
        pt, ph, pw = self.patch_size
        patch_dim = channels * pt * ph * pw
        x = hidden_states.reshape(batch, channels, frames // pt, pt, height // ph, ph, width // pw, pw)
        x = x.transpose(0, 2, 4, 6, 1, 3, 5, 7).reshape(batch, -1, patch_dim)
        return linear(x, self.weights["patch_embedding.weight"], self.weights.get("patch_embedding.bias"))

    def condition(self, timestep, encoder_hidden_states):
        t_freq = timestep_embedding(timestep, self.freq_dim).astype(
            weight_dtype(self.weights["condition_embedder.time_embedder.linear_1.weight"]))
        temb = linear(
            t_freq,
            self.weights["condition_embedder.time_embedder.linear_1.weight"],
            self.weights["condition_embedder.time_embedder.linear_1.bias"],
        )
        temb = silu(temb)
        temb = linear(
            temb,
            self.weights["condition_embedder.time_embedder.linear_2.weight"],
            self.weights["condition_embedder.time_embedder.linear_2.bias"],
        )
        timestep_proj = silu(temb)
        timestep_proj = linear(
            timestep_proj,
            self.weights["condition_embedder.time_proj.weight"],
            self.weights["condition_embedder.time_proj.bias"],
        ).reshape(timestep.shape[0], 6, self.hidden_size)

        encoder_hidden_states = linear(
            encoder_hidden_states,
            self.weights["condition_embedder.text_embedder.linear_1.weight"],
            self.weights["condition_embedder.text_embedder.linear_1.bias"],
        )
        encoder_hidden_states = gelu_tanh(encoder_hidden_states)
        encoder_hidden_states = linear(
            encoder_hidden_states,
            self.weights["condition_embedder.text_embedder.linear_2.weight"],
            self.weights["condition_embedder.text_embedder.linear_2.bias"],
        )
        return temb, timestep_proj, encoder_hidden_states

    def output(self, hidden_states, temb, *, batch: int, frames: int, height: int, width: int):
        pt, ph, pw = self.patch_size
        post_patch_frames = frames // pt
        post_patch_height = height // ph
        post_patch_width = width // pw
        shift, scale = mx_split_two(self.weights["scale_shift_table"] + temb[:, None, :], axis=1)
        hidden_states = layer_norm(hidden_states, eps=float(self.config["eps"])) * (1.0 + scale) + shift
        hidden_states = hidden_states.astype(weight_dtype(self.weights["proj_out.weight"]))
        hidden_states = linear(hidden_states, self.weights["proj_out.weight"], self.weights["proj_out.bias"])
        hidden_states = hidden_states.reshape(
            batch,
            post_patch_frames,
            post_patch_height,
            post_patch_width,
            pt,
            ph,
            pw,
            self.out_channels,
        )
        hidden_states = hidden_states.transpose(0, 7, 1, 4, 2, 5, 3, 6)
        return hidden_states.reshape(batch, self.out_channels, frames, height, width)

    def _forward(self, hidden_states, encoder_hidden_states, timestep, cos, sin):
        """Pure forward used both eagerly and as the mx.compile target.

        ``cos``/``sin`` are passed as separate array args (rather than a tuple)
        so the function traces cleanly under mx.compile.
        """
        batch, _, frames, height, width = hidden_states.shape
        freqs_cis = (cos, sin) if cos is not None else None
        hidden_states = self.patch_embed(hidden_states)
        temb, timestep_proj, encoder_hidden_states = self.condition(timestep, encoder_hidden_states)
        for block in self.blocks:
            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis=freqs_cis)
        return self.output(hidden_states, temb, batch=batch, frames=frames, height=height, width=width)

    def __call__(self, hidden_states, encoder_hidden_states, timestep, freqs_cis):
        cos, sin = freqs_cis if freqs_cis is not None else (None, None)
        if self._enable_compile and cos is not None:
            import mlx.core as mx

            # mx.compile keeps one traced graph per input signature, and each
            # graph pins its own materialization of the quantized weights. The
            # two-pass modes (--refine) call the DiT at a second resolution, so
            # keeping both graphs alive doubles resident DiT memory: 14B refine
            # peaked at 34.7 GiB instead of 20.8 GiB. Retire the previous graph
            # when the signature changes; the retrace costs far less than a
            # second copy of the weights.
            signature = (hidden_states.shape, encoder_hidden_states.shape, timestep.shape)
            if self._compiled_forward is not None and signature != self._compiled_signature:
                self._compiled_forward = None
                self._compiled_signature = None
                mx.clear_cache()
            if self._compiled_forward is None:
                self._compiled_forward = mx.compile(self._forward)
                self._compiled_signature = signature
            compiled_forward = self._compiled_forward
            try:
                return compiled_forward(hidden_states, encoder_hidden_states, timestep, cos, sin)
            except Exception as exc:  # noqa: BLE001 - some quant graphs may not trace; fall back to eager.
                logger.warning("mx.compile forward failed (%s); falling back to eager execution.", exc)
                self._enable_compile = False
                self._compiled_forward = None
                self._compiled_signature = None
        return self._forward(hidden_states, encoder_hidden_states, timestep, cos, sin)


def mx_split_two(x, *, axis: int):
    import mlx.core as mx

    left, right = mx.split(x, 2, axis=axis)
    return left, right


def _load_safetensor_value(handle, name: str):
    return handle.get_tensor(name)


def _load_mx_array_from_safetensor(handle, name: str, dtype):
    """Load a safetensors value and cast before creating the MLX array.

    The FastWan Diffusers checkpoint is fp32. Creating an MLX array first and
    then casting it to fp16 briefly materializes a large fp32 MLX allocation.
    Casting the CPU tensor before crossing into MLX keeps the transient GPU-side
    footprint lower.
    """
    import mlx.core as mx
    import torch

    tensor = handle.get_tensor(name)
    if dtype == mx.float16:
        tensor = tensor.to(torch.float16)
    elif dtype == mx.float32:
        tensor = tensor.to(torch.float32)
    elif dtype == mx.bfloat16:
        # NumPy has no bfloat16, so bridge through fp32 and cast on-device below.
        tensor = tensor.to(torch.float32)
    array = mx.array(tensor.numpy())
    del tensor
    if dtype is not None and array.dtype != dtype:
        array = array.astype(dtype)
    mx.eval(array)
    return array


def _eval_loaded_weight(value) -> None:
    import mlx.core as mx

    if isinstance(value, QuantizedMatrix):
        eval_args = [value.weight, value.scales]
        if value.biases is not None:
            eval_args.append(value.biases)
        mx.eval(*eval_args)
    else:
        mx.eval(value)


# Diffusers-to-FastVideo key mapping for WanTransformerBlock weights.
# Shared by both MLX and torch block loaders to keep mappings synchronized.
_WAN_BLOCK_KEY_MAP = {
    "scale_shift_table": "scale_shift_table",
    "attn1.to_q.weight": "to_q.weight",
    "attn1.to_q.bias": "to_q.bias",
    "attn1.to_k.weight": "to_k.weight",
    "attn1.to_k.bias": "to_k.bias",
    "attn1.to_v.weight": "to_v.weight",
    "attn1.to_v.bias": "to_v.bias",
    "attn1.to_out.0.weight": "to_out.weight",
    "attn1.to_out.0.bias": "to_out.bias",
    "attn1.norm_q.weight": "norm_q.weight",
    "attn1.norm_k.weight": "norm_k.weight",
    "attn2.to_q.weight": "attn2.to_q.weight",
    "attn2.to_q.bias": "attn2.to_q.bias",
    "attn2.to_k.weight": "attn2.to_k.weight",
    "attn2.to_k.bias": "attn2.to_k.bias",
    "attn2.to_v.weight": "attn2.to_v.weight",
    "attn2.to_v.bias": "attn2.to_v.bias",
    "attn2.to_out.0.weight": "attn2.to_out.weight",
    "attn2.to_out.0.bias": "attn2.to_out.bias",
    "attn2.norm_q.weight": "attn2.norm_q.weight",
    "attn2.norm_k.weight": "attn2.norm_k.weight",
    "ffn.net.0.proj.weight": "ffn.fc_in.weight",
    "ffn.net.0.proj.bias": "ffn.fc_in.bias",
    "ffn.net.2.weight": "ffn.fc_out.weight",
    "ffn.net.2.bias": "ffn.fc_out.bias",
    "norm2.weight": "self_attn_residual_norm.norm.weight",
    "norm2.bias": "self_attn_residual_norm.norm.bias",
}


def mlx_block_weights_from_diffusers_safetensors(
    checkpoint_path: str | Path,
    *,
    block_index: int = 0,
    quantization: str | MLXQuantizationSpec | None = None,
    dtype=None,
) -> dict[str, mx.array]:
    """Load one Diffusers-format Wan block into the MLX dense-block key layout."""
    from safetensors import safe_open

    prefix = f"blocks.{block_index}."
    key_map = _WAN_BLOCK_KEY_MAP

    spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None
                                                           or isinstance(quantization, str)) else quantization
    ensure_quantization_supported(spec)
    matrix_targets = {target for target in key_map.values() if target.endswith(".weight") and "norm" not in target}
    weights = {}
    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
        available = set(handle.keys())
        for source_name, target_name in key_map.items():
            full = prefix + source_name
            if full not in available:
                # Biases are optional: e.g. Wan2.1-14B has bias-free attention/FFN.
                # The block forward already fetches biases via ``.get(...)``.
                if source_name.endswith(".bias"):
                    continue
                raise KeyError(f"missing required block weight: {full}")
            array = _load_mx_array_from_safetensor(handle, full, dtype)
            loaded = quantize_matrix(array, spec) if target_name in matrix_targets else array
            _eval_loaded_weight(loaded)
            weights[target_name] = loaded
            del array
    return weights


def mlx_dit_from_diffusers_safetensors(
    checkpoint_path: str | Path,
    config_path: str | Path,
    *,
    dtype: str = "fp16",
    num_blocks: int | None = None,
    quantization: str | MLXQuantizationSpec | None = None,
    compile: bool = False,
) -> MLXWanDiT:
    import mlx.core as mx
    from safetensors import safe_open

    from fastvideo.mlx_runtime.checkpoint_compat import raise_if_unsupported_mlx_checkpoint

    raise_if_unsupported_mlx_checkpoint(checkpoint_path, config_path)

    config = json.loads(Path(config_path).read_text())
    total_blocks = int(config["num_layers"])
    if num_blocks is None:
        num_blocks = total_blocks
    cast_dtype = {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[dtype]
    spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None
                                                           or isinstance(quantization, str)) else quantization
    ensure_quantization_supported(spec)

    top_level_names = [
        "patch_embedding.weight",
        "patch_embedding.bias",
        "condition_embedder.time_embedder.linear_1.weight",
        "condition_embedder.time_embedder.linear_1.bias",
        "condition_embedder.time_embedder.linear_2.weight",
        "condition_embedder.time_embedder.linear_2.bias",
        "condition_embedder.time_proj.weight",
        "condition_embedder.time_proj.bias",
        "condition_embedder.text_embedder.linear_1.weight",
        "condition_embedder.text_embedder.linear_1.bias",
        "condition_embedder.text_embedder.linear_2.weight",
        "condition_embedder.text_embedder.linear_2.bias",
        "scale_shift_table",
        "proj_out.weight",
        "proj_out.bias",
    ]
    weights = {}
    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
        available = set(handle.keys())
        for name in top_level_names:
            if name not in available:
                if name.endswith(".bias"):
                    continue
                raise KeyError(f"missing required weight: {name}")
            array = _load_mx_array_from_safetensor(handle, name, cast_dtype)
            if name == "patch_embedding.weight":
                array = array.reshape(int(config["num_attention_heads"]) * int(config["attention_head_dim"]), -1)
            if name.endswith(".weight") and name not in {"scale_shift_table"}:
                loaded = quantize_matrix(array, spec)
            else:
                loaded = array
            _eval_loaded_weight(loaded)
            weights[name] = loaded
            del array

    blocks = []
    for block_index in range(num_blocks):
        block_weights = mlx_block_weights_from_diffusers_safetensors(
            checkpoint_path,
            block_index=block_index,
            quantization=spec,
            dtype=cast_dtype,
        )
        block_weights = {
            name: (value if isinstance(value, QuantizedMatrix) else value.astype(cast_dtype))
            for name, value in block_weights.items()
        }
        for value in block_weights.values():
            _eval_loaded_weight(value)
        blocks.append(
            MLXWanTransformerBlock(
                block_weights,
                dim=int(config["num_attention_heads"]) * int(config["attention_head_dim"]),
                ffn_dim=int(config["ffn_dim"]),
                num_heads=int(config["num_attention_heads"]),
                eps=float(config["eps"]),
            ))
    return MLXWanDiT(weights, blocks, config, compile=compile)


def torch_block_state_from_diffusers_safetensors(
    checkpoint_path: str | Path,
    *,
    block_index: int = 0,
) -> dict[str, torch.Tensor]:
    """Load one Diffusers-format Wan block into FastVideo's dense block keys."""
    from safetensors import safe_open

    prefix = f"blocks.{block_index}."
    key_map = _WAN_BLOCK_KEY_MAP

    state = {}
    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
        available = set(handle.keys())
        for source_name, target_name in key_map.items():
            full = prefix + source_name
            if full not in available:
                # Biases are optional: e.g. Wan2.1-14B has bias-free attention/FFN.
                if source_name.endswith(".bias"):
                    continue
                raise KeyError(f"missing required block weight: {full}")
            state[target_name] = handle.get_tensor(full).float()
    return state