981 lines
37 KiB
Python
981 lines
37 KiB
Python
# 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 == "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,
|
|
)
|
|
|
|
|
|
def linear(x, weight, bias=None):
|
|
import mlx.core as mx
|
|
|
|
if isinstance(weight, QuantizedMatrix):
|
|
y = mx.quantized_matmul(
|
|
x,
|
|
weight.weight,
|
|
weight.scales,
|
|
weight.biases,
|
|
transpose=True,
|
|
group_size=weight.spec.group_size,
|
|
bits=weight.spec.bits,
|
|
mode=weight.spec.mode,
|
|
).astype(x.dtype)
|
|
else:
|
|
y = 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.
|
|
|
|
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.
|
|
"""
|
|
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
|
|
|
|
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
|