287 lines
13 KiB
Python
287 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""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``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
from collections.abc import Callable
|
|
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.mlx_runtime.fastwan import (
|
|
MLXWanT2VCrossAttention,
|
|
gelu_tanh,
|
|
layer_norm,
|
|
linear,
|
|
mlx_dit_from_diffusers_safetensors,
|
|
rms_norm,
|
|
silu,
|
|
timestep_embedding,
|
|
weight_dtype,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
import mlx.core as mx
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class MLXWan22TransformerBlock:
|
|
"""Dense Wan block with per-token (``[B, L, dim]``) timestep modulation."""
|
|
|
|
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, timestep_proj, cos, sin) -> mx.array:
|
|
import mlx.core as mx
|
|
|
|
orig_dtype = hidden_states.dtype
|
|
batch = hidden_states.shape[0]
|
|
|
|
# timestep_proj: [B, L, 6, dim] -> six per-token [B, L, dim] modulations.
|
|
e = self.weights["scale_shift_table"][None] + timestep_proj.astype(mx.float32)
|
|
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = [
|
|
part.squeeze(2) for part in mx.split(e, 6, axis=2)
|
|
]
|
|
|
|
# 1. Self-attention (dense, bidirectional) with per-token modulation.
|
|
norm_hidden = layer_norm(hidden_states.astype(mx.float32), eps=self.eps)
|
|
norm_hidden = (norm_hidden * (1.0 + scale_msa) + shift_msa).astype(orig_dtype)
|
|
|
|
query = linear(norm_hidden, self.weights["to_q.weight"], self.weights.get("to_q.bias"))
|
|
key = linear(norm_hidden, self.weights["to_k.weight"], self.weights.get("to_k.bias"))
|
|
value = linear(norm_hidden, self.weights["to_v.weight"], self.weights.get("to_v.bias"))
|
|
|
|
query = rms_norm(query, self.weights["norm_q.weight"],
|
|
eps=self.eps).reshape(batch, -1, self.num_heads, self.head_dim)
|
|
key = rms_norm(key, self.weights["norm_k.weight"], eps=self.eps).reshape(batch, -1, self.num_heads,
|
|
self.head_dim)
|
|
value = value.reshape(batch, -1, self.num_heads, self.head_dim)
|
|
|
|
from fastvideo.mlx_runtime.fastwan import apply_rotary_emb
|
|
query = apply_rotary_emb(query, cos, sin, is_neox_style=False)
|
|
key = apply_rotary_emb(key, cos, sin, is_neox_style=False)
|
|
|
|
attn = mx.fast.scaled_dot_product_attention(
|
|
query.transpose(0, 2, 1, 3),
|
|
key.transpose(0, 2, 1, 3),
|
|
value.transpose(0, 2, 1, 3),
|
|
scale=self.head_dim**-0.5,
|
|
).transpose(0, 2, 1, 3)
|
|
attn = attn.reshape(batch, -1, self.dim)
|
|
attn = linear(attn, self.weights["to_out.weight"], self.weights.get("to_out.bias"))
|
|
|
|
hidden_states = hidden_states + (attn * gate_msa).astype(orig_dtype)
|
|
norm_hidden = layer_norm(hidden_states.astype(mx.float32),
|
|
weight=self.weights["self_attn_residual_norm.norm.weight"],
|
|
bias=self.weights["self_attn_residual_norm.norm.bias"],
|
|
eps=self.eps).astype(orig_dtype)
|
|
|
|
# 2. Cross-attention, then per-token shift/scale modulation.
|
|
cross = self.attn2(norm_hidden, encoder_hidden_states)
|
|
hidden_states = hidden_states + cross
|
|
norm_hidden = layer_norm(hidden_states.astype(mx.float32), eps=self.eps)
|
|
norm_hidden = (norm_hidden * (1.0 + c_scale_msa) + c_shift_msa).astype(orig_dtype)
|
|
|
|
# 3. Feed-forward with per-token gate.
|
|
ff = linear(norm_hidden, self.weights["ffn.fc_in.weight"], self.weights.get("ffn.fc_in.bias"))
|
|
ff = gelu_tanh(ff)
|
|
ff = linear(ff, self.weights["ffn.fc_out.weight"], self.weights.get("ffn.fc_out.bias"))
|
|
hidden_states = hidden_states + (ff * c_gate_msa).astype(orig_dtype)
|
|
return hidden_states.astype(orig_dtype)
|
|
|
|
|
|
class MLXWan22DiT:
|
|
"""Wan2.2-TI2V-5B dense DiT with per-token timestep conditioning."""
|
|
|
|
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
|
|
|
|
def _patch_embed(self, hidden_states) -> mx.array:
|
|
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) -> tuple:
|
|
"""Per-token conditioning. ``timestep`` is ``[B, L]`` (one level per token)."""
|
|
batch, seq = timestep.shape
|
|
t_freq = timestep_embedding(timestep.reshape(-1), 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 = linear(silu(temb), self.weights["condition_embedder.time_proj.weight"],
|
|
self.weights["condition_embedder.time_proj.bias"])
|
|
timestep_proj = timestep_proj.reshape(batch, seq, 6, self.hidden_size)
|
|
|
|
ehs = linear(encoder_hidden_states, self.weights["condition_embedder.text_embedder.linear_1.weight"],
|
|
self.weights["condition_embedder.text_embedder.linear_1.bias"])
|
|
ehs = gelu_tanh(ehs)
|
|
ehs = linear(ehs, self.weights["condition_embedder.text_embedder.linear_2.weight"],
|
|
self.weights["condition_embedder.text_embedder.linear_2.bias"])
|
|
temb_out = temb.reshape(batch, seq, self.hidden_size)
|
|
return temb_out, timestep_proj, ehs
|
|
|
|
def _output(self, hidden_states, temb_out, *, batch, frames, height, width) -> mx.array:
|
|
import mlx.core as mx
|
|
|
|
pt, ph, pw = self.patch_size
|
|
post_pt, post_ph, post_pw = frames // pt, height // ph, width // pw
|
|
# Per-token output modulation: scale_shift_table[1,2,dim] + temb[B,L,1,dim].
|
|
e = self.weights["scale_shift_table"][None] + temb_out[:, :, None, :].astype(mx.float32)
|
|
shift, scale = [part.squeeze(2) for part in mx.split(e, 2, axis=2)]
|
|
norm = layer_norm(hidden_states.astype(mx.float32), eps=self.eps)
|
|
norm = (norm * (1.0 + scale) + shift).astype(weight_dtype(self.weights["proj_out.weight"]))
|
|
out = linear(norm, self.weights["proj_out.weight"], self.weights["proj_out.bias"])
|
|
out = out.reshape(batch, post_pt, post_ph, post_pw, pt, ph, pw, self.out_channels)
|
|
out = out.transpose(0, 7, 1, 4, 2, 5, 3, 6)
|
|
return out.reshape(batch, self.out_channels, frames, height, width)
|
|
|
|
def _forward(self, hidden_states, encoder_hidden_states, timestep, cos, sin) -> mx.array:
|
|
batch, _, frames, height, width = hidden_states.shape
|
|
hidden = self._patch_embed(hidden_states)
|
|
temb_out, timestep_proj, ehs = self._condition(timestep, encoder_hidden_states)
|
|
for block in self.blocks:
|
|
hidden = block(hidden, ehs, timestep_proj, cos, sin)
|
|
return self._output(hidden, temb_out, batch=batch, frames=frames, height=height, width=width)
|
|
|
|
def __call__(self, hidden_states, encoder_hidden_states, timestep, freqs_cis) -> mx.array:
|
|
cos, sin = freqs_cis
|
|
if self._enable_compile and cos is not None:
|
|
import mlx.core as mx
|
|
|
|
# One traced graph per input signature, and each pins its own copy of
|
|
# the quantized weights. --refine denoises at two resolutions, so
|
|
# keeping both alive doubles resident DiT memory. Retire the previous
|
|
# graph when the signature changes.
|
|
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 graphs may not trace; fall back to eager.
|
|
logger.warning(
|
|
"Wan2.2 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 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)
|
|
|
|
|
|
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)
|