From a7ca642e5c901cf5c24c06b51a7e3380e35fca72 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Mon, 31 Aug 2026 12:15:47 -0700 Subject: [PATCH] [feat]: add opt-in CUDA TAEH3 preview decode for FastH3 Skip the 9.7 GiB video VAE on T2VA so GB10 can reconstruct alpine 768 in a couple of seconds instead of a full VAE pass. Independent of lazy-module-load; uses sequential start already on main. --- .../inference_schema_parity_inventory.yaml | 3 + .../installation/spark_performance.md | 6 + examples/inference/basic/basic_fasth3.py | 9 + fastvideo/fastvideo_args.py | 26 +++ fastvideo/models/vaes/minimax_h3_taeh3.py | 198 ++++++++++++++++++ .../basic/minimax_h3/minimax_h3_pipeline.py | 86 +++++--- .../pipelines/basic/minimax_h3/packing.py | 15 +- .../minimax_h3/stages/minimax_h3_decoding.py | 29 ++- .../stages/minimax_h3_latent_preparation.py | 24 ++- .../inference/test_basic_fasth3_profile.py | 8 + .../test_minimax_h3_sequential_start.py | 51 +++++ .../stages/test_minimax_h3_vae_streaming.py | 33 ++- fastvideo/tests/vaes/test_minimax_h3_taeh3.py | 22 ++ .../benchmark_minimax_h3_video_vae_memory.py | 14 +- 14 files changed, 469 insertions(+), 55 deletions(-) create mode 100644 fastvideo/models/vaes/minimax_h3_taeh3.py create mode 100644 fastvideo/tests/vaes/test_minimax_h3_taeh3.py diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 196b305d56..38ac52ea1b 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -81,6 +81,9 @@ surfaces: inference_torch_compile: "Regional inference compile opt-in currently carried through PipelineSelection.experimental rather than CompileConfig." vae_parallel_decode: "MiniMax-H3 sequence-parallel VAE decode opt-in; model-specific optimization not yet represented in the typed public schema." h3_sequential_load: "MiniMax-H3 sequential text-encoder then DiT/VAE load; model-specific optimization not yet represented in the typed public schema." + video_decode_backend: "MiniMax-H3 video decoder selection (full VAE vs TAEH3 preview); model-specific optimization not yet represented in the typed public schema." + taeh3_checkpoint: "Optional local TAEH3 safetensors path; model-specific optimization not yet represented in the typed public schema." + taeh3_chunk_size: "TAEH3 temporal chunk length; model-specific optimization not yet represented in the typed public schema." vae_parallel_encode: "MiniMax-H3 sequence-parallel reference VAE encode opt-in; model-specific optimization not yet represented in the typed public schema." vae_parallel_decode_strategy: "Chunk-transport collective for vae_parallel_decode; model-specific optimization not yet represented in the typed public schema." attention_backend: "Process-wide default attention-backend request applied per component at load time; kernel-selection knob not yet represented in the typed public schema." diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index da92c0084b..490df05cb3 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -170,6 +170,12 @@ is power-cycled. To avoid it: encodes first, releases the encoder, then loads DiT and VAEs onto the accelerator (`to_cpu` follows `cpu_offload`, which is off here). See [Offloading](../../inference/offloading.md). +- **FastH3 TAEH3** (`--video-decode-backend taeh3`) is an opt-in preview decoder. + T2VA never materializes the 9.7 GiB video VAE (DiT still loads after Qwen via + sequential start). On this box, alpine 768×1344×124 decoded in **2.4 s** versus + **68 s** for the full VAE, and one T2VA generation finished in **224 s** + end-to-end. Reconstruction is approximate, not lossless. FL2VA/Ref2VA still + need the full VAE to encode references. ## Gotchas specific to the GB10 diff --git a/examples/inference/basic/basic_fasth3.py b/examples/inference/basic/basic_fasth3.py index 2e240c6fa0..d2ecfd125c 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/examples/inference/basic/basic_fasth3.py @@ -104,6 +104,11 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser: default=None, help="encode with Qwen3-VL, release it, then load DiT/VAEs. Default auto: on for " "unified-memory devices (GB10), off on discrete GPUs") + parser.add_argument("--video-decode-backend", + choices=("h3-vae", "taeh3"), + default="h3-vae", + help="h3-vae is the full MiniMax VAE; taeh3 is the fast approximate preview decoder") + parser.add_argument("--taeh3-checkpoint", default=None, help="local taeh3.safetensors; unset uses the pinned cache") parser.add_argument("--replicated-dit", action=argparse.BooleanOptionalAction, default=True, @@ -236,6 +241,10 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig: } if args.h3_sequential_load is not None: experimental["h3_sequential_load"] = args.h3_sequential_load + if args.video_decode_backend != "h3-vae": + experimental["video_decode_backend"] = args.video_decode_backend + if args.taeh3_checkpoint is not None: + experimental["taeh3_checkpoint"] = args.taeh3_checkpoint if use_vsa: experimental.update({ "VSA_sparsity": args.vsa_sparsity, diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index ff812437f8..dd6dbb1810 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -166,6 +166,13 @@ class FastVideoArgs: # False overrides the probe. Training never defers. h3_sequential_load: bool | None = None + # MiniMax-H3 video reconstruction. ``h3-vae`` is the full ViT decoder. + # ``taeh3`` is Ollin Boer Bohan's tiny preview decoder; it changes quality + # and is opt-in. T2VA with TAEH3 does not need the video VAE weights. + video_decode_backend: str = "h3-vae" + taeh3_checkpoint: str | None = None + taeh3_chunk_size: int = 5 + # Sequence-parallel MiniMax-H3 VAE (opt-in, default off). With SP > 1 the # video VAE's temporal chunks (decode) and clips (reference encode) are # round-robined across the sequence-parallel ranks and reassembled @@ -729,6 +736,25 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "Omit for auto (on for unified-memory devices such as GB10; off on discrete GPUs). " "Pass --no-h3-sequential-load to keep the encoder resident for later generate() calls.", ) + parser.add_argument( + "--video-decode-backend", + type=str, + choices=("h3-vae", "taeh3"), + default=FastVideoArgs.video_decode_backend, + help="MiniMax-H3 video decoder. taeh3 is a fast approximate preview decoder; h3-vae is the full VAE.", + ) + parser.add_argument( + "--taeh3-checkpoint", + type=str, + default=None, + help="Local taeh3.safetensors path. Unset downloads the pinned upstream weights into the cache.", + ) + parser.add_argument( + "--taeh3-chunk-size", + type=int, + default=FastVideoArgs.taeh3_chunk_size, + help="TAEH3 latent frames per execution chunk.", + ) parser.add_argument( "--vae-parallel-decode", action=StoreBoolean, diff --git a/fastvideo/models/vaes/minimax_h3_taeh3.py b/fastvideo/models/vaes/minimax_h3_taeh3.py new file mode 100644 index 0000000000..0de71db813 --- /dev/null +++ b/fastvideo/models/vaes/minimax_h3_taeh3.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Optional, approximate MiniMax H3 tiny decoder for CUDA/CPU PyTorch. + +Architecture and temporal mapping adapted from madebyollin/taehv at +62f7591f59dfbb4c3c02b7a621d180a9eeaba26c (MIT, Ollin Boer Bohan). +This decoder consumes normalized diffusion latents; it does not use the full +H3 VAE's latent mean/std or pixel denormalization. +""" + +from __future__ import annotations + +import hashlib +import tempfile +import urllib.request +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +TAEH3_REVISION = "62f7591f59dfbb4c3c02b7a621d180a9eeaba26c" +TAEH3_URL = f"https://raw.githubusercontent.com/madebyollin/taehv/{TAEH3_REVISION}/safetensors/taeh3.safetensors" +TAEH3_SHA256 = "4fd022bfcab08772fe0536b17ea1a3bbb5625be11e397868d1c5d891863d4c13" + +_EXPECTED_SHAPES: dict[str, tuple[int, ...]] = { + "decoder.1.weight": (256, 24, 3, 3), + "decoder.1.bias": (256, ), + "decoder.7.conv.weight": (256, 256, 1, 1), + "decoder.8.weight": (128, 256, 3, 3), + "decoder.13.conv.weight": (256, 128, 1, 1), + "decoder.14.weight": (64, 128, 3, 3), + "decoder.19.conv.weight": (128, 64, 1, 1), + "decoder.20.weight": (64, 64, 3, 3), + "decoder.22.weight": (12, 64, 3, 3), + "decoder.22.bias": (12, ), +} +for _index, _channels in ((3, 256), (4, 256), (5, 256), (9, 128), (10, 128), (11, 128), (15, 64), (16, 64), (17, 64)): + for _layer in (0, 2, 4): + _prefix = f"decoder.{_index}.conv.{_layer}" + _EXPECTED_SHAPES[f"{_prefix}.weight"] = (_channels, _channels * (2 if _layer == 0 else 1), 3, 3) + _EXPECTED_SHAPES[f"{_prefix}.bias"] = (_channels, ) + + +def ensure_taeh3_checkpoint(checkpoint_path: str | Path | None = None) -> Path: + """Fetch only pinned weights, atomically; never download executable code.""" + if checkpoint_path is not None: + path = Path(checkpoint_path).expanduser() + if not path.is_file(): + raise FileNotFoundError(f"TAEH3 checkpoint not found: {path}") + if path.suffix != ".safetensors": + raise ValueError("The TAEH3 decoder requires a .safetensors checkpoint.") + return path + path = Path.home() / ".cache/fastvideo/taehv/taeh3.safetensors" + + def verify(candidate: Path) -> None: + hasher = hashlib.sha256() + with candidate.open("rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + hasher.update(chunk) + digest = hasher.hexdigest() + if digest != TAEH3_SHA256: + raise RuntimeError(f"TAEH3 checkpoint failed SHA-256 verification: {candidate}") + + if path.exists(): + verify(path) + return path + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".safetensors", delete=False) as temporary_file: + temporary = Path(temporary_file.name) + try: + with urllib.request.urlopen(TAEH3_URL, timeout=60) as response, temporary.open("wb") as handle: + while chunk := response.read(1 << 20): + handle.write(chunk) + verify(temporary) + temporary.replace(path) + finally: + temporary.unlink(missing_ok=True) + logger.info("Cached TAEH3 checkpoint at %s", path) + return path + + +class TorchTAEH3Decoder: + """Decode H3 NTCHW latents with bounded temporal feature memory.""" + + def __init__(self, checkpoint_path: str | Path, *, dtype: torch.dtype = torch.float32) -> None: + from safetensors.torch import load_file + + raw = load_file(str(checkpoint_path)) + actual = {key for key in raw if key.startswith("decoder.")} + if actual != set(_EXPECTED_SHAPES): + raise ValueError(f"TAEH3 decoder keys mismatch: missing={set(_EXPECTED_SHAPES) - actual}, " + f"unexpected={actual - set(_EXPECTED_SHAPES)}") + self.dtype = dtype + self.weights: dict[str, torch.Tensor] = {} + for key, shape in _EXPECTED_SHAPES.items(): + value = raw[key] + if tuple(value.shape) != shape: + raise ValueError(f"TAEH3 weight {key} has shape {tuple(value.shape)}, expected {shape}") + self.weights[key] = value.detach().to(dtype=dtype).contiguous() + + def to(self, device: torch.device) -> "TorchTAEH3Decoder": + self.weights = {key: value.to(device=device, non_blocking=True) for key, value in self.weights.items()} + return self + + def _conv(self, x: torch.Tensor, name: str) -> torch.Tensor: + weight = self.weights[f"{name}.weight"] + bias = self.weights.get(f"{name}.bias") + padding = weight.shape[-1] // 2 + return F.conv2d(x, weight, bias, padding=padding) + + def _chunk(self, x: torch.Tensor, memory: dict[int, torch.Tensor]) -> torch.Tensor: + n, t, c, h, w = x.shape + x = F.relu(self._conv(torch.tanh(x.reshape(n * t, c, h, w) / 3.0) * 3.0, "decoder.1")) + for indices, grow, projection, stride in (((3, 4, 5), 7, 8, 1), ((9, 10, 11), 13, 14, 2), ((15, 16, 17), 19, 20, + 2)): + for index in indices: + nt, c, h, w = x.shape + sequence = x.reshape(n, -1, c, h, w) + previous = memory.get(index) + if previous is None: + previous = torch.zeros_like(sequence[:, :1]) + past = torch.cat([previous, sequence[:, :-1]], dim=1).reshape_as(x) + memory[index] = sequence[:, -1:].contiguous() + y = torch.cat([x, past], dim=1) + for layer in (0, 2, 4): + y = self._conv(y, f"decoder.{index}.conv.{layer}") + if layer != 4: + y = F.relu(y) + x = F.relu(x + y) + nt, c_pre, h, w = x.shape + x = F.interpolate(x, scale_factor=2, mode="nearest") + x = self._conv(x, f"decoder.{grow}.conv") + x = x.reshape(nt, stride, c_pre, h * 2, w * 2).reshape(nt * stride, c_pre, h * 2, w * 2) + x = self._conv(x, f"decoder.{projection}") + x = self._conv(F.relu(x), "decoder.22") + x = F.pixel_shuffle(x, 2).clamp(0, 1) + nt, c, h, w = x.shape + return x.reshape(n, -1, c, h, w) + + def decode_ntchw(self, latents: torch.Tensor, *, chunk_size: int = 5) -> torch.Tensor: + """Return NTCHW RGB in [0, 1] for H3's valid 5*k-3 latent lengths.""" + if latents.ndim != 5 or latents.shape[2] != 24 or min(latents.shape) <= 0: + raise ValueError(f"Expected nonempty NTCHW H3 latents with 24 channels, got {tuple(latents.shape)}") + if latents.shape[1] % 5 != 2: + raise ValueError("H3 latent time must be 5*k-3, for example 2, 7, or 37.") + if chunk_size < 1: + raise ValueError("TAEH3 chunk_size must be positive.") + x = latents.to(dtype=self.dtype) + memory: dict[int, torch.Tensor] = {} + frames: list[torch.Tensor] = [] + for start in range(0, x.shape[1], chunk_size): + decoded = self._chunk(x[:, start:start + chunk_size], memory) + keep = [i for i in range(decoded.shape[1]) if (start * 4 + i) % 20 >= 3] + frames.append(decoded[:, keep]) + return torch.cat(frames, dim=1) + + +_DECODER_CACHE: dict[tuple[str, str, str], TorchTAEH3Decoder] = {} + + +def decode_ncthw_latents_taeh3( + latents: torch.Tensor, + *, + device: torch.device, + checkpoint_path: str | Path | None = None, + chunk_size: int = 5, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Decode normalized NCTHW diffusion latents into NCTHW RGB in [0, 1].""" + if latents.ndim != 5: + raise ValueError(f"Expected NCTHW latents, got {tuple(latents.shape)}") + checkpoint = ensure_taeh3_checkpoint(checkpoint_path) + cache_key = (str(checkpoint), str(device), str(dtype)) + decoder = _DECODER_CACHE.get(cache_key) + if decoder is None: + decoder = TorchTAEH3Decoder(checkpoint, dtype=dtype).to(device) + _DECODER_CACHE[cache_key] = decoder + ntchw = latents.to(device=device, dtype=dtype).permute(0, 2, 1, 3, 4).contiguous() + rgb = decoder.decode_ntchw(ntchw, chunk_size=chunk_size) + return rgb.permute(0, 2, 1, 3, 4).contiguous() + + +def taeh3_decoded_pixel_shape(latent_shape: tuple[int, ...] | torch.Size) -> tuple[int, int, int, int, int]: + """Return NCTHW pixel shape for H3 TAEH3 (16x spatial, drop 3 of every 20 raw frames).""" + if len(latent_shape) != 5: + raise ValueError(f"MiniMax-H3 latents must be five-dimensional, got shape {tuple(latent_shape)}.") + batch, channels, latent_frames, latent_height, latent_width = map(int, latent_shape) + if channels != 24: + raise ValueError(f"TAEH3 latents must have 24 channels, got {channels}.") + if latent_frames % 5 != 2: + raise ValueError(f"H3 latent time must be 5*k-3, got {latent_frames}.") + raw_frames = latent_frames * 4 + kept = sum(1 for index in range(raw_frames) if index % 20 >= 3) + return (batch, 3, kept, latent_height * 16, latent_width * 16) diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index e1dceab49c..02d618497c 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -5,15 +5,16 @@ import gc from dataclasses import dataclass +from pathlib import Path from typing import Any import torch from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEArchConfig -from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEArchConfig from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger +from fastvideo.models.hf_transformer_utils import get_diffusers_config from fastvideo.pipelines.basic.minimax_h3.stages import ( MiniMaxH3AudioDecodingStage, MiniMaxH3ConditioningStage, @@ -34,10 +35,28 @@ _DENOISE_MODULE_NAMES = ("vae", "audio_vae", "transformer") -@dataclass(frozen=True) -class _H3VideoGeometry: - spatial_compression_ratio: int - latent_channels: int +def _apply_h3_checkpoint_arch_configs(model_path: str, fastvideo_args: FastVideoArgs, + extra_config_module_map: dict[str, str]) -> None: + """Overlay checkpoint config.json onto pipeline configs without loading weights.""" + root = Path(model_path) + vae_dir = root / "vae" + if (vae_dir / "config.json").is_file(): + fastvideo_args.pipeline_config.vae_config.update_model_arch(get_diffusers_config(str(vae_dir))) + transformer_dir = root / extra_config_module_map.get("transformer", "transformer") + if (transformer_dir / "config.json").is_file(): + fastvideo_args.pipeline_config.dit_config.update_model_arch(get_diffusers_config(str(transformer_dir))) + dit_arch = getattr(fastvideo_args.pipeline_config.dit_config, "arch_config", None) + vae_arch = getattr(fastvideo_args.pipeline_config.vae_config, "arch_config", None) + logger.info( + "MiniMax-H3 geometry from config: patch_size=%s spatial_compression_ratio=%s latent_channels=%s", + getattr(dit_arch, "patch_size", None), + getattr(vae_arch, "spatial_compression_ratio", None), + getattr(vae_arch, "latent_channels", None), + ) + + +def _use_taeh3_t2va(fastvideo_args: FastVideoArgs | None, *, ref2va: bool) -> bool: + return (not ref2va) and getattr(fastvideo_args, "video_decode_backend", "h3-vae") == "taeh3" @dataclass(frozen=True) @@ -45,14 +64,6 @@ class _H3AudioGeometry: sampling_rate: int -def _default_video_geometry() -> _H3VideoGeometry: - arch = MiniMaxH3VideoVAEArchConfig() - return _H3VideoGeometry( - spatial_compression_ratio=int(arch.spatial_compression_ratio), - latent_channels=int(arch.latent_channels), - ) - - def _default_audio_geometry() -> _H3AudioGeometry: return _H3AudioGeometry(sampling_rate=int(MiniMaxH3AudioVAEArchConfig().sampling_rate)) @@ -82,6 +93,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): ] pipeline_config_cls: type[MiniMaxH3PipelineConfig] = MiniMaxH3PipelineConfig + _ref2va_default = False _required_config_modules = [ "text_encoder", "tokenizer", @@ -94,7 +106,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): ] def __init__(self, *args: Any, **kwargs: Any) -> None: - self._ref2va = False + self._ref2va = getattr(self, "_ref2va_default", False) self._denoise_stages_ready = False super().__init__(*args, **kwargs) @@ -103,7 +115,7 @@ def get_hf_download_component_dirs(cls) -> tuple[str, ...]: return tuple(sorted(cls._extra_config_module_map.get(name, name) for name in cls._required_config_modules)) def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: - del fastvideo_args + _apply_h3_checkpoint_arch_configs(self.model_path, fastvideo_args, self._extra_config_module_map) for module_name, modality, expected_shift in ( ("scheduler", "video", 12.0), ("audio_scheduler", "audio", 3.0), @@ -129,19 +141,35 @@ def _defer_denoise_modules(self, fastvideo_args: FastVideoArgs) -> bool: logger.info("MiniMax-H3 sequential module load auto=%s (unified_memory=%s)", unified, unified) return unified + def _denoise_module_names(self, fastvideo_args: FastVideoArgs | None = None) -> tuple[str, ...]: + args = fastvideo_args if fastvideo_args is not None else getattr(self, "fastvideo_args", None) + if _use_taeh3_t2va(args, ref2va=self._ref2va): + return tuple(name for name in _DENOISE_MODULE_NAMES if name != "vae") + return _DENOISE_MODULE_NAMES + def _denoise_modules_loaded(self) -> bool: - return all(self.get_module(name) is not None for name in _DENOISE_MODULE_NAMES) + return all(self.get_module(name) is not None for name in self._denoise_module_names()) def load_modules(self, fastvideo_args: FastVideoArgs, loaded_modules: dict[str, torch.nn.Module] | None = None) -> dict[str, Any]: """Load the Qwen3-VL conditioner first; defer DiT and VAEs until after encode.""" if not self._defer_denoise_modules(fastvideo_args): + if _use_taeh3_t2va(fastvideo_args, ref2va=self._ref2va): + saved = list(self.required_config_modules) + self._required_config_modules = [name for name in saved if name != "vae"] + try: + return super().load_modules(fastvideo_args, loaded_modules) + finally: + self._required_config_modules = saved return super().load_modules(fastvideo_args, loaded_modules) - if loaded_modules is not None and all(name in loaded_modules for name in _DENOISE_MODULE_NAMES): + if loaded_modules is not None and all(name in loaded_modules + for name in self._denoise_module_names(fastvideo_args)): return super().load_modules(fastvideo_args, loaded_modules) saved = list(self.required_config_modules) + # Always defer the full denoise set on the first load. TAEH3 T2VA then + # omits the video VAE from the second load via `_denoise_module_names`. self._required_config_modules = [name for name in saved if name not in _DENOISE_MODULE_NAMES] try: logger.info("Loading MiniMax-H3 condition modules first: %s", self._required_config_modules) @@ -153,10 +181,13 @@ def _load_denoise_modules(self, fastvideo_args: FastVideoArgs) -> None: if self._denoise_modules_loaded(): return saved = list(self.required_config_modules) + denoise_names = self._denoise_module_names(fastvideo_args) self._required_config_modules = [name for name in saved if name != "text_encoder"] + if _use_taeh3_t2va(fastvideo_args, ref2va=self._ref2va): + self._required_config_modules = [name for name in self._required_config_modules if name != "vae"] try: logger.info("Loading MiniMax-H3 denoise modules after releasing the text encoder: %s", - [name for name in self._required_config_modules if name in _DENOISE_MODULE_NAMES]) + [name for name in self._required_config_modules if name in denoise_names]) loaded = super().load_modules(fastvideo_args, loaded_modules=self.modules) for name, module in loaded.items(): self.add_module(name, module) @@ -177,7 +208,10 @@ def _release_text_encoder(self) -> None: torch.cuda.empty_cache() def _input_vae(self) -> Any: - return self.get_module("vae") or _default_video_geometry() + live = self.get_module("vae") + if live is not None: + return live + return self.fastvideo_args.pipeline_config.vae_config.arch_config def _input_audio_vae(self, *, ref2va: bool) -> Any | None: if not ref2va: @@ -209,13 +243,16 @@ def _add_denoise_stages(self, *, ref2va: bool) -> None: audio_vae = self.get_module("audio_vae") scheduler = self.get_module("scheduler") audio_scheduler = self.get_module("audio_scheduler") - if transformer is None or vae is None or audio_vae is None: - raise RuntimeError("MiniMax-H3 denoise stages require transformer, vae, and audio_vae to be loaded.") + use_taeh3 = _use_taeh3_t2va(getattr(self, "fastvideo_args", None), ref2va=ref2va) + if transformer is None or audio_vae is None: + raise RuntimeError("MiniMax-H3 denoise stages require transformer and audio_vae to be loaded.") + if not use_taeh3 and vae is None: + raise RuntimeError("MiniMax-H3 full-VAE decode requires the video VAE to be loaded.") + encode_vae = vae if vae is not None else self._input_vae() self.add_stage( "latent_preparation_stage", MiniMaxH3LatentPreparationStage( - transformer=transformer, - vae=vae, + vae=encode_vae, audio_vae=audio_vae, scheduler=scheduler, ref2va=ref2va, @@ -229,7 +266,7 @@ def _add_denoise_stages(self, *, ref2va: bool) -> None: audio_scheduler=audio_scheduler, ), ) - self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=vae, transformer=transformer)) + self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=None if use_taeh3 else vae)) self.add_stage("audio_decoding_stage", MiniMaxH3AudioDecodingStage(audio_vae=audio_vae)) self._denoise_stages_ready = True @@ -274,6 +311,7 @@ class MiniMaxH3RefPipeline(MiniMaxH3BasePipeline): """Ordered-reference joint video/stereo-audio pipeline for Ref2VA.""" _extra_config_module_map = {"transformer": "transformer_ref"} + _ref2va_default = True def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: del fastvideo_args diff --git a/fastvideo/pipelines/basic/minimax_h3/packing.py b/fastvideo/pipelines/basic/minimax_h3/packing.py index 7ef3468512..1c67c49a35 100644 --- a/fastvideo/pipelines/basic/minimax_h3/packing.py +++ b/fastvideo/pipelines/basic/minimax_h3/packing.py @@ -4,7 +4,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -38,6 +38,19 @@ MINIMAX_H3_KEYFRAME_NOISE_AUG = 0.999 MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42 + +def h3_dit_patch_size(fastvideo_args: Any) -> tuple[int, int, int]: + """Read DiT patch size from pipeline config, not live transformer weights.""" + dit_config = getattr(getattr(fastvideo_args, "pipeline_config", None), "dit_config", None) + patch_size = getattr(dit_config, "patch_size", None) + if patch_size is None: + raise ValueError("MiniMax-H3 requires pipeline_config.dit_config.patch_size.") + values = tuple(int(axis) for axis in patch_size) + if len(values) != 3 or min(values) <= 0: + raise ValueError(f"MiniMax-H3 patch_size must be three positive ints, got {patch_size!r}.") + return values + + MINIMAX_H3_ROPE_FRAME_RESCALE = 5.0 / 3.0 MINIMAX_H3_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) _ROPE_SPATIAL_SCALE = 32 diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py index 7cbc0d6e28..1eae7005e5 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -16,6 +16,7 @@ from fastvideo.profiler import nvtx_range from fastvideo.pipelines.basic.minimax_h3.packing import ( MiniMaxH3PackedLayout, + h3_dit_patch_size, unpack_audio_tokens, unpatchify_video_tokens, ) @@ -58,10 +59,9 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): performance_component_metric = "vae_decode_time_s" - def __init__(self, vae: AutoencoderKLMiniMaxH3, transformer: Any) -> None: + def __init__(self, vae: AutoencoderKLMiniMaxH3 | None) -> None: super().__init__() self.vae = vae - self.transformer = transformer def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() @@ -97,9 +97,32 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward latent_height, latent_width, channels, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), ) device = get_local_torch_device() + backend = getattr(fastvideo_args, "video_decode_backend", "h3-vae") + if backend == "taeh3": + from fastvideo.models.vaes.minimax_h3_taeh3 import decode_ncthw_latents_taeh3, taeh3_decoded_pixel_shape + + if fastvideo_args.output_type == "latent": + batch.output = latents.detach().float().cpu() if is_output_rank else placeholder + return batch + expected = taeh3_decoded_pixel_shape(tuple(latents.shape)) + logger.info("MiniMax-H3 video decode: TAEH3 preview (%s -> %s)", tuple(latents.shape), expected) + with nvtx_range("minimax_h3.taeh3"): + pixels = decode_ncthw_latents_taeh3( + latents, + device=device, + checkpoint_path=getattr(fastvideo_args, "taeh3_checkpoint", None), + chunk_size=int(getattr(fastvideo_args, "taeh3_chunk_size", 5) or 5), + ) + batch.output = pixels.float().cpu() if is_output_rank else placeholder + if is_output_rank and tuple(batch.output.shape) != expected: + raise RuntimeError(f"TAEH3 wrote {tuple(batch.output.shape)}, expected {expected}.") + return batch + + if self.vae is None: + raise RuntimeError("MiniMax-H3 full VAE decode requires a loaded video VAE.") self.vae.to(device) try: latents = self.vae.denormalize_latents(latents.to(device=device, dtype=torch.float32)) diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py index b30e599f7e..2d3b90c6a7 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py @@ -21,6 +21,7 @@ audio_latent_num_frames, build_packed_sequence, build_ref2va_packed_sequence, + h3_dit_patch_size, keyframe_condition_noise, patchify_video_latents, ) @@ -65,7 +66,6 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): def __init__( self, - transformer: Any, vae: Any, audio_vae: Any, scheduler: Any, @@ -73,7 +73,6 @@ def __init__( ref2va: bool = False, ) -> None: super().__init__() - self.transformer = transformer self.vae = vae self.audio_vae = audio_vae self.scheduler = scheduler @@ -111,7 +110,7 @@ def _encode_visual_rows( device: torch.device, fastvideo_args: FastVideoArgs, ) -> list[torch.Tensor]: - patch_size = self.transformer.patch_size + patch_size = h3_dit_patch_size(fastvideo_args) # Reference encode runs on every rank (all ranks hold identical # prepared references), so clip-parallel encode keeps participation # uniform by construction: each rank encodes a clip subset and the @@ -176,6 +175,9 @@ def _encode_fl2va_conditions( raise TypeError("MiniMax-H3 keyframes must be a list.") if not keyframes: return None, None + if not hasattr(self.vae, "encode_keyframe"): + raise RuntimeError("TAEH3 T2VA preview decode does not load the video VAE. " + "FL2VA keyframes still need --video-decode-backend h3-vae.") vae_device = get_local_torch_device() self.vae.to(vae_device) @@ -184,7 +186,7 @@ def _encode_fl2va_conditions( for image in keyframes: clean_rows.append( patchify_video_latents(self._encode_keyframe_latents(image, vae_device), - self.transformer.patch_size)) + h3_dit_patch_size(fastvideo_args))) finally: if fastvideo_args.vae_cpu_offload: self.vae.to("cpu") @@ -193,7 +195,7 @@ def _encode_fl2va_conditions( shapes = ((1, latent_height, latent_width), ) * len(keyframes) noise = keyframe_condition_noise( shapes, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), self.vae.latent_channels, generator=batch.generator, device=device, @@ -241,7 +243,7 @@ def _encode_ref2va_conditions( for reference in references if reference.media_type != "audio") noise = keyframe_condition_noise( shapes, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), self.vae.latent_channels, generator=batch.generator, device=device, @@ -259,7 +261,7 @@ def _encode_ref2va_conditions( reference.waveform = None return video_conditions, audio_conditions - def _build_layout(self, batch: ForwardBatch) -> MiniMaxH3PackedLayout: + def _build_layout(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> MiniMaxH3PackedLayout: text_token_tags = batch.extra.get(MINIMAX_H3_TEXT_TOKEN_TAGS_KEY) if not isinstance(text_token_tags, torch.Tensor): raise ValueError("MiniMax-H3 conditioning must produce text token tags.") @@ -276,7 +278,7 @@ def _build_layout(self, batch: ForwardBatch) -> MiniMaxH3PackedLayout: height, width, num_audio_latents, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), ) anchors = batch.extra.get(MINIMAX_H3_KEYFRAME_ANCHORS_KEY, ()) if not isinstance(anchors, tuple): @@ -287,7 +289,7 @@ def _build_layout(self, batch: ForwardBatch) -> MiniMaxH3PackedLayout: height, width, num_audio_latents, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), anchors, ) @@ -301,7 +303,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward else: condition_video, condition_audio = self._encode_fl2va_conditions(batch, fastvideo_args, device) - layout = self._build_layout(batch) + layout = self._build_layout(batch, fastvideo_args) video_channels, num_frames, height, width = _video_geometry(batch) expected_video_shape = (1, video_channels, num_frames, height, width) if video_noise is None: @@ -315,7 +317,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward raise ValueError(f"MiniMax-H3 injected video latents must have shape {expected_video_shape}, " f"got {tuple(video_noise.shape)}.") video_rows = patchify_video_latents(video_noise.to(device=device, dtype=torch.float32), - self.transformer.patch_size) + h3_dit_patch_size(fastvideo_args)) num_audio_latents = layout.num_audio_latents expected_audio_shape = (MINIMAX_H3_AUDIO_CHANNELS, self.audio_vae.latent_channels, num_audio_latents) diff --git a/fastvideo/tests/inference/test_basic_fasth3_profile.py b/fastvideo/tests/inference/test_basic_fasth3_profile.py index 5c2a2baea4..af9b28fae4 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_profile.py +++ b/fastvideo/tests/inference/test_basic_fasth3_profile.py @@ -247,3 +247,11 @@ def from_config(cls, config): assert "Measured E2E wall times (n=3, warmup excluded): [6.0, 7.0, 8.0]" in output assert "Median E2E wall time: 7.000s" in output assert "Median denoising time: 2.500s" in output + + +def test_taeh3_backend_is_opt_in_experimental(): + default = fasth3.build_generator_config(_args()) + taeh3 = fasth3.build_generator_config(_args("--video-decode-backend", "taeh3")) + + assert "video_decode_backend" not in default.pipeline.experimental + assert taeh3.pipeline.experimental["video_decode_backend"] == "taeh3" diff --git a/fastvideo/tests/stages/test_minimax_h3_sequential_start.py b/fastvideo/tests/stages/test_minimax_h3_sequential_start.py index 6ddfc2d91b..8f31973789 100644 --- a/fastvideo/tests/stages/test_minimax_h3_sequential_start.py +++ b/fastvideo/tests/stages/test_minimax_h3_sequential_start.py @@ -197,3 +197,54 @@ def test_cli_tri_state_h3_sequential_load() -> None: assert parser.parse_args([]).h3_sequential_load is None assert parser.parse_args(["--h3-sequential-load"]).h3_sequential_load is True assert parser.parse_args(["--no-h3-sequential-load"]).h3_sequential_load is False + + +def test_taeh3_t2va_skips_video_vae_on_the_deferred_load(monkeypatch) -> None: + events: list = [] + _patch_pipeline_construction(monkeypatch, events) + loads: list[list[str]] = [] + + def fake_load(self, fastvideo_args, loaded_modules=None): + del fastvideo_args + requested = list(self.required_config_modules) + loads.append(requested) + modules = dict(loaded_modules or {}) + for name in requested: + modules.setdefault(name, _stub_module(name)) + return modules + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_load) + args = FastVideoArgs( + model_path="unused/for-this-test", + enable_stage_verification=False, + h3_sequential_load=True, + video_decode_backend="taeh3", + ) + pipeline = MiniMaxH3Pipeline("unused/for-this-test", args) + pipeline.post_init() + + condition_stage = pipeline._stage_name_mapping["conditioning_stage"] + passthrough = lambda batch, _args: batch + monkeypatch.setattr(pipeline._stage_name_mapping["input_preparation_stage"], "forward", passthrough) + monkeypatch.setattr(condition_stage, "forward", passthrough) + original_add_denoise = pipeline._add_denoise_stages + + def fake_add_denoise(*, ref2va: bool) -> None: + original_add_denoise(ref2va=ref2va) + for name in ( + "latent_preparation_stage", + "denoising_stage", + "video_decoding_stage", + "audio_decoding_stage", + ): + monkeypatch.setattr(pipeline._stage_name_mapping[name], "forward", passthrough) + + monkeypatch.setattr(pipeline, "_add_denoise_stages", fake_add_denoise) + pipeline.forward(ForwardBatch(data_type="video", prompt="alpine dancer"), args) + + assert "text_encoder" in loads[0] + assert all(name not in loads[0] for name in _DENOISE_MODULE_NAMES) + assert "transformer" in loads[1] + assert "vae" not in loads[1] + assert pipeline.get_module("vae") is None + assert pipeline.get_module("transformer") is not None diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py index 4e35bf66fe..6202e11d02 100644 --- a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -51,7 +51,6 @@ def normalize_latents(self, latents): return latents stage = MiniMaxH3LatentPreparationStage( - transformer=SimpleNamespace(patch_size=(1, 1, 1)), vae=VAE(), audio_vae=None, scheduler=None, @@ -61,7 +60,10 @@ def normalize_latents(self, latents): media_type="video", frames=np.zeros((22, 16, 16, 3), dtype=np.uint8), ) - args = SimpleNamespace(vae_parallel_encode=False) + args = SimpleNamespace( + vae_parallel_encode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))), + ) rows = stage._encode_visual_rows([reference], torch.device("cpu"), args) assert observed["pixels"].dtype == torch.uint8 @@ -95,9 +97,15 @@ def decode_to_pixels(self, decoded_latents, output): output.fill_(0.25) monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu")) - result = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward( + result = MiniMaxH3VideoDecodingStage(VAE()).forward( batch, - SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=False, vae_parallel_decode=False), + SimpleNamespace( + output_type="pil", + pin_cpu_memory=False, + vae_cpu_offload=False, + vae_parallel_decode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))), + ), ) torch.testing.assert_close(observed["latents"], latents) @@ -120,7 +128,7 @@ def to(self, device): monkeypatch.setattr(minimax_h3_decoding, "get_world_group", lambda: SimpleNamespace(is_first_rank=False)) args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True, vae_parallel_decode=False) - video = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace()).forward(ForwardBatch(data_type="video"), args) + video = MiniMaxH3VideoDecodingStage(VAE()).forward(ForwardBatch(data_type="video"), args) assert video.output.shape == (0, 3, 0, 0, 0) audio_batch = ForwardBatch(data_type="audio", latents=torch.zeros(1), audio_latents=torch.zeros(1)) @@ -161,11 +169,14 @@ def fake_parallel(vae, latents, output, group, strategy): monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu")) monkeypatch.setattr(minimax_h3_decoding, "model_parallel_is_initialized", lambda: True) monkeypatch.setattr(minimax_h3_decoding, "decode_to_pixels_parallel", fake_parallel) - args = SimpleNamespace(output_type="pil", - pin_cpu_memory=False, - vae_cpu_offload=False, - vae_parallel_decode=True, - vae_parallel_decode_strategy="gather") + args = SimpleNamespace( + output_type="pil", + pin_cpu_memory=False, + vae_cpu_offload=False, + vae_parallel_decode=True, + vae_parallel_decode_strategy="gather", + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))), + ) for rank, is_first in ((0, True), (2, False)): monkeypatch.setattr( @@ -175,7 +186,7 @@ def fake_parallel(vae, latents, output, group, strategy): rank_in_group=rank)) batch = ForwardBatch(data_type="video", latents=rows.clone(), raw_latent_shape=latent_shape) batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape) - result = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward(batch, args) + result = MiniMaxH3VideoDecodingStage(VAE()).forward(batch, args) if is_first: assert result.output.shape == (1, 3, 5, 16, 16) assert torch.all(result.output == 0.5) diff --git a/fastvideo/tests/vaes/test_minimax_h3_taeh3.py b/fastvideo/tests/vaes/test_minimax_h3_taeh3.py new file mode 100644 index 0000000000..66f1aedcb9 --- /dev/null +++ b/fastvideo/tests/vaes/test_minimax_h3_taeh3.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU contracts for the CUDA TAEH3 preview decoder.""" +from __future__ import annotations + +import torch + +from fastvideo.models.vaes.minimax_h3_taeh3 import TorchTAEH3Decoder, ensure_taeh3_checkpoint, taeh3_decoded_pixel_shape + + +def test_taeh3_pixel_shape_matches_h3_temporal_contract() -> None: + assert taeh3_decoded_pixel_shape((1, 24, 37, 30, 52)) == (1, 3, 124, 480, 832) + assert taeh3_decoded_pixel_shape((1, 24, 2, 4, 4)) == (1, 3, 5, 64, 64) + + +def test_taeh3_chunk_sizes_agree_on_cpu() -> None: + checkpoint = ensure_taeh3_checkpoint() + decoder = TorchTAEH3Decoder(checkpoint, dtype=torch.float32) + latents = torch.randn(1, 7, 24, 4, 4) + full = decoder.decode_ntchw(latents, chunk_size=7) + chunked = decoder.decode_ntchw(latents, chunk_size=3) + torch.testing.assert_close(full, chunked, atol=1e-5, rtol=1e-5) + assert full.shape == (1, 22, 3, 64, 64) diff --git a/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py b/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py index 57e68a7426..2bd39b9c0a 100644 --- a/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py +++ b/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py @@ -94,8 +94,13 @@ def _build_operation(args, vae, device): from fastvideo.pipelines.pipeline_batch_info import ForwardBatch patch_size = MiniMaxH3Config().arch_config.patch_size - transformer = SimpleNamespace(patch_size=patch_size) - runtime_args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True) + runtime_args = SimpleNamespace( + output_type="pil", + pin_cpu_memory=False, + vae_cpu_offload=True, + vae_parallel_encode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=patch_size)), + ) if args.operation == "encode": if int(os.environ.get("WORLD_SIZE", "1")) != 1: @@ -107,7 +112,6 @@ def _build_operation(args, vae, device): dtype=np.uint8, ) stage = MiniMaxH3LatentPreparationStage( - transformer=transformer, vae=vae, audio_vae=None, scheduler=None, @@ -118,7 +122,7 @@ def run_once(): reference = MiniMaxH3PreparedReference(media_type="video", frames=frames) vae.to(device) try: - return stage._encode_visual_rows([reference], device)[0] + return stage._encode_visual_rows([reference], device, runtime_args)[0] finally: vae.to("cpu") @@ -139,7 +143,7 @@ def run_once(): latents = torch.randn(latent_shape, generator=generator, device=device, dtype=torch.float32) rows = patchify_video_latents(latents, patch_size) layout = _make_layout(rows, latent_shape) - stage = MiniMaxH3VideoDecodingStage(vae, transformer) + stage = MiniMaxH3VideoDecodingStage(vae) def run_once(): batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape)