From a31d853951a8f4461b17f71db026feac5ef9c31c Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Mon, 31 Aug 2026 01:19:56 -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. --- .../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 | 20 +- .../minimax_h3/stages/minimax_h3_decoding.py | 25 ++- .../stages/minimax_h3_latent_preparation.py | 3 + .../inference/test_basic_fasth3_profile.py | 8 + fastvideo/tests/vaes/test_minimax_h3_taeh3.py | 22 ++ 10 files changed, 311 insertions(+), 9 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 55a4bac033..f1123327da 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -81,6 +81,9 @@ surfaces: VSA_tile_size: "VSA-H3 tile geometry request; model-specific optimization not yet represented in the typed public schema." 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." + 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 96284efe89..21bddfdb2d 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -167,6 +167,12 @@ is power-cycled. To avoid it: preferred). The CUDA pipeline now 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** (`--lazy-module-load --video-decode-backend taeh3`) is an + opt-in preview decoder. T2VA never materializes the 9.7 GiB video VAE. On this + box, alpine 768×1344×124 decoded in **2.4 s** versus **68 s** for the full VAE, + and one lazy-load T2VA generation finished in **224 s** end-to-end. 1761 plus + the full VAE still OOMs that canvas. 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 38c407cb99..828a80240f 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/examples/inference/basic/basic_fasth3.py @@ -54,6 +54,11 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser: "needs it, so peak memory is the largest overlapping set instead of the sum of every " "component. Enable when the model does not fit at load time; costs a reload per " "generation, so leave it off when it does fit") + 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("--profile", choices=("all", "strict"), default="all", @@ -235,6 +240,10 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig: "vae_parallel_decode": args.parallel_vae, "vae_parallel_decode_strategy": "gather", } + 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 74c9244de1..d8333e85b0 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -168,6 +168,13 @@ class FastVideoArgs: # fit. Inference only; training keeps every component resident. lazy_module_load: bool = False + # 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 @@ -723,6 +730,25 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "so peak memory is the largest overlapping set of components instead of their sum. Enable when a " "model does not fit at load time. Costs a reload per generation, so leave it off when it does fit.", ) + 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( "--pin-cpu-memory", 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 25dc999856..536a59098d 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -3,6 +3,7 @@ from __future__ import annotations +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.pipelines.basic.minimax_h3.stages import ( @@ -71,17 +72,22 @@ def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: if shift is None or float(shift) != expected_shift: raise ValueError(f"MiniMax-H3 {modality} scheduler must expose shift={expected_shift:g}, got {shift}.") - def _add_stages(self, *, ref2va: bool) -> None: + def _add_stages(self, fastvideo_args: FastVideoArgs, *, ref2va: bool) -> None: transformer = self.get_module("transformer") vae = self.get_module("vae") audio_vae = self.get_module("audio_vae") scheduler = self.get_module("scheduler") audio_scheduler = self.get_module("audio_scheduler") + use_taeh3 = getattr(fastvideo_args, "video_decode_backend", "h3-vae") == "taeh3" + # T2VA preview with TAEH3 only needs VAE geometry scalars. Passing the + # arch config keeps the 9.7 GiB ViT decoder from materializing. + video_geometry = MiniMaxH3VideoVAEArchConfig() if use_taeh3 and not ref2va else vae + decode_vae = None if use_taeh3 else vae self.add_stage( "input_preparation_stage", MiniMaxH3InputPreparationStage( - vae=vae, + vae=video_geometry, audio_vae=audio_vae if ref2va else None, ref2va=ref2va, ), @@ -99,7 +105,7 @@ def _add_stages(self, *, ref2va: bool) -> None: "latent_preparation_stage", MiniMaxH3LatentPreparationStage( transformer=transformer, - vae=vae, + vae=video_geometry, audio_vae=audio_vae, scheduler=scheduler, ref2va=ref2va, @@ -113,7 +119,7 @@ def _add_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=decode_vae, transformer=transformer)) self.add_stage("audio_decoding_stage", MiniMaxH3AudioDecodingStage(audio_vae=audio_vae)) @@ -121,8 +127,7 @@ class MiniMaxH3Pipeline(MiniMaxH3BasePipeline): """One-request joint video/stereo-audio pipeline for T2VA and FL2VA.""" def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - del fastvideo_args - self._add_stages(ref2va=False) + self._add_stages(fastvideo_args, ref2va=False) class MiniMaxH3RefPipeline(MiniMaxH3BasePipeline): @@ -131,8 +136,7 @@ class MiniMaxH3RefPipeline(MiniMaxH3BasePipeline): _extra_config_module_map = {"transformer": "transformer_ref"} def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - del fastvideo_args - self._add_stages(ref2va=True) + self._add_stages(fastvideo_args, ref2va=True) class MiniMaxH3ModularPipeline(MiniMaxH3Pipeline): 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..5c9ff5ec2c 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -58,7 +58,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): performance_component_metric = "vae_decode_time_s" - def __init__(self, vae: AutoencoderKLMiniMaxH3, transformer: Any) -> None: + def __init__(self, vae: AutoencoderKLMiniMaxH3 | None, transformer: Any) -> None: super().__init__() self.vae = vae self.transformer = transformer @@ -100,6 +100,29 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward self.transformer.patch_size, ) 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..e43929c818 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 @@ -176,6 +176,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) diff --git a/fastvideo/tests/inference/test_basic_fasth3_profile.py b/fastvideo/tests/inference/test_basic_fasth3_profile.py index e4f8bc2508..3e68a96bb4 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_profile.py +++ b/fastvideo/tests/inference/test_basic_fasth3_profile.py @@ -236,3 +236,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(tmp_path): + 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/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)