fastwan ¶

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)

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