ac48bb3cd1
Co-authored-by: Aryan Kumar <aryank@Aryans-Mac-Studio.local> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: CodeRabbit <noreply@coderabbit.ai>
284 lines
10 KiB
Python
284 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tiny random-weight MiniMax H3 fixtures shared by the MLX runtime tests.
|
|
|
|
Builds a miniature ``MiniMaxH3Transformer3DModel`` (upstream torch reference,
|
|
merged in #1674) plus the conversion into
|
|
``fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT``. All matmul dims are
|
|
multiples of the int8 group size (64) so the quantized variants of the tests
|
|
reuse the same model; ``rope_freq_dim`` keeps the rotary width (6 * dim)
|
|
inside the head dim, as in the real config (96 <= 128).
|
|
"""
|
|
|
|
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.minimax_h3 import ( # noqa: E402
|
|
MiniMaxH3ArchConfig,
|
|
MiniMaxH3Config,
|
|
)
|
|
from fastvideo.forward_context import set_forward_context # noqa: E402
|
|
from fastvideo.mlx_runtime.minimax_h3 import ( # noqa: E402
|
|
MLXMiniMaxH3DiT,
|
|
MiniMaxH3SchedulerState,
|
|
MINIMAX_H3_AUDIO_SHIFT,
|
|
MINIMAX_H3_VIDEO_SHIFT,
|
|
build_packed_layout,
|
|
build_row_timesteps,
|
|
patchify_video_latents,
|
|
)
|
|
from fastvideo.models.dits.minimax_h3 import MiniMaxH3Transformer3DModel # noqa: E402
|
|
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch # noqa: E402
|
|
|
|
SEED = 2027
|
|
|
|
TINY_ARCH = dict(
|
|
num_attention_heads=8,
|
|
attention_head_dim=64, # inner 512 > hidden 256, like the real 7168 > 5376
|
|
hidden_size=256,
|
|
num_layers=2,
|
|
num_refiner_layers=1,
|
|
ffn_dim=384,
|
|
in_channels=8,
|
|
audio_in_channels=4,
|
|
patch_size=(1, 2, 2),
|
|
text_dim=128,
|
|
freq_dim=64,
|
|
time_embed_hidden_dim=256,
|
|
time_embed_dim=192,
|
|
rope_freq_dim=8, # rotary width 48 <= 64 head dim
|
|
rope_theta=10000.0,
|
|
norm_eps=1e-5,
|
|
qk_norm_eps=1e-5,
|
|
final_norm_eps=1e-5,
|
|
)
|
|
|
|
# Geometry: 2 latent frames of 4x4, patch (1,2,2) -> 4 rows/frame -> 8 video
|
|
# rows; 6 text rows; 3 audio latents -> 6 audio rows. Sequence length 20.
|
|
NUM_TEXT_TOKENS = 6
|
|
NUM_LATENT_FRAMES = 2
|
|
LATENT_HEIGHT = 4
|
|
LATENT_WIDTH = 4
|
|
NUM_AUDIO_LATENTS = 3
|
|
VIDEO_TIMESTEP = 0.7
|
|
AUDIO_TIMESTEP = 0.4
|
|
|
|
|
|
def build_tiny_h3_config() -> MiniMaxH3Config:
|
|
return MiniMaxH3Config(arch_config=MiniMaxH3ArchConfig(**TINY_ARCH))
|
|
|
|
|
|
def build_hf_config(config: MiniMaxH3Config) -> dict[str, object]:
|
|
arch = config.arch_config
|
|
return {
|
|
"num_attention_heads": arch.num_attention_heads,
|
|
"attention_head_dim": arch.attention_head_dim,
|
|
"hidden_size": arch.hidden_size,
|
|
"num_layers": arch.num_layers,
|
|
"num_refiner_layers": arch.num_refiner_layers,
|
|
"ffn_dim": arch.ffn_dim,
|
|
"in_channels": arch.in_channels,
|
|
"audio_in_channels": arch.audio_in_channels,
|
|
"patch_size": list(arch.patch_size),
|
|
"text_dim": arch.text_dim,
|
|
"freq_dim": arch.freq_dim,
|
|
"time_embed_hidden_dim": arch.time_embed_hidden_dim,
|
|
"time_embed_dim": arch.time_embed_dim,
|
|
"rope_freq_dim": arch.rope_freq_dim,
|
|
"rope_theta": arch.rope_theta,
|
|
"norm_eps": arch.norm_eps,
|
|
"qk_norm_eps": arch.qk_norm_eps,
|
|
"final_norm_eps": arch.final_norm_eps,
|
|
}
|
|
|
|
|
|
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.
|
|
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() -> MiniMaxH3Transformer3DModel:
|
|
config = build_tiny_h3_config()
|
|
model = MiniMaxH3Transformer3DModel(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_layout():
|
|
return build_packed_layout(
|
|
NUM_TEXT_TOKENS,
|
|
NUM_LATENT_FRAMES,
|
|
LATENT_HEIGHT,
|
|
LATENT_WIDTH,
|
|
NUM_AUDIO_LATENTS,
|
|
patch_size=tuple(TINY_ARCH["patch_size"]),
|
|
)
|
|
|
|
|
|
def build_inputs():
|
|
"""Random rows + the packed layout, in both torch and numpy form."""
|
|
layout = build_layout()
|
|
generator = torch.Generator(device="cpu").manual_seed(SEED + 1)
|
|
channels = TINY_ARCH["in_channels"]
|
|
latents = torch.randn(
|
|
1,
|
|
channels,
|
|
NUM_LATENT_FRAMES,
|
|
LATENT_HEIGHT,
|
|
LATENT_WIDTH,
|
|
generator=generator,
|
|
dtype=torch.float32,
|
|
)
|
|
video_rows_np = patchify_video_latents(latents.numpy(), tuple(TINY_ARCH["patch_size"]))
|
|
patch_dim = video_rows_np.shape[-1]
|
|
audio_rows_np = np.random.default_rng(SEED + 2).standard_normal(
|
|
(NUM_AUDIO_LATENTS * 2, TINY_ARCH["audio_in_channels"])).astype(np.float32)
|
|
text_rows_np = np.random.default_rng(SEED + 5).standard_normal(
|
|
(NUM_TEXT_TOKENS, TINY_ARCH["text_dim"])).astype(np.float32)
|
|
|
|
unique, inverse = build_row_timesteps(layout, VIDEO_TIMESTEP, AUDIO_TIMESTEP)
|
|
|
|
torch_inputs = dict(
|
|
hidden_states=torch.from_numpy(video_rows_np)[None],
|
|
audio_hidden_states=torch.from_numpy(audio_rows_np)[None],
|
|
encoder_hidden_states=torch.from_numpy(text_rows_np)[None],
|
|
timestep=torch.from_numpy(unique),
|
|
timestep_indices=torch.from_numpy(inverse),
|
|
token_tags=torch.from_numpy(layout.token_tags),
|
|
position_ids=torch.from_numpy(layout.position_ids),
|
|
video_indices=torch.from_numpy(layout.video_indices),
|
|
audio_indices=torch.from_numpy(layout.audio_indices),
|
|
text_indices=torch.from_numpy(layout.text_indices),
|
|
)
|
|
mlx_inputs = dict(
|
|
video_rows=video_rows_np,
|
|
audio_rows=audio_rows_np,
|
|
text_rows=text_rows_np,
|
|
timesteps=unique,
|
|
timestep_indices=inverse,
|
|
)
|
|
return layout, torch_inputs, mlx_inputs
|
|
|
|
|
|
# torch (FastVideo module names) -> MLX (released checkpoint / diffusers names,
|
|
# per MiniMaxH3ArchConfig.param_names_mapping read in reverse).
|
|
def torch_key_to_mlx(name: str) -> str:
|
|
if name.startswith("time_embedder.fc_in."):
|
|
return name.replace("time_embedder.fc_in.", "time_embedder.linear_1.", 1)
|
|
if name.startswith("time_embedder.fc_out."):
|
|
return name.replace("time_embedder.fc_out.", "time_embedder.linear_2.", 1)
|
|
name = name.replace(".attn.to_out.", ".attn.to_out.0.")
|
|
name = name.replace(".ff.fc_in.", ".ff.net.0.proj.")
|
|
name = name.replace(".ff.fc_out.", ".ff.net.2.")
|
|
return name
|
|
|
|
|
|
def mlx_dit_from_torch_model(
|
|
model: MiniMaxH3Transformer3DModel,
|
|
hf_config: dict[str, object],
|
|
*,
|
|
quantization=None,
|
|
) -> MLXMiniMaxH3DiT:
|
|
import mlx.core as mx
|
|
|
|
from fastvideo.mlx_runtime.minimax_h3 import _is_quantizable, quantize_matrix
|
|
|
|
weights: dict[str, object] = {}
|
|
blocks: list[dict[str, object]] = [{} for _ in range(TINY_ARCH["num_layers"])]
|
|
refiner: list[dict[str, object]] = [{} for _ in range(TINY_ARCH["num_refiner_layers"])]
|
|
|
|
for torch_name, param in model.state_dict().items():
|
|
name = torch_key_to_mlx(torch_name)
|
|
if name.startswith("rope."):
|
|
continue # non-persistent analytic buffer, rebuilt on the fly
|
|
array = mx.array(param.detach().float().numpy())
|
|
if quantization is not None and _is_quantizable(name):
|
|
value = quantize_matrix(array, quantization)
|
|
else:
|
|
value = array
|
|
if name.startswith("transformer_blocks."):
|
|
_, index_str, sub = name.split(".", 2)
|
|
blocks[int(index_str)][sub] = value
|
|
elif name.startswith("token_refiner.refiner_blocks."):
|
|
_, _, index_str, sub = name.split(".", 3)
|
|
refiner[int(index_str)][sub] = value
|
|
else:
|
|
weights[name] = value
|
|
return MLXMiniMaxH3DiT(weights, blocks, refiner, dict(hf_config))
|
|
|
|
|
|
def torch_reference_output(model, torch_inputs) -> tuple[np.ndarray, np.ndarray]:
|
|
with torch.no_grad(), set_forward_context(
|
|
current_timestep=0,
|
|
attn_metadata=None,
|
|
forward_batch=ForwardBatch(data_type="dummy"),
|
|
):
|
|
video_output, audio_output = model(**torch_inputs)
|
|
return video_output[0].float().numpy(), audio_output[0].float().numpy()
|
|
|
|
|
|
def mlx_output(dit: MLXMiniMaxH3DiT, layout, mlx_inputs) -> tuple[np.ndarray, np.ndarray]:
|
|
import mlx.core as mx
|
|
|
|
video_output, audio_output = dit.forward(
|
|
mx.array(mlx_inputs["video_rows"]),
|
|
mx.array(mlx_inputs["audio_rows"]),
|
|
mx.array(mlx_inputs["text_rows"]),
|
|
position_ids=mx.array(layout.position_ids.astype(np.float32)),
|
|
token_tags=mx.array(layout.token_tags),
|
|
timestep_indices=mx.array(mlx_inputs["timestep_indices"]),
|
|
timesteps=mx.array(mlx_inputs["timesteps"]),
|
|
video_indices=mx.array(layout.video_indices),
|
|
audio_indices=mx.array(layout.audio_indices),
|
|
text_indices=mx.array(layout.text_indices),
|
|
)
|
|
mx.eval(video_output, audio_output)
|
|
return np.asarray(video_output), np.asarray(audio_output)
|
|
|
|
|
|
def mlx_cache_output(dit: MLXMiniMaxH3DiT, layout, mlx_inputs) -> tuple[np.ndarray, np.ndarray]:
|
|
import mlx.core as mx
|
|
|
|
union = np.unique(
|
|
np.concatenate([
|
|
mlx_inputs["timesteps"],
|
|
np.array([1.0], dtype=np.float32), # condition rows, if any
|
|
]))
|
|
dit.precompute_adaln(union, drop_weights=True)
|
|
video_output, audio_output = dit.forward_with_cache(
|
|
mx.array(mlx_inputs["video_rows"]),
|
|
mx.array(mlx_inputs["audio_rows"]),
|
|
mx.array(mlx_inputs["text_rows"]),
|
|
layout=layout,
|
|
step_timesteps=mlx_inputs["timesteps"],
|
|
row_timestep_inverse=mlx_inputs["timestep_indices"],
|
|
)
|
|
mx.eval(video_output, audio_output)
|
|
return np.asarray(video_output), np.asarray(audio_output)
|
|
|
|
|
|
def build_schedulers(num_denoise_steps: int = 6) -> tuple[MiniMaxH3SchedulerState, MiniMaxH3SchedulerState]:
|
|
return (
|
|
MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_denoise_steps),
|
|
MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_denoise_steps),
|
|
)
|