From 5e1e06013e49bb47449b99c39ad5ae42eb4d0b09 Mon Sep 17 00:00:00 2001 From: Ang Li Date: Thu, 13 Aug 2026 21:33:12 +0000 Subject: [PATCH 1/6] Add MiniMax H3 integration Signed-off-by: Ang Li --- .../minimax_h3/minimax_h3/__init__.py | 20 + integrations/minimax_h3/minimax_h3/config.py | 67 ++ .../minimax_h3/minimax_h3/constants.py | 52 ++ integrations/minimax_h3/minimax_h3/lora.py | 222 +++++ integrations/minimax_h3/minimax_h3/model.py | 200 +++++ .../minimax_h3/minimax_h3/pipeline.py | 778 ++++++++++++++++++ .../minimax_h3/minimax_h3/references.py | 106 +++ integrations/minimax_h3/minimax_h3/runner.py | 225 +++++ .../minimax_h3/minimax_h3/scheduler.py | 112 +++ .../minimax_h3/minimax_h3/transformer.py | 500 +++++++++++ integrations/minimax_h3/pyproject.toml | 55 ++ integrations/minimax_h3/tests/test_smoke.py | 341 ++++++++ pyproject.toml | 2 + uv.lock | 89 +- 14 files changed, 2743 insertions(+), 26 deletions(-) create mode 100644 integrations/minimax_h3/minimax_h3/__init__.py create mode 100644 integrations/minimax_h3/minimax_h3/config.py create mode 100644 integrations/minimax_h3/minimax_h3/constants.py create mode 100644 integrations/minimax_h3/minimax_h3/lora.py create mode 100644 integrations/minimax_h3/minimax_h3/model.py create mode 100644 integrations/minimax_h3/minimax_h3/pipeline.py create mode 100644 integrations/minimax_h3/minimax_h3/references.py create mode 100644 integrations/minimax_h3/minimax_h3/runner.py create mode 100644 integrations/minimax_h3/minimax_h3/scheduler.py create mode 100644 integrations/minimax_h3/minimax_h3/transformer.py create mode 100644 integrations/minimax_h3/pyproject.toml create mode 100644 integrations/minimax_h3/tests/test_smoke.py diff --git a/integrations/minimax_h3/minimax_h3/__init__.py b/integrations/minimax_h3/minimax_h3/__init__.py new file mode 100644 index 000000000..8e740516e --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/__init__.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MiniMax H3 integration for the FlashDreams runtime.""" + +from minimax_h3.pipeline import MiniMaxH3Pipeline, MiniMaxH3PipelineConfig + +__all__ = ["MiniMaxH3Pipeline", "MiniMaxH3PipelineConfig"] diff --git a/integrations/minimax_h3/minimax_h3/config.py b/integrations/minimax_h3/minimax_h3/config.py new file mode 100644 index 000000000..790dd543a --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/config.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Registered MiniMax H3 workflow and runner configs.""" + +from __future__ import annotations + +from flashdreams.infra.runner import RunnerConfig +from minimax_h3.model import MiniMaxH3DiffusionModelConfig +from minimax_h3.pipeline import MiniMaxH3PipelineConfig +from minimax_h3.runner import ( + MiniMaxH3FL2VARunnerConfig, + MiniMaxH3Ref2VARunnerConfig, + MiniMaxH3T2VARunnerConfig, +) +from minimax_h3.transformer import ( + H3_REF_TRANSFORMER_CHECKPOINT, + MiniMaxH3TransformerConfig, +) + +PIPELINE_MINIMAX_H3_T2VA = MiniMaxH3PipelineConfig( + name="minimax-h3-t2va", + workflow="t2va", +) +PIPELINE_MINIMAX_H3_FL2VA = MiniMaxH3PipelineConfig( + name="minimax-h3-fl2va", + workflow="fl2va", +) +PIPELINE_MINIMAX_H3_REF2VA = MiniMaxH3PipelineConfig( + name="minimax-h3-ref2va", + workflow="ref2va", + diffusion_model=MiniMaxH3DiffusionModelConfig( + transformer=MiniMaxH3TransformerConfig( + checkpoint_path=H3_REF_TRANSFORMER_CHECKPOINT, + device="cuda", + execution_device="cuda", + sequential_cpu_offload=False, + ) + ), +) + +RUNNER_MINIMAX_H3_T2VA = MiniMaxH3T2VARunnerConfig( + runner_name=PIPELINE_MINIMAX_H3_T2VA.name, + description="MiniMax H3 prompt-to-video generation with low-host-RAM staging.", + pipeline=PIPELINE_MINIMAX_H3_T2VA, +) +RUNNER_MINIMAX_H3_FL2VA = MiniMaxH3FL2VARunnerConfig( + runner_name=PIPELINE_MINIMAX_H3_FL2VA.name, + description=( + "MiniMax H3 first-frame, last-frame, or dual-keyframe video generation." + ), + pipeline=PIPELINE_MINIMAX_H3_FL2VA, +) +RUNNER_MINIMAX_H3_REF2VA = MiniMaxH3Ref2VARunnerConfig( + runner_name=PIPELINE_MINIMAX_H3_REF2VA.name, + description="MiniMax H3 ordered image, video, and audio reference generation.", + pipeline=PIPELINE_MINIMAX_H3_REF2VA, +) + +RUNNER_CONFIGS: dict[str, RunnerConfig] = { + config.runner_name: config + for config in ( + RUNNER_MINIMAX_H3_T2VA, + RUNNER_MINIMAX_H3_FL2VA, + RUNNER_MINIMAX_H3_REF2VA, + ) +} diff --git a/integrations/minimax_h3/minimax_h3/constants.py b/integrations/minimax_h3/minimax_h3/constants.py new file mode 100644 index 000000000..9c7145cd5 --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/constants.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MiniMax H3 geometry, timing, and model constants.""" + +from __future__ import annotations + +import math + +MODEL_ID = "MiniMaxAI/MiniMax-H3" +FPS = 24 +MIN_DURATION = 5.0 +MAX_DURATION = 15.0 +FRAME_CHUNK = 17 +FRAME_REMAINDER = 5 +CANVAS_MULTIPLE = 32 + + +def align_num_frames(duration: float) -> int: + """Convert seconds to H3's next decodable frame count.""" + if not math.isfinite(duration) or not MIN_DURATION <= duration <= MAX_DURATION: + raise ValueError( + f"duration must be between {MIN_DURATION:g} and {MAX_DURATION:g} seconds" + ) + frames = math.ceil(duration * FPS) + while frames % FRAME_CHUNK != FRAME_REMAINDER: + frames += 1 + if frames / FPS > MAX_DURATION: + raise ValueError("duration aligns beyond MiniMax H3's 15-second maximum") + return frames + + +def validate_canvas(width: int, height: int) -> None: + """Validate an H3 output canvas.""" + if width <= 0 or height <= 0: + raise ValueError("width and height must be positive") + if width % CANVAS_MULTIPLE or height % CANVAS_MULTIPLE: + raise ValueError(f"width and height must be multiples of {CANVAS_MULTIPLE}") + if not 0.25 <= width / height <= 4: + raise ValueError("aspect ratio must be between 1:4 and 4:1") diff --git a/integrations/minimax_h3/minimax_h3/lora.py b/integrations/minimax_h3/minimax_h3/lora.py new file mode 100644 index 000000000..3997afade --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/lora.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Musubi LoRA conversion and merging for the native H3 transformer.""" + +from __future__ import annotations + +import hashlib +import os +import re +from pathlib import Path +from typing import Any + +import torch + +CONVERSION_VERSION = "5" +MUSUBI_KEY = re.compile( + r"^lora_unet_blocks_(?P\d+)_" + r"(?Pattn_qkv_proj|attn_out_proj|mlp_fc1|mlp_fc2)" + r"\.(?Palpha|lora_down\.weight|lora_up\.weight)$" +) +DIRECT_TARGETS = { + "attn_out_proj": "attn.to_out.0", + "mlp_fc1": "ff.net.0.proj", + "mlp_fc2": "ff.net.2", +} + + +def _lora_cache_root() -> Path: + configured = os.environ.get("MINIMAX_H3_LORA_CACHE") + if configured: + return Path(configured).expanduser().resolve() + return ( + Path( + os.environ.get("FLASHDREAMS_CACHE_DIR", "~/.cache/flashdreams") + ).expanduser() + / "minimax_h3" + / "lora" + ) + + +def resolve_lora(source: str, weight_name: str | None = None) -> tuple[Path, str]: + """Resolve a local Musubi adapter or download one Hub file.""" + local = Path(source).expanduser() + if local.is_file(): + return local.resolve(), "local" + if local.exists(): + raise ValueError(f"LoRA source is not a file: {local}") + + from huggingface_hub import HfApi, hf_hub_download + + info = HfApi().model_info(source) + if weight_name is None: + candidates = sorted( + sibling.rfilename + for sibling in info.siblings or [] + if sibling.rfilename.endswith(".safetensors") + ) + if len(candidates) != 1: + raise ValueError( + f"LoRA repository {source!r} contains {len(candidates)} safetensors " + "files; pass --lora-weight-name explicitly" + ) + weight_name = candidates[0] + path = hf_hub_download(source, filename=weight_name, revision=info.sha) + return Path(path), info.sha or "unversioned" + + +def _converted_path(source: Path, revision: str) -> Path: + identity = ( + f"{source.resolve()}:{source.stat().st_size}:" + f"{source.stat().st_mtime_ns}:{revision}:{CONVERSION_VERSION}" + ) + digest = hashlib.sha256(identity.encode()).hexdigest()[:16] + return _lora_cache_root() / f"{source.stem}-{digest}.flashdreams.safetensors" + + +def _source_groups(handle: Any) -> dict[tuple[int, str], dict[str, str]]: + groups: dict[tuple[int, str], dict[str, str]] = {} + unknown: list[str] = [] + source_keys = handle.keys() + for key in source_keys: + match = MUSUBI_KEY.fullmatch(key) + if match is None: + unknown.append(key) + continue + group = (int(match["block"]), match["module"]) + groups.setdefault(group, {})[match["part"]] = key + if unknown: + raise ValueError(f"unsupported MiniMax H3 LoRA tensors: {unknown[:3]}") + modules = {"attn_qkv_proj", *DIRECT_TARGETS} + expected = {(block, module) for block in range(50) for module in modules} + if set(groups) != expected: + missing = sorted(expected - set(groups)) + extra = sorted(set(groups) - expected) + raise ValueError( + "LoRA does not cover the expected 50 H3 blocks " + f"(missing={missing[:3]}, extra={extra[:3]})" + ) + for group, parts in groups.items(): + if set(parts) != {"alpha", "lora_down.weight", "lora_up.weight"}: + raise ValueError(f"incomplete LoRA tensor group {group}: {sorted(parts)}") + return groups + + +def convert_musubi_lora(source: Path, output: Path) -> Path: + """Translate Musubi adapter names to the native H3 modules.""" + from safetensors import safe_open + from safetensors.torch import save_file + + converted: dict[str, torch.Tensor] = {} + output_metadata = { + "format": "pt", + "source_format": "musubi-minimax-h3", + "source_file": source.name, + "conversion": "split fused qkv; fold alpha/rank into lora_B", + "conversion_version": CONVERSION_VERSION, + } + with safe_open(source, framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + architecture = metadata.get("modelspec.architecture") + if architecture not in {None, "MiniMax-H3/lora"}: + raise ValueError( + f"expected a MiniMax-H3 LoRA, found architecture {architecture!r}" + ) + if training_mode := metadata.get("ss_h3_training_mode"): + output_metadata["training_mode"] = training_mode + groups = _source_groups(handle) + for (block, module), parts in sorted(groups.items()): + down = handle.get_tensor(parts["lora_down.weight"]) + up = handle.get_tensor(parts["lora_up.weight"]) + alpha = float(handle.get_tensor(parts["alpha"]).item()) + if down.ndim != 2 or up.ndim != 2 or down.shape[0] != up.shape[1]: + raise ValueError( + f"invalid LoRA shapes for block {block} {module}: " + f"{down.shape}, {up.shape}" + ) + scaled_up = up * (alpha / down.shape[0]) + block_prefix = f"transformer.transformer_blocks.{block}" + if module == "attn_qkv_proj": + if scaled_up.shape[0] % 3: + raise ValueError( + f"QKV LoRA output is not divisible by three in block {block}: " + f"{scaled_up.shape}" + ) + for target, target_up in zip( + ("attn.to_q", "attn.to_k", "attn.to_v"), + scaled_up.chunk(3, dim=0), + strict=True, + ): + prefix = f"{block_prefix}.{target}" + converted[f"{prefix}.lora_A.weight"] = down.clone() + converted[f"{prefix}.lora_B.weight"] = target_up + else: + prefix = f"{block_prefix}.{DIRECT_TARGETS[module]}" + converted[f"{prefix}.lora_A.weight"] = down + converted[f"{prefix}.lora_B.weight"] = scaled_up + + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_name(output.name + ".tmp") + save_file(converted, str(temporary), metadata=output_metadata) + os.replace(temporary, output) + return output + + +def prepare_lora(source: str, weight_name: str | None = None) -> Path: + """Return a converted, cached native adapter path.""" + resolved, revision = resolve_lora(source, weight_name) + output = _converted_path(resolved, revision) + if not output.is_file(): + convert_musubi_lora(resolved, output) + return output + + +def load_lora( + transformer: Any, + source: str, + scale: float, + weight_name: str | None = None, +) -> Path: + """Merge a converted Musubi adapter into the native BF16 transformer.""" + if not 0 <= scale <= 4: + raise ValueError("LoRA scale must be between 0 and 4") + converted = prepare_lora(source, weight_name) + from safetensors import safe_open + + with safe_open(converted, framework="pt", device="cpu") as handle: + keys = set(handle.keys()) + a_keys = sorted(key for key in keys if key.endswith(".lora_A.weight")) + expected_b = { + key.removesuffix(".lora_A.weight") + ".lora_B.weight" for key in a_keys + } + if expected_b != {key for key in keys if key.endswith(".lora_B.weight")}: + raise ValueError(f"incomplete converted LoRA pairs in {converted}") + with torch.no_grad(): + for a_key in a_keys: + prefix = a_key.removeprefix("transformer.").removesuffix( + ".lora_A.weight" + ) + a = handle.get_tensor(a_key) + b = handle.get_tensor( + a_key.removesuffix(".lora_A.weight") + ".lora_B.weight" + ) + module = transformer.get_submodule(prefix) + if not hasattr(module, "weight"): + raise ValueError(f"LoRA target has no weight: {prefix}") + down = a.to(device=module.weight.device, dtype=module.weight.dtype) + up = b.to(device=module.weight.device, dtype=module.weight.dtype) + module.weight.addmm_(up, down, beta=1.0, alpha=scale) + return converted diff --git a/integrations/minimax_h3/minimax_h3/model.py b/integrations/minimax_h3/minimax_h3/model.py new file mode 100644 index 000000000..83f8ad7b9 --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/model.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FlashDreams diffusion model for MiniMax H3's paired latent streams.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, cast + +import torch +from torch import Tensor, nn + +from flashdreams.infra.diffusion.model import DiffusionModel, DiffusionModelConfig +from flashdreams.infra.diffusion.transformer import Transformer, TransformerConfig +from minimax_h3.scheduler import MiniMaxH3Scheduler, MiniMaxH3SchedulerConfig +from minimax_h3.transformer import ( + MiniMaxH3TransformerCache, + MiniMaxH3TransformerConfig, +) + + +@dataclass(kw_only=True) +class MiniMaxH3DiffusionModelConfig(DiffusionModelConfig): + """Native H3 transformer plus separate video and audio schedules.""" + + _target: type[MiniMaxH3DiffusionModel] = field( + default_factory=lambda: MiniMaxH3DiffusionModel + ) + transformer: TransformerConfig = field( + default_factory=lambda: MiniMaxH3TransformerConfig( + device="cuda", + execution_device="cuda", + sequential_cpu_offload=False, + ) + ) + scheduler: MiniMaxH3SchedulerConfig = field( + default_factory=MiniMaxH3SchedulerConfig + ) + audio_scheduler: MiniMaxH3SchedulerConfig = field( + default_factory=lambda: MiniMaxH3SchedulerConfig(shift=3.0) + ) + + +@dataclass(kw_only=True) +class MiniMaxH3DenoiseState: + """Packed conditioning, noise, and layout produced before denoising.""" + + latents: Tensor + audio_latents: Tensor + prompt_embeds: Tensor + position_ids: Tensor + token_tags: Tensor + video_indices: Tensor + audio_indices: Tensor + text_indices: Tensor + num_condition_video_rows: int + num_condition_audio_rows: int + num_latent_frames: int + latent_height: int + latent_width: int + + +class MiniMaxH3DiffusionModel(DiffusionModel[MiniMaxH3TransformerCache]): + """Run H3's joint forward under two FlashDreams-owned schedulers.""" + + config: MiniMaxH3DiffusionModelConfig + transformer: Transformer[MiniMaxH3TransformerCache] + scheduler: MiniMaxH3Scheduler + audio_scheduler: MiniMaxH3Scheduler + + def __init__(self, config: MiniMaxH3DiffusionModelConfig) -> None: + nn.Module.__init__(self) + self.config = config + self.transformer = config.transformer.setup() + self.scheduler = config.scheduler.setup() + self.audio_scheduler = config.audio_scheduler.setup() + + @staticmethod + def _row_timesteps( + state: MiniMaxH3DenoiseState, + video_timestep: Tensor, + audio_timestep: Tensor, + ) -> tuple[Tensor, Tensor]: + sequence_length = ( + state.video_indices.numel() + + state.audio_indices.numel() + + state.text_indices.numel() + ) + row_timesteps = torch.full( + (sequence_length,), + float(video_timestep), + dtype=torch.float32, + device=state.video_indices.device, + ) + video_condition = state.video_indices[: state.num_condition_video_rows] + audio_condition = state.audio_indices[: state.num_condition_audio_rows] + audio_target = state.audio_indices[state.num_condition_audio_rows :] + row_timesteps[video_condition] = max(float(video_timestep), 0.999) + row_timesteps[audio_target] = audio_timestep + row_timesteps[audio_condition] = 1.0 + return torch.unique(row_timesteps, sorted=True, return_inverse=True) + + @torch.no_grad() + def generate_joint(self, state: MiniMaxH3DenoiseState) -> Tensor: + """Denoise both streams and return only unpacked video latents.""" + device = self.transformer.device + video = state.latents.to(device) + audio = state.audio_latents.to(device) + state.prompt_embeds = state.prompt_embeds.to(device) + state.position_ids = state.position_ids.to(device) + state.token_tags = state.token_tags.to(device) + state.video_indices = state.video_indices.to(device) + state.audio_indices = state.audio_indices.to(device) + state.text_indices = state.text_indices.to(device) + + video_sigmas, video_timesteps = self.scheduler.schedule(device) + audio_sigmas, audio_timesteps = self.audio_scheduler.schedule(device) + if len(video_timesteps) != len(audio_timesteps): + raise RuntimeError("H3 video and audio schedules must have equal length") + + cache = MiniMaxH3TransformerCache( + audio_hidden_states=audio[None], + encoder_hidden_states=state.prompt_embeds, + timestep=torch.empty(0, device=device), + timestep_indices=torch.empty(0, dtype=torch.long, device=device), + token_tags=state.token_tags, + position_ids=state.position_ids, + video_indices=state.video_indices, + audio_indices=state.audio_indices, + text_indices=state.text_indices, + ) + for index, (video_timestep, audio_timestep) in enumerate( + zip(video_timesteps, audio_timesteps, strict=True) + ): + print( + f"MiniMax H3 denoise step {index + 1}/{len(video_timesteps)}", + flush=True, + ) + cache.timestep, cache.timestep_indices = self._row_timesteps( + state, video_timestep, audio_timestep + ) + cache.audio_hidden_states = audio[None] + video_flow = self.transformer.predict_flow( + video[None], video_timestep, cache + )[0] + if cache.last_audio_flow is None: + raise RuntimeError("H3 transformer did not produce an audio flow") + audio_flow = cache.last_audio_flow[0] + + video_start = state.num_condition_video_rows + audio_start = state.num_condition_audio_rows + video[video_start:] = self.scheduler.step( + video[video_start:], + video_flow[video_start:].float(), + video_timestep, + video_sigmas[index], + video_sigmas[index + 1], + ) + audio[audio_start:] = self.audio_scheduler.step( + audio[audio_start:], + audio_flow[audio_start:].float(), + audio_timestep, + audio_sigmas[index], + audio_sigmas[index + 1], + ) + + rows = video[state.num_condition_video_rows :] + transformer_config = cast(Any, self.config.transformer) + patch_t, patch_h, patch_w = transformer_config.patch_size + channels = transformer_config.in_channels + rows = rows.reshape( + -1, + state.num_latent_frames // patch_t, + state.latent_height // patch_h, + state.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, + state.num_latent_frames, + state.latent_height, + state.latent_width, + ) + .contiguous() + .cpu() + ) + + +__all__ = [ + "MiniMaxH3DenoiseState", + "MiniMaxH3DiffusionModel", + "MiniMaxH3DiffusionModelConfig", +] diff --git a/integrations/minimax_h3/minimax_h3/pipeline.py b/integrations/minimax_h3/minimax_h3/pipeline.py new file mode 100644 index 000000000..ca7853895 --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/pipeline.py @@ -0,0 +1,778 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Crash-safe MiniMax H3 FL2VA pipeline for the FlashDreams runtime.""" + +from __future__ import annotations + +import gc +import hashlib +import json +import os +import time +from dataclasses import dataclass, field, replace +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Literal, cast + +import numpy as np +import torch +from torch import nn + +from flashdreams.infra.pipeline import ( + StreamInferencePipeline, + StreamInferencePipelineConfig, +) +from minimax_h3.constants import MODEL_ID, align_num_frames, validate_canvas +from minimax_h3.lora import load_lora +from minimax_h3.model import ( + MiniMaxH3DenoiseState, + MiniMaxH3DiffusionModelConfig, +) +from minimax_h3.references import MiniMaxH3ReferenceSpec, load_references +from minimax_h3.transformer import MiniMaxH3TransformerConfig + +MiniMaxH3Workflow = Literal["t2va", "fl2va", "ref2va"] + + +@dataclass(kw_only=True) +class MiniMaxH3PipelineConfig(StreamInferencePipelineConfig): + """Config for the FlashDreams-native H3 denoising pipeline.""" + + _target: type[MiniMaxH3Pipeline] = field(default_factory=lambda: MiniMaxH3Pipeline) + + diffusion_model: MiniMaxH3DiffusionModelConfig = field( + default_factory=MiniMaxH3DiffusionModelConfig + ) + """Native joint transformer and paired scheduler configuration.""" + + model_id: str = MODEL_ID + """Hugging Face model repository or local snapshot path.""" + + cache_dir: Path | None = None + """Optional Hugging Face model cache root.""" + + workflow: MiniMaxH3Workflow = "fl2va" + """Released checkpoint workflow selected by this registered pipeline.""" + + +@dataclass(frozen=True) +class _ReferenceLayout: + """Reference properties needed after its encoded media is checkpointed.""" + + kind: str + has_audio: bool + + +@dataclass(kw_only=True) +class MiniMaxH3PipelineCache: + """Per-rollout H3 request, checkpoints, and runtime metrics.""" + + prompt: str + workflow: MiniMaxH3Workflow + image_path: Path | None + last_image_path: Path | None + references: tuple[MiniMaxH3ReferenceSpec, ...] + output_path: Path + width: int + height: int + duration: float + steps: int + seed: int + low_ram: bool + restart: bool + attention: str + lora: str | None + lora_weight_name: str | None + lora_scale: float + latent_checkpoint: Path + conditioning_checkpoint: Path + generated: bool = False + elapsed_seconds: float = 0.0 + conditioning_seconds: float = 0.0 + denoise_seconds: float = 0.0 + decode_seconds: float = 0.0 + peak_gpu_memory_gib: float = 0.0 + attention_backend: str = "default" + resumed_stage: str | None = None + + +def _replace_blocks(pipe: Any, names: tuple[str, ...]) -> None: + from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks + + selected: dict[str, Any] = {} + for requested in names: + if requested in pipe._blocks.sub_blocks: + selected[requested] = pipe._blocks.sub_blocks[requested] + continue + prefix = requested + "." + for actual, block in pipe._blocks.sub_blocks.items(): + if actual.startswith(prefix): + selected[actual.removeprefix(prefix)] = block + if not selected: + raise KeyError(f"none of the requested pipeline stages exist: {names}") + pipe._blocks = SequentialPipelineBlocks.from_blocks_dict(selected) + + +def _release_pipeline(pipe: Any) -> None: + for component in pipe.components.values(): + if isinstance(component, nn.Module): + try: + component.to_empty(device="cpu") + except (AttributeError, RuntimeError): + pass + del pipe + gc.collect() + torch.cuda.empty_cache() + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + ".tmp") + temporary.write_text(json.dumps(payload, indent=2) + "\n") + os.replace(temporary, path) + + +def _write_status(checkpoint: Path, stage: str, **details: Any) -> None: + _atomic_json( + checkpoint.with_suffix(checkpoint.suffix + ".status.json"), + { + "stage": stage, + "updated_at": datetime.now(timezone.utc).isoformat(), + **details, + }, + ) + + +def _file_identity(path: Path | None) -> dict[str, str | int] | None: + if path is None: + return None + resolved = path.resolve() + stat = resolved.stat() + return { + "path": str(resolved), + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + } + + +def _conditioning_manifest( + cache: MiniMaxH3PipelineCache, model_id: str +) -> dict[str, Any]: + return { + "workflow": cache.workflow, + "prompt": cache.prompt, + "image": _file_identity(cache.image_path), + "last_image": _file_identity(cache.last_image_path), + "references": [reference.manifest() for reference in cache.references], + "width": cache.width, + "height": cache.height, + "duration": cache.duration, + "model_id": model_id, + } + + +def _generation_manifest( + cache: MiniMaxH3PipelineCache, model_id: str +) -> dict[str, Any]: + return { + **_conditioning_manifest(cache, model_id), + "steps": cache.steps, + "seed": cache.seed, + "attention": cache.attention, + "lora": cache.lora, + "lora_weight_name": cache.lora_weight_name, + "lora_scale": cache.lora_scale, + } + + +def _signature(manifest: dict[str, Any]) -> str: + encoded = json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _save_conditioning( + cache: MiniMaxH3PipelineCache, + model_id: str, + values: dict[str, Any], +) -> None: + from safetensors.torch import save_file + + manifest = _conditioning_manifest(cache, model_id) + path = cache.conditioning_checkpoint + temporary = path.with_name(path.name + ".tmp") + condition_latents = values["condition_latents"] + audio_condition_latents = values["audio_condition_latents"] + tensors = { + "prompt_embeds": values["prompt_embeds"].detach().cpu().contiguous(), + "text_token_tags": values["text_token_tags"].detach().cpu().contiguous(), + **{ + f"condition_latents.{index}": latent.detach().cpu().contiguous() + for index, latent in enumerate(condition_latents) + }, + **{ + f"audio_condition_latents.{index}": latent.detach().cpu().contiguous() + for index, latent in enumerate(audio_condition_latents) + }, + } + save_file( + tensors, + str(temporary), + metadata={ + "stage": "conditioned", + "manifest": json.dumps(manifest, sort_keys=True), + "signature": _signature(manifest), + "height": str(values["height"]), + "width": str(values["width"]), + "num_frames": str(values["num_frames"]), + "keyframe_anchors": json.dumps(list(values["keyframe_anchors"])), + "condition_count": str(len(condition_latents)), + "audio_condition_count": str(len(audio_condition_latents)), + "reference_layout": json.dumps( + [ + {"kind": reference.kind, "has_audio": reference.has_audio} + for reference in values["normalized_references"] + ] + ), + }, + ) + os.replace(temporary, path) + _write_status(path, "conditioned") + + +def _load_conditioning(cache: MiniMaxH3PipelineCache, model_id: str) -> dict[str, Any]: + from safetensors import safe_open + from safetensors.torch import load_file + + path = cache.conditioning_checkpoint + with safe_open(path, framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + expected = _signature(_conditioning_manifest(cache, model_id)) + if metadata.get("stage") != "conditioned" or metadata.get("signature") != expected: + raise ValueError( + f"conditioning checkpoint does not match this request: {path}; " + "use --restart" + ) + tensors = load_file(path, device="cpu") + condition_count = int(metadata["condition_count"]) + audio_condition_count = int(metadata.get("audio_condition_count", "0")) + reference_layout = json.loads(metadata.get("reference_layout", "[]")) + return { + "prompt_embeds": tensors["prompt_embeds"], + "text_token_tags": tensors["text_token_tags"], + "condition_latents": [ + tensors[f"condition_latents.{index}"] for index in range(condition_count) + ], + "audio_condition_latents": [ + tensors[f"audio_condition_latents.{index}"] + for index in range(audio_condition_count) + ], + "normalized_references": [ + _ReferenceLayout(kind=reference["kind"], has_audio=reference["has_audio"]) + for reference in reference_layout + ], + "height": int(metadata["height"]), + "width": int(metadata["width"]), + "num_frames": int(metadata["num_frames"]), + "keyframe_anchors": tuple(json.loads(metadata["keyframe_anchors"])), + } + + +def _save_latents( + cache: MiniMaxH3PipelineCache, + model_id: str, + latents: torch.Tensor, +) -> None: + from safetensors.torch import save_file + + path = cache.latent_checkpoint + temporary = path.with_name(path.name + ".tmp") + manifest = _generation_manifest(cache, model_id) + save_file( + {"latents": latents.detach().cpu().contiguous()}, + str(temporary), + metadata={ + "stage": "denoised", + "manifest": json.dumps(manifest, sort_keys=True), + "signature": _signature(manifest), + }, + ) + os.replace(temporary, path) + _write_status(path, "denoised") + + +def _load_latents(cache: MiniMaxH3PipelineCache, model_id: str) -> torch.Tensor: + from safetensors import safe_open + from safetensors.torch import load_file + + path = cache.latent_checkpoint + with safe_open(path, framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + expected = _signature(_generation_manifest(cache, model_id)) + if metadata.get("stage") != "denoised" or metadata.get("signature") != expected: + raise ValueError( + f"latent checkpoint does not match this request: {path}; use --restart" + ) + tensors = load_file(path, device="cpu") + return tensors["latents"] + + +class MiniMaxH3Pipeline(StreamInferencePipeline[Any, Any, Any]): + """FlashDreams H3 runtime with staged third-party conditioning and decode.""" + + config: MiniMaxH3PipelineConfig + + def __init__(self, config: MiniMaxH3PipelineConfig) -> None: + nn.Module.__init__(self) + self.config = config + self.register_buffer("_device_anchor", torch.empty(0), persistent=False) + + @property + def device(self) -> torch.device: + return cast(torch.Tensor, self._device_anchor).device + + def initialize_cache( + self, + *, + prompt: str, + image_path: Path | None, + last_image_path: Path | None, + references: tuple[MiniMaxH3ReferenceSpec, ...], + output_path: Path, + width: int, + height: int, + duration: float, + steps: int, + seed: int, + low_ram: bool, + restart: bool, + attention: str, + lora: str | None, + lora_weight_name: str | None, + lora_scale: float, + ) -> MiniMaxH3PipelineCache: + """Build a validated, checkpoint-aware workflow cache.""" + if not prompt.strip(): + raise ValueError("prompt cannot be empty") + workflow = self.config.workflow + if workflow == "t2va" and ( + image_path is not None or last_image_path is not None or references + ): + raise ValueError("t2va does not accept keyframes or references") + if workflow == "fl2va" and not (image_path or last_image_path): + raise ValueError("fl2va requires --image-path and/or --last-image-path") + if workflow == "ref2va" and not references: + raise ValueError("ref2va requires at least one --reference") + if workflow != "ref2va" and references: + raise ValueError(f"{workflow} does not accept ordered references") + for label, path in ( + ("first-frame", image_path), + ("last-frame", last_image_path), + ): + if path is not None and not path.is_file(): + raise FileNotFoundError(f"{label} image not found: {path}") + if steps < 2: + raise ValueError("steps must be at least 2 scheduler points") + if attention not in {"auto", "flash", "default"}: + raise ValueError(f"unsupported attention backend: {attention}") + if not 0 <= lora_scale <= 4: + raise ValueError("LoRA scale must be between 0 and 4") + validate_canvas(width, height) + align_num_frames(duration) + output_path.parent.mkdir(parents=True, exist_ok=True) + latent_checkpoint = output_path.with_suffix( + output_path.suffix + ".latents.safetensors" + ) + conditioning_checkpoint = output_path.with_suffix( + output_path.suffix + ".conditioning.safetensors" + ) + return MiniMaxH3PipelineCache( + prompt=prompt, + workflow=workflow, + image_path=image_path, + last_image_path=last_image_path, + references=references, + output_path=output_path, + width=width, + height=height, + duration=duration, + steps=steps, + seed=seed, + low_ram=low_ram, + restart=restart, + attention=attention, + lora=lora, + lora_weight_name=lora_weight_name, + lora_scale=lora_scale, + latent_checkpoint=latent_checkpoint, + conditioning_checkpoint=conditioning_checkpoint, + ) + + @torch.no_grad() + def generate( + self, + autoregressive_index: int, + cache: MiniMaxH3PipelineCache, + input: Any = None, + ) -> torch.Tensor: + """Generate and video-decode the single non-streaming H3 step.""" + del input + if autoregressive_index != 0 or cache.generated: + raise ValueError("MiniMax H3 supports exactly one runtime step per cache") + started = time.monotonic() + torch.cuda.reset_peak_memory_stats() + + if cache.latent_checkpoint.is_file() and not cache.restart: + cache.resumed_stage = "decode" + latents = _load_latents(cache, self.config.model_id) + else: + denoise_started = time.monotonic() + if cache.low_ram: + latents = self._generate_low_ram(cache) + else: + latents = self._generate_standard(cache) + cache.denoise_seconds = time.monotonic() - denoise_started + _save_latents(cache, self.config.model_id, latents) + + decode_started = time.monotonic() + frames = self._decode_video(cache, latents) + cache.decode_seconds = time.monotonic() - decode_started + cache.elapsed_seconds = time.monotonic() - started + cache.peak_gpu_memory_gib = torch.cuda.max_memory_allocated() / 2**30 + cache.generated = True + _write_status(cache.latent_checkpoint, "decoded-video") + return frames + + def mark_complete(self, cache: MiniMaxH3PipelineCache) -> None: + """Record completion only after the runtime output target closes.""" + if not cache.generated or not cache.output_path.is_file(): + raise RuntimeError("cannot complete H3 job before its MP4 is written") + _write_status( + cache.latent_checkpoint, "complete", output=str(cache.output_path) + ) + + def finalize( + self, + autoregressive_index: int, + cache: MiniMaxH3PipelineCache, + ) -> dict[str, float]: + """Return runtime metrics for the completed H3 rollout.""" + if autoregressive_index != 0 or not cache.generated: + raise ValueError("finalize requires the completed H3 runtime step") + return { + "conditioning_seconds": cache.conditioning_seconds, + "denoise_seconds": cache.denoise_seconds, + "decode_seconds": cache.decode_seconds, + "total_seconds": cache.elapsed_seconds, + "peak_gpu_memory_gib": cache.peak_gpu_memory_gib, + } + + def _cache_dir(self) -> str | None: + return str(self.config.cache_dir) if self.config.cache_dir is not None else None + + def _apply_lora(self, transformer: Any, cache: MiniMaxH3PipelineCache) -> None: + if cache.lora is None: + return + converted = load_lora( + transformer, + cache.lora, + cache.lora_scale, + cache.lora_weight_name, + ) + print(f"Loaded LoRA {converted} at scale {cache.lora_scale:g}", flush=True) + + def _generate_low_ram(self, cache: MiniMaxH3PipelineCache) -> torch.Tensor: + if cache.conditioning_checkpoint.is_file() and not cache.restart: + cache.resumed_stage = "denoise" + conditioned = _load_conditioning(cache, self.config.model_id) + else: + conditioning_started = time.monotonic() + _write_status(cache.latent_checkpoint, "conditioning") + conditioned = self._condition(cache) + _save_conditioning(cache, self.config.model_id, conditioned) + cache.conditioning_seconds = time.monotonic() - conditioning_started + return self._run_native_denoise(cache, conditioned) + + def _condition(self, cache: MiniMaxH3PipelineCache) -> dict[str, Any]: + from diffusers.utils import load_image + + num_frames = align_num_frames(cache.duration) + media: dict[str, Any] + if cache.workflow == "t2va": + media = { + "height": cache.height, + "width": cache.width, + "num_frames": num_frames, + "keyframe_anchors": (), + "normalized_references": [], + } + text_inputs = {"prompt": cache.prompt} + encoded = { + "condition_latents": [], + "audio_condition_latents": [], + } + elif cache.workflow == "fl2va": + resize_inputs: dict[str, Any] = { + "height": cache.height, + "width": cache.width, + } + if cache.image_path is not None: + resize_inputs["image"] = load_image(str(cache.image_path)) + if cache.last_image_path is not None: + resize_inputs["last_image"] = load_image(str(cache.last_image_path)) + media = self._run_conditioning_stage( + cache.workflow, + ("before_encode",), + resize_inputs, + ["height", "width", "keyframes", "keyframe_anchors"], + ) + media["num_frames"] = num_frames + media["normalized_references"] = [] + text_inputs = {"prompt": cache.prompt, "keyframes": media["keyframes"]} + encoded = self._run_conditioning_stage( + cache.workflow, + ("vae_encoder",), + {"keyframes": media["keyframes"]}, + ["condition_latents"], + cuda_components=("vae",), + ) + encoded["audio_condition_latents"] = [] + else: + references = load_references(cache.references) + media = self._run_conditioning_stage( + cache.workflow, + ("before_encode",), + { + "references": references, + "height": cache.height, + "width": cache.width, + "num_frames": num_frames, + }, + ["height", "width", "num_frames", "normalized_references"], + ) + media["keyframe_anchors"] = () + text_inputs = { + "prompt": cache.prompt, + "normalized_references": media["normalized_references"], + } + encoded = self._run_conditioning_stage( + cache.workflow, + ("vae_encoder",), + {"normalized_references": media["normalized_references"]}, + ["condition_latents", "audio_condition_latents"], + cuda_components=("vae", "audio_vae"), + ) + + text = self._run_conditioning_stage( + cache.workflow, + ("text_encoder",), + text_inputs, + ["prompt_embeds", "text_token_tags"], + cuda_components=("text_encoder",), + ) + return { + **text, + **encoded, + "height": media["height"], + "width": media["width"], + "keyframe_anchors": media["keyframe_anchors"], + "normalized_references": media["normalized_references"], + "num_frames": media["num_frames"], + } + + def _run_conditioning_stage( + self, + workflow: MiniMaxH3Workflow, + blocks: tuple[str, ...], + inputs: dict[str, Any], + outputs: list[str], + *, + cuda_components: tuple[str, ...] = (), + ) -> dict[str, Any]: + from diffusers.modular_pipelines.modular_pipeline import ModularPipeline + + pipe = ModularPipeline.from_pretrained( + self.config.model_id, + workflow=workflow, + cache_dir=self._cache_dir(), + ) + _replace_blocks(pipe, blocks) + required = list( + dict.fromkeys(spec.name for spec in pipe._blocks.expected_components) + ) + cpu_components = [name for name in required if name not in cuda_components] + pipe.load_components(names=cpu_components, dtype=torch.bfloat16) + for name in cuda_components: + pipe.load_components( + names=[name], + dtype=torch.bfloat16, + device_map="cuda", + low_cpu_mem_usage=True, + ) + try: + return dict(pipe(**inputs, output=outputs)) + finally: + _release_pipeline(pipe) + + def _build_prepare_pipeline(self, workflow: MiniMaxH3Workflow) -> Any: + from diffusers.modular_pipelines.modular_pipeline import ModularPipeline + + pipe = ModularPipeline.from_pretrained( + self.config.model_id, + workflow=workflow, + cache_dir=self._cache_dir(), + ) + blocks = { + "t2va": ( + "denoise.no_keyframe_anchors", + "denoise.prepare_layout", + "denoise.prepare_latents", + ), + "fl2va": ( + "denoise.prepare_layout", + "denoise.prepare_condition_latents", + "denoise.prepare_latents", + "denoise.prepare_latents_fl2va", + ), + "ref2va": ( + "denoise.prepare_layout", + "denoise.prepare_condition_latents", + "denoise.prepare_latents", + "denoise.prepare_latents_ref2va", + ), + }[workflow] + _replace_blocks(pipe, blocks) + if workflow != "t2va": + pipe.load_components(names=["scheduler"], dtype=torch.bfloat16) + return pipe + + def _prepare_denoise_state( + self, + cache: MiniMaxH3PipelineCache, + conditioned: dict[str, Any], + ) -> MiniMaxH3DenoiseState: + from diffusers.modular_pipelines.modular_pipeline import PipelineState + + pipe = self._build_prepare_pipeline(cache.workflow) + state = PipelineState() + for name, value in conditioned.items(): + state.set(name, value) + state.set("generator", torch.Generator(device="cpu").manual_seed(cache.seed)) + fields = [ + "latents", + "audio_latents", + "prompt_embeds", + "position_ids", + "token_tags", + "video_indices", + "audio_indices", + "text_indices", + "num_condition_video_rows", + "num_condition_audio_rows", + "num_latent_frames", + "latent_height", + "latent_width", + ] + try: + results = pipe(state=state, output=fields) + finally: + _release_pipeline(pipe) + return MiniMaxH3DenoiseState(**results) + + def _run_native_denoise( + self, cache: MiniMaxH3PipelineCache, conditioned: dict[str, Any] + ) -> torch.Tensor: + _write_status(cache.latent_checkpoint, "denoising-native-flashdreams") + state = self._prepare_denoise_state(cache, conditioned) + backend = "cudnn" if cache.attention == "default" else "flash" + cache.attention_backend = backend + base_transformer = cast( + MiniMaxH3TransformerConfig, self.config.diffusion_model.transformer + ) + transformer_config = replace( + base_transformer, + attention_backend=backend, + device="cuda", + execution_device="cuda", + sequential_cpu_offload=False, + ) + model_config = replace( + self.config.diffusion_model, + transformer=transformer_config, + scheduler=replace( + self.config.diffusion_model.scheduler, + num_inference_steps=cache.steps, + ), + audio_scheduler=replace( + self.config.diffusion_model.audio_scheduler, + num_inference_steps=cache.steps, + ), + seed=cache.seed, + ) + model = model_config.setup() + try: + self._apply_lora(model.transformer, cache) + return model.generate_joint(state) + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + def _generate_standard(self, cache: MiniMaxH3PipelineCache) -> torch.Tensor: + conditioning_started = time.monotonic() + conditioned = self._condition(cache) + cache.conditioning_seconds = time.monotonic() - conditioning_started + return self._run_native_denoise(cache, conditioned) + + def _decode_video( + self, cache: MiniMaxH3PipelineCache, latents: torch.Tensor + ) -> torch.Tensor: + from diffusers.modular_pipelines.modular_pipeline import ( + ModularPipeline, + PipelineState, + SequentialPipelineBlocks, + ) + + _write_status(cache.latent_checkpoint, "decoding-video") + pipe = ModularPipeline.from_pretrained( + self.config.model_id, + workflow=cache.workflow, + cache_dir=self._cache_dir(), + ) + video_block = pipe._blocks.sub_blocks.get("decode.video") + if video_block is None: + raise RuntimeError("MiniMax H3 workflow has no video decode block") + pipe._blocks = SequentialPipelineBlocks.from_blocks_dict({"video": video_block}) + pipe.load_components( + names=["vae", "video_processor"], + dtype={"vae": torch.float32}, + ) + pipe.vae.encoder.to_empty(device="cpu") + pipe.vae.quant_conv.to_empty(device="cpu") + pipe.vae.post_quant_conv.to("cuda") + pipe.vae.decoder.to("cuda") + first_encoder_parameter = next(pipe.vae.encoder.parameters()) + first_encoder_parameter.data = torch.empty_like( + first_encoder_parameter, device="cuda" + ) + + state = PipelineState() + state.set("latents", latents.to("cuda")) + state.set("output_type", "np") + results = pipe(state=state, output=["videos"]) + video = np.asarray(results["videos"][0]) + frames = torch.from_numpy(video).permute(0, 3, 1, 2).float().mul(2).sub(1) + _release_pipeline(pipe) + return frames.contiguous() diff --git a/integrations/minimax_h3/minimax_h3/references.py b/integrations/minimax_h3/minimax_h3/references.py new file mode 100644 index 000000000..613dfe09b --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/references.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Ordered local-media references for MiniMax H3 ref2va requests.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, cast + +import torch + +ReferenceKind = Literal["image", "video", "audio"] + + +@dataclass(frozen=True) +class MiniMaxH3ReferenceSpec: + """One validated ``kind:path`` reference, preserving request order.""" + + kind: ReferenceKind + path: Path + + def manifest(self) -> dict[str, str | int]: + """Return the source identity used by restart-safe checkpoints.""" + stat = self.path.stat() + return { + "kind": self.kind, + "path": str(self.path), + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + } + + +def parse_reference_specs( + entries: Sequence[str], +) -> tuple[MiniMaxH3ReferenceSpec, ...]: + """Parse and enforce H3's documented ordered-reference limits.""" + specs: list[MiniMaxH3ReferenceSpec] = [] + for entry in entries: + kind, separator, path_value = entry.partition(":") + if not separator or kind not in {"image", "video", "audio"}: + raise ValueError( + f"invalid reference {entry!r}; expected image:path, video:path, or audio:path" + ) + path = Path(path_value).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"reference file not found: {path}") + specs.append(MiniMaxH3ReferenceSpec(kind=cast(ReferenceKind, kind), path=path)) + + if not specs: + raise ValueError("ref2va requires at least one --reference") + limits = {"image": 9, "video": 3, "audio": 3} + for kind, limit in limits.items(): + count = sum(spec.kind == kind for spec in specs) + if count > limit: + raise ValueError(f"MiniMax H3 accepts at most {limit} {kind} references") + if len(specs) > 12: + raise ValueError("MiniMax H3 accepts at most 12 references in total") + if all(spec.kind == "audio" for spec in specs): + raise ValueError( + "an audio reference must be paired with an image or video reference" + ) + return tuple(specs) + + +def load_references(specs: tuple[MiniMaxH3ReferenceSpec, ...]) -> list[Any]: + """Decode references through Diffusers' official H3 media containers.""" + from diffusers.modular_pipelines.minimax_h3 import ( + MiniMaxH3AudioReference, + MiniMaxH3ImageReference, + MiniMaxH3VideoReference, + ) + + classes = { + "image": MiniMaxH3ImageReference, + "video": MiniMaxH3VideoReference, + "audio": MiniMaxH3AudioReference, + } + references: list[Any] = [classes[spec.kind].from_file(spec.path) for spec in specs] + for reference in references: + if not reference.has_audio or reference.sample_rate in {None, 32000}: + continue + import av + import numpy as np + + waveform = reference.audio.detach().cpu().to(torch.float32).numpy() + layout = "mono" if waveform.shape[0] == 1 else "stereo" + frame = av.AudioFrame.from_ndarray(waveform, format="fltp", layout=layout) + frame.sample_rate = reference.sample_rate + resampler = av.AudioResampler(format="fltp", layout=layout, rate=32000) + resampled = [*resampler.resample(frame), *resampler.resample(None)] + reference.audio = torch.from_numpy( + np.concatenate([chunk.to_ndarray() for chunk in resampled], axis=-1) + ).to(torch.float32) + reference.sample_rate = 32000 + return references + + +__all__ = [ + "MiniMaxH3ReferenceSpec", + "ReferenceKind", + "load_references", + "parse_reference_specs", +] diff --git a/integrations/minimax_h3/minimax_h3/runner.py b/integrations/minimax_h3/minimax_h3/runner.py new file mode 100644 index 000000000..4e8609409 --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/runner.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 text, keyframe, and ordered-reference runners.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Annotated, Literal + +from loguru import logger +from tyro.conf import UseAppendAction + +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import runner_artifact_path, write_runner_stats +from flashdreams.runtime.video_output import Mp4VideoOutputTarget +from minimax_h3.pipeline import MiniMaxH3Pipeline, MiniMaxH3PipelineCache +from minimax_h3.references import parse_reference_specs + + +@dataclass(kw_only=True) +class MiniMaxH3RunnerConfig(RunnerConfig): + """Options shared by all released MiniMax H3 workflows.""" + + _target: type[MiniMaxH3Runner] = field(default_factory=lambda: MiniMaxH3Runner) + + prompt: str = "Animate this scene with coherent natural motion." + """Text description of the desired video motion and appearance.""" + + pixel_height: int = 768 + """Output video height in pixels.""" + + pixel_width: int = 768 + """Output video width in pixels.""" + + duration: float = 5.0 + """Requested duration before H3 frame-grid alignment.""" + + steps: int = 30 + """Number of scheduler grid points.""" + + seed: int = 42 + """CPU generator seed used by both H3 schedulers.""" + + low_ram: bool = True + """Split conditioning, denoising, and decoding into checkpointed stages.""" + + restart: bool = False + """Ignore matching stage checkpoints and regenerate the full rollout.""" + + attention: Literal["auto", "flash", "default"] = "auto" + """FlashDreams native SDPA backend selection.""" + + lora: str | None = None + """Local Musubi adapter path or Hugging Face repository ID.""" + + lora_weight_name: str | None = None + """Adapter filename override for Hugging Face repositories.""" + + lora_scale: float = 1.0 + """LoRA adapter strength.""" + + fps: int = 24 + """H3's fixed output frame rate.""" + + postprocess_output_layout: VideoTensorLayout | None = "tchw" + """Decoded H3 frame layout used by FlashDreams runtime output.""" + + +@dataclass(kw_only=True) +class MiniMaxH3T2VARunnerConfig(MiniMaxH3RunnerConfig): + """Runner config for prompt-only generation.""" + + _target: type[MiniMaxH3T2VARunner] = field( + default_factory=lambda: MiniMaxH3T2VARunner + ) + + +@dataclass(kw_only=True) +class MiniMaxH3FL2VARunnerConfig(MiniMaxH3RunnerConfig): + """Runner config for first-frame, last-frame, or dual-keyframe generation.""" + + _target: type[MiniMaxH3FL2VARunner] = field( + default_factory=lambda: MiniMaxH3FL2VARunner + ) + + image_path: Path | None = None + """Optional first-frame image path.""" + + last_image_path: Path | None = None + """Optional last-frame image path.""" + + +@dataclass(kw_only=True) +class MiniMaxH3Ref2VARunnerConfig(MiniMaxH3RunnerConfig): + """Runner config for ordered image, video, and audio references.""" + + _target: type[MiniMaxH3Ref2VARunner] = field( + default_factory=lambda: MiniMaxH3Ref2VARunner + ) + + reference: Annotated[list[str], UseAppendAction] = field(default_factory=list) + """Ordered ``image:path``, ``video:path``, or ``audio:path`` references.""" + + +class MiniMaxH3Runner(Runner[MiniMaxH3RunnerConfig, MiniMaxH3Pipeline]): + """Drive one H3 workflow and persist its video-only artifact.""" + + config: MiniMaxH3RunnerConfig + pipeline: MiniMaxH3Pipeline + + def _initialize_cache(self, output_path: Path) -> MiniMaxH3PipelineCache: + raise NotImplementedError + + def _initialize_common( + self, + output_path: Path, + *, + image_path: Path | None = None, + last_image_path: Path | None = None, + reference: list[str] | tuple[str, ...] = (), + ) -> MiniMaxH3PipelineCache: + config = self.config + return self.pipeline.initialize_cache( + prompt=config.prompt, + image_path=image_path, + last_image_path=last_image_path, + references=parse_reference_specs(reference) if reference else (), + output_path=output_path, + width=config.pixel_width, + height=config.pixel_height, + duration=config.duration, + steps=config.steps, + seed=config.seed, + low_ram=config.low_ram, + restart=config.restart, + attention=config.attention, + lora=config.lora, + lora_weight_name=config.lora_weight_name, + lora_scale=config.lora_scale, + ) + + def run(self) -> None: + """Generate the single H3 step and write a video-only MP4.""" + config = self.config + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + cache = self._initialize_cache(video_path) + output_stream = self.create_video_output_stream(fps=config.fps) + output_target = Mp4VideoOutputTarget( + output_path=video_path, + fps=config.fps, + output_layout=output_stream.output_layout, + enabled=self.is_rank_zero, + ) + output_target.open() + frames = self.pipeline.generate(0, cache) + metrics = self.pipeline.finalize(0, cache) + output_target.write( + output_stream.process(frames, autoregressive_index=0, metrics=metrics) + ) + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if not artifacts: + return + self.pipeline.mark_complete(cache) + video_artifact = artifacts[0] + logger.info( + f"[{config.runner_name}] wrote {tuple(frames.shape)} video to " + f"{Path(video_artifact.uri).resolve()}" + ) + stats_history = video_artifact.metadata["stats_history"] + if stats_history: + stats_path = write_runner_stats( + config.output_dir, + config.runner_name, + list(stats_history), + ) + logger.info(f"[{config.runner_name}] wrote stats to {stats_path.resolve()}") + + +class MiniMaxH3T2VARunner(MiniMaxH3Runner): + """Prompt-only H3 runner.""" + + config: MiniMaxH3T2VARunnerConfig + + def _initialize_cache(self, output_path: Path) -> MiniMaxH3PipelineCache: + return self._initialize_common(output_path) + + +class MiniMaxH3FL2VARunner(MiniMaxH3Runner): + """First-frame, last-frame, or dual-keyframe H3 runner.""" + + config: MiniMaxH3FL2VARunnerConfig + + def _initialize_cache(self, output_path: Path) -> MiniMaxH3PipelineCache: + return self._initialize_common( + output_path, + image_path=self.config.image_path, + last_image_path=self.config.last_image_path, + ) + + +class MiniMaxH3Ref2VARunner(MiniMaxH3Runner): + """Ordered-reference H3 runner.""" + + config: MiniMaxH3Ref2VARunnerConfig + + def _initialize_cache(self, output_path: Path) -> MiniMaxH3PipelineCache: + return self._initialize_common(output_path, reference=self.config.reference) + + +__all__ = [ + "MiniMaxH3FL2VARunner", + "MiniMaxH3FL2VARunnerConfig", + "MiniMaxH3Ref2VARunner", + "MiniMaxH3Ref2VARunnerConfig", + "MiniMaxH3Runner", + "MiniMaxH3RunnerConfig", + "MiniMaxH3T2VARunner", + "MiniMaxH3T2VARunnerConfig", +] diff --git a/integrations/minimax_h3/minimax_h3/scheduler.py b/integrations/minimax_h3/minimax_h3/scheduler.py new file mode 100644 index 000000000..6a5b01ac0 --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/scheduler.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FlashDreams scheduler for MiniMax H3's data-ward velocity.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch +from torch import Tensor + +from flashdreams.infra.diffusion.scheduler import ( + FlowPredictor, + Scheduler, + SchedulerConfig, +) + + +@dataclass(kw_only=True) +class MiniMaxH3SchedulerConfig(SchedulerConfig): + """Configuration for one of H3's modality-specific schedules.""" + + _target: type[MiniMaxH3Scheduler] = field( + default_factory=lambda: MiniMaxH3Scheduler + ) + num_inference_steps: int = 30 + shift: float = 12.0 + + +class MiniMaxH3Scheduler(Scheduler): + """Rectified-flow Euler schedule used by the released H3 checkpoint.""" + + config: MiniMaxH3SchedulerConfig + + def __init__(self, config: MiniMaxH3SchedulerConfig) -> None: + super().__init__(config) + if config.num_inference_steps < 2: + raise ValueError("num_inference_steps must be at least 2") + if config.shift <= 0: + raise ValueError("shift must be positive") + + def schedule(self, device: torch.device | str) -> tuple[Tensor, Tensor]: + """Return the shifted sigma grid and its H3 timesteps.""" + base = torch.linspace( + 1.0, 0.0, self.config.num_inference_steps, dtype=torch.float32 + ) + shift = self.config.shift + sigmas = torch.unique_consecutive(shift * base / (1 + (shift - 1) * base)) + sigmas = sigmas.to(device) + return sigmas, 1.0 - sigmas[:-1] + + @staticmethod + def step( + sample: Tensor, + flow: Tensor, + timestep: Tensor, + sigma: Tensor, + sigma_next: Tensor, + ) -> Tensor: + """Take one deterministic data-ward Euler step.""" + sigma_from_timestep = 1 - timestep.to(device=sample.device, dtype=sample.dtype) + denoised = sample + sigma_from_timestep * flow + compute_dtype = ( + torch.float32 + if sample.dtype in (torch.float16, torch.bfloat16) + else sample.dtype + ) + ratio = sigma_next.to(sample.device, compute_dtype) / sigma.to( + sample.device, compute_dtype + ) + previous = ratio * sample.to(compute_dtype) + (1 - ratio) * denoised.to( + compute_dtype + ) + return previous.to(sample.dtype) + + @torch.no_grad() + def sample( + self, + initial_noise: Tensor, + predict_flow: FlowPredictor, + rng: torch.Generator | None = None, + ) -> Tensor: + """Denoise one stream with the H3 schedule.""" + del rng + sigmas, timesteps = self.schedule(initial_noise.device) + sample = initial_noise + for index, timestep in enumerate(timesteps): + flow = predict_flow(sample, timestep) + sample = self.step(sample, flow, timestep, sigmas[index], sigmas[index + 1]) + return sample + + def add_noise( + self, + clean_input: Tensor, + timestep: Tensor, + rng: torch.Generator | None = None, + ) -> Tensor: + """Mix clean input with Gaussian noise under H3's time convention.""" + noise = torch.randn( + clean_input.shape, + dtype=clean_input.dtype, + device=clean_input.device, + generator=rng, + ) + time = timestep.to(clean_input.device, clean_input.dtype) + while time.ndim < clean_input.ndim: + time = time.unsqueeze(-1) + return time * clean_input + (1 - time) * noise + + +__all__ = ["MiniMaxH3Scheduler", "MiniMaxH3SchedulerConfig"] diff --git a/integrations/minimax_h3/minimax_h3/transformer.py b/integrations/minimax_h3/minimax_h3/transformer.py new file mode 100644 index 000000000..befa71a53 --- /dev/null +++ b/integrations/minimax_h3/minimax_h3/transformer.py @@ -0,0 +1,500 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FlashDreams-native MiniMax H3 joint video/audio transformer.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Any, Literal, cast + +import torch +from torch import Tensor, nn + +from flashdreams.core.attention import NativeAttention +from flashdreams.core.checkpoint.load import load_checkpoint +from flashdreams.infra.diffusion.transformer import ( + Transformer, + TransformerAutoregressiveCache, + TransformerConfig, +) + +H3_TRANSFORMER_CHECKPOINT = ( + "https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/" + "42ed227ee7df40d41602854ae760620d6eb651fe/transformer/" + "diffusion_pytorch_model.safetensors.index.json" +) +H3_REF_TRANSFORMER_CHECKPOINT = ( + "https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/" + "42ed227ee7df40d41602854ae760620d6eb651fe/transformer_ref/" + "diffusion_pytorch_model.safetensors.index.json" +) +MODALITY_COUNT = 3 + + +@dataclass(kw_only=True) +class MiniMaxH3TransformerCache(TransformerAutoregressiveCache): + """Non-autoregressive request state required by ``predict_flow``.""" + + audio_hidden_states: Tensor + encoder_hidden_states: Tensor + timestep: Tensor + timestep_indices: Tensor + token_tags: Tensor + position_ids: Tensor + video_indices: Tensor + audio_indices: Tensor + text_indices: Tensor + last_audio_flow: Tensor | None = None + + +@dataclass(kw_only=True) +class MiniMaxH3TransformerConfig(TransformerConfig): + """Architecture and checkpoint configuration for MiniMax H3 FL2VA.""" + + _target: type[MiniMaxH3Transformer] = field( + default_factory=lambda: MiniMaxH3Transformer + ) + checkpoint_path: str | None = H3_TRANSFORMER_CHECKPOINT + checkpoint_min_free_gb: float | None = None + device: str = "cpu" + execution_device: str = "cuda" + sequential_cpu_offload: bool = True + dtype: torch.dtype = torch.bfloat16 + attention_backend: Literal["flash", "cudnn", "efficient", "math"] = "flash" + 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 _module_dtype(module: nn.Module) -> torch.dtype: + return next(module.parameters()).dtype + + +class _SwiGLU(nn.Module): + def __init__(self, dim: int, inner_dim: int, **factory: Any) -> None: + super().__init__() + self.proj = nn.Linear(dim, inner_dim * 2, bias=False, **factory) + + def forward(self, hidden_states: Tensor) -> Tensor: + value, gate = self.proj(hidden_states).chunk(2, dim=-1) + return value * nn.functional.silu(gate) + + +class _FeedForward(nn.Module): + def __init__(self, dim: int, inner_dim: int, **factory: Any) -> None: + super().__init__() + self.net = nn.ModuleList( + [ + _SwiGLU(dim, inner_dim, **factory), + nn.Dropout(0.0), + nn.Linear(inner_dim, dim, bias=False, **factory), + ] + ) + + def forward(self, hidden_states: Tensor) -> Tensor: + for layer in self.net: + hidden_states = layer(hidden_states) + return hidden_states + + +class _TimestepEmbedding(nn.Module): + def __init__( + self, in_dim: int, hidden_dim: int, out_dim: int, **factory: Any + ) -> None: + super().__init__() + self.linear_1 = nn.Linear(in_dim, hidden_dim, **factory) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(hidden_dim, out_dim, **factory) + + def forward(self, sample: Tensor) -> Tensor: + return self.linear_2(self.act(self.linear_1(sample))) + + +class _RotaryEmbedding(nn.Module): + def __init__( + self, + freq_dim: int, + theta: float, + *, + device: torch.device, + ) -> None: + super().__init__() + inv_freq = 1.0 / ( + theta + ** ( + torch.arange(0, 2 * freq_dim, 2, dtype=torch.float32, device=device) + / (2 * freq_dim) + ) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, position_ids: Tensor) -> tuple[Tensor, Tensor]: + inv_freq = cast(Tensor, self.inv_freq) + frequencies = position_ids.float().unsqueeze(-1) * inv_freq.view(1, 1, -1) + frequencies = torch.cat(frequencies.unbind(dim=1), dim=-1) + frequencies = torch.cat((frequencies, frequencies), dim=-1) + return frequencies.cos(), frequencies.sin() + + +def _apply_rotary(hidden_states: Tensor, cos: Tensor, sin: Tensor) -> Tensor: + rotary_dim = cos.shape[-1] + rotary, passthrough = ( + hidden_states[..., :rotary_dim], + hidden_states[..., rotary_dim:], + ) + cos = cos.to(hidden_states.dtype)[None, :, None, :] + sin = sin.to(hidden_states.dtype)[None, :, None, :] + first, second = rotary.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + return torch.cat((rotary * cos + rotated * sin, passthrough), dim=-1).contiguous() + + +class _Attention(nn.Module): + def __init__( + self, + hidden_size: int, + heads: int, + head_dim: int, + eps: float, + backend: Literal["flash", "cudnn", "efficient", "math"], + **factory: Any, + ) -> None: + super().__init__() + inner_dim = heads * head_dim + self.heads = heads + self.to_q = nn.Linear(hidden_size, inner_dim, bias=False, **factory) + self.to_k = nn.Linear(hidden_size, inner_dim, bias=False, **factory) + self.to_v = nn.Linear(hidden_size, inner_dim, bias=False, **factory) + self.norm_q = nn.RMSNorm(head_dim, eps=eps, **factory) + self.norm_k = nn.RMSNorm(head_dim, eps=eps, **factory) + self.to_out = nn.ModuleList( + [nn.Linear(inner_dim, hidden_size, bias=False, **factory), nn.Dropout(0.0)] + ) + self.attn_op = NativeAttention(qkv_format="bshd", backend=backend) + + def forward( + self, hidden_states: Tensor, rotary: tuple[Tensor, Tensor] | None = None + ) -> Tensor: + query = self.norm_q(self.to_q(hidden_states).unflatten(-1, (self.heads, -1))) + key = self.norm_k(self.to_k(hidden_states).unflatten(-1, (self.heads, -1))) + value = self.to_v(hidden_states).unflatten(-1, (self.heads, -1)) + if rotary is not None: + query = _apply_rotary(query, *rotary) + key = _apply_rotary(key, *rotary) + output = self.attn_op(query, key, value).flatten(2, 3).type_as(query) + return self.to_out[1](self.to_out[0](output)) + + +class _RefinerBlock(nn.Module): + def __init__(self, config: MiniMaxH3TransformerConfig, **factory: Any) -> None: + super().__init__() + self.norm1 = nn.RMSNorm(config.hidden_size, eps=config.norm_eps, **factory) + self.attn = _Attention( + config.hidden_size, + config.num_attention_heads, + config.attention_head_dim, + config.qk_norm_eps, + config.attention_backend, + **factory, + ) + self.norm2 = nn.RMSNorm(config.hidden_size, eps=config.norm_eps, **factory) + self.ff = _FeedForward(config.hidden_size, config.ffn_dim, **factory) + + def forward(self, hidden_states: Tensor) -> Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states)) + return hidden_states + self.ff(self.norm2(hidden_states)) + + +class _TokenRefiner(nn.Module): + def __init__(self, config: MiniMaxH3TransformerConfig, **factory: Any) -> None: + super().__init__() + self.refiner_blocks = nn.ModuleList( + [_RefinerBlock(config, **factory) for _ in range(config.num_refiner_layers)] + ) + self.final_norm = nn.RMSNorm( + config.hidden_size, eps=config.final_norm_eps, **factory + ) + + def forward(self, hidden_states: Tensor) -> Tensor: + for block in self.refiner_blocks: + hidden_states = block(hidden_states) + return self.final_norm(hidden_states) + + +class _AdaLNProjection(nn.Module): + def __init__(self, config: MiniMaxH3TransformerConfig, **factory: Any) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.linear = nn.Linear( + config.time_embed_dim, + 6 * config.hidden_size * MODALITY_COUNT, + **factory, + ) + + def forward(self, temb: Tensor) -> tuple[Tensor, ...]: + output = self.linear(nn.functional.silu(temb).to(_module_dtype(self.linear))) + return output.view(-1, 6 * self.hidden_size).chunk(6, dim=-1) + + +class _TransformerBlock(nn.Module): + def __init__(self, config: MiniMaxH3TransformerConfig, **factory: Any) -> None: + super().__init__() + self.norm1 = nn.RMSNorm(config.hidden_size, eps=config.norm_eps, **factory) + self.attn = _Attention( + config.hidden_size, + config.num_attention_heads, + config.attention_head_dim, + config.qk_norm_eps, + config.attention_backend, + **factory, + ) + self.norm2 = nn.RMSNorm(config.hidden_size, eps=config.norm_eps, **factory) + self.ff = _FeedForward(config.hidden_size, config.ffn_dim, **factory) + self.adaln_proj = _AdaLNProjection(config, **factory) + + def forward( + self, + hidden_states: Tensor, + temb: Tensor, + adaln_indices: Tensor, + rotary: tuple[Tensor, Tensor], + ) -> Tensor: + shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = self.adaln_proj(temb) + normalized = self.norm1(hidden_states) + normalized = normalized * (1 + scale_a[adaln_indices]) + shift_a[adaln_indices] + hidden_states = hidden_states + gate_a[adaln_indices] * self.attn( + normalized, rotary + ) + normalized = self.norm2(hidden_states) + normalized = normalized * (1 + scale_m[adaln_indices]) + shift_m[adaln_indices] + return hidden_states + gate_m[adaln_indices] * self.ff(normalized) + + +class _OutputNorm(nn.Module): + def __init__(self, config: MiniMaxH3TransformerConfig, **factory: Any) -> None: + super().__init__() + self.norm = nn.RMSNorm(config.hidden_size, eps=config.final_norm_eps, **factory) + self.linear = nn.Linear( + config.time_embed_dim, 2 * config.hidden_size, **factory + ) + + def forward( + self, hidden_states: Tensor, temb: Tensor, timestep_indices: Tensor + ) -> Tensor: + shift, scale = self.linear( + nn.functional.silu(temb).to(_module_dtype(self.linear)) + ).chunk(2, dim=-1) + hidden_states = self.norm(hidden_states) + return hidden_states * (1 + scale[timestep_indices]) + shift[timestep_indices] + + +class MiniMaxH3Transformer(Transformer[MiniMaxH3TransformerCache]): + """Native FlashDreams transformer for H3's packed multimodal sequence.""" + + config: MiniMaxH3TransformerConfig + + def __init__(self, config: MiniMaxH3TransformerConfig) -> None: + super().__init__(config) + self.config = config + device = torch.device(config.device) + low_precision: dict[str, Any] = {"device": device, "dtype": config.dtype} + full_precision: dict[str, Any] = { + "device": device, + "dtype": torch.float32, + } + video_dim = config.in_channels * math.prod(config.patch_size) + + self.proj_in = nn.Linear(video_dim, config.hidden_size, **full_precision) + self.audio_proj_in = nn.Linear( + config.audio_in_channels, config.hidden_size, **full_precision + ) + self.context_embedder = nn.Linear( + config.text_dim, config.hidden_size, **low_precision + ) + self.time_embedder = _TimestepEmbedding( + config.freq_dim, + config.time_embed_hidden_dim, + config.time_embed_dim, + **full_precision, + ) + self.rope = _RotaryEmbedding( + config.rope_freq_dim, config.rope_theta, device=device + ) + self.token_refiner = _TokenRefiner(config, **low_precision) + self.transformer_blocks = nn.ModuleList( + [ + _TransformerBlock(config, **low_precision) + for _ in range(config.num_layers) + ] + ) + self.norm_out = _OutputNorm(config, **low_precision) + self.proj_out = nn.Linear(config.hidden_size, video_dim, **full_precision) + self.audio_proj_out = nn.Linear( + config.hidden_size, config.audio_in_channels, **full_precision + ) + if config.checkpoint_path is not None: + load_checkpoint( + config.checkpoint_path, + model=self, + checkpoint_min_free_gb=config.checkpoint_min_free_gb, + ) + self.eval() + + @property + def latent_shape(self) -> tuple[int, ...]: + return () + + @staticmethod + def _time_projection(timestep: Tensor, dim: int) -> Tensor: + half = dim // 2 + exponent = ( + -math.log(10000) + * torch.arange(half, dtype=torch.float32, device=timestep.device) + / half + ) + angles = timestep[:, None].float() * torch.exp(exponent)[None] + embedding = torch.cat((angles.cos(), angles.sin()), dim=-1) + if dim % 2: + embedding = nn.functional.pad(embedding, (0, 1)) + return embedding + + def _run_on_execution_device(self, module: nn.Module, *args: Any) -> Any: + """Run one weight group on the accelerator, then return it to CPU.""" + execution_device = torch.device(self.config.execution_device) + if not self.config.sequential_cpu_offload: + return module(*args) + module.to(execution_device) + try: + return module(*args) + finally: + torch.cuda.synchronize(execution_device) + module.to("cpu") + torch.cuda.empty_cache() + + def forward_joint( + self, + hidden_states: Tensor, + audio_hidden_states: Tensor, + encoder_hidden_states: Tensor, + timestep: Tensor, + timestep_indices: Tensor, + token_tags: Tensor, + position_ids: Tensor, + video_indices: Tensor, + audio_indices: Tensor, + text_indices: Tensor, + ) -> tuple[Tensor, Tensor]: + """Predict video and audio velocity for one packed H3 denoising step.""" + execution_device = torch.device(self.config.execution_device) + if self.config.sequential_cpu_offload: + hidden_states = hidden_states.to(execution_device) + audio_hidden_states = audio_hidden_states.to(execution_device) + encoder_hidden_states = encoder_hidden_states.to(execution_device) + timestep = timestep.to(execution_device) + timestep_indices = timestep_indices.to(execution_device) + token_tags = token_tags.to(execution_device) + position_ids = position_ids.to(execution_device) + video_indices = video_indices.to(execution_device) + audio_indices = audio_indices.to(execution_device) + text_indices = text_indices.to(execution_device) + sequence_length = position_ids.shape[0] + if position_ids.shape != (sequence_length, 3): + raise ValueError("position_ids must have shape [sequence_length, 3]") + self.rope.to(execution_device) + rotary = self.rope(position_ids) + video = self._run_on_execution_device( + self.proj_in, hidden_states.to(_module_dtype(self.proj_in)) + ) + audio = self._run_on_execution_device( + self.audio_proj_in, + audio_hidden_states.to(_module_dtype(self.audio_proj_in)), + ) + text = self._run_on_execution_device( + self.context_embedder, + encoder_hidden_states.to(_module_dtype(self.context_embedder)), + ) + text = self._run_on_execution_device(self.token_refiner, text) + packed = text.new_zeros((text.shape[0], sequence_length, text.shape[-1])) + packed = packed.index_copy(1, text_indices, text) + packed = packed.index_copy(1, video_indices, video.to(text.dtype)) + packed = packed.index_copy(1, audio_indices, audio.to(text.dtype)) + + temb = self._run_on_execution_device( + self.time_embedder, + self._time_projection(timestep, self.config.freq_dim).to( + _module_dtype(self.time_embedder) + ), + ) + adaln_indices = timestep_indices * MODALITY_COUNT + token_tags + for block in self.transformer_blocks: + packed = self._run_on_execution_device( + block, packed, temb, adaln_indices, rotary + ) + packed = self._run_on_execution_device( + self.norm_out, packed, temb, timestep_indices + ).to(_module_dtype(self.proj_out)) + return ( + self._run_on_execution_device(self.proj_out, packed).index_select( + 1, video_indices + ), + self._run_on_execution_device(self.audio_proj_out, packed).index_select( + 1, audio_indices + ), + ) + + def predict_flow( + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: MiniMaxH3TransformerCache, + input: object = None, + ) -> Tensor: + """Adapt the joint forward to FlashDreams' transformer interface.""" + del timestep, input + video_flow, cache.last_audio_flow = self.forward_joint( + noisy_latent, + cache.audio_hidden_states, + cache.encoder_hidden_states, + cache.timestep, + cache.timestep_indices, + cache.token_tags, + cache.position_ids, + cache.video_indices, + cache.audio_indices, + cache.text_indices, + ) + return video_flow + + def patchify_and_maybe_split_cp(self, x: object) -> object: + """H3 inputs arrive already packed into sequence rows.""" + return x + + def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor: + """H3's custom pipeline unpacks rows after paired denoising.""" + return x + + +__all__ = [ + "H3_REF_TRANSFORMER_CHECKPOINT", + "H3_TRANSFORMER_CHECKPOINT", + "MiniMaxH3Transformer", + "MiniMaxH3TransformerCache", + "MiniMaxH3TransformerConfig", +] diff --git a/integrations/minimax_h3/pyproject.toml b/integrations/minimax_h3/pyproject.toml new file mode 100644 index 000000000..91c49711a --- /dev/null +++ b/integrations/minimax_h3/pyproject.toml @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "flashdreams-minimax-h3" +version = "0.1.0" +description = "MiniMax H3 video generation workflows for the FlashDreams runtime." +requires-python = ">=3.10" +dependencies = [ + "accelerate>=1.12", + "av>=16", + "diffusers @ git+https://github.com/huggingface/diffusers.git@175fe6b2419a01db9c2ceabd01ec37d2c0305fc2", + "flashdreams", + "huggingface-hub>=0.33", + "numpy>=1.24,<2.5", + "pillow>=11", + "safetensors>=0.4", + "sentencepiece>=0.2", + "torch>=2.9", + "transformers @ git+https://github.com/huggingface/transformers.git@d1123114da1ab4395198146f4f84dae7fe8b693e", +] + +[tool.uv.sources] +flashdreams = { workspace = true } + +[project.optional-dependencies] +dev = ["pytest>=8.0", "tomli>=2.0"] + +[project.entry-points."flashdreams.runner_configs"] +"minimax-h3-t2va" = "minimax_h3.config:RUNNER_MINIMAX_H3_T2VA" +"minimax-h3-fl2va" = "minimax_h3.config:RUNNER_MINIMAX_H3_FL2VA" +"minimax-h3-ref2va" = "minimax_h3.config:RUNNER_MINIMAX_H3_REF2VA" + +[tool.setuptools.packages.find] +include = ["minimax_h3*"] +exclude = ["tests"] + +[tool.uv] +managed = true diff --git a/integrations/minimax_h3/tests/test_smoke.py b/integrations/minimax_h3/tests/test_smoke.py new file mode 100644 index 000000000..a39111667 --- /dev/null +++ b/integrations/minimax_h3/tests/test_smoke.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU smoke tests for the MiniMax H3 runner plugin.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, cast + +import pytest +import tomli as tomllib +import torch +from minimax_h3 import config as config_mod +from minimax_h3.config import ( + PIPELINE_MINIMAX_H3_FL2VA, + PIPELINE_MINIMAX_H3_REF2VA, + PIPELINE_MINIMAX_H3_T2VA, + RUNNER_CONFIGS, + RUNNER_MINIMAX_H3_FL2VA, + RUNNER_MINIMAX_H3_REF2VA, + RUNNER_MINIMAX_H3_T2VA, +) +from minimax_h3.constants import align_num_frames, validate_canvas +from minimax_h3.lora import convert_musubi_lora +from minimax_h3.pipeline import MiniMaxH3Pipeline +from minimax_h3.references import parse_reference_specs +from minimax_h3.runner import ( + MiniMaxH3FL2VARunner, + MiniMaxH3Ref2VARunner, + MiniMaxH3RunnerConfig, + MiniMaxH3T2VARunner, +) +from minimax_h3.scheduler import MiniMaxH3SchedulerConfig +from minimax_h3.transformer import MiniMaxH3TransformerConfig + +from flashdreams.infra.diffusion.transformer import Transformer +from flashdreams.infra.runner import RunnerConfig + +pytestmark = pytest.mark.ci_cpu + +ENTRY_POINT_GROUP = "flashdreams.runner_configs" + + +def test_runners_dict_is_non_empty() -> None: + """Plugin must expose at least one runner.""" + assert RUNNER_CONFIGS, "RUNNER_CONFIGS is empty" + + +def test_runner_name_mirrors_pipeline_name() -> None: + """Runner names must match pipeline names for CLI discovery.""" + drifted = { + slug: (config.runner_name, config.pipeline.name) + for slug, config in RUNNER_CONFIGS.items() + if config.runner_name != config.pipeline.name + } + assert not drifted, f"runner_name != pipeline.name: {drifted}" + + +def test_runners_have_descriptions() -> None: + """Every registered runner must have a CLI description.""" + empty = [ + slug + for slug, config in RUNNER_CONFIGS.items() + if not config.description.strip() + ] + assert not empty, f"runners missing description: {empty}" + + +def test_entry_points_match_module_literals() -> None: + """Package entry points must resolve to every registered runner literal.""" + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + with pyproject.open("rb") as handle: + metadata = tomllib.load(handle) + entries = metadata["project"]["entry-points"][ENTRY_POINT_GROUP] + assert set(entries) == set(RUNNER_CONFIGS) + + for slug, target in entries.items(): + module_name, attribute = target.split(":", 1) + assert module_name == "minimax_h3.config" + config = cast(RunnerConfig, getattr(config_mod, attribute)) + assert config.runner_name == slug + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="entry-point discovery test relies on importlib.metadata 3.10+ shape", +) +def test_entry_points_discoverable_when_installed() -> None: + """Installed plugin entry points must expose every registered runner.""" + from importlib.metadata import entry_points + + entries = entry_points(group=ENTRY_POINT_GROUP) + discovered = { + entry.name for entry in entries if entry.value.startswith("minimax_h3.") + } + if not discovered: + pytest.skip("plugin not installed; run uv sync from the repository root") + assert discovered == set(RUNNER_CONFIGS) + + +def test_pipeline_config_constructs_without_loading_weights() -> None: + """Construct every runtime pipeline without network or checkpoint access.""" + pipelines = { + config.workflow: config.setup() + for config in ( + PIPELINE_MINIMAX_H3_T2VA, + PIPELINE_MINIMAX_H3_FL2VA, + PIPELINE_MINIMAX_H3_REF2VA, + ) + } + assert set(pipelines) == {"t2va", "fl2va", "ref2va"} + assert all( + isinstance(pipeline, MiniMaxH3Pipeline) for pipeline in pipelines.values() + ) + assert all( + pipeline.config.model_id == "MiniMaxAI/MiniMax-H3" + for pipeline in pipelines.values() + ) + assert RUNNER_MINIMAX_H3_T2VA._target is MiniMaxH3T2VARunner + assert RUNNER_MINIMAX_H3_FL2VA._target is MiniMaxH3FL2VARunner + assert RUNNER_MINIMAX_H3_REF2VA._target is MiniMaxH3Ref2VARunner + for config in ( + PIPELINE_MINIMAX_H3_T2VA, + PIPELINE_MINIMAX_H3_FL2VA, + PIPELINE_MINIMAX_H3_REF2VA, + ): + assert issubclass(config.diffusion_model.transformer._target, Transformer) + + +def test_low_ram_is_an_explicit_default_flag() -> None: + """Default to crash-safe staging without enabling a third-party LoRA.""" + for runner in RUNNER_CONFIGS.values(): + assert isinstance(runner, MiniMaxH3RunnerConfig) + assert runner.low_ram is True + assert runner.lora is None + + +def test_native_bf16_gpu_path_is_default() -> None: + """Keep the quality-preserving low-host-RAM path on the accelerator.""" + transformer = PIPELINE_MINIMAX_H3_FL2VA.diffusion_model.transformer + assert isinstance(transformer, MiniMaxH3TransformerConfig) + assert transformer.device == "cuda" + assert transformer.sequential_cpu_offload is False + + +def test_musubi_lora_conversion_targets_native_layers(tmp_path: Path) -> None: + """Convert all H3 block targets without introducing a default adapter.""" + from safetensors.torch import load_file, save_file + + source = tmp_path / "adapter.safetensors" + tensors: dict[str, torch.Tensor] = {} + for block in range(50): + for module in ("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2"): + prefix = f"lora_unet_blocks_{block}_{module}" + tensors[f"{prefix}.alpha"] = torch.tensor(2.0) + tensors[f"{prefix}.lora_down.weight"] = torch.ones(2, 3) + out_features = 6 if module == "attn_qkv_proj" else 4 + tensors[f"{prefix}.lora_up.weight"] = torch.ones(out_features, 2) + save_file(tensors, source) + + converted = load_file( + convert_musubi_lora(source, tmp_path / "converted.safetensors") + ) + assert len(converted) == 600 + assert "transformer.transformer_blocks.0.attn.to_q.lora_A.weight" in converted + assert "transformer.transformer_blocks.0.attn.to_v.lora_B.weight" in converted + assert "transformer.transformer_blocks.49.ff.net.2.lora_B.weight" in converted + + +def test_duration_and_canvas_contracts() -> None: + """Align five seconds and reject non-H3 canvas dimensions.""" + assert align_num_frames(5.0) == 124 + validate_canvas(576, 768) + with pytest.raises(ValueError, match="multiples of 32"): + validate_canvas(577, 768) + + +def test_runtime_cache_uses_stage_specific_checkpoints(tmp_path: Path) -> None: + """Derive conditioning and denoised checkpoints beside the output.""" + image = tmp_path / "image.png" + image.write_bytes(b"test") + pipeline = PIPELINE_MINIMAX_H3_FL2VA.setup() + cache = pipeline.initialize_cache( + prompt="animate", + image_path=image, + last_image_path=None, + references=(), + output_path=tmp_path / "out.mp4", + width=576, + height=768, + duration=5.0, + steps=30, + seed=42, + low_ram=True, + restart=False, + attention="auto", + lora=None, + lora_weight_name=None, + lora_scale=1.0, + ) + assert cache.conditioning_checkpoint.name == "out.mp4.conditioning.safetensors" + assert cache.latent_checkpoint.name == "out.mp4.latents.safetensors" + + +def test_reference_parser_preserves_order_and_enforces_limits(tmp_path: Path) -> None: + """Keep semantic reference order while rejecting unsupported requests.""" + image = tmp_path / "subject.png" + video = tmp_path / "motion.mp4" + audio = tmp_path / "voice.wav" + for path in (image, video, audio): + path.write_bytes(b"test") + parsed = parse_reference_specs( + (f"image:{image}", f"audio:{audio}", f"video:{video}") + ) + assert [reference.kind for reference in parsed] == ["image", "audio", "video"] + with pytest.raises(ValueError, match="paired"): + parse_reference_specs((f"audio:{audio}",)) + with pytest.raises(ValueError, match="at most 3 video"): + parse_reference_specs(tuple(f"video:{video}" for _ in range(4))) + + +def test_registered_workflows_validate_their_inputs(tmp_path: Path) -> None: + """Reject cross-workflow media instead of silently selecting another model.""" + image = tmp_path / "image.png" + image.write_bytes(b"test") + common = { + "prompt": "animate", + "output_path": tmp_path / "out.mp4", + "width": 512, + "height": 768, + "duration": 5.0, + "steps": 30, + "seed": 42, + "low_ram": True, + "restart": False, + "attention": "auto", + "lora": None, + "lora_weight_name": None, + "lora_scale": 1.0, + } + t2va = PIPELINE_MINIMAX_H3_T2VA.setup() + cache = t2va.initialize_cache( + image_path=None, last_image_path=None, references=(), **common + ) + assert cache.workflow == "t2va" + fl2va = PIPELINE_MINIMAX_H3_FL2VA.setup() + cache = fl2va.initialize_cache( + image_path=None, last_image_path=image, references=(), **common + ) + assert cache.workflow == "fl2va" + with pytest.raises(ValueError, match="requires --image-path"): + fl2va.initialize_cache( + image_path=None, last_image_path=None, references=(), **common + ) + + +def test_native_scheduler_matches_official_h3_euler() -> None: + """Match the released H3 schedule and data-ward Euler update exactly.""" + from diffusers.schedulers.scheduling_minimax_h3 import ( + MiniMaxH3Scheduler as OfficialScheduler, + ) + + official: Any = OfficialScheduler(shift=12.0) + official.set_timesteps(7, device="cpu") + native = MiniMaxH3SchedulerConfig(num_inference_steps=7, shift=12.0).setup() + sigmas, timesteps = native.schedule("cpu") + torch.testing.assert_close(sigmas, official.sigmas) + torch.testing.assert_close(timesteps, official.timesteps) + + sample = torch.randn(2, 3) + flow = torch.randn_like(sample) + expected = official.step(flow, official.timesteps[0], sample).prev_sample + actual = native.step(sample, flow, timesteps[0], sigmas[0], sigmas[1]) + torch.testing.assert_close(actual, expected) + + +def test_native_transformer_matches_official_h3_forward() -> None: + """Prove state-dict and numerical compatibility on a tiny CPU model.""" + from diffusers.models.transformers.transformer_minimax_h3 import ( + MiniMaxH3Transformer3DModel, + ) + + architecture: dict[str, Any] = { + "num_attention_heads": 2, + "attention_head_dim": 12, + "hidden_size": 16, + "num_layers": 2, + "num_refiner_layers": 1, + "ffn_dim": 32, + "in_channels": 2, + "audio_in_channels": 4, + "patch_size": (1, 1, 1), + "text_dim": 10, + "freq_dim": 8, + "time_embed_hidden_dim": 16, + "time_embed_dim": 8, + "rope_freq_dim": 2, + } + torch.manual_seed(1) + official: Any = MiniMaxH3Transformer3DModel(**architecture) + native = MiniMaxH3TransformerConfig( + checkpoint_path=None, + device="cpu", + execution_device="cpu", + sequential_cpu_offload=False, + dtype=torch.float32, + attention_backend="math", + **architecture, + ).setup() + native.load_state_dict(official.state_dict(), strict=True) + inputs = { + "hidden_states": torch.randn(1, 4, 2), + "audio_hidden_states": torch.randn(1, 2, 4), + "encoder_hidden_states": torch.randn(1, 3, 10), + "timestep": torch.tensor([0.1, 0.5, 0.999]), + "timestep_indices": torch.tensor([1, 1, 1, 2, 0, 1, 1, 0, 2]), + "token_tags": torch.tensor([1, 1, 1, 0, 0, 2, 2, 0, 0]), + "position_ids": torch.randn(9, 3), + "video_indices": torch.tensor([3, 4, 7, 8]), + "audio_indices": torch.tensor([5, 6]), + "text_indices": torch.tensor([0, 1, 2]), + } + with torch.no_grad(): + expected = official(**inputs, return_dict=False) + actual = native.forward_joint(**inputs) + for native_output, official_output in zip(actual, expected, strict=True): + torch.testing.assert_close(native_output, official_output) diff --git a/pyproject.toml b/pyproject.toml index 02f62aa87..5b38633df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ extraPaths = [ "integrations/fastvideo_causal_wan22", "integrations/hy_worldplay", "integrations/lingbot", + "integrations/minimax_h3", "integrations/sana", "integrations/self_forcing", "integrations/wan21", @@ -78,6 +79,7 @@ extra-paths = [ "integrations/fastvideo_causal_wan22", "integrations/hy_worldplay", "integrations/lingbot", + "integrations/minimax_h3", "integrations/sana", "integrations/self_forcing", "integrations/wan21", diff --git a/uv.lock b/uv.lock index 7f4aba302..8dd23450e 100644 --- a/uv.lock +++ b/uv.lock @@ -26,6 +26,7 @@ members = [ "flashdreams-flashvsr", "flashdreams-hy-worldplay", "flashdreams-lingbot", + "flashdreams-minimax-h3", "flashdreams-omnidreams", "flashdreams-sana-wm", "flashdreams-self-forcing", @@ -871,8 +872,8 @@ wheels = [ [[package]] name = "diffusers" -version = "0.38.0" -source = { registry = "https://pypi.org/simple" } +version = "0.40.0.dev0" +source = { git = "https://github.com/huggingface/diffusers.git?rev=175fe6b2419a01db9c2ceabd01ec37d2c0305fc2#175fe6b2419a01db9c2ceabd01ec37d2c0305fc2" } dependencies = [ { name = "filelock" }, { name = "httpx" }, @@ -885,10 +886,6 @@ dependencies = [ { name = "requests" }, { name = "safetensors" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/01/ed/255d3dfd4a2271dffc8f1895f9d2720b3bf1beaecf02148bb5604439e594/diffusers-0.38.0.tar.gz", hash = "sha256:1e094ec5c16f18c42fb89d37f07a94cf9aab3ebbe527ab059c609597b8857626", size = 4328401, upload-time = "2026-05-01T05:42:15.276Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/42/c0/3237566ea6e3a542f3c0669a253d62fe75f27b84b3d7bd4fb3b5ee89d73c/diffusers-0.38.0-py3-none-any.whl", hash = "sha256:18e53f9e539096320470f62c6360a6fd5727ff28cffda566265316e13fcdb612", size = 5245919, upload-time = "2026-05-01T05:42:12.779Z" }, -] [[package]] name = "distlib" @@ -1275,6 +1272,51 @@ requires-dist = [ ] provides-extras = ["dev"] +[[package]] +name = "flashdreams-minimax-h3" +version = "0.1.0" +source = { editable = "integrations/minimax_h3" } +dependencies = [ + { name = "accelerate" }, + { name = "av" }, + { name = "diffusers" }, + { name = "flashdreams" }, + { name = "huggingface-hub" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "pillow" }, + { name = "safetensors" }, + { name = "sentencepiece" }, + { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'win32' and extra == 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'win32' and extra == 'extra-11-flashdreams-dev') or (sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "transformers" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, + { name = "tomli" }, +] + +[package.metadata] +requires-dist = [ + { name = "accelerate", specifier = ">=1.12" }, + { name = "av", specifier = ">=16" }, + { name = "diffusers", git = "https://github.com/huggingface/diffusers.git?rev=175fe6b2419a01db9c2ceabd01ec37d2c0305fc2" }, + { name = "flashdreams", editable = "flashdreams" }, + { name = "huggingface-hub", specifier = ">=0.33" }, + { name = "numpy", specifier = ">=1.24,<2.5" }, + { name = "pillow", specifier = ">=11" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "safetensors", specifier = ">=0.4" }, + { name = "sentencepiece", specifier = ">=0.2" }, + { name = "tomli", marker = "extra == 'dev'", specifier = ">=2.0" }, + { name = "torch", specifier = ">=2.9" }, + { name = "transformers", git = "https://github.com/huggingface/transformers.git?rev=d1123114da1ab4395198146f4f84dae7fe8b693e" }, +] +provides-extras = ["dev"] + [[package]] name = "flashdreams-omnidreams" version = "0.1.0" @@ -1733,18 +1775,18 @@ wheels = [ [[package]] name = "hf-xet" -version = "1.5.1" +version = "1.6.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/4b/2d/57fd21d84d93efb4bd0b962383790e19dd1bc053501b4264c97903b4e83e/hf_xet-1.5.1.tar.gz", hash = "sha256:51ef4500dab3764b41135ee1381a4b62ce56fc54d4c92b719b59e597d6df5bf6", size = 876636, upload-time = "2026-06-08T23:02:53.897Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1b/ab/522a2ab67f27971a9d48ca666d4fca85ef7d5282d142e31fd087e27b1bbe/hf_xet-1.6.0.tar.gz", hash = "sha256:2e58454a340b3556dfa4972d5451aff4fba8dd42a236600ba1a1d2b1514f0fef", size = 920527, upload-time = "2026-08-03T22:33:13.243Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/d8/5e54cf37434759d1f4f2ba9b66077ff9d4c4e1f37b6bd7975da5c40d94ab/hf_xet-1.5.1-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:6abd35c3221eff63836618ddfb954dcf84798603f71d8e33e3ed7b04acfdbe6e", size = 4077794, upload-time = "2026-06-08T23:02:40.656Z" }, - { url = "https://files.pythonhosted.org/packages/35/94/4b2ecfbad8f8b04701a23aefb62f540b9137d058b7e1dbef16a32676f0e9/hf_xet-1.5.1-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:94e761bbd266bf4c03cee73753916062665ce8365aa40ed321f45afcb934b41e", size = 3845354, upload-time = "2026-06-08T23:02:42.702Z" }, - { url = "https://files.pythonhosted.org/packages/de/cc/f99f4bc7295023d7bd9ebbfd51f75cc530ca262c1227666268b8208f4b77/hf_xet-1.5.1-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:892e3a3a3aecc12aded8b93cf4f9cd059282c7de0732f7d55026f3abdf474350", size = 4514864, upload-time = "2026-06-08T23:02:44.497Z" }, - { url = "https://files.pythonhosted.org/packages/cd/6e/21f7e5a2381278bd3b7b7a5a4d90038518bb6308a0c1daf5d9f8268bb178/hf_xet-1.5.1-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:a93df2039190502835b1db8cd7e178b0b7b889fe9ab51299d5ced26e0dd879a4", size = 4303784, upload-time = "2026-06-08T23:02:46.203Z" }, - { url = "https://files.pythonhosted.org/packages/35/0e/f992bb6927ac1cb30ef74e62268f551f338bc32b2191f7c96a44c6f7283e/hf_xet-1.5.1-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0c97106032ef70467b4f6bc2d0ccc266d7613ee076afc56516c502f87ce1c4a6", size = 4500703, upload-time = "2026-06-08T23:02:47.628Z" }, - { url = "https://files.pythonhosted.org/packages/fb/d1/90a498d05447980b977b1669246eeeeae4cfb0ea3e7a286eaba627f91bf9/hf_xet-1.5.1-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6208adb15d192b90e4c2ad2a27ed864359b2cb0f2494eb6d7c7f3699ac02e2bf", size = 4719498, upload-time = "2026-06-08T23:02:49.268Z" }, - { url = "https://files.pythonhosted.org/packages/6d/b6/20f99cfe97cc663a711f7b33cc21d4793e51968e9a26125b4afcd77315ba/hf_xet-1.5.1-cp37-abi3-win_amd64.whl", hash = "sha256:f7b3002f95d1c13e24bcb4537baa8f0eb3838957067c91bb4959bc004a6435f5", size = 4026419, upload-time = "2026-06-08T23:02:50.829Z" }, - { url = "https://files.pythonhosted.org/packages/f9/fa/77453694888f03e5a8c8852d1514a0894d8e81c622d39edbaf308ea0dcf4/hf_xet-1.5.1-cp37-abi3-win_arm64.whl", hash = "sha256:93d090b57b211133f6c0dab0205ef5cb6d89162979ba75a74845045cc3063b8e", size = 3855178, upload-time = "2026-06-08T23:02:52.452Z" }, + { url = "https://files.pythonhosted.org/packages/a2/50/7afa2c9c787405864fc47a0d1bbc02c62e9101947ed43c1f43899fc7d91d/hf_xet-1.6.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:633dc0cd71d32da58ab8c03ad38e2fac452c15c2b0a2866ebf6ededfe0a5061d", size = 4071729, upload-time = "2026-08-03T22:33:00.721Z" }, + { url = "https://files.pythonhosted.org/packages/4b/69/55b8dcf636142ae660fec1869fcac14c4da2e8412e14d6eee1523be77e9f/hf_xet-1.6.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:f0906082d9932ae0c0057fa194041c22b4e2cdb46b2592ef3b91f020d62a081a", size = 3876287, upload-time = "2026-08-03T22:33:02.251Z" }, + { url = "https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d62671bb130879cef0ee4c9ebe47a14af6c66ec53e6d84dc15936e5ffdfac82f", size = 4464663, upload-time = "2026-08-03T22:33:03.802Z" }, + { url = "https://files.pythonhosted.org/packages/9a/69/1f0cbc2fb22ae6082d094f743d1b8945a3f36f6089cb95f42b7ee348cda7/hf_xet-1.6.0-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0e6e21fa3cdfcdcd76748564bf593870a5e013f47d97cf10aed63aa222cff5b7", size = 4262538, upload-time = "2026-08-03T22:33:05.287Z" }, + { url = "https://files.pythonhosted.org/packages/d1/3a/4f4f2301ade26e404462d3336fa11f7958d914cabbabdd6e03c3c5d5658c/hf_xet-1.6.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:4fc74352a17015bd0ee90038bc9efe38db894cde45f268b6712b04fce8cd0acb", size = 4460520, upload-time = "2026-08-03T22:33:06.81Z" }, + { url = "https://files.pythonhosted.org/packages/ab/5f/311725e2a905534dfee2dcb5b08414f249147f1f12252bfc2bd24caa075c/hf_xet-1.6.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8fb4f71cba6129110c3374a33f919001ff130488fc23553698e34cc1c2a1198c", size = 4675937, upload-time = "2026-08-03T22:33:08.616Z" }, + { url = "https://files.pythonhosted.org/packages/98/b7/8c59a66d15205024662f1d66968136f13893f96df1ddc5087e2e281fc95f/hf_xet-1.6.0-cp38-abi3-win_amd64.whl", hash = "sha256:fb4fadde1b2b70bf4c0c14a6dccbe7194b1c28947fefd5bbe3fed9d940676c3b", size = 4033128, upload-time = "2026-08-03T22:33:10.171Z" }, + { url = "https://files.pythonhosted.org/packages/73/63/ca511b6f802f28cf3489b280fe77475bcca8de85e81a6299d7916b5b5555/hf_xet-1.6.0-cp38-abi3-win_arm64.whl", hash = "sha256:3dc3e35441ba395006af5aaacc40ef2e603c51ef46c3530b9156185f00935ea3", size = 3859359, upload-time = "2026-08-03T22:33:11.725Z" }, ] [[package]] @@ -1777,7 +1819,7 @@ wheels = [ [[package]] name = "huggingface-hub" -version = "1.20.1" +version = "1.27.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -1788,12 +1830,11 @@ dependencies = [ { name = "packaging" }, { name = "pyyaml" }, { name = "tqdm" }, - { name = "typer" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e6/7e/fad82ad491b226e832d2da90a1a59f36acd4526cda8c726f639834754aa4/huggingface_hub-1.20.1.tar.gz", hash = "sha256:9f6d63bfbeab2d2a8357200a9bc4f18cd2c8bfac9579f792f5922e77bf6471d0", size = 859910, upload-time = "2026-06-18T22:06:53.348Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/9b/ddf3d02a8681f1b9ce52fda03d755dad6b74c4f8172304c4c8d2975450f9/huggingface_hub-1.27.0.tar.gz", hash = "sha256:c1fed40ea82a6b41b477f5243546549b792ae0a93abcea608cff66089bf8f8df", size = 942668, upload-time = "2026-08-07T12:48:05.161Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8e/b5/ff8516e74b459da3dce9567540c39f2d305ee7a2655109f6802873ff1588/huggingface_hub-1.20.1-py3-none-any.whl", hash = "sha256:274448a45c1ba6f112fe2fb168ead05574c654faa156904157a84085cfae14bd", size = 719837, upload-time = "2026-06-18T22:06:51.486Z" }, + { url = "https://files.pythonhosted.org/packages/de/d8/95b735e183957c1f26d94c52977f09d466d55119cbbc1558ea4975e4c216/huggingface_hub-1.27.0-py3-none-any.whl", hash = "sha256:7df6827c2f956c60fbaa64646e979e566db76f619dd0a9729dfb8c5a3eb4f68d", size = 784926, upload-time = "2026-08-07T12:48:02.905Z" }, ] [[package]] @@ -4823,8 +4864,8 @@ sdist = { url = "https://files.pythonhosted.org/packages/9d/f1/ad613157261ac7126 [[package]] name = "transformers" -version = "5.12.1" -source = { registry = "https://pypi.org/simple" } +version = "5.16.0.dev0" +source = { git = "https://github.com/huggingface/transformers.git?rev=d1123114da1ab4395198146f4f84dae7fe8b693e#d1123114da1ab4395198146f4f84dae7fe8b693e" } dependencies = [ { name = "huggingface-hub" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, @@ -4837,10 +4878,6 @@ dependencies = [ { name = "tqdm" }, { name = "typer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/aa/7c/8240f612819718100a9346dc28dea6a11370c3ca9c8c6eabadd3dea4ef29/transformers-5.12.1.tar.gz", hash = "sha256:679ee731c8225347889ad4fb3b2c926a62e9da3b7d284e9d12c791da7272466b", size = 8924054, upload-time = "2026-06-15T17:27:50.604Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/df/56/bbd60dd8668055803bf8ba55a81f9b8a8b31497f620109a9671d26a2076d/transformers-5.12.1-py3-none-any.whl", hash = "sha256:2a5e109d2021265df7098ffbb738295acaf5ad256f12cbc586db2ea4dcbb1a8a", size = 11150587, upload-time = "2026-06-15T17:27:46.679Z" }, -] [[package]] name = "triton" From d6e1145113a1f9ae4e306a85715ce88d9bfca987 Mon Sep 17 00:00:00 2001 From: Ang Li Date: Thu, 13 Aug 2026 23:00:00 +0000 Subject: [PATCH 2/6] Update logging --- integrations/minimax_h3/minimax_h3/model.py | 8 +++++--- integrations/minimax_h3/minimax_h3/pipeline.py | 3 ++- integrations/minimax_h3/minimax_h3/runner.py | 10 +++++++--- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/integrations/minimax_h3/minimax_h3/model.py b/integrations/minimax_h3/minimax_h3/model.py index 83f8ad7b9..aa6466c92 100644 --- a/integrations/minimax_h3/minimax_h3/model.py +++ b/integrations/minimax_h3/minimax_h3/model.py @@ -9,6 +9,7 @@ from typing import Any, cast import torch +from loguru import logger from torch import Tensor, nn from flashdreams.infra.diffusion.model import DiffusionModel, DiffusionModelConfig @@ -133,9 +134,10 @@ def generate_joint(self, state: MiniMaxH3DenoiseState) -> Tensor: for index, (video_timestep, audio_timestep) in enumerate( zip(video_timesteps, audio_timesteps, strict=True) ): - print( - f"MiniMax H3 denoise step {index + 1}/{len(video_timesteps)}", - flush=True, + logger.info( + "MiniMax H3 denoise step {}/{}", + index + 1, + len(video_timesteps), ) cache.timestep, cache.timestep_indices = self._row_timesteps( state, video_timestep, audio_timestep diff --git a/integrations/minimax_h3/minimax_h3/pipeline.py b/integrations/minimax_h3/minimax_h3/pipeline.py index ca7853895..862ebb3b5 100644 --- a/integrations/minimax_h3/minimax_h3/pipeline.py +++ b/integrations/minimax_h3/minimax_h3/pipeline.py @@ -29,6 +29,7 @@ import numpy as np import torch +from loguru import logger from torch import nn from flashdreams.infra.pipeline import ( @@ -491,7 +492,7 @@ def _apply_lora(self, transformer: Any, cache: MiniMaxH3PipelineCache) -> None: cache.lora_scale, cache.lora_weight_name, ) - print(f"Loaded LoRA {converted} at scale {cache.lora_scale:g}", flush=True) + logger.info("Loaded LoRA {} at scale {:g}", converted, cache.lora_scale) def _generate_low_ram(self, cache: MiniMaxH3PipelineCache) -> torch.Tensor: if cache.conditioning_checkpoint.is_file() and not cache.restart: diff --git a/integrations/minimax_h3/minimax_h3/runner.py b/integrations/minimax_h3/minimax_h3/runner.py index 4e8609409..d1809ebf6 100644 --- a/integrations/minimax_h3/minimax_h3/runner.py +++ b/integrations/minimax_h3/minimax_h3/runner.py @@ -169,8 +169,10 @@ def run(self) -> None: self.pipeline.mark_complete(cache) video_artifact = artifacts[0] logger.info( - f"[{config.runner_name}] wrote {tuple(frames.shape)} video to " - f"{Path(video_artifact.uri).resolve()}" + "[{}] wrote {} video to {}", + config.runner_name, + tuple(frames.shape), + Path(video_artifact.uri).resolve(), ) stats_history = video_artifact.metadata["stats_history"] if stats_history: @@ -179,7 +181,9 @@ def run(self) -> None: config.runner_name, list(stats_history), ) - logger.info(f"[{config.runner_name}] wrote stats to {stats_path.resolve()}") + logger.info( + "[{}] wrote stats to {}", config.runner_name, stats_path.resolve() + ) class MiniMaxH3T2VARunner(MiniMaxH3Runner): From 8192a5189442e9f247500b508b76c6e76db0a6f3 Mon Sep 17 00:00:00 2001 From: Ang Li Date: Fri, 14 Aug 2026 17:57:22 +0000 Subject: [PATCH 3/6] Optimize MiniMax H3 GPU handoff --- integrations/minimax_h3/minimax_h3/model.py | 33 ++-- .../minimax_h3/minimax_h3/pipeline.py | 82 +++++++++- integrations/minimax_h3/tests/test_smoke.py | 141 ++++++++++++++++++ 3 files changed, 232 insertions(+), 24 deletions(-) diff --git a/integrations/minimax_h3/minimax_h3/model.py b/integrations/minimax_h3/minimax_h3/model.py index aa6466c92..f8fe0a64a 100644 --- a/integrations/minimax_h3/minimax_h3/model.py +++ b/integrations/minimax_h3/minimax_h3/model.py @@ -88,23 +88,24 @@ def _row_timesteps( + state.audio_indices.numel() + state.text_indices.numel() ) - row_timesteps = torch.full( - (sequence_length,), - float(video_timestep), - dtype=torch.float32, - device=state.video_indices.device, + video_timestep = video_timestep.to( + device=state.video_indices.device, dtype=torch.float32 ) + audio_timestep = audio_timestep.to( + device=state.video_indices.device, dtype=torch.float32 + ) + row_timesteps = video_timestep.expand(sequence_length).clone() video_condition = state.video_indices[: state.num_condition_video_rows] audio_condition = state.audio_indices[: state.num_condition_audio_rows] audio_target = state.audio_indices[state.num_condition_audio_rows :] - row_timesteps[video_condition] = max(float(video_timestep), 0.999) + row_timesteps[video_condition] = video_timestep.clamp_min(0.999) row_timesteps[audio_target] = audio_timestep row_timesteps[audio_condition] = 1.0 return torch.unique(row_timesteps, sorted=True, return_inverse=True) @torch.no_grad() def generate_joint(self, state: MiniMaxH3DenoiseState) -> Tensor: - """Denoise both streams and return only unpacked video latents.""" + """Denoise both streams and return video latents on the execution device.""" device = self.transformer.device video = state.latents.to(device) audio = state.audio_latents.to(device) @@ -182,17 +183,13 @@ def generate_joint(self, state: MiniMaxH3DenoiseState) -> Tensor: patch_w, ) rows = rows.permute(0, 4, 1, 5, 2, 6, 3, 7) - return ( - rows.reshape( - -1, - channels, - state.num_latent_frames, - state.latent_height, - state.latent_width, - ) - .contiguous() - .cpu() - ) + return rows.reshape( + -1, + channels, + state.num_latent_frames, + state.latent_height, + state.latent_width, + ).contiguous() __all__ = [ diff --git a/integrations/minimax_h3/minimax_h3/pipeline.py b/integrations/minimax_h3/minimax_h3/pipeline.py index 862ebb3b5..7e0b1513a 100644 --- a/integrations/minimax_h3/minimax_h3/pipeline.py +++ b/integrations/minimax_h3/minimax_h3/pipeline.py @@ -21,7 +21,9 @@ import hashlib import json import os +import threading import time +from concurrent.futures import Future from dataclasses import dataclass, field, replace from datetime import datetime, timezone from pathlib import Path @@ -104,6 +106,12 @@ class MiniMaxH3PipelineCache: elapsed_seconds: float = 0.0 conditioning_seconds: float = 0.0 denoise_seconds: float = 0.0 + denoise_prepare_seconds: float = 0.0 + transformer_load_seconds: float = 0.0 + denoise_compute_seconds: float = 0.0 + denoise_cleanup_seconds: float = 0.0 + latent_checkpoint_seconds: float = 0.0 + latent_checkpoint_future: Future[float] | None = field(default=None, repr=False) decode_seconds: float = 0.0 peak_gpu_memory_gib: float = 0.0 attention_backend: str = "default" @@ -314,6 +322,42 @@ def _save_latents( _write_status(path, "denoised") +def _save_latents_async( + cache: MiniMaxH3PipelineCache, + model_id: str, + latents: torch.Tensor, +) -> Future[float]: + """Persist the recovery checkpoint without blocking video decoding.""" + future: Future[float] = Future() + + def persist() -> None: + if not future.set_running_or_notify_cancel(): + return + started = time.monotonic() + try: + _save_latents(cache, model_id, latents) + except BaseException as exc: + future.set_exception(exc) + else: + future.set_result(time.monotonic() - started) + + threading.Thread( + target=persist, + name="minimax-h3-latent-checkpoint", + daemon=False, + ).start() + return future + + +def _finish_latent_checkpoint(cache: MiniMaxH3PipelineCache, *, wait: bool) -> None: + """Collect a completed checkpoint write, optionally waiting for it.""" + future = cache.latent_checkpoint_future + if future is None or (not wait and not future.done()): + return + cache.latent_checkpoint_seconds = future.result() + cache.latent_checkpoint_future = None + + def _load_latents(cache: MiniMaxH3PipelineCache, model_id: str) -> torch.Tensor: from safetensors import safe_open from safetensors.torch import load_file @@ -439,27 +483,34 @@ def generate( cache.resumed_stage = "decode" latents = _load_latents(cache, self.config.model_id) else: - denoise_started = time.monotonic() if cache.low_ram: latents = self._generate_low_ram(cache) else: latents = self._generate_standard(cache) - cache.denoise_seconds = time.monotonic() - denoise_started - _save_latents(cache, self.config.model_id, latents) + cache.latent_checkpoint_future = _save_latents_async( + cache, self.config.model_id, latents + ) decode_started = time.monotonic() - frames = self._decode_video(cache, latents) + try: + frames = self._decode_video(cache, latents) + except BaseException: + try: + _finish_latent_checkpoint(cache, wait=True) + except BaseException: + logger.exception("Latent checkpoint also failed after decode failure") + raise cache.decode_seconds = time.monotonic() - decode_started cache.elapsed_seconds = time.monotonic() - started cache.peak_gpu_memory_gib = torch.cuda.max_memory_allocated() / 2**30 cache.generated = True - _write_status(cache.latent_checkpoint, "decoded-video") return frames def mark_complete(self, cache: MiniMaxH3PipelineCache) -> None: """Record completion only after the runtime output target closes.""" if not cache.generated or not cache.output_path.is_file(): raise RuntimeError("cannot complete H3 job before its MP4 is written") + _finish_latent_checkpoint(cache, wait=True) _write_status( cache.latent_checkpoint, "complete", output=str(cache.output_path) ) @@ -472,9 +523,15 @@ def finalize( """Return runtime metrics for the completed H3 rollout.""" if autoregressive_index != 0 or not cache.generated: raise ValueError("finalize requires the completed H3 runtime step") + _finish_latent_checkpoint(cache, wait=False) return { "conditioning_seconds": cache.conditioning_seconds, "denoise_seconds": cache.denoise_seconds, + "denoise_prepare_seconds": cache.denoise_prepare_seconds, + "transformer_load_seconds": cache.transformer_load_seconds, + "denoise_compute_seconds": cache.denoise_compute_seconds, + "denoise_cleanup_seconds": cache.denoise_cleanup_seconds, + "latent_checkpoint_seconds": cache.latent_checkpoint_seconds, "decode_seconds": cache.decode_seconds, "total_seconds": cache.elapsed_seconds, "peak_gpu_memory_gib": cache.peak_gpu_memory_gib, @@ -695,8 +752,11 @@ def _prepare_denoise_state( def _run_native_denoise( self, cache: MiniMaxH3PipelineCache, conditioned: dict[str, Any] ) -> torch.Tensor: + denoise_started = time.monotonic() _write_status(cache.latent_checkpoint, "denoising-native-flashdreams") + prepare_started = time.monotonic() state = self._prepare_denoise_state(cache, conditioned) + cache.denoise_prepare_seconds = time.monotonic() - prepare_started backend = "cudnn" if cache.attention == "default" else "flash" cache.attention_backend = backend base_transformer = cast( @@ -722,14 +782,24 @@ def _run_native_denoise( ), seed=cache.seed, ) + transformer_load_started = time.monotonic() model = model_config.setup() + cache.transformer_load_seconds = time.monotonic() - transformer_load_started try: self._apply_lora(model.transformer, cache) - return model.generate_joint(state) + compute_started = time.monotonic() + latents = model.generate_joint(state) + if latents.is_cuda: + torch.cuda.synchronize(latents.device) + cache.denoise_compute_seconds = time.monotonic() - compute_started finally: + cleanup_started = time.monotonic() del model gc.collect() torch.cuda.empty_cache() + cache.denoise_cleanup_seconds = time.monotonic() - cleanup_started + cache.denoise_seconds = time.monotonic() - denoise_started + return latents def _generate_standard(self, cache: MiniMaxH3PipelineCache) -> torch.Tensor: conditioning_started = time.monotonic() diff --git a/integrations/minimax_h3/tests/test_smoke.py b/integrations/minimax_h3/tests/test_smoke.py index a39111667..99c24b32f 100644 --- a/integrations/minimax_h3/tests/test_smoke.py +++ b/integrations/minimax_h3/tests/test_smoke.py @@ -18,6 +18,7 @@ from __future__ import annotations import sys +import threading from pathlib import Path from typing import Any, cast @@ -25,6 +26,7 @@ import tomli as tomllib import torch from minimax_h3 import config as config_mod +from minimax_h3 import pipeline as h3_pipeline from minimax_h3.config import ( PIPELINE_MINIMAX_H3_FL2VA, PIPELINE_MINIMAX_H3_REF2VA, @@ -36,6 +38,7 @@ ) from minimax_h3.constants import align_num_frames, validate_canvas from minimax_h3.lora import convert_musubi_lora +from minimax_h3.model import MiniMaxH3DenoiseState, MiniMaxH3DiffusionModel from minimax_h3.pipeline import MiniMaxH3Pipeline from minimax_h3.references import parse_reference_specs from minimax_h3.runner import ( @@ -216,6 +219,117 @@ def test_runtime_cache_uses_stage_specific_checkpoints(tmp_path: Path) -> None: assert cache.latent_checkpoint.name == "out.mp4.latents.safetensors" +def test_generate_preserves_non_overlapping_stage_metrics( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Keep conditioning separate from native denoise and checkpoint timings.""" + image = tmp_path / "image.png" + image.write_bytes(b"test") + pipeline = PIPELINE_MINIMAX_H3_FL2VA.setup() + cache = pipeline.initialize_cache( + prompt="animate", + image_path=image, + last_image_path=None, + references=(), + output_path=tmp_path / "out.mp4", + width=576, + height=768, + duration=5.0, + steps=30, + seed=42, + low_ram=True, + restart=True, + attention="auto", + lora=None, + lora_weight_name=None, + lora_scale=1.0, + ) + + def generate_low_ram(_: Any) -> torch.Tensor: + cache.conditioning_seconds = 10.0 + cache.denoise_seconds = 20.0 + cache.denoise_prepare_seconds = 1.0 + cache.transformer_load_seconds = 4.0 + cache.denoise_compute_seconds = 14.0 + cache.denoise_cleanup_seconds = 1.0 + return torch.zeros(1) + + monkeypatch.setattr(pipeline, "_generate_low_ram", generate_low_ram) + monkeypatch.setattr( + pipeline, + "_decode_video", + lambda _cache, _latents: torch.zeros(1, 3, 1, 1), + ) + monkeypatch.setattr(h3_pipeline, "_save_latents", lambda *_args: None) + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) + + pipeline.generate(0, cache) + metrics = pipeline.finalize(0, cache) + + assert metrics["conditioning_seconds"] == 10.0 + assert metrics["denoise_seconds"] == 20.0 + assert metrics["denoise_prepare_seconds"] == 1.0 + assert metrics["transformer_load_seconds"] == 4.0 + assert metrics["denoise_compute_seconds"] == 14.0 + assert metrics["denoise_cleanup_seconds"] == 1.0 + assert metrics["latent_checkpoint_seconds"] >= 0.0 + + +def test_generate_overlaps_latent_checkpoint_with_decode( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Do not put recovery-checkpoint I/O on generate's critical path.""" + image = tmp_path / "image.png" + image.write_bytes(b"test") + pipeline = PIPELINE_MINIMAX_H3_FL2VA.setup() + cache = pipeline.initialize_cache( + prompt="animate", + image_path=image, + last_image_path=None, + references=(), + output_path=tmp_path / "out.mp4", + width=576, + height=768, + duration=5.0, + steps=30, + seed=42, + low_ram=True, + restart=True, + attention="auto", + lora=None, + lora_weight_name=None, + lora_scale=1.0, + ) + checkpoint_started = threading.Event() + release_checkpoint = threading.Event() + + def save_latents(*_args: Any) -> None: + checkpoint_started.set() + assert release_checkpoint.wait(timeout=5) + + def decode_video(*_args: Any) -> torch.Tensor: + assert checkpoint_started.wait(timeout=5) + return torch.zeros(1, 3, 1, 1) + + monkeypatch.setattr(pipeline, "_generate_low_ram", lambda _: torch.zeros(1)) + monkeypatch.setattr(pipeline, "_decode_video", decode_video) + monkeypatch.setattr(h3_pipeline, "_save_latents", save_latents) + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) + + try: + pipeline.generate(0, cache) + assert cache.latent_checkpoint_future is not None + assert not cache.latent_checkpoint_future.done() + finally: + release_checkpoint.set() + + cache.output_path.write_bytes(b"mp4") + pipeline.mark_complete(cache) + assert cache.latent_checkpoint_future is None + + def test_reference_parser_preserves_order_and_enforces_limits(tmp_path: Path) -> None: """Keep semantic reference order while rejecting unsupported requests.""" image = tmp_path / "subject.png" @@ -288,6 +402,33 @@ def test_native_scheduler_matches_official_h3_euler() -> None: torch.testing.assert_close(actual, expected) +def test_row_timestep_plan_preserves_device_and_conditioning_levels() -> None: + """Build packed row timesteps without materializing accelerator scalars.""" + state = MiniMaxH3DenoiseState( + latents=torch.empty(0), + audio_latents=torch.empty(0), + prompt_embeds=torch.empty(0), + position_ids=torch.empty(0), + token_tags=torch.empty(0), + video_indices=torch.tensor([0, 1, 4]), + audio_indices=torch.tensor([2, 3]), + text_indices=torch.tensor([5]), + num_condition_video_rows=1, + num_condition_audio_rows=1, + num_latent_frames=0, + latent_height=0, + latent_width=0, + ) + + timesteps, indices = MiniMaxH3DiffusionModel._row_timesteps( + state, torch.tensor(0.5), torch.tensor(0.25) + ) + + assert timesteps.device == state.video_indices.device + torch.testing.assert_close(timesteps, torch.tensor([0.25, 0.5, 0.999, 1.0])) + torch.testing.assert_close(indices, torch.tensor([2, 1, 3, 0, 1, 1])) + + def test_native_transformer_matches_official_h3_forward() -> None: """Prove state-dict and numerical compatibility on a tiny CPU model.""" from diffusers.models.transformers.transformer_minimax_h3 import ( From 3893e460d519e7ed62460f7de8afe04c8248cfd6 Mon Sep 17 00:00:00 2001 From: Ang Li Date: Fri, 14 Aug 2026 17:57:22 +0000 Subject: [PATCH 4/6] Optimize MiniMax H3 GPU handoff --- .../minimax_h3/minimax_h3/pipeline.py | 3 +-- integrations/minimax_h3/tests/test_smoke.py | 19 ++++++++++++++----- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/integrations/minimax_h3/minimax_h3/pipeline.py b/integrations/minimax_h3/minimax_h3/pipeline.py index 7e0b1513a..eb4d952b3 100644 --- a/integrations/minimax_h3/minimax_h3/pipeline.py +++ b/integrations/minimax_h3/minimax_h3/pipeline.py @@ -319,7 +319,6 @@ def _save_latents( }, ) os.replace(temporary, path) - _write_status(path, "denoised") def _save_latents_async( @@ -523,7 +522,7 @@ def finalize( """Return runtime metrics for the completed H3 rollout.""" if autoregressive_index != 0 or not cache.generated: raise ValueError("finalize requires the completed H3 runtime step") - _finish_latent_checkpoint(cache, wait=False) + _finish_latent_checkpoint(cache, wait=True) return { "conditioning_seconds": cache.conditioning_seconds, "denoise_seconds": cache.denoise_seconds, diff --git a/integrations/minimax_h3/tests/test_smoke.py b/integrations/minimax_h3/tests/test_smoke.py index 99c24b32f..9b827b7f5 100644 --- a/integrations/minimax_h3/tests/test_smoke.py +++ b/integrations/minimax_h3/tests/test_smoke.py @@ -267,6 +267,7 @@ def generate_low_ram(_: Any) -> torch.Tensor: pipeline.generate(0, cache) metrics = pipeline.finalize(0, cache) + assert cache.latent_checkpoint_future is None assert metrics["conditioning_seconds"] == 10.0 assert metrics["denoise_seconds"] == 20.0 assert metrics["denoise_prepare_seconds"] == 1.0 @@ -325,9 +326,10 @@ def decode_video(*_args: Any) -> torch.Tensor: finally: release_checkpoint.set() + pipeline.finalize(0, cache) + assert cache.latent_checkpoint_future is None cache.output_path.write_bytes(b"mp4") pipeline.mark_complete(cache) - assert cache.latent_checkpoint_future is None def test_reference_parser_preserves_order_and_enforces_limits(tmp_path: Path) -> None: @@ -402,7 +404,9 @@ def test_native_scheduler_matches_official_h3_euler() -> None: torch.testing.assert_close(actual, expected) -def test_row_timestep_plan_preserves_device_and_conditioning_levels() -> None: +def test_row_timestep_plan_preserves_device_and_conditioning_levels( + monkeypatch: pytest.MonkeyPatch, +) -> None: """Build packed row timesteps without materializing accelerator scalars.""" state = MiniMaxH3DenoiseState( latents=torch.empty(0), @@ -420,9 +424,14 @@ def test_row_timestep_plan_preserves_device_and_conditioning_levels() -> None: latent_width=0, ) - timesteps, indices = MiniMaxH3DiffusionModel._row_timesteps( - state, torch.tensor(0.5), torch.tensor(0.25) - ) + def reject_scalar_conversion(_tensor: torch.Tensor) -> float: + raise AssertionError("row timestep construction converted a tensor to float") + + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "__float__", reject_scalar_conversion) + timesteps, indices = MiniMaxH3DiffusionModel._row_timesteps( + state, torch.tensor(0.5), torch.tensor(0.25) + ) assert timesteps.device == state.video_indices.device torch.testing.assert_close(timesteps, torch.tensor([0.25, 0.5, 0.999, 1.0])) From 6bf7f0455b21d24f3c8f7a2dca8ff305dc1b0650 Mon Sep 17 00:00:00 2001 From: Ang Li Date: Wed, 19 Aug 2026 17:11:08 +0000 Subject: [PATCH 5/6] Use flashdreams RoPE --- flashdreams/tests/test_rope_kernel.py | 41 ++++++++++- .../minimax_h3/minimax_h3/transformer.py | 35 ++++----- integrations/minimax_h3/tests/test_smoke.py | 58 +-------------- .../minimax_h3/tests/test_transformer_cuda.py | 72 +++++++++++++++++++ 4 files changed, 127 insertions(+), 79 deletions(-) create mode 100644 integrations/minimax_h3/tests/test_transformer_cuda.py diff --git a/flashdreams/tests/test_rope_kernel.py b/flashdreams/tests/test_rope_kernel.py index fa5f8f71e..420a1d8fe 100644 --- a/flashdreams/tests/test_rope_kernel.py +++ b/flashdreams/tests/test_rope_kernel.py @@ -33,10 +33,9 @@ import pytest import torch -from torch import Tensor - from flashdreams.core.attention.rope import apply_rope_freqs from flashdreams.core.attention.rope_kernel import apply_rotary_pos_emb +from torch import Tensor def _load_te_apply_rope() -> Callable[..., Tensor] | None: @@ -191,6 +190,44 @@ def test_zero_freqs_is_identity(cuda_device): torch.testing.assert_close(out, x) +@pytest.mark.parametrize("x_dtype", _DTYPES) +def test_partial_prefix_view_matches_torch( + cuda_device: torch.device, + x_dtype: torch.dtype, +) -> None: + """A sliced head prefix rotates in place without touching trailing channels.""" + B, S, H, D = 2, 67, 7, 128 + rotary_dim = 96 + x = torch.randn( + B, + S, + H, + D, + device=cuda_device, + dtype=x_dtype, + ) + frequencies = _expanded_freqs( + S, rotary_dim, interleaved=False, device=cuda_device, seed=7 + ) + cos = frequencies[:, 0, 0].cos().to(x_dtype)[None, :, None, :] + sin = frequencies[:, 0, 0].sin().to(x_dtype)[None, :, None, :] + + rotary = x[..., :rotary_dim] + first, second = rotary.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + expected = x.clone() + expected[..., :rotary_dim] = rotary * cos + rotated * sin + + actual = x.clone() + prefix = actual[..., :rotary_dim] + returned = apply_rope_freqs(prefix, frequencies) + + atol, rtol = _parity_tol(x_dtype) + assert returned.data_ptr() == prefix.data_ptr() + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + assert torch.equal(actual[..., rotary_dim:], x[..., rotary_dim:]) + + @_requires_te def test_non_contiguous_x(cuda_device): """The kernel respects arbitrary strides on the B / S / H axes.""" diff --git a/integrations/minimax_h3/minimax_h3/transformer.py b/integrations/minimax_h3/minimax_h3/transformer.py index befa71a53..2de92f8e2 100644 --- a/integrations/minimax_h3/minimax_h3/transformer.py +++ b/integrations/minimax_h3/minimax_h3/transformer.py @@ -10,15 +10,14 @@ from typing import Any, Literal, cast import torch -from torch import Tensor, nn - -from flashdreams.core.attention import NativeAttention +from flashdreams.core.attention import NativeAttention, apply_rope_freqs from flashdreams.core.checkpoint.load import load_checkpoint from flashdreams.infra.diffusion.transformer import ( Transformer, TransformerAutoregressiveCache, TransformerConfig, ) +from torch import Tensor, nn H3_TRANSFORMER_CHECKPOINT = ( "https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/" @@ -145,25 +144,21 @@ def __init__( ) self.register_buffer("inv_freq", inv_freq, persistent=False) - def forward(self, position_ids: Tensor) -> tuple[Tensor, Tensor]: + def forward(self, position_ids: Tensor) -> Tensor: inv_freq = cast(Tensor, self.inv_freq) frequencies = position_ids.float().unsqueeze(-1) * inv_freq.view(1, 1, -1) frequencies = torch.cat(frequencies.unbind(dim=1), dim=-1) - frequencies = torch.cat((frequencies, frequencies), dim=-1) - return frequencies.cos(), frequencies.sin() + return torch.cat((frequencies, frequencies), dim=-1) -def _apply_rotary(hidden_states: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - rotary_dim = cos.shape[-1] - rotary, passthrough = ( +def _apply_rotary(hidden_states: Tensor, frequencies: Tensor) -> Tensor: + """Apply H3's partial RoPE in place with FlashDreams' CUDA kernel.""" + rotary_dim = frequencies.shape[-1] + apply_rope_freqs( hidden_states[..., :rotary_dim], - hidden_states[..., rotary_dim:], + frequencies[:, None, None, :], ) - cos = cos.to(hidden_states.dtype)[None, :, None, :] - sin = sin.to(hidden_states.dtype)[None, :, None, :] - first, second = rotary.chunk(2, dim=-1) - rotated = torch.cat((-second, first), dim=-1) - return torch.cat((rotary * cos + rotated * sin, passthrough), dim=-1).contiguous() + return hidden_states class _Attention(nn.Module): @@ -189,15 +184,13 @@ def __init__( ) self.attn_op = NativeAttention(qkv_format="bshd", backend=backend) - def forward( - self, hidden_states: Tensor, rotary: tuple[Tensor, Tensor] | None = None - ) -> Tensor: + def forward(self, hidden_states: Tensor, rotary: Tensor | None = None) -> Tensor: query = self.norm_q(self.to_q(hidden_states).unflatten(-1, (self.heads, -1))) key = self.norm_k(self.to_k(hidden_states).unflatten(-1, (self.heads, -1))) value = self.to_v(hidden_states).unflatten(-1, (self.heads, -1)) if rotary is not None: - query = _apply_rotary(query, *rotary) - key = _apply_rotary(key, *rotary) + query = _apply_rotary(query, rotary) + key = _apply_rotary(key, rotary) output = self.attn_op(query, key, value).flatten(2, 3).type_as(query) return self.to_out[1](self.to_out[0](output)) @@ -274,7 +267,7 @@ def forward( hidden_states: Tensor, temb: Tensor, adaln_indices: Tensor, - rotary: tuple[Tensor, Tensor], + rotary: Tensor, ) -> Tensor: shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = self.adaln_proj(temb) normalized = self.norm1(hidden_states) diff --git a/integrations/minimax_h3/tests/test_smoke.py b/integrations/minimax_h3/tests/test_smoke.py index 9b827b7f5..31b7d288e 100644 --- a/integrations/minimax_h3/tests/test_smoke.py +++ b/integrations/minimax_h3/tests/test_smoke.py @@ -25,6 +25,8 @@ import pytest import tomli as tomllib import torch +from flashdreams.infra.diffusion.transformer import Transformer +from flashdreams.infra.runner import RunnerConfig from minimax_h3 import config as config_mod from minimax_h3 import pipeline as h3_pipeline from minimax_h3.config import ( @@ -50,9 +52,6 @@ from minimax_h3.scheduler import MiniMaxH3SchedulerConfig from minimax_h3.transformer import MiniMaxH3TransformerConfig -from flashdreams.infra.diffusion.transformer import Transformer -from flashdreams.infra.runner import RunnerConfig - pytestmark = pytest.mark.ci_cpu ENTRY_POINT_GROUP = "flashdreams.runner_configs" @@ -436,56 +435,3 @@ def reject_scalar_conversion(_tensor: torch.Tensor) -> float: assert timesteps.device == state.video_indices.device torch.testing.assert_close(timesteps, torch.tensor([0.25, 0.5, 0.999, 1.0])) torch.testing.assert_close(indices, torch.tensor([2, 1, 3, 0, 1, 1])) - - -def test_native_transformer_matches_official_h3_forward() -> None: - """Prove state-dict and numerical compatibility on a tiny CPU model.""" - from diffusers.models.transformers.transformer_minimax_h3 import ( - MiniMaxH3Transformer3DModel, - ) - - architecture: dict[str, Any] = { - "num_attention_heads": 2, - "attention_head_dim": 12, - "hidden_size": 16, - "num_layers": 2, - "num_refiner_layers": 1, - "ffn_dim": 32, - "in_channels": 2, - "audio_in_channels": 4, - "patch_size": (1, 1, 1), - "text_dim": 10, - "freq_dim": 8, - "time_embed_hidden_dim": 16, - "time_embed_dim": 8, - "rope_freq_dim": 2, - } - torch.manual_seed(1) - official: Any = MiniMaxH3Transformer3DModel(**architecture) - native = MiniMaxH3TransformerConfig( - checkpoint_path=None, - device="cpu", - execution_device="cpu", - sequential_cpu_offload=False, - dtype=torch.float32, - attention_backend="math", - **architecture, - ).setup() - native.load_state_dict(official.state_dict(), strict=True) - inputs = { - "hidden_states": torch.randn(1, 4, 2), - "audio_hidden_states": torch.randn(1, 2, 4), - "encoder_hidden_states": torch.randn(1, 3, 10), - "timestep": torch.tensor([0.1, 0.5, 0.999]), - "timestep_indices": torch.tensor([1, 1, 1, 2, 0, 1, 1, 0, 2]), - "token_tags": torch.tensor([1, 1, 1, 0, 0, 2, 2, 0, 0]), - "position_ids": torch.randn(9, 3), - "video_indices": torch.tensor([3, 4, 7, 8]), - "audio_indices": torch.tensor([5, 6]), - "text_indices": torch.tensor([0, 1, 2]), - } - with torch.no_grad(): - expected = official(**inputs, return_dict=False) - actual = native.forward_joint(**inputs) - for native_output, official_output in zip(actual, expected, strict=True): - torch.testing.assert_close(native_output, official_output) diff --git a/integrations/minimax_h3/tests/test_transformer_cuda.py b/integrations/minimax_h3/tests/test_transformer_cuda.py new file mode 100644 index 000000000..e7285fe79 --- /dev/null +++ b/integrations/minimax_h3/tests/test_transformer_cuda.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CUDA parity tests for the MiniMax H3 transformer.""" + +from __future__ import annotations + +from typing import Any + +import pytest +import torch +from diffusers.models.transformers.transformer_minimax_h3 import ( + MiniMaxH3Transformer3DModel, +) +from minimax_h3.transformer import MiniMaxH3TransformerConfig + +pytestmark = pytest.mark.ci_gpu + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_native_transformer_matches_official_h3_forward() -> None: + """Prove state-dict and numerical compatibility on a tiny CUDA model.""" + architecture: dict[str, Any] = { + "num_attention_heads": 2, + "attention_head_dim": 16, + "hidden_size": 16, + "num_layers": 2, + "num_refiner_layers": 1, + "ffn_dim": 32, + "in_channels": 2, + "audio_in_channels": 4, + "patch_size": (1, 1, 1), + "text_dim": 10, + "freq_dim": 8, + "time_embed_hidden_dim": 16, + "time_embed_dim": 8, + "rope_freq_dim": 2, + } + device = torch.device("cuda") + torch.manual_seed(1) + official: Any = MiniMaxH3Transformer3DModel(**architecture) + official.to(device) + native = MiniMaxH3TransformerConfig( + checkpoint_path=None, + device="cuda", + execution_device="cuda", + sequential_cpu_offload=False, + dtype=torch.float32, + attention_backend="math", + **architecture, + ).setup() + native.load_state_dict(official.state_dict(), strict=True) + inputs = { + "hidden_states": torch.randn(1, 4, 2, device=device), + "audio_hidden_states": torch.randn(1, 2, 4, device=device), + "encoder_hidden_states": torch.randn(1, 3, 10, device=device), + "timestep": torch.tensor([0.1, 0.5, 0.999], device=device), + "timestep_indices": torch.tensor([1, 1, 1, 2, 0, 1, 1, 0, 2], device=device), + "token_tags": torch.tensor([1, 1, 1, 0, 0, 2, 2, 0, 0], device=device), + "position_ids": torch.randn(9, 3, device=device), + "video_indices": torch.tensor([3, 4, 7, 8], device=device), + "audio_indices": torch.tensor([5, 6], device=device), + "text_indices": torch.tensor([0, 1, 2], device=device), + } + with ( + torch.no_grad(), + torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH), + ): + expected = official(**inputs, return_dict=False) + actual = native.forward_joint(**inputs) + for native_output, official_output in zip(actual, expected, strict=True): + torch.testing.assert_close(native_output, official_output) From 2213f4066bd7606e89290689f442de913cbdd2eb Mon Sep 17 00:00:00 2001 From: Ang Li Date: Wed, 9 Sep 2026 03:41:15 +0000 Subject: [PATCH 6/6] Migrate to v2 API --- LICENSE | 5 + NOTICE | 4 + REUSE.toml | 15 + THIRD-PARTY-NOTICES | 27 +- apps/t2v/t2v/application.py | 16 + apps/t2v/t2v/session.py | 22 +- apps/t2v/tests/test_session.py | 27 + .../multi_head_attention/__init__.py | 5 +- .../multi_head_attention/optimized.py | 166 ++-- .../accelerated/multi_head_attention/sdpa.py | 114 +++ .../accelerated/multi_head_attention/torch.py | 63 +- .../flashdreams/core/checkpoint/load.py | 65 +- .../infra/acceleration/__init__.py | 2 + .../infra/acceleration/encoder_lifecycle.py | 8 +- .../infra/diffusion/scheduler/__init__.py | 12 + .../diffusion/scheduler/data_flow_euler.py | 102 +++ .../infra/diffusion/scheduler/synchronized.py | 74 ++ .../infra/encoder/text/qwen3_vl.py | 108 +++ .../test_mha_optimized.py | 44 +- .../multi_head_attention/test_mha_torch.py | 84 +- flashdreams/tests/test_checkpoint_loading.py | 86 ++ .../tests/test_synchronized_sampling.py | 75 ++ .../minimax_h3/minimax_h3/__init__.py | 20 - integrations/minimax_h3/minimax_h3/config.py | 67 -- integrations/minimax_h3/minimax_h3/model.py | 199 ---- .../minimax_h3/minimax_h3/pipeline.py | 848 ------------------ .../minimax_h3/minimax_h3/references.py | 106 --- integrations/minimax_h3/minimax_h3/runner.py | 229 ----- .../minimax_h3/minimax_h3/scheduler.py | 112 --- integrations/minimax_h3/pyproject.toml | 55 -- integrations/minimax_h3/tests/test_smoke.py | 437 --------- .../minimax_h3/tests/test_transformer_cuda.py | 72 -- integrations_v2/README.md | 2 + integrations_v2/minimax_h3/README.md | 80 ++ integrations_v2/minimax_h3/__init__.py | 4 + integrations_v2/minimax_h3/apps/__init__.py | 4 + integrations_v2/minimax_h3/apps/t2v/README.md | 5 + .../minimax_h3/apps/t2v/__init__.py | 4 + .../minimax_h3/apps/t2v/adapter.py | 159 ++++ integrations_v2/minimax_h3/config.py | 26 + integrations_v2/minimax_h3/impl/__init__.py | 4 + .../minimax_h3/impl/audio_encoder.py | 296 ++++++ .../minimax_h3/impl/conditioning.py | 490 ++++++++++ .../minimax_h3/impl}/constants.py | 2 + integrations_v2/minimax_h3/impl/layout.py | 523 +++++++++++ .../minimax_h3/impl}/lora.py | 0 integrations_v2/minimax_h3/impl/pipeline.py | 345 +++++++ integrations_v2/minimax_h3/impl/references.py | 156 ++++ .../minimax_h3/impl}/transformer.py | 311 +++---- integrations_v2/minimax_h3/impl/video_vae.py | 819 +++++++++++++++++ integrations_v2/minimax_h3/impl/weights.py | 88 ++ integrations_v2/minimax_h3/pyproject.toml | 40 + .../minimax_h3/tests/test_codecs.py | 295 ++++++ .../minimax_h3/tests/test_conditioning.py | 444 +++++++++ integrations_v2/minimax_h3/tests/test_lora.py | 50 ++ .../minimax_h3/tests/test_pipeline.py | 214 +++++ .../minimax_h3/tests/test_transformer.py | 116 +++ pyproject.toml | 2 - uv.lock | 16 +- 59 files changed, 5331 insertions(+), 2433 deletions(-) create mode 100644 flashdreams/flashdreams/accelerated/multi_head_attention/sdpa.py create mode 100644 flashdreams/flashdreams/infra/diffusion/scheduler/data_flow_euler.py create mode 100644 flashdreams/flashdreams/infra/diffusion/scheduler/synchronized.py create mode 100644 flashdreams/flashdreams/infra/encoder/text/qwen3_vl.py create mode 100644 flashdreams/tests/test_synchronized_sampling.py delete mode 100644 integrations/minimax_h3/minimax_h3/__init__.py delete mode 100644 integrations/minimax_h3/minimax_h3/config.py delete mode 100644 integrations/minimax_h3/minimax_h3/model.py delete mode 100644 integrations/minimax_h3/minimax_h3/pipeline.py delete mode 100644 integrations/minimax_h3/minimax_h3/references.py delete mode 100644 integrations/minimax_h3/minimax_h3/runner.py delete mode 100644 integrations/minimax_h3/minimax_h3/scheduler.py delete mode 100644 integrations/minimax_h3/pyproject.toml delete mode 100644 integrations/minimax_h3/tests/test_smoke.py delete mode 100644 integrations/minimax_h3/tests/test_transformer_cuda.py create mode 100644 integrations_v2/minimax_h3/README.md create mode 100644 integrations_v2/minimax_h3/__init__.py create mode 100644 integrations_v2/minimax_h3/apps/__init__.py create mode 100644 integrations_v2/minimax_h3/apps/t2v/README.md create mode 100644 integrations_v2/minimax_h3/apps/t2v/__init__.py create mode 100644 integrations_v2/minimax_h3/apps/t2v/adapter.py create mode 100644 integrations_v2/minimax_h3/config.py create mode 100644 integrations_v2/minimax_h3/impl/__init__.py create mode 100644 integrations_v2/minimax_h3/impl/audio_encoder.py create mode 100644 integrations_v2/minimax_h3/impl/conditioning.py rename {integrations/minimax_h3/minimax_h3 => integrations_v2/minimax_h3/impl}/constants.py (93%) create mode 100644 integrations_v2/minimax_h3/impl/layout.py rename {integrations/minimax_h3/minimax_h3 => integrations_v2/minimax_h3/impl}/lora.py (100%) create mode 100644 integrations_v2/minimax_h3/impl/pipeline.py create mode 100644 integrations_v2/minimax_h3/impl/references.py rename {integrations/minimax_h3/minimax_h3 => integrations_v2/minimax_h3/impl}/transformer.py (63%) create mode 100644 integrations_v2/minimax_h3/impl/video_vae.py create mode 100644 integrations_v2/minimax_h3/impl/weights.py create mode 100644 integrations_v2/minimax_h3/pyproject.toml create mode 100644 integrations_v2/minimax_h3/tests/test_codecs.py create mode 100644 integrations_v2/minimax_h3/tests/test_conditioning.py create mode 100644 integrations_v2/minimax_h3/tests/test_lora.py create mode 100644 integrations_v2/minimax_h3/tests/test_pipeline.py create mode 100644 integrations_v2/minimax_h3/tests/test_transformer.py diff --git a/LICENSE b/LICENSE index 1e0105f47..22556846a 100644 --- a/LICENSE +++ b/LICENSE @@ -19,6 +19,11 @@ FlashDreams is open-source software. Its licensing posture is: is licensed under the Zlib license. Full text: LICENSES/Zlib.txt. + * MiniMax H3 codec and conditioning ports under + integrations_v2/minimax_h3/impl/ retain their upstream Apache-2.0 + attribution. Full text: LICENSES/Apache-2.0.txt; source details + are recorded in THIRD-PARTY-NOTICES. + Each first-party source file carries an inline SPDX license identifier; the REUSE 3.3 manifest at REUSE.toml fills the gaps for files (configuration, assets, generated outputs) that cannot diff --git a/NOTICE b/NOTICE index 8903fdebb..7c4cc9112 100644 --- a/NOTICE +++ b/NOTICE @@ -15,6 +15,10 @@ LICENSES/: cudaraster/framework/3rdparty/lodepng/{lodepng.h,lodepng.cpp} Zlib (see LICENSES/Zlib.txt) +The MiniMax/HuggingFace H3 codec and conditioning ports under +integrations_v2/minimax_h3/impl/ retain their upstream Apache-2.0 +attribution (see LICENSES/Apache-2.0.txt). + Third-party software attributions, source-level redistribution disclosures, and the full per-dependency license inventory are documented in THIRD-PARTY-NOTICES at the repository root. diff --git a/REUSE.toml b/REUSE.toml index 0c8a7f83d..ca3798b9d 100644 --- a/REUSE.toml +++ b/REUSE.toml @@ -66,6 +66,21 @@ SPDX-FileCopyrightText = [ ] SPDX-License-Identifier = "Apache-2.0" +# --- MiniMax/HuggingFace inference mathematics adapted to native APIs. --- +[[annotations]] +path = [ + "integrations_v2/minimax_h3/impl/video_vae.py", + "integrations_v2/minimax_h3/impl/audio_encoder.py", + "integrations_v2/minimax_h3/impl/conditioning.py", + "integrations_v2/minimax_h3/impl/layout.py", +] +precedence = "aggregate" +SPDX-FileCopyrightText = [ + "Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.", + "Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.", +] +SPDX-License-Identifier = "Apache-2.0" + # --- Generated protobuf stubs. The .proto sources are authoritative; the # *_pb2.py / *_pb2.pyi / *_pb2_grpc.py outputs inherit the same license # because the regeneration script (compile_protos.sh) is project-owned. --- diff --git a/THIRD-PARTY-NOTICES b/THIRD-PARTY-NOTICES index 641e9ebe3..d0de3dbc5 100644 --- a/THIRD-PARTY-NOTICES +++ b/THIRD-PARTY-NOTICES @@ -147,6 +147,30 @@ The following third-party source is physically present in this repository (not just consumed at runtime). Each file retains its original copyright notice inline; license texts are reproduced under LICENSES/. +-------------------------------------------------------------------------------- +MiniMax H3 native codecs and conditioning +-------------------------------------------------------------------------------- + +Path: integrations_v2/minimax_h3/impl/{video_vae,audio_encoder,conditioning,layout}.py +License: Apache-2.0 (see LICENSES/Apache-2.0.txt) + Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved. +Upstream: https://github.com/huggingface/diffusers/tree/175fe6b2419a01db9c2ceabd01ec37d2c0305fc2/src/diffusers + +The inference mathematics are adapted from the MiniMax H3 implementation. +FlashDreams replaces pipeline blocks, model mixins, attention processors, +checkpoint loading, and runtime orchestration with its native APIs. The +audio port contains only reference encoding, not generated-audio decoding. +Diffusers is not a runtime dependency of this integration. Model weights +are not redistributed. Native checkpoint assets are pinned to +MiniMaxAI/MiniMax-H3 revision 42ed227ee7df40d41602854ae760620d6eb651fe. + +The integration dynamically uses Transformers and Accelerate (Apache-2.0; +https://github.com/huggingface/transformers and +https://github.com/huggingface/accelerate), PyAV (BSD-3-Clause; +https://github.com/PyAV-Org/PyAV), and Pillow (MIT-CMU; +https://github.com/python-pillow/Pillow), alongside the shared runtime +dependencies listed above. These dependencies are not vendored. + -------------------------------------------------------------------------------- HPG-2011 NVIDIA CUDA Rasterizer port -------------------------------------------------------------------------------- @@ -176,8 +200,7 @@ License: Zlib (see LICENSES/Zlib.txt) Upstream: https://lodev.org/lodepng/ LodePNG is a PNG codec implementation embedded by the upstream -cudaraster framework. It is the only non-NVIDIA source physically -redistributed in this repository. The inline notice in lodepng.h +cudaraster framework. The inline notice in lodepng.h preserves the upstream copyright notice as required by the Zlib license; LICENSES/Zlib.txt reproduces the license text in full. diff --git a/apps/t2v/t2v/application.py b/apps/t2v/t2v/application.py index 2c722ce2d..e48402e95 100644 --- a/apps/t2v/t2v/application.py +++ b/apps/t2v/t2v/application.py @@ -180,6 +180,15 @@ def create_session(self, session_desc: SessionDesc) -> ISession: if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()): raise ValueError("A session prompt must be non-empty text.") self._validate_frame_size(session_desc, pipeline) + cache_kwargs = self._cache_initialization_kwargs(session_desc) + if cache_kwargs: + return self.session_type( + pipeline, + prompt, + session_desc, + config.total_blocks, + cache_init_kwargs=cache_kwargs, + ) return self.session_type(pipeline, prompt, session_desc, config.total_blocks) def close(self) -> None: @@ -193,6 +202,13 @@ def close(self) -> None: ## Integration hooks + def _cache_initialization_kwargs(self, session_desc: SessionDesc) -> dict[str, Any]: + """Return model-specific request inputs retained across session resets. + + Standard text, image, height, and width inputs remain framework-owned. + """ + return {} + def _configure_argument_parser(self, parser: argparse.ArgumentParser) -> None: """Add arguments this integration takes beyond the shared ones.""" diff --git a/apps/t2v/t2v/session.py b/apps/t2v/t2v/session.py index c168897ca..2ac64a6c9 100644 --- a/apps/t2v/t2v/session.py +++ b/apps/t2v/t2v/session.py @@ -3,7 +3,7 @@ """One text-to-video rollout: a prompt in, a chunk of frames per step out.""" -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from flashdreams.api_v2.loop import IModelLoop @@ -25,6 +25,9 @@ class T2VModelState: image: Any = None """Optional first-frame tensor for image-conditioned generation.""" + cache_init_kwargs: dict[str, Any] = field(default_factory=dict) + """Integration-specific request inputs reused when resetting the cache.""" + blocks_generated: int = 0 cache: Any = None @@ -91,6 +94,7 @@ def __init__( total_blocks: int, *, image: Any = None, + cache_init_kwargs: dict[str, Any] | None = None, ) -> None: """ Args: @@ -101,12 +105,20 @@ def __init__( against what the model can produce. total_blocks: Blocks this rollout generates before it is finished. image: Optional first-frame tensor retained across session resets. + cache_init_kwargs: Additional request inputs, excluding standard + ``text``, ``image``, ``height``, and ``width`` arguments. """ self._pipeline = pipeline self._prompt = prompt self._session_desc = session_desc self._total_blocks = total_blocks self._image = image + self._cache_init_kwargs = dict(cache_init_kwargs or {}) + reserved = {"text", "image", "height", "width"} & self._cache_init_kwargs.keys() + if reserved: + raise ValueError( + f"Cache inputs cannot override framework arguments: {sorted(reserved)}" + ) def init(self) -> None: """Encode the prompt and prepare the rollout's cache. @@ -119,6 +131,7 @@ def init(self) -> None: session_desc=self._session_desc, total_blocks=self._total_blocks, image=self._image, + cache_init_kwargs=dict(self._cache_init_kwargs), ) if state.prompt is not None: state.cache = _new_cache(state) @@ -139,10 +152,16 @@ def _new_cache(state: T2VModelState) -> Any: """Encode the prompt into a cache for one rollout.""" if state.prompt is None: raise RuntimeError("Cannot initialize a text-to-video cache without a prompt.") + reserved = {"text", "image", "height", "width"} & state.cache_init_kwargs.keys() + if reserved: + raise ValueError( + f"Cache inputs cannot override framework arguments: {sorted(reserved)}" + ) if state.image is not None: return state.pipeline.initialize_cache( text=[state.prompt], image=state.image, + **state.cache_init_kwargs, ) ratio = state.pipeline.decoder.spatial_compression_ratio return state.pipeline.initialize_cache( @@ -150,4 +169,5 @@ def _new_cache(state: T2VModelState) -> Any: image=None, height=state.session_desc.video_height // ratio, width=state.session_desc.video_width // ratio, + **state.cache_init_kwargs, ) diff --git a/apps/t2v/tests/test_session.py b/apps/t2v/tests/test_session.py index e31c40102..39f211082 100644 --- a/apps/t2v/tests/test_session.py +++ b/apps/t2v/tests/test_session.py @@ -22,6 +22,33 @@ pytestmark = pytest.mark.ci_cpu + +def test_request_kwargs_survive_reset_and_cannot_override_standard_fields(): + from t2v.session import T2VSession + + pipeline = _stand_in() + session = T2VSession( + pipeline, + _PROMPT, + _session_desc(), + 1, + cache_init_kwargs={"reference_paths": ("first", "last")}, + ) + session.init() + session.model_loop.reset() + assert len(pipeline.caches) == 2 + assert pipeline.caches[0]["reference_paths"] == ("first", "last") + assert pipeline.caches[1]["reference_paths"] == ("first", "last") + with pytest.raises(ValueError, match="framework arguments"): + T2VSession( + pipeline, + _PROMPT, + _session_desc(), + 1, + cache_init_kwargs={"text": ["override"]}, + ) + + _WIDTH = 128 """Frame width the stand-in generates. Not square, so a transposed frame cannot pass unnoticed, and a whole number of latents across.""" diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py b/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py index ca8766268..f257ae0da 100644 --- a/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py @@ -271,7 +271,7 @@ def compute_kv( def forward( self, x: Tensor, - kv_cache: KVCacheT, + kv_cache: KVCacheT | None = None, rope_freqs: Tensor | None = None, ) -> Tensor: """Apply the configured attention type to ``x`` and ``kv_cache``. @@ -281,7 +281,8 @@ def forward( Args: x: Query tokens, shape ``[..., L, query_dim]``. - kv_cache: Streaming cache for self-attention or precomputed static + kv_cache: ``None`` for cacheless bidirectional self-attention, + streaming cache for self-attention or precomputed static cache for cross-attention. A streaming cache must already be in its current-chunk update phase. rope_freqs: Optional positional data. Before-cache RoPE expects the diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py b/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py index 7e2118c16..6115840fd 100644 --- a/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py @@ -24,8 +24,6 @@ from enum import Enum import torch -from torch import Tensor, nn - from flashdreams.accelerated.common.non_persistent_linear import ( NonPersistentLinear, ) @@ -37,9 +35,9 @@ RoPEScope, RoPEStyle, ) -from flashdreams.accelerated.multi_head_attention.cudnn import ( - native_cudnn_fp8_sdpa, - torch_cudnn_sdpa, +from flashdreams.accelerated.multi_head_attention.sdpa import ( + SDPABackend, + scaled_dot_product_attention, ) from flashdreams.accelerated.multi_head_attention.triton import ( flash_attention_2, @@ -57,16 +55,7 @@ ) from flashdreams.core.attention import BlockKVCache from flashdreams.core.attention.rope_kernel import apply_rotary_pos_emb - - -class SDPABackend(str, Enum): - """Scaled-dot-product attention implementation.""" - - CUDNN = "cudnn" - """Use Torch cuDNN for FP16/BF16 and native cuDNN Frontend for FP8.""" - - FA2 = "fa2" - """Use Triton FlashAttention2 (FA2).""" +from torch import Tensor, nn class QKVFusionOption(str, Enum): @@ -160,6 +149,8 @@ def __post_init__(self) -> None: ) if not isinstance(self.use_tma, bool): raise TypeError(f"use_tma must be a bool; got {self.use_tma!r}") + if self.sdpa_backend is SDPABackend.TORCH and self.quantization.quantized_sdpa: + raise ValueError("Torch SDPA does not support quantized SDPA") class OptimizedMultiHeadAttention(MultiHeadAttention[BlockKVCache]): @@ -316,6 +307,11 @@ def _new_quantized_projection( dtype, ) + @torch.no_grad() + def refresh_derived_weights(self) -> None: + """Refresh fused and quantized projections after in-place weight edits.""" + self._refresh_derived_weights() + @torch.no_grad() def _refresh_derived_weights(self, *args: object) -> None: """Rebuild fused projection modules from checkpoint parameters. @@ -535,7 +531,7 @@ def compute_kv( def forward( self, x: Tensor, - kv_cache: BlockKVCache, + kv_cache: BlockKVCache | None = None, rope_freqs: Tensor | None = None, ) -> Tensor: """Apply self- or cross-attention using the configured cache lifecycle. @@ -546,7 +542,8 @@ def forward( Args: x: Query tokens shaped ``[..., L, Q]``. - kv_cache: Prepared rolling cache for self-attention or precomputed + kv_cache: ``None`` for cacheless self-attention, or prepared rolling + cache for self-attention or precomputed static cache for cross-attention. rope_freqs: Optional positional data. Before-cache RoPE expects the current chunk. After-cache RoPE expects positions covering the @@ -556,27 +553,37 @@ def forward( Returns: Output-projected tokens with the same shape and dtype as ``x``. """ - query_rope_freqs, key_rope_freqs = self._slice_rope_freqs( - rope_freqs, kv_cache, x.shape[-2] - ) - if self.attention_type is AttentionType.SELF_ATTENTION: - query = self._update_kv_and_compute_query(x, kv_cache, query_rope_freqs) + if kv_cache is None: + if self.attention_type is not AttentionType.SELF_ATTENTION: + raise ValueError("cross-attention requires a K/V cache") + self._validate_tokens(x, self.attention_config.query_dim, "x") + if self.qkv_fusion_option is QKVFusionOption.FULL: + query, key, value = self._project_qkv(x) + query = self._apply_qk_norm(query, self.query_norm) + key = self._apply_qk_norm(key, self.key_norm) + else: + query = self._project_query(x) + key, value = self._project_kv(x) + if rope_freqs is not None: + query = self._apply_rope(query, rope_freqs) + key = self._apply_rope(key, rope_freqs) else: - query = self._compute_query(x, query_rope_freqs) - self._validate_cache(kv_cache, x) - - # ``cached_k/v`` expose only the valid prefix while a rolling cache fills, - # and the complete fixed-size buffer after it reaches steady state. - key = kv_cache.cached_k() - if ( - self.attention_config.rope_config is not None - and self.attention_config.rope_config.scope is RoPEScope.AFTER_KV_CACHE - and key_rope_freqs is not None - ): - # The shared RoPE kernel is in-place; keep cache storage unrotated so - # rolling positions can be applied again on the next attention call. - key = self._apply_rope(key.to(x.dtype, copy=True), key_rope_freqs) - value = kv_cache.cached_v() + query_rope_freqs, key_rope_freqs = self._slice_rope_freqs( + rope_freqs, kv_cache, x.shape[-2] + ) + if self.attention_type is AttentionType.SELF_ATTENTION: + query = self._update_kv_and_compute_query(x, kv_cache, query_rope_freqs) + else: + query = self._compute_query(x, query_rope_freqs) + self._validate_cache(kv_cache, x) + key = kv_cache.cached_k() + if ( + self.attention_config.rope_config is not None + and self.attention_config.rope_config.scope is RoPEScope.AFTER_KV_CACHE + and key_rope_freqs is not None + ): + key = self._apply_rope(key.to(x.dtype, copy=True), key_rope_freqs) + value = kv_cache.cached_v() if self.optimized_impl_config.quantization.quantized_sdpa: query = query.to(torch.float8_e4m3fn) key = key.to(torch.float8_e4m3fn) @@ -730,33 +737,14 @@ def _attention( Returns: Attention output with shape ``[B, L, H, D]``. """ - if self.sdpa_backend is SDPABackend.CUDNN: - # The module and Triton kernel use token-major ``[B, L/S, H, D]``. - # PyTorch SDPA instead interprets its two middle axes as ``[H, L/S]``. - # These transposes change only shape/stride metadata. - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - - # PyTorch's public dispatcher rejects FP8 inputs, so use a cuDNN - # Frontend FP8 graph for e4m3 attention. - if query.dtype is torch.float8_e4m3fn: - output = native_cudnn_fp8_sdpa(query, key, value) - else: - output = torch_cudnn_sdpa(query, key, value) - - # Restore the module-wide ``[B, L, H, D]`` contract for head merging. - output = output.transpose(1, 2) - return output if output_dtype is None else output.to(output_dtype) - - attention = ( - flash_attention_2_tma - if self.use_tma and is_tma_flash_attention_supported(query, key, value) - else flash_attention_2 + return scaled_dot_product_attention( + query, + key, + value, + backend=self.sdpa_backend, + use_tma=self.use_tma, + output_dtype=output_dtype, ) - if output_dtype is None: - return attention(query, key, value) - return attention(query, key, value, output_dtype=output_dtype) def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: """Apply the shared RoPE kernel to token-major head features. @@ -771,12 +759,24 @@ def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: rope_config = self.attention_config.rope_config if rope_config is None: return x - return apply_rotary_pos_emb( - x, + if rope_freqs.ndim != 4: + raise ValueError("rope_freqs must have shape [L, 1, 1, rotary_dim]") + rotary_dim = rope_freqs.shape[-1] + if rotary_dim <= 0 or rotary_dim % 2 or rotary_dim > x.shape[-1]: + raise ValueError( + "RoPE width must be positive, even and no larger than head_dim" + ) + if tuple(rope_freqs.shape) != (x.shape[1], 1, 1, rotary_dim): + raise ValueError("rope_freqs must have shape [L, 1, 1, rotary_dim]") + if rope_freqs.device != x.device: + raise RuntimeError("rope_freqs and x must be on the same device") + apply_rotary_pos_emb( + x[..., :rotary_dim], rope_freqs, interleaved=rope_config.style is RoPEStyle.INTERLEAVED, inplace=True, ) + return x # ------------------------ Validation ------------------------ # @@ -787,7 +787,7 @@ def _validate_cuda_device(self, device: torch.device | str) -> None: device: Device used by attention inputs and cache storage. Raises: - RuntimeError: The CUDA device predates Hopper. + RuntimeError: The device predates Ampere, or FP8 requires Hopper. """ device = torch.device(device) if device.type != "cuda": @@ -797,9 +797,19 @@ def _validate_cuda_device(self, device: torch.device | str) -> None: ) if self._validated_cuda_device_index == device_index: return - if torch.cuda.get_device_capability(device_index)[0] < 9: + quantization = self.optimized_impl_config.quantization + fp8 = (torch.float8_e4m3fn, torch.float8_e5m2) + needs_hopper = ( + quantization.quantized_sdpa + or quantization.projection in fp8 + or quantization.output_projection in fp8 + ) + if torch.cuda.get_device_capability(device_index)[0] < ( + 9 if needs_hopper else 8 + ): raise RuntimeError( - "OptimizedMultiHeadAttention requires compute capability 9.0 or newer" + "OptimizedMultiHeadAttention requires compute capability " + + ("9.0 for FP8" if needs_hopper else "8.0 or newer") ) self._validated_cuda_device_index = device_index @@ -813,7 +823,8 @@ def _validate_tokens(self, x: Tensor, feature_dim: int, name: str) -> None: Raises: ValueError: ``x`` lacks sequence/feature axes or has the wrong width. - RuntimeError: ``x`` is not CUDA FP16/BF16 or the GPU predates Hopper. + RuntimeError: ``x`` is not CUDA FP16/BF16 or the device does not + support the configured precision. """ if x.ndim < 2: raise ValueError( @@ -974,7 +985,18 @@ def _validate_fused_update_inputs( if self.attention_config.rope_config is not None and rope_freqs is not None: # RoPE coefficients cover this ``L``-token chunk and broadcast across # flattened batches and ``H`` heads inside the shared kernel. - expected_rope_shape = (x.shape[-2], 1, 1, self.attention_config.head_dim) + if rope_freqs.ndim != 4: + raise ValueError("rope_freqs must have shape [L, 1, 1, rotary_dim]") + rotary_dim = rope_freqs.shape[-1] + if ( + rotary_dim <= 0 + or rotary_dim % 2 + or rotary_dim > self.attention_config.head_dim + ): + raise ValueError( + "RoPE width must be positive, even and no larger than head_dim" + ) + expected_rope_shape = (x.shape[-2], 1, 1, rotary_dim) if tuple(rope_freqs.shape) != expected_rope_shape: raise ValueError( f"rope_freqs must have shape {expected_rope_shape}; " @@ -1151,11 +1173,11 @@ def _project_output(self, x: Tensor) -> Tensor: __all__ = [ + "OptimizedImplConfig", + "OptimizedMultiHeadAttention", "QKVFusionOption", "QuantizationOption", "SDPABackend", - "OptimizedImplConfig", - "OptimizedMultiHeadAttention", "flash_attention_2", "flash_attention_2_tma", "is_tma_flash_attention_supported", diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/sdpa.py b/flashdreams/flashdreams/accelerated/multi_head_attention/sdpa.py new file mode 100644 index 000000000..7717613a2 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/sdpa.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared scaled-dot-product attention over projected token-major tensors.""" + +from enum import Enum + +import torch +import torch.nn.functional as F +from torch import Tensor + + +class SDPABackend(str, Enum): + """Scaled-dot-product attention implementation.""" + + TORCH = "torch" + """Use PyTorch's native device- and dtype-aware dispatcher.""" + + CUDNN = "cudnn" + """Use Torch cuDNN or native cuDNN Frontend for FP8.""" + + FA2 = "fa2" + """Use Triton FlashAttention2.""" + + +def scaled_dot_product_attention( + query: Tensor, + key: Tensor, + value: Tensor, + *, + is_causal: bool = False, + backend: SDPABackend = SDPABackend.TORCH, + use_tma: bool = False, + output_dtype: torch.dtype | None = None, +) -> Tensor: + """Attend to projected Q/K/V in ``[B, L/S, H, D]`` layout. + + Args: + query: Projected query tokens. + key: Projected key tokens. + value: Projected value tokens. + is_causal: Apply a causal mask; supported by the native Torch backend. + backend: Implementation policy, independent of model projections. + use_tma: Prefer TMA when the FA2 backend and hardware support it. + output_dtype: Optional output storage dtype. + + Returns: + Attention output in token-major layout. + + Raises: + ValueError: Tensor geometry or backend policy is invalid. + """ + if not isinstance(backend, SDPABackend): + raise TypeError(f"unsupported SDPA backend: {backend!r}") + if any(x.ndim != 4 for x in (query, key, value)): + raise ValueError("Q/K/V must have shape [B, L/S, H, D]") + if ( + key.shape != value.shape + or query.shape[0] != key.shape[0] + or query.shape[2:] != key.shape[2:] + ): + raise ValueError("Q/K/V batch, head and feature dimensions must match") + if is_causal and backend is not SDPABackend.TORCH: + raise ValueError("causal attention requires the Torch SDPA backend") + if backend is SDPABackend.TORCH: + if query.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + raise ValueError("Torch SDPA does not support FP8 inputs") + output = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + dropout_p=0.0, + is_causal=is_causal, + ).transpose(1, 2) + elif backend is SDPABackend.CUDNN: + from flashdreams.accelerated.multi_head_attention.cudnn import ( + native_cudnn_fp8_sdpa, + torch_cudnn_sdpa, + ) + + attention = ( + native_cudnn_fp8_sdpa + if query.dtype is torch.float8_e4m3fn + else torch_cudnn_sdpa + ) + output = attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) + ).transpose(1, 2) + else: + from flashdreams.accelerated.multi_head_attention.triton import ( + flash_attention_2, + flash_attention_2_tma, + is_tma_flash_attention_supported, + ) + + attention = ( + flash_attention_2_tma + if use_tma and is_tma_flash_attention_supported(query, key, value) + else flash_attention_2 + ) + return attention(query, key, value, output_dtype=output_dtype) + return output if output_dtype is None else output.to(output_dtype) diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py b/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py index e28001905..8dd0b8889 100644 --- a/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py @@ -20,9 +20,6 @@ from abc import abstractmethod import torch -import torch.nn.functional as F -from torch import Tensor, nn - from flashdreams.accelerated.multi_head_attention import ( AttentionConfig, AttentionType, @@ -31,7 +28,11 @@ RoPEScope, RoPEStyle, ) +from flashdreams.accelerated.multi_head_attention.sdpa import ( + scaled_dot_product_attention, +) from flashdreams.core.attention import BlockKVCache +from torch import Tensor, nn class TorchMultiHeadAttention(MultiHeadAttention[BlockKVCache]): @@ -178,14 +179,15 @@ def compute_kv( def forward( self, x: Tensor, - kv_cache: BlockKVCache, + kv_cache: BlockKVCache | None = None, rope_freqs: Tensor | None = None, ) -> Tensor: """Apply self- or cross-attention using the configured cache lifecycle. Args: x: Query tokens, shape ``[..., L, query_dim]``. - kv_cache: Prepared rolling cache for self-attention or precomputed + kv_cache: ``None`` for cacheless self-attention, or prepared rolling + cache for self-attention or precomputed static cache for cross-attention. rope_freqs: Optional positional data. Before-cache RoPE expects the current chunk. After-cache RoPE expects positions covering the @@ -195,6 +197,17 @@ def forward( Returns: Output-projected tokens with the same shape as ``x``. """ + if kv_cache is None: + if self.attention_type is not AttentionType.SELF_ATTENTION: + raise ValueError("cross-attention requires a K/V cache") + query = self._project_query(x) + key, value = self._project_kv(x) + if rope_freqs is not None: + query = self._apply_rope(query, rope_freqs) + key = self._apply_rope(key, rope_freqs) + output = self._output_projection(self._attention(query, key, value)) + return output.reshape(x.shape[:-2] + output.shape[-2:]) + query_rope_freqs, key_rope_freqs = self._slice_rope_freqs( rope_freqs, kv_cache, x.shape[-2] ) @@ -467,12 +480,17 @@ def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: """ if self.attention_config.rope_config is None: return x - if x.shape[-1] % 2 != 0: - raise ValueError(f"RoPE requires an even head_dim; got {x.shape[-1]}") + if rope_freqs.ndim != 4: + raise ValueError("rope_freqs must have shape [L, 1, 1, rotary_dim]") + rotary_dim = rope_freqs.shape[-1] + if rotary_dim <= 0 or rotary_dim % 2 or rotary_dim > x.shape[-1]: + raise ValueError( + "RoPE width must be positive, even and no larger than head_dim" + ) # RoPE lookup shape is ``[L, 1, 1, D]`` for an input shaped # ``[..., L, H, D]``. - expected_shape = (x.shape[-3], 1, 1, x.shape[-1]) + expected_shape = (x.shape[-3], 1, 1, rotary_dim) if tuple(rope_freqs.shape) != expected_shape: raise ValueError( f"rope_freqs must have shape {expected_shape}; " @@ -482,8 +500,9 @@ def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: # Broadcast positions over leading dimensions and heads: # ``[L, 1, 1, D] -> [..., L, 1, D]``. freqs = rope_freqs[:, 0, 0, :].reshape( - (1,) * (x.ndim - 3) + (x.shape[-3], 1, x.shape[-1]) + (1,) * (x.ndim - 3) + (x.shape[-3], 1, rotary_dim) ) + prefix, tail = x[..., :rotary_dim], x[..., rotary_dim:] # Materialize rotation coefficients in activation precision so the # elementwise rotation neither promotes projected tensors nor cache data. @@ -491,14 +510,16 @@ def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: sin_freqs = torch.sin(freqs).to(dtype=x.dtype) if self.attention_config.rope_config.style is RoPEStyle.INTERLEAVED: # Rotate adjacent feature pairs; shape stays ``[..., L, H, D]``. - rotated = torch.stack((-x[..., 1::2], x[..., 0::2]), dim=-1).flatten(-2) + rotated = torch.stack( + (-prefix[..., 1::2], prefix[..., 0::2]), dim=-1 + ).flatten(-2) else: # Rotate matching half-split features; shape stays ``[..., L, H, D]``. - first, second = x.chunk(2, dim=-1) + first, second = prefix.chunk(2, dim=-1) rotated = torch.cat((-second, first), dim=-1) # Apply the elementwise complex rotation: ``[..., L, H, D]``. - return x * cos_freqs + rotated * sin_freqs + return torch.cat((prefix * cos_freqs + rotated * sin_freqs, tail), dim=-1) def _attention(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor: """Apply non-causal scaled dot-product attention over visible K/V. @@ -514,23 +535,7 @@ def _attention(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor: # Move heads before tokens for SDPA: # Q ``[..., L, H, D] -> [..., H, L, D]`` and # K/V ``[..., S, H, D] -> [..., H, S, D]``. - query_heads = query.transpose(-3, -2) - key_heads = key.transpose(-3, -2) - value_heads = value.transpose(-3, -2) - - # Let PyTorch dispatch the available SDPA backend so the reference works - # on CPU and CUDA. Cache visibility defines the allowed context, while - # zero dropout and a non-causal mask make inference deterministic. - output = F.scaled_dot_product_attention( - query_heads, - key_heads, - value_heads, - dropout_p=0.0, - is_causal=False, - ) - - # Restore token-major layout: ``[..., H, L, D] -> [..., L, H, D]``. - return output.transpose(-3, -2) + return scaled_dot_product_attention(query, key, value) def _output_projection(self, x: Tensor) -> Tensor: """Concatenate attention heads and project back to query features. diff --git a/flashdreams/flashdreams/core/checkpoint/load.py b/flashdreams/flashdreams/core/checkpoint/load.py index 4432e1be5..00355543b 100644 --- a/flashdreams/flashdreams/core/checkpoint/load.py +++ b/flashdreams/flashdreams/core/checkpoint/load.py @@ -765,12 +765,17 @@ def _copy_checkpoint_tensor(destination: torch.Tensor, source: torch.Tensor) -> def _stream_safetensors_into_model( model: torch.nn.Module, path: str, + include_prefixes: tuple[str, ...] | None = None, ) -> torch.nn.Module: """Copy a safetensors checkpoint into a model with bounded host residency.""" model_state = model.state_dict() with safe_open(path, framework="pt", device="cpu") as source: - checkpoint_keys = set(source.keys()) + checkpoint_keys = { + key + for key in source.keys() + if include_prefixes is None or key.startswith(include_prefixes) + } model_keys = set(model_state) missing = sorted(model_keys - checkpoint_keys) unexpected = sorted(checkpoint_keys - model_keys) @@ -941,6 +946,7 @@ def _stream_sharded_safetensors_index_into_model( *, model: torch.nn.Module, checkpoint_min_free_gb: float | None, + include_prefixes: tuple[str, ...] | None = None, ) -> torch.nn.Module | None: """Stream a safetensors index checkpoint into ``model`` without merging.""" if checkpoint_path.startswith("s3://"): @@ -991,6 +997,7 @@ def _stream_sharded_safetensors_index_into_model( f"Invalid or empty weight_map in safetensors index: {index_local}" ) + weight_map = _select_checkpoint_prefixes(weight_map, model, include_prefixes) unique_shards = sorted(set(weight_map.values())) shard_to_path = _parallel_hf_hub_download_shards( repo_id=repo_id, @@ -1022,6 +1029,7 @@ def resolve_shard_path(shard_file: str) -> str: raise ValueError( f"Invalid or empty weight_map in safetensors index: {checkpoint_path}" ) + weight_map = _select_checkpoint_prefixes(weight_map, model, include_prefixes) base_dir = os.path.dirname(os.path.abspath(checkpoint_path)) def resolve_shard_path(shard_file: str) -> str: @@ -1034,6 +1042,29 @@ def resolve_shard_path(shard_file: str) -> str: ) +def _select_checkpoint_prefixes( + weight_map: dict[str, str], + model: torch.nn.Module, + include_prefixes: tuple[str, ...] | None, +) -> dict[str, str]: + """Select and validate component keys before downloading their shards.""" + if include_prefixes is None: + return weight_map + selected = { + name: shard + for name, shard in weight_map.items() + if name.startswith(include_prefixes) + } + model_keys = set(model.state_dict()) + if set(selected) != model_keys: + raise RuntimeError( + "Selected checkpoint components do not match model: " + f"missing={sorted(model_keys - selected.keys())[:20]}, " + f"unexpected={sorted(selected.keys() - model_keys)[:20]}" + ) + return selected + + def _resolve_streamable_safetensors_path( checkpoint_path: str, *, @@ -1126,6 +1157,8 @@ def load_checkpoint( map_location: str | torch.device = "cpu", check_success: bool = False, checkpoint_min_free_gb: float | None = None, + *, + include_prefixes: tuple[str, ...] | None = None, ) -> dict[str, torch.Tensor]: ... @@ -1139,6 +1172,8 @@ def load_checkpoint( map_location: str | torch.device = "cpu", check_success: bool = False, checkpoint_min_free_gb: float | None = None, + *, + include_prefixes: tuple[str, ...] | None = None, ) -> torch.nn.Module: ... @@ -1151,6 +1186,8 @@ def load_checkpoint( map_location: str | torch.device = "cpu", check_success: bool = False, checkpoint_min_free_gb: float | None = None, + *, + include_prefixes: tuple[str, ...] | None = None, ) -> dict[str, torch.Tensor] | torch.nn.Module: """Load checkpoints from S3, local disk, or Hugging Face. @@ -1172,6 +1209,9 @@ def load_checkpoint( checkpoint_min_free_gb: Optional first-run free-space requirement in GiB for Hugging Face checkpoint downloads. The ``FLASHDREAMS_MIN_CACHE_FREE_GB`` environment override still wins. + include_prefixes: Explicit component prefixes, including trailing dots. + Selected keys keep their names and must exactly match ``model``. + Only local/Hugging Face safetensors model loads support selection. Returns: State dict if ``model`` is ``None``, otherwise ``model`` with weights @@ -1186,6 +1226,26 @@ def load_checkpoint( >>> state = load_checkpoint("s3://bucket/foo.safetensors") >>> model = load_checkpoint("s3://bucket/dcp_dir/", model=my_model) """ + if include_prefixes is not None: + if not include_prefixes or any( + not prefix or not prefix.endswith(".") for prefix in include_prefixes + ): + raise ValueError( + "include_prefixes must contain nonempty module prefixes ending in '.'" + ) + if ( + model is None + or checkpoint_path.startswith("s3://") + or not ( + _is_sharded_safetensors_index_checkpoint(checkpoint_path) + or _get_checkpoint_extension(checkpoint_path) == ".safetensors" + ) + or checkpoint_type == "distributed" + ): + raise ValueError( + "Prefix selection requires a model and local/Hugging Face safetensors" + ) + # Auto-detect checkpoint type if checkpoint_type == "auto": if _is_sharded_safetensors_index_checkpoint(checkpoint_path): @@ -1204,6 +1264,7 @@ def load_checkpoint( checkpoint_path, model=model, checkpoint_min_free_gb=checkpoint_min_free_gb, + include_prefixes=include_prefixes, ) if streamed_model is not None: logger.info(f"Streamed checkpoint into model: {checkpoint_path}") @@ -1214,7 +1275,7 @@ def load_checkpoint( checkpoint_min_free_gb=checkpoint_min_free_gb, ) if stream_path is not None: - _stream_safetensors_into_model(model, stream_path) + _stream_safetensors_into_model(model, stream_path, include_prefixes) logger.info(f"Streamed checkpoint into model: {checkpoint_path}") return model state_dict = load_single_checkpoint( diff --git a/flashdreams/flashdreams/infra/acceleration/__init__.py b/flashdreams/flashdreams/infra/acceleration/__init__.py index c92c244cc..9f469611e 100644 --- a/flashdreams/flashdreams/infra/acceleration/__init__.py +++ b/flashdreams/flashdreams/infra/acceleration/__init__.py @@ -13,6 +13,7 @@ move_tensors_to_cpu, release_one_shot_encoder_references, run_one_shot_encoder_stage, + run_one_shot_stage, setup_one_shot_encoder, ) from flashdreams.infra.acceleration.frame_prefetch import ( @@ -57,6 +58,7 @@ "release_one_shot_encoder_references", "run_prewarm_sequence", "run_one_shot_encoder_stage", + "run_one_shot_stage", "run_timed_prewarm", "setup_one_shot_encoder", ] diff --git a/flashdreams/flashdreams/infra/acceleration/encoder_lifecycle.py b/flashdreams/flashdreams/infra/acceleration/encoder_lifecycle.py index 5fd405b85..376b8e50b 100644 --- a/flashdreams/flashdreams/infra/acceleration/encoder_lifecycle.py +++ b/flashdreams/flashdreams/infra/acceleration/encoder_lifecycle.py @@ -134,14 +134,14 @@ def move_tensors_to_cpu(value: Any, *, torch_module: Any | None = None) -> Any: return value -def run_one_shot_encoder_stage( +def run_one_shot_stage( stage: Callable[[], Any], *, release: Callable[[], Any] | None = None, cpu_result: bool = True, torch_module: Any | None = None, ) -> Any: - """Run an encoder-only stage under ``no_grad`` and release encoders after it.""" + """Run a stage under ``no_grad`` and release owned modules even on failure.""" torch = torch_module if torch_module is not None else _maybe_import_torch() no_grad = getattr(torch, "no_grad", None) context = no_grad() if callable(no_grad) else nullcontext() @@ -157,6 +157,10 @@ def run_one_shot_encoder_stage( del release_result +run_one_shot_encoder_stage = run_one_shot_stage +"""Backward-compatible name for encoder-only callers.""" + + def _maybe_import_torch() -> Any | None: try: return importlib.import_module("torch") diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/__init__.py b/flashdreams/flashdreams/infra/diffusion/scheduler/__init__.py index 6df37c806..03b8a3705 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/__init__.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/__init__.py @@ -20,6 +20,14 @@ Scheduler, SchedulerConfig, ) +from flashdreams.infra.diffusion.scheduler.data_flow_euler import ( + DataFlowEulerScheduler, + DataFlowEulerSchedulerConfig, +) +from flashdreams.infra.diffusion.scheduler.synchronized import ( + StepScheduler, + sample_synchronized, +) from flashdreams.infra.diffusion.scheduler.fm import ( FlowMatchScheduler, FlowMatchSchedulerConfig, @@ -34,6 +42,10 @@ ) __all__ = [ + "DataFlowEulerScheduler", + "DataFlowEulerSchedulerConfig", + "StepScheduler", + "sample_synchronized", "FlowPredictor", "Scheduler", "SchedulerConfig", diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/data_flow_euler.py b/flashdreams/flashdreams/infra/diffusion/scheduler/data_flow_euler.py new file mode 100644 index 000000000..6ed9a3475 --- /dev/null +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/data_flow_euler.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Data-ward rectified-flow Euler sampling on a shifted endpoint-inclusive grid.""" + +from dataclasses import dataclass, field +import math + +import torch +from torch import Tensor + +from flashdreams.infra.diffusion.scheduler.base import ( + FlowPredictor, + Scheduler, + SchedulerConfig, +) +from flashdreams.infra.diffusion.scheduler.synchronized import sample_synchronized + + +@dataclass(kw_only=True) +class DataFlowEulerSchedulerConfig(SchedulerConfig): + """Shifted sigma grid with a velocity pointing toward clean data.""" + + _target: type["DataFlowEulerScheduler"] = field( + default_factory=lambda: DataFlowEulerScheduler + ) + num_inference_steps: int = 30 + """Number of grid points; there are one fewer model evaluations.""" + shift: float = 12.0 + """Positive rational warp of the endpoint-inclusive sigma grid.""" + + +class DataFlowEulerScheduler(Scheduler): + """Blend the current sample and its data prediction in sigma space.""" + + def __init__(self, config: DataFlowEulerSchedulerConfig) -> None: + super().__init__(config) + if config.num_inference_steps < 2: + raise ValueError("num_inference_steps must be at least 2 grid points") + if not math.isfinite(config.shift) or config.shift <= 0: + raise ValueError("shift must be finite and positive") + base = torch.linspace( + 1.0, 0.0, config.num_inference_steps, dtype=torch.float32, device="cpu" + ) + sigmas = torch.unique_consecutive( + config.shift * base / (1 + (config.shift - 1) * base) + ) + self.register_buffer("sigmas", sigmas, persistent=False) + self.register_buffer("timesteps", 1.0 - sigmas[:-1], persistent=False) + + def _apply(self, fn, recurse=True): + grids = {name: getattr(self, name) for name in ("sigmas", "timesteps")} + super()._apply(fn, recurse=recurse) + for name, original in grids.items(): + setattr(self, name, original.to(device=getattr(self, name).device)) + return self + + def step(self, sample: Tensor, flow: Tensor, index: int) -> Tensor: + """Advance one step, preserving data-ward blend rounding.""" + time = self.timesteps[index].to(sample.device, sample.dtype) + denoised = sample + (1 - time) * flow + compute_dtype = ( + torch.float32 + if sample.dtype in (torch.float16, torch.bfloat16) + else sample.dtype + ) + ratio = self.sigmas[index + 1].to(sample.device, compute_dtype) / self.sigmas[ + index + ].to(sample.device, compute_dtype) + return ( + ratio * sample.to(compute_dtype) + (1 - ratio) * denoised.to(compute_dtype) + ).to(sample.dtype) + + @torch.no_grad() + def sample( + self, + initial_noise: Tensor, + predict_flow: FlowPredictor, + rng: torch.Generator | None = None, + ) -> Tensor: + """Sample one stream through the shared synchronized loop.""" + del rng + return sample_synchronized( + (initial_noise,), + (self,), + lambda samples, times: (predict_flow(samples[0], times[0]),), + )[0] + + def add_noise( + self, clean_input: Tensor, timestep: Tensor, rng: torch.Generator | None = None + ) -> Tensor: + """Mix clean data and Gaussian noise under the data-ward time convention.""" + noise = torch.randn( + clean_input.shape, + dtype=clean_input.dtype, + device=clean_input.device, + generator=rng, + ) + time = timestep.to(clean_input.device, clean_input.dtype) + while time.ndim < clean_input.ndim: + time = time.unsqueeze(-1) + return time * clean_input + (1 - time) * noise diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/synchronized.py b/flashdreams/flashdreams/infra/diffusion/scheduler/synchronized.py new file mode 100644 index 000000000..8d94fe3f0 --- /dev/null +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/synchronized.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Synchronized sampling of coupled diffusion streams.""" + +from collections.abc import Callable +from typing import Protocol + +from torch import Tensor + + +class StepScheduler(Protocol): + """Indexed deterministic schedule for one member of a joint prediction.""" + + timesteps: Tensor + """One-dimensional grid with one entry per model evaluation.""" + + def step(self, sample: Tensor, flow: Tensor, index: int) -> Tensor: + """Advance one sample using the prediction at ``index``.""" + ... + + +def sample_synchronized( + initial_samples: tuple[Tensor, ...], + schedulers: tuple[StepScheduler, ...], + predict_flow: Callable[ + [tuple[Tensor, ...], tuple[Tensor, ...]], tuple[Tensor, ...] + ], +) -> tuple[Tensor, ...]: + """Advance coupled streams with exactly one joint prediction per step. + + Args: + initial_samples: Generated samples, excluding immutable conditioning. + schedulers: One schedule per sample, all with the same number of steps. + predict_flow: Joint predictor returning one flow of each sample's shape. + + Returns: + Final samples in their original order, shapes, devices, and dtypes. + + Raises: + ValueError: Stream counts, schedule lengths, or tensor geometry disagree. + """ + if not initial_samples or len(initial_samples) != len(schedulers): + raise ValueError("Provide a nonempty sample tuple and one scheduler per sample") + grids = tuple(scheduler.timesteps for scheduler in schedulers) + if any(grid.ndim != 1 or grid.numel() == 0 for grid in grids): + raise ValueError("Schedules must be nonempty one-dimensional timestep grids") + if any(grid.numel() != grids[0].numel() for grid in grids): + raise ValueError("Coupled schedules must have equal lengths") + samples = initial_samples + for index in range(grids[0].numel()): + times = tuple( + grid[index].to(sample.device) + for grid, sample in zip(grids, samples, strict=True) + ) + flows = predict_flow(samples, times) + if not isinstance(flows, tuple) or len(flows) != len(samples): + raise ValueError("Joint predictor must return one flow per sample") + for sample, flow in zip(samples, flows, strict=True): + if flow.shape != sample.shape or flow.device != sample.device: + raise ValueError("Predicted flows must match sample shapes and devices") + if not flow.is_floating_point(): + raise ValueError("Predicted flows must be floating-point tensors") + advanced = tuple( + scheduler.step(sample, flow, index) + for scheduler, sample, flow in zip(schedulers, samples, flows, strict=True) + ) + if any( + new.shape != old.shape or new.device != old.device or new.dtype != old.dtype + for new, old in zip(advanced, samples, strict=True) + ): + raise ValueError("Schedulers must preserve sample shape, device, and dtype") + samples = advanced + return samples diff --git a/flashdreams/flashdreams/infra/encoder/text/qwen3_vl.py b/flashdreams/flashdreams/infra/encoder/text/qwen3_vl.py new file mode 100644 index 000000000..251f5d460 --- /dev/null +++ b/flashdreams/flashdreams/infra/encoder/text/qwen3_vl.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Headless Qwen3-VL hidden-state conditioning through Transformers.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from flashdreams.infra.encoder import Encoder, EncoderConfig + + +@dataclass(kw_only=True) +class Qwen3VLEncoderConfig(EncoderConfig): + """Checkpoint source and raw hidden-state selection.""" + + _target: type["Qwen3VLEncoder"] = field(default_factory=lambda: Qwen3VLEncoder) + model_name: str = "MiniMaxAI/MiniMax-H3" + """Repository or local checkpoint directory.""" + revision: str | None = None + """Pinned checkpoint revision when using the Hub.""" + cache_dir: str | None = None + """Optional Hugging Face cache directory.""" + hidden_layer: int = 50 + """Raw intermediate hidden state; must precede the final normalized state.""" + dtype: torch.dtype = torch.bfloat16 + """Weight and output precision.""" + subfolder: str = "text_encoder" + """Checkpoint partition containing the conditioner.""" + processor_subfolder: str = "processor" + """Checkpoint partition containing Qwen's media processor.""" + tokenizer_subfolder: str = "tokenizer" + """Checkpoint partition containing the presentation tokenizer.""" + local_files_only: bool = False + """Require all checkpoint files to be present locally.""" + loading_device: str | None = None + """Place checkpoint tensors directly on this device to avoid a CPU weight copy.""" + + +class Qwen3VLEncoder(Encoder): + """Selected Qwen3-VL state without a language-model projection.""" + + def __init__(self, config: Qwen3VLEncoderConfig) -> None: + super().__init__(config) + from transformers import Qwen2TokenizerFast, Qwen3VLModel, Qwen3VLProcessor + + kwargs = dict( + revision=config.revision, + cache_dir=config.cache_dir, + local_files_only=config.local_files_only, + ) + self.model = ( + Qwen3VLModel.from_pretrained( + config.model_name, + subfolder=config.subfolder, + dtype=config.dtype, + device_map=config.loading_device, + **kwargs, + ) + .eval() + .requires_grad_(False) + ) + self.processor = Qwen3VLProcessor.from_pretrained( + config.model_name, subfolder=config.processor_subfolder, **kwargs + ) + self.tokenizer = Qwen2TokenizerFast.from_pretrained( + config.model_name, subfolder=config.tokenizer_subfolder, **kwargs + ) + if ( + not 0 + <= config.hidden_layer + < self.model.config.text_config.num_hidden_layers + ): + raise ValueError( + "Qwen conditioning must select a raw intermediate hidden layer" + ) + + @torch.no_grad() + def forward(self, input: dict[str, Any]) -> torch.Tensor: + """Encode token IDs and optional processor-produced vision tensors.""" + device = self.model.device + token_ids = input["token_ids"] + input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) + modalities = torch.tensor( + self.processor.create_mm_token_type_ids([token_ids]), + dtype=torch.long, + device=device, + ) + vision = { + name: value.to(device=device, dtype=self.model.dtype) + if name.startswith("pixel_") + else value.to(device=device) + for name, value in input.get("vision_inputs", {}).items() + } + output = self.model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + mm_token_type_ids=modalities, + use_cache=False, + output_hidden_states=True, + **vision, + ) + return output.hidden_states[self.config.hidden_layer].to( + dtype=self.config.dtype + ) diff --git a/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py b/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py index 5413e68ec..e323a9ed2 100644 --- a/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py +++ b/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py @@ -19,8 +19,6 @@ import pytest import torch -from torch import Tensor - from flashdreams.accelerated.multi_head_attention import ( AttentionConfig, AttentionType, @@ -43,6 +41,7 @@ DTYPE_MAX, ) from flashdreams.core.attention import BlockKVCache +from torch import Tensor pytestmark = pytest.mark.ci_gpu @@ -577,7 +576,7 @@ def test_mha_optimized_quantized_projections_match_torch( ) @pytest.mark.parametrize("rope_scope", tuple(RoPEScope), ids=lambda value: value.value) @pytest.mark.parametrize( - "sdpa_backend", tuple(SDPABackend), ids=lambda value: value.value + "sdpa_backend", (SDPABackend.CUDNN, SDPABackend.FA2), ids=lambda value: value.value ) @pytest.mark.parametrize("use_tma", (False, True), ids=("no-tma", "tma")) @torch.inference_mode() @@ -631,6 +630,45 @@ def test_mha_optimized_quantized_sdpa_matches_torch( ) +@torch.inference_mode() +@pytest.mark.parametrize("fusion", (QKVFusionOption.NONE, QKVFusionOption.FULL)) +@pytest.mark.parametrize("rotary_dim", (48, 96, 128)) +def test_cacheless_partial_rope_on_ampere( + cuda_device: torch.device, + fusion: QKVFusionOption, + rotary_dim: int, +) -> None: + """Compare cacheless native attention without Hopper-only features.""" + if torch.cuda.get_device_capability(cuda_device)[0] < 8: + pytest.skip("Ampere or newer required") + config = AttentionConfig( + query_dim=128, + n_heads=2, + head_dim=128, + rope_config=RoPEConfig(style=RoPEStyle.SPLIT), + ) + reference = _TorchMHA(AttentionType.SELF_ATTENTION, config) + actual = _OptimizedMHA( + AttentionType.SELF_ATTENTION, + config, + OptimizedImplConfig( + qkv_fusion_option=fusion, + sdpa_backend=SDPABackend.TORCH, + use_tma=False, + ), + ) + actual.load_state_dict(reference.state_dict(), strict=True) + reference.to(device=cuda_device, dtype=torch.bfloat16).eval() + actual.to(device=cuda_device, dtype=torch.bfloat16).eval() + x = torch.randn(1, 7, 128, device=cuda_device, dtype=torch.bfloat16) + half = torch.randn(7, 1, 1, rotary_dim // 2, device=cuda_device) + frequencies = torch.cat((half, half), dim=-1) + _assert_close( + actual(x, rope_freqs=frequencies), reference(x, rope_freqs=frequencies) + ) + _assert_close(actual(x), reference(x)) + + @torch.inference_mode() def test_native_cudnn_fp8_sdpa_returns_independent_outputs( cuda_device: torch.device, diff --git a/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py b/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py index 6a0e80ed4..e3ace0ea3 100644 --- a/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py +++ b/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py @@ -20,8 +20,6 @@ import pytest import torch import torch.nn.functional as F -from torch import Tensor, nn - from flashdreams.accelerated.multi_head_attention import ( AttentionConfig, AttentionType, @@ -31,10 +29,92 @@ RoPEStyle, ) from flashdreams.accelerated.multi_head_attention.torch import TorchMultiHeadAttention +from torch import Tensor, nn pytestmark = pytest.mark.ci_cpu +@pytest.mark.parametrize("rope_style", tuple(RoPEStyle)) +def test_cacheless_partial_rope_matches_manual_attention(rope_style: RoPEStyle) -> None: + """Preserve unrotated features and cacheless full-sequence semantics.""" + module = _IdentityMHA(rope_style) + x = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) / 10 + frequencies = torch.tensor([0.1, 0.4, 0.9]).reshape(3, 1, 1, 1).expand(3, 1, 1, 2) + heads = x.unsqueeze(2) + rotated = torch.cat( + ( + _apply_rope(heads[..., :2], frequencies, rope_style), + heads[..., 2:], + ), + dim=-1, + ) + actual_rotated = module._apply_rope(heads, frequencies) + torch.testing.assert_close(actual_rotated, rotated) + assert torch.equal(actual_rotated[..., 2:], heads[..., 2:]) + expected = ( + F.scaled_dot_product_attention( + rotated.transpose(1, 2), + rotated.transpose(1, 2), + heads.transpose(1, 2), + ) + .transpose(1, 2) + .flatten(-2) + ) + torch.testing.assert_close(module(x, rope_freqs=frequencies), expected) + torch.testing.assert_close( + module(x), + F.scaled_dot_product_attention( + heads.transpose(1, 2), + heads.transpose(1, 2), + heads.transpose(1, 2), + ) + .transpose(1, 2) + .flatten(-2), + ) + + +def test_cacheless_cross_attention_rejected() -> None: + module = _IdentityMHA(RoPEStyle.SPLIT, AttentionType.CROSS_ATTENTION) + with pytest.raises(ValueError, match="requires a K/V cache"): + module(torch.zeros(1, 2, 4)) + + +def test_partial_rope_rounds_coefficients_before_half_precision_products() -> None: + """Preserve video-VAE coefficient precision and the untouched feature tail.""" + module = _IdentityMHA(RoPEStyle.SPLIT) + x = torch.tensor([0.17, -0.41, 0.83, 0.32], dtype=torch.float16).reshape(1, 1, 1, 4) + freqs = torch.tensor([0.71, 0.71]).reshape(1, 1, 1, 2) + prefix = x[..., :2] + rotated = torch.cat((-prefix[..., 1:], prefix[..., :1]), dim=-1) + expected = prefix * freqs.cos().half() + rotated * freqs.sin().half() + actual = module._apply_rope(x, freqs) + assert torch.equal(actual[..., :2], expected) + assert torch.equal(actual[..., 2:], x[..., 2:]) + + +def test_projected_causal_sdpa_matches_torch() -> None: + """Keep causal audio attention and FP32 computation on the shared path.""" + from flashdreams.accelerated.multi_head_attention.sdpa import ( + SDPABackend, + scaled_dot_product_attention, + ) + + query = torch.randn(2, 5, 3, 8) + expected = F.scaled_dot_product_attention( + query.transpose(1, 2), + query.transpose(1, 2), + query.transpose(1, 2), + is_causal=True, + ).transpose(1, 2) + torch.testing.assert_close( + scaled_dot_product_attention(query, query, query, is_causal=True), expected + ) + with pytest.raises(ValueError, match="causal attention requires"): + scaled_dot_product_attention( + query, query, query, is_causal=True, backend=SDPABackend.FA2 + ) + + class _IdentityMHA(TorchMultiHeadAttention): """Provide identity projections for direct attention comparisons.""" diff --git a/flashdreams/tests/test_checkpoint_loading.py b/flashdreams/tests/test_checkpoint_loading.py index f3b129bd1..70116cbbd 100644 --- a/flashdreams/tests/test_checkpoint_loading.py +++ b/flashdreams/tests/test_checkpoint_loading.py @@ -17,6 +17,92 @@ pytestmark = pytest.mark.ci_cpu +@pytest.mark.parametrize("sharded", [False, True]) +def test_scoped_model_load_is_strict_and_ignores_unselected_shards(tmp_path, sharded): + """Select codec components without opening unrelated weight shards.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + model = torch.nn.Module() + model.decoder = torch.nn.Linear(2, 2, bias=False) + expected = torch.arange(4, dtype=torch.float32).reshape(2, 2) + shard = tmp_path / "decoder.safetensors" + if sharded: + save_safetensors_file({"decoder.weight": expected}, shard) + checkpoint = tmp_path / "model.safetensors.index.json" + checkpoint.write_text( + json.dumps( + { + "weight_map": { + "decoder.weight": shard.name, + "encoder.weight": "absent-unused-shard.safetensors", + } + } + ) + ) + else: + checkpoint = shard + save_safetensors_file( + {"decoder.weight": expected, "encoder.weight": torch.zeros(1)}, shard + ) + checkpoint_load.load_checkpoint( + str(checkpoint), model=model, include_prefixes=("decoder.",) + ) + torch.testing.assert_close(model.decoder.weight, expected) + with pytest.raises(RuntimeError, match="match"): + checkpoint_load.load_checkpoint( + str(checkpoint), model=model, include_prefixes=("encoder.",) + ) + with pytest.raises(RuntimeError, match="match"): + checkpoint_load.load_checkpoint(str(checkpoint), model=model) + + +def test_scoped_remote_load_filters_before_downloading(monkeypatch, tmp_path): + """Download only selected shards from a Hub index.""" + module = importlib.import_module("flashdreams.core.checkpoint.load") + index = tmp_path / "model.safetensors.index.json" + index.write_text( + json.dumps( + { + "weight_map": { + "decoder.weight": "wanted.safetensors", + "encoder.weight": "unwanted.safetensors", + } + } + ) + ) + shard = tmp_path / "wanted.safetensors" + save_safetensors_file({"decoder.weight": torch.ones(2, 2)}, shard) + monkeypatch.setattr(module, "hf_hub_download", lambda **kwargs: str(index)) + monkeypatch.setattr( + module, "_preflight_checkpoint_cache_requirement", lambda **kwargs: None + ) + monkeypatch.setattr(module, "_preflight_hf_cache", lambda **kwargs: 0) + downloaded = [] + + def fetch(**kwargs): + downloaded.extend(kwargs["shard_files"]) + return {"wanted.safetensors": str(shard)} + + monkeypatch.setattr(module, "_parallel_hf_hub_download_shards", fetch) + model = torch.nn.Module() + model.decoder = torch.nn.Linear(2, 2, bias=False) + module.load_checkpoint( + "https://huggingface.co/example/model/blob/abc/model.safetensors.index.json", + model=model, + include_prefixes=("decoder.",), + ) + assert downloaded == ["wanted.safetensors"] + + +@pytest.mark.parametrize("prefixes", [(), ("",), ("decoder",)]) +def test_scoped_load_rejects_ambiguous_prefixes(prefixes): + from flashdreams.core.checkpoint.load import load_checkpoint + + with pytest.raises(ValueError, match="prefixes"): + load_checkpoint( + "unused.safetensors", model=torch.nn.Linear(2, 2), include_prefixes=prefixes + ) + + def test_local_safetensors_uses_file_backed_loader( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/flashdreams/tests/test_synchronized_sampling.py b/flashdreams/tests/test_synchronized_sampling.py new file mode 100644 index 000000000..de9131ee8 --- /dev/null +++ b/flashdreams/tests/test_synchronized_sampling.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU contracts for synchronized data-ward flow sampling.""" + +import pytest +import torch + +from flashdreams.infra.diffusion.scheduler import ( + DataFlowEulerSchedulerConfig, + sample_synchronized, +) + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_joint_matches_scalar_reference(dtype): + samples = ( + torch.linspace(-1, 1, 12).reshape(3, 4).to(dtype), + torch.zeros(2, 3, dtype=dtype), + ) + schedulers = tuple( + DataFlowEulerSchedulerConfig(shift=shift).setup() for shift in (12, 3) + ) + calls = [] + + def predict(values, times): + calls.append(times) + return tuple(value.float() * 0.1 + 0.2 for value in values) + + result = sample_synchronized(samples, schedulers, predict) + assert len(calls) == 29 + for initial, scheduler, actual in zip(samples, schedulers, result, strict=True): + expected = initial + for index, timestep in enumerate(scheduler.timesteps): + flow = expected.float() * 0.1 + 0.2 + denoised = expected + (1 - timestep.to(dtype)) * flow + ratio = scheduler.sigmas[index + 1] / scheduler.sigmas[index] + expected = (ratio * expected.float() + (1 - ratio) * denoised.float()).to( + dtype + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + scalar = scheduler.sample( + initial, lambda value, time: value.float() * 0.1 + 0.2 + ) + torch.testing.assert_close(actual, scalar, rtol=0, atol=0) + + +def test_mismatched_schedules_fail_before_prediction(): + schedulers = tuple( + DataFlowEulerSchedulerConfig(num_inference_steps=n).setup() for n in (3, 4) + ) + with pytest.raises(ValueError, match="equal lengths"): + sample_synchronized( + (torch.zeros(1), torch.zeros(1)), + schedulers, + lambda *_: pytest.fail("must not predict"), + ) + + +def test_reject_invalid_prediction_shape(): + scheduler = DataFlowEulerSchedulerConfig(num_inference_steps=2).setup() + with pytest.raises(ValueError, match="shapes"): + sample_synchronized( + (torch.zeros(2),), (scheduler,), lambda *_: (torch.zeros(3),) + ) + + +def test_grid_survives_dtype_conversion(): + scheduler = DataFlowEulerSchedulerConfig().setup() + expected = scheduler.sigmas.clone() + scheduler.to(torch.bfloat16) + assert scheduler.timesteps.dtype == torch.float32 + torch.testing.assert_close(scheduler.sigmas, expected, rtol=0, atol=0) diff --git a/integrations/minimax_h3/minimax_h3/__init__.py b/integrations/minimax_h3/minimax_h3/__init__.py deleted file mode 100644 index 8e740516e..000000000 --- a/integrations/minimax_h3/minimax_h3/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""MiniMax H3 integration for the FlashDreams runtime.""" - -from minimax_h3.pipeline import MiniMaxH3Pipeline, MiniMaxH3PipelineConfig - -__all__ = ["MiniMaxH3Pipeline", "MiniMaxH3PipelineConfig"] diff --git a/integrations/minimax_h3/minimax_h3/config.py b/integrations/minimax_h3/minimax_h3/config.py deleted file mode 100644 index 790dd543a..000000000 --- a/integrations/minimax_h3/minimax_h3/config.py +++ /dev/null @@ -1,67 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Registered MiniMax H3 workflow and runner configs.""" - -from __future__ import annotations - -from flashdreams.infra.runner import RunnerConfig -from minimax_h3.model import MiniMaxH3DiffusionModelConfig -from minimax_h3.pipeline import MiniMaxH3PipelineConfig -from minimax_h3.runner import ( - MiniMaxH3FL2VARunnerConfig, - MiniMaxH3Ref2VARunnerConfig, - MiniMaxH3T2VARunnerConfig, -) -from minimax_h3.transformer import ( - H3_REF_TRANSFORMER_CHECKPOINT, - MiniMaxH3TransformerConfig, -) - -PIPELINE_MINIMAX_H3_T2VA = MiniMaxH3PipelineConfig( - name="minimax-h3-t2va", - workflow="t2va", -) -PIPELINE_MINIMAX_H3_FL2VA = MiniMaxH3PipelineConfig( - name="minimax-h3-fl2va", - workflow="fl2va", -) -PIPELINE_MINIMAX_H3_REF2VA = MiniMaxH3PipelineConfig( - name="minimax-h3-ref2va", - workflow="ref2va", - diffusion_model=MiniMaxH3DiffusionModelConfig( - transformer=MiniMaxH3TransformerConfig( - checkpoint_path=H3_REF_TRANSFORMER_CHECKPOINT, - device="cuda", - execution_device="cuda", - sequential_cpu_offload=False, - ) - ), -) - -RUNNER_MINIMAX_H3_T2VA = MiniMaxH3T2VARunnerConfig( - runner_name=PIPELINE_MINIMAX_H3_T2VA.name, - description="MiniMax H3 prompt-to-video generation with low-host-RAM staging.", - pipeline=PIPELINE_MINIMAX_H3_T2VA, -) -RUNNER_MINIMAX_H3_FL2VA = MiniMaxH3FL2VARunnerConfig( - runner_name=PIPELINE_MINIMAX_H3_FL2VA.name, - description=( - "MiniMax H3 first-frame, last-frame, or dual-keyframe video generation." - ), - pipeline=PIPELINE_MINIMAX_H3_FL2VA, -) -RUNNER_MINIMAX_H3_REF2VA = MiniMaxH3Ref2VARunnerConfig( - runner_name=PIPELINE_MINIMAX_H3_REF2VA.name, - description="MiniMax H3 ordered image, video, and audio reference generation.", - pipeline=PIPELINE_MINIMAX_H3_REF2VA, -) - -RUNNER_CONFIGS: dict[str, RunnerConfig] = { - config.runner_name: config - for config in ( - RUNNER_MINIMAX_H3_T2VA, - RUNNER_MINIMAX_H3_FL2VA, - RUNNER_MINIMAX_H3_REF2VA, - ) -} diff --git a/integrations/minimax_h3/minimax_h3/model.py b/integrations/minimax_h3/minimax_h3/model.py deleted file mode 100644 index f8fe0a64a..000000000 --- a/integrations/minimax_h3/minimax_h3/model.py +++ /dev/null @@ -1,199 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""FlashDreams diffusion model for MiniMax H3's paired latent streams.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, cast - -import torch -from loguru import logger -from torch import Tensor, nn - -from flashdreams.infra.diffusion.model import DiffusionModel, DiffusionModelConfig -from flashdreams.infra.diffusion.transformer import Transformer, TransformerConfig -from minimax_h3.scheduler import MiniMaxH3Scheduler, MiniMaxH3SchedulerConfig -from minimax_h3.transformer import ( - MiniMaxH3TransformerCache, - MiniMaxH3TransformerConfig, -) - - -@dataclass(kw_only=True) -class MiniMaxH3DiffusionModelConfig(DiffusionModelConfig): - """Native H3 transformer plus separate video and audio schedules.""" - - _target: type[MiniMaxH3DiffusionModel] = field( - default_factory=lambda: MiniMaxH3DiffusionModel - ) - transformer: TransformerConfig = field( - default_factory=lambda: MiniMaxH3TransformerConfig( - device="cuda", - execution_device="cuda", - sequential_cpu_offload=False, - ) - ) - scheduler: MiniMaxH3SchedulerConfig = field( - default_factory=MiniMaxH3SchedulerConfig - ) - audio_scheduler: MiniMaxH3SchedulerConfig = field( - default_factory=lambda: MiniMaxH3SchedulerConfig(shift=3.0) - ) - - -@dataclass(kw_only=True) -class MiniMaxH3DenoiseState: - """Packed conditioning, noise, and layout produced before denoising.""" - - latents: Tensor - audio_latents: Tensor - prompt_embeds: Tensor - position_ids: Tensor - token_tags: Tensor - video_indices: Tensor - audio_indices: Tensor - text_indices: Tensor - num_condition_video_rows: int - num_condition_audio_rows: int - num_latent_frames: int - latent_height: int - latent_width: int - - -class MiniMaxH3DiffusionModel(DiffusionModel[MiniMaxH3TransformerCache]): - """Run H3's joint forward under two FlashDreams-owned schedulers.""" - - config: MiniMaxH3DiffusionModelConfig - transformer: Transformer[MiniMaxH3TransformerCache] - scheduler: MiniMaxH3Scheduler - audio_scheduler: MiniMaxH3Scheduler - - def __init__(self, config: MiniMaxH3DiffusionModelConfig) -> None: - nn.Module.__init__(self) - self.config = config - self.transformer = config.transformer.setup() - self.scheduler = config.scheduler.setup() - self.audio_scheduler = config.audio_scheduler.setup() - - @staticmethod - def _row_timesteps( - state: MiniMaxH3DenoiseState, - video_timestep: Tensor, - audio_timestep: Tensor, - ) -> tuple[Tensor, Tensor]: - sequence_length = ( - state.video_indices.numel() - + state.audio_indices.numel() - + state.text_indices.numel() - ) - video_timestep = video_timestep.to( - device=state.video_indices.device, dtype=torch.float32 - ) - audio_timestep = audio_timestep.to( - device=state.video_indices.device, dtype=torch.float32 - ) - row_timesteps = video_timestep.expand(sequence_length).clone() - video_condition = state.video_indices[: state.num_condition_video_rows] - audio_condition = state.audio_indices[: state.num_condition_audio_rows] - audio_target = state.audio_indices[state.num_condition_audio_rows :] - row_timesteps[video_condition] = video_timestep.clamp_min(0.999) - row_timesteps[audio_target] = audio_timestep - row_timesteps[audio_condition] = 1.0 - return torch.unique(row_timesteps, sorted=True, return_inverse=True) - - @torch.no_grad() - def generate_joint(self, state: MiniMaxH3DenoiseState) -> Tensor: - """Denoise both streams and return video latents on the execution device.""" - device = self.transformer.device - video = state.latents.to(device) - audio = state.audio_latents.to(device) - state.prompt_embeds = state.prompt_embeds.to(device) - state.position_ids = state.position_ids.to(device) - state.token_tags = state.token_tags.to(device) - state.video_indices = state.video_indices.to(device) - state.audio_indices = state.audio_indices.to(device) - state.text_indices = state.text_indices.to(device) - - video_sigmas, video_timesteps = self.scheduler.schedule(device) - audio_sigmas, audio_timesteps = self.audio_scheduler.schedule(device) - if len(video_timesteps) != len(audio_timesteps): - raise RuntimeError("H3 video and audio schedules must have equal length") - - cache = MiniMaxH3TransformerCache( - audio_hidden_states=audio[None], - encoder_hidden_states=state.prompt_embeds, - timestep=torch.empty(0, device=device), - timestep_indices=torch.empty(0, dtype=torch.long, device=device), - token_tags=state.token_tags, - position_ids=state.position_ids, - video_indices=state.video_indices, - audio_indices=state.audio_indices, - text_indices=state.text_indices, - ) - for index, (video_timestep, audio_timestep) in enumerate( - zip(video_timesteps, audio_timesteps, strict=True) - ): - logger.info( - "MiniMax H3 denoise step {}/{}", - index + 1, - len(video_timesteps), - ) - cache.timestep, cache.timestep_indices = self._row_timesteps( - state, video_timestep, audio_timestep - ) - cache.audio_hidden_states = audio[None] - video_flow = self.transformer.predict_flow( - video[None], video_timestep, cache - )[0] - if cache.last_audio_flow is None: - raise RuntimeError("H3 transformer did not produce an audio flow") - audio_flow = cache.last_audio_flow[0] - - video_start = state.num_condition_video_rows - audio_start = state.num_condition_audio_rows - video[video_start:] = self.scheduler.step( - video[video_start:], - video_flow[video_start:].float(), - video_timestep, - video_sigmas[index], - video_sigmas[index + 1], - ) - audio[audio_start:] = self.audio_scheduler.step( - audio[audio_start:], - audio_flow[audio_start:].float(), - audio_timestep, - audio_sigmas[index], - audio_sigmas[index + 1], - ) - - rows = video[state.num_condition_video_rows :] - transformer_config = cast(Any, self.config.transformer) - patch_t, patch_h, patch_w = transformer_config.patch_size - channels = transformer_config.in_channels - rows = rows.reshape( - -1, - state.num_latent_frames // patch_t, - state.latent_height // patch_h, - state.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, - state.num_latent_frames, - state.latent_height, - state.latent_width, - ).contiguous() - - -__all__ = [ - "MiniMaxH3DenoiseState", - "MiniMaxH3DiffusionModel", - "MiniMaxH3DiffusionModelConfig", -] diff --git a/integrations/minimax_h3/minimax_h3/pipeline.py b/integrations/minimax_h3/minimax_h3/pipeline.py deleted file mode 100644 index eb4d952b3..000000000 --- a/integrations/minimax_h3/minimax_h3/pipeline.py +++ /dev/null @@ -1,848 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Crash-safe MiniMax H3 FL2VA pipeline for the FlashDreams runtime.""" - -from __future__ import annotations - -import gc -import hashlib -import json -import os -import threading -import time -from concurrent.futures import Future -from dataclasses import dataclass, field, replace -from datetime import datetime, timezone -from pathlib import Path -from typing import Any, Literal, cast - -import numpy as np -import torch -from loguru import logger -from torch import nn - -from flashdreams.infra.pipeline import ( - StreamInferencePipeline, - StreamInferencePipelineConfig, -) -from minimax_h3.constants import MODEL_ID, align_num_frames, validate_canvas -from minimax_h3.lora import load_lora -from minimax_h3.model import ( - MiniMaxH3DenoiseState, - MiniMaxH3DiffusionModelConfig, -) -from minimax_h3.references import MiniMaxH3ReferenceSpec, load_references -from minimax_h3.transformer import MiniMaxH3TransformerConfig - -MiniMaxH3Workflow = Literal["t2va", "fl2va", "ref2va"] - - -@dataclass(kw_only=True) -class MiniMaxH3PipelineConfig(StreamInferencePipelineConfig): - """Config for the FlashDreams-native H3 denoising pipeline.""" - - _target: type[MiniMaxH3Pipeline] = field(default_factory=lambda: MiniMaxH3Pipeline) - - diffusion_model: MiniMaxH3DiffusionModelConfig = field( - default_factory=MiniMaxH3DiffusionModelConfig - ) - """Native joint transformer and paired scheduler configuration.""" - - model_id: str = MODEL_ID - """Hugging Face model repository or local snapshot path.""" - - cache_dir: Path | None = None - """Optional Hugging Face model cache root.""" - - workflow: MiniMaxH3Workflow = "fl2va" - """Released checkpoint workflow selected by this registered pipeline.""" - - -@dataclass(frozen=True) -class _ReferenceLayout: - """Reference properties needed after its encoded media is checkpointed.""" - - kind: str - has_audio: bool - - -@dataclass(kw_only=True) -class MiniMaxH3PipelineCache: - """Per-rollout H3 request, checkpoints, and runtime metrics.""" - - prompt: str - workflow: MiniMaxH3Workflow - image_path: Path | None - last_image_path: Path | None - references: tuple[MiniMaxH3ReferenceSpec, ...] - output_path: Path - width: int - height: int - duration: float - steps: int - seed: int - low_ram: bool - restart: bool - attention: str - lora: str | None - lora_weight_name: str | None - lora_scale: float - latent_checkpoint: Path - conditioning_checkpoint: Path - generated: bool = False - elapsed_seconds: float = 0.0 - conditioning_seconds: float = 0.0 - denoise_seconds: float = 0.0 - denoise_prepare_seconds: float = 0.0 - transformer_load_seconds: float = 0.0 - denoise_compute_seconds: float = 0.0 - denoise_cleanup_seconds: float = 0.0 - latent_checkpoint_seconds: float = 0.0 - latent_checkpoint_future: Future[float] | None = field(default=None, repr=False) - decode_seconds: float = 0.0 - peak_gpu_memory_gib: float = 0.0 - attention_backend: str = "default" - resumed_stage: str | None = None - - -def _replace_blocks(pipe: Any, names: tuple[str, ...]) -> None: - from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks - - selected: dict[str, Any] = {} - for requested in names: - if requested in pipe._blocks.sub_blocks: - selected[requested] = pipe._blocks.sub_blocks[requested] - continue - prefix = requested + "." - for actual, block in pipe._blocks.sub_blocks.items(): - if actual.startswith(prefix): - selected[actual.removeprefix(prefix)] = block - if not selected: - raise KeyError(f"none of the requested pipeline stages exist: {names}") - pipe._blocks = SequentialPipelineBlocks.from_blocks_dict(selected) - - -def _release_pipeline(pipe: Any) -> None: - for component in pipe.components.values(): - if isinstance(component, nn.Module): - try: - component.to_empty(device="cpu") - except (AttributeError, RuntimeError): - pass - del pipe - gc.collect() - torch.cuda.empty_cache() - - -def _atomic_json(path: Path, payload: dict[str, Any]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - temporary = path.with_name(path.name + ".tmp") - temporary.write_text(json.dumps(payload, indent=2) + "\n") - os.replace(temporary, path) - - -def _write_status(checkpoint: Path, stage: str, **details: Any) -> None: - _atomic_json( - checkpoint.with_suffix(checkpoint.suffix + ".status.json"), - { - "stage": stage, - "updated_at": datetime.now(timezone.utc).isoformat(), - **details, - }, - ) - - -def _file_identity(path: Path | None) -> dict[str, str | int] | None: - if path is None: - return None - resolved = path.resolve() - stat = resolved.stat() - return { - "path": str(resolved), - "size": stat.st_size, - "mtime_ns": stat.st_mtime_ns, - } - - -def _conditioning_manifest( - cache: MiniMaxH3PipelineCache, model_id: str -) -> dict[str, Any]: - return { - "workflow": cache.workflow, - "prompt": cache.prompt, - "image": _file_identity(cache.image_path), - "last_image": _file_identity(cache.last_image_path), - "references": [reference.manifest() for reference in cache.references], - "width": cache.width, - "height": cache.height, - "duration": cache.duration, - "model_id": model_id, - } - - -def _generation_manifest( - cache: MiniMaxH3PipelineCache, model_id: str -) -> dict[str, Any]: - return { - **_conditioning_manifest(cache, model_id), - "steps": cache.steps, - "seed": cache.seed, - "attention": cache.attention, - "lora": cache.lora, - "lora_weight_name": cache.lora_weight_name, - "lora_scale": cache.lora_scale, - } - - -def _signature(manifest: dict[str, Any]) -> str: - encoded = json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode() - return hashlib.sha256(encoded).hexdigest() - - -def _save_conditioning( - cache: MiniMaxH3PipelineCache, - model_id: str, - values: dict[str, Any], -) -> None: - from safetensors.torch import save_file - - manifest = _conditioning_manifest(cache, model_id) - path = cache.conditioning_checkpoint - temporary = path.with_name(path.name + ".tmp") - condition_latents = values["condition_latents"] - audio_condition_latents = values["audio_condition_latents"] - tensors = { - "prompt_embeds": values["prompt_embeds"].detach().cpu().contiguous(), - "text_token_tags": values["text_token_tags"].detach().cpu().contiguous(), - **{ - f"condition_latents.{index}": latent.detach().cpu().contiguous() - for index, latent in enumerate(condition_latents) - }, - **{ - f"audio_condition_latents.{index}": latent.detach().cpu().contiguous() - for index, latent in enumerate(audio_condition_latents) - }, - } - save_file( - tensors, - str(temporary), - metadata={ - "stage": "conditioned", - "manifest": json.dumps(manifest, sort_keys=True), - "signature": _signature(manifest), - "height": str(values["height"]), - "width": str(values["width"]), - "num_frames": str(values["num_frames"]), - "keyframe_anchors": json.dumps(list(values["keyframe_anchors"])), - "condition_count": str(len(condition_latents)), - "audio_condition_count": str(len(audio_condition_latents)), - "reference_layout": json.dumps( - [ - {"kind": reference.kind, "has_audio": reference.has_audio} - for reference in values["normalized_references"] - ] - ), - }, - ) - os.replace(temporary, path) - _write_status(path, "conditioned") - - -def _load_conditioning(cache: MiniMaxH3PipelineCache, model_id: str) -> dict[str, Any]: - from safetensors import safe_open - from safetensors.torch import load_file - - path = cache.conditioning_checkpoint - with safe_open(path, framework="pt", device="cpu") as handle: - metadata = handle.metadata() or {} - expected = _signature(_conditioning_manifest(cache, model_id)) - if metadata.get("stage") != "conditioned" or metadata.get("signature") != expected: - raise ValueError( - f"conditioning checkpoint does not match this request: {path}; " - "use --restart" - ) - tensors = load_file(path, device="cpu") - condition_count = int(metadata["condition_count"]) - audio_condition_count = int(metadata.get("audio_condition_count", "0")) - reference_layout = json.loads(metadata.get("reference_layout", "[]")) - return { - "prompt_embeds": tensors["prompt_embeds"], - "text_token_tags": tensors["text_token_tags"], - "condition_latents": [ - tensors[f"condition_latents.{index}"] for index in range(condition_count) - ], - "audio_condition_latents": [ - tensors[f"audio_condition_latents.{index}"] - for index in range(audio_condition_count) - ], - "normalized_references": [ - _ReferenceLayout(kind=reference["kind"], has_audio=reference["has_audio"]) - for reference in reference_layout - ], - "height": int(metadata["height"]), - "width": int(metadata["width"]), - "num_frames": int(metadata["num_frames"]), - "keyframe_anchors": tuple(json.loads(metadata["keyframe_anchors"])), - } - - -def _save_latents( - cache: MiniMaxH3PipelineCache, - model_id: str, - latents: torch.Tensor, -) -> None: - from safetensors.torch import save_file - - path = cache.latent_checkpoint - temporary = path.with_name(path.name + ".tmp") - manifest = _generation_manifest(cache, model_id) - save_file( - {"latents": latents.detach().cpu().contiguous()}, - str(temporary), - metadata={ - "stage": "denoised", - "manifest": json.dumps(manifest, sort_keys=True), - "signature": _signature(manifest), - }, - ) - os.replace(temporary, path) - - -def _save_latents_async( - cache: MiniMaxH3PipelineCache, - model_id: str, - latents: torch.Tensor, -) -> Future[float]: - """Persist the recovery checkpoint without blocking video decoding.""" - future: Future[float] = Future() - - def persist() -> None: - if not future.set_running_or_notify_cancel(): - return - started = time.monotonic() - try: - _save_latents(cache, model_id, latents) - except BaseException as exc: - future.set_exception(exc) - else: - future.set_result(time.monotonic() - started) - - threading.Thread( - target=persist, - name="minimax-h3-latent-checkpoint", - daemon=False, - ).start() - return future - - -def _finish_latent_checkpoint(cache: MiniMaxH3PipelineCache, *, wait: bool) -> None: - """Collect a completed checkpoint write, optionally waiting for it.""" - future = cache.latent_checkpoint_future - if future is None or (not wait and not future.done()): - return - cache.latent_checkpoint_seconds = future.result() - cache.latent_checkpoint_future = None - - -def _load_latents(cache: MiniMaxH3PipelineCache, model_id: str) -> torch.Tensor: - from safetensors import safe_open - from safetensors.torch import load_file - - path = cache.latent_checkpoint - with safe_open(path, framework="pt", device="cpu") as handle: - metadata = handle.metadata() or {} - expected = _signature(_generation_manifest(cache, model_id)) - if metadata.get("stage") != "denoised" or metadata.get("signature") != expected: - raise ValueError( - f"latent checkpoint does not match this request: {path}; use --restart" - ) - tensors = load_file(path, device="cpu") - return tensors["latents"] - - -class MiniMaxH3Pipeline(StreamInferencePipeline[Any, Any, Any]): - """FlashDreams H3 runtime with staged third-party conditioning and decode.""" - - config: MiniMaxH3PipelineConfig - - def __init__(self, config: MiniMaxH3PipelineConfig) -> None: - nn.Module.__init__(self) - self.config = config - self.register_buffer("_device_anchor", torch.empty(0), persistent=False) - - @property - def device(self) -> torch.device: - return cast(torch.Tensor, self._device_anchor).device - - def initialize_cache( - self, - *, - prompt: str, - image_path: Path | None, - last_image_path: Path | None, - references: tuple[MiniMaxH3ReferenceSpec, ...], - output_path: Path, - width: int, - height: int, - duration: float, - steps: int, - seed: int, - low_ram: bool, - restart: bool, - attention: str, - lora: str | None, - lora_weight_name: str | None, - lora_scale: float, - ) -> MiniMaxH3PipelineCache: - """Build a validated, checkpoint-aware workflow cache.""" - if not prompt.strip(): - raise ValueError("prompt cannot be empty") - workflow = self.config.workflow - if workflow == "t2va" and ( - image_path is not None or last_image_path is not None or references - ): - raise ValueError("t2va does not accept keyframes or references") - if workflow == "fl2va" and not (image_path or last_image_path): - raise ValueError("fl2va requires --image-path and/or --last-image-path") - if workflow == "ref2va" and not references: - raise ValueError("ref2va requires at least one --reference") - if workflow != "ref2va" and references: - raise ValueError(f"{workflow} does not accept ordered references") - for label, path in ( - ("first-frame", image_path), - ("last-frame", last_image_path), - ): - if path is not None and not path.is_file(): - raise FileNotFoundError(f"{label} image not found: {path}") - if steps < 2: - raise ValueError("steps must be at least 2 scheduler points") - if attention not in {"auto", "flash", "default"}: - raise ValueError(f"unsupported attention backend: {attention}") - if not 0 <= lora_scale <= 4: - raise ValueError("LoRA scale must be between 0 and 4") - validate_canvas(width, height) - align_num_frames(duration) - output_path.parent.mkdir(parents=True, exist_ok=True) - latent_checkpoint = output_path.with_suffix( - output_path.suffix + ".latents.safetensors" - ) - conditioning_checkpoint = output_path.with_suffix( - output_path.suffix + ".conditioning.safetensors" - ) - return MiniMaxH3PipelineCache( - prompt=prompt, - workflow=workflow, - image_path=image_path, - last_image_path=last_image_path, - references=references, - output_path=output_path, - width=width, - height=height, - duration=duration, - steps=steps, - seed=seed, - low_ram=low_ram, - restart=restart, - attention=attention, - lora=lora, - lora_weight_name=lora_weight_name, - lora_scale=lora_scale, - latent_checkpoint=latent_checkpoint, - conditioning_checkpoint=conditioning_checkpoint, - ) - - @torch.no_grad() - def generate( - self, - autoregressive_index: int, - cache: MiniMaxH3PipelineCache, - input: Any = None, - ) -> torch.Tensor: - """Generate and video-decode the single non-streaming H3 step.""" - del input - if autoregressive_index != 0 or cache.generated: - raise ValueError("MiniMax H3 supports exactly one runtime step per cache") - started = time.monotonic() - torch.cuda.reset_peak_memory_stats() - - if cache.latent_checkpoint.is_file() and not cache.restart: - cache.resumed_stage = "decode" - latents = _load_latents(cache, self.config.model_id) - else: - if cache.low_ram: - latents = self._generate_low_ram(cache) - else: - latents = self._generate_standard(cache) - cache.latent_checkpoint_future = _save_latents_async( - cache, self.config.model_id, latents - ) - - decode_started = time.monotonic() - try: - frames = self._decode_video(cache, latents) - except BaseException: - try: - _finish_latent_checkpoint(cache, wait=True) - except BaseException: - logger.exception("Latent checkpoint also failed after decode failure") - raise - cache.decode_seconds = time.monotonic() - decode_started - cache.elapsed_seconds = time.monotonic() - started - cache.peak_gpu_memory_gib = torch.cuda.max_memory_allocated() / 2**30 - cache.generated = True - return frames - - def mark_complete(self, cache: MiniMaxH3PipelineCache) -> None: - """Record completion only after the runtime output target closes.""" - if not cache.generated or not cache.output_path.is_file(): - raise RuntimeError("cannot complete H3 job before its MP4 is written") - _finish_latent_checkpoint(cache, wait=True) - _write_status( - cache.latent_checkpoint, "complete", output=str(cache.output_path) - ) - - def finalize( - self, - autoregressive_index: int, - cache: MiniMaxH3PipelineCache, - ) -> dict[str, float]: - """Return runtime metrics for the completed H3 rollout.""" - if autoregressive_index != 0 or not cache.generated: - raise ValueError("finalize requires the completed H3 runtime step") - _finish_latent_checkpoint(cache, wait=True) - return { - "conditioning_seconds": cache.conditioning_seconds, - "denoise_seconds": cache.denoise_seconds, - "denoise_prepare_seconds": cache.denoise_prepare_seconds, - "transformer_load_seconds": cache.transformer_load_seconds, - "denoise_compute_seconds": cache.denoise_compute_seconds, - "denoise_cleanup_seconds": cache.denoise_cleanup_seconds, - "latent_checkpoint_seconds": cache.latent_checkpoint_seconds, - "decode_seconds": cache.decode_seconds, - "total_seconds": cache.elapsed_seconds, - "peak_gpu_memory_gib": cache.peak_gpu_memory_gib, - } - - def _cache_dir(self) -> str | None: - return str(self.config.cache_dir) if self.config.cache_dir is not None else None - - def _apply_lora(self, transformer: Any, cache: MiniMaxH3PipelineCache) -> None: - if cache.lora is None: - return - converted = load_lora( - transformer, - cache.lora, - cache.lora_scale, - cache.lora_weight_name, - ) - logger.info("Loaded LoRA {} at scale {:g}", converted, cache.lora_scale) - - def _generate_low_ram(self, cache: MiniMaxH3PipelineCache) -> torch.Tensor: - if cache.conditioning_checkpoint.is_file() and not cache.restart: - cache.resumed_stage = "denoise" - conditioned = _load_conditioning(cache, self.config.model_id) - else: - conditioning_started = time.monotonic() - _write_status(cache.latent_checkpoint, "conditioning") - conditioned = self._condition(cache) - _save_conditioning(cache, self.config.model_id, conditioned) - cache.conditioning_seconds = time.monotonic() - conditioning_started - return self._run_native_denoise(cache, conditioned) - - def _condition(self, cache: MiniMaxH3PipelineCache) -> dict[str, Any]: - from diffusers.utils import load_image - - num_frames = align_num_frames(cache.duration) - media: dict[str, Any] - if cache.workflow == "t2va": - media = { - "height": cache.height, - "width": cache.width, - "num_frames": num_frames, - "keyframe_anchors": (), - "normalized_references": [], - } - text_inputs = {"prompt": cache.prompt} - encoded = { - "condition_latents": [], - "audio_condition_latents": [], - } - elif cache.workflow == "fl2va": - resize_inputs: dict[str, Any] = { - "height": cache.height, - "width": cache.width, - } - if cache.image_path is not None: - resize_inputs["image"] = load_image(str(cache.image_path)) - if cache.last_image_path is not None: - resize_inputs["last_image"] = load_image(str(cache.last_image_path)) - media = self._run_conditioning_stage( - cache.workflow, - ("before_encode",), - resize_inputs, - ["height", "width", "keyframes", "keyframe_anchors"], - ) - media["num_frames"] = num_frames - media["normalized_references"] = [] - text_inputs = {"prompt": cache.prompt, "keyframes": media["keyframes"]} - encoded = self._run_conditioning_stage( - cache.workflow, - ("vae_encoder",), - {"keyframes": media["keyframes"]}, - ["condition_latents"], - cuda_components=("vae",), - ) - encoded["audio_condition_latents"] = [] - else: - references = load_references(cache.references) - media = self._run_conditioning_stage( - cache.workflow, - ("before_encode",), - { - "references": references, - "height": cache.height, - "width": cache.width, - "num_frames": num_frames, - }, - ["height", "width", "num_frames", "normalized_references"], - ) - media["keyframe_anchors"] = () - text_inputs = { - "prompt": cache.prompt, - "normalized_references": media["normalized_references"], - } - encoded = self._run_conditioning_stage( - cache.workflow, - ("vae_encoder",), - {"normalized_references": media["normalized_references"]}, - ["condition_latents", "audio_condition_latents"], - cuda_components=("vae", "audio_vae"), - ) - - text = self._run_conditioning_stage( - cache.workflow, - ("text_encoder",), - text_inputs, - ["prompt_embeds", "text_token_tags"], - cuda_components=("text_encoder",), - ) - return { - **text, - **encoded, - "height": media["height"], - "width": media["width"], - "keyframe_anchors": media["keyframe_anchors"], - "normalized_references": media["normalized_references"], - "num_frames": media["num_frames"], - } - - def _run_conditioning_stage( - self, - workflow: MiniMaxH3Workflow, - blocks: tuple[str, ...], - inputs: dict[str, Any], - outputs: list[str], - *, - cuda_components: tuple[str, ...] = (), - ) -> dict[str, Any]: - from diffusers.modular_pipelines.modular_pipeline import ModularPipeline - - pipe = ModularPipeline.from_pretrained( - self.config.model_id, - workflow=workflow, - cache_dir=self._cache_dir(), - ) - _replace_blocks(pipe, blocks) - required = list( - dict.fromkeys(spec.name for spec in pipe._blocks.expected_components) - ) - cpu_components = [name for name in required if name not in cuda_components] - pipe.load_components(names=cpu_components, dtype=torch.bfloat16) - for name in cuda_components: - pipe.load_components( - names=[name], - dtype=torch.bfloat16, - device_map="cuda", - low_cpu_mem_usage=True, - ) - try: - return dict(pipe(**inputs, output=outputs)) - finally: - _release_pipeline(pipe) - - def _build_prepare_pipeline(self, workflow: MiniMaxH3Workflow) -> Any: - from diffusers.modular_pipelines.modular_pipeline import ModularPipeline - - pipe = ModularPipeline.from_pretrained( - self.config.model_id, - workflow=workflow, - cache_dir=self._cache_dir(), - ) - blocks = { - "t2va": ( - "denoise.no_keyframe_anchors", - "denoise.prepare_layout", - "denoise.prepare_latents", - ), - "fl2va": ( - "denoise.prepare_layout", - "denoise.prepare_condition_latents", - "denoise.prepare_latents", - "denoise.prepare_latents_fl2va", - ), - "ref2va": ( - "denoise.prepare_layout", - "denoise.prepare_condition_latents", - "denoise.prepare_latents", - "denoise.prepare_latents_ref2va", - ), - }[workflow] - _replace_blocks(pipe, blocks) - if workflow != "t2va": - pipe.load_components(names=["scheduler"], dtype=torch.bfloat16) - return pipe - - def _prepare_denoise_state( - self, - cache: MiniMaxH3PipelineCache, - conditioned: dict[str, Any], - ) -> MiniMaxH3DenoiseState: - from diffusers.modular_pipelines.modular_pipeline import PipelineState - - pipe = self._build_prepare_pipeline(cache.workflow) - state = PipelineState() - for name, value in conditioned.items(): - state.set(name, value) - state.set("generator", torch.Generator(device="cpu").manual_seed(cache.seed)) - fields = [ - "latents", - "audio_latents", - "prompt_embeds", - "position_ids", - "token_tags", - "video_indices", - "audio_indices", - "text_indices", - "num_condition_video_rows", - "num_condition_audio_rows", - "num_latent_frames", - "latent_height", - "latent_width", - ] - try: - results = pipe(state=state, output=fields) - finally: - _release_pipeline(pipe) - return MiniMaxH3DenoiseState(**results) - - def _run_native_denoise( - self, cache: MiniMaxH3PipelineCache, conditioned: dict[str, Any] - ) -> torch.Tensor: - denoise_started = time.monotonic() - _write_status(cache.latent_checkpoint, "denoising-native-flashdreams") - prepare_started = time.monotonic() - state = self._prepare_denoise_state(cache, conditioned) - cache.denoise_prepare_seconds = time.monotonic() - prepare_started - backend = "cudnn" if cache.attention == "default" else "flash" - cache.attention_backend = backend - base_transformer = cast( - MiniMaxH3TransformerConfig, self.config.diffusion_model.transformer - ) - transformer_config = replace( - base_transformer, - attention_backend=backend, - device="cuda", - execution_device="cuda", - sequential_cpu_offload=False, - ) - model_config = replace( - self.config.diffusion_model, - transformer=transformer_config, - scheduler=replace( - self.config.diffusion_model.scheduler, - num_inference_steps=cache.steps, - ), - audio_scheduler=replace( - self.config.diffusion_model.audio_scheduler, - num_inference_steps=cache.steps, - ), - seed=cache.seed, - ) - transformer_load_started = time.monotonic() - model = model_config.setup() - cache.transformer_load_seconds = time.monotonic() - transformer_load_started - try: - self._apply_lora(model.transformer, cache) - compute_started = time.monotonic() - latents = model.generate_joint(state) - if latents.is_cuda: - torch.cuda.synchronize(latents.device) - cache.denoise_compute_seconds = time.monotonic() - compute_started - finally: - cleanup_started = time.monotonic() - del model - gc.collect() - torch.cuda.empty_cache() - cache.denoise_cleanup_seconds = time.monotonic() - cleanup_started - cache.denoise_seconds = time.monotonic() - denoise_started - return latents - - def _generate_standard(self, cache: MiniMaxH3PipelineCache) -> torch.Tensor: - conditioning_started = time.monotonic() - conditioned = self._condition(cache) - cache.conditioning_seconds = time.monotonic() - conditioning_started - return self._run_native_denoise(cache, conditioned) - - def _decode_video( - self, cache: MiniMaxH3PipelineCache, latents: torch.Tensor - ) -> torch.Tensor: - from diffusers.modular_pipelines.modular_pipeline import ( - ModularPipeline, - PipelineState, - SequentialPipelineBlocks, - ) - - _write_status(cache.latent_checkpoint, "decoding-video") - pipe = ModularPipeline.from_pretrained( - self.config.model_id, - workflow=cache.workflow, - cache_dir=self._cache_dir(), - ) - video_block = pipe._blocks.sub_blocks.get("decode.video") - if video_block is None: - raise RuntimeError("MiniMax H3 workflow has no video decode block") - pipe._blocks = SequentialPipelineBlocks.from_blocks_dict({"video": video_block}) - pipe.load_components( - names=["vae", "video_processor"], - dtype={"vae": torch.float32}, - ) - pipe.vae.encoder.to_empty(device="cpu") - pipe.vae.quant_conv.to_empty(device="cpu") - pipe.vae.post_quant_conv.to("cuda") - pipe.vae.decoder.to("cuda") - first_encoder_parameter = next(pipe.vae.encoder.parameters()) - first_encoder_parameter.data = torch.empty_like( - first_encoder_parameter, device="cuda" - ) - - state = PipelineState() - state.set("latents", latents.to("cuda")) - state.set("output_type", "np") - results = pipe(state=state, output=["videos"]) - video = np.asarray(results["videos"][0]) - frames = torch.from_numpy(video).permute(0, 3, 1, 2).float().mul(2).sub(1) - _release_pipeline(pipe) - return frames.contiguous() diff --git a/integrations/minimax_h3/minimax_h3/references.py b/integrations/minimax_h3/minimax_h3/references.py deleted file mode 100644 index 613dfe09b..000000000 --- a/integrations/minimax_h3/minimax_h3/references.py +++ /dev/null @@ -1,106 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Ordered local-media references for MiniMax H3 ref2va requests.""" - -from __future__ import annotations - -from collections.abc import Sequence -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Literal, cast - -import torch - -ReferenceKind = Literal["image", "video", "audio"] - - -@dataclass(frozen=True) -class MiniMaxH3ReferenceSpec: - """One validated ``kind:path`` reference, preserving request order.""" - - kind: ReferenceKind - path: Path - - def manifest(self) -> dict[str, str | int]: - """Return the source identity used by restart-safe checkpoints.""" - stat = self.path.stat() - return { - "kind": self.kind, - "path": str(self.path), - "size": stat.st_size, - "mtime_ns": stat.st_mtime_ns, - } - - -def parse_reference_specs( - entries: Sequence[str], -) -> tuple[MiniMaxH3ReferenceSpec, ...]: - """Parse and enforce H3's documented ordered-reference limits.""" - specs: list[MiniMaxH3ReferenceSpec] = [] - for entry in entries: - kind, separator, path_value = entry.partition(":") - if not separator or kind not in {"image", "video", "audio"}: - raise ValueError( - f"invalid reference {entry!r}; expected image:path, video:path, or audio:path" - ) - path = Path(path_value).expanduser().resolve() - if not path.is_file(): - raise FileNotFoundError(f"reference file not found: {path}") - specs.append(MiniMaxH3ReferenceSpec(kind=cast(ReferenceKind, kind), path=path)) - - if not specs: - raise ValueError("ref2va requires at least one --reference") - limits = {"image": 9, "video": 3, "audio": 3} - for kind, limit in limits.items(): - count = sum(spec.kind == kind for spec in specs) - if count > limit: - raise ValueError(f"MiniMax H3 accepts at most {limit} {kind} references") - if len(specs) > 12: - raise ValueError("MiniMax H3 accepts at most 12 references in total") - if all(spec.kind == "audio" for spec in specs): - raise ValueError( - "an audio reference must be paired with an image or video reference" - ) - return tuple(specs) - - -def load_references(specs: tuple[MiniMaxH3ReferenceSpec, ...]) -> list[Any]: - """Decode references through Diffusers' official H3 media containers.""" - from diffusers.modular_pipelines.minimax_h3 import ( - MiniMaxH3AudioReference, - MiniMaxH3ImageReference, - MiniMaxH3VideoReference, - ) - - classes = { - "image": MiniMaxH3ImageReference, - "video": MiniMaxH3VideoReference, - "audio": MiniMaxH3AudioReference, - } - references: list[Any] = [classes[spec.kind].from_file(spec.path) for spec in specs] - for reference in references: - if not reference.has_audio or reference.sample_rate in {None, 32000}: - continue - import av - import numpy as np - - waveform = reference.audio.detach().cpu().to(torch.float32).numpy() - layout = "mono" if waveform.shape[0] == 1 else "stereo" - frame = av.AudioFrame.from_ndarray(waveform, format="fltp", layout=layout) - frame.sample_rate = reference.sample_rate - resampler = av.AudioResampler(format="fltp", layout=layout, rate=32000) - resampled = [*resampler.resample(frame), *resampler.resample(None)] - reference.audio = torch.from_numpy( - np.concatenate([chunk.to_ndarray() for chunk in resampled], axis=-1) - ).to(torch.float32) - reference.sample_rate = 32000 - return references - - -__all__ = [ - "MiniMaxH3ReferenceSpec", - "ReferenceKind", - "load_references", - "parse_reference_specs", -] diff --git a/integrations/minimax_h3/minimax_h3/runner.py b/integrations/minimax_h3/minimax_h3/runner.py deleted file mode 100644 index d1809ebf6..000000000 --- a/integrations/minimax_h3/minimax_h3/runner.py +++ /dev/null @@ -1,229 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""MiniMax H3 text, keyframe, and ordered-reference runners.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Annotated, Literal - -from loguru import logger -from tyro.conf import UseAppendAction - -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.infra.runner_io import runner_artifact_path, write_runner_stats -from flashdreams.runtime.video_output import Mp4VideoOutputTarget -from minimax_h3.pipeline import MiniMaxH3Pipeline, MiniMaxH3PipelineCache -from minimax_h3.references import parse_reference_specs - - -@dataclass(kw_only=True) -class MiniMaxH3RunnerConfig(RunnerConfig): - """Options shared by all released MiniMax H3 workflows.""" - - _target: type[MiniMaxH3Runner] = field(default_factory=lambda: MiniMaxH3Runner) - - prompt: str = "Animate this scene with coherent natural motion." - """Text description of the desired video motion and appearance.""" - - pixel_height: int = 768 - """Output video height in pixels.""" - - pixel_width: int = 768 - """Output video width in pixels.""" - - duration: float = 5.0 - """Requested duration before H3 frame-grid alignment.""" - - steps: int = 30 - """Number of scheduler grid points.""" - - seed: int = 42 - """CPU generator seed used by both H3 schedulers.""" - - low_ram: bool = True - """Split conditioning, denoising, and decoding into checkpointed stages.""" - - restart: bool = False - """Ignore matching stage checkpoints and regenerate the full rollout.""" - - attention: Literal["auto", "flash", "default"] = "auto" - """FlashDreams native SDPA backend selection.""" - - lora: str | None = None - """Local Musubi adapter path or Hugging Face repository ID.""" - - lora_weight_name: str | None = None - """Adapter filename override for Hugging Face repositories.""" - - lora_scale: float = 1.0 - """LoRA adapter strength.""" - - fps: int = 24 - """H3's fixed output frame rate.""" - - postprocess_output_layout: VideoTensorLayout | None = "tchw" - """Decoded H3 frame layout used by FlashDreams runtime output.""" - - -@dataclass(kw_only=True) -class MiniMaxH3T2VARunnerConfig(MiniMaxH3RunnerConfig): - """Runner config for prompt-only generation.""" - - _target: type[MiniMaxH3T2VARunner] = field( - default_factory=lambda: MiniMaxH3T2VARunner - ) - - -@dataclass(kw_only=True) -class MiniMaxH3FL2VARunnerConfig(MiniMaxH3RunnerConfig): - """Runner config for first-frame, last-frame, or dual-keyframe generation.""" - - _target: type[MiniMaxH3FL2VARunner] = field( - default_factory=lambda: MiniMaxH3FL2VARunner - ) - - image_path: Path | None = None - """Optional first-frame image path.""" - - last_image_path: Path | None = None - """Optional last-frame image path.""" - - -@dataclass(kw_only=True) -class MiniMaxH3Ref2VARunnerConfig(MiniMaxH3RunnerConfig): - """Runner config for ordered image, video, and audio references.""" - - _target: type[MiniMaxH3Ref2VARunner] = field( - default_factory=lambda: MiniMaxH3Ref2VARunner - ) - - reference: Annotated[list[str], UseAppendAction] = field(default_factory=list) - """Ordered ``image:path``, ``video:path``, or ``audio:path`` references.""" - - -class MiniMaxH3Runner(Runner[MiniMaxH3RunnerConfig, MiniMaxH3Pipeline]): - """Drive one H3 workflow and persist its video-only artifact.""" - - config: MiniMaxH3RunnerConfig - pipeline: MiniMaxH3Pipeline - - def _initialize_cache(self, output_path: Path) -> MiniMaxH3PipelineCache: - raise NotImplementedError - - def _initialize_common( - self, - output_path: Path, - *, - image_path: Path | None = None, - last_image_path: Path | None = None, - reference: list[str] | tuple[str, ...] = (), - ) -> MiniMaxH3PipelineCache: - config = self.config - return self.pipeline.initialize_cache( - prompt=config.prompt, - image_path=image_path, - last_image_path=last_image_path, - references=parse_reference_specs(reference) if reference else (), - output_path=output_path, - width=config.pixel_width, - height=config.pixel_height, - duration=config.duration, - steps=config.steps, - seed=config.seed, - low_ram=config.low_ram, - restart=config.restart, - attention=config.attention, - lora=config.lora, - lora_weight_name=config.lora_weight_name, - lora_scale=config.lora_scale, - ) - - def run(self) -> None: - """Generate the single H3 step and write a video-only MP4.""" - config = self.config - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - cache = self._initialize_cache(video_path) - output_stream = self.create_video_output_stream(fps=config.fps) - output_target = Mp4VideoOutputTarget( - output_path=video_path, - fps=config.fps, - output_layout=output_stream.output_layout, - enabled=self.is_rank_zero, - ) - output_target.open() - frames = self.pipeline.generate(0, cache) - metrics = self.pipeline.finalize(0, cache) - output_target.write( - output_stream.process(frames, autoregressive_index=0, metrics=metrics) - ) - tail = output_stream.finish() - if tail is not None: - output_target.write(tail) - artifacts = output_target.close() - if not artifacts: - return - self.pipeline.mark_complete(cache) - video_artifact = artifacts[0] - logger.info( - "[{}] wrote {} video to {}", - config.runner_name, - tuple(frames.shape), - Path(video_artifact.uri).resolve(), - ) - stats_history = video_artifact.metadata["stats_history"] - if stats_history: - stats_path = write_runner_stats( - config.output_dir, - config.runner_name, - list(stats_history), - ) - logger.info( - "[{}] wrote stats to {}", config.runner_name, stats_path.resolve() - ) - - -class MiniMaxH3T2VARunner(MiniMaxH3Runner): - """Prompt-only H3 runner.""" - - config: MiniMaxH3T2VARunnerConfig - - def _initialize_cache(self, output_path: Path) -> MiniMaxH3PipelineCache: - return self._initialize_common(output_path) - - -class MiniMaxH3FL2VARunner(MiniMaxH3Runner): - """First-frame, last-frame, or dual-keyframe H3 runner.""" - - config: MiniMaxH3FL2VARunnerConfig - - def _initialize_cache(self, output_path: Path) -> MiniMaxH3PipelineCache: - return self._initialize_common( - output_path, - image_path=self.config.image_path, - last_image_path=self.config.last_image_path, - ) - - -class MiniMaxH3Ref2VARunner(MiniMaxH3Runner): - """Ordered-reference H3 runner.""" - - config: MiniMaxH3Ref2VARunnerConfig - - def _initialize_cache(self, output_path: Path) -> MiniMaxH3PipelineCache: - return self._initialize_common(output_path, reference=self.config.reference) - - -__all__ = [ - "MiniMaxH3FL2VARunner", - "MiniMaxH3FL2VARunnerConfig", - "MiniMaxH3Ref2VARunner", - "MiniMaxH3Ref2VARunnerConfig", - "MiniMaxH3Runner", - "MiniMaxH3RunnerConfig", - "MiniMaxH3T2VARunner", - "MiniMaxH3T2VARunnerConfig", -] diff --git a/integrations/minimax_h3/minimax_h3/scheduler.py b/integrations/minimax_h3/minimax_h3/scheduler.py deleted file mode 100644 index 6a5b01ac0..000000000 --- a/integrations/minimax_h3/minimax_h3/scheduler.py +++ /dev/null @@ -1,112 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""FlashDreams scheduler for MiniMax H3's data-ward velocity.""" - -from __future__ import annotations - -from dataclasses import dataclass, field - -import torch -from torch import Tensor - -from flashdreams.infra.diffusion.scheduler import ( - FlowPredictor, - Scheduler, - SchedulerConfig, -) - - -@dataclass(kw_only=True) -class MiniMaxH3SchedulerConfig(SchedulerConfig): - """Configuration for one of H3's modality-specific schedules.""" - - _target: type[MiniMaxH3Scheduler] = field( - default_factory=lambda: MiniMaxH3Scheduler - ) - num_inference_steps: int = 30 - shift: float = 12.0 - - -class MiniMaxH3Scheduler(Scheduler): - """Rectified-flow Euler schedule used by the released H3 checkpoint.""" - - config: MiniMaxH3SchedulerConfig - - def __init__(self, config: MiniMaxH3SchedulerConfig) -> None: - super().__init__(config) - if config.num_inference_steps < 2: - raise ValueError("num_inference_steps must be at least 2") - if config.shift <= 0: - raise ValueError("shift must be positive") - - def schedule(self, device: torch.device | str) -> tuple[Tensor, Tensor]: - """Return the shifted sigma grid and its H3 timesteps.""" - base = torch.linspace( - 1.0, 0.0, self.config.num_inference_steps, dtype=torch.float32 - ) - shift = self.config.shift - sigmas = torch.unique_consecutive(shift * base / (1 + (shift - 1) * base)) - sigmas = sigmas.to(device) - return sigmas, 1.0 - sigmas[:-1] - - @staticmethod - def step( - sample: Tensor, - flow: Tensor, - timestep: Tensor, - sigma: Tensor, - sigma_next: Tensor, - ) -> Tensor: - """Take one deterministic data-ward Euler step.""" - sigma_from_timestep = 1 - timestep.to(device=sample.device, dtype=sample.dtype) - denoised = sample + sigma_from_timestep * flow - compute_dtype = ( - torch.float32 - if sample.dtype in (torch.float16, torch.bfloat16) - else sample.dtype - ) - ratio = sigma_next.to(sample.device, compute_dtype) / sigma.to( - sample.device, compute_dtype - ) - previous = ratio * sample.to(compute_dtype) + (1 - ratio) * denoised.to( - compute_dtype - ) - return previous.to(sample.dtype) - - @torch.no_grad() - def sample( - self, - initial_noise: Tensor, - predict_flow: FlowPredictor, - rng: torch.Generator | None = None, - ) -> Tensor: - """Denoise one stream with the H3 schedule.""" - del rng - sigmas, timesteps = self.schedule(initial_noise.device) - sample = initial_noise - for index, timestep in enumerate(timesteps): - flow = predict_flow(sample, timestep) - sample = self.step(sample, flow, timestep, sigmas[index], sigmas[index + 1]) - return sample - - def add_noise( - self, - clean_input: Tensor, - timestep: Tensor, - rng: torch.Generator | None = None, - ) -> Tensor: - """Mix clean input with Gaussian noise under H3's time convention.""" - noise = torch.randn( - clean_input.shape, - dtype=clean_input.dtype, - device=clean_input.device, - generator=rng, - ) - time = timestep.to(clean_input.device, clean_input.dtype) - while time.ndim < clean_input.ndim: - time = time.unsqueeze(-1) - return time * clean_input + (1 - time) * noise - - -__all__ = ["MiniMaxH3Scheduler", "MiniMaxH3SchedulerConfig"] diff --git a/integrations/minimax_h3/pyproject.toml b/integrations/minimax_h3/pyproject.toml deleted file mode 100644 index 91c49711a..000000000 --- a/integrations/minimax_h3/pyproject.toml +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -[build-system] -requires = ["setuptools>=69", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "flashdreams-minimax-h3" -version = "0.1.0" -description = "MiniMax H3 video generation workflows for the FlashDreams runtime." -requires-python = ">=3.10" -dependencies = [ - "accelerate>=1.12", - "av>=16", - "diffusers @ git+https://github.com/huggingface/diffusers.git@175fe6b2419a01db9c2ceabd01ec37d2c0305fc2", - "flashdreams", - "huggingface-hub>=0.33", - "numpy>=1.24,<2.5", - "pillow>=11", - "safetensors>=0.4", - "sentencepiece>=0.2", - "torch>=2.9", - "transformers @ git+https://github.com/huggingface/transformers.git@d1123114da1ab4395198146f4f84dae7fe8b693e", -] - -[tool.uv.sources] -flashdreams = { workspace = true } - -[project.optional-dependencies] -dev = ["pytest>=8.0", "tomli>=2.0"] - -[project.entry-points."flashdreams.runner_configs"] -"minimax-h3-t2va" = "minimax_h3.config:RUNNER_MINIMAX_H3_T2VA" -"minimax-h3-fl2va" = "minimax_h3.config:RUNNER_MINIMAX_H3_FL2VA" -"minimax-h3-ref2va" = "minimax_h3.config:RUNNER_MINIMAX_H3_REF2VA" - -[tool.setuptools.packages.find] -include = ["minimax_h3*"] -exclude = ["tests"] - -[tool.uv] -managed = true diff --git a/integrations/minimax_h3/tests/test_smoke.py b/integrations/minimax_h3/tests/test_smoke.py deleted file mode 100644 index 31b7d288e..000000000 --- a/integrations/minimax_h3/tests/test_smoke.py +++ /dev/null @@ -1,437 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""CPU smoke tests for the MiniMax H3 runner plugin.""" - -from __future__ import annotations - -import sys -import threading -from pathlib import Path -from typing import Any, cast - -import pytest -import tomli as tomllib -import torch -from flashdreams.infra.diffusion.transformer import Transformer -from flashdreams.infra.runner import RunnerConfig -from minimax_h3 import config as config_mod -from minimax_h3 import pipeline as h3_pipeline -from minimax_h3.config import ( - PIPELINE_MINIMAX_H3_FL2VA, - PIPELINE_MINIMAX_H3_REF2VA, - PIPELINE_MINIMAX_H3_T2VA, - RUNNER_CONFIGS, - RUNNER_MINIMAX_H3_FL2VA, - RUNNER_MINIMAX_H3_REF2VA, - RUNNER_MINIMAX_H3_T2VA, -) -from minimax_h3.constants import align_num_frames, validate_canvas -from minimax_h3.lora import convert_musubi_lora -from minimax_h3.model import MiniMaxH3DenoiseState, MiniMaxH3DiffusionModel -from minimax_h3.pipeline import MiniMaxH3Pipeline -from minimax_h3.references import parse_reference_specs -from minimax_h3.runner import ( - MiniMaxH3FL2VARunner, - MiniMaxH3Ref2VARunner, - MiniMaxH3RunnerConfig, - MiniMaxH3T2VARunner, -) -from minimax_h3.scheduler import MiniMaxH3SchedulerConfig -from minimax_h3.transformer import MiniMaxH3TransformerConfig - -pytestmark = pytest.mark.ci_cpu - -ENTRY_POINT_GROUP = "flashdreams.runner_configs" - - -def test_runners_dict_is_non_empty() -> None: - """Plugin must expose at least one runner.""" - assert RUNNER_CONFIGS, "RUNNER_CONFIGS is empty" - - -def test_runner_name_mirrors_pipeline_name() -> None: - """Runner names must match pipeline names for CLI discovery.""" - drifted = { - slug: (config.runner_name, config.pipeline.name) - for slug, config in RUNNER_CONFIGS.items() - if config.runner_name != config.pipeline.name - } - assert not drifted, f"runner_name != pipeline.name: {drifted}" - - -def test_runners_have_descriptions() -> None: - """Every registered runner must have a CLI description.""" - empty = [ - slug - for slug, config in RUNNER_CONFIGS.items() - if not config.description.strip() - ] - assert not empty, f"runners missing description: {empty}" - - -def test_entry_points_match_module_literals() -> None: - """Package entry points must resolve to every registered runner literal.""" - pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" - with pyproject.open("rb") as handle: - metadata = tomllib.load(handle) - entries = metadata["project"]["entry-points"][ENTRY_POINT_GROUP] - assert set(entries) == set(RUNNER_CONFIGS) - - for slug, target in entries.items(): - module_name, attribute = target.split(":", 1) - assert module_name == "minimax_h3.config" - config = cast(RunnerConfig, getattr(config_mod, attribute)) - assert config.runner_name == slug - - -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="entry-point discovery test relies on importlib.metadata 3.10+ shape", -) -def test_entry_points_discoverable_when_installed() -> None: - """Installed plugin entry points must expose every registered runner.""" - from importlib.metadata import entry_points - - entries = entry_points(group=ENTRY_POINT_GROUP) - discovered = { - entry.name for entry in entries if entry.value.startswith("minimax_h3.") - } - if not discovered: - pytest.skip("plugin not installed; run uv sync from the repository root") - assert discovered == set(RUNNER_CONFIGS) - - -def test_pipeline_config_constructs_without_loading_weights() -> None: - """Construct every runtime pipeline without network or checkpoint access.""" - pipelines = { - config.workflow: config.setup() - for config in ( - PIPELINE_MINIMAX_H3_T2VA, - PIPELINE_MINIMAX_H3_FL2VA, - PIPELINE_MINIMAX_H3_REF2VA, - ) - } - assert set(pipelines) == {"t2va", "fl2va", "ref2va"} - assert all( - isinstance(pipeline, MiniMaxH3Pipeline) for pipeline in pipelines.values() - ) - assert all( - pipeline.config.model_id == "MiniMaxAI/MiniMax-H3" - for pipeline in pipelines.values() - ) - assert RUNNER_MINIMAX_H3_T2VA._target is MiniMaxH3T2VARunner - assert RUNNER_MINIMAX_H3_FL2VA._target is MiniMaxH3FL2VARunner - assert RUNNER_MINIMAX_H3_REF2VA._target is MiniMaxH3Ref2VARunner - for config in ( - PIPELINE_MINIMAX_H3_T2VA, - PIPELINE_MINIMAX_H3_FL2VA, - PIPELINE_MINIMAX_H3_REF2VA, - ): - assert issubclass(config.diffusion_model.transformer._target, Transformer) - - -def test_low_ram_is_an_explicit_default_flag() -> None: - """Default to crash-safe staging without enabling a third-party LoRA.""" - for runner in RUNNER_CONFIGS.values(): - assert isinstance(runner, MiniMaxH3RunnerConfig) - assert runner.low_ram is True - assert runner.lora is None - - -def test_native_bf16_gpu_path_is_default() -> None: - """Keep the quality-preserving low-host-RAM path on the accelerator.""" - transformer = PIPELINE_MINIMAX_H3_FL2VA.diffusion_model.transformer - assert isinstance(transformer, MiniMaxH3TransformerConfig) - assert transformer.device == "cuda" - assert transformer.sequential_cpu_offload is False - - -def test_musubi_lora_conversion_targets_native_layers(tmp_path: Path) -> None: - """Convert all H3 block targets without introducing a default adapter.""" - from safetensors.torch import load_file, save_file - - source = tmp_path / "adapter.safetensors" - tensors: dict[str, torch.Tensor] = {} - for block in range(50): - for module in ("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2"): - prefix = f"lora_unet_blocks_{block}_{module}" - tensors[f"{prefix}.alpha"] = torch.tensor(2.0) - tensors[f"{prefix}.lora_down.weight"] = torch.ones(2, 3) - out_features = 6 if module == "attn_qkv_proj" else 4 - tensors[f"{prefix}.lora_up.weight"] = torch.ones(out_features, 2) - save_file(tensors, source) - - converted = load_file( - convert_musubi_lora(source, tmp_path / "converted.safetensors") - ) - assert len(converted) == 600 - assert "transformer.transformer_blocks.0.attn.to_q.lora_A.weight" in converted - assert "transformer.transformer_blocks.0.attn.to_v.lora_B.weight" in converted - assert "transformer.transformer_blocks.49.ff.net.2.lora_B.weight" in converted - - -def test_duration_and_canvas_contracts() -> None: - """Align five seconds and reject non-H3 canvas dimensions.""" - assert align_num_frames(5.0) == 124 - validate_canvas(576, 768) - with pytest.raises(ValueError, match="multiples of 32"): - validate_canvas(577, 768) - - -def test_runtime_cache_uses_stage_specific_checkpoints(tmp_path: Path) -> None: - """Derive conditioning and denoised checkpoints beside the output.""" - image = tmp_path / "image.png" - image.write_bytes(b"test") - pipeline = PIPELINE_MINIMAX_H3_FL2VA.setup() - cache = pipeline.initialize_cache( - prompt="animate", - image_path=image, - last_image_path=None, - references=(), - output_path=tmp_path / "out.mp4", - width=576, - height=768, - duration=5.0, - steps=30, - seed=42, - low_ram=True, - restart=False, - attention="auto", - lora=None, - lora_weight_name=None, - lora_scale=1.0, - ) - assert cache.conditioning_checkpoint.name == "out.mp4.conditioning.safetensors" - assert cache.latent_checkpoint.name == "out.mp4.latents.safetensors" - - -def test_generate_preserves_non_overlapping_stage_metrics( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - """Keep conditioning separate from native denoise and checkpoint timings.""" - image = tmp_path / "image.png" - image.write_bytes(b"test") - pipeline = PIPELINE_MINIMAX_H3_FL2VA.setup() - cache = pipeline.initialize_cache( - prompt="animate", - image_path=image, - last_image_path=None, - references=(), - output_path=tmp_path / "out.mp4", - width=576, - height=768, - duration=5.0, - steps=30, - seed=42, - low_ram=True, - restart=True, - attention="auto", - lora=None, - lora_weight_name=None, - lora_scale=1.0, - ) - - def generate_low_ram(_: Any) -> torch.Tensor: - cache.conditioning_seconds = 10.0 - cache.denoise_seconds = 20.0 - cache.denoise_prepare_seconds = 1.0 - cache.transformer_load_seconds = 4.0 - cache.denoise_compute_seconds = 14.0 - cache.denoise_cleanup_seconds = 1.0 - return torch.zeros(1) - - monkeypatch.setattr(pipeline, "_generate_low_ram", generate_low_ram) - monkeypatch.setattr( - pipeline, - "_decode_video", - lambda _cache, _latents: torch.zeros(1, 3, 1, 1), - ) - monkeypatch.setattr(h3_pipeline, "_save_latents", lambda *_args: None) - monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", lambda: None) - monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) - - pipeline.generate(0, cache) - metrics = pipeline.finalize(0, cache) - - assert cache.latent_checkpoint_future is None - assert metrics["conditioning_seconds"] == 10.0 - assert metrics["denoise_seconds"] == 20.0 - assert metrics["denoise_prepare_seconds"] == 1.0 - assert metrics["transformer_load_seconds"] == 4.0 - assert metrics["denoise_compute_seconds"] == 14.0 - assert metrics["denoise_cleanup_seconds"] == 1.0 - assert metrics["latent_checkpoint_seconds"] >= 0.0 - - -def test_generate_overlaps_latent_checkpoint_with_decode( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - """Do not put recovery-checkpoint I/O on generate's critical path.""" - image = tmp_path / "image.png" - image.write_bytes(b"test") - pipeline = PIPELINE_MINIMAX_H3_FL2VA.setup() - cache = pipeline.initialize_cache( - prompt="animate", - image_path=image, - last_image_path=None, - references=(), - output_path=tmp_path / "out.mp4", - width=576, - height=768, - duration=5.0, - steps=30, - seed=42, - low_ram=True, - restart=True, - attention="auto", - lora=None, - lora_weight_name=None, - lora_scale=1.0, - ) - checkpoint_started = threading.Event() - release_checkpoint = threading.Event() - - def save_latents(*_args: Any) -> None: - checkpoint_started.set() - assert release_checkpoint.wait(timeout=5) - - def decode_video(*_args: Any) -> torch.Tensor: - assert checkpoint_started.wait(timeout=5) - return torch.zeros(1, 3, 1, 1) - - monkeypatch.setattr(pipeline, "_generate_low_ram", lambda _: torch.zeros(1)) - monkeypatch.setattr(pipeline, "_decode_video", decode_video) - monkeypatch.setattr(h3_pipeline, "_save_latents", save_latents) - monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", lambda: None) - monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0) - - try: - pipeline.generate(0, cache) - assert cache.latent_checkpoint_future is not None - assert not cache.latent_checkpoint_future.done() - finally: - release_checkpoint.set() - - pipeline.finalize(0, cache) - assert cache.latent_checkpoint_future is None - cache.output_path.write_bytes(b"mp4") - pipeline.mark_complete(cache) - - -def test_reference_parser_preserves_order_and_enforces_limits(tmp_path: Path) -> None: - """Keep semantic reference order while rejecting unsupported requests.""" - image = tmp_path / "subject.png" - video = tmp_path / "motion.mp4" - audio = tmp_path / "voice.wav" - for path in (image, video, audio): - path.write_bytes(b"test") - parsed = parse_reference_specs( - (f"image:{image}", f"audio:{audio}", f"video:{video}") - ) - assert [reference.kind for reference in parsed] == ["image", "audio", "video"] - with pytest.raises(ValueError, match="paired"): - parse_reference_specs((f"audio:{audio}",)) - with pytest.raises(ValueError, match="at most 3 video"): - parse_reference_specs(tuple(f"video:{video}" for _ in range(4))) - - -def test_registered_workflows_validate_their_inputs(tmp_path: Path) -> None: - """Reject cross-workflow media instead of silently selecting another model.""" - image = tmp_path / "image.png" - image.write_bytes(b"test") - common = { - "prompt": "animate", - "output_path": tmp_path / "out.mp4", - "width": 512, - "height": 768, - "duration": 5.0, - "steps": 30, - "seed": 42, - "low_ram": True, - "restart": False, - "attention": "auto", - "lora": None, - "lora_weight_name": None, - "lora_scale": 1.0, - } - t2va = PIPELINE_MINIMAX_H3_T2VA.setup() - cache = t2va.initialize_cache( - image_path=None, last_image_path=None, references=(), **common - ) - assert cache.workflow == "t2va" - fl2va = PIPELINE_MINIMAX_H3_FL2VA.setup() - cache = fl2va.initialize_cache( - image_path=None, last_image_path=image, references=(), **common - ) - assert cache.workflow == "fl2va" - with pytest.raises(ValueError, match="requires --image-path"): - fl2va.initialize_cache( - image_path=None, last_image_path=None, references=(), **common - ) - - -def test_native_scheduler_matches_official_h3_euler() -> None: - """Match the released H3 schedule and data-ward Euler update exactly.""" - from diffusers.schedulers.scheduling_minimax_h3 import ( - MiniMaxH3Scheduler as OfficialScheduler, - ) - - official: Any = OfficialScheduler(shift=12.0) - official.set_timesteps(7, device="cpu") - native = MiniMaxH3SchedulerConfig(num_inference_steps=7, shift=12.0).setup() - sigmas, timesteps = native.schedule("cpu") - torch.testing.assert_close(sigmas, official.sigmas) - torch.testing.assert_close(timesteps, official.timesteps) - - sample = torch.randn(2, 3) - flow = torch.randn_like(sample) - expected = official.step(flow, official.timesteps[0], sample).prev_sample - actual = native.step(sample, flow, timesteps[0], sigmas[0], sigmas[1]) - torch.testing.assert_close(actual, expected) - - -def test_row_timestep_plan_preserves_device_and_conditioning_levels( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Build packed row timesteps without materializing accelerator scalars.""" - state = MiniMaxH3DenoiseState( - latents=torch.empty(0), - audio_latents=torch.empty(0), - prompt_embeds=torch.empty(0), - position_ids=torch.empty(0), - token_tags=torch.empty(0), - video_indices=torch.tensor([0, 1, 4]), - audio_indices=torch.tensor([2, 3]), - text_indices=torch.tensor([5]), - num_condition_video_rows=1, - num_condition_audio_rows=1, - num_latent_frames=0, - latent_height=0, - latent_width=0, - ) - - def reject_scalar_conversion(_tensor: torch.Tensor) -> float: - raise AssertionError("row timestep construction converted a tensor to float") - - with monkeypatch.context() as patch: - patch.setattr(torch.Tensor, "__float__", reject_scalar_conversion) - timesteps, indices = MiniMaxH3DiffusionModel._row_timesteps( - state, torch.tensor(0.5), torch.tensor(0.25) - ) - - assert timesteps.device == state.video_indices.device - torch.testing.assert_close(timesteps, torch.tensor([0.25, 0.5, 0.999, 1.0])) - torch.testing.assert_close(indices, torch.tensor([2, 1, 3, 0, 1, 1])) diff --git a/integrations/minimax_h3/tests/test_transformer_cuda.py b/integrations/minimax_h3/tests/test_transformer_cuda.py deleted file mode 100644 index e7285fe79..000000000 --- a/integrations/minimax_h3/tests/test_transformer_cuda.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""CUDA parity tests for the MiniMax H3 transformer.""" - -from __future__ import annotations - -from typing import Any - -import pytest -import torch -from diffusers.models.transformers.transformer_minimax_h3 import ( - MiniMaxH3Transformer3DModel, -) -from minimax_h3.transformer import MiniMaxH3TransformerConfig - -pytestmark = pytest.mark.ci_gpu - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_native_transformer_matches_official_h3_forward() -> None: - """Prove state-dict and numerical compatibility on a tiny CUDA model.""" - architecture: dict[str, Any] = { - "num_attention_heads": 2, - "attention_head_dim": 16, - "hidden_size": 16, - "num_layers": 2, - "num_refiner_layers": 1, - "ffn_dim": 32, - "in_channels": 2, - "audio_in_channels": 4, - "patch_size": (1, 1, 1), - "text_dim": 10, - "freq_dim": 8, - "time_embed_hidden_dim": 16, - "time_embed_dim": 8, - "rope_freq_dim": 2, - } - device = torch.device("cuda") - torch.manual_seed(1) - official: Any = MiniMaxH3Transformer3DModel(**architecture) - official.to(device) - native = MiniMaxH3TransformerConfig( - checkpoint_path=None, - device="cuda", - execution_device="cuda", - sequential_cpu_offload=False, - dtype=torch.float32, - attention_backend="math", - **architecture, - ).setup() - native.load_state_dict(official.state_dict(), strict=True) - inputs = { - "hidden_states": torch.randn(1, 4, 2, device=device), - "audio_hidden_states": torch.randn(1, 2, 4, device=device), - "encoder_hidden_states": torch.randn(1, 3, 10, device=device), - "timestep": torch.tensor([0.1, 0.5, 0.999], device=device), - "timestep_indices": torch.tensor([1, 1, 1, 2, 0, 1, 1, 0, 2], device=device), - "token_tags": torch.tensor([1, 1, 1, 0, 0, 2, 2, 0, 0], device=device), - "position_ids": torch.randn(9, 3, device=device), - "video_indices": torch.tensor([3, 4, 7, 8], device=device), - "audio_indices": torch.tensor([5, 6], device=device), - "text_indices": torch.tensor([0, 1, 2], device=device), - } - with ( - torch.no_grad(), - torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH), - ): - expected = official(**inputs, return_dict=False) - actual = native.forward_joint(**inputs) - for native_output, official_output in zip(actual, expected, strict=True): - torch.testing.assert_close(native_output, official_output) diff --git a/integrations_v2/README.md b/integrations_v2/README.md index 049e11003..4d6c5a5d7 100644 --- a/integrations_v2/README.md +++ b/integrations_v2/README.md @@ -35,6 +35,8 @@ follows is already done for you. `apps/v2v` video-to-video application. - `null_model` — not an application. A v1 pipeline the framework tests use as a fixture. +- `minimax_h3` — native joint video/audio inference with video-only T2VA, + first/last-keyframe, and ordered-reference bindings to `apps/t2v`. ## The layout diff --git a/integrations_v2/minimax_h3/README.md b/integrations_v2/minimax_h3/README.md new file mode 100644 index 000000000..804963c4f --- /dev/null +++ b/integrations_v2/minimax_h3/README.md @@ -0,0 +1,80 @@ + + + +# MiniMax H3 + +Native, video-only MiniMax H3 inference through FlashDreams v2. The three +workflows share FlashDreams' T2V application, synchronized diffusion sampler, +accelerated attention, and scoped checkpoint loader. There is no Diffusers +runtime dependency or fallback pipeline. + +## Run + +```bash +uv sync --package flashdreams-minimax-h3 --extra dev --inexact +uv run --no-sync flashdreams-run-v2 t2v-minimax-h3-t2va --output-path clip.mp4 -- \ + --prompt "A cat surfing" --duration 5 --steps 30 --seed 42 +uv run --no-sync flashdreams-run-v2 t2v-minimax-h3-fl2va --output-path clip.mp4 -- \ + --prompt "The camera moves through the scene" --image-path first.png --last-image-path last.png +uv run --no-sync flashdreams-run-v2 t2v-minimax-h3-ref2va --output-path clip.mp4 -- \ + --prompt "A cinematic scene" --reference image:subject.png --reference audio:reference.wav +``` + +Runtime options precede `--`; H3 options follow it. Use `-- --help` without +loading any weights. Resolution defaults to 768×768; runtime +`--pixel-width`/`--pixel-height` must be multiples of 32 and have aspect ratio +between 1:4 and 4:1. FPS is fixed at 24. Duration is 5–15 seconds, rounded +up to H3's `17n+5` frame grid, without exceeding 15 seconds. One block generates +the complete video; 30 scheduler points mean 29 joint model predictions. + +FL2VA accepts first, last, or both keyframes. REF2VA keeps reference order and +supports images, videos, and audio (at least one visual reference). Audio +conditioning and audio denoising remain active, but generated audio is not +decoded or muxed into the output. `--mode webrtc` uses the shared prompt UI. + +LoRA options are `--lora PATH_OR_REPO`, `--lora-weight-name FILE`, and +`--lora-scale NUMBER`. Musubi conversion is retained; a new network is loaded +for each request, so LoRA weights do not accumulate across resets. + +## Memory and acceleration + +Stage-scoped loading is always enabled: conditioning, denoising, and decoding +do not retain each other's weights. This replaces the old `--low-ram` switch. +Only requested codec subtrees are constructed and loaded. The joint transformer +still needs to fit on one GPU; blockwise CPU offload and distributed inference +are not implemented. A100 80 GB is the initial validation target, not a measured +guarantee that every resolution/duration fits. + +Default attention is FlashDreams `OptimizedMultiHeadAttention`, native BF16 +SDPA, no quantization, TMA, or duplicated QKV weights. `--attention torch` selects +FlashDreams' reference implementation. `--fuse-qkv` opts into additional fused +weight storage (roughly 11 GiB for the transformer). `--compile` uses the shared +compile helper and is opt-in. These options are **unvalidated performance +candidates**, not published speedups. CUDA graphs and quantization are deferred. + +The video codec retains FP32 weights, FP16 decode autocast, reference tiling and +overlap, and precision-preserving Torch attention. Audio references use FP32 +encoding. Checkpoint assets are pinned to +`MiniMaxAI/MiniMax-H3@42ed227ee7df40d41602854ae760620d6eb651fe`; `--model-id` accepts +a compatible local snapshot or repository, and `--revision` overrides the pin. +Transformers/Accelerate remain dependencies for headless Qwen3-VL loading; they +do not orchestrate H3 diffusion. + +## Migration and checks + +The old `flashdreams-run minimax-h3-*` commands and recovery checkpoint files +are not used. Replace them with the three v2 commands above. Existing output +and recovery files are left untouched. No background checkpoint writer exists. + +```bash +PYTHONPATH=flashdreams:apps/t2v:integrations_v2 .venv/bin/python -m pytest \ + integrations_v2/minimax_h3/tests -m ci_cpu +``` + +Native CPU checks require no checkpoint downloads. Optional reference checks +use an already-installed pinned Diffusers oracle and cached checkpoint headers; +they are skipped when unavailable. Diffusers is not installed by this package. +GPU tests are separate (`-m ci_gpu`); real-checkpoint generation/parity must be +requested explicitly. CPU correctness checks do not establish GPU video parity +or a speedup. Stage timings and peak allocated GPU memory are reported through +the shared v2 `--stats-path` output. diff --git a/integrations_v2/minimax_h3/__init__.py b/integrations_v2/minimax_h3/__init__.py new file mode 100644 index 000000000..4fd63b3c4 --- /dev/null +++ b/integrations_v2/minimax_h3/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native MiniMax H3 video generation for the FlashDreams v2 runtime.""" diff --git a/integrations_v2/minimax_h3/apps/__init__.py b/integrations_v2/minimax_h3/apps/__init__.py new file mode 100644 index 000000000..c93a64165 --- /dev/null +++ b/integrations_v2/minimax_h3/apps/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 application bindings.""" diff --git a/integrations_v2/minimax_h3/apps/t2v/README.md b/integrations_v2/minimax_h3/apps/t2v/README.md new file mode 100644 index 000000000..1eb97f90d --- /dev/null +++ b/integrations_v2/minimax_h3/apps/t2v/README.md @@ -0,0 +1,5 @@ +# MiniMax H3 T2V application + +Use `flashdreams-run-v2 t2v-minimax-h3-t2va`, +`t2v-minimax-h3-fl2va`, or `t2v-minimax-h3-ref2va`. +See the [integration guide](../../README.md) for inputs, migration, and validation. diff --git a/integrations_v2/minimax_h3/apps/t2v/__init__.py b/integrations_v2/minimax_h3/apps/t2v/__init__.py new file mode 100644 index 000000000..4afedde68 --- /dev/null +++ b/integrations_v2/minimax_h3/apps/t2v/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 bindings for the shared T2V application.""" diff --git a/integrations_v2/minimax_h3/apps/t2v/adapter.py b/integrations_v2/minimax_h3/apps/t2v/adapter.py new file mode 100644 index 000000000..41c2ce2fd --- /dev/null +++ b/integrations_v2/minimax_h3/apps/t2v/adapter.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 workflows over the shared FlashDreams v2 T2V application.""" + +import argparse +from dataclasses import replace +from pathlib import Path +from typing import Any + +from t2v import T2VApplication, T2VApplicationDefaults + +from flashdreams.accelerated.multi_head_attention.optimized import QKVFusionOption +from flashdreams.api_v2.application import IApplication +from flashdreams.infra.config import derive_config +from flashdreams.runtime_v2.session_desc import SessionDesc +from minimax_h3.config import ( + PIPELINE_MINIMAX_H3_FL2VA, + PIPELINE_MINIMAX_H3_REF2VA, + PIPELINE_MINIMAX_H3_T2VA, +) +from minimax_h3.impl.constants import FPS, align_num_frames, validate_canvas +from minimax_h3.impl.pipeline import MiniMaxH3PipelineConfig +from minimax_h3.impl.references import parse_reference_specs + + +class MiniMaxH3Application(T2VApplication): + """One-block joint audio/video inference with video-only v2 output.""" + + def __init__( + self, pipeline_config: MiniMaxH3PipelineConfig = PIPELINE_MINIMAX_H3_T2VA + ) -> None: + super().__init__( + defaults=T2VApplicationDefaults( + pipeline_config=pipeline_config, + total_blocks=1, + pixel_width=768, + pixel_height=768, + fps=FPS, + ) + ) + self._request_inputs: dict[str, Any] = {} + + def _configure_argument_parser(self, parser: argparse.ArgumentParser) -> None: + parser.add_argument("--duration", type=float, default=5.0) + parser.add_argument( + "--steps", + type=int, + default=30, + help="Scheduler grid points (30 means 29 joint predictions).", + ) + parser.add_argument("--image-path", type=Path) + parser.add_argument("--last-image-path", type=Path) + parser.add_argument( + "--reference", + action="append", + default=[], + help="Ordered image:path, video:path, or audio:path input.", + ) + parser.add_argument("--lora") + parser.add_argument("--lora-weight-name") + parser.add_argument("--lora-scale", type=float, default=1.0) + parser.add_argument( + "--attention", choices=("optimized", "torch"), default="optimized" + ) + parser.add_argument( + "--fuse-qkv", + action="store_true", + help="Opt in to extra fused-QKV weight storage; increases peak memory.", + ) + parser.add_argument("--model-id", default=self.pipeline_config.model_id) + parser.add_argument("--revision", default=self.pipeline_config.revision) + + def _apply_parsed_arguments(self, args: argparse.Namespace) -> None: + align_num_frames(args.duration) + if args.steps < 2: + raise ValueError("--steps must be at least 2 scheduler points") + if not 0 <= args.lora_scale <= 4: + raise ValueError("--lora-scale must be between 0 and 4") + references = parse_reference_specs(args.reference) if args.reference else () + workflow = self.pipeline_config.workflow + if workflow == "t2va" and ( + args.image_path or args.last_image_path or references + ): + raise ValueError("t2va does not accept keyframes or references") + if workflow == "fl2va" and ( + not (args.image_path or args.last_image_path) or references + ): + raise ValueError("fl2va requires first/last keyframes and no references") + if workflow == "ref2va" and ( + not references or args.image_path or args.last_image_path + ): + raise ValueError("ref2va requires ordered references and no keyframes") + for path in (args.image_path, args.last_image_path): + if path is not None and not path.is_file(): + raise FileNotFoundError(path) + self._request_inputs = dict( + duration=args.duration, + image_path=args.image_path, + last_image_path=args.last_image_path, + references=references, + lora=args.lora, + lora_weight_name=args.lora_weight_name, + lora_scale=args.lora_scale, + ) + config = self.pipeline_config + optimized = replace( + config.transformer.optimized_impl, + qkv_fusion_option=QKVFusionOption.FULL + if args.fuse_qkv + else QKVFusionOption.NONE, + ) + self._pipeline_config = replace( + config, + model_id=args.model_id, + revision=args.revision, + transformer=replace( + config.transformer, + attention_backend=args.attention, + optimized_impl=optimized, + ), + scheduler=replace(config.scheduler, num_inference_steps=args.steps), + audio_scheduler=replace( + config.audio_scheduler, num_inference_steps=args.steps + ), + ) + + def _cache_initialization_kwargs(self, session_desc: SessionDesc) -> dict[str, Any]: + return dict(self._request_inputs) + + def _validate_total_blocks(self, total_blocks: int) -> None: + if total_blocks != 1: + raise ValueError("MiniMax H3 generates its complete clip in one block") + + def _validate_frame_size(self, session_desc: SessionDesc, pipeline: Any) -> None: + validate_canvas(session_desc.video_width, session_desc.video_height) + if session_desc.frames_per_second_for_step != FPS: + raise ValueError("MiniMax H3 requires 24 fps") + + def _apply_compile_override(self, pipeline_config: Any, enabled: bool) -> Any: + return derive_config(pipeline_config, compile_network=enabled) + + def _apply_seed_override(self, pipeline_config: Any, seed: int) -> Any: + return derive_config(pipeline_config, seed=seed) + + +def create_app() -> IApplication: + """Create the prompt-only H3 application without loading weights.""" + return MiniMaxH3Application() + + +def create_app_fl2va() -> IApplication: + """Create the first/last-keyframe H3 application.""" + return MiniMaxH3Application(PIPELINE_MINIMAX_H3_FL2VA) + + +def create_app_ref2va() -> IApplication: + """Create the ordered-reference H3 application.""" + return MiniMaxH3Application(PIPELINE_MINIMAX_H3_REF2VA) diff --git a/integrations_v2/minimax_h3/config.py b/integrations_v2/minimax_h3/config.py new file mode 100644 index 000000000..303f692a7 --- /dev/null +++ b/integrations_v2/minimax_h3/config.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native MiniMax H3 workflow configurations.""" + +from flashdreams.infra.config import derive_config +from minimax_h3.impl.pipeline import MiniMaxH3PipelineConfig + +PIPELINE_MINIMAX_H3_T2VA = MiniMaxH3PipelineConfig( + name="minimax-h3-t2va", workflow="t2va" +) +PIPELINE_MINIMAX_H3_FL2VA = derive_config( + PIPELINE_MINIMAX_H3_T2VA, name="minimax-h3-fl2va", workflow="fl2va" +) +PIPELINE_MINIMAX_H3_REF2VA = derive_config( + PIPELINE_MINIMAX_H3_T2VA, name="minimax-h3-ref2va", workflow="ref2va" +) + +MINIMAX_H3_CONFIGS = { + config.name: config + for config in ( + PIPELINE_MINIMAX_H3_T2VA, + PIPELINE_MINIMAX_H3_FL2VA, + PIPELINE_MINIMAX_H3_REF2VA, + ) +} diff --git a/integrations_v2/minimax_h3/impl/__init__.py b/integrations_v2/minimax_h3/impl/__init__.py new file mode 100644 index 000000000..01a8560e8 --- /dev/null +++ b/integrations_v2/minimax_h3/impl/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 model components and request conditioning.""" diff --git a/integrations_v2/minimax_h3/impl/audio_encoder.py b/integrations_v2/minimax_h3/impl/audio_encoder.py new file mode 100644 index 000000000..6c825e85b --- /dev/null +++ b/integrations_v2/minimax_h3/impl/audio_encoder.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright 2025 The MiniMax authors and The HuggingFace Team. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MiniMax H3 reference-audio encoder without decoder or posterior wrappers.""" + +# Adapted from Diffusers' MiniMax H3 audio autoencoder: encoder-only modules +# with native checkpoint names and shared FlashDreams causal attention. + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass, fields +from typing import Any + +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from torch.nn.utils import weight_norm + +from flashdreams.accelerated.multi_head_attention.sdpa import ( + scaled_dot_product_attention, +) + + +@dataclass(frozen=True, kw_only=True) +class AudioEncoderConfig: + """Encoder-only checkpoint geometry and per-channel latent normalization.""" + + encoder_dim: int = 64 + """Initial waveform-encoder feature width.""" + encoder_rates: tuple[int, ...] = (2, 4, 4, 5, 5) + """Strides of the channel-doubling encoder blocks.""" + latent_dim: int = 2048 + """Encoder trunk width before the attention projection.""" + latent_channels: int = 32 + """Channels in the raw audio latent domain.""" + num_attention_heads: int = 8 + """Causal projection heads, mean-pooled before latent projection.""" + sampling_rate: int = 32000 + """Reference waveform sample rate.""" + latents_mean: tuple[float, ...] = (0.0,) * 32 + """Raw latent means applied by the conditioning facade.""" + latents_std: tuple[float, ...] = (1.0,) * 32 + """Raw latent standard deviations applied by the conditioning facade.""" + + @classmethod + def from_dict(cls, values: Mapping[str, Any]) -> AudioEncoderConfig: + """Read a full audio VAE config while excluding decoder-only settings.""" + names = {item.name for item in fields(cls)} + decoder_names = { + "decoder_dim", + "decoder_rates", + "decoder_kernel_sizes", + "resblock_kernel_sizes", + "resblock_dilation_sizes", + } + unknown = ( + {key for key in values if not key.startswith("_")} - names - decoder_names + ) + if unknown: + raise ValueError( + f"Unknown audio VAE configuration fields: {sorted(unknown)}" + ) + return cls(**{key: value for key, value in values.items() if key in names}) + + @property + def hop_length(self) -> int: + """Return waveform samples per encoded latent.""" + return math.prod(self.encoder_rates) + + +class _Snake(nn.Module): + """DAC encoder activation with checkpoint-native per-channel frequency.""" + + def __init__(self, channels: int): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x: Tensor) -> Tensor: + """Add the learned periodic residual to the waveform features.""" + return x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2) + + +class _ResidualUnit(nn.Module): + """Weight-normalized DAC residual unit.""" + + def __init__(self, dim: int, dilation: int): + super().__init__() + self.block = nn.Sequential( + _Snake(dim), + weight_norm( + nn.Conv1d(dim, dim, 7, dilation=dilation, padding=3 * dilation) + ), + _Snake(dim), + weight_norm(nn.Conv1d(dim, dim, 1)), + ) + + def forward(self, x: Tensor) -> Tensor: + """Add the residual, center-cropping its shortcut when required.""" + residual = self.block(x) + pad = (x.shape[-1] - residual.shape[-1]) // 2 + if pad > 0: + x = x[..., pad:-pad] + return x + residual + + +class _EncoderBlock(nn.Module): + """Three residual units followed by strided channel doubling.""" + + def __init__(self, dim: int, stride: int): + super().__init__() + self.block = nn.Sequential( + *[_ResidualUnit(dim // 2, dilation) for dilation in (1, 3, 9)], + _Snake(dim // 2), + weight_norm( + nn.Conv1d( + dim // 2, + dim, + 2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) + ), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class _Encoder(nn.Module): + """Mono waveform encoder retaining checkpoint-native sequential indices.""" + + def __init__(self, config: AudioEncoderConfig): + super().__init__() + dim = config.encoder_dim + blocks: list[nn.Module] = [weight_norm(nn.Conv1d(1, dim, 7, padding=3))] + for stride in config.encoder_rates: + dim *= 2 + blocks.append(_EncoderBlock(dim, stride)) + blocks.extend( + [_Snake(dim), weight_norm(nn.Conv1d(dim, config.latent_dim, 3, padding=1))] + ) + self.block = nn.Sequential(*blocks) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class _CausalAttention(nn.Module): + """Causal attention with head-mean and adaptive feature pooling.""" + + def __init__(self, in_dim: int, out_dim: int, num_heads: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = in_dim // num_heads + self.out_dim = out_dim + self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False) + 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 = nn.Linear(out_dim, out_dim) + + def forward(self, x: Tensor) -> Tensor: + """Attend causally, then pool heads and feature width independently.""" + b, length, _ = x.shape + qkv = F.linear( + x, self.qkv.weight, torch.cat([self.q_bias, self.zero_k_bias, self.v_bias]) + ) + q, k, v = ( + qkv.reshape(b, length, 3, self.num_heads, self.head_dim) + .permute(2, 0, 1, 3, 4) + .unbind(0) + ) + x = scaled_dot_product_attention(q, k, v, is_causal=True) + x = F.adaptive_avg_pool1d(x.mean(dim=2), self.out_dim) + return self.proj(x) + + +class _GeGLU(nn.Module): + """Pre-normalized GeGLU projection with checkpoint-native names.""" + + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.w0 = nn.Linear(dim, hidden_dim) + self.w1 = nn.Linear(dim, hidden_dim) + self.w2 = nn.Linear(hidden_dim, dim) + + def forward(self, x: Tensor) -> Tensor: + x = self.norm(x) + return self.w2(F.gelu(self.w0(x), approximate="tanh") * self.w1(x)) + + +class _AttentionProjection(nn.Module): + """Residual causal-attention projection from trunk to latent width.""" + + def __init__(self, in_dim: int, out_dim: int, num_heads: int): + super().__init__() + self.norm1 = nn.LayerNorm(in_dim) + self.attn = _CausalAttention(in_dim, out_dim, num_heads) + self.proj = nn.Linear(in_dim, out_dim) + self.norm3 = nn.LayerNorm(in_dim) + self.norm2 = nn.LayerNorm(out_dim) + self.mlp = _GeGLU(out_dim, out_dim * 2) + + def forward(self, x: Tensor) -> Tensor: + x = self.proj(self.norm3(x)) + self.attn(self.norm1(x)) + return x + self.mlp(self.norm2(x)) + + +class MiniMaxH3AudioEncoder(nn.Module): + """FP32 reference-audio encoder returning only the consumed posterior mean.""" + + def __init__(self, config: AudioEncoderConfig): + super().__init__() + if any( + value <= 0 + for value in ( + config.encoder_dim, + config.latent_dim, + config.latent_channels, + config.num_attention_heads, + config.sampling_rate, + ) + ): + raise ValueError( + "Audio encoder widths, head count and sampling rate must be positive" + ) + if ( + len(config.latents_mean) != config.latent_channels + or len(config.latents_std) != config.latent_channels + ): + raise ValueError("Audio latent normalization must match latent_channels") + if any(not math.isfinite(value) for value in config.latents_mean) or any( + not math.isfinite(value) or value <= 0 for value in config.latents_std + ): + raise ValueError( + "Audio latent means must be finite and standard deviations positive" + ) + if ( + config.latent_dim % config.latent_channels + or config.latent_dim % config.num_attention_heads + ): + raise ValueError( + "Audio trunk width must be divisible by latent channels and attention heads" + ) + if not config.encoder_rates or any(rate <= 0 for rate in config.encoder_rates): + raise ValueError("Audio encoder strides must be positive") + self.config = config + self.encoder = _Encoder(config) + self.pre_block = _AttentionProjection( + config.latent_dim, config.latent_channels, config.num_attention_heads + ) + self.mean_proj = nn.Conv1d(config.latent_channels, config.latent_channels, 1) + + def encode(self, waveform: Tensor) -> Tensor: + """Return raw posterior means for mono waveforms shaped ``[B,1,S]``. + + Stereo references occupy two batch items. Right-pad to a whole encoder + hop; the caller applies the checkpoint's per-channel mean and std. + """ + if ( + waveform.ndim != 3 + or waveform.shape[1] != 1 + or any(size <= 0 for size in waveform.shape) + ): + raise ValueError( + f"Expected nonempty mono waveform [B,1,S], got {tuple(waveform.shape)}" + ) + if next(self.parameters()).dtype != torch.float32: + raise ValueError("H3 audio encoder weights must remain float32") + with torch.autocast(device_type=waveform.device.type, enabled=False): + waveform = F.pad( + waveform.float(), (0, (-waveform.shape[-1]) % self.config.hop_length) + ) + x = self.encoder(waveform) + x = self.pre_block(x.transpose(1, 2)).transpose(1, 2) + return self.mean_proj(x) + + def forward(self, waveform: Tensor) -> Tensor: + """Return raw audio posterior means.""" + return self.encode(waveform) diff --git a/integrations_v2/minimax_h3/impl/conditioning.py b/integrations_v2/minimax_h3/impl/conditioning.py new file mode 100644 index 000000000..d3ffef194 --- /dev/null +++ b/integrations_v2/minimax_h3/impl/conditioning.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 MiniMax and HuggingFace Teams +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 presentation packing and staged native conditioning.""" + +from __future__ import annotations + +import math +from collections.abc import Callable +from pathlib import Path +from typing import Any +import numpy as np +import torch +from PIL import Image, ImageOps +from flashdreams.infra.acceleration.encoder_lifecycle import ( + run_one_shot_stage, + collect_and_release_cuda_memory, +) +from .references import MiniMaxH3Reference, MiniMaxH3ReferenceSpec, load_references + +MINIMAX_H3_MIN_ASPECT_RATIO = 1 / 4 +MINIMAX_H3_MAX_ASPECT_RATIO = 4 + + +def _stage(factory: Callable, operation: Callable) -> Any: + holder = [] + + def compute(): + holder.append(factory()) + return operation(holder[0]) + + def release(): + holder.clear() + collect_and_release_cuda_memory() + + return run_one_shot_stage(compute, release=release) + + +def normalize_references( + references: list[MiniMaxH3Reference], num_frames: int +) -> list[MiniMaxH3Reference]: + """Normalize reference media at the released image, video and audio rates.""" + normalized = [] + for ref in references: + audio = ref.audio + if audio is not None: + if ref.sample_rate != 32000: + raise ValueError("Reference loader must resample audio to 32000 Hz") + audio = audio.float()[:, : int(num_frames / 24 * 32000)] + if audio.shape[0] == 1: + audio = audio.expand(2, -1).contiguous() + if audio.shape[0] != 2 or not audio.shape[1]: + raise ValueError( + "Reference audio must contain nonempty mono or stereo samples" + ) + if ref.kind == "image": + image = ref.image + width, height = image.size + if not 1 / 4 <= width / height <= 4: + raise ValueError( + "Reference image aspect ratio must be between 1:4 and 4:1" + ) + scale = 2048 / min(width, height) + size = ( + max(32, round(width * scale / 32) * 32), + max(32, round(height * scale / 32) * 32), + ) + normalized.append( + MiniMaxH3Reference( + kind="image", image=image.resize(size, Image.Resampling.LANCZOS) + ) + ) + elif ref.kind == "video": + normalized.append( + MiniMaxH3Reference( + kind="video", + frames=_normalize_video_condition( + ref.frames, ref.fps, num_frames, 32, 768, 768 * 1344, 24 + ), + fps=24, + audio=audio, + sample_rate=32000 if audio is not None else None, + ) + ) + else: + normalized.append( + MiniMaxH3Reference(kind="audio", audio=audio, sample_rate=32000) + ) + return normalized + + +def prepare_keyframes( + image_path: Path | None, last_image_path: Path | None, width: int, height: int +) -> tuple[list[Image.Image], tuple[str, ...]]: + """Stretch the geometry anchor and cover-crop its optional follower.""" + frames, anchors = [], [] + for anchor, path in (("first", image_path), ("last", last_image_path)): + if path is None: + continue + with Image.open(path) as source: + frame = ImageOps.exif_transpose(source).convert("RGB") + if not frames: + frame = frame.resize((width, height), Image.Resampling.LANCZOS) + else: + scale = max(width / frame.width, height / frame.height) + size = ( + max(width, round(frame.width * scale)), + max(height, round(frame.height * scale)), + ) + left, top = (size[0] - width) // 2, (size[1] - height) // 2 + frame = frame.resize(size, Image.Resampling.LANCZOS).crop( + (left, top, left + width, top + height) + ) + frames.append(frame) + anchors.append(anchor) + return frames, tuple(anchors) + + +def encode_visual_condition(encoder: Any, pixels: torch.Tensor) -> torch.Tensor: + """Sample seed-42 visual conditioning and apply the released fp16 rounding.""" + mean = pixels.new_tensor((0.485, 0.456, 0.406), dtype=torch.float32).view( + 1, 3, 1, 1, 1 + ) + std = pixels.new_tensor((0.229, 0.224, 0.225), dtype=torch.float32).view( + 1, 3, 1, 1, 1 + ) + pixels = (pixels.float() / 255 - mean) / std + latents = encoder.sample(pixels, generator=torch.Generator().manual_seed(42)) + latents = latents.half().float().cpu() + mean = torch.tensor(encoder.config.latents_mean).view(1, -1, 1, 1, 1) + std = torch.tensor(encoder.config.latents_std).view(1, -1, 1, 1, 1) + return (latents - mean) / std + + +def condition_request( + *, + prompt: str, + workflow: str, + width: int, + height: int, + num_frames: int, + qwen_encoder_factory: Callable, + image_path: Path | None = None, + last_image_path: Path | None = None, + references: tuple[MiniMaxH3ReferenceSpec, ...] = (), + video_encoder_factory: Callable | None = None, + audio_encoder_factory: Callable | None = None, + device: str | torch.device = "cpu", +) -> dict[str, Any]: + """Prepare one request with only one heavyweight encoder resident at a time.""" + if not isinstance(prompt, str) or not prompt.strip(): + raise ValueError("Prompt must be nonempty text") + if min(width, height) <= 0 or width % 32 or height % 32: + raise ValueError("Canvas dimensions must be positive multiples of 32") + if not 1 / 4 <= width / height <= 4: + raise ValueError("Canvas aspect ratio must be between 1:4 and 4:1") + if num_frames % 17 != 5 or not 5 <= num_frames / 24 <= 15: + raise ValueError( + "Frame count must be 17*n+5 and duration between 5 and 15 seconds" + ) + if workflow not in {"t2va", "fl2va", "ref2va"}: + raise ValueError(f"Unsupported H3 workflow: {workflow}") + if workflow != "fl2va" and (image_path is not None or last_image_path is not None): + raise ValueError("Keyframes require fl2va") + if workflow != "ref2va" and references: + raise ValueError("Ordered references require ref2va") + keyframes, anchors = prepare_keyframes(image_path, last_image_path, width, height) + if workflow == "fl2va" and not keyframes: + raise ValueError("fl2va requires a first or last keyframe") + refs = ( + normalize_references(load_references(references), num_frames) + if workflow == "ref2va" + else [] + ) + visual_refs = ( + [MiniMaxH3Reference(kind="image", image=image) for image in keyframes] + if keyframes + else refs + ) + + def encode_video(encoder): + conditions = [] + for reference in visual_refs: + if reference.kind == "image": + pixels = torch.from_numpy(np.array(reference.image)).permute(2, 0, 1)[ + None, :, None + ] + elif reference.kind == "video": + count = max(1, (len(reference.frames) - 5) // 17) * 17 + 5 + pixels = torch.from_numpy(reference.frames[:count].copy()).permute( + 3, 0, 1, 2 + )[None] + else: + continue + conditions.append(encode_visual_condition(encoder, pixels.to(device))) + return conditions + + has_visual = any(ref.kind in {"image", "video"} for ref in visual_refs) + if has_visual and video_encoder_factory is None: + raise ValueError("Visual conditioning requires a video encoder factory") + conditions = _stage(video_encoder_factory, encode_video) if has_visual else [] + + def encode_audio(encoder): + mean = torch.tensor(encoder.config.latents_mean).view(1, 1, -1) + std = torch.tensor(encoder.config.latents_std).view(1, 1, -1) + return [ + ( + ( + encoder.encode(ref.audio.to(device)[:, None]) + .float() + .cpu() + .transpose(1, 2) + - mean + ) + / std + ).reshape(-1, 32) + for ref in refs + if ref.has_audio + ] + + has_audio = any(ref.has_audio for ref in refs) + if has_audio and audio_encoder_factory is None: + raise ValueError("Audio references require an audio encoder factory") + audio_conditions = _stage(audio_encoder_factory, encode_audio) if has_audio else [] + + def encode_text(encoder): + vision, image_counts, video_counts, timestamps = _gather_vision_features( + encoder.processor, visual_refs, 24 + ) + ids, tags = _build_presentation( + encoder.tokenizer, + prompt, + visual_refs, + image_counts, + video_counts, + timestamps, + ) + embeddings = encoder({"token_ids": ids, "vision_inputs": vision}) + return { + "prompt_embeds": embeddings, + "text_token_tags": torch.tensor(tags, dtype=torch.long), + } + + return { + **_stage(qwen_encoder_factory, encode_text), + "condition_latents": conditions, + "audio_condition_latents": audio_conditions, + "height": height, + "width": width, + "num_frames": num_frames, + "keyframe_anchors": anchors, + "normalized_references": refs, + } + + +def resolve_canvas_size( + aspect_width: float, + aspect_height: float, + canvas_multiple: int, + short_edge: int, + max_pixels: int, + min_aspect_ratio: float = MINIMAX_H3_MIN_ASPECT_RATIO, + max_aspect_ratio: float = MINIMAX_H3_MAX_ASPECT_RATIO, +) -> tuple[int, int]: + """Resolve a display aspect ratio into a MiniMax-H3 canvas.""" + 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 min_aspect_ratio <= ratio <= max_aspect_ratio: + raise ValueError( + f"MiniMax-H3 supports aspect ratios from 1:{1 / min_aspect_ratio:g} to {max_aspect_ratio:g}:1, got " + f"{aspect_width}:{aspect_height} ({ratio:g})." + ) + + if ratio >= 1.0: + width, height = short_edge * ratio, float(short_edge) + else: + width, height = float(short_edge), short_edge / ratio + + area = width * height + if area > max_pixels: + scale = (max_pixels / area) ** 0.5 + width, height = width * scale, height * scale + + multiple = canvas_multiple + return max(multiple, round(height / multiple) * multiple), max( + multiple, round(width / multiple) * multiple + ) + + +def _normalize_video_condition( + frames, + fps: float, + num_frames: int, + canvas_multiple: int, + canvas_short_edge: int, + canvas_max_pixels: int, + target_fps: float, +) -> np.ndarray: + """Normalize a video reference's frames: any accepted layout, onto `uint8` at 24 fps, truncated to the generated""" + # Any accepted layout onto `uint8` THWC. A `torch.Tensor` is channels-first, as everywhere else in + # diffusers, and a `np.ndarray` channels-last; floating point values are read over `[0, 1]`. + if isinstance(frames, list): + frames = np.stack([np.asarray(frame.convert("RGB")) for frame in frames]) + if isinstance(frames, torch.Tensor): + frames = frames.movedim(-3, -1).cpu().numpy() + frames = np.asarray(frames) + if frames.dtype != np.uint8: + frames = (frames * 255.0).round().clip(0, 255).astype(np.uint8) + if frames.ndim != 4 or frames.shape[3] != 3: + raise ValueError( + f"A reference video must be `(num_frames, height, width, 3)` RGB frames, got {tuple(frames.shape)}." + ) + + # Onto MiniMax-H3's 24 fps grid: every frame is held until the slot of the next one, and the last one until + # the slot the stream's end rounds to. + if not math.isfinite(fps) or fps <= 0: + raise ValueError( + f"A reference video must have a positive frame rate, got {fps}." + ) + if fps != target_fps: + scale = target_fps / fps + slots = np.floor(np.arange(frames.shape[0]) * scale + 0.5).astype(np.int64) + frames = np.repeat( + frames, + np.diff(slots, append=math.floor(frames.shape[0] * scale + 0.5)), + axis=0, + ) + + # Truncated to the generated frame count and put on the canvas of its *own* aspect ratio — the same rule the + # target canvas follows, unlike an image reference. + frames = frames[:num_frames] + if not len(frames): + raise ValueError("Reference video is too short to contain a frame at 24 fps") + height, width = resolve_canvas_size( + frames.shape[2], + frames.shape[1], + canvas_multiple, + canvas_short_edge, + canvas_max_pixels, + ) + 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_video_condition_frames( + frames: np.ndarray, fps: float, sample_fps: float, temporal_patch: int +) -> tuple[list[np.ndarray], list[float]]: + """Sample the frames the conditioner sees from a normalized reference video, and label their vision blocks.""" + stride = fps / sample_fps + indices, cursor = [], 0.0 + while round(cursor) < frames.shape[0]: + if not indices or round(cursor) > indices[-1]: + indices.append(round(cursor)) + cursor += stride + if len(indices) < temporal_patch: + minimum = round((temporal_patch - 1) * stride) + 1 + raise ValueError( + f"A reference video is read at {sample_fps:g} fps and its sampled frames are merged in groups of " + f"{temporal_patch}, so it must run at least {minimum} frames at {fps:g} fps " + f"({minimum / fps:.2g} seconds), got {frames.shape[0]}." + ) + + timestamps = [index / sample_fps for index in range(len(indices))] + timestamps += [timestamps[-1]] * (-len(timestamps) % temporal_patch) + block_timestamps = [ + (timestamps[index] + timestamps[index + temporal_patch - 1]) / 2 + for index in range(0, len(timestamps), temporal_patch) + ] + return [frames[index] for index in indices], block_timestamps + + +def _gather_vision_features( + processor, references: list[MiniMaxH3Reference], fps: float +) -> tuple[dict, list[int], list[int], list[list[float]]]: + """Run the references' pixels through the conditioner's processors, batched per modality.""" + merge_size = processor.image_processor.merge_size**2 + vision_inputs = {} + + image_token_counts = [] + images = [reference.image for reference in references if reference.kind == "image"] + if images: + image_features = processor.image_processor(images=images, return_tensors="pt") + vision_inputs["pixel_values"] = image_features["pixel_values"] + vision_inputs["image_grid_thw"] = image_features["image_grid_thw"] + image_token_counts = [ + int(grid.prod()) // merge_size for grid in image_features["image_grid_thw"] + ] + + video_block_token_counts, video_block_timestamps = [], [] + videos = [reference for reference in references if reference.kind == "video"] + if videos: + temporal_patch = processor.video_processor.temporal_patch_size + sampled = [ + _sample_video_condition_frames(reference.frames, fps, 2.0, temporal_patch) + for reference in videos + ] + video_block_timestamps = [timestamps for _, timestamps in sampled] + video_features = processor.video_processor( + videos=[np.stack(frames) for frames, _ in sampled], + do_sample_frames=False, + return_tensors="pt", + ) + vision_inputs["pixel_values_videos"] = video_features["pixel_values_videos"] + vision_inputs["video_grid_thw"] = video_features["video_grid_thw"] + video_block_token_counts = [ + int(grid[1]) * int(grid[2]) // merge_size + for grid in video_features["video_grid_thw"] + ] + for timestamps, grid in zip( + video_block_timestamps, video_features["video_grid_thw"] + ): + if int(grid[0]) != len(timestamps): + raise ValueError( + f"The processor merged a reference video into {int(grid[0])} vision blocks, but MiniMax-H3 " + f"labels {len(timestamps)} of them." + ) + + return ( + vision_inputs, + image_token_counts, + video_block_token_counts, + video_block_timestamps, + ) + + +def _build_presentation( + tokenizer, + prompt: str, + references: list[MiniMaxH3Reference], + image_token_counts: list[int], + video_block_token_counts: list[int], + video_block_timestamps: list[list[float]], + text_tag: int = 1, + video_tag: int = 0, +) -> tuple[list[int], list[int]]: + """Tokenize MiniMax-H3's presentation of a `ref2va` request.""" + + def text(value: str) -> tuple[list[int], list[int]]: + token_ids = tokenizer(value, add_special_tokens=False)["input_ids"] + return token_ids, [text_tag] * len(token_ids) + + def vision(pad_token: str, num_tokens: int) -> tuple[list[int], list[int]]: + token_ids = ( + [tokenizer.convert_tokens_to_ids("<|vision_start|>")] + + [tokenizer.convert_tokens_to_ids(pad_token)] * num_tokens + + [tokenizer.convert_tokens_to_ids("<|vision_end|>")] + ) + return token_ids, [video_tag] * len(token_ids) + + token_ids, token_tags = [], [] + + 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"