291 lines
11 KiB
Python
291 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tiny random-weight Wan DiT fixtures shared by the MLX runtime tests.
|
|
|
|
Builds a miniature ``WanTransformer3DModel`` (all matmul dims are multiples of
|
|
the int8 group size 64, so the same model exercises the quantized paths) plus
|
|
the conversions into ``fastvideo.mlx_runtime.fastwan.MLXWanDiT``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
|
|
os.environ.setdefault("MASTER_ADDR", "localhost")
|
|
os.environ.setdefault("MASTER_PORT", "29513")
|
|
|
|
from fastvideo.configs.models.dits.wanvideo import ( # noqa: E402
|
|
WanVideoArchConfig,
|
|
WanVideoConfig,
|
|
)
|
|
from fastvideo.forward_context import set_forward_context # noqa: E402
|
|
from fastvideo.layers.rotary_embedding import get_rotary_pos_embed # noqa: E402
|
|
from fastvideo.mlx_runtime.fastwan import ( # noqa: E402
|
|
MLXQuantizationSpec,
|
|
MLXWanDiT,
|
|
MLXWanTransformerBlock,
|
|
mlx_block_weights_from_torch,
|
|
quantize_matrix,
|
|
)
|
|
from fastvideo.models.dits.wanvideo import WanTransformer3DModel # noqa: E402
|
|
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch # noqa: E402
|
|
|
|
SEED = 2026
|
|
|
|
# All matmul dims must be multiples of the int8 group size (64) so the
|
|
# quantized variants of the tests can reuse the same tiny model.
|
|
TINY_ARCH = dict(
|
|
num_attention_heads=4,
|
|
attention_head_dim=16,
|
|
in_channels=16,
|
|
out_channels=16,
|
|
text_dim=64,
|
|
freq_dim=64,
|
|
ffn_dim=128,
|
|
num_layers=2,
|
|
patch_size=(1, 2, 2),
|
|
rope_max_seq_len=64,
|
|
)
|
|
|
|
|
|
def build_tiny_wan_config() -> WanVideoConfig:
|
|
"""Build the deterministic miniature Wan video configuration.
|
|
|
|
Returns:
|
|
WanVideoConfig: Configuration created from the tiny test architecture.
|
|
"""
|
|
return WanVideoConfig(arch_config=WanVideoArchConfig(**TINY_ARCH))
|
|
|
|
|
|
def build_hf_config(config: WanVideoConfig) -> dict[str, object]:
|
|
"""
|
|
Builds the configuration dictionary required by the Hugging Face-style Wan model representation.
|
|
|
|
Parameters:
|
|
config (WanVideoConfig): Wan video configuration to convert.
|
|
|
|
Returns:
|
|
dict[str, object]: Configuration values for the Hugging Face-style model.
|
|
"""
|
|
return {
|
|
"num_attention_heads": config.num_attention_heads,
|
|
"attention_head_dim": config.attention_head_dim,
|
|
"in_channels": config.in_channels,
|
|
"out_channels": config.out_channels,
|
|
"text_dim": config.text_dim,
|
|
"freq_dim": config.freq_dim,
|
|
"ffn_dim": config.ffn_dim,
|
|
"num_layers": config.num_layers,
|
|
"patch_size": config.patch_size,
|
|
"text_len": config.text_len,
|
|
"rope_max_seq_len": config.rope_max_seq_len,
|
|
"eps": 1e-6,
|
|
}
|
|
|
|
|
|
def initialize_model_parameters(model: torch.nn.Module) -> None:
|
|
# ReplicatedLinear parameters are allocated with torch.empty and need an
|
|
# explicit initialization in tests to avoid undefined values.
|
|
"""Initialize model parameters deterministically for reproducible tests.
|
|
|
|
Parameters:
|
|
model (torch.nn.Module): Model whose parameters will be initialized.
|
|
"""
|
|
torch.manual_seed(SEED + 3)
|
|
with torch.no_grad():
|
|
for name, param in model.named_parameters():
|
|
if param.ndim <= 1:
|
|
if name.endswith("weight") and "norm" in name:
|
|
param.fill_(1.0)
|
|
else:
|
|
param.normal_(mean=0.0, std=0.02)
|
|
continue
|
|
torch.nn.init.xavier_uniform_(param)
|
|
|
|
|
|
def build_torch_model() -> WanTransformer3DModel:
|
|
"""
|
|
Build a deterministically initialized CPU float32 Wan transformer model.
|
|
|
|
Returns:
|
|
WanTransformer3DModel: The initialized model in evaluation mode.
|
|
"""
|
|
config = build_tiny_wan_config()
|
|
model = WanTransformer3DModel(config=config, hf_config=build_hf_config(config))
|
|
model = model.to(device="cpu", dtype=torch.float32)
|
|
initialize_model_parameters(model)
|
|
model.eval()
|
|
return model
|
|
|
|
|
|
def build_inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Create deterministic latent, text-encoder, and timestep inputs for the tiny Wan model.
|
|
|
|
Returns:
|
|
tuple[torch.Tensor, torch.Tensor, torch.Tensor]: The latent hidden states, text-encoder hidden states, and timestep."""
|
|
generator = torch.Generator(device="cpu").manual_seed(SEED + 1)
|
|
hidden_states = torch.randn(1, TINY_ARCH["in_channels"], 4, 8, 8, generator=generator, dtype=torch.float32)
|
|
encoder_hidden_states = torch.randn(1, 8, TINY_ARCH["text_dim"], generator=generator, dtype=torch.float32)
|
|
timestep = torch.tensor([10], dtype=torch.long)
|
|
return hidden_states, encoder_hidden_states, timestep
|
|
|
|
|
|
# Top-level weights: the MLX runtime uses the Diffusers key layout; the torch
|
|
# model uses FastVideo's module names (see WanVideoConfig.param_names_mapping).
|
|
TOP_LEVEL_KEY_MAP = {
|
|
"patch_embedding.weight": "patch_embedding.proj.weight",
|
|
"patch_embedding.bias": "patch_embedding.proj.bias",
|
|
"condition_embedder.time_embedder.linear_1.weight": "condition_embedder.time_embedder.mlp.fc_in.weight",
|
|
"condition_embedder.time_embedder.linear_1.bias": "condition_embedder.time_embedder.mlp.fc_in.bias",
|
|
"condition_embedder.time_embedder.linear_2.weight": "condition_embedder.time_embedder.mlp.fc_out.weight",
|
|
"condition_embedder.time_embedder.linear_2.bias": "condition_embedder.time_embedder.mlp.fc_out.bias",
|
|
"condition_embedder.time_proj.weight": "condition_embedder.time_modulation.linear.weight",
|
|
"condition_embedder.time_proj.bias": "condition_embedder.time_modulation.linear.bias",
|
|
"condition_embedder.text_embedder.linear_1.weight": "condition_embedder.text_embedder.fc_in.weight",
|
|
"condition_embedder.text_embedder.linear_1.bias": "condition_embedder.text_embedder.fc_in.bias",
|
|
"condition_embedder.text_embedder.linear_2.weight": "condition_embedder.text_embedder.fc_out.weight",
|
|
"condition_embedder.text_embedder.linear_2.bias": "condition_embedder.text_embedder.fc_out.bias",
|
|
"scale_shift_table": "scale_shift_table",
|
|
"proj_out.weight": "proj_out.weight",
|
|
"proj_out.bias": "proj_out.bias",
|
|
}
|
|
|
|
|
|
def mlx_dit_from_torch_model(
|
|
model: WanTransformer3DModel,
|
|
hf_config: dict[str, object],
|
|
*,
|
|
quantization: MLXQuantizationSpec | None = None,
|
|
) -> MLXWanDiT:
|
|
"""
|
|
Convert a PyTorch Wan DiT model to an MLX Wan DiT model.
|
|
|
|
Parameters:
|
|
model (WanTransformer3DModel): PyTorch model whose parameters are converted.
|
|
hf_config (dict[str, object]): Model configuration used to construct the MLX representation.
|
|
quantization (MLXQuantizationSpec | None): Optional quantization specification for eligible weight matrices.
|
|
|
|
Returns:
|
|
MLXWanDiT: The converted MLX model.
|
|
"""
|
|
import mlx.core as mx
|
|
|
|
state = {name: value.detach().float() for name, value in model.state_dict().items()}
|
|
inner_dim = int(hf_config["num_attention_heads"]) * int(hf_config["attention_head_dim"]) # type: ignore[arg-type]
|
|
|
|
weights = {}
|
|
for mlx_name, torch_name in TOP_LEVEL_KEY_MAP.items():
|
|
tensor = state[torch_name]
|
|
if mlx_name == "patch_embedding.weight":
|
|
tensor = tensor.reshape(inner_dim, -1)
|
|
array = mx.array(tensor.numpy())
|
|
if quantization is not None and mlx_name.endswith(".weight") and mlx_name != "scale_shift_table":
|
|
weights[mlx_name] = quantize_matrix(array, quantization)
|
|
else:
|
|
weights[mlx_name] = array
|
|
|
|
blocks = []
|
|
for torch_block in model.blocks:
|
|
block_weights = mlx_block_weights_from_torch(torch_block)
|
|
if quantization is not None:
|
|
block_weights = {
|
|
name: (quantize_matrix(value, quantization)
|
|
if name.endswith(".weight") and "norm" not in name and len(value.shape) >= 2 else value)
|
|
for name, value in block_weights.items()
|
|
}
|
|
blocks.append(
|
|
MLXWanTransformerBlock(
|
|
block_weights,
|
|
dim=inner_dim,
|
|
ffn_dim=int(hf_config["ffn_dim"]), # type: ignore[arg-type]
|
|
num_heads=int(hf_config["num_attention_heads"]), # type: ignore[arg-type]
|
|
eps=float(hf_config["eps"]), # type: ignore[arg-type]
|
|
))
|
|
return MLXWanDiT(weights, blocks, dict(hf_config))
|
|
|
|
|
|
def mlx_rotary_embeddings(hidden_states: torch.Tensor):
|
|
"""
|
|
Generate MLX rotary-position embedding tables matching the tiny Wan model.
|
|
|
|
Parameters:
|
|
hidden_states (torch.Tensor): Input latent tensor whose spatial and temporal dimensions determine the embedding grid.
|
|
|
|
Returns:
|
|
Tuple of MLX arrays containing the cosine and sine rotary-position embeddings.
|
|
"""
|
|
import mlx.core as mx
|
|
|
|
_, _, frames, height, width = hidden_states.shape
|
|
p_t, p_h, p_w = TINY_ARCH["patch_size"]
|
|
head_dim = TINY_ARCH["attention_head_dim"]
|
|
hidden_size = TINY_ARCH["num_attention_heads"] * head_dim
|
|
rope_dim_list = [head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)]
|
|
freqs_cos, freqs_sin = get_rotary_pos_embed(
|
|
(frames // p_t, height // p_h, width // p_w),
|
|
hidden_size,
|
|
TINY_ARCH["num_attention_heads"],
|
|
rope_dim_list,
|
|
dtype=torch.float64,
|
|
rope_theta=10000,
|
|
)
|
|
return (
|
|
mx.array(freqs_cos.float().numpy()).astype(mx.float32),
|
|
mx.array(freqs_sin.float().numpy()).astype(mx.float32),
|
|
)
|
|
|
|
|
|
def torch_reference_output(model, hidden_states, encoder_hidden_states, timestep) -> np.ndarray:
|
|
"""
|
|
Run the PyTorch model on the provided inputs and return its output as a NumPy array.
|
|
|
|
Parameters:
|
|
model: The PyTorch Wan transformer model.
|
|
hidden_states: The latent hidden states provided to the model.
|
|
encoder_hidden_states: The encoder hidden states provided to the model.
|
|
timestep: The diffusion timestep provided to the model.
|
|
|
|
Returns:
|
|
The model output as a CPU float32 NumPy array.
|
|
"""
|
|
with torch.no_grad(), set_forward_context(
|
|
current_timestep=0,
|
|
attn_metadata=None,
|
|
forward_batch=ForwardBatch(data_type="dummy"),
|
|
):
|
|
output = model(
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
timestep=timestep,
|
|
)
|
|
return output.detach().float().cpu().numpy()
|
|
|
|
|
|
def mlx_output(dit, hidden_states, encoder_hidden_states, timestep, freqs_cis) -> np.ndarray:
|
|
"""
|
|
Run the MLX DiT model on the provided inputs.
|
|
|
|
Parameters:
|
|
dit: The MLX DiT model to evaluate.
|
|
hidden_states: Input latent states.
|
|
encoder_hidden_states: Text encoder states.
|
|
timestep: Diffusion timestep input.
|
|
freqs_cis: Rotary-position embedding frequencies.
|
|
|
|
Returns:
|
|
np.ndarray: The model output as a float32 NumPy array.
|
|
"""
|
|
import mlx.core as mx
|
|
|
|
out = dit(
|
|
mx.array(hidden_states.numpy()),
|
|
mx.array(encoder_hidden_states.numpy()),
|
|
mx.array(timestep.float().numpy()),
|
|
freqs_cis,
|
|
)
|
|
mx.eval(out)
|
|
return np.array(out.astype(mx.float32))
|