#!/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()