Add MiniMax-H3 Turbo 4-step adapter and inference

6786963 verified about 1 month ago

20.3 kB

#!/usr/bin/env python3
"""Diffusers demo for MiniMax-H3 FL2VA (T2VA and FL2VA generation)."""

from __future__ import annotations

import argparse
import json
from contextlib import nullcontext
from pathlib import Path
from types import MethodType

import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download, snapshot_download
from peft import LoraConfig
from peft.utils import set_peft_model_state_dict
from PIL import Image
from safetensors import safe_open
from torch import nn

from diffusers import ComponentsManager, ModularPipeline
from diffusers.models.modeling_utils import get_parameter_dtype
from diffusers.models.transformers.transformer_minimax_h3 import (
    MiniMaxH3Transformer3DModel,
    MiniMaxH3TransformerOutput,
)
from diffusers.utils import encode_video

OFFICIAL_REPO = "MiniMaxAI/MiniMax-H3"
TURBO_REPO = "Wuli-art/MiniMax-H3-Turbo"
FOUR_STEP_FILENAME = "Wuli-MiniMax-H3-Turbo-4steps-V1.0-bf16.safetensors"
FOUR_STEP_FORMAT = "wuli_minimax_h3_turbo_v1"
BASE_COMPONENTS = (
    "text_encoder", "tokenizer", "processor", "vae", "audio_vae",
    "scheduler", "audio_scheduler",
)


# Four-step compact adapter

def resolve_four_step_checkpoint(model_path: str | None, repo_id: str) -> Path:
    if model_path is None:
        return Path(hf_hub_download(repo_id, FOUR_STEP_FILENAME, subfolder="turbo-4step"))
    local = Path(model_path).expanduser()
    for candidate in (local, local / FOUR_STEP_FILENAME, local / "turbo-4step" / FOUR_STEP_FILENAME):
        if candidate.is_file():
            return candidate
    raise FileNotFoundError(
        f"Could not find {FOUR_STEP_FILENAME} in {local} or {local / 'turbo-4step'}"
    )


def _shifted_sigma(sigma: torch.Tensor, shift: float) -> torch.Tensor:
    return shift * sigma / (1 + (shift - 1) * sigma)


class _FusedProjection(nn.Module):
    def __init__(self, weights: list[torch.Tensor], biases: list[torch.Tensor]) -> None:
        super().__init__()
        self.weights = nn.Parameter(torch.stack(weights), requires_grad=False)
        self.biases = nn.Parameter(torch.stack(biases), requires_grad=False)
        self.step = 0

    @property
    def weight(self) -> torch.Tensor:
        return self.weights[self.step]

    @property
    def bias(self) -> torch.Tensor:
        return self.biases[self.step]

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return F.linear(hidden_states, self.weight, self.bias)


class _FourStepController:
    def __init__(self, transformer: nn.Module, adapter_names: list[str]) -> None:
        self.transformer = transformer
        self.adapter_names = adapter_names
        self.index = 0
        self.reset()

    def _select(self, index: int) -> None:
        self.transformer.proj_out.step = index
        self.transformer.audio_proj_out.step = index
        self.transformer.set_adapter(self.adapter_names[index])

    def reset(self) -> None:
        self.index = 0
        self._select(0)

    def __call__(self, _module, _args, output):
        self.index = (self.index + 1) % len(self.adapter_names)
        self._select(self.index)
        return output


def _install_four_step_lora(
    transformer: nn.Module, state: dict[str, torch.Tensor], saved: dict
) -> None:
    name = "generator"
    transformer.add_adapter(
        LoraConfig(
            r=int(saved["r"]),
            lora_alpha=float(saved.get("lora_alpha") or saved["r"]),
            lora_dropout=float(saved.get("lora_dropout") or 0),
            target_modules=list(saved["target_modules"]),
        ),
        adapter_name=name,
    )
    incompatible = set_peft_model_state_dict(transformer, state, adapter_name=name)
    missing = [
        key for key in (getattr(incompatible, "missing_keys", None) or [])
        if "lora_" in key and f".{name}." in key
    ]
    unexpected = getattr(incompatible, "unexpected_keys", None) or []
    if missing or unexpected:
        raise ValueError(f"Incomplete four-step LoRA: missing={missing[:4]}, unexpected={unexpected[:4]}")


def apply_four_step_adapter(transformer: nn.Module, checkpoint: str | Path) -> dict:
    with safe_open(checkpoint, framework="pt", device="cpu") as handle:
        metadata = handle.metadata() or {}
        if "pdd_config" not in metadata or "lora_adapter_metadata" not in metadata:
            raise ValueError("The four-step checkpoint is missing required metadata")
        config = json.loads(metadata["pdd_config"])
        adapter_config = json.loads(metadata["lora_adapter_metadata"])
        keys = list(handle.keys())
        lora_state = {
            key.removeprefix("lora.generator."): handle.get_tensor(key)
            for key in keys if key.startswith("lora.generator.")
        }
        heads = {
            key.removeprefix("pdd."): handle.get_tensor(key)
            for key in keys if key.startswith("pdd.")
        }

    if config.get("format") != FOUR_STEP_FORMAT or int(config.get("nfe", 0)) != 4:
        raise ValueError("Expected a Wuli MiniMax-H3 Turbo V1 four-step checkpoint")
    grid_size = int(config["grid_size"])
    if grid_size != 256:
        raise ValueError(f"The four-step V1 checkpoint requires grid_size=256, got {grid_size}")
    adapter_names = list(config["step_adapter_names"])
    if adapter_names != ["generator"] * 4:
        raise ValueError("The four-step V1 checkpoint requires the shared generator adapter")
    _install_four_step_lora(transformer, lora_state, adapter_config)

    block_size = grid_size // 4
    base_sigma = torch.linspace(1, 0, grid_size + 1, dtype=torch.float32)
    for modality, attribute, shift_key in (
        ("video", "proj_out", "video_flow_shift"),
        ("audio", "audio_proj_out", "audio_flow_shift"),
    ):
        sigma = _shifted_sigma(base_sigma, float(config[shift_key]))
        spans = [sigma[i * block_size] - sigma[(i + 1) * block_size] for i in range(4)]
        weights = [heads[f"{modality}.{i}.weight"] / spans[i] for i in range(4)]
        biases = [heads[f"{modality}.{i}.bias"] / spans[i] for i in range(4)]
        source = getattr(transformer, attribute)
        if weights[0].shape != source.weight.shape or biases[0].shape != source.bias.shape:
            raise ValueError(f"{modality} fused head shape does not match MiniMax-H3")
        setattr(transformer, attribute, _FusedProjection(weights, biases))

    controller = _FourStepController(transformer, adapter_names)
    transformer.register_forward_hook(controller)
    transformer._four_step_controller = controller
    transformer.requires_grad_(False).eval()
    return config


def reset_four_step_adapter(transformer: nn.Module) -> None:
    controller = getattr(transformer, "_four_step_controller", None)
    if controller is None:
        raise RuntimeError("The four-step adapter is not installed")
    controller.reset()


def validate_four_step_call(transformer: nn.Module) -> None:
    controller = getattr(transformer, "_four_step_controller", None)
    if controller is None or controller.index != 0:
        raise RuntimeError("The pipeline did not execute exactly four transformer calls")


def ensure_processor_compat(processor) -> None:
    if hasattr(processor, "create_mm_token_type_ids"):
        return

    def create_mm_token_type_ids(instance, batches):
        return [[
            1 if token == instance.image_token_id else
            2 if token == instance.video_token_id else 0
            for token in row
        ] for row in batches]

    processor.create_mm_token_type_ids = MethodType(create_mm_token_type_ids, processor)


class _BlockTimeEmbedding(nn.Module):
    def __init__(self, rows: int, modalities: int, hidden_size: int) -> None:
        super().__init__()
        self.register_buffer(
            "weight",
            torch.empty(rows, modalities, 6, hidden_size, dtype=torch.bfloat16),
        )


class _OutputTimeEmbedding(nn.Module):
    def __init__(self, rows: int, hidden_size: int) -> None:
        super().__init__()
        self.register_buffer(
            "weight", torch.empty(rows, 2, hidden_size, dtype=torch.bfloat16)
        )


class FixedMiniMaxH3Transformer3DModel(MiniMaxH3Transformer3DModel):
    """Official Diffusers H3 trunk with fixed video/audio time embeddings."""

    _supports_gradient_checkpointing = False
    _keep_in_fp32_modules = ["proj_in", "audio_proj_in", "proj_out", "audio_proj_out", "rope"]

    def __init__(
        self,
        num_attention_heads: int = 56,
        attention_head_dim: int = 128,
        hidden_size: int = 5376,
        num_layers: int = 50,
        num_refiner_layers: int = 2,
        ffn_dim: int = 14336,
        in_channels: int = 24,
        audio_in_channels: int = 32,
        patch_size: tuple[int, int, int] = (1, 2, 2),
        text_dim: int = 5120,
        freq_dim: int = 256,
        time_embed_hidden_dim: int = 5376,
        time_embed_dim: int = 2688,
        rope_freq_dim: int = 16,
        rope_theta: float = 10000.0,
        norm_eps: float = 1e-5,
        qk_norm_eps: float = 1e-5,
        final_norm_eps: float = 1e-5,
        fixed_num_steps: int = 8,
    ) -> None:
        super().__init__(
            num_attention_heads=num_attention_heads,
            attention_head_dim=attention_head_dim,
            hidden_size=hidden_size,
            num_layers=num_layers,
            num_refiner_layers=num_refiner_layers,
            ffn_dim=ffn_dim,
            in_channels=in_channels,
            audio_in_channels=audio_in_channels,
            patch_size=patch_size,
            text_dim=text_dim,
            freq_dim=freq_dim,
            time_embed_hidden_dim=time_embed_hidden_dim,
            time_embed_dim=time_embed_dim,
            rope_freq_dim=rope_freq_dim,
            rope_theta=rope_theta,
            norm_eps=norm_eps,
            qk_norm_eps=qk_norm_eps,
            final_norm_eps=final_norm_eps,
        )
        fixed_num_steps = int(fixed_num_steps)
        self.register_to_config(fixed_num_steps=fixed_num_steps)
        hidden_size = int(self.config.hidden_size)
        self.video_time_embedding = nn.ModuleList(
            [_BlockTimeEmbedding(fixed_num_steps + 2, 2, hidden_size) for _ in self.transformer_blocks]
        )
        self.audio_time_embedding = nn.ModuleList(
            [_BlockTimeEmbedding(fixed_num_steps, 1, hidden_size) for _ in self.transformer_blocks]
        )
        self.video_output_time_embedding = _OutputTimeEmbedding(fixed_num_steps + 2, hidden_size)
        self.audio_output_time_embedding = _OutputTimeEmbedding(fixed_num_steps, hidden_size)
        self.register_buffer("video_timesteps", torch.empty(fixed_num_steps, dtype=torch.float32))
        self.register_buffer("audio_timesteps", torch.empty(fixed_num_steps, dtype=torch.float32))
        self.register_buffer("visual_condition_timestep", torch.tensor(0.999, dtype=torch.float32))

        self.time_proj = None
        self.time_embedder = None
        for block in self.transformer_blocks:
            block.adaln_proj = None
        self.norm_out.linear = None

    @staticmethod
    def _step(value: torch.Tensor, schedule: torch.Tensor) -> int:
        error = (schedule - value.float()).abs()
        index = int(error.argmin())
        if float(error[index]) > 2e-6:
            raise ValueError(
                f"unsupported timestep {float(value):.9f}; "
                f"use the checkpoint's fixed {schedule.numel()}-step schedule"
            )
        return index

    def _row_plan(self, timestep, timestep_indices, token_tags):
        row_t = timestep.index_select(0, timestep_indices).float()
        is_video = token_tags == 0
        is_text = token_tags == 1
        is_audio = token_tags == 2
        is_condition = is_video & torch.isclose(
            row_t, self.visual_condition_timestep, rtol=0.0, atol=2e-6
        )
        probe = row_t[is_text]
        if probe.numel() == 0:
            probe = row_t[is_video & ~is_condition]
        if probe.numel() == 0:
            raise ValueError("cannot resolve the current video step")
        step = self._step(probe[0], self.video_timesteps)
        if is_audio.any() and self._step(row_t[is_audio][0], self.audio_timesteps) != step:
            raise ValueError("video and audio schedules are on different steps")
        return step, is_video, is_text, is_audio, is_condition

    def forward(
        self,
        hidden_states: torch.Tensor,
        audio_hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        timestep: torch.Tensor,
        timestep_indices: torch.Tensor,
        token_tags: torch.Tensor,
        position_ids: torch.Tensor,
        video_indices: torch.Tensor,
        audio_indices: torch.Tensor,
        text_indices: torch.Tensor,
        attention_kwargs: dict | None = None,
        return_dict: bool = True,
    ):
        del attention_kwargs
        rotary_emb = self.rope(position_ids)
        video = self.proj_in(hidden_states.to(get_parameter_dtype(self.proj_in)))
        audio = self.audio_proj_in(audio_hidden_states.to(get_parameter_dtype(self.audio_proj_in)))
        text = self.context_embedder(encoder_hidden_states.to(get_parameter_dtype(self.context_embedder)))
        text = self.token_refiner(text)
        packed = text.new_zeros((text.shape[0], position_ids.shape[0], text.shape[-1]))
        packed = packed.index_copy(1, text_indices, text)
        packed = packed.index_copy(1, video_indices, video.to(text.dtype))
        packed = packed.index_copy(1, audio_indices, audio.to(text.dtype))

        step, is_video, is_text, is_audio, is_condition = self._row_plan(
            timestep, timestep_indices, token_tags
        )
        condition_row = self.config.fixed_num_steps + int(step == self.config.fixed_num_steps - 1)
        for layer, block in enumerate(self.transformer_blocks):
            table = packed.new_empty((token_tags.numel(), 6, packed.shape[-1]))
            table[is_video & ~is_condition] = self.video_time_embedding[layer].weight[step, 0]
            table[is_condition] = self.video_time_embedding[layer].weight[condition_row, 0]
            table[is_text] = self.video_time_embedding[layer].weight[step, 1]
            table[is_audio] = self.audio_time_embedding[layer].weight[step, 0]
            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = table.unbind(dim=1)

            residual = packed
            value = block.norm1(packed)
            value = value * (1.0 + scale_msa[None]) + shift_msa[None]
            packed = residual + gate_msa[None] * block.attn(value, rotary_emb)
            residual = packed
            value = block.norm2(packed)
            value = value * (1.0 + scale_mlp[None]) + shift_mlp[None]
            packed = residual + gate_mlp[None] * block.ff(value)

        output_table = packed.new_empty((token_tags.numel(), 2, packed.shape[-1]))
        output_table[is_video & ~is_condition] = self.video_output_time_embedding.weight[step]
        output_table[is_condition] = self.video_output_time_embedding.weight[condition_row]
        output_table[is_text] = self.video_output_time_embedding.weight[step]
        output_table[is_audio] = self.audio_output_time_embedding.weight[step]
        shift, scale = output_table.unbind(dim=1)
        packed = self.norm_out.norm(packed) * (1.0 + scale[None]) + shift[None]
        packed = packed.to(get_parameter_dtype(self.proj_out))
        video = self.proj_out(packed).index_select(1, video_indices)
        audio = self.audio_proj_out(packed).index_select(1, audio_indices)
        if not return_dict:
            return video, audio
        return MiniMaxH3TransformerOutput(sample=video, audio_sample=audio)


def load_pipeline(
    model_path: str | None = None,
    workflow: str = "auto",
    components_path: str = OFFICIAL_REPO,
    steps: int = 8,
) -> ModularPipeline:
    del workflow
    if steps not in (4, 6, 8):
        raise ValueError("steps must be 4, 6, or 8")
    component_names = list(BASE_COMPONENTS)
    if steps == 4:
        component_names.append("transformer")
    base_root = Path(components_path)
    if not base_root.is_dir():
        base_root = Path(snapshot_download(
            repo_id=components_path,
            allow_patterns=["model_index.json", "modular_model_index.json"]
            + [f"{name}/*" for name in component_names],
        ))

    manager = ComponentsManager()
    manager.enable_auto_cpu_offload(device="cuda", memory_reserve_margin="12GB")
    pipe = ModularPipeline.from_pretrained(base_root, components_manager=manager)
    pipe.load_components(
        names=component_names,
        dtype=torch.bfloat16,
        pretrained_model_name_or_path=str(base_root),
    )

    if steps == 4:
        checkpoint = resolve_four_step_checkpoint(model_path, TURBO_REPO)
        config = apply_four_step_adapter(pipe.transformer, checkpoint)
        ensure_processor_compat(pipe.processor)
        pipe.scheduler.set_shift(float(config["video_flow_shift"]))
        pipe.audio_scheduler.set_shift(float(config["audio_flow_shift"]))
        return pipe

    if model_path is None:
        variant = f"turbo-{steps}step"
        model_root = Path(snapshot_download(
            repo_id=TURBO_REPO,
            allow_patterns=[f"{variant}/transformer/*"],
        )) / variant
    else:
        model_root = Path(model_path)
    transformer_path = model_root / "transformer"
    transformer = FixedMiniMaxH3Transformer3DModel.from_pretrained(
        transformer_path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True
    )
    index = json.loads(
        (transformer_path / "diffusion_pytorch_model.safetensors.index.json").read_text()
    )["weight_map"]
    for name in ("video_timesteps", "audio_timesteps", "visual_condition_timestep"):
        with safe_open(transformer_path / index[name], framework="pt", device="cpu") as handle:
            setattr(transformer, name, handle.get_tensor(name))
    pipe.update_components(transformer=transformer)
    return pipe


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--steps", type=int, choices=(4, 6, 8), default=8)
    parser.add_argument("--model", help="Local weight directory instead of the private Hub repo")
    parser.add_argument("--components", default=OFFICIAL_REPO)
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--image", help="First frame; enables FL2VA")
    parser.add_argument("--last-image", help="Optional last frame; enables FL2VA")
    parser.add_argument("--output", default="output.mp4")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--num-frames", type=int, default=124)
    parser.add_argument("--height", type=int)
    parser.add_argument("--width", type=int)
    parser.add_argument(
        "--disable-cudnn",
        action="store_true",
        help="Disable cuDNN only when its optional runtime libraries cannot be loaded",
    )
    args = parser.parse_args()
    if (args.height is None) != (args.width is None):
        parser.error("--height and --width must be provided together")

    workflow = "fl2va" if args.image or args.last_image else "t2va"
    pipe = load_pipeline(args.model, workflow, args.components, args.steps)
    request = {
        "prompt": args.prompt,
        "num_inference_steps": args.steps + 1,
        "num_frames": args.num_frames,
        "generator": torch.Generator(device="cuda").manual_seed(args.seed),
        "output_type": "np",
    }
    if args.height is not None:
        request.update(height=args.height, width=args.width)
    if args.image:
        request["image"] = Image.open(args.image).convert("RGB")
    if args.last_image:
        request["last_image"] = Image.open(args.last_image).convert("RGB")
    if args.steps == 4:
        reset_four_step_adapter(pipe.transformer)
    if args.disable_cudnn:
        torch.backends.cuda.enable_cudnn_sdp(False)
    cudnn_context = torch.backends.cudnn.flags(enabled=False) if args.disable_cudnn else nullcontext()
    with cudnn_context:
        result = pipe(**request)
    if args.steps == 4:
        validate_four_step_call(pipe.transformer)
    encode_video(
        result.videos[0], fps=24, output_path=args.output,
        audio=result.audio[0], audio_sample_rate=result.sampling_rate,
    )


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Diffusers demo for MiniMax-H3 FL2VA (T2VA and FL2VA generation)."""

	from __future__ import annotations

	import argparse
	import json
	from contextlib import nullcontext
	from pathlib import Path
	from types import MethodType

	import torch
	import torch.nn.functional as F
	from huggingface_hub import hf_hub_download, snapshot_download
	from peft import LoraConfig
	from peft.utils import set_peft_model_state_dict
	from PIL import Image
	from safetensors import safe_open
	from torch import nn

	from diffusers import ComponentsManager, ModularPipeline
	from diffusers.models.modeling_utils import get_parameter_dtype
	from diffusers.models.transformers.transformer_minimax_h3 import (
	MiniMaxH3Transformer3DModel,
	MiniMaxH3TransformerOutput,
	)
	from diffusers.utils import encode_video

	OFFICIAL_REPO = "MiniMaxAI/MiniMax-H3"
	TURBO_REPO = "Wuli-art/MiniMax-H3-Turbo"
	FOUR_STEP_FILENAME = "Wuli-MiniMax-H3-Turbo-4steps-V1.0-bf16.safetensors"
	FOUR_STEP_FORMAT = "wuli_minimax_h3_turbo_v1"
	BASE_COMPONENTS = (
	"text_encoder", "tokenizer", "processor", "vae", "audio_vae",
	"scheduler", "audio_scheduler",
	)


	# Four-step compact adapter

	def resolve_four_step_checkpoint(model_path: str \| None, repo_id: str) -> Path:
	if model_path is None:
	return Path(hf_hub_download(repo_id, FOUR_STEP_FILENAME, subfolder="turbo-4step"))
	local = Path(model_path).expanduser()
	for candidate in (local, local / FOUR_STEP_FILENAME, local / "turbo-4step" / FOUR_STEP_FILENAME):
	if candidate.is_file():
	return candidate
	raise FileNotFoundError(
	f"Could not find {FOUR_STEP_FILENAME} in {local} or {local / 'turbo-4step'}"
	)


	def _shifted_sigma(sigma: torch.Tensor, shift: float) -> torch.Tensor:
	return shift * sigma / (1 + (shift - 1) * sigma)


	class _FusedProjection(nn.Module):
	def __init__(self, weights: list[torch.Tensor], biases: list[torch.Tensor]) -> None:
	super().__init__()
	self.weights = nn.Parameter(torch.stack(weights), requires_grad=False)
	self.biases = nn.Parameter(torch.stack(biases), requires_grad=False)
	self.step = 0

	@property
	def weight(self) -> torch.Tensor:
	return self.weights[self.step]

	@property
	def bias(self) -> torch.Tensor:
	return self.biases[self.step]

	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	return F.linear(hidden_states, self.weight, self.bias)


	class _FourStepController:
	def __init__(self, transformer: nn.Module, adapter_names: list[str]) -> None:
	self.transformer = transformer
	self.adapter_names = adapter_names
	self.index = 0
	self.reset()

	def _select(self, index: int) -> None:
	self.transformer.proj_out.step = index
	self.transformer.audio_proj_out.step = index
	self.transformer.set_adapter(self.adapter_names[index])

	def reset(self) -> None:
	self.index = 0
	self._select(0)

	def __call__(self, _module, _args, output):
	self.index = (self.index + 1) % len(self.adapter_names)
	self._select(self.index)
	return output


	def _install_four_step_lora(
	transformer: nn.Module, state: dict[str, torch.Tensor], saved: dict
	) -> None:
	name = "generator"
	transformer.add_adapter(
	LoraConfig(
	r=int(saved["r"]),
	lora_alpha=float(saved.get("lora_alpha") or saved["r"]),
	lora_dropout=float(saved.get("lora_dropout") or 0),
	target_modules=list(saved["target_modules"]),
	),
	adapter_name=name,
	)
	incompatible = set_peft_model_state_dict(transformer, state, adapter_name=name)
	missing = [
	key for key in (getattr(incompatible, "missing_keys", None) or [])
	if "lora_" in key and f".{name}." in key
	]
	unexpected = getattr(incompatible, "unexpected_keys", None) or []
	if missing or unexpected:
	raise ValueError(f"Incomplete four-step LoRA: missing={missing[:4]}, unexpected={unexpected[:4]}")


	def apply_four_step_adapter(transformer: nn.Module, checkpoint: str \| Path) -> dict:
	with safe_open(checkpoint, framework="pt", device="cpu") as handle:
	metadata = handle.metadata() or {}
	if "pdd_config" not in metadata or "lora_adapter_metadata" not in metadata:
	raise ValueError("The four-step checkpoint is missing required metadata")
	config = json.loads(metadata["pdd_config"])
	adapter_config = json.loads(metadata["lora_adapter_metadata"])
	keys = list(handle.keys())
	lora_state = {
	key.removeprefix("lora.generator."): handle.get_tensor(key)
	for key in keys if key.startswith("lora.generator.")
	}
	heads = {
	key.removeprefix("pdd."): handle.get_tensor(key)
	for key in keys if key.startswith("pdd.")
	}

	if config.get("format") != FOUR_STEP_FORMAT or int(config.get("nfe", 0)) != 4:
	raise ValueError("Expected a Wuli MiniMax-H3 Turbo V1 four-step checkpoint")
	grid_size = int(config["grid_size"])
	if grid_size != 256:
	raise ValueError(f"The four-step V1 checkpoint requires grid_size=256, got {grid_size}")
	adapter_names = list(config["step_adapter_names"])
	if adapter_names != ["generator"] * 4:
	raise ValueError("The four-step V1 checkpoint requires the shared generator adapter")
	_install_four_step_lora(transformer, lora_state, adapter_config)

	block_size = grid_size // 4
	base_sigma = torch.linspace(1, 0, grid_size + 1, dtype=torch.float32)
	for modality, attribute, shift_key in (
	("video", "proj_out", "video_flow_shift"),
	("audio", "audio_proj_out", "audio_flow_shift"),
	):
	sigma = _shifted_sigma(base_sigma, float(config[shift_key]))
	spans = [sigma[i * block_size] - sigma[(i + 1) * block_size] for i in range(4)]
	weights = [heads[f"{modality}.{i}.weight"] / spans[i] for i in range(4)]
	biases = [heads[f"{modality}.{i}.bias"] / spans[i] for i in range(4)]
	source = getattr(transformer, attribute)
	if weights[0].shape != source.weight.shape or biases[0].shape != source.bias.shape:
	raise ValueError(f"{modality} fused head shape does not match MiniMax-H3")
	setattr(transformer, attribute, _FusedProjection(weights, biases))

	controller = _FourStepController(transformer, adapter_names)
	transformer.register_forward_hook(controller)
	transformer._four_step_controller = controller
	transformer.requires_grad_(False).eval()
	return config


	def reset_four_step_adapter(transformer: nn.Module) -> None:
	controller = getattr(transformer, "_four_step_controller", None)
	if controller is None:
	raise RuntimeError("The four-step adapter is not installed")
	controller.reset()


	def validate_four_step_call(transformer: nn.Module) -> None:
	controller = getattr(transformer, "_four_step_controller", None)
	if controller is None or controller.index != 0:
	raise RuntimeError("The pipeline did not execute exactly four transformer calls")


	def ensure_processor_compat(processor) -> None:
	if hasattr(processor, "create_mm_token_type_ids"):
	return

	def create_mm_token_type_ids(instance, batches):
	return [[
	1 if token == instance.image_token_id else
	2 if token == instance.video_token_id else 0
	for token in row
	] for row in batches]

	processor.create_mm_token_type_ids = MethodType(create_mm_token_type_ids, processor)


	class _BlockTimeEmbedding(nn.Module):
	def __init__(self, rows: int, modalities: int, hidden_size: int) -> None:
	super().__init__()
	self.register_buffer(
	"weight",
	torch.empty(rows, modalities, 6, hidden_size, dtype=torch.bfloat16),
	)


	class _OutputTimeEmbedding(nn.Module):
	def __init__(self, rows: int, hidden_size: int) -> None:
	super().__init__()
	self.register_buffer(
	"weight", torch.empty(rows, 2, hidden_size, dtype=torch.bfloat16)
	)


	class FixedMiniMaxH3Transformer3DModel(MiniMaxH3Transformer3DModel):
	"""Official Diffusers H3 trunk with fixed video/audio time embeddings."""

	_supports_gradient_checkpointing = False
	_keep_in_fp32_modules = ["proj_in", "audio_proj_in", "proj_out", "audio_proj_out", "rope"]

	def __init__(
	self,
	num_attention_heads: int = 56,
	attention_head_dim: int = 128,
	hidden_size: int = 5376,
	num_layers: int = 50,
	num_refiner_layers: int = 2,
	ffn_dim: int = 14336,
	in_channels: int = 24,
	audio_in_channels: int = 32,
	patch_size: tuple[int, int, int] = (1, 2, 2),
	text_dim: int = 5120,
	freq_dim: int = 256,
	time_embed_hidden_dim: int = 5376,
	time_embed_dim: int = 2688,
	rope_freq_dim: int = 16,
	rope_theta: float = 10000.0,
	norm_eps: float = 1e-5,
	qk_norm_eps: float = 1e-5,
	final_norm_eps: float = 1e-5,
	fixed_num_steps: int = 8,
	) -> None:
	super().__init__(
	num_attention_heads=num_attention_heads,
	attention_head_dim=attention_head_dim,
	hidden_size=hidden_size,
	num_layers=num_layers,
	num_refiner_layers=num_refiner_layers,
	ffn_dim=ffn_dim,
	in_channels=in_channels,
	audio_in_channels=audio_in_channels,
	patch_size=patch_size,
	text_dim=text_dim,
	freq_dim=freq_dim,
	time_embed_hidden_dim=time_embed_hidden_dim,
	time_embed_dim=time_embed_dim,
	rope_freq_dim=rope_freq_dim,
	rope_theta=rope_theta,
	norm_eps=norm_eps,
	qk_norm_eps=qk_norm_eps,
	final_norm_eps=final_norm_eps,
	)
	fixed_num_steps = int(fixed_num_steps)
	self.register_to_config(fixed_num_steps=fixed_num_steps)
	hidden_size = int(self.config.hidden_size)
	self.video_time_embedding = nn.ModuleList(
	[_BlockTimeEmbedding(fixed_num_steps + 2, 2, hidden_size) for _ in self.transformer_blocks]
	)
	self.audio_time_embedding = nn.ModuleList(
	[_BlockTimeEmbedding(fixed_num_steps, 1, hidden_size) for _ in self.transformer_blocks]
	)
	self.video_output_time_embedding = _OutputTimeEmbedding(fixed_num_steps + 2, hidden_size)
	self.audio_output_time_embedding = _OutputTimeEmbedding(fixed_num_steps, hidden_size)
	self.register_buffer("video_timesteps", torch.empty(fixed_num_steps, dtype=torch.float32))
	self.register_buffer("audio_timesteps", torch.empty(fixed_num_steps, dtype=torch.float32))
	self.register_buffer("visual_condition_timestep", torch.tensor(0.999, dtype=torch.float32))

	self.time_proj = None
	self.time_embedder = None
	for block in self.transformer_blocks:
	block.adaln_proj = None
	self.norm_out.linear = None

	@staticmethod
	def _step(value: torch.Tensor, schedule: torch.Tensor) -> int:
	error = (schedule - value.float()).abs()
	index = int(error.argmin())
	if float(error[index]) > 2e-6:
	raise ValueError(
	f"unsupported timestep {float(value):.9f}; "
	f"use the checkpoint's fixed {schedule.numel()}-step schedule"
	)
	return index

	def _row_plan(self, timestep, timestep_indices, token_tags):
	row_t = timestep.index_select(0, timestep_indices).float()
	is_video = token_tags == 0
	is_text = token_tags == 1
	is_audio = token_tags == 2
	is_condition = is_video & torch.isclose(
	row_t, self.visual_condition_timestep, rtol=0.0, atol=2e-6
	)
	probe = row_t[is_text]
	if probe.numel() == 0:
	probe = row_t[is_video & ~is_condition]
	if probe.numel() == 0:
	raise ValueError("cannot resolve the current video step")
	step = self._step(probe[0], self.video_timesteps)
	if is_audio.any() and self._step(row_t[is_audio][0], self.audio_timesteps) != step:
	raise ValueError("video and audio schedules are on different steps")
	return step, is_video, is_text, is_audio, is_condition

	def forward(
	self,
	hidden_states: torch.Tensor,
	audio_hidden_states: torch.Tensor,
	encoder_hidden_states: torch.Tensor,
	timestep: torch.Tensor,
	timestep_indices: torch.Tensor,
	token_tags: torch.Tensor,
	position_ids: torch.Tensor,
	video_indices: torch.Tensor,
	audio_indices: torch.Tensor,
	text_indices: torch.Tensor,
	attention_kwargs: dict \| None = None,
	return_dict: bool = True,
	):
	del attention_kwargs
	rotary_emb = self.rope(position_ids)
	video = self.proj_in(hidden_states.to(get_parameter_dtype(self.proj_in)))
	audio = self.audio_proj_in(audio_hidden_states.to(get_parameter_dtype(self.audio_proj_in)))
	text = self.context_embedder(encoder_hidden_states.to(get_parameter_dtype(self.context_embedder)))
	text = self.token_refiner(text)
	packed = text.new_zeros((text.shape[0], position_ids.shape[0], text.shape[-1]))
	packed = packed.index_copy(1, text_indices, text)
	packed = packed.index_copy(1, video_indices, video.to(text.dtype))
	packed = packed.index_copy(1, audio_indices, audio.to(text.dtype))

	step, is_video, is_text, is_audio, is_condition = self._row_plan(
	timestep, timestep_indices, token_tags
	)
	condition_row = self.config.fixed_num_steps + int(step == self.config.fixed_num_steps - 1)
	for layer, block in enumerate(self.transformer_blocks):
	table = packed.new_empty((token_tags.numel(), 6, packed.shape[-1]))
	table[is_video & ~is_condition] = self.video_time_embedding[layer].weight[step, 0]
	table[is_condition] = self.video_time_embedding[layer].weight[condition_row, 0]
	table[is_text] = self.video_time_embedding[layer].weight[step, 1]
	table[is_audio] = self.audio_time_embedding[layer].weight[step, 0]
	shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = table.unbind(dim=1)

	residual = packed
	value = block.norm1(packed)
	value = value * (1.0 + scale_msa[None]) + shift_msa[None]
	packed = residual + gate_msa[None] * block.attn(value, rotary_emb)
	residual = packed
	value = block.norm2(packed)
	value = value * (1.0 + scale_mlp[None]) + shift_mlp[None]
	packed = residual + gate_mlp[None] * block.ff(value)

	output_table = packed.new_empty((token_tags.numel(), 2, packed.shape[-1]))
	output_table[is_video & ~is_condition] = self.video_output_time_embedding.weight[step]
	output_table[is_condition] = self.video_output_time_embedding.weight[condition_row]
	output_table[is_text] = self.video_output_time_embedding.weight[step]
	output_table[is_audio] = self.audio_output_time_embedding.weight[step]
	shift, scale = output_table.unbind(dim=1)
	packed = self.norm_out.norm(packed) * (1.0 + scale[None]) + shift[None]
	packed = packed.to(get_parameter_dtype(self.proj_out))
	video = self.proj_out(packed).index_select(1, video_indices)
	audio = self.audio_proj_out(packed).index_select(1, audio_indices)
	if not return_dict:
	return video, audio
	return MiniMaxH3TransformerOutput(sample=video, audio_sample=audio)


	def load_pipeline(
	model_path: str \| None = None,
	workflow: str = "auto",
	components_path: str = OFFICIAL_REPO,
	steps: int = 8,
	) -> ModularPipeline:
	del workflow
	if steps not in (4, 6, 8):
	raise ValueError("steps must be 4, 6, or 8")
	component_names = list(BASE_COMPONENTS)
	if steps == 4:
	component_names.append("transformer")
	base_root = Path(components_path)
	if not base_root.is_dir():
	base_root = Path(snapshot_download(
	repo_id=components_path,
	allow_patterns=["model_index.json", "modular_model_index.json"]
	+ [f"{name}/*" for name in component_names],
	))

	manager = ComponentsManager()
	manager.enable_auto_cpu_offload(device="cuda", memory_reserve_margin="12GB")
	pipe = ModularPipeline.from_pretrained(base_root, components_manager=manager)
	pipe.load_components(
	names=component_names,
	dtype=torch.bfloat16,
	pretrained_model_name_or_path=str(base_root),
	)

	if steps == 4:
	checkpoint = resolve_four_step_checkpoint(model_path, TURBO_REPO)
	config = apply_four_step_adapter(pipe.transformer, checkpoint)
	ensure_processor_compat(pipe.processor)
	pipe.scheduler.set_shift(float(config["video_flow_shift"]))
	pipe.audio_scheduler.set_shift(float(config["audio_flow_shift"]))
	return pipe

	if model_path is None:
	variant = f"turbo-{steps}step"
	model_root = Path(snapshot_download(
	repo_id=TURBO_REPO,
	allow_patterns=[f"{variant}/transformer/*"],
	)) / variant
	else:
	model_root = Path(model_path)
	transformer_path = model_root / "transformer"
	transformer = FixedMiniMaxH3Transformer3DModel.from_pretrained(
	transformer_path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True
	)
	index = json.loads(
	(transformer_path / "diffusion_pytorch_model.safetensors.index.json").read_text()
	)["weight_map"]
	for name in ("video_timesteps", "audio_timesteps", "visual_condition_timestep"):
	with safe_open(transformer_path / index[name], framework="pt", device="cpu") as handle:
	setattr(transformer, name, handle.get_tensor(name))
	pipe.update_components(transformer=transformer)
	return pipe


	def main() -> None:
	parser = argparse.ArgumentParser()
	parser.add_argument("--steps", type=int, choices=(4, 6, 8), default=8)
	parser.add_argument("--model", help="Local weight directory instead of the private Hub repo")
	parser.add_argument("--components", default=OFFICIAL_REPO)
	parser.add_argument("--prompt", required=True)
	parser.add_argument("--image", help="First frame; enables FL2VA")
	parser.add_argument("--last-image", help="Optional last frame; enables FL2VA")
	parser.add_argument("--output", default="output.mp4")
	parser.add_argument("--seed", type=int, default=42)
	parser.add_argument("--num-frames", type=int, default=124)
	parser.add_argument("--height", type=int)
	parser.add_argument("--width", type=int)
	parser.add_argument(
	"--disable-cudnn",
	action="store_true",
	help="Disable cuDNN only when its optional runtime libraries cannot be loaded",
	)
	args = parser.parse_args()
	if (args.height is None) != (args.width is None):
	parser.error("--height and --width must be provided together")

	workflow = "fl2va" if args.image or args.last_image else "t2va"
	pipe = load_pipeline(args.model, workflow, args.components, args.steps)
	request = {
	"prompt": args.prompt,
	"num_inference_steps": args.steps + 1,
	"num_frames": args.num_frames,
	"generator": torch.Generator(device="cuda").manual_seed(args.seed),
	"output_type": "np",
	}
	if args.height is not None:
	request.update(height=args.height, width=args.width)
	if args.image:
	request["image"] = Image.open(args.image).convert("RGB")
	if args.last_image:
	request["last_image"] = Image.open(args.last_image).convert("RGB")
	if args.steps == 4:
	reset_four_step_adapter(pipe.transformer)
	if args.disable_cudnn:
	torch.backends.cuda.enable_cudnn_sdp(False)
	cudnn_context = torch.backends.cudnn.flags(enabled=False) if args.disable_cudnn else nullcontext()
	with cudnn_context:
	result = pipe(**request)
	if args.steps == 4:
	validate_four_step_call(pipe.transformer)
	encode_video(
	result.videos[0], fps=24, output_path=args.output,
	audio=result.audio[0], audio_sample_rate=result.sampling_rate,
	)


	if __name__ == "__main__":
	main()