From e8b0e4c61e58aeff33b0e331fc21a75ee99c49ce Mon Sep 17 00:00:00 2001 From: Kaiqin Kong Date: Tue, 4 Aug 2026 13:54:39 -0700 Subject: [PATCH] [feat] Add MiniMax H3 (#1674) --- .../inference_schema_parity_inventory.yaml | 5 +- docs/inference/support_matrix.md | 5 + .../inference/basic/basic_minimax_h3_fl2va.py | 91 ++ .../basic/basic_minimax_h3_ref2va.py | 97 ++ .../inference/basic/basic_minimax_h3_t2v.py | 83 ++ fastvideo/api/sampling_param.py | 3 + fastvideo/api/schema.py | 4 + fastvideo/configs/models/dits/__init__.py | 3 +- fastvideo/configs/models/dits/minimax_h3.py | 75 ++ fastvideo/configs/models/encoders/__init__.py | 4 +- .../models/encoders/minimax_h3_qwen3_vl.py | 175 ++++ fastvideo/configs/models/vaes/__init__.py | 9 + .../configs/models/vaes/minimax_h3_audio.py | 51 ++ .../configs/models/vaes/minimax_h3_video.py | 71 ++ fastvideo/configs/pipelines/minimax_h3.py | 38 + fastvideo/models/dits/minimax_h3.py | 661 ++++++++++++++ .../models/encoders/minimax_h3_qwen3_vl.py | 704 +++++++++++++++ fastvideo/models/loader/component_loader.py | 36 +- .../schedulers/scheduling_minimax_h3.py | 147 +++ fastvideo/models/vaes/minimax_h3_audio.py | 530 +++++++++++ fastvideo/models/vaes/minimax_h3_video.py | 847 ++++++++++++++++++ .../pipelines/basic/minimax_h3/__init__.py | 15 + .../basic/minimax_h3/minimax_h3_pipeline.py | 125 +++ .../pipelines/basic/minimax_h3/packing.py | 492 ++++++++++ .../pipelines/basic/minimax_h3/presets.py | 74 ++ .../pipelines/basic/minimax_h3/reference.py | 381 ++++++++ .../basic/minimax_h3/stages/__init__.py | 19 + .../stages/minimax_h3_conditioning.py | 329 +++++++ .../minimax_h3/stages/minimax_h3_decoding.py | 150 ++++ .../minimax_h3/stages/minimax_h3_denoising.py | 149 +++ .../stages/minimax_h3_input_preparation.py | 238 +++++ .../stages/minimax_h3_latent_preparation.py | 340 +++++++ fastvideo/pipelines/pipeline_batch_info.py | 3 + fastvideo/registry.py | 22 + fastvideo/tests/api/test_parser.py | 4 + .../tests/ssim/test_minimax_h3_similarity.py | 67 ++ fastvideo/utils.py | 65 +- .../test_minimax_h3_qwen3_vl_parity.py | 240 +++++ tests/local_tests/minimax_h3/PORT_STATUS.md | 41 + tests/local_tests/minimax_h3/README.md | 55 ++ tests/local_tests/minimax_h3/_reference.py | 42 + .../minimax_h3/test_minimax_h3_packing.py | 85 ++ .../test_minimax_h3_ref2va_media.py | 237 +++++ .../test_minimax_h3_ref2va_packing.py | 181 ++++ .../test_minimax_h3_scheduler_parity.py | 59 ++ .../test_minimax_h3_pipeline_smoke.py | 95 ++ .../test_minimax_h3_transformer_parity.py | 244 +++++ .../vaes/test_minimax_h3_audio_vae_parity.py | 174 ++++ .../vaes/test_minimax_h3_video_vae_parity.py | 175 ++++ 49 files changed, 7707 insertions(+), 33 deletions(-) create mode 100644 examples/inference/basic/basic_minimax_h3_fl2va.py create mode 100644 examples/inference/basic/basic_minimax_h3_ref2va.py create mode 100644 examples/inference/basic/basic_minimax_h3_t2v.py create mode 100644 fastvideo/configs/models/dits/minimax_h3.py create mode 100644 fastvideo/configs/models/encoders/minimax_h3_qwen3_vl.py create mode 100644 fastvideo/configs/models/vaes/minimax_h3_audio.py create mode 100644 fastvideo/configs/models/vaes/minimax_h3_video.py create mode 100644 fastvideo/configs/pipelines/minimax_h3.py create mode 100644 fastvideo/models/dits/minimax_h3.py create mode 100644 fastvideo/models/encoders/minimax_h3_qwen3_vl.py create mode 100644 fastvideo/models/schedulers/scheduling_minimax_h3.py create mode 100644 fastvideo/models/vaes/minimax_h3_audio.py create mode 100644 fastvideo/models/vaes/minimax_h3_video.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/__init__.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/packing.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/presets.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/reference.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/stages/__init__.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py create mode 100644 fastvideo/tests/ssim/test_minimax_h3_similarity.py create mode 100644 tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py create mode 100644 tests/local_tests/minimax_h3/PORT_STATUS.md create mode 100644 tests/local_tests/minimax_h3/README.md create mode 100644 tests/local_tests/minimax_h3/_reference.py create mode 100644 tests/local_tests/minimax_h3/test_minimax_h3_packing.py create mode 100644 tests/local_tests/minimax_h3/test_minimax_h3_ref2va_media.py create mode 100644 tests/local_tests/minimax_h3/test_minimax_h3_ref2va_packing.py create mode 100644 tests/local_tests/minimax_h3/test_minimax_h3_scheduler_parity.py create mode 100644 tests/local_tests/pipelines/test_minimax_h3_pipeline_smoke.py create mode 100644 tests/local_tests/transformers/test_minimax_h3_transformer_parity.py create mode 100644 tests/local_tests/vaes/test_minimax_h3_audio_vae_parity.py create mode 100644 tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index d836df5ab..eda1577fc 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -445,7 +445,11 @@ surfaces: moved: image_path: request.inputs.image_path pil_image: request.inputs.pil_image + last_image: request.inputs.last_image + references: request.inputs.references video_path: request.inputs.video_path + latents: request.inputs.latents + audio_latents: request.inputs.audio_latents mouse_cond: request.inputs.mouse_cond keyboard_cond: request.inputs.keyboard_cond grid_sizes: request.inputs.grid_sizes @@ -525,7 +529,6 @@ surfaces: inpaint_mask: request.extensions.stable_audio.inpaint_mask internal_only: data_type: "Derived from the request shape and not a public input." - latents: "Pre-generated diffusion latents supplied by parity/debug harnesses; not a public input." sampling_param_extensions: {} diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 36c7b7331..39f664a79 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -66,6 +66,7 @@ column links a runnable script in `examples/inference/basic/` where one exists. | ltx2 | `Lightricks/LTX-2`
`FastVideo/LTX2-base`
`FastVideo/LTX2-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) | | matrixgame | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers`
`FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers`
`FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers`
`FastVideo/Matrix-Game-2.0-Base-Diffusers`
`FastVideo/Matrix-Game-2.0-GTA-Diffusers`
`FastVideo/Matrix-Game-2.0-TempleRun-Diffusers`
`mignonjia/mg_longtuning_distilled_zelda`
`mignonjia/mg_sf_distilled_zelda_1k_steps`
`mignonjia/mg_sf_distilled_zelda`
`mignonjia/mg_causal_zelda`
`mignonjia/mg_bidirectional_zelda` | I2V | [basic_matrixgame2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame2.py) | | matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) | +| minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)
[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)
[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) | | sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) | | stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) | | stable_audio | `FastVideo/stable-audio-open-small-Diffusers` | T2V | [basic_stable_audio_small.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio_small.py) | @@ -93,6 +94,10 @@ column links a runnable script in `examples/inference/basic/` where one exists. (`StableAudioT2AConfig` / `StableAudioOpenSmallConfig`); they are registered under the generic T2V workload option in the registry. +**Note (MiniMax H3)**: T2VA, FL2VA, and Ref2VA all generate video with stereo +audio. Use the Ref2VA example when passing ordered image, video, or audio +references. + **Note (Wan-VACE)**: not currently supported — no VACE pipeline or registered model ID exists on `main` ([#1435](https://github.com/hao-ai-lab/FastVideo/issues/1435)). The closest diff --git a/examples/inference/basic/basic_minimax_h3_fl2va.py b/examples/inference/basic/basic_minimax_h3_fl2va.py new file mode 100644 index 000000000..25c15d4d8 --- /dev/null +++ b/examples/inference/basic/basic_minimax_h3_fl2va.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generate synchronized video/audio from a first frame with MiniMax H3.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from PIL import Image + +from fastvideo import VideoGenerator +from fastvideo.api import ( + EngineConfig, + GenerationRequest, + GeneratorConfig, + InputConfig, + OffloadConfig, + OutputConfig, + ParallelismConfig, + SamplingConfig, +) +from fastvideo.pipelines.basic.minimax_h3.packing import resolve_canvas_size + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3") + parser.add_argument("--image", required=True, help="First-frame image path.") + parser.add_argument("--last-image", help="Optional last-frame image path.") + parser.add_argument("--output", default="outputs/minimax_h3_fl2va") + parser.add_argument("--prompt", required=True) + parser.add_argument("--num-frames", type=int, default=192, help="192 frames is exactly 8 seconds at 24 fps.") + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--num-gpus", type=int, default=4) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + first_image = Image.open(args.image).convert("RGB") + last_image = Image.open(args.last_image).convert("RGB") if args.last_image else None + height, width = resolve_canvas_size(*first_image.size) + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path=args.model_path, + engine=EngineConfig( + num_gpus=args.num_gpus, + use_fsdp_inference=args.num_gpus > 1, + parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus), + offload=OffloadConfig( + dit=False, + dit_layerwise=False, + text_encoder=True, + vae=True, + pin_cpu_memory=False, + ), + ), + )) + try: + result = generator.generate( + GenerationRequest( + prompt=args.prompt, + negative_prompt="", + inputs=InputConfig(pil_image=first_image, last_image=last_image), + sampling=SamplingConfig( + height=height, + width=width, + num_frames=args.num_frames, + fps=24, + num_inference_steps=args.steps, + guidance_scale=1.0, + batch_cfg=False, + seed=args.seed, + ), + output=OutputConfig( + output_path=str(output_dir / "minimax_h3_fl2va.mp4"), + save_video=True, + return_frames=False, + ), + )) + print(f"Output written to: {result.video_path}") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_minimax_h3_ref2va.py b/examples/inference/basic/basic_minimax_h3_ref2va.py new file mode 100644 index 000000000..1786dddb0 --- /dev/null +++ b/examples/inference/basic/basic_minimax_h3_ref2va.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generate synchronized video/audio from ordered references with MiniMax H3.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from fastvideo import VideoGenerator +from fastvideo.api import ( + ComponentConfig, + EngineConfig, + GenerationRequest, + GeneratorConfig, + InputConfig, + OffloadConfig, + OutputConfig, + ParallelismConfig, + PipelineSelection, + SamplingConfig, +) +from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3") + parser.add_argument("--reference-video", required=True) + parser.add_argument("--reference-audio", help="Optional additional audio reference.") + parser.add_argument("--output", default="outputs/minimax_h3_ref2va") + parser.add_argument("--prompt", required=True) + parser.add_argument("--height", type=int, default=768) + parser.add_argument("--width", type=int, default=1344) + parser.add_argument("--num-frames", type=int, default=124) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--num-gpus", type=int, default=4) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + references = [MiniMaxH3Reference(source=args.reference_video, media_type="video")] + if args.reference_audio: + references.append(MiniMaxH3Reference(source=args.reference_audio, media_type="audio")) + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path=args.model_path, + engine=EngineConfig( + num_gpus=args.num_gpus, + use_fsdp_inference=args.num_gpus > 1, + parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus), + offload=OffloadConfig( + dit=False, + dit_layerwise=False, + text_encoder=True, + vae=True, + pin_cpu_memory=False, + ), + ), + pipeline=PipelineSelection( + workload_type="i2v", + components=ComponentConfig(override_pipeline_cls_name="MiniMaxH3Ref2VAModularPipeline"), + ), + )) + try: + result = generator.generate( + GenerationRequest( + prompt=args.prompt, + negative_prompt="", + inputs=InputConfig(references=references), + sampling=SamplingConfig( + height=args.height, + width=args.width, + num_frames=args.num_frames, + fps=24, + num_inference_steps=args.steps, + guidance_scale=1.0, + batch_cfg=False, + seed=args.seed, + ), + output=OutputConfig( + output_path=str(output_dir / "minimax_h3_ref2va.mp4"), + save_video=True, + return_frames=False, + ), + )) + print(f"Output written to: {result.video_path}") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_minimax_h3_t2v.py b/examples/inference/basic/basic_minimax_h3_t2v.py new file mode 100644 index 000000000..1b6bd1749 --- /dev/null +++ b/examples/inference/basic/basic_minimax_h3_t2v.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generate video and audio from text with MiniMax H3.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from fastvideo import VideoGenerator +from fastvideo.api import ( + EngineConfig, + GenerationRequest, + GeneratorConfig, + OffloadConfig, + OutputConfig, + ParallelismConfig, + SamplingConfig, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3") + parser.add_argument("--prompt", required=True) + parser.add_argument("--output", default="outputs/minimax_h3_t2v") + parser.add_argument("--height", type=int, default=768) + parser.add_argument("--width", type=int, default=1344) + parser.add_argument("--num-frames", type=int, default=124) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--num-gpus", type=int, default=4) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path=args.model_path, + engine=EngineConfig( + num_gpus=args.num_gpus, + use_fsdp_inference=args.num_gpus > 1, + parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus), + offload=OffloadConfig( + dit=False, + dit_layerwise=False, + text_encoder=True, + vae=True, + pin_cpu_memory=False, + ), + ), + )) + try: + result = generator.generate( + GenerationRequest( + prompt=args.prompt, + negative_prompt="", + sampling=SamplingConfig( + height=args.height, + width=args.width, + num_frames=args.num_frames, + fps=24, + num_inference_steps=args.steps, + guidance_scale=1.0, + batch_cfg=False, + seed=args.seed, + ), + output=OutputConfig( + output_path=str(output_dir / "minimax_h3_t2v.mp4"), + save_video=True, + return_frames=False, + ), + )) + print(f"Output written to: {result.video_path}") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index d5d9d28b0..464d7f0fa 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -25,6 +25,8 @@ class SamplingParam: # Image inputs image_path: str | None = None pil_image: Any | None = None + last_image: Any | None = None + references: list[Any] | None = None # Video inputs video_path: str | None = None @@ -32,6 +34,7 @@ class SamplingParam: # Optional pre-generated diffusion latents. Used by parity/debug harnesses # and advanced callers that need deterministic latent reuse. latents: Any | None = None + audio_latents: Any | None = None # Action control inputs (Matrix-Game) mouse_cond: Any | None = None # Shape: (B, T, 2) diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index 8eede2fdc..0dd3adfed 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -125,6 +125,10 @@ class InputConfig: image_path: str | list[str] | None = None video_path: str | list[str] | None = None pil_image: Any | None = None + last_image: Any | None = None + references: list[Any] | None = None + latents: Any | None = None + audio_latents: Any | None = None pose: str | None = None mouse_cond: Any | None = None keyboard_cond: Any | None = None diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 1c2f97a1a..9c31b3a1f 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -10,6 +10,7 @@ from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config from fastvideo.configs.models.dits.longcat import LongCatVideoConfig from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config from fastvideo.configs.models.dits.stable_audio import StableAudioConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig from fastvideo.configs.models.dits.zimage import ZImageDiTConfig @@ -23,5 +24,5 @@ __all__ = [ "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig", - "ZImageDiTConfig" + "MiniMaxH3Config", "ZImageDiTConfig" ] diff --git a/fastvideo/configs/models/dits/minimax_h3.py b/fastvideo/configs/models/dits/minimax_h3.py new file mode 100644 index 000000000..05df8e079 --- /dev/null +++ b/fastvideo/configs/models/dits/minimax_h3.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Architecture configuration for the MiniMax H3 joint audio-video DiT.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.platforms import AttentionBackendEnum + + +def _is_minimax_h3_block(name: str, module: object) -> bool: + """Select the main and text-refiner transformer blocks for FSDP.""" + del module + parts = name.split(".") + return ((len(parts) == 2 and parts[0] == "transformer_blocks" and parts[1].isdigit()) + or (len(parts) == 3 and parts[:2] == ["token_refiner", "refiner_blocks"] and parts[2].isdigit())) + + +@dataclass +class MiniMaxH3ArchConfig(DiTArchConfig): + """One-to-one representation of the released transformer config.""" + + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_minimax_h3_block]) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.FLASH_ATTN, + ) + + param_names_mapping: dict = field( + default_factory=lambda: { + r"^time_embedder\.linear_1\.(.*)$": r"time_embedder.fc_in.\1", + r"^time_embedder\.linear_2\.(.*)$": r"time_embedder.fc_out.\1", + r"^(.*)\.attn\.to_out\.0\.(.*)$": r"\1.attn.to_out.\2", + r"^(.*)\.ff\.net\.0\.proj\.(.*)$": r"\1.ff.fc_in.\2", + r"^(.*)\.ff\.net\.2\.(.*)$": r"\1.ff.fc_out.\2", + }) + 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 + + def __post_init__(self) -> None: + super().__post_init__() + if len(self.patch_size) != 3: + raise ValueError(f"MiniMax H3 patch_size must have three axes, got {self.patch_size}.") + self.patch_size = (self.patch_size[0], self.patch_size[1], self.patch_size[2]) + self.num_channels_latents = self.in_channels + self.out_channels = self.in_channels + rotary_dim = 2 * 3 * self.rope_freq_dim + if rotary_dim > self.attention_head_dim or rotary_dim % 2: + raise ValueError(f"MiniMax H3 rotary width must be even and no larger than the head width; got " + f"rotary_dim={rotary_dim}, attention_head_dim={self.attention_head_dim}.") + + +@dataclass +class MiniMaxH3Config(DiTConfig): + """FastVideo component configuration for MiniMax H3 transformers.""" + + arch_config: MiniMaxH3ArchConfig = field(default_factory=MiniMaxH3ArchConfig) + prefix: str = "minimax_h3" diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index 26c43ee0c..cdedd9ca4 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -9,6 +9,7 @@ from fastvideo.configs.models.encoders.siglip import SiglipVisionConfig from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig +from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (MiniMaxH3Qwen3VLArchConfig, MiniMaxH3Qwen3VLConfig) from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig, @@ -20,5 +21,6 @@ __all__ = [ "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig", "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig", - "LingBotWorld2UMT5ArchConfig", "LingBotWorld2UMT5Config", "LingBotVideoQwen3VLTextConfig" + "LingBotWorld2UMT5ArchConfig", "LingBotWorld2UMT5Config", "LingBotVideoQwen3VLTextConfig", + "MiniMaxH3Qwen3VLArchConfig", "MiniMaxH3Qwen3VLConfig" ] diff --git a/fastvideo/configs/models/encoders/minimax_h3_qwen3_vl.py b/fastvideo/configs/models/encoders/minimax_h3_qwen3_vl.py new file mode 100644 index 000000000..7f962d6b2 --- /dev/null +++ b/fastvideo/configs/models/encoders/minimax_h3_qwen3_vl.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3-VL conditioner configuration for MiniMax H3.""" + +from dataclasses import dataclass, field +from typing import Any + +from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig + + +def _is_language_transformer_layer(name: str, module: Any) -> bool: + parts = name.split(".") + return len(parts) >= 2 and parts[-2] == "layers" and parts[-1].isdigit() + + +def _is_vision_transformer_layer(name: str, module: Any) -> bool: + parts = name.split(".") + return len(parts) >= 2 and parts[-2] == "blocks" and parts[-1].isdigit() + + +def _is_embeddings(name: str, module: Any) -> bool: + return name.endswith(("embed_tokens", "patch_embed")) + + +def _is_vision_merger(name: str, module: Any) -> bool: + parts = name.split(".") + return name.endswith("visual.merger") or (len(parts) >= 2 and parts[-2] == "deepstack_merger_list" + and parts[-1].isdigit()) + + +def _is_final_norm(name: str, module: Any) -> bool: + return name == "norm" or name.endswith("language_model.norm") + + +_VISION_CONFIG_MAPPING = { + "depth": "vision_depth", + "hidden_size": "vision_hidden_size", + "hidden_act": "vision_hidden_act", + "intermediate_size": "vision_intermediate_size", + "num_heads": "vision_num_heads", + "in_channels": "vision_in_channels", + "patch_size": "vision_patch_size", + "spatial_merge_size": "vision_spatial_merge_size", + "temporal_patch_size": "vision_temporal_patch_size", + "out_hidden_size": "vision_out_hidden_size", + "num_position_embeddings": "vision_num_position_embeddings", + "initializer_range": "vision_initializer_range", + "deepstack_visual_indexes": "vision_deepstack_visual_indexes", +} + +_OFFICIAL_ARCHITECTURES = { + "Qwen3VLForConditionalGeneration", + "Qwen3VLModel", +} + + +@dataclass +class MiniMaxH3Qwen3VLArchConfig(TextEncoderArchConfig): + """Architecture of the Qwen3-VL encoder used by MiniMax H3.""" + + architectures: list[str] = field(default_factory=lambda: ["MiniMaxH3Qwen3VLConditioner"]) + vocab_size: int = 151936 + hidden_size: int = 5120 + intermediate_size: int = 25600 + num_hidden_layers: int = 64 + num_attention_heads: int = 64 + num_key_value_heads: int = 8 + head_dim: int = 128 + hidden_act: str = "silu" + max_position_embeddings: int = 262144 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + attention_bias: bool = False + attention_dropout: float = 0.0 + rope_theta: float = 5000000.0 + rope_scaling: dict[str, Any] | None = field(default_factory=lambda: { + "mrope_interleaved": True, + "mrope_section": [24, 20, 20], + "rope_type": "default", + }) + mrope_interleaved: bool = True + mrope_section: tuple[int, int, int] = (24, 20, 20) + + bos_token_id: int = 151643 + eos_token_id: int = 151645 + pad_token_id: int | None = None + vision_start_token_id: int = 151652 + vision_end_token_id: int = 151653 + image_token_id: int = 151655 + video_token_id: int = 151656 + tie_word_embeddings: bool = False + + vision_depth: int = 27 + vision_hidden_size: int = 1152 + vision_hidden_act: str = "gelu_pytorch_tanh" + vision_intermediate_size: int = 4304 + vision_num_heads: int = 16 + vision_in_channels: int = 3 + vision_patch_size: int = 16 + vision_spatial_merge_size: int = 2 + vision_temporal_patch_size: int = 2 + vision_out_hidden_size: int = 5120 + vision_num_position_embeddings: int = 2304 + vision_initializer_range: float = 0.02 + vision_deepstack_visual_indexes: tuple[int, ...] = (8, 16, 24) + + output_hidden_states: bool = True + stacked_params_mapping: list[tuple[str, str, str | int]] = field(default_factory=list) + _fsdp_shard_conditions: list = field(default_factory=lambda: [ + _is_language_transformer_layer, + _is_vision_transformer_layer, + _is_embeddings, + _is_vision_merger, + _is_final_norm, + ]) + + def __post_init__(self) -> None: + rope_scaling = dict(self.rope_scaling or {}) + self.mrope_interleaved = bool(rope_scaling.get("mrope_interleaved", self.mrope_interleaved)) + if not self.mrope_interleaved: + raise ValueError("MiniMax H3 Qwen3-VL only supports interleaved multimodal RoPE.") + section = rope_scaling.get("mrope_section", self.mrope_section) + if not isinstance(section, list | tuple) or len(section) != 3: + raise ValueError("MiniMax H3 Qwen3-VL mRoPE requires exactly three section sizes.") + self.mrope_section = (int(section[0]), int(section[1]), int(section[2])) + if sum(self.mrope_section) * 2 != self.head_dim: + raise ValueError("MiniMax H3 Qwen3-VL mRoPE sections must cover exactly half of each attention head.") + rope_scaling["mrope_interleaved"] = self.mrope_interleaved + rope_scaling["mrope_section"] = list(self.mrope_section) + rope_scaling.setdefault("rope_type", "default") + self.rope_scaling = rope_scaling + + if self.vision_out_hidden_size != self.hidden_size: + raise ValueError("MiniMax H3 Qwen3-VL vision_out_hidden_size must match the language hidden_size " + "for visual-token and DeepStack injection.") + + # H3 builds its presentation in the pipeline stage and tokenizes each + # segment verbatim without adding tokenizer-owned presentation tokens. + self.tokenizer_kwargs = {"add_special_tokens": False} + + +@dataclass +class MiniMaxH3Qwen3VLConfig(TextEncoderConfig): + """FastVideo loader config for H3's base Qwen3-VL encoder.""" + + arch_config: TextEncoderArchConfig = field(default_factory=MiniMaxH3Qwen3VLArchConfig) + prefix: str = "minimax_h3_qwen3_vl" + is_chat_model: bool = False + + def update_model_arch(self, source_model_dict: dict[str, Any]) -> None: + flattened = dict(source_model_dict) + + text_config = flattened.pop("text_config", {}) + flattened.update(dict(text_config)) + + vision_config = flattened.pop("vision_config", {}) + for source_name, target_name in _VISION_CONFIG_MAPPING.items(): + if source_name in vision_config: + flattened[target_name] = vision_config[source_name] + + architectures = flattened.get("architectures") + if isinstance(architectures, str): + architectures = [architectures] + if isinstance(architectures, list | tuple) and any(architecture in _OFFICIAL_ARCHITECTURES + for architecture in architectures): + flattened["architectures"] = ["MiniMaxH3Qwen3VLConditioner"] + + section = flattened.get("mrope_section") + if isinstance(section, list): + flattened["mrope_section"] = tuple(section) + deepstack_indexes = flattened.get("vision_deepstack_visual_indexes") + if isinstance(deepstack_indexes, list): + flattened["vision_deepstack_visual_indexes"] = tuple(deepstack_indexes) + + super().update_model_arch(flattened) diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index 6f3ca9768..9c72bb731 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -6,6 +6,11 @@ from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig +from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEArchConfig, MiniMaxH3AudioVAEConfig +from fastvideo.configs.models.vaes.minimax_h3_video import ( + MiniMaxH3VideoVAEArchConfig, + MiniMaxH3VideoVAEConfig, +) from fastvideo.configs.models.vaes.oobleck import OobleckVAEArchConfig, OobleckVAEConfig from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig from fastvideo.configs.models.vaes.wanvae import WanVAEConfig @@ -19,6 +24,10 @@ __all__ = [ "Gen3CVAEConfig", "Hunyuan15VAEConfig", "LTX2VAEConfig", + "MiniMaxH3AudioVAEArchConfig", + "MiniMaxH3AudioVAEConfig", + "MiniMaxH3VideoVAEArchConfig", + "MiniMaxH3VideoVAEConfig", "OobleckVAEArchConfig", "OobleckVAEConfig", "Flux2VAEConfig", diff --git a/fastvideo/configs/models/vaes/minimax_h3_audio.py b/fastvideo/configs/models/vaes/minimax_h3_audio.py new file mode 100644 index 000000000..29723c556 --- /dev/null +++ b/fastvideo/configs/models/vaes/minimax_h3_audio.py @@ -0,0 +1,51 @@ +# Copyright 2025 MiniMax authors and HuggingFace Team +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class MiniMaxH3AudioVAEArchConfig(VAEArchConfig): + """Architecture of the MiniMax H3 waveform autoencoder.""" + + _class_name: str = "AutoencoderKLMiniMaxH3Audio" + architectures: list[str] = field(default_factory=lambda: ["AutoencoderKLMiniMaxH3Audio"]) + + encoder_dim: int = 64 + encoder_rates: tuple[int, ...] | list[int] = (2, 4, 4, 5, 5) + latent_dim: int = 2048 + latent_channels: int = 32 + num_attention_heads: int = 8 + + decoder_dim: int = 1024 + decoder_rates: tuple[int, ...] | list[int] = (5, 5, 2, 2, 2, 2, 2) + decoder_kernel_sizes: tuple[int, ...] | list[int] = (9, 9, 4, 4, 4, 4, 4) + resblock_kernel_sizes: tuple[int, ...] | list[int] = (3, 7, 11) + resblock_dilation_sizes: tuple[tuple[int, ...], ...] | list[list[int]] = ( + (1, 3, 5), + (1, 3, 5), + (1, 3, 5), + ) + + sampling_rate: int = 32000 + latents_mean: tuple[float, ...] | list[float] | None = None + latents_std: tuple[float, ...] | list[float] | None = None + + scaling_factor: float = 1.0 + temporal_compression_ratio: int = 1 + spatial_compression_ratio: int = 1 + + +@dataclass +class MiniMaxH3AudioVAEConfig(VAEConfig): + """FastVideo loader configuration for the full H3 audio VAE.""" + + arch_config: MiniMaxH3AudioVAEArchConfig = field(default_factory=MiniMaxH3AudioVAEArchConfig) + + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False diff --git a/fastvideo/configs/models/vaes/minimax_h3_video.py b/fastvideo/configs/models/vaes/minimax_h3_video.py new file mode 100644 index 000000000..f9002d7a2 --- /dev/null +++ b/fastvideo/configs/models/vaes/minimax_h3_video.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-H3 video VAE configuration.""" + +import math +from dataclasses import dataclass, field + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class MiniMaxH3VideoVAEArchConfig(VAEArchConfig): + """Architecture fields from ``AutoencoderKLMiniMaxH3``.""" + + _class_name: str = "AutoencoderKLMiniMaxH3" + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 24 + block_out_channels: tuple[int, ...] = (128, 256, 256, 512, 512, 1024) + layers_per_block: int = 2 + spatial_downsample_factors: tuple[int, ...] = (2, 2, 2, 2, 1, 1) + temporal_downsample_factors: tuple[int, ...] = (1, 2, 2, 1, 1, 1) + norm_num_groups: int = 32 + norm_eps: float = 1e-6 + spatial_padding_mode: str = "reflect" + + decoder_num_layers: int = 36 + decoder_num_attention_heads: int = 32 + decoder_attention_head_dim: int = 64 + decoder_num_register_tokens: int = 4 + decoder_ffn_mult: int = 4 + decoder_rope_theta: float = 100.0 + decoder_rope_dim_ratio: float = 0.75 + decoder_norm_eps: float = 1e-5 + + clip_length: int = 17 + token_drop: int = 3 + latents_mean: tuple[float, ...] = (0.0, ) * 24 + latents_std: tuple[float, ...] = (1.0, ) * 24 + + # MiniMax-H3 uses per-channel mean/std instead of a scalar latent scale. + scaling_factor: float = 1.0 + temporal_compression_ratio: int = 4 + spatial_compression_ratio: int = 16 + + def __post_init__(self) -> None: + num_blocks = len(self.block_out_channels) + if len(self.spatial_downsample_factors) != num_blocks: + raise ValueError("`spatial_downsample_factors` must have one entry per encoder block.") + if len(self.temporal_downsample_factors) != num_blocks: + raise ValueError("`temporal_downsample_factors` must have one entry per encoder block.") + if len(self.latents_mean) != self.latent_channels: + raise ValueError("`latents_mean` must have one entry per latent channel.") + if len(self.latents_std) != self.latent_channels: + raise ValueError("`latents_std` must have one entry per latent channel.") + + self.spatial_compression_ratio = math.prod(self.spatial_downsample_factors) + self.temporal_compression_ratio = math.prod(self.temporal_downsample_factors) + + rotary_dim = int(self.decoder_attention_head_dim * self.decoder_rope_dim_ratio) + if rotary_dim % 6 != 0: + raise ValueError("MiniMax-H3 decoder rotary dimensions must be divisible by 6.") + + +@dataclass +class MiniMaxH3VideoVAEConfig(VAEConfig): + """FastVideo runtime controls for the MiniMax-H3 video VAE.""" + + arch_config: MiniMaxH3VideoVAEArchConfig = field(default_factory=MiniMaxH3VideoVAEArchConfig) + + tile_sample_min_overlap_height: int = 64 + tile_sample_min_overlap_width: int = 64 diff --git a/fastvideo/configs/pipelines/minimax_h3.py b/fastvideo/configs/pipelines/minimax_h3.py new file mode 100644 index 000000000..ed033bae2 --- /dev/null +++ b/fastvideo/configs/pipelines/minimax_h3.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline configuration for MiniMax H3 joint video/audio generation.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models import EncoderConfig, VAEConfig +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config +from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConfig +from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEConfig +from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + + +@dataclass +class MiniMaxH3PipelineConfig(PipelineConfig): + """Component and precision policy shared by T2VA, FL2VA, and Ref2VA.""" + + dit_config: MiniMaxH3Config = field(default_factory=MiniMaxH3Config) + vae_config: VAEConfig = field(default_factory=MiniMaxH3VideoVAEConfig) + audio_vae_config: VAEConfig = field(default_factory=MiniMaxH3AudioVAEConfig) + text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (MiniMaxH3Qwen3VLConfig(), )) + + flow_shift: float | None = None + embedded_cfg_scale: float | None = None + dit_precision: str = "bf16" + vae_precision: str = "fp32" + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", )) + vae_sp: bool = False + + def check_pipeline_config(self) -> None: + super().check_pipeline_config() + if self.flow_shift is not None: + raise ValueError("MiniMax-H3 uses fixed video/audio scheduler shifts; flow_shift must remain unset.") + + +__all__ = ["MiniMaxH3PipelineConfig"] diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py new file mode 100644 index 000000000..e20877071 --- /dev/null +++ b/fastvideo/models/dits/minimax_h3.py @@ -0,0 +1,661 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo-native MiniMax H3 joint audio-video diffusion transformer.""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.attention import DistributedAttention +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config +from fastvideo.distributed.communication_op import ( + sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard, +) +from fastvideo.distributed.parallel_state import get_sp_world_size, model_parallel_is_initialized +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.quantization import QuantizationConfig +from fastvideo.layers.visual_embedding import Timesteps +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum + + +MINIMAX_H3_MODALITY_NUM = 3 +_CFG = MiniMaxH3Config() + + +class MiniMaxH3RotaryPosEmbed(nn.Module): + """Three-axis rotary frequencies over packed `(t, h, w)` coordinates.""" + + def __init__(self, rope_freq_dim: int, rope_theta: float) -> None: + super().__init__() + inv_freq = 1.0 / ( + rope_theta ** (torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32) / (2 * rope_freq_dim)) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + position_ids = position_ids.to(torch.float32) + freqs = position_ids.unsqueeze(-1) * self.inv_freq.view(1, 1, -1) + freqs_t, freqs_h, freqs_w = freqs.unbind(dim=1) + freqs = torch.cat((freqs_t, freqs_h, freqs_w), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +class MiniMaxH3TimestepEmbedding(nn.Module): + """Linear-SiLU-linear timestep MLP.""" + + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.fc_in = ReplicatedLinear( + input_dim, + hidden_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.fc_in", + ) + self.fc_out = ReplicatedLinear( + hidden_dim, + output_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.fc_out", + ) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.fc_in(sample) + hidden_states, _ = self.fc_out(F.silu(hidden_states)) + return hidden_states + + +class MiniMaxH3FeedForward(nn.Module): + """Bias-free H3 SwiGLU with value-first packed halves.""" + + def __init__( + self, + hidden_size: int, + ffn_dim: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.fc_in = ReplicatedLinear( + hidden_size, + 2 * ffn_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fc_in", + ) + self.fc_out = ReplicatedLinear( + ffn_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fc_out", + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.fc_in(hidden_states) + hidden_states, gate = hidden_states.chunk(2, dim=-1) + hidden_states = hidden_states * F.silu(gate) + hidden_states, _ = self.fc_out(hidden_states) + return hidden_states + + +class MiniMaxH3Attention(nn.Module): + """Full self-attention over one sequence-parallel packed document.""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + attention_head_dim: int, + qk_norm_eps: float, + supported_attention_backends: tuple[AttentionBackendEnum, ...], + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + self.to_q = ReplicatedLinear( + hidden_size, + inner_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_q", + ) + self.to_k = ReplicatedLinear( + hidden_size, + inner_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_k", + ) + self.to_v = ReplicatedLinear( + hidden_size, + inner_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_v", + ) + self.norm_q = nn.RMSNorm(attention_head_dim, eps=qk_norm_eps) + self.norm_k = nn.RMSNorm(attention_head_dim, eps=qk_norm_eps) + self.to_out = ReplicatedLinear( + inner_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_out", + ) + self.distributed_attention = DistributedAttention( + num_heads=num_attention_heads, + head_size=attention_head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=prefix, + ) + + @staticmethod + def _apply_rotary_emb( + hidden_states: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + """Rotate the RoPE prefix while preserving the rest of each attention head.""" + cos, sin = rotary_emb + rotary_dim = cos.shape[-1] + hidden_states_rotary = hidden_states[..., :rotary_dim] + hidden_states_pass = hidden_states[..., rotary_dim:] + cos = cos.to(hidden_states.dtype)[None, :, None, :] + sin = sin.to(hidden_states.dtype)[None, :, None, :] + first_half, second_half = hidden_states_rotary.chunk(2, dim=-1) + hidden_states_rotated = torch.cat((-second_half, first_half), dim=-1) + hidden_states_rotary = hidden_states_rotary * cos + hidden_states_rotated * sin + return torch.cat((hidden_states_rotary, hidden_states_pass), dim=-1).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None, + original_seq_len: int, + ) -> torch.Tensor: + query, _ = self.to_q(hidden_states) + key, _ = self.to_k(hidden_states) + value, _ = self.to_v(hidden_states) + query = query.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) + key = key.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) + value = value.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) + query = self.norm_q(query) + key = self.norm_k(key) + if rotary_emb is not None: + query = self._apply_rotary_emb(query, rotary_emb) + key = self._apply_rotary_emb(key, rotary_emb) + + # H3 rotates only 96/128 channels, which the generic `freqs_cis` + # branch cannot express. Apply it above, then pass no RoPE here. + hidden_states, _ = self.distributed_attention( + query, + key, + value, + original_seq_len=original_seq_len, + freqs_cis=None, + ) + hidden_states = hidden_states.flatten(2, 3).type_as(query) + hidden_states, _ = self.to_out(hidden_states) + return hidden_states + + +class MiniMaxH3TokenRefinerBlock(nn.Module): + """Plain pre-norm Transformer block for the projected text stream.""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + attention_head_dim: int, + ffn_dim: int, + norm_eps: float, + qk_norm_eps: float, + supported_attention_backends: tuple[AttentionBackendEnum, ...], + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps) + self.attn = MiniMaxH3Attention( + hidden_size, + num_attention_heads, + attention_head_dim, + qk_norm_eps, + supported_attention_backends, + quant_config, + prefix=f"{prefix}.attn", + ) + self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps) + self.ff = MiniMaxH3FeedForward( + hidden_size, + ffn_dim, + quant_config=quant_config, + prefix=f"{prefix}.ff", + ) + + def forward(self, hidden_states: torch.Tensor, original_seq_len: int) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states), None, original_seq_len) + return hidden_states + self.ff(self.norm2(hidden_states)) + + +class MiniMaxH3TokenRefiner(nn.Module): + """Two-block text refiner used before packing the modalities.""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + attention_head_dim: int, + ffn_dim: int, + num_layers: int, + norm_eps: float, + qk_norm_eps: float, + final_norm_eps: float, + supported_attention_backends: tuple[AttentionBackendEnum, ...], + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.refiner_blocks = nn.ModuleList([ + MiniMaxH3TokenRefinerBlock( + hidden_size, + num_attention_heads, + attention_head_dim, + ffn_dim, + norm_eps, + qk_norm_eps, + supported_attention_backends, + quant_config, + prefix=f"{prefix}.refiner_blocks.{index}", + ) for index in range(num_layers) + ]) + self.final_norm = nn.RMSNorm(hidden_size, eps=final_norm_eps) + + def forward(self, hidden_states: torch.Tensor, original_seq_len: int) -> torch.Tensor: + for block in self.refiner_blocks: + hidden_states = block(hidden_states, original_seq_len) + return self.final_norm(hidden_states) + + +class MiniMaxH3AdaLayerNormModulation(nn.Module): + """Produce six modulation tables for every `(timestep, modality)` pair.""" + + def __init__( + self, + time_embed_dim: int, + hidden_size: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.linear = ReplicatedLinear( + time_embed_dim, + 6 * hidden_size * MINIMAX_H3_MODALITY_NUM, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.linear", + ) + + def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]: + temb, _ = self.linear(F.silu(temb).to(self.linear.weight.dtype)) + return temb.view(-1, 6 * self.hidden_size).chunk(6, dim=-1) + + +class MiniMaxH3AdaLayerNormOut(nn.Module): + """Final RMSNorm with per-timestep row modulation.""" + + def __init__( + self, + hidden_size: int, + time_embed_dim: int, + eps: float, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.norm = nn.RMSNorm(hidden_size, eps=eps) + self.linear = ReplicatedLinear( + time_embed_dim, + 2 * hidden_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.linear", + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + timestep_indices: torch.Tensor, + ) -> torch.Tensor: + shift_scale, _ = self.linear(F.silu(temb).to(self.linear.weight.dtype)) + shift, scale = shift_scale.chunk(2, dim=-1) + hidden_states = self.norm(hidden_states) + return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + +class MiniMaxH3TransformerBlock(nn.Module): + """Packed self-attention and feed-forward branches with row-indexed AdaLN.""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + attention_head_dim: int, + ffn_dim: int, + time_embed_dim: int, + norm_eps: float, + qk_norm_eps: float, + supported_attention_backends: tuple[AttentionBackendEnum, ...], + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps) + self.attn = MiniMaxH3Attention( + hidden_size, + num_attention_heads, + attention_head_dim, + qk_norm_eps, + supported_attention_backends, + quant_config, + prefix=f"{prefix}.attn", + ) + self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps) + self.ff = MiniMaxH3FeedForward( + hidden_size, + ffn_dim, + quant_config=quant_config, + prefix=f"{prefix}.ff", + ) + self.adaln_proj = MiniMaxH3AdaLayerNormModulation( + time_embed_dim, + hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.adaln_proj", + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + adaln_indices: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor], + original_seq_len: int, + ) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb) + + residual = hidden_states + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * ( + 1.0 + scale_msa.index_select(0, adaln_indices) + ) + shift_msa.index_select(0, adaln_indices) + attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len) + hidden_states = residual + gate_msa.index_select(0, adaln_indices) * attention_output + + residual = hidden_states + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * ( + 1.0 + scale_mlp.index_select(0, adaln_indices) + ) + shift_mlp.index_select(0, adaln_indices) + feed_forward_output = self.ff(norm_hidden_states) + return residual + gate_mlp.index_select(0, adaln_indices) * feed_forward_output + + +class MiniMaxH3Transformer3DModel(BaseDiT): + """Joint H3 Transformer over one padless text/audio/video document. + + The layout builder validates semantic rows before denoising. Sequence- + parallel padding is transport-only and `DistributedAttention` trims it + before attention. + """ + + _fsdp_shard_conditions = _CFG._fsdp_shard_conditions + _compile_conditions = _CFG._compile_conditions + _supported_attention_backends = _CFG._supported_attention_backends + param_names_mapping = _CFG.param_names_mapping + reverse_param_names_mapping = _CFG.reverse_param_names_mapping + lora_param_names_mapping = _CFG.lora_param_names_mapping + _keep_in_fp32_modules = frozenset({ + "proj_in", + "audio_proj_in", + "time_embedder", + "proj_out", + "audio_proj_out", + "rope", + }) + + def _get_parameter_dtype(self, name: str, default_dtype: torch.dtype) -> torch.dtype: + """Keep the released input, timestep, and output projections in FP32.""" + return torch.float32 if name.split(".", 1)[0] in self._keep_in_fp32_modules else default_dtype + + def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None: + super().__init__(config, hf_config) + arch = config.arch_config + sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1 + if arch.num_attention_heads % sp_world_size: + raise ValueError( + f"MiniMax H3 attention heads ({arch.num_attention_heads}) must be divisible by " + f"sequence parallel size ({sp_world_size})." + ) + + self.hidden_size = arch.hidden_size + self.num_attention_heads = arch.num_attention_heads + self.num_channels_latents = arch.in_channels + self.patch_size = tuple(int(value) for value in arch.patch_size) + video_patch_dim = arch.in_channels * math.prod(arch.patch_size) + + self.proj_in = ReplicatedLinear( + video_patch_dim, + arch.hidden_size, + bias=True, + quant_config=config.quant_config, + prefix=f"{config.prefix}.proj_in", + ) + self.audio_proj_in = ReplicatedLinear( + arch.audio_in_channels, + arch.hidden_size, + bias=True, + quant_config=config.quant_config, + prefix=f"{config.prefix}.audio_proj_in", + ) + self.context_embedder = ReplicatedLinear( + arch.text_dim, + arch.hidden_size, + bias=True, + quant_config=config.quant_config, + prefix=f"{config.prefix}.context_embedder", + ) + self.time_proj = Timesteps( + num_channels=arch.freq_dim, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + self.time_embedder = MiniMaxH3TimestepEmbedding( + arch.freq_dim, + arch.time_embed_hidden_dim, + arch.time_embed_dim, + quant_config=config.quant_config, + prefix=f"{config.prefix}.time_embedder", + ) + self.rope = MiniMaxH3RotaryPosEmbed(arch.rope_freq_dim, arch.rope_theta) + self.token_refiner = MiniMaxH3TokenRefiner( + arch.hidden_size, + arch.num_attention_heads, + arch.attention_head_dim, + arch.ffn_dim, + arch.num_refiner_layers, + arch.norm_eps, + arch.qk_norm_eps, + arch.final_norm_eps, + self.supported_attention_backends, + config.quant_config, + prefix=f"{config.prefix}.token_refiner", + ) + self.transformer_blocks = nn.ModuleList([ + MiniMaxH3TransformerBlock( + arch.hidden_size, + arch.num_attention_heads, + arch.attention_head_dim, + arch.ffn_dim, + arch.time_embed_dim, + arch.norm_eps, + arch.qk_norm_eps, + self.supported_attention_backends, + config.quant_config, + prefix=f"{config.prefix}.transformer_blocks.{index}", + ) for index in range(arch.num_layers) + ]) + self.norm_out = MiniMaxH3AdaLayerNormOut( + arch.hidden_size, + arch.time_embed_dim, + arch.final_norm_eps, + quant_config=config.quant_config, + prefix=f"{config.prefix}.norm_out", + ) + self.proj_out = ReplicatedLinear( + arch.hidden_size, + video_patch_dim, + bias=True, + quant_config=config.quant_config, + prefix=f"{config.prefix}.proj_out", + ) + self.audio_proj_out = ReplicatedLinear( + arch.hidden_size, + arch.audio_in_channels, + bias=True, + quant_config=config.quant_config, + prefix=f"{config.prefix}.audio_proj_out", + ) + self.__post_init__() + + def materialize_non_persistent_buffers( + self, + device: torch.device, + dtype: torch.dtype | None = None, + ) -> None: + """Rebuild analytic RoPE state after meta-device checkpoint loading.""" + del dtype + if self.rope.inv_freq.is_meta: + arch = self.config.arch_config + inv_freq = 1.0 / ( + arch.rope_theta + ** ( + torch.arange(0, 2 * arch.rope_freq_dim, 2, device=device, dtype=torch.float32) + / (2 * arch.rope_freq_dim) + ) + ) + self.rope._buffers["inv_freq"] = inv_freq + + 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, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Predict video and audio velocities from one caller-defined packed layout.""" + if position_ids.ndim != 2 or position_ids.shape[-1] != 3: + raise ValueError(f"position_ids must have shape (seq_len, 3), got {tuple(position_ids.shape)}.") + sequence_length = position_ids.shape[0] + if token_tags.shape != (sequence_length,) or timestep_indices.shape != (sequence_length,): + raise ValueError("token_tags and timestep_indices must both match the packed sequence length.") + if hidden_states.shape[1] != video_indices.numel(): + raise ValueError("hidden_states row count must match video_indices.") + if audio_hidden_states.shape[1] != audio_indices.numel(): + raise ValueError("audio_hidden_states row count must match audio_indices.") + if encoder_hidden_states.shape[1] != text_indices.numel(): + raise ValueError("encoder_hidden_states row count must match text_indices.") + + rotary_emb = self.rope(position_ids) + video_embeds, _ = self.proj_in(hidden_states.to(self.proj_in.weight.dtype)) + audio_embeds, _ = self.audio_proj_in(audio_hidden_states.to(self.audio_proj_in.weight.dtype)) + text_embeds, _ = self.context_embedder(encoder_hidden_states.to(self.context_embedder.weight.dtype)) + + text_original_seq_len = text_embeds.shape[1] + sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1 + if sp_world_size > 1: + text_embeds, _ = sequence_model_parallel_shard(text_embeds, dim=1) + text_embeds = self.token_refiner(text_embeds, text_original_seq_len) + if sp_world_size > 1: + text_embeds = sequence_model_parallel_all_gather_with_unpad( + text_embeds, + text_original_seq_len, + dim=1, + ) + + packed_hidden_states = text_embeds.new_zeros( + (text_embeds.shape[0], sequence_length, text_embeds.shape[-1]) + ) + packed_hidden_states = packed_hidden_states.index_copy(1, text_indices, text_embeds) + packed_hidden_states = packed_hidden_states.index_copy(1, video_indices, video_embeds.to(text_embeds.dtype)) + packed_hidden_states = packed_hidden_states.index_copy(1, audio_indices, audio_embeds.to(text_embeds.dtype)) + + temb = self.time_proj(timestep) + temb = self.time_embedder(temb.to(self.time_embedder.fc_in.weight.dtype)) + adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags + local_timestep_indices = timestep_indices + original_seq_len = sequence_length + + if sp_world_size > 1: + packed_hidden_states, _ = sequence_model_parallel_shard(packed_hidden_states, dim=1) + rotary_cos, _ = sequence_model_parallel_shard(rotary_emb[0], dim=0) + rotary_sin, _ = sequence_model_parallel_shard(rotary_emb[1], dim=0) + adaln_indices, _ = sequence_model_parallel_shard(adaln_indices, dim=0) + local_timestep_indices, _ = sequence_model_parallel_shard(local_timestep_indices, dim=0) + rotary_emb = (rotary_cos, rotary_sin) + + for block in self.transformer_blocks: + packed_hidden_states = block( + packed_hidden_states, + temb, + adaln_indices, + rotary_emb, + original_seq_len, + ) + + packed_hidden_states = self.norm_out( + packed_hidden_states, + temb, + local_timestep_indices, + ).to(self.proj_out.weight.dtype) + video_output, _ = self.proj_out(packed_hidden_states) + audio_output, _ = self.audio_proj_out(packed_hidden_states) + if sp_world_size > 1: + video_output = sequence_model_parallel_all_gather_with_unpad(video_output, original_seq_len, dim=1) + audio_output = sequence_model_parallel_all_gather_with_unpad(audio_output, original_seq_len, dim=1) + video_output = video_output.index_select(1, video_indices) + audio_output = audio_output.index_select(1, audio_indices) + + return video_output, audio_output + + +EntryClass = MiniMaxH3Transformer3DModel diff --git a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py new file mode 100644 index 000000000..6c9d72691 --- /dev/null +++ b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py @@ -0,0 +1,704 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Native Qwen3-VL conditioner used by MiniMax H3.""" + +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConfig +from fastvideo.distributed import get_tp_world_size +from fastvideo.layers.layernorm import RMSNorm +from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear +from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.loader.weight_utils import default_weight_loader + + +def _rotate_half(tensor: torch.Tensor) -> torch.Tensor: + first, second = tensor.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + +class MiniMaxH3Qwen3VLTextRotaryEmbedding(nn.Module): + """Shared Qwen3-VL interleaved temporal/height/width rotary embedding.""" + + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: + super().__init__() + head_dim = config.head_dim + inv_freq = 1.0 / (config.rope_theta**(torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.mrope_section = tuple(config.mrope_section) + + def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if position_ids.ndim == 2: + position_ids = position_ids[None].expand(3, -1, -1) + inv_freq = self.inv_freq[None, None, :, None].expand(3, position_ids.shape[1], -1, 1) + positions = position_ids[:, :, None, :].float() + device_type = hidden_states.device.type if hidden_states.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + frequencies = (inv_freq.float() @ positions).transpose(2, 3) + interleaved = frequencies[0] + for dimension, offset in enumerate((1, 2), start=1): + stop = self.mrope_section[dimension] * 3 + interleaved[..., offset:stop:3] = frequencies[dimension, ..., offset:stop:3] + embedding = torch.cat((interleaved, interleaved), dim=-1) + cos = embedding.cos() + sin = embedding.sin() + return cos.to(hidden_states.dtype), sin.to(hidden_states.dtype) + + +class MiniMaxH3Qwen3VLTextAttention(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig, prefix: str) -> None: + super().__init__() + tp_size = get_tp_world_size() + self.total_num_heads = config.num_attention_heads + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_heads % tp_size: + raise ValueError(f"Qwen3-VL attention heads {self.total_num_heads} are not divisible by TP={tp_size}") + if tp_size > self.total_num_kv_heads: + raise ValueError( + f"Qwen3-VL native K/V projections require TP={tp_size} to be no larger than " + f"the {self.total_num_kv_heads} KV heads" + ) + if self.total_num_kv_heads % tp_size: + raise ValueError(f"Qwen3-VL KV heads {self.total_num_kv_heads} are not divisible by TP={tp_size}") + + self.num_heads = self.total_num_heads // tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.head_dim + self.scaling = self.head_dim**-0.5 + quant_config = getattr(config, "quant_config", None) + self.q_proj = ColumnParallelLinear( + input_size=config.hidden_size, + output_size=self.total_num_heads * self.head_dim, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.k_proj = ColumnParallelLinear( + input_size=config.hidden_size, + output_size=self.total_num_kv_heads * self.head_dim, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.k_proj", + ) + self.v_proj = ColumnParallelLinear( + input_size=config.hidden_size, + output_size=self.total_num_kv_heads * self.head_dim, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.v_proj", + ) + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=config.hidden_size, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + ) -> torch.Tensor: + query, _ = self.q_proj(hidden_states) + key, _ = self.k_proj(hidden_states) + value, _ = self.v_proj(hidden_states) + batch_size, sequence_length = hidden_states.shape[:2] + query = self.q_norm(query.view(batch_size, sequence_length, self.num_heads, self.head_dim)).transpose(1, 2) + key = self.k_norm(key.view(batch_size, sequence_length, self.num_kv_heads, self.head_dim)).transpose(1, 2) + value = value.view(batch_size, sequence_length, self.num_kv_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + query = query * cos + _rotate_half(query) * sin + key = key * cos + _rotate_half(key) * sin + + if attention_mask is None: + output = F.scaled_dot_product_attention( + query, + key, + value, + dropout_p=0.0, + is_causal=sequence_length > 1, + scale=self.scaling, + enable_gqa=self.num_heads != self.num_kv_heads, + ) + else: + if self.num_heads != self.num_kv_heads: + groups = self.num_heads // self.num_kv_heads + key = key[:, :, None].expand(-1, -1, groups, -1, -1).reshape( + batch_size, self.num_heads, sequence_length, self.head_dim) + value = value[:, :, None].expand(-1, -1, groups, -1, -1).reshape( + batch_size, self.num_heads, sequence_length, self.head_dim) + causal = torch.ones(sequence_length, sequence_length, device=query.device, dtype=torch.bool).tril() + key_mask = attention_mask.to(device=query.device, dtype=torch.bool) + mask = causal[None, None] & key_mask[:, None, None] + output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=mask, + dropout_p=0.0, + is_causal=False, + scale=self.scaling, + ) + output = output.transpose(1, 2).reshape(batch_size, sequence_length, -1) + output, _ = self.o_proj(output) + return output + + +class MiniMaxH3Qwen3VLTextMLP(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig, prefix: str) -> None: + super().__init__() + quant_config = getattr(config, "quant_config", None) + self.gate_proj = ColumnParallelLinear( + input_size=config.hidden_size, + output_size=config.intermediate_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_proj", + ) + self.up_proj = ColumnParallelLinear( + input_size=config.hidden_size, + output_size=config.intermediate_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + gate, _ = self.gate_proj(hidden_states) + up, _ = self.up_proj(hidden_states) + hidden_states = F.silu(gate) * up + hidden_states, _ = self.down_proj(hidden_states) + return hidden_states + + +class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig, prefix: str) -> None: + super().__init__() + self.self_attn = MiniMaxH3Qwen3VLTextAttention(config, prefix=f"{prefix}.self_attn") + self.mlp = MiniMaxH3Qwen3VLTextMLP(config, prefix=f"{prefix}.mlp") + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_embeddings, attention_mask) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class MiniMaxH3Qwen3VLLanguageModel(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: + super().__init__() + quant_config = getattr(config, "quant_config", None) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + quant_config=quant_config, + ) + self.layers = nn.ModuleList( + MiniMaxH3Qwen3VLTextDecoderLayer(config, prefix=f"{config.prefix}.language_model.layers.{index}") + for index in range(config.num_hidden_layers) + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = MiniMaxH3Qwen3VLTextRotaryEmbedding(config) + + def forward( + self, + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor | None, + output_hidden_states: bool, + visual_pos_masks: torch.Tensor | None, + deepstack_visual_embeds: list[torch.Tensor] | None, + ) -> BaseEncoderOutput: + if attention_mask is not None and bool(attention_mask.to(torch.bool).all()): + attention_mask = None + position_embeddings = self.rotary_emb(inputs_embeds, position_ids) + hidden_states = inputs_embeds + all_hidden_states: tuple[torch.Tensor, ...] | None = () if output_hidden_states else None + for layer_index, layer in enumerate(self.layers): + if all_hidden_states is not None: + all_hidden_states += (hidden_states,) + hidden_states = layer(hidden_states, position_embeddings, attention_mask) + if deepstack_visual_embeds is not None and layer_index < len(deepstack_visual_embeds): + if visual_pos_masks is None: + raise ValueError("Qwen3-VL DeepStack features require visual token positions") + mask = visual_pos_masks.to(hidden_states.device) + visual = deepstack_visual_embeds[layer_index].to(hidden_states.device, hidden_states.dtype) + updated = hidden_states[mask].clone() + visual + hidden_states[mask] = updated + hidden_states = self.norm(hidden_states) + if all_hidden_states is not None: + all_hidden_states += (hidden_states,) + return BaseEncoderOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states) + + +class MiniMaxH3Qwen3VLVisionPatchEmbed(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: + super().__init__() + self.in_channels = config.vision_in_channels + self.temporal_patch_size = config.vision_temporal_patch_size + self.patch_size = config.vision_patch_size + self.hidden_size = config.vision_hidden_size + kernel = (self.temporal_patch_size, self.patch_size, self.patch_size) + self.proj = nn.Conv3d(self.in_channels, self.hidden_size, kernel_size=kernel, stride=kernel, bias=True) + + def forward(self, pixels: torch.Tensor) -> torch.Tensor: + pixels = pixels.view(-1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size) + return self.proj(pixels.to(self.proj.weight.dtype)).view(-1, self.hidden_size) + + +class MiniMaxH3Qwen3VLVisionRotaryEmbedding(nn.Module): + def __init__(self, dimension: int) -> None: + super().__init__() + # Transformers initializes this non-persistent constant on CPU before + # moving the model. FastVideo constructs native encoders inside a CUDA + # device context, where evaluating the power directly differs by a few + # fp32 ULPs and is amplified by the 64-layer language model. + target_device = torch.empty(0).device + exponents = torch.arange(0, dimension, 2, dtype=torch.float32, device="cpu") / dimension + inv_freq = (1.0 / (10000.0**exponents)).to(target_device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, sequence_length: int) -> torch.Tensor: + positions = torch.arange(sequence_length, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + return torch.outer(positions, self.inv_freq) + + +class MiniMaxH3Qwen3VLVisionPatchMerger(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig, use_postshuffle_norm: bool) -> None: + super().__init__() + self.hidden_size = config.vision_hidden_size * config.vision_spatial_merge_size**2 + self.use_postshuffle_norm = use_postshuffle_norm + norm_size = self.hidden_size if use_postshuffle_norm else config.vision_hidden_size + self.norm = nn.LayerNorm(norm_size, eps=1e-6) + self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size) + self.linear_fc2 = nn.Linear(self.hidden_size, config.vision_out_hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.use_postshuffle_norm: + hidden_states = self.norm(hidden_states.view(-1, self.hidden_size)) + else: + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.view(-1, self.hidden_size) + hidden_states = F.gelu(self.linear_fc1(hidden_states)) + return self.linear_fc2(hidden_states) + + +class MiniMaxH3Qwen3VLVisionMLP(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: + super().__init__() + self.linear_fc1 = nn.Linear(config.vision_hidden_size, config.vision_intermediate_size, bias=True) + self.linear_fc2 = nn.Linear(config.vision_intermediate_size, config.vision_hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.linear_fc2(F.gelu(self.linear_fc1(hidden_states), approximate="tanh")) + + +class MiniMaxH3Qwen3VLVisionAttention(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: + super().__init__() + self.num_heads = config.vision_num_heads + self.head_dim = config.vision_hidden_size // self.num_heads + self.scaling = self.head_dim**-0.5 + self.qkv = nn.Linear(config.vision_hidden_size, config.vision_hidden_size * 3, bias=True) + self.proj = nn.Linear(config.vision_hidden_size, config.vision_hidden_size, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + sequence_lengths: list[int], + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + sequence_length = hidden_states.shape[0] + query, key, value = self.qkv(hidden_states).view(sequence_length, 3, self.num_heads, self.head_dim).unbind(1) + cos, sin = position_embeddings + query_dtype, key_dtype = query.dtype, key.dtype + cos = cos.unsqueeze(-2).float() + sin = sin.unsqueeze(-2).float() + query = (query.float() * cos + _rotate_half(query.float()) * sin).to(query_dtype) + key = (key.float() * cos + _rotate_half(key.float()) * sin).to(key_dtype) + + query_chunks = query.split(sequence_lengths, dim=0) + key_chunks = key.split(sequence_lengths, dim=0) + value_chunks = value.split(sequence_lengths, dim=0) + outputs = [] + for query_chunk, key_chunk, value_chunk in zip(query_chunks, key_chunks, value_chunks, strict=True): + output = F.scaled_dot_product_attention( + query_chunk.transpose(0, 1).unsqueeze(0), + key_chunk.transpose(0, 1).unsqueeze(0), + value_chunk.transpose(0, 1).unsqueeze(0), + dropout_p=0.0, + is_causal=False, + scale=self.scaling, + ) + outputs.append(output.transpose(1, 2)) + hidden_states = torch.cat(outputs, dim=1).reshape(sequence_length, -1) + return self.proj(hidden_states) + + +class MiniMaxH3Qwen3VLVisionBlock(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.vision_hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.vision_hidden_size, eps=1e-6) + self.attn = MiniMaxH3Qwen3VLVisionAttention(config) + self.mlp = MiniMaxH3Qwen3VLVisionMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + sequence_lengths: list[int], + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states), sequence_lengths, position_embeddings) + return hidden_states + self.mlp(self.norm2(hidden_states)) + + +class MiniMaxH3Qwen3VLVisionModel(nn.Module): + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: + super().__init__() + self.config = config + self.spatial_merge_size = config.vision_spatial_merge_size + self.patch_embed = MiniMaxH3Qwen3VLVisionPatchEmbed(config) + self.pos_embed = nn.Embedding(config.vision_num_position_embeddings, config.vision_hidden_size) + self.num_grid_per_side = int(config.vision_num_position_embeddings**0.5) + head_dim = config.vision_hidden_size // config.vision_num_heads + self.rotary_pos_emb = MiniMaxH3Qwen3VLVisionRotaryEmbedding(head_dim // 2) + self.blocks = nn.ModuleList(MiniMaxH3Qwen3VLVisionBlock(config) for _ in range(config.vision_depth)) + self.merger = MiniMaxH3Qwen3VLVisionPatchMerger(config, use_postshuffle_norm=False) + self.deepstack_visual_indexes = tuple(config.vision_deepstack_visual_indexes) + self.deepstack_merger_list = nn.ModuleList( + MiniMaxH3Qwen3VLVisionPatchMerger(config, use_postshuffle_norm=True) + for _ in self.deepstack_visual_indexes + ) + + def _rotary_positions(self, grid_thw: torch.Tensor) -> torch.Tensor: + max_height_width = int(grid_thw[:, 1:].max().item()) + frequency_table = self.rotary_pos_emb(max_height_width) + total_tokens = int(torch.prod(grid_thw, dim=1).sum().item()) + position_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=frequency_table.device) + offset = 0 + merge = self.spatial_merge_size + for frames_tensor, height_tensor, width_tensor in grid_thw: + frames, height, width = int(frames_tensor), int(height_tensor), int(width_tensor) + merged_height, merged_width = height // merge, width // merge + block_rows = torch.arange(merged_height, device=frequency_table.device) + block_cols = torch.arange(merged_width, device=frequency_table.device) + intra_rows = torch.arange(merge, device=frequency_table.device) + intra_cols = torch.arange(merge, device=frequency_table.device) + rows = block_rows[:, None, None, None] * merge + intra_rows[None, None, :, None] + cols = block_cols[None, :, None, None] * merge + intra_cols[None, None, None, :] + rows = rows.expand(merged_height, merged_width, merge, merge).reshape(-1) + cols = cols.expand(merged_height, merged_width, merge, merge).reshape(-1) + coordinates = torch.stack((rows, cols), dim=-1).repeat(frames, 1) + position_ids[offset:offset + coordinates.shape[0]] = coordinates + offset += coordinates.shape[0] + return frequency_table[position_ids].flatten(1) + + def _interpolate_position_embeddings(self, grid_thw: torch.Tensor) -> torch.Tensor: + index_lists: list[list[int]] = [[] for _ in range(4)] + weight_lists: list[list[float]] = [[] for _ in range(4)] + merge = self.spatial_merge_size + patch_counts: list[int] = [] + grids: list[tuple[int, int, int]] = [] + for frames_tensor, height_tensor, width_tensor in grid_thw: + frames, height, width = int(frames_tensor), int(height_tensor), int(width_tensor) + grids.append((frames, height, width)) + patch_counts.append(height * width) + height_positions = torch.linspace(0, self.num_grid_per_side - 1, height) + width_positions = torch.linspace(0, self.num_grid_per_side - 1, width) + height_floor = height_positions.int() + width_floor = width_positions.int() + height_ceil = (height_floor + 1).clip(max=self.num_grid_per_side - 1) + width_ceil = (width_floor + 1).clip(max=self.num_grid_per_side - 1) + delta_height = height_positions - height_floor + delta_width = width_positions - width_floor + base_height = height_floor * self.num_grid_per_side + base_height_ceil = height_ceil * self.num_grid_per_side + indices = ( + (base_height[:, None] + width_floor[None]).flatten(), + (base_height[:, None] + width_ceil[None]).flatten(), + (base_height_ceil[:, None] + width_floor[None]).flatten(), + (base_height_ceil[:, None] + width_ceil[None]).flatten(), + ) + weights = ( + ((1 - delta_height)[:, None] * (1 - delta_width)[None]).flatten(), + ((1 - delta_height)[:, None] * delta_width[None]).flatten(), + (delta_height[:, None] * (1 - delta_width)[None]).flatten(), + (delta_height[:, None] * delta_width[None]).flatten(), + ) + for index in range(4): + index_lists[index].extend(indices[index].tolist()) + weight_lists[index].extend(weights[index].tolist()) + + index_tensor = torch.tensor(index_lists, dtype=torch.long, device=self.pos_embed.weight.device) + weight_tensor = torch.tensor(weight_lists, dtype=self.pos_embed.weight.dtype, device=self.pos_embed.weight.device) + embeddings = self.pos_embed(index_tensor) * weight_tensor[:, :, None] + embeddings = (embeddings[0] + embeddings[1] + embeddings[2] + embeddings[3]).split(patch_counts) + permuted = [] + for embedding, (frames, height, width) in zip(embeddings, grids, strict=True): + embedding = embedding.repeat(frames, 1) + embedding = embedding.view(frames, height // merge, merge, width // merge, merge, -1) + permuted.append(embedding.permute(0, 1, 3, 2, 4, 5).flatten(0, 4)) + return torch.cat(permuted) + + def forward(self, pixels: torch.Tensor, grid_thw: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: + hidden_states = self.patch_embed(pixels) + hidden_states = hidden_states + self._interpolate_position_embeddings(grid_thw) + rotary = self._rotary_positions(grid_thw).reshape(hidden_states.shape[0], -1) + embedding = torch.cat((rotary, rotary), dim=-1) + position_embeddings = (embedding.cos(), embedding.sin()) + sequence_lengths = [int(value) for value in torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + )] + deepstack_features = [] + for layer_index, block in enumerate(self.blocks): + hidden_states = block(hidden_states, sequence_lengths, position_embeddings) + if layer_index in self.deepstack_visual_indexes: + merger_index = self.deepstack_visual_indexes.index(layer_index) + deepstack_features.append(self.deepstack_merger_list[merger_index](hidden_states)) + return self.merger(hidden_states), deepstack_features + + +class MiniMaxH3Qwen3VLConditioner(TextEncoder): + """FastVideo-native Qwen3-VL body without the unused language-model head.""" + + supports_hf_from_pretrained = False + + def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: + super().__init__(config) + self.visual = MiniMaxH3Qwen3VLVisionModel(config) + self.language_model = MiniMaxH3Qwen3VLLanguageModel(config) + + @property + def dtype(self) -> torch.dtype: + parameter = next(self.parameters(), None) + return torch.float32 if parameter is None else parameter.dtype + + @property + def num_hidden_layers(self) -> int: + return self.config.num_hidden_layers + + def _get_rope_index( + self, + input_ids: torch.Tensor, + image_grid_thw: torch.Tensor | None, + video_grid_thw: torch.Tensor | None, + attention_mask: torch.Tensor | None, + ) -> torch.Tensor: + if image_grid_thw is None and video_grid_thw is None: + if attention_mask is None: + return torch.arange(input_ids.shape[1], device=input_ids.device).view(1, 1, -1).expand( + 3, input_ids.shape[0], -1) + positions = attention_mask.long().cumsum(-1) - 1 + positions.masked_fill_(attention_mask == 0, 1) + return positions.unsqueeze(0).expand(3, -1, -1) + + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0).clone() + video_grid_thw[:, 0] = 1 + mask = torch.ones_like(input_ids) if attention_mask is None else attention_mask + position_ids = torch.ones(3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device) + image_index = 0 + video_index = 0 + merge = self.config.vision_spatial_merge_size + for batch_index, sequence in enumerate(input_ids): + sequence = sequence[mask[batch_index] == 1] + starts = torch.argwhere(sequence == self.config.vision_start_token_id).squeeze(1) + vision_tokens = sequence[starts + 1] + image_count = int((vision_tokens == self.config.image_token_id).sum()) + video_count = int((vision_tokens == self.config.video_token_id).sum()) + tokens = sequence.tolist() + pieces = [] + start = 0 + remaining_images, remaining_videos = image_count, video_count + for _ in range(image_count + video_count): + image_end = tokens.index(self.config.image_token_id, start) if remaining_images else len(tokens) + 1 + video_end = tokens.index(self.config.video_token_id, start) if remaining_videos else len(tokens) + 1 + if image_end < video_end: + if image_grid_thw is None: + raise ValueError("Qwen3-VL image tokens require image_grid_thw") + frames, height, width = (int(value) for value in image_grid_thw[image_index]) + image_index += 1 + remaining_images -= 1 + end = image_end + else: + if video_grid_thw is None: + raise ValueError("Qwen3-VL video tokens require video_grid_thw") + frames, height, width = (int(value) for value in video_grid_thw[video_index]) + video_index += 1 + remaining_videos -= 1 + end = video_end + grid_frames, grid_height, grid_width = frames, height // merge, width // merge + text_length = end - start + offset = int(pieces[-1].max()) + 1 if pieces else 0 + pieces.append(torch.arange(text_length).view(1, -1).expand(3, -1) + offset) + temporal = torch.arange(grid_frames).view(-1, 1).expand(-1, grid_height * grid_width).flatten() + rows = torch.arange(grid_height).view(1, -1, 1).expand(grid_frames, -1, grid_width).flatten() + columns = torch.arange(grid_width).view(1, 1, -1).expand(grid_frames, grid_height, -1).flatten() + pieces.append(torch.stack((temporal, rows, columns)) + text_length + offset) + start = end + grid_frames * grid_height * grid_width + if start < len(tokens): + offset = int(pieces[-1].max()) + 1 if pieces else 0 + pieces.append(torch.arange(len(tokens) - start).view(1, -1).expand(3, -1) + offset) + positions = torch.cat(pieces, dim=1).to(position_ids.device) + position_ids[:, batch_index, mask[batch_index] == 1] = positions + return position_ids + + def _visual_features( + self, + pixels: torch.Tensor, + grid_thw: torch.Tensor, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + return self.visual(pixels.to(self.visual.patch_embed.proj.weight.dtype), grid_thw) + + @staticmethod + def _placeholder_mask( + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor, + token_id: int, + features: torch.Tensor, + label: str, + ) -> torch.Tensor: + mask = input_ids == token_id + expanded = mask.unsqueeze(-1).expand_as(inputs_embeds) + if inputs_embeds[expanded].numel() != features.numel(): + raise ValueError( + f"Qwen3-VL {label} features and placeholder tokens do not match: " + f"tokens={int(mask.sum())}, features={features.shape[0]}" + ) + return mask + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + mm_token_type_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> BaseEncoderOutput: + del mm_token_type_ids, kwargs + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Exactly one of input_ids or inputs_embeds is required") + if inputs_embeds is None: + assert input_ids is not None + inputs_embeds = self.language_model.embed_tokens(input_ids) + if input_ids is None and (pixel_values is not None or pixel_values_videos is not None): + raise ValueError("Multimodal Qwen3-VL inputs require input_ids for placeholder matching") + + image_mask = None + video_mask = None + image_deepstack = None + video_deepstack = None + if pixel_values is not None: + if input_ids is None or image_grid_thw is None: + raise ValueError("pixel_values require input_ids and image_grid_thw") + image_features, image_deepstack = self._visual_features(pixel_values, image_grid_thw) + image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + image_mask = self._placeholder_mask( + input_ids, inputs_embeds, self.config.image_token_id, image_features, "image") + inputs_embeds = inputs_embeds.masked_scatter(image_mask.unsqueeze(-1), image_features) + if pixel_values_videos is not None: + if input_ids is None or video_grid_thw is None: + raise ValueError("pixel_values_videos require input_ids and video_grid_thw") + video_features, video_deepstack = self._visual_features(pixel_values_videos, video_grid_thw) + video_features = video_features.to(inputs_embeds.device, inputs_embeds.dtype) + video_mask = self._placeholder_mask( + input_ids, inputs_embeds, self.config.video_token_id, video_features, "video") + inputs_embeds = inputs_embeds.masked_scatter(video_mask.unsqueeze(-1), video_features) + + visual_mask = None + deepstack_features = None + if image_mask is not None and video_mask is not None: + visual_mask = image_mask | video_mask + deepstack_features = [] + image_joint = image_mask[visual_mask] + video_joint = video_mask[visual_mask] + assert image_deepstack is not None and video_deepstack is not None + for image_features, video_features in zip(image_deepstack, video_deepstack, strict=True): + combined = image_features.new_zeros((int(visual_mask.sum()), image_features.shape[-1])) + combined[image_joint] = image_features + combined[video_joint] = video_features + deepstack_features.append(combined) + elif image_mask is not None: + visual_mask = image_mask + deepstack_features = image_deepstack + elif video_mask is not None: + visual_mask = video_mask + deepstack_features = video_deepstack + + if position_ids is None: + if input_ids is None: + sequence_length = inputs_embeds.shape[1] + position_ids = torch.arange(sequence_length, device=inputs_embeds.device).view(1, 1, -1).expand( + 3, inputs_embeds.shape[0], -1) + else: + position_ids = self._get_rope_index(input_ids, image_grid_thw, video_grid_thw, attention_mask) + output_hidden_states = self.config.output_hidden_states if output_hidden_states is None else output_hidden_states + outputs = self.language_model( + inputs_embeds, + position_ids, + attention_mask, + output_hidden_states, + visual_mask, + deepstack_features, + ) + outputs.attention_mask = attention_mask + return outputs + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + parameters = dict(self.named_parameters()) + loaded: set[str] = set() + + for source_name, tensor in weights: + if source_name == "lm_head.weight": + continue + name = source_name[6:] if source_name.startswith("model.") else source_name + if name not in parameters: + raise ValueError(f"Unexpected MiniMax-H3 Qwen3-VL checkpoint key: {source_name}") + parameter = parameters[name] + loader = getattr(parameter, "weight_loader", default_weight_loader) + loader(parameter, tensor) + loaded.add(name) + return loaded + + +EntryClass = MiniMaxH3Qwen3VLConditioner + +__all__ = ["MiniMaxH3Qwen3VLConditioner"] diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 48218c823..dba975819 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -84,7 +84,9 @@ class ComponentLoader(ABC): # Map of module types to their loader classes and expected library module_loaders = { "scheduler": (SchedulerLoader, "diffusers"), + "audio_scheduler": (SchedulerLoader, "diffusers"), "transformer": (TransformerLoader, "diffusers"), + "transformer_ref": (TransformerLoader, "diffusers"), "sr_transformer": (TransformerLoader, "diffusers"), "transformer_2": (TransformerLoader, "diffusers"), "transformer_3": (TransformerLoader, "diffusers"), @@ -911,7 +913,7 @@ class VAELoader(ComponentLoader): # Diffusers-format AutoencoderKL checkpoints should match exactly; load # strictly so missing/unexpected keys are surfaced early. - strict_load = class_name == "AutoencoderKL" + strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3"} vae.load_state_dict(loaded, strict=strict_load) if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False) @@ -923,15 +925,36 @@ class VAELoader(ComponentLoader): class AudioDecoderLoader(ComponentLoader): - """Loader for LTX-2 audio decoder (audio_vae component).""" + """Loader for full audio VAEs and legacy decode-only audio components.""" def load(self, model_path: str, fastvideo_args: FastVideoArgs): config = get_diffusers_config(model=model_path) class_name = config.pop("_class_name", None) or "LTX2AudioDecoder" - model_cls, _ = ModelRegistry.resolve_model_cls(class_name) target_device = get_local_torch_device() + if class_name == "AutoencoderKLMiniMaxH3Audio": + from fastvideo.platforms import current_platform + + configured_audio_vae = getattr(fastvideo_args.pipeline_config, "audio_vae_config", None) + if configured_audio_vae is None: + raise ValueError("MiniMax H3 requires audio_vae_config.") + config.pop("_name_or_path", None) + audio_vae_config = deepcopy(configured_audio_vae) + audio_vae_config.update_model_arch(config) + if getattr(fastvideo_args, "vae_cpu_offload", False): + target_device = torch.device("mps") if current_platform.is_mps() else torch.device("cpu") + with set_default_torch_dtype(torch.float32): + audio_vae = model_cls(audio_vae_config).to(device=target_device, dtype=torch.float32) + safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) + if not safetensors_list: + raise ValueError(f"No safetensors files found in {model_path}") + loaded: dict[str, torch.Tensor] = {} + for sf_file in safetensors_list: + loaded.update(safetensors_load_file(sf_file)) + audio_vae.load_state_dict(loaded, strict=True) + return audio_vae.eval() + precision = getattr( fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16" ) @@ -1132,9 +1155,10 @@ class TransformerLoader(ComponentLoader): total_params = sum(p.numel() for p in model.parameters()) logger.info("Loaded model with %.2fB parameters", total_params / 1e9) - assert next(model.parameters()).dtype == default_dtype, ( - "Model dtype does not match default dtype" - ) + first_name, first_parameter = next(model.named_parameters()) + dtype_selector = getattr(model, "_get_parameter_dtype", None) + expected_dtype = dtype_selector(first_name, default_dtype) if callable(dtype_selector) else default_dtype + assert first_parameter.dtype == expected_dtype, "Model dtype does not match the configured parameter dtype" model = model.eval() diff --git a/fastvideo/models/schedulers/scheduling_minimax_h3.py b/fastvideo/models/schedulers/scheduling_minimax_h3.py new file mode 100644 index 000000000..44cf99ff7 --- /dev/null +++ b/fastvideo/models/schedulers/scheduling_minimax_h3.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The MiniMax authors and The HuggingFace Team. + +from dataclasses import dataclass + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + + +@dataclass +class MiniMaxH3SchedulerOutput(BaseOutput): + prev_sample: torch.FloatTensor + + +class MiniMaxH3Scheduler(SchedulerMixin, ConfigMixin): + """Rectified-flow scheduler using MiniMax-H3's clean-time convention.""" + + _compatibles: list[str] = [] + order = 1 + + @register_to_config + def __init__(self, shift: float = 12.0) -> None: + if shift <= 0: + raise ValueError(f"`shift` must be positive, got {shift}.") + self.num_inference_steps: int | None = None + self.sigmas: torch.Tensor | None = None + self.timesteps: torch.Tensor | None = None + self._shift = float(shift) + self._step_index: int | None = None + self._begin_index: int | None = None + + @property + def shift(self) -> float: + return self._shift + + @property + def step_index(self) -> int | None: + return self._step_index + + @property + def begin_index(self) -> int | None: + return self._begin_index + + def set_begin_index(self, begin_index: int = 0) -> None: + self._begin_index = begin_index + + def set_shift(self, shift: float) -> None: + if shift <= 0: + raise ValueError(f"`shift` must be positive, got {shift}.") + self._shift = float(shift) + + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + sigmas: list[float] | torch.Tensor | None = None, + ) -> None: + if sigmas is None: + if num_inference_steps is None or num_inference_steps < 2: + raise ValueError( + "`set_timesteps` requires explicit `sigmas` or " + f"`num_inference_steps` >= 2, got {num_inference_steps}." + ) + base = torch.linspace(1.0, 0.0, int(num_inference_steps), dtype=torch.float32) + sigma_tensor = self._shift * base / (1 + (self._shift - 1) * base) + sigma_tensor = torch.unique_consecutive(sigma_tensor) + else: + sigma_tensor = torch.as_tensor(sigmas, dtype=torch.float32).flatten().cpu() + is_valid = ( + sigma_tensor.numel() >= 2 + and bool((sigma_tensor[1:] < sigma_tensor[:-1]).all()) + and sigma_tensor[-1].item() == 0.0 + ) + if not is_valid: + raise ValueError("`sigmas` must hold at least two strictly decreasing values ending at 0.0.") + + self.sigmas = sigma_tensor.to(device=device) + self.timesteps = (1.0 - sigma_tensor[:-1]).to(device=device) + self.num_inference_steps = int(self.timesteps.numel()) + self._step_index = None + self._begin_index = None + + def index_for_timestep(self, timestep: float | torch.Tensor) -> int: + if self.timesteps is None: + raise ValueError("Call `set_timesteps` before looking up a timestep.") + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + indices = (self.timesteps == timestep).nonzero() + if len(indices) == 0: + raise ValueError("Passed `timestep` is not in `self.timesteps`.") + return indices[0].item() + + def scale_model_input(self, sample: torch.Tensor, timestep: float | torch.Tensor | None = None) -> torch.Tensor: + return sample + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: float | torch.FloatTensor, + noise: torch.FloatTensor, + ) -> torch.FloatTensor: + if not isinstance(timestep, torch.Tensor): + timestep = torch.tensor(timestep, dtype=sample.dtype, device=sample.device) + timestep = timestep.to(device=sample.device, dtype=sample.dtype) + while timestep.ndim < sample.ndim: + timestep = timestep.unsqueeze(-1) + return timestep * sample + (1.0 - timestep) * noise + + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor, + sample: torch.FloatTensor, + return_dict: bool = True, + ) -> MiniMaxH3SchedulerOutput | tuple[torch.FloatTensor]: + if isinstance(timestep, int) or (isinstance(timestep, torch.Tensor) and not timestep.is_floating_point()): + raise ValueError("Pass a floating-point value from `scheduler.timesteps`, not an integer index.") + if self.sigmas is None or self.timesteps is None: + raise ValueError("Call `set_timesteps` before `step`.") + if self._step_index is None: + self._step_index = self.index_for_timestep(timestep) if self._begin_index is None else self._begin_index + + if not isinstance(timestep, torch.Tensor): + timestep = torch.tensor(timestep, dtype=sample.dtype) + # H3 deliberately derives x0's sigma from the transformer timestep, while + # the Euler ratio below uses the stored grid. Keep the two float32 paths separate. + sigma_from_timestep = 1 - timestep.to(device=sample.device, dtype=sample.dtype) + while sigma_from_timestep.ndim < sample.ndim: + sigma_from_timestep = sigma_from_timestep.unsqueeze(-1) + denoised = sample + sigma_from_timestep * model_output + + compute_dtype = torch.float32 if sample.dtype in (torch.float16, torch.bfloat16) else sample.dtype + sigma = self.sigmas[self._step_index].to(device=sample.device, dtype=compute_dtype) + sigma_next = self.sigmas[self._step_index + 1].to(device=sample.device, dtype=compute_dtype) + ratio = sigma_next / sigma + prev_sample = ratio * sample.to(compute_dtype) + (1.0 - ratio) * denoised.to(compute_dtype) + prev_sample = prev_sample.to(sample.dtype) + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + return MiniMaxH3SchedulerOutput(prev_sample=prev_sample) + + +EntryClass = MiniMaxH3Scheduler diff --git a/fastvideo/models/vaes/minimax_h3_audio.py b/fastvideo/models/vaes/minimax_h3_audio.py new file mode 100644 index 000000000..8b07dd731 --- /dev/null +++ b/fastvideo/models/vaes/minimax_h3_audio.py @@ -0,0 +1,530 @@ +# Copyright 2025 MiniMax authors and HuggingFace Team +# SPDX-License-Identifier: Apache-2.0 + +"""Native MiniMax H3 waveform autoencoder.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils import weight_norm + +from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEConfig +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class MiniMaxH3AudioDiagonalGaussianDistribution: + """Diagonal Gaussian parameterized by mean and log standard deviation.""" + + def __init__(self, mean: torch.Tensor, logs: torch.Tensor): + self.mean = mean + self.logs = logs + self.std = torch.exp(logs) + + def mode(self) -> torch.Tensor: + return self.mean + + def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: + noise_device = self.mean.device + if generator is not None and generator.device.type == "cpu": + noise_device = torch.device("cpu") + noise = torch.randn( + self.mean.shape, + generator=generator, + device=noise_device, + dtype=self.mean.dtype, + ).to(self.mean.device) + return self.mean + self.std * noise + + +@dataclass +class MiniMaxH3AudioEncoderOutput: + latent_dist: MiniMaxH3AudioDiagonalGaussianDistribution + + +@dataclass +class MiniMaxH3AudioDecoderOutput: + sample: torch.Tensor + + +def _wn_conv1d(*args, **kwargs) -> nn.Module: + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def _linear(layer: ReplicatedLinear, hidden_states: torch.Tensor) -> torch.Tensor: + return layer(hidden_states)[0] + + +def _module_dtype(module: nn.Module) -> torch.dtype: + return next(module.parameters()).dtype + + +def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + """Build the persistent low-pass filter used by alias-free activations.""" + + half_size = kernel_size // 2 + attenuation = 2.285 * (half_size - 1) * math.pi * (4 * half_width) + 7.95 + if attenuation > 50.0: + beta = 0.1102 * (attenuation - 8.7) + elif attenuation >= 21.0: + beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) + else: + beta = 0.0 + + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32) + if kernel_size % 2 == 0: + time = torch.arange(-half_size, half_size, dtype=torch.float32) + 0.5 + else: + time = torch.arange(kernel_size, dtype=torch.float32) - half_size + + filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time) + filter_ /= filter_.sum() + return filter_.view(1, 1, kernel_size) + + +class MiniMaxH3AudioSnake1d(nn.Module): + + def __init__(self, channels: int): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * hidden_states).pow(2) + + +class MiniMaxH3AudioSnakeBeta(nn.Module): + + def __init__(self, channels: int): + super().__init__() + self.alpha = nn.Parameter(torch.zeros(channels)) + self.beta = nn.Parameter(torch.zeros(channels)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + alpha = torch.exp(self.alpha.unsqueeze(0).unsqueeze(-1)) + beta = torch.exp(self.beta.unsqueeze(0).unsqueeze(-1)) + return hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2) + + +class MiniMaxH3AudioLowPassFilter1d(nn.Module): + + def __init__(self, cutoff: float, half_width: float, stride: int, kernel_size: int): + super().__init__() + even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_channels = hidden_states.shape[1] + hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate") + return F.conv1d( + hidden_states, + self.filter.expand(num_channels, -1, -1), + stride=self.stride, + groups=num_channels, + ) + + +class MiniMaxH3AudioUpSample1d(nn.Module): + + def __init__(self, ratio: int, kernel_size: int): + super().__init__() + self.ratio = ratio + self.stride = ratio + self.pad = kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (kernel_size - self.stride + 1) // 2 + self.register_buffer( + "filter", + kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=kernel_size), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_channels = hidden_states.shape[1] + hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate") + hidden_states = self.ratio * F.conv_transpose1d( + hidden_states, + self.filter.expand(num_channels, -1, -1), + stride=self.stride, + groups=num_channels, + ) + return hidden_states[..., self.pad_left:-self.pad_right] + + +class MiniMaxH3AudioDownSample1d(nn.Module): + + def __init__(self, ratio: int, kernel_size: int): + super().__init__() + self.lowpass = MiniMaxH3AudioLowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=kernel_size, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lowpass(hidden_states) + + +class MiniMaxH3AudioActivation1d(nn.Module): + + def __init__(self, activation: nn.Module, ratio: int = 2, kernel_size: int = 12): + super().__init__() + self.act = activation + self.upsample = MiniMaxH3AudioUpSample1d(ratio, kernel_size) + self.downsample = MiniMaxH3AudioDownSample1d(ratio, kernel_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.upsample(hidden_states) + hidden_states = self.act(hidden_states) + return self.downsample(hidden_states) + + +class MiniMaxH3AudioResidualUnit(nn.Module): + + def __init__(self, dim: int, dilation: int): + super().__init__() + self.block = nn.Sequential( + MiniMaxH3AudioSnake1d(dim), + _wn_conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=((7 - 1) * dilation) // 2), + MiniMaxH3AudioSnake1d(dim), + _wn_conv1d(dim, dim, kernel_size=1), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = self.block(hidden_states) + pad = (hidden_states.shape[-1] - residual.shape[-1]) // 2 + if pad > 0: + hidden_states = hidden_states[..., pad:-pad] + return hidden_states + residual + + +class MiniMaxH3AudioEncoderBlock(nn.Module): + + def __init__(self, dim: int, stride: int): + super().__init__() + self.block = nn.Sequential( + MiniMaxH3AudioResidualUnit(dim // 2, dilation=1), + MiniMaxH3AudioResidualUnit(dim // 2, dilation=3), + MiniMaxH3AudioResidualUnit(dim // 2, dilation=9), + MiniMaxH3AudioSnake1d(dim // 2), + _wn_conv1d( + dim // 2, + dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.block(hidden_states) + + +class MiniMaxH3AudioEncoder(nn.Module): + + def __init__(self, d_model: int, strides: tuple[int, ...], d_latent: int): + super().__init__() + block: list[nn.Module] = [_wn_conv1d(1, d_model, kernel_size=7, padding=3)] + for stride in strides: + d_model *= 2 + block.append(MiniMaxH3AudioEncoderBlock(d_model, stride=stride)) + block += [ + MiniMaxH3AudioSnake1d(d_model), + _wn_conv1d(d_model, d_latent, kernel_size=3, padding=1), + ] + self.block = nn.Sequential(*block) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.block(hidden_states) + + +class MiniMaxH3AudioGeGluMlp(nn.Module): + + def __init__(self, in_features: int, hidden_features: int): + super().__init__() + self.norm = nn.LayerNorm(in_features) + self.act = nn.GELU(approximate="tanh") + self.w0 = ReplicatedLinear(in_features, hidden_features, params_dtype=torch.float32) + self.w1 = ReplicatedLinear(in_features, hidden_features, params_dtype=torch.float32) + self.w2 = ReplicatedLinear(hidden_features, in_features, params_dtype=torch.float32) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm(hidden_states) + hidden_states = self.act(_linear(self.w0, hidden_states)) * _linear(self.w1, hidden_states) + return _linear(self.w2, hidden_states) + + +class MiniMaxH3AudioCausalAttention(nn.Module): + """Causal projection from the encoder trunk width to the latent width.""" + + def __init__(self, in_dim: int, out_dim: int, num_heads: int): + super().__init__() + self.out_dim = out_dim + self.num_heads = num_heads + self.head_dim = in_dim // num_heads + self.qkv = ReplicatedLinear(in_dim, in_dim * 3, bias=False, params_dtype=torch.float32) + self.q_bias = nn.Parameter(torch.zeros(in_dim)) + self.v_bias = nn.Parameter(torch.zeros(in_dim)) + self.register_buffer("zero_k_bias", torch.zeros(in_dim)) + self.proj = ReplicatedLinear(out_dim, out_dim, params_dtype=torch.float32) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + qkv = F.linear( + hidden_states, + self.qkv.weight, + torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)), + ) + query, key, value = ( + qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim).permute(2, 0, 1, 3, 4).unbind(0) + ) + + # H3's projection is an unusual causal flat stream, not a VAE spatial-attention block. + hidden_states = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + dropout_p=0.0, + is_causal=True, + ).transpose(1, 2) + hidden_states = torch.mean(hidden_states, dim=2) + hidden_states = F.adaptive_avg_pool1d(hidden_states, self.out_dim) + return _linear(self.proj, hidden_states) + + +class MiniMaxH3AudioAttnProjection(nn.Module): + + def __init__(self, in_dim: int, out_dim: int, num_heads: int, mlp_ratio: int = 2): + super().__init__() + self.norm1 = nn.LayerNorm(in_dim) + self.attn = MiniMaxH3AudioCausalAttention(in_dim, out_dim, num_heads) + self.proj = ReplicatedLinear(in_dim, out_dim, params_dtype=torch.float32) + self.norm3 = nn.LayerNorm(in_dim) + self.norm2 = nn.LayerNorm(out_dim) + self.mlp = MiniMaxH3AudioGeGluMlp(in_features=out_dim, hidden_features=out_dim * mlp_ratio) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = _linear(self.proj, self.norm3(hidden_states)) + self.attn(self.norm1(hidden_states)) + return hidden_states + self.mlp(self.norm2(hidden_states)) + + +class MiniMaxH3AudioAMPBlock(nn.Module): + + def __init__(self, channels: int, kernel_size: int, dilation: tuple[int, ...]): + super().__init__() + self.convs1 = nn.ModuleList([ + _wn_conv1d(channels, channels, kernel_size, dilation=d, padding=(kernel_size * d - d) // 2) + for d in dilation + ]) + self.convs2 = nn.ModuleList([ + _wn_conv1d(channels, channels, kernel_size, dilation=1, padding=(kernel_size - 1) // 2) + for _ in dilation + ]) + self.activations = nn.ModuleList([ + MiniMaxH3AudioActivation1d(activation=MiniMaxH3AudioSnakeBeta(channels)) + for _ in range(2 * len(dilation)) + ]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + acts1, acts2 = self.activations[::2], self.activations[1::2] + for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2): + residual = conv1(act1(hidden_states)) + residual = conv2(act2(residual)) + hidden_states = residual + hidden_states + return hidden_states + + +class MiniMaxH3AudioBigVGANDecoder(nn.Module): + + def __init__( + self, + in_channels: int, + upsample_initial_channel: int, + upsample_rates: tuple[int, ...], + upsample_kernel_sizes: tuple[int, ...], + resblock_kernel_sizes: tuple[int, ...], + resblock_dilation_sizes: tuple[tuple[int, ...], ...], + ): + super().__init__() + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self.conv_pre = _wn_conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3) + + self.ups = nn.ModuleList() + for i, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append(nn.ModuleList([ + weight_norm(nn.ConvTranspose1d( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2 ** (i + 1)), + kernel, + rate, + padding=(kernel - rate) // 2, + )) + ])) + + self.resblocks = nn.ModuleList() + for i in range(self.num_upsamples): + channels = upsample_initial_channel // (2 ** (i + 1)) + for kernel, dilation in zip(resblock_kernel_sizes, resblock_dilation_sizes): + self.resblocks.append(MiniMaxH3AudioAMPBlock(channels, kernel, tuple(dilation))) + + self.activation_post = MiniMaxH3AudioActivation1d(activation=MiniMaxH3AudioSnakeBeta(channels)) + self.conv_post = _wn_conv1d(channels, 1, 7, 1, padding=3, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_pre(hidden_states) + for i in range(self.num_upsamples): + hidden_states = self.ups[i][0](hidden_states) + residual = None + for j in range(self.num_kernels): + block = self.resblocks[i * self.num_kernels + j](hidden_states) + residual = block if residual is None else residual + block + hidden_states = residual / self.num_kernels + + hidden_states = self.activation_post(hidden_states) + hidden_states = self.conv_post(hidden_states) + return torch.clamp(hidden_states, min=-1.0, max=1.0) + + +class MiniMaxH3AudioVAE(nn.Module): + """DAC encoder plus BigVGAN decoder for mono 32 kHz waveforms.""" + + def __init__(self, config: MiniMaxH3AudioVAEConfig): + super().__init__() + self.config = config + self.fastvideo_config = config + arch = config.arch_config + + encoder_rates = tuple(int(rate) for rate in arch.encoder_rates) + decoder_rates = tuple(int(rate) for rate in arch.decoder_rates) + self.hop_length = math.prod(encoder_rates) + self.sampling_rate = int(arch.sampling_rate) + self.latent_channels = int(arch.latent_channels) + self.audio_channels = 1 + latents_mean = arch.latents_mean if arch.latents_mean is not None else [0.0] * self.latent_channels + latents_std = arch.latents_std if arch.latents_std is not None else [1.0] * self.latent_channels + self.register_buffer( + "latents_mean", + torch.tensor(latents_mean, dtype=torch.float32).view(1, -1, 1), + persistent=False, + ) + self.register_buffer( + "latents_std", + torch.tensor(latents_std, dtype=torch.float32).view(1, -1, 1), + persistent=False, + ) + + if math.prod(decoder_rates) != self.hop_length: + raise ValueError( + f"`decoder_rates` must upsample by the encoder hop length {self.hop_length}, got " + f"{math.prod(decoder_rates)}." + ) + if arch.latent_dim % arch.latent_channels != 0: + raise ValueError( + f"`latent_dim` ({arch.latent_dim}) must be a multiple of `latent_channels` " + f"({arch.latent_channels})." + ) + + self.encoder = MiniMaxH3AudioEncoder( + d_model=arch.encoder_dim, + strides=encoder_rates, + d_latent=arch.latent_dim, + ) + self.pre_block = MiniMaxH3AudioAttnProjection( + arch.latent_dim, + arch.latent_channels, + num_heads=arch.num_attention_heads, + ) + self.mean_proj = nn.Conv1d(arch.latent_channels, arch.latent_channels, 1) + self.logs_proj = nn.Conv1d(arch.latent_channels, arch.latent_channels, 1) + + self.dec_in_proj = nn.Conv1d(arch.latent_channels, arch.latent_dim, 1) + self.decoder = MiniMaxH3AudioBigVGANDecoder( + in_channels=arch.latent_dim, + upsample_initial_channel=arch.decoder_dim, + upsample_rates=decoder_rates, + upsample_kernel_sizes=tuple(int(kernel) for kernel in arch.decoder_kernel_sizes), + resblock_kernel_sizes=tuple(int(kernel) for kernel in arch.resblock_kernel_sizes), + resblock_dilation_sizes=tuple( + tuple(int(dilation) for dilation in group) for group in arch.resblock_dilation_sizes + ), + ) + + # The H3 checkpoint and waveform numerics require this component to stay FP32. + self.float() + + def encode( + self, + sample: torch.Tensor, + return_dict: bool = True, + ) -> MiniMaxH3AudioEncoderOutput | tuple[MiniMaxH3AudioDiagonalGaussianDistribution]: + if sample.ndim != 3 or sample.shape[1] != 1: + raise ValueError(f"`sample` must have shape [batch_size, 1, samples], got {tuple(sample.shape)}.") + + right_pad = math.ceil(sample.shape[-1] / self.hop_length) * self.hop_length - sample.shape[-1] + if right_pad > 0: + sample = F.pad(sample, (0, right_pad)) + + encoder_dtype = _module_dtype(self.encoder) + hidden_states = self.encoder(sample.to(encoder_dtype)) + hidden_states = self.pre_block(hidden_states.transpose(1, 2)).transpose(1, 2) + mean = self.mean_proj(hidden_states) + logs = self.logs_proj(hidden_states) + if encoder_dtype != torch.float32: + mean = mean.float() + logs = logs.float() + + posterior = MiniMaxH3AudioDiagonalGaussianDistribution(mean, logs) + if not return_dict: + return (posterior,) + return MiniMaxH3AudioEncoderOutput(latent_dist=posterior) + + def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor: + return (latents - self.latents_mean) / self.latents_std + + def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor: + return latents * self.latents_std + self.latents_mean + + def decode( + self, + latents: torch.Tensor, + return_dict: bool = True, + ) -> MiniMaxH3AudioDecoderOutput | tuple[torch.Tensor]: + if latents.ndim != 3: + raise ValueError( + f"`latents` must have shape [batch_size, latent_channels, num_frames], got {tuple(latents.shape)}." + ) + + decoder_dtype = _module_dtype(self.decoder) + decoded = self.decoder(self.dec_in_proj(latents.to(decoder_dtype))) + if decoder_dtype != torch.float32: + decoded = decoded.float() + + if not return_dict: + return (decoded,) + return MiniMaxH3AudioDecoderOutput(sample=decoded) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> MiniMaxH3AudioDecoderOutput | tuple[torch.Tensor]: + posterior = self.encode(sample).latent_dist + latents = posterior.sample(generator=generator) if sample_posterior else posterior.mode() + return self.decode(latents, return_dict=return_dict) + + +# Keep the Diffusers checkpoint class name loadable through FastVideo's registry. +AutoencoderKLMiniMaxH3Audio = MiniMaxH3AudioVAE +EntryClass = [MiniMaxH3AudioVAE, AutoencoderKLMiniMaxH3Audio] diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py new file mode 100644 index 000000000..d925a70e5 --- /dev/null +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -0,0 +1,847 @@ +# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Native MiniMax-H3 video VAE. + +The encoder is a causal 3D CNN while the decoder is a full-attention ViT. +This module intentionally uses only PyTorch and FastVideo configuration types. +""" + +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.checkpoint import checkpoint + +from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig + + +class DiagonalGaussianDistribution: + """Diagonal Gaussian posterior used by the KL encoder.""" + + def __init__(self, parameters: torch.Tensor, deterministic: bool = False) -> None: + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if deterministic: + self.std = torch.zeros_like(self.mean) + self.var = torch.zeros_like(self.mean) + + def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: + noise_device = self.mean.device + if generator is not None and generator.device.type == "cpu": + noise_device = torch.device("cpu") + noise = torch.randn( + self.mean.shape, + generator=generator, + device=noise_device, + dtype=self.mean.dtype, + ).to(self.mean.device) + return self.mean + self.std * noise + + def mode(self) -> torch.Tensor: + return self.mean + + +@dataclass +class AutoencoderKLOutput: + latent_dist: DiagonalGaussianDistribution + + +@dataclass +class DecoderOutput: + sample: torch.Tensor + + +class MiniMaxH3VideoCausalConv3d(nn.Conv3d): + """3D convolution with reflect spatial padding and causal temporal padding.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int], + stride: int | tuple[int, int, int] = 1, + spatial_padding: int = 0, + temporal_padding: int = 0, + spatial_padding_mode: str = "reflect", + ) -> None: + super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=0) + self.spatial_padding = spatial_padding + self.temporal_padding = temporal_padding + self.spatial_padding_mode = spatial_padding_mode + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.spatial_padding > 0: + padding = self.spatial_padding + hidden_states = F.pad( + hidden_states, + (padding, padding, padding, padding, 0, 0), + mode=self.spatial_padding_mode, + ) + if self.temporal_padding > 0: + hidden_states = F.pad(hidden_states, (0, 0, 0, 0, self.temporal_padding, 0), mode="constant") + return F.conv3d(hidden_states, self.weight, self.bias, stride=self.stride, padding=0, dilation=self.dilation) + + +class MiniMaxH3VideoGroupNorm(nn.GroupNorm): + """GroupNorm with each temporal frame normalized independently.""" + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 1, 3, 4).contiguous() + hidden_states = hidden_states.view(batch_size * num_frames, num_channels, 1, height, width) + hidden_states = super().forward(hidden_states) + hidden_states = hidden_states.view(batch_size, num_frames, num_channels, height, width) + return hidden_states.permute(0, 2, 1, 3, 4).contiguous() + + +class MiniMaxH3VideoResnetBlock3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + spatial_padding_mode: str = "reflect", + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.norm1 = MiniMaxH3VideoGroupNorm(norm_num_groups, in_channels, eps=norm_eps, affine=True) + self.conv1 = MiniMaxH3VideoCausalConv3d( + in_channels, + out_channels, + kernel_size=3, + spatial_padding=1, + temporal_padding=2, + spatial_padding_mode=spatial_padding_mode, + ) + self.norm2 = MiniMaxH3VideoGroupNorm(norm_num_groups, out_channels, eps=norm_eps, affine=True) + self.conv2 = MiniMaxH3VideoCausalConv3d( + out_channels, + out_channels, + kernel_size=3, + spatial_padding=1, + temporal_padding=2, + spatial_padding_mode=spatial_padding_mode, + ) + self.conv_shortcut = None + if in_channels != out_channels: + self.conv_shortcut = MiniMaxH3VideoCausalConv3d(in_channels, out_channels, kernel_size=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = hidden_states + hidden_states = self.conv1(F.silu(self.norm1(hidden_states))) + hidden_states = self.conv2(F.silu(self.norm2(hidden_states))) + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + return residual + hidden_states + + +class MiniMaxH3VideoDownsample3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temporal_stride: int = 1, + spatial_stride: int = 2, + spatial_padding_mode: str = "reflect", + ) -> None: + super().__init__() + self.spatial_stride = spatial_stride + self.spatial_padding_mode = spatial_padding_mode + self.conv = MiniMaxH3VideoCausalConv3d( + in_channels, + out_channels, + kernel_size=3, + stride=(temporal_stride, spatial_stride, spatial_stride), + temporal_padding=2, + spatial_padding_mode=spatial_padding_mode, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.spatial_stride == 2: + hidden_states = F.pad(hidden_states, (0, 1, 0, 1, 0, 0), mode=self.spatial_padding_mode) + return self.conv(hidden_states) + + +class MiniMaxH3VideoDownBlock3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int, + temporal_downsample_factor: int, + spatial_downsample_factor: int, + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + spatial_padding_mode: str = "reflect", + ) -> None: + super().__init__() + self.resnets = nn.ModuleList( + [ + MiniMaxH3VideoResnetBlock3d( + in_channels=in_channels if index == 0 else out_channels, + out_channels=out_channels, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + spatial_padding_mode=spatial_padding_mode, + ) + for index in range(num_layers) + ] + ) + self.downsamplers = None + if temporal_downsample_factor * spatial_downsample_factor > 1: + self.downsamplers = nn.ModuleList( + [ + MiniMaxH3VideoDownsample3d( + out_channels, + out_channels, + temporal_stride=temporal_downsample_factor, + spatial_stride=spatial_downsample_factor, + spatial_padding_mode=spatial_padding_mode, + ) + ] + ) + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for resnet in self.resnets: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = checkpoint(resnet, hidden_states, use_reentrant=False) + else: + hidden_states = resnet(hidden_states) + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + return hidden_states + + +class MiniMaxH3VideoEncoder3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + block_out_channels: tuple[int, ...], + layers_per_block: int, + spatial_downsample_factors: tuple[int, ...], + temporal_downsample_factors: tuple[int, ...], + norm_num_groups: int, + norm_eps: float, + spatial_padding_mode: str, + ) -> None: + super().__init__() + self.conv_in = MiniMaxH3VideoCausalConv3d( + in_channels, + block_out_channels[0], + kernel_size=3, + spatial_padding=1, + temporal_padding=2, + spatial_padding_mode=spatial_padding_mode, + ) + block_in_channels = (block_out_channels[0],) + tuple(block_out_channels[:-1]) + self.down_blocks = nn.ModuleList( + [ + MiniMaxH3VideoDownBlock3d( + in_channels=block_in_channels[index], + out_channels=block_out_channels[index], + num_layers=layers_per_block, + temporal_downsample_factor=temporal_downsample_factors[index], + spatial_downsample_factor=spatial_downsample_factors[index], + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + spatial_padding_mode=spatial_padding_mode, + ) + for index in range(len(block_out_channels)) + ] + ) + self.norm_out = MiniMaxH3VideoGroupNorm(norm_num_groups, block_out_channels[-1], eps=norm_eps, affine=True) + self.conv_out = MiniMaxH3VideoCausalConv3d( + block_out_channels[-1], + out_channels, + kernel_size=3, + spatial_padding=1, + temporal_padding=2, + spatial_padding_mode=spatial_padding_mode, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(hidden_states) + for down_block in self.down_blocks: + hidden_states = down_block(hidden_states) + return self.conv_out(F.silu(self.norm_out(hidden_states))) + + +class MiniMaxH3VideoRotaryPosEmbed(nn.Module): + def __init__(self, dim: int, theta: float = 100.0, num_axes: int = 3) -> None: + super().__init__() + if dim % (2 * num_axes) != 0: + raise ValueError(f"`dim` {dim} must be divisible by `2 * num_axes` {2 * num_axes}.") + inv_freq = 1.0 / theta ** torch.arange(0, 1, 2 * num_axes / dim, dtype=torch.float32) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + angles = 2.0 * math.pi * position_ids[:, :, :, None] * self.inv_freq[None, None, None, :] + angles = angles.flatten(2, 3).tile(2).unsqueeze(2) + return angles.cos(), angles.sin() + + +class MiniMaxH3VideoAttention(nn.Module): + def __init__(self, dim: int, heads: int, dim_head: int, eps: float = 1e-5, bias: bool = True) -> None: + super().__init__() + self.heads = heads + self.dim_head = dim_head + self.use_bias = bias + inner_dim = heads * dim_head + self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=False) + self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=False) + self.to_q = nn.Linear(dim, inner_dim, bias=bias) + self.to_k = nn.Linear(dim, inner_dim, bias=bias) + self.to_v = nn.Linear(dim, inner_dim, bias=bias) + self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=bias), nn.Dropout(0.0)]) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + query = self.to_q(hidden_states).unflatten(2, (self.heads, -1)) + key = self.to_k(hidden_states).unflatten(2, (self.heads, -1)) + value = self.to_v(hidden_states).unflatten(2, (self.heads, -1)) + + query = self.norm_q(query.float()).to(query.dtype) + key = self.norm_k(key.float()).to(key.dtype) + if rotary_emb is not None: + cos, sin = rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + rotary_dim = cos.shape[-1] + query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:] + key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:] + query_first, query_second = query_rotary.chunk(2, dim=-1) + key_first, key_second = key_rotary.chunk(2, dim=-1) + query_rotated = torch.cat([-query_second, query_first], dim=-1) + key_rotated = torch.cat([-key_second, key_first], dim=-1) + query = torch.cat([query_rotary * cos + query_rotated * sin, query_pass], dim=-1) + key = torch.cat([key_rotary * cos + key_rotated * sin, key_pass], dim=-1) + + query, key, value = (tensor.permute(0, 2, 1, 3) for tensor in (query, key, value)) + hidden_states = F.scaled_dot_product_attention(query, key, value) + hidden_states = hidden_states.permute(0, 2, 1, 3).flatten(2, 3) + return self.to_out[0](hidden_states) + + +class MiniMaxH3VideoSwiGLU(nn.Module): + def __init__(self, dim_in: int, dim_out: int, bias: bool = True) -> None: + super().__init__() + self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias) + self.activation = nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1) + return hidden_states * self.activation(gate) + + +class MiniMaxH3VideoFeedForward(nn.Module): + def __init__(self, dim: int, mult: int = 4, bias: bool = True) -> None: + super().__init__() + inner_dim = int(dim * mult) + self.net = nn.ModuleList( + [ + MiniMaxH3VideoSwiGLU(dim, inner_dim, bias=bias), + nn.Dropout(0.0), + nn.Linear(inner_dim, dim, bias=bias), + ] + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +class MiniMaxH3VideoTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + heads: int, + dim_head: int, + ffn_mult: int = 4, + eps: float = 1e-5, + bias: bool = True, + ) -> None: + super().__init__() + self.norm1 = nn.RMSNorm(dim, eps=eps, elementwise_affine=True) + self.attn = MiniMaxH3VideoAttention(dim=dim, heads=heads, dim_head=dim_head, eps=eps, bias=bias) + self.scale1 = nn.Parameter(torch.zeros(dim)) + self.norm2 = nn.RMSNorm(dim, eps=eps, elementwise_affine=True) + self.ff = MiniMaxH3VideoFeedForward(dim, mult=ffn_mult, bias=bias) + self.scale2 = nn.Parameter(torch.zeros(dim)) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + norm_hidden_states = self.norm1(hidden_states.float()).to(hidden_states.dtype) + hidden_states = hidden_states + self.attn(norm_hidden_states, rotary_emb) * self.scale1 + norm_hidden_states = self.norm2(hidden_states.float()).to(hidden_states.dtype) + return hidden_states + self.ff(norm_hidden_states) * self.scale2 + + +class MiniMaxH3VideoViTDecoder3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + patch_size: int, + patch_size_t: int, + num_layers: int, + num_attention_heads: int, + attention_head_dim: int, + num_register_tokens: int, + ffn_mult: int, + rope_theta: float, + rope_dim_ratio: float, + norm_eps: float, + ) -> None: + super().__init__() + dim = num_attention_heads * attention_head_dim + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.out_channels = out_channels + self.num_register_tokens = num_register_tokens + self.rope = MiniMaxH3VideoRotaryPosEmbed(int(attention_head_dim * rope_dim_ratio), theta=rope_theta) + self.proj_in = nn.Linear(in_channels, dim) + self.register_tokens = nn.Parameter(torch.zeros(1, num_register_tokens, dim)) + self.transformer_blocks = nn.ModuleList( + [ + MiniMaxH3VideoTransformerBlock( + dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + ffn_mult=ffn_mult, + eps=norm_eps, + ) + for _ in range(num_layers) + ] + ) + self.norm_out = nn.LayerNorm(dim, elementwise_affine=True, eps=norm_eps) + self.proj_out = nn.Linear(dim, out_channels * patch_size_t * patch_size * patch_size) + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape( + batch_size, + num_frames * height * width, + num_channels, + ) + hidden_states = self.proj_in(hidden_states) + num_patches = hidden_states.shape[1] + register_tokens = self.register_tokens.expand(batch_size, -1, -1) + cls_token = torch.zeros_like(hidden_states[:, :1, :]) + hidden_states = torch.cat([hidden_states, register_tokens, cls_token], dim=1) + + grids = [ + 2.0 * (torch.arange(0.5, size, dtype=torch.float32, device=hidden_states.device) / size) - 1.0 + for size in (num_frames, height, width) + ] + position_ids = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=-1).flatten(0, 2) + position_ids = position_ids.unsqueeze(0).expand(batch_size, -1, -1) + suffix_ids = position_ids.new_zeros((batch_size, self.num_register_tokens + 1, 3)) + rotary_emb = self.rope(torch.cat([position_ids, suffix_ids], dim=1)) + + for block in self.transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = checkpoint(block, hidden_states, rotary_emb, use_reentrant=False) + else: + hidden_states = block(hidden_states, rotary_emb) + + hidden_states = self.proj_out(self.norm_out(hidden_states))[:, :num_patches, :] + patch_size, patch_size_t = self.patch_size, self.patch_size_t + hidden_states = hidden_states.view( + batch_size, + num_frames, + height, + width, + self.out_channels, + patch_size_t, + patch_size, + patch_size, + ) + hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous() + return hidden_states.reshape( + batch_size, + self.out_channels, + num_frames * patch_size_t, + height * patch_size, + width * patch_size, + ) + + +class AutoencoderKLMiniMaxH3(nn.Module): + """MiniMax-H3 causal encoder and ViT decoder with exact release geometry.""" + + _supports_gradient_checkpointing = True + _no_split_modules = ["MiniMaxH3VideoResnetBlock3d", "MiniMaxH3VideoTransformerBlock"] + _repeated_blocks = ["MiniMaxH3VideoTransformerBlock"] + _keep_in_fp32_modules = ["encoder", "decoder", "quant_conv", "post_quant_conv"] + + def __init__(self, config: MiniMaxH3VideoVAEConfig) -> None: + super().__init__() + self.model_config = config + self.config = config.arch_config + arch = config.arch_config + self.latent_channels = int(arch.latent_channels) + self.spatial_compression_ratio = math.prod(arch.spatial_downsample_factors) + self.temporal_compression_ratio = math.prod(arch.temporal_downsample_factors) + + self.encoder = MiniMaxH3VideoEncoder3d( + in_channels=arch.in_channels, + out_channels=2 * arch.latent_channels, + block_out_channels=tuple(arch.block_out_channels), + layers_per_block=arch.layers_per_block, + spatial_downsample_factors=tuple(arch.spatial_downsample_factors), + temporal_downsample_factors=tuple(arch.temporal_downsample_factors), + norm_num_groups=arch.norm_num_groups, + norm_eps=arch.norm_eps, + spatial_padding_mode=arch.spatial_padding_mode, + ) + self.quant_conv = nn.Conv3d(2 * arch.latent_channels, 2 * arch.latent_channels, kernel_size=1) + self.post_quant_conv = nn.Conv3d(arch.latent_channels, arch.latent_channels, kernel_size=1) + self.decoder = MiniMaxH3VideoViTDecoder3d( + in_channels=arch.latent_channels, + out_channels=arch.out_channels, + patch_size=self.spatial_compression_ratio, + patch_size_t=self.temporal_compression_ratio, + num_layers=arch.decoder_num_layers, + num_attention_heads=arch.decoder_num_attention_heads, + attention_head_dim=arch.decoder_attention_head_dim, + num_register_tokens=arch.decoder_num_register_tokens, + ffn_mult=arch.decoder_ffn_mult, + rope_theta=arch.decoder_rope_theta, + rope_dim_ratio=arch.decoder_rope_dim_ratio, + norm_eps=arch.decoder_norm_eps, + ) + + self.frame_pre_padding = (-arch.clip_length) % self.temporal_compression_ratio + self.tokens_chunk_size = math.ceil(arch.clip_length / self.temporal_compression_ratio) + self.token_overlap = (-arch.token_drop) % self.tokens_chunk_size + self.frame_overlap = max( + self.token_overlap * self.temporal_compression_ratio - self.frame_pre_padding, + 0, + ) + self.use_slicing = False + self.use_tiling = config.use_tiling + self.tile_sample_min_height = config.tile_sample_min_height + self.tile_sample_min_width = config.tile_sample_min_width + self.tile_sample_min_overlap_height = config.tile_sample_min_overlap_height + self.tile_sample_min_overlap_width = config.tile_sample_min_overlap_width + + self.register_buffer( + "latents_mean", + torch.tensor(arch.latents_mean, dtype=torch.float32).view(1, -1, 1, 1, 1), + persistent=False, + ) + self.register_buffer( + "latents_std", + torch.tensor(arch.latents_std, dtype=torch.float32).view(1, -1, 1, 1, 1), + persistent=False, + ) + self.register_buffer( + "pixel_mean", + torch.tensor((0.485, 0.456, 0.406), dtype=torch.float32).view(1, -1, 1, 1, 1), + persistent=False, + ) + self.register_buffer( + "pixel_std", + torch.tensor((0.229, 0.224, 0.225), dtype=torch.float32).view(1, -1, 1, 1, 1), + persistent=False, + ) + # The released encoder and decoder stay in FP32 for both weights and compute. + self.float() + + def to(self, *args, **kwargs) -> "AutoencoderKLMiniMaxH3": + device, dtype, non_blocking, memory_format = torch._C._nn._parse_to(*args, **kwargs) + if dtype is None or dtype == torch.float32: + return super().to(*args, **kwargs) + if not torch.is_floating_point(torch.empty((), dtype=dtype)): + return super().to(*args, **kwargs) + + to_kwargs = { + "device": device, + "dtype": torch.float32, + "non_blocking": non_blocking, + } + if memory_format is not None: + to_kwargs["memory_format"] = memory_format + return super().to(**to_kwargs) + + def half(self) -> "AutoencoderKLMiniMaxH3": + return self + + def bfloat16(self) -> "AutoencoderKLMiniMaxH3": + return self + + def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor: + return (latents - self.latents_mean) / self.latents_std + + def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor: + return latents * self.latents_std + self.latents_mean + + def normalize_pixels(self, pixels: torch.Tensor) -> torch.Tensor: + return (pixels - self.pixel_mean) / self.pixel_std + + def denormalize_pixels(self, sample: torch.Tensor) -> torch.Tensor: + return sample * self.pixel_std + self.pixel_mean + + def enable_slicing(self) -> None: + self.use_slicing = True + + def disable_slicing(self) -> None: + self.use_slicing = False + + def enable_tiling( + self, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_min_overlap_height: int | None = None, + tile_sample_min_overlap_width: int | None = None, + ) -> None: + self.use_tiling = True + self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height + self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width + self.tile_sample_min_overlap_height = tile_sample_min_overlap_height or self.tile_sample_min_overlap_height + self.tile_sample_min_overlap_width = tile_sample_min_overlap_width or self.tile_sample_min_overlap_width + + def disable_tiling(self) -> None: + self.use_tiling = False + + def _split_tiles(self, length: int, tile_size: int, min_overlap: int) -> tuple[list[int], list[int], list[int]]: + if tile_size >= length: + return [0], [length], [] + num_tiles = math.ceil(length / tile_size) + while tile_size * num_tiles - min_overlap * (num_tiles - 1) - length < 0: + num_tiles += 1 + overlaps = [min_overlap] * (num_tiles - 1) + remaining = tile_size * num_tiles - sum(overlaps) - length + for index in range(remaining // self.spatial_compression_ratio): + overlaps[index % (num_tiles - 1)] += self.spatial_compression_ratio + tile_start_indices = [0] + for index in range(num_tiles - 1): + tile_start_indices.append(tile_start_indices[-1] + tile_size - overlaps[index]) + return tile_start_indices, [tile_size] * num_tiles, overlaps + + @staticmethod + def _blend(a: torch.Tensor, b: torch.Tensor, blend_extent: int, dim: int) -> torch.Tensor: + blend_extent = min(a.shape[dim], b.shape[dim], blend_extent) + positions = torch.arange(blend_extent, device=b.device, dtype=b.dtype) + shape = [1] * a.ndim + shape[dim] = blend_extent + weight_a = (1 - positions / blend_extent).view(shape) + weight_b = (positions / blend_extent).view(shape) + slice_a = [slice(None)] * a.ndim + slice_a[dim] = slice(-blend_extent, None) + slice_b = [slice(None)] * b.ndim + slice_b[dim] = slice(0, blend_extent) + blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b + if blend_extent == b.shape[dim]: + return blended + slice_rest = [slice(None)] * b.ndim + slice_rest[dim] = slice(blend_extent, None) + return torch.cat([blended, b[tuple(slice_rest)]], dim=dim) + + def _stitch_tiles( + self, + tiles: list[list[torch.Tensor]], + height_overlaps: list[int], + width_overlaps: list[int], + ) -> torch.Tensor: + result_rows = [] + for row_index, row in enumerate(tiles): + result_row = [] + for column_index, tile in enumerate(row): + if row_index > 0: + tile = self._blend(tiles[row_index - 1][column_index], tile, height_overlaps[row_index - 1], -2) + if column_index > 0: + tile = self._blend(row[column_index - 1], tile, width_overlaps[column_index - 1], -1) + if row_index < len(tiles) - 1: + tile = tile[..., : -height_overlaps[row_index], :] + if column_index < len(row) - 1: + tile = tile[..., :, : -width_overlaps[column_index]] + result_row.append(tile) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + def _encode_clip(self, x: torch.Tensor) -> torch.Tensor: + if not self.use_tiling: + return self.quant_conv(self.encoder(x)) + height, width = x.shape[-2:] + y_indices, y_lengths, y_overlaps = self._split_tiles( + height, + self.tile_sample_min_height, + self.tile_sample_min_overlap_height, + ) + x_indices, x_lengths, x_overlaps = self._split_tiles( + width, + self.tile_sample_min_width, + self.tile_sample_min_overlap_width, + ) + rows = [] + for y_position, y_length in zip(y_indices, y_lengths): + row = [] + for x_position, x_length in zip(x_indices, x_lengths): + tile = x[..., y_position:y_position + y_length, x_position:x_position + x_length] + row.append(self.quant_conv(self.encoder(tile))) + rows.append(row) + latent_y_overlaps = [overlap // self.spatial_compression_ratio for overlap in y_overlaps] + latent_x_overlaps = [overlap // self.spatial_compression_ratio for overlap in x_overlaps] + return self._stitch_tiles(rows, latent_y_overlaps, latent_x_overlaps) + + def _decode_clip(self, z: torch.Tensor) -> torch.Tensor: + if not self.use_tiling: + return self.decoder(self.post_quant_conv(z)) + height = z.shape[-2] * self.spatial_compression_ratio + width = z.shape[-1] * self.spatial_compression_ratio + y_indices, y_lengths, y_overlaps = self._split_tiles( + height, + self.tile_sample_min_height, + self.tile_sample_min_overlap_height, + ) + x_indices, x_lengths, x_overlaps = self._split_tiles( + width, + self.tile_sample_min_width, + self.tile_sample_min_overlap_width, + ) + ratio = self.spatial_compression_ratio + rows = [] + for y_position, y_length in zip(y_indices, y_lengths): + row = [] + for x_position, x_length in zip(x_indices, x_lengths): + tile = z[ + ..., + y_position // ratio:y_position // ratio + y_length // ratio, + x_position // ratio:x_position // ratio + x_length // ratio, + ] + row.append(self.decoder(self.post_quant_conv(tile))) + rows.append(row) + return self._stitch_tiles(rows, y_overlaps, x_overlaps) + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + clip_length = self.config.clip_length + num_frames = x.shape[2] + if num_frames % clip_length != 0: + pad_frames = x[:, :, -1:].repeat(1, 1, (-num_frames) % clip_length, 1, 1) + x = torch.cat([x, pad_frames], dim=2) + moments = torch.cat( + [ + self._encode_clip(x[:, :, index * clip_length:(index + 1) * clip_length]) + for index in range(x.shape[2] // clip_length) + ], + dim=2, + ) + if self.config.token_drop > 0: + moments = moments[:, :, :-self.config.token_drop] + return moments + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + tokens_chunk_size = self.tokens_chunk_size + token_drop = self.config.token_drop + temporal_ratio = self.temporal_compression_ratio + chunk_num_frames = tokens_chunk_size * temporal_ratio + num_tokens = z.shape[2] + token_drop + pad_tokens = (-num_tokens) % tokens_chunk_size + num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0) + if pad_tokens > 0: + z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2) + + decoded_chunks = [] + overlap = None + for index in range(num_chunks): + start = index * tokens_chunk_size + clip = self._decode_clip(z[:, :, start:start + tokens_chunk_size + self.token_overlap]) + for overlap_index in range(int(token_drop > 0) + 1): + frame_start = overlap_index * chunk_num_frames + chunk = clip[:, :, frame_start:frame_start + chunk_num_frames] + chunk = chunk[:, :, self.frame_pre_padding:] + if overlap_index == 0: + if overlap is not None: + chunk = self._blend(overlap, chunk, self.frame_overlap, dim=-3) + decoded_chunks.append(chunk) + else: + overlap = chunk + if overlap is not None: + decoded_chunks.append(overlap) + decoded = torch.cat(decoded_chunks, dim=2) + + if pad_tokens > 0: + intra_tail = self.config.clip_length % temporal_ratio + num_tokens_before_pad = z.shape[2] - pad_tokens + pad_frames = sum( + intra_tail + if intra_tail and (num_tokens_before_pad + offset) % tokens_chunk_size == 0 + else temporal_ratio + for offset in range(pad_tokens) + ) + decoded = decoded[:, :, :-pad_frames] + return decoded + + def encode( + self, + x: torch.Tensor, + return_dict: bool = True, + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: + if self.use_slicing and x.shape[0] > 1: + moments = torch.cat([self._encode(x_slice) for x_slice in x.split(1)]) + else: + moments = self._encode(x) + posterior = DiagonalGaussianDistribution(moments) + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def encode_keyframe( + self, + x: torch.Tensor, + return_dict: bool = True, + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: + """Encode one-frame conditioning inputs without video chunk padding.""" + if x.ndim != 5 or x.shape[2] != 1: + raise ValueError(f"`x` must contain exactly one video frame, got shape {tuple(x.shape)}.") + if self.use_slicing and x.shape[0] > 1: + moments = torch.cat([self._encode_clip(x_slice) for x_slice in x.split(1)]) + else: + moments = self._encode_clip(x) + posterior = DiagonalGaussianDistribution(moments) + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: + if self.use_slicing and z.shape[0] > 1: + decoded = torch.cat([self._decode(z_slice) for z_slice in z.split(1)]) + else: + decoded = self._decode(z) + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + generator: torch.Generator | None = None, + return_dict: bool = True, + ) -> DecoderOutput | tuple[torch.Tensor]: + posterior = self.encode(sample).latent_dist + latents = posterior.sample(generator=generator) if sample_posterior else posterior.mode() + return self.decode(latents, return_dict=return_dict) + + +EntryClass = AutoencoderKLMiniMaxH3 diff --git a/fastvideo/pipelines/basic/minimax_h3/__init__.py b/fastvideo/pipelines/basic/minimax_h3/__init__.py new file mode 100644 index 000000000..64af6a722 --- /dev/null +++ b/fastvideo/pipelines/basic/minimax_h3/__init__.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import ( + EntryClass, + MiniMaxH3ModularPipeline, + MiniMaxH3Ref2VAModularPipeline, +) +from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3Reference + +__all__ = [ + "EntryClass", + "MiniMaxH3ModularPipeline", + "MiniMaxH3Ref2VAModularPipeline", + "MiniMaxH3Reference", +] diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py new file mode 100644 index 000000000..a8bc11b9e --- /dev/null +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo composed pipelines for MiniMax H3.""" + +from __future__ import annotations + +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.basic.minimax_h3.stages import ( + MiniMaxH3AudioDecodingStage, + MiniMaxH3ConditioningStage, + MiniMaxH3DenoisingStage, + MiniMaxH3InputPreparationStage, + MiniMaxH3LatentPreparationStage, + MiniMaxH3VideoDecodingStage, +) +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase + + +class MiniMaxH3BasePipeline(ComposedPipelineBase): + """Shared loading and target-generation path for MiniMax H3.""" + + pipeline_config_cls: type[MiniMaxH3PipelineConfig] = MiniMaxH3PipelineConfig + _required_config_modules = [ + "text_encoder", + "tokenizer", + "processor", + "vae", + "audio_vae", + "transformer", + "scheduler", + "audio_scheduler", + ] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + del fastvideo_args + for module_name, modality, expected_shift in ( + ("scheduler", "video", 12.0), + ("audio_scheduler", "audio", 3.0), + ): + shift = getattr(self.get_module(module_name), "shift", None) + if shift is None or float(shift) != expected_shift: + raise ValueError(f"MiniMax-H3 {modality} scheduler must expose shift={expected_shift:g}, got {shift}.") + + def _add_stages(self, *, ref2va: bool) -> None: + transformer = self.get_module("transformer") + vae = self.get_module("vae") + audio_vae = self.get_module("audio_vae") + scheduler = self.get_module("scheduler") + audio_scheduler = self.get_module("audio_scheduler") + + self.add_stage( + "input_preparation_stage", + MiniMaxH3InputPreparationStage( + vae=vae, + audio_vae=audio_vae if ref2va else None, + ref2va=ref2va, + ), + ) + self.add_stage( + "conditioning_stage", + MiniMaxH3ConditioningStage( + conditioner=self.get_module("text_encoder"), + tokenizer=self.get_module("tokenizer"), + processor=self.get_module("processor"), + ref2va=ref2va, + ), + ) + self.add_stage( + "latent_preparation_stage", + MiniMaxH3LatentPreparationStage( + transformer=transformer, + vae=vae, + audio_vae=audio_vae, + scheduler=scheduler, + ref2va=ref2va, + ), + ) + self.add_stage( + "denoising_stage", + MiniMaxH3DenoisingStage( + transformer=transformer, + scheduler=scheduler, + audio_scheduler=audio_scheduler, + ), + ) + self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=vae, transformer=transformer)) + self.add_stage("audio_decoding_stage", MiniMaxH3AudioDecodingStage(audio_vae=audio_vae)) + + +class MiniMaxH3Pipeline(MiniMaxH3BasePipeline): + """One-request joint video/stereo-audio pipeline for T2VA and FL2VA.""" + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + del fastvideo_args + self._add_stages(ref2va=False) + + +class MiniMaxH3RefPipeline(MiniMaxH3BasePipeline): + """Ordered-reference joint video/stereo-audio pipeline for Ref2VA.""" + + _extra_config_module_map = {"transformer": "transformer_ref"} + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + del fastvideo_args + self._add_stages(ref2va=True) + + +class MiniMaxH3ModularPipeline(MiniMaxH3Pipeline): + """Public T2VA/FL2VA entry matching the official manifest class name.""" + + +class MiniMaxH3Ref2VAModularPipeline(MiniMaxH3RefPipeline): + """Public Ref2VA entry using the checkpoint's ``transformer_ref`` partition.""" + + +EntryClass = [MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline] + +__all__ = [ + "EntryClass", + "MiniMaxH3BasePipeline", + "MiniMaxH3ModularPipeline", + "MiniMaxH3Pipeline", + "MiniMaxH3Ref2VAModularPipeline", + "MiniMaxH3RefPipeline", +] diff --git a/fastvideo/pipelines/basic/minimax_h3/packing.py b/fastvideo/pipelines/basic/minimax_h3/packing.py new file mode 100644 index 000000000..4fdfd9eea --- /dev/null +++ b/fastvideo/pipelines/basic/minimax_h3/packing.py @@ -0,0 +1,492 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The MiniMax and HuggingFace Teams. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +import torch +from diffusers.utils.torch_utils import randn_tensor +from PIL import Image + +if TYPE_CHECKING: + from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference + +MINIMAX_H3_VIDEO_TAG = 0 +MINIMAX_H3_TEXT_TAG = 1 +MINIMAX_H3_AUDIO_TAG = 2 + +MINIMAX_H3_FPS = 24 +MINIMAX_H3_SHORT_EDGE = 768 +MINIMAX_H3_MAX_PIXELS = 768 * 1344 +MINIMAX_H3_CANVAS_MULTIPLE = 32 +MINIMAX_H3_MIN_ASPECT_RATIO = 1 / 4 +MINIMAX_H3_MAX_ASPECT_RATIO = 4 +MINIMAX_H3_MIN_DURATION = 5.0 +MINIMAX_H3_MAX_DURATION = 15.0 +MINIMAX_H3_FRAMES_PER_CHUNK = 17 +MINIMAX_H3_LATENTS_PER_CHUNK = 5 +MINIMAX_H3_TEXT_ENCODER_LAYER = 50 +MINIMAX_H3_VISION_START_TOKEN = "<|vision_start|>" +MINIMAX_H3_IMAGE_PAD_TOKEN = "<|image_pad|>" +MINIMAX_H3_VIDEO_PAD_TOKEN = "<|video_pad|>" +MINIMAX_H3_VISION_END_TOKEN = "<|vision_end|>" +MINIMAX_H3_AUDIO_LATENTS_PER_SECOND = 40 +MINIMAX_H3_AUDIO_CHANNELS = 2 +MINIMAX_H3_KEYFRAME_NOISE_AUG = 0.999 +MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42 + +MINIMAX_H3_ROPE_FRAME_RESCALE = 5.0 / 3.0 +MINIMAX_H3_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) +_ROPE_SPATIAL_SCALE = 32 + + +@dataclass(frozen=True) +class MiniMaxH3PackedLayout: + """One packed joint sequence and the geometry needed to interpret it.""" + + sequence_length: int + position_ids: torch.Tensor + token_tags: torch.Tensor + video_indices: torch.Tensor + audio_indices: torch.Tensor + text_indices: torch.Tensor + num_condition_video_rows: int + num_condition_audio_rows: int + num_video_latent_frames: int + latent_height: int + latent_width: int + num_audio_latents: int + + +def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: + if aspect_width <= 0 or aspect_height <= 0: + raise ValueError(f"The aspect ratio must be positive, got {aspect_width}:{aspect_height}.") + ratio = aspect_width / aspect_height + if not MINIMAX_H3_MIN_ASPECT_RATIO <= ratio <= MINIMAX_H3_MAX_ASPECT_RATIO: + raise ValueError(f"MiniMax-H3 supports aspect ratios from 1:4 to 4:1, got {aspect_width}:{aspect_height}.") + + if ratio >= 1: + width, height = MINIMAX_H3_SHORT_EDGE * ratio, float(MINIMAX_H3_SHORT_EDGE) + else: + width, height = float(MINIMAX_H3_SHORT_EDGE), MINIMAX_H3_SHORT_EDGE / ratio + area = width * height + if area > MINIMAX_H3_MAX_PIXELS: + scale = (MINIMAX_H3_MAX_PIXELS / area)**0.5 + width, height = width * scale, height * scale + multiple = MINIMAX_H3_CANVAS_MULTIPLE + return max(multiple, round(height / multiple) * multiple), max(multiple, round(width / multiple) * multiple) + + +def align_num_frames(num_frames: int) -> int: + if num_frames < 1: + raise ValueError(f"`num_frames` must be positive, got {num_frames}.") + while num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK: + num_frames += 1 + return num_frames + + +def video_latent_num_frames(num_frames: int) -> int: + if num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK: + raise ValueError(f"`num_frames` must be of the form 17 * n + 5, got {num_frames}.") + return (num_frames - MINIMAX_H3_LATENTS_PER_CHUNK) // MINIMAX_H3_FRAMES_PER_CHUNK * MINIMAX_H3_LATENTS_PER_CHUNK + 2 + + +def audio_latent_num_frames(num_frames: int) -> int: + return int(round(num_frames / MINIMAX_H3_FPS * MINIMAX_H3_AUDIO_LATENTS_PER_SECOND)) + + +def prepare_keyframe_image(image: Image.Image, height: int, width: int, stretch: bool) -> Image.Image: + if image.size == (width, height): + return image + if stretch: + return image.resize((width, height), Image.Resampling.LANCZOS) + scale = max(width / image.size[0], height / image.size[1]) + resized_size = (max(width, round(image.size[0] * scale)), max(height, round(image.size[1] * scale))) + left = max(0, (resized_size[0] - width) // 2) + top = max(0, (resized_size[1] - height) // 2) + resized = image.resize(resized_size, Image.Resampling.LANCZOS) + return resized.crop((left, top, left + width, top + height)) + + +def patchify_video_latents(latents: torch.Tensor, patch_size: tuple[int, int, int]) -> torch.Tensor: + patch_t, patch_h, patch_w = patch_size + batch_size, channels, num_frames, height, width = latents.shape + if num_frames % patch_t or height % patch_h or width % patch_w: + raise ValueError(f"Latents of shape {tuple(latents.shape)} are not divisible by the patch {patch_size}.") + latents = latents.reshape( + batch_size, + channels, + num_frames // patch_t, + patch_t, + height // patch_h, + patch_h, + width // patch_w, + patch_w, + ) + # Channel-major patch features are a checkpoint contract; do not use einops's common patch-major order. + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7) + return latents.reshape(-1, channels * patch_t * patch_h * patch_w).contiguous() + + +def unpatchify_video_tokens( + rows: torch.Tensor, + num_latent_frames: int, + latent_height: int, + latent_width: int, + channels: int, + patch_size: tuple[int, int, int], +) -> torch.Tensor: + patch_t, patch_h, patch_w = patch_size + rows = rows.reshape( + -1, + num_latent_frames // patch_t, + latent_height // patch_h, + latent_width // patch_w, + channels, + patch_t, + patch_h, + patch_w, + ) + rows = rows.permute(0, 4, 1, 5, 2, 6, 3, 7) + return rows.reshape(-1, channels, num_latent_frames, latent_height, latent_width).contiguous() + + +def unpack_audio_tokens(rows: torch.Tensor, num_audio_latents: int) -> torch.Tensor: + rows = rows.reshape(MINIMAX_H3_AUDIO_CHANNELS, num_audio_latents, rows.shape[-1]) + return rows.permute(0, 2, 1).contiguous() + + +def spatial_position_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: + ratio = dim / sqrt_area + left = (1.0 - ratio) / 2.0 + # NumPy endpoint=False and float64 are part of the reference RoPE coordinate arithmetic. + grid = np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE + return torch.from_numpy(grid).to(torch.float64) + + +def temporal_position_grid(num_latent_frames: int, origin: float) -> torch.Tensor: + spans = torch.tensor( + [ + MINIMAX_H3_ROPE_FRAME_RESCALE * + MINIMAX_H3_ROPE_FRAMES_PER_LATENT[index % len(MINIMAX_H3_ROPE_FRAMES_PER_LATENT)] + for index in range(num_latent_frames) + ], + dtype=torch.float64, + ) + return origin + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)]) + + +def _temporal_position_span(num_latent_frames: int) -> float: + # Preserve NumPy's pairwise summation: sequential torch summation diverges in the final ulp for longer clips. + spans = np.ones(num_latent_frames, dtype=np.float64) * MINIMAX_H3_ROPE_FRAME_RESCALE + for index, frames in enumerate(MINIMAX_H3_ROPE_FRAMES_PER_LATENT): + spans[index::len(MINIMAX_H3_ROPE_FRAMES_PER_LATENT)] *= frames + return float(spans.sum()) + + +def build_packed_sequence( + text_token_tags: torch.Tensor, + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + patch_size: tuple[int, int, int], + keyframe_anchors: tuple[str, ...] = (), +) -> MiniMaxH3PackedLayout: + if text_token_tags.ndim != 1: + raise ValueError(f"text_token_tags must be one-dimensional, got {tuple(text_token_tags.shape)}.") + valid_text_tags = (text_token_tags == MINIMAX_H3_TEXT_TAG) | (text_token_tags == MINIMAX_H3_VIDEO_TAG) + if not bool(valid_text_tags.all()): + raise ValueError("text_token_tags may contain only text and vision tags; semantic padding is not allowed.") + + _, patch_h, patch_w = patch_size + rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w) + num_text_tokens = text_token_tags.shape[0] + num_condition_rows = len(keyframe_anchors) * rows_per_frame + num_audio_rows = num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS + num_video_rows = num_latent_frames * rows_per_frame + sequence_length = num_text_tokens + num_condition_rows + num_audio_rows + num_video_rows + + condition_start = num_text_tokens + audio_start = condition_start + num_condition_rows + video_start = audio_start + num_audio_rows + position_ids = torch.zeros(sequence_length, 3, dtype=torch.float64) + position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) + + sqrt_area = np.sqrt(latent_height * latent_width) + height_grid = spatial_position_grid(latent_height, patch_h, sqrt_area) + width_grid = spatial_position_grid(latent_width, patch_w, sqrt_area) + frame_grid = torch.stack([axis.reshape(-1) for axis in torch.meshgrid(height_grid, width_grid, indexing="ij")], -1) + + for index, anchor in enumerate(keyframe_anchors): + if anchor == "first": + anchor_time = float(num_text_tokens) + elif anchor == "last": + anchor_time = (float(num_text_tokens) + _temporal_position_span(num_latent_frames) - + MINIMAX_H3_ROPE_FRAME_RESCALE) + else: + raise ValueError(f"A keyframe anchor must be 'first' or 'last', got {anchor!r}.") + rows = slice(condition_start + index * rows_per_frame, condition_start + (index + 1) * rows_per_frame) + position_ids[rows, 0] = anchor_time + position_ids[rows, 1:] = frame_grid + + audio_time = float(num_text_tokens) + torch.arange(num_audio_latents, dtype=torch.float64) + position_ids[audio_start:video_start, 0] = audio_time.repeat(MINIMAX_H3_AUDIO_CHANNELS) + position_ids[audio_start:video_start, 2] = torch.cat([ + torch.full((num_audio_latents, ), float(width_grid[0]), dtype=torch.float64), + torch.full((num_audio_rows - num_audio_latents, ), float(width_grid[-1]), dtype=torch.float64), + ]) + + video_positions = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64) + video_positions[:, :, 0] = temporal_position_grid(num_latent_frames, float(num_text_tokens))[:, None] + video_positions[:, :, 1:] = frame_grid[None] + position_ids[video_start:] = video_positions.reshape(-1, 3) + + video_indices = torch.cat([torch.arange(condition_start, audio_start), torch.arange(video_start, sequence_length)]) + audio_indices = torch.arange(audio_start, video_start) + text_indices = torch.arange(num_text_tokens) + token_tags = torch.empty(sequence_length, dtype=torch.long) + token_tags[text_indices] = text_token_tags.to(torch.long) + token_tags[audio_indices] = MINIMAX_H3_AUDIO_TAG + token_tags[video_indices] = MINIMAX_H3_VIDEO_TAG + return MiniMaxH3PackedLayout( + sequence_length=sequence_length, + position_ids=position_ids, + token_tags=token_tags, + video_indices=video_indices, + audio_indices=audio_indices, + text_indices=text_indices, + num_condition_video_rows=num_condition_rows, + num_condition_audio_rows=0, + num_video_latent_frames=num_latent_frames, + latent_height=latent_height, + latent_width=latent_width, + num_audio_latents=num_audio_latents, + ) + + +def _reference_temporal_span(num_latent_frames: int) -> float: + """Advance the Ref2VA clock using its required sequential summation order.""" + return sum(MINIMAX_H3_ROPE_FRAME_RESCALE * + MINIMAX_H3_ROPE_FRAMES_PER_LATENT[index % len(MINIMAX_H3_ROPE_FRAMES_PER_LATENT)] + for index in range(num_latent_frames)) + + +def _frame_position_grid( + latent_height: int, + latent_width: int, + patch_h: int, + patch_w: int, +) -> tuple[torch.Tensor, torch.Tensor]: + sqrt_area = np.sqrt(latent_height * latent_width) + height_grid = spatial_position_grid(latent_height, patch_h, sqrt_area) + width_grid = spatial_position_grid(latent_width, patch_w, sqrt_area) + grids = torch.meshgrid(height_grid, width_grid, indexing="ij") + return torch.stack([grid.reshape(-1) for grid in grids], dim=-1), width_grid + + +def _fill_audio_positions( + position_ids: torch.Tensor, + rows: slice, + num_audio_latents: int, + rotary_time: float, + width_grid: torch.Tensor, +) -> None: + time = rotary_time + torch.arange(num_audio_latents, dtype=torch.float64) + position_ids[rows, 0] = time.repeat(MINIMAX_H3_AUDIO_CHANNELS) + position_ids[rows, 2] = torch.cat([ + torch.full((num_audio_latents, ), float(width_grid[0]), dtype=torch.float64), + torch.full((num_audio_latents, ), float(width_grid[-1]), dtype=torch.float64), + ]) + + +def _num_video_rows(reference: MiniMaxH3PreparedReference, patch_size: tuple[int, int, int]) -> int: + patch_t, patch_h, patch_w = patch_size + geometry = (reference.num_latent_frames, reference.latent_height, reference.latent_width) + if any(value <= 0 for value in geometry): + raise ValueError(f"Incomplete visual reference geometry: {geometry}.") + if any(value % patch for value, patch in zip(geometry, patch_size, strict=True)): + raise ValueError(f"Visual reference geometry {geometry} is not divisible by patch {patch_size}.") + return (reference.num_latent_frames // patch_t) * (reference.latent_height // patch_h) * (reference.latent_width // + patch_w) + + +def build_ref2va_packed_sequence( + text_token_tags: torch.Tensor, + references: list[MiniMaxH3PreparedReference], + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + patch_size: tuple[int, int, int], +) -> MiniMaxH3PackedLayout: + """Build `[text | ordered references | target audio | target video]`.""" + if patch_size != (1, 2, 2): + raise ValueError(f"MiniMax-H3 Ref2VA requires patch_size=(1, 2, 2), got {patch_size}.") + if text_token_tags.ndim != 1: + raise ValueError(f"text_token_tags must be one-dimensional, got {tuple(text_token_tags.shape)}.") + valid_text_tags = (text_token_tags == MINIMAX_H3_TEXT_TAG) | (text_token_tags == MINIMAX_H3_VIDEO_TAG) + if not bool(valid_text_tags.all()): + raise ValueError("text_token_tags may contain only text and vision tags.") + if not references: + raise ValueError("Ref2VA requires at least one prepared reference.") + + patch_t, patch_h, patch_w = patch_size + target_geometry = (num_latent_frames, latent_height, latent_width) + if any(value <= 0 for value in target_geometry) or num_audio_latents <= 0: + raise ValueError("Ref2VA target latent geometry must be positive.") + if any(value % patch for value, patch in zip(target_geometry, patch_size, strict=True)): + raise ValueError(f"Target geometry {target_geometry} is not divisible by patch {patch_size}.") + + visual_row_counts = [ + 0 if reference.media_type == "audio" else _num_video_rows(reference, patch_size) for reference in references + ] + for reference in references: + if reference.media_type == "audio" and not reference.has_audio: + raise ValueError("An audio reference must carry decoded waveform latents.") + if reference.has_audio and reference.num_audio_latents <= 0: + raise ValueError("An audio-bearing reference has no resolved audio latents.") + + num_text_tokens = int(text_token_tags.shape[0]) + num_target_video_rows = (num_latent_frames // patch_t) * (latent_height // patch_h) * (latent_width // patch_w) + num_target_audio_rows = num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS + num_reference_video_rows = sum(visual_row_counts) + num_reference_audio_rows = sum(reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS for reference in references + if reference.has_audio) + sequence_length = (num_text_tokens + num_reference_video_rows + num_reference_audio_rows + num_target_audio_rows + + num_target_video_rows) + + position_ids = torch.zeros(sequence_length, 3, dtype=torch.float64) + position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) + target_frame_grid, target_width_grid = _frame_position_grid(latent_height, latent_width, patch_h, patch_w) + + video_indices: list[torch.Tensor] = [] + audio_indices: list[torch.Tensor] = [] + cursor = num_text_tokens + rotary_time = float(num_text_tokens) + for reference, visual_row_count in zip(references, visual_row_counts, strict=True): + if reference.media_type == "image": + rows = slice(cursor, cursor + visual_row_count) + cursor = rows.stop + video_indices.append(torch.arange(rows.start, rows.stop)) + frame_grid, _ = _frame_position_grid( + reference.latent_height, + reference.latent_width, + patch_h, + patch_w, + ) + position_ids[rows, 0] = rotary_time + position_ids[rows, 1:] = frame_grid + rotary_time += 1.0 + elif reference.media_type == "audio": + count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS + rows = slice(cursor, cursor + count) + cursor = rows.stop + audio_indices.append(torch.arange(rows.start, rows.stop)) + _fill_audio_positions(position_ids, rows, reference.num_audio_latents, rotary_time, target_width_grid) + rotary_time += float(reference.num_audio_latents) + elif reference.media_type == "video": + audio_count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS if reference.has_audio else 0 + audio_rows = slice(cursor, cursor + audio_count) + video_rows = slice(audio_rows.stop, audio_rows.stop + visual_row_count) + cursor = video_rows.stop + if audio_count: + audio_indices.append(torch.arange(audio_rows.start, audio_rows.stop)) + video_indices.append(torch.arange(video_rows.start, video_rows.stop)) + + frame_grid, width_grid = _frame_position_grid( + reference.latent_height, + reference.latent_width, + patch_h, + patch_w, + ) + if audio_count: + _fill_audio_positions( + position_ids, + audio_rows, + reference.num_audio_latents, + rotary_time, + width_grid, + ) + frame_time = temporal_position_grid(reference.num_latent_frames, rotary_time) + rows_per_frame = frame_grid.shape[0] + position_ids[video_rows, 0] = frame_time.repeat_interleave(rows_per_frame) + position_ids[video_rows, 1:] = frame_grid.repeat(reference.num_latent_frames // patch_t, 1) + rotary_time += max( + float(reference.num_audio_latents if reference.has_audio else 0), + _reference_temporal_span(reference.num_latent_frames), + ) + else: + raise ValueError(f"Unsupported prepared reference type: {reference.media_type!r}.") + + audio_start = cursor + video_start = audio_start + num_target_audio_rows + _fill_audio_positions(position_ids, slice(audio_start, video_start), num_audio_latents, rotary_time, + target_width_grid) + frame_time = temporal_position_grid(num_latent_frames, rotary_time) + position_ids[video_start:, 0] = frame_time.repeat_interleave(target_frame_grid.shape[0]) + position_ids[video_start:, 1:] = target_frame_grid.repeat(num_latent_frames // patch_t, 1) + + video_indices.append(torch.arange(video_start, sequence_length)) + audio_indices.append(torch.arange(audio_start, video_start)) + packed_video_indices = torch.cat(video_indices) + packed_audio_indices = torch.cat(audio_indices) + text_indices = torch.arange(num_text_tokens) + token_tags = torch.empty(sequence_length, dtype=torch.long) + token_tags[text_indices] = text_token_tags.to(torch.long) + token_tags[packed_audio_indices] = MINIMAX_H3_AUDIO_TAG + token_tags[packed_video_indices] = MINIMAX_H3_VIDEO_TAG + + return MiniMaxH3PackedLayout( + sequence_length=sequence_length, + position_ids=position_ids, + token_tags=token_tags, + video_indices=packed_video_indices, + audio_indices=packed_audio_indices, + text_indices=text_indices, + num_condition_video_rows=num_reference_video_rows, + num_condition_audio_rows=num_reference_audio_rows, + num_video_latent_frames=num_latent_frames, + latent_height=latent_height, + latent_width=latent_width, + num_audio_latents=num_audio_latents, + ) + + +def build_row_timesteps( + layout: MiniMaxH3PackedLayout, + video_timestep: float, + audio_timestep: float, + condition_video_timestep: float, + condition_audio_timestep: float, +) -> tuple[torch.Tensor, torch.Tensor]: + row_timesteps = torch.full((layout.sequence_length, ), video_timestep, dtype=torch.float32) + row_timesteps[layout.video_indices[:layout.num_condition_video_rows]] = condition_video_timestep + row_timesteps[layout.audio_indices[layout.num_condition_audio_rows:]] = audio_timestep + row_timesteps[layout.audio_indices[:layout.num_condition_audio_rows]] = condition_audio_timestep + return torch.unique(row_timesteps, sorted=True, return_inverse=True) + + +def keyframe_condition_noise( + condition_latent_shapes: tuple[tuple[int, int, int], ...], + patch_size: tuple[int, int, int], + latent_channels: int, + generator: torch.Generator | list[torch.Generator] | None = None, + device: torch.device | None = None, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + rows = [] + for num_latent_frames, latent_height, latent_width in condition_latent_shapes: + noise = randn_tensor( + (1, latent_channels, num_latent_frames, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + rows.append(patchify_video_latents(noise, patch_size)) + if not rows: + return torch.empty((0, latent_channels * int(np.prod(patch_size))), device=device, dtype=dtype) + return torch.cat(rows) diff --git a/fastvideo/pipelines/basic/minimax_h3/presets.py b/fastvideo/pipelines/basic/minimax_h3/presets.py new file mode 100644 index 000000000..0412d951f --- /dev/null +++ b/fastvideo/pipelines/basic/minimax_h3/presets.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inference presets for MiniMax H3 joint video/audio generation.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Joint video/audio flow-matching denoising", + allowed_overrides=frozenset({"num_inference_steps"}), +) + +_SHARED_DEFAULTS = { + "fps": 24, + "guidance_scale": 1.0, + "batch_cfg": False, + "negative_prompt": "", + "num_inference_steps": 50, + "seed": 0, +} + +MINIMAX_H3_T2VA = InferencePreset( + name="minimax_h3_t2va", + version=1, + model_family="minimax_h3", + description="MiniMax H3 text-to-video with synchronized stereo audio at 768p", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + **_SHARED_DEFAULTS, + "height": 768, + "width": 1344, + "num_frames": 124, + }, +) + +MINIMAX_H3_FL2VA = InferencePreset( + name="minimax_h3_fl2va", + version=1, + model_family="minimax_h3", + description="MiniMax H3 first/last-frame-to-video with synchronized stereo audio", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + **_SHARED_DEFAULTS, + "height": 768, + "width": 1344, + "num_frames": 192, + }, +) + +MINIMAX_H3_REF2VA = InferencePreset( + name="minimax_h3_ref2va", + version=1, + model_family="minimax_h3", + description="MiniMax H3 ordered image/video/audio references to joint video and stereo audio", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + **_SHARED_DEFAULTS, + "height": 768, + "width": 1344, + "num_frames": 124, + }, +) + +ALL_PRESETS = (MINIMAX_H3_T2VA, MINIMAX_H3_FL2VA, MINIMAX_H3_REF2VA) + +__all__ = [ + "ALL_PRESETS", + "MINIMAX_H3_FL2VA", + "MINIMAX_H3_REF2VA", + "MINIMAX_H3_T2VA", +] diff --git a/fastvideo/pipelines/basic/minimax_h3/reference.py b/fastvideo/pipelines/basic/minimax_h3/reference.py new file mode 100644 index 000000000..55a35712a --- /dev/null +++ b/fastvideo/pipelines/basic/minimax_h3/reference.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 reference inputs and media preparation helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import math +from numbers import Integral, Real +import os +from typing import Any, Literal + +import numpy as np +from PIL import Image, ImageOps +import torch + +from fastvideo.models.vision_utils import load_image +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_AUDIO_CHANNELS, + MINIMAX_H3_CANVAS_MULTIPLE, + MINIMAX_H3_FPS, + MINIMAX_H3_FRAMES_PER_CHUNK, + MINIMAX_H3_LATENTS_PER_CHUNK, + resolve_canvas_size, +) + +MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE = 2048 +MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS = 2.0 +MINIMAX_H3_QWEN_TEMPORAL_PATCH = 2 +MINIMAX_H3_MAX_REFERENCE_IMAGES = 9 +MINIMAX_H3_MAX_REFERENCE_VIDEOS = 3 +MINIMAX_H3_MAX_REFERENCE_AUDIOS = 3 +MINIMAX_H3_MAX_REFERENCES = 12 + + +@dataclass(frozen=True) +class MiniMaxH3Reference: + """One ordered Ref2VA input; decoding is deferred until request preparation.""" + + source: Any + media_type: Literal["image", "video", "audio"] = "image" + soundtrack: Any | None = None + fps: float | None = None + sample_rate: int | None = None + + def __post_init__(self) -> None: + if self.source is None: + raise ValueError("A MiniMax-H3 reference requires a media source.") + if self.media_type not in ("image", "video", "audio"): + raise ValueError(f"Unsupported MiniMax-H3 reference type: {self.media_type!r}.") + if self.soundtrack is not None and self.media_type != "video": + raise ValueError("Only a video reference may carry a separate soundtrack.") + if self.fps is not None and (self.media_type != "video" or isinstance(self.fps, bool) + or not isinstance(self.fps, Real) or self.fps <= 0): + raise ValueError("Reference `fps` must be positive and is only valid for video.") + if self.sample_rate is not None and (self.media_type == "image" or isinstance(self.sample_rate, bool) + or not isinstance(self.sample_rate, Integral) or self.sample_rate <= 0): + raise ValueError("Reference `sample_rate` must be positive and is only valid for audio-bearing media.") + + +@dataclass +class MiniMaxH3PreparedReference: + """Decoded Ref2VA media and the latent geometry resolved by condition encoding.""" + + media_type: Literal["image", "video", "audio"] + has_audio: bool = False + image: Image.Image | None = None + frames: np.ndarray | None = None + waveform: torch.Tensor | None = None + block_timestamps: list[float] = field(default_factory=list) + num_latent_frames: int = 1 + latent_height: int = 0 + latent_width: int = 0 + num_audio_latents: int = 0 + + +def validate_references(references: list[Any]) -> list[MiniMaxH3Reference]: + """Validate the ordered Ref2VA media list and its per-modality limits.""" + if not references: + raise ValueError("MiniMax-H3 Ref2VA requires at least one reference.") + if not all(isinstance(reference, MiniMaxH3Reference) for reference in references): + raise TypeError("Every Ref2VA entry must be a MiniMaxH3Reference.") + + typed = list(references) + counts = { + media_type: sum(reference.media_type == media_type for reference in typed) + for media_type in ("image", "video", "audio") + } + limits = { + "image": MINIMAX_H3_MAX_REFERENCE_IMAGES, + "video": MINIMAX_H3_MAX_REFERENCE_VIDEOS, + "audio": MINIMAX_H3_MAX_REFERENCE_AUDIOS, + } + for media_type, limit in limits.items(): + if counts[media_type] > limit: + raise ValueError(f"MiniMax-H3 accepts at most {limit} {media_type} references.") + if len(typed) > MINIMAX_H3_MAX_REFERENCES: + raise ValueError(f"MiniMax-H3 accepts at most {MINIMAX_H3_MAX_REFERENCES} references.") + if counts["audio"] == len(typed): + raise ValueError("Audio references must be paired with at least one image or video reference.") + return typed + + +def _import_av() -> Any: + try: + import av + except ImportError as error: + raise ImportError("Decoding MiniMax-H3 video or audio references requires PyAV.") from error + return av + + +def _decode_audio_stream(av_module: Any, container: Any, stream: Any) -> tuple[torch.Tensor, int]: + sample_rate = int(stream.codec_context.sample_rate) + resampler = av_module.audio.resampler.AudioResampler(format="fltp", layout=stream.layout, rate=sample_rate) + chunks: list[torch.Tensor] = [] + for frame in container.decode(stream): + chunks.extend(torch.from_numpy(item.to_ndarray()) for item in resampler.resample(frame)) + chunks.extend(torch.from_numpy(item.to_ndarray()) for item in resampler.resample(None)) + if not chunks: + raise ValueError("The reference audio stream contains no samples.") + return torch.cat(chunks, dim=-1).to(torch.float32), sample_rate + + +def decode_reference_audio(source: str | os.PathLike[str]) -> tuple[torch.Tensor, int]: + """Decode the first audio stream without imposing H3's target sample rate.""" + av_module = _import_av() + with av_module.open(str(source)) as container: + if not container.streams.audio: + raise ValueError(f"Reference media {source!s} contains no audio stream.") + return _decode_audio_stream(av_module, container, container.streams.audio[0]) + + +def decode_reference_video(source: str | os.PathLike[str]) -> tuple[np.ndarray, float, tuple[torch.Tensor, int] | None]: + """Decode RGB frames, display rotation, frame rate, and an optional soundtrack.""" + av_module = _import_av() + with av_module.open(str(source)) as container: + if not container.streams.video: + raise ValueError(f"Reference media {source!s} contains no video stream.") + stream = container.streams.video[0] + frames = [] + rotation = 0.0 + for frame in container.decode(stream): + rotation = float(getattr(frame, "rotation", 0.0) or 0.0) + frames.append(frame.to_ndarray(format="rgb24")) + if not frames: + raise ValueError(f"Reference media {source!s} contains no video frames.") + rate = stream.average_rate or getattr(stream, "guessed_rate", None) + if rate is None: + raise ValueError(f"Reference media {source!s} does not expose a frame rate.") + + soundtrack = None + if container.streams.audio: + container.seek(0) + soundtrack = _decode_audio_stream(av_module, container, container.streams.audio[0]) + + result = np.stack(frames) + turns = round(rotation / 90.0) % 4 + if turns: + result = np.ascontiguousarray(np.rot90(result, k=-turns, axes=(1, 2))) + return result, float(rate), soundtrack + + +def resolve_reference_image_size(width: int, height: int) -> tuple[int, int]: + """Resolve the released 2048-short-edge reference-image canvas.""" + if width <= 0 or height <= 0: + raise ValueError(f"A reference image must have a positive size, got {width}x{height}.") + if width > 4 * height or height > 4 * width: + raise ValueError(f"A reference image must be within 1:4 and 4:1, got {width}x{height}.") + scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height) + multiple = MINIMAX_H3_CANVAS_MULTIPLE + return ( + max(multiple, + round(height * scale / multiple) * multiple), + max(multiple, + round(width * scale / multiple) * multiple), + ) + + +def reference_media_to_uint8(media: Any) -> np.ndarray: + """Normalize PIL, tensor, NumPy, or frame-list media to channel-last uint8.""" + if isinstance(media, list): + if not media: + raise ValueError("A reference video must contain at least one frame.") + return np.stack([reference_media_to_uint8(item) for item in media]) + if isinstance(media, Image.Image): + return np.asarray(media.convert("RGB")) + if isinstance(media, torch.Tensor): + media = media.movedim(-3, -1).detach().cpu().numpy() + media = np.asarray(media) + if media.dtype != np.uint8: + media = (media * 255.0).round().clip(0, 255).astype(np.uint8) + return media + + +def prepare_reference_image(image: Image.Image, height: int, width: int) -> Image.Image: + """Resize an RGB reference image to its resolved condition canvas.""" + return image if image.size == (width, height) else image.resize((width, height), Image.Resampling.LANCZOS) + + +def resample_reference_frames(frames: np.ndarray, fps: float) -> np.ndarray: + """Nearest-resample reference frames onto H3's fixed 24-fps timeline.""" + if frames.ndim != 4 or frames.shape[-1] != 3 or frames.shape[0] == 0: + raise ValueError(f"A reference video must be non-empty RGB frames, got {tuple(frames.shape)}.") + if fps <= 0: + raise ValueError(f"A reference video must have a positive frame rate, got {fps}.") + if fps == MINIMAX_H3_FPS: + return frames + scale = MINIMAX_H3_FPS / fps + slots = np.floor(np.arange(frames.shape[0]) * scale + 0.5).astype(np.int64) + repeats = np.diff(slots, append=math.floor(frames.shape[0] * scale + 0.5)) + return np.repeat(frames, repeats, axis=0) + + +def prepare_reference_frames(frames: np.ndarray, num_frames: int) -> np.ndarray: + """Trim a reference video to the request and resize every RGB frame.""" + if frames.ndim != 4 or frames.shape[-1] != 3 or frames.shape[0] == 0: + raise ValueError(f"A reference video must be non-empty RGB frames, got {tuple(frames.shape)}.") + frames = frames[:num_frames] + height, width = resolve_canvas_size(frames.shape[2], frames.shape[1]) + if frames.shape[1:3] == (height, width): + return frames + return np.stack( + [np.asarray(Image.fromarray(frame).resize((width, height), Image.Resampling.LANCZOS)) for frame in frames]) + + +def sample_reference_video_frames(frames: np.ndarray) -> tuple[list[np.ndarray], list[float]]: + """Sample the released 2-fps Qwen presentation and its block timestamps.""" + if frames.ndim != 4 or frames.shape[0] == 0: + raise ValueError("A prepared reference video must contain frames.") + stride = MINIMAX_H3_FPS / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS + indices: list[int] = [] + cursor = 0.0 + while round(cursor) < frames.shape[0]: + if not indices or round(cursor) > indices[-1]: + indices.append(round(cursor)) + cursor += stride + timestamps = [index / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS for index in range(len(indices))] + timestamps += [timestamps[-1]] * (-len(timestamps) % MINIMAX_H3_QWEN_TEMPORAL_PATCH) + block_timestamps = [(timestamps[index] + timestamps[index + MINIMAX_H3_QWEN_TEMPORAL_PATCH - 1]) / 2 + for index in range(0, len(timestamps), MINIMAX_H3_QWEN_TEMPORAL_PATCH)] + return [frames[index] for index in indices], block_timestamps + + +def prepare_reference_waveform( + waveform: torch.Tensor, + sample_rate: int, + target_sample_rate: int, + max_duration: float, +) -> torch.Tensor: + """Normalize reference audio to stereo float32 at the audio-VAE rate.""" + waveform = torch.as_tensor(waveform).detach().cpu() + if waveform.ndim != 2 or waveform.shape[0] not in (1, MINIMAX_H3_AUDIO_CHANNELS): + raise ValueError( + f"A reference soundtrack must be mono or stereo [channels, samples], got {tuple(waveform.shape)}.") + if waveform.shape[-1] == 0: + raise ValueError("A reference soundtrack must contain samples.") + if sample_rate <= 0 or target_sample_rate <= 0: + raise ValueError("Reference audio sample rates must be positive.") + waveform = waveform.to(torch.float32)[:, :int(max_duration * sample_rate)] + if waveform.shape[0] == 1: + waveform = waveform.expand(MINIMAX_H3_AUDIO_CHANNELS, -1).contiguous() + if sample_rate == target_sample_rate: + return waveform + try: + import torchaudio + except (ImportError, OSError): + from math import gcd + + from scipy.signal import resample_poly + + divisor = gcd(sample_rate, target_sample_rate) + resampled = resample_poly( + waveform.numpy(), + up=target_sample_rate // divisor, + down=sample_rate // divisor, + axis=-1, + ) + return torch.from_numpy(np.asarray(resampled, dtype=np.float32)) + return torchaudio.transforms.Resample(sample_rate, target_sample_rate)(waveform) + + +def trim_reference_num_frames(num_frames: int) -> int: + """Trim video references to the causal VAE's complete chunk geometry.""" + if num_frames < 1: + raise ValueError(f"A reference video must have at least one frame, got {num_frames}.") + return ( + max(1, + (num_frames - MINIMAX_H3_LATENTS_PER_CHUNK) // MINIMAX_H3_FRAMES_PER_CHUNK) * MINIMAX_H3_FRAMES_PER_CHUNK + + MINIMAX_H3_LATENTS_PER_CHUNK) + + +def _resolve_audio_source(source: Any) -> tuple[torch.Tensor, int | None]: + if isinstance(source, str | os.PathLike): + return decode_reference_audio(source) + return torch.as_tensor(source), None + + +def prepare_reference( + reference: MiniMaxH3Reference, + num_frames: int, + target_sample_rate: int, +) -> MiniMaxH3PreparedReference: + """Decode and normalize one deferred Ref2VA medium.""" + prepared = MiniMaxH3PreparedReference(media_type=reference.media_type) + if reference.media_type == "image": + source = reference.source + image = load_image(str(source)) if isinstance(source, str | os.PathLike) else source + if not isinstance(image, Image.Image): + pixels = reference_media_to_uint8(image) + if pixels.ndim != 3 or pixels.shape[-1] != 3: + raise ValueError(f"An image reference must be RGB, got {tuple(pixels.shape)}.") + image = Image.fromarray(pixels) + image = ImageOps.exif_transpose(image).convert("RGB") + height, width = resolve_reference_image_size(*image.size) + prepared.image = prepare_reference_image(image, height, width) + return prepared + + if reference.media_type == "video": + decoded_soundtrack = None + decoded_sample_rate = None + if isinstance(reference.source, str | os.PathLike): + frames, decoded_fps, soundtrack = decode_reference_video(reference.source) + if soundtrack is not None: + decoded_soundtrack, decoded_sample_rate = soundtrack + else: + frames = reference_media_to_uint8(reference.source) + decoded_fps = float(MINIMAX_H3_FPS) + fps = float(reference.fps if reference.fps is not None else decoded_fps) + prepared.frames = prepare_reference_frames(resample_reference_frames(frames, fps), num_frames) + + if reference.soundtrack is not None: + waveform, source_rate = _resolve_audio_source(reference.soundtrack) + decoded_soundtrack = waveform + decoded_sample_rate = source_rate + if decoded_soundtrack is not None: + sample_rate = reference.sample_rate or decoded_sample_rate or target_sample_rate + prepared.waveform = prepare_reference_waveform( + decoded_soundtrack, + int(sample_rate), + target_sample_rate, + max_duration=num_frames / MINIMAX_H3_FPS, + ) + prepared.has_audio = True + elif reference.sample_rate is not None: + raise ValueError("A silent video reference cannot specify `sample_rate`.") + return prepared + + waveform, decoded_sample_rate = _resolve_audio_source(reference.source) + sample_rate = reference.sample_rate or decoded_sample_rate or target_sample_rate + prepared.waveform = prepare_reference_waveform( + waveform, + int(sample_rate), + target_sample_rate, + max_duration=num_frames / MINIMAX_H3_FPS, + ) + prepared.has_audio = True + return prepared + + +__all__ = [ + "MINIMAX_H3_MAX_REFERENCES", + "MINIMAX_H3_MAX_REFERENCE_AUDIOS", + "MINIMAX_H3_MAX_REFERENCE_IMAGES", + "MINIMAX_H3_MAX_REFERENCE_VIDEOS", + "MINIMAX_H3_QWEN_TEMPORAL_PATCH", + "MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS", + "MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE", + "MiniMaxH3PreparedReference", + "MiniMaxH3Reference", + "decode_reference_audio", + "decode_reference_video", + "prepare_reference", + "prepare_reference_frames", + "prepare_reference_image", + "prepare_reference_waveform", + "reference_media_to_uint8", + "resample_reference_frames", + "resolve_reference_image_size", + "sample_reference_video_frames", + "trim_reference_num_frames", + "validate_references", +] diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/__init__.py b/fastvideo/pipelines/basic/minimax_h3/stages/__init__.py new file mode 100644 index 000000000..4a8abb303 --- /dev/null +++ b/fastvideo/pipelines/basic/minimax_h3/stages/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_conditioning import MiniMaxH3ConditioningStage +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_decoding import ( + MiniMaxH3AudioDecodingStage, + MiniMaxH3VideoDecodingStage, +) +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import MiniMaxH3DenoisingStage +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_input_preparation import MiniMaxH3InputPreparationStage +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MiniMaxH3LatentPreparationStage + +__all__ = [ + "MiniMaxH3AudioDecodingStage", + "MiniMaxH3ConditioningStage", + "MiniMaxH3DenoisingStage", + "MiniMaxH3InputPreparationStage", + "MiniMaxH3LatentPreparationStage", + "MiniMaxH3VideoDecodingStage", +] diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py new file mode 100644 index 000000000..9b4e64d03 --- /dev/null +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3-VL conditioning for MiniMax H3 FL2VA and Ref2VA.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from torch.distributed.tensor import DTensor + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_IMAGE_PAD_TOKEN, + MINIMAX_H3_TEXT_ENCODER_LAYER, + MINIMAX_H3_TEXT_TAG, + MINIMAX_H3_VIDEO_PAD_TOKEN, + MINIMAX_H3_VIDEO_TAG, + MINIMAX_H3_VISION_END_TOKEN, + MINIMAX_H3_VISION_START_TOKEN, +) +from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference, sample_reference_video_frames +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_input_preparation import MINIMAX_H3_KEYFRAMES_KEY +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +MINIMAX_H3_TEXT_TOKEN_TAGS_KEY = "minimax_h3_text_token_tags" + + +def _token_ids(tokenized: Any) -> list[int]: + input_ids = tokenized["input_ids"] if isinstance(tokenized, dict) else tokenized.input_ids + if isinstance(input_ids, torch.Tensor): + input_ids = input_ids.tolist() + if input_ids and isinstance(input_ids[0], list): + if len(input_ids) != 1: + raise ValueError("MiniMax H3 tokenization must produce exactly one sequence.") + input_ids = input_ids[0] + return [int(token_id) for token_id in input_ids] + + +def _module_dtype(module: Any) -> torch.dtype: + dtype = getattr(module, "dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + parameter = next(module.parameters(), None) + return torch.float32 if parameter is None else parameter.dtype + + +def _create_mm_token_type_ids(processor: Any, token_ids: list[int]) -> list[list[int]]: + """Build Qwen3-VL modality IDs across old and new Transformers releases.""" + create_ids = getattr(processor, "create_mm_token_type_ids", None) + if callable(create_ids): + return create_ids([token_ids]) + + modality_ids = [0] * len(token_ids) + for modality, modality_type in (("image", 1), ("video", 2), ("audio", 3)): + special_ids = getattr(processor, f"{modality}_token_ids", None) + if special_ids is None: + special_id = getattr(processor, f"{modality}_token_id", None) + special_ids = [] if special_id is None else [special_id] + resolved_ids = {int(special_id) for special_id in special_ids if special_id is not None} + for index, token_id in enumerate(token_ids): + if token_id in resolved_ids: + modality_ids[index] = modality_type + return [modality_ids] + + +def build_ref2va_presentation( + tokenizer: Any, + prompt: str, + references: list[MiniMaxH3PreparedReference], + image_token_counts: list[int], + video_block_token_counts: list[int], +) -> tuple[list[int], list[int]]: + """Tokenize ordered reference labels, vision blocks, then the prompt.""" + + def text(value: str) -> tuple[list[int], list[int]]: + ids = _token_ids(tokenizer(value, add_special_tokens=False)) + return ids, [MINIMAX_H3_TEXT_TAG] * len(ids) + + def vision(pad_token: str, count: int) -> tuple[list[int], list[int]]: + ids = ([int(tokenizer.convert_tokens_to_ids(MINIMAX_H3_VISION_START_TOKEN))] + + [int(tokenizer.convert_tokens_to_ids(pad_token))] * count + + [int(tokenizer.convert_tokens_to_ids(MINIMAX_H3_VISION_END_TOKEN))]) + return ids, [MINIMAX_H3_VIDEO_TAG] * len(ids) + + token_ids: list[int] = [] + token_tags: list[int] = [] + + def emit(segment: tuple[list[int], list[int]]) -> None: + token_ids.extend(segment[0]) + token_tags.extend(segment[1]) + + counts = {"image": 0, "video": 0, "audio": 0} + for reference in references: + if reference.has_audio: + counts["audio"] += 1 + emit(text(f"