From a28f2bab4b41c2e714119ca704db1098a483e4fb Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Mon, 31 Aug 2026 05:08:06 -0700 Subject: [PATCH] [feat] Add an optional MLX TAEH3 preview decoder (#1794) Co-authored-by: Aryan Kumar --- docs/inference/mlx_taeh3.md | 157 ++++++++++++++ examples/inference/basic/mlx_fasth3.py | 10 + .../basic/mlx_h3_decode_benchmark.py | 104 ++++++++++ fastvideo/mlx_runtime/minimax_h3_pipeline.py | 44 +++- fastvideo/mlx_runtime/minimax_h3_taeh3.py | 195 ++++++++++++++++++ .../mlx/test_mlx_minimax_h3_fast_mode.py | 1 + .../mlx/test_mlx_minimax_h3_fast_spatial.py | 1 + .../test_mlx_minimax_h3_vsa_regressions.py | 1 + fastvideo/tests/mlx/test_mlx_taeh3.py | 132 ++++++++++++ 9 files changed, 642 insertions(+), 3 deletions(-) create mode 100644 docs/inference/mlx_taeh3.md create mode 100644 examples/inference/basic/mlx_h3_decode_benchmark.py create mode 100644 fastvideo/mlx_runtime/minimax_h3_taeh3.py create mode 100644 fastvideo/tests/mlx/test_mlx_taeh3.py diff --git a/docs/inference/mlx_taeh3.md b/docs/inference/mlx_taeh3.md new file mode 100644 index 000000000..b85626b2f --- /dev/null +++ b/docs/inference/mlx_taeh3.md @@ -0,0 +1,157 @@ +# Decode H3 previews with TAEH3 on MLX + +TAEH3 is an optional tiny video decoder for MiniMax H3. It replaces only +video reconstruction. The denoiser, sampler, resolution, frame count, and +audio decoder stay unchanged. The full H3 VAE remains the default. + +TAEH3 produces a different reconstruction. Fine fur, fabric, vegetation, +and surface textures can look softer. Use it for previews or when you accept +that tradeoff. It is not a lossless acceleration of the full VAE. + +## Generate a video + +Use your existing MLX FastH3 environment and converted checkpoint: + +```bash +python examples/inference/basic/mlx_fasth3.py \ + --model-root ~/models/FastH3-Preview-v0.2 \ + --mlx-checkpoint ~/models/FastH3-MLX/int6 \ + --prompt 'A red panda beside a mountain lake at sunrise.' \ + --height 480 --width 832 --num-frames 124 --steps 4 --seed 2027 \ + --video-decode-backend taeh3 --vae-dtype fp32 \ + --output-path video_samples/taeh3_preview.mp4 +``` + +The first run downloads a 22.7 MB safetensors checkpoint from an immutable +upstream revision and verifies its SHA-256 digest. No remote Python code runs. +The cache is `~/.cache/fastvideo/taehv/taeh3.safetensors`. +Use `--taeh3-checkpoint /path/to/taeh3.safetensors` for offline use or a custom +trained checkpoint. Custom files must match the decoder architecture; they +are not required to match the upstream digest. + +| Option | Default | Behavior | +| --- | --- | --- | +| `--video-decode-backend` | `h3-vae` | Select `taeh3` for approximate decoding. | +| `--taeh3-checkpoint` | Unset | Use the pinned upstream checkpoint from cache. | +| `--taeh3-chunk-size` | `5` | Latent frames per execution chunk. Smaller chunks reduce feature memory. | +| `--vae-dtype` | `fp32` | Decoder computation dtype. FP16 and BF16 are separate numerical tradeoffs. | + +`--tiled-video-decode` controls the full VAE only. TAEH3 uses the whole spatial +canvas and bounded temporal chunks. Its memory blocks carry state across +chunks. The pipeline reports `video_decode_backend`, decode timing, and MLX +peak memory alongside the existing generation metrics. + +The mode composes with `--fast-spatial` and temporal `--fast`. Those options +change the denoising workload and have additional quality costs. A TAEH3-only +measurement does not establish the quality or speed of a combined mode. + +## Latent contract + +The native decoder reads normalized diffusion latents in NCTHW layout through +`decode_latents_taeh3_mlx`. Do not apply the full VAE's mean, standard deviation, +or pixel denormalization. Its 24 latent channels reconstruct RGB at 16 times +the latent spatial dimensions. + +H3 uses latent lengths `5*k-3`, such as 2, 7, and 37. TAEH3 removes three raw +frames from each group of 20 decoder outputs. Thus 37 latent frames produce +124 RGB frames. The port validates that contract before returning output. + +## Provenance and validation + +Architecture and weights come from Ollin Boer Bohan's MIT-licensed +[TAEHV H3 implementation](https://github.com/madebyollin/taehv/commit/62f7591f59dfbb4c3c02b7a621d180a9eeaba26c). +The [Aryan fork](https://github.com/aryan5v/taehv/tree/aryan/first-class-taeh3) +adds an explicit `TAEH3` API and checkpoint-loading tests. The fork is not an +official MiniMax release and does not contain newly trained H3 weights. + +Run the numerical tests against a local TAEHV checkout containing the released +weights: + +```bash +TAEH3_REFERENCE_DIR=/path/to/taehv \ + python -m pytest fastvideo/tests/mlx/test_mlx_taeh3.py -q +``` + +Tests compare MLX FP32 with upstream sequential FP32 and parallel FP64 at +`atol=1e-5, rtol=1e-5`. The initial parallel CPU FP32 comparison failed that +strict gate, reaching about `4e-5` maximum error on a 37-latent small fixture. +CPU convolution rounding changes with its batch size. The FP64 reference +and sequential FP32 checks distinguish this from a temporal chunking error. +The original failed comparison is not reported as a pass. + +Passing these tests means the MLX port agrees with the tiny decoder within +the specified tolerance. It does not mean TAEH3 matches the full H3 VAE. + +## Compare decoders without another denoising run + +Save the normalized packed `video_rows` returned by `pipeline.denoise` as +`np.savez("latents.npz", video=video_rows)`. Then run: + +```bash +python examples/inference/basic/mlx_h3_decode_benchmark.py \ + --latents latents.npz \ + --model-root ~/models/FastH3-Preview-v0.2 \ + --mlx-checkpoint ~/models/FastH3-MLX/int6 \ + --height 480 --width 832 --num-frames 124 \ + --output-dir outputs/taeh3_comparison +``` + +Use a fresh output directory. The benchmark writes the first decoded frame +arrays and a JSON report with the input digest, MLX version, device, per-trial +latency, MLX peak active memory, lifetime process peak RSS, and swap snapshots. +Decoder loading is included; first-time checkpoint downloading is excluded. +Only run one MLX workload at a time. `--repeats` reverses the decoder order on +alternate trials. Do not treat two memory counters as additive or infer zero +page-outs from unchanged swap snapshots. + +## Measured decoder results + +On an Apple M4 Max with 36 GB unified memory, MLX 0.32.2, FP32 decoding, +and five-latent execution chunks: + +| Workload | Full H3 VAE | TAEH3 | +| --- | --- | --- | +| Saved 37-frame latents to 124 RGB frames, 832x480 | 107.90 s | 1.44 s | +| MLX peak active memory for that decode | 11.03 GiB | 3.62 GiB | + +These are one matched pair using the same seed-2027 production latent file, +including decoder loading. Both swap snapshots stayed unchanged. The decoded +images differ: PSNR against the full VAE was 29.86 dB, and inspected frames +showed softer fine detail. The approximately 75x ratio applies only to this +decoder comparison, not the entire generation pipeline. + +Eight additional TAEH3-only decodes measured 0.96 s on first use and a 0.98 s median across seven warm trials, ranging from 0.96 to 0.99 s. Decoder construction and weight loading were included in each trial. The full VAE was not repeated eight times. + +### Native resolution with temporal fast + +A separate uncached run with `--fast --video-decode-backend taeh3`, without +spatial fast, completed in **205.47 s wall time**. It kept the native 832x480 +canvas, denoised 73 source frames, and used RIFE to produce 124 output frames +with full-duration audio. Seed 2027, four steps, dense attention, INT6. + +Prompt encoding took 16.66 s, denoising 181.14 s, TAEH3 decoding 0.64 s, RIFE +5.23 s, audio decoding 0.73 s, and muxing 0.40 s. Peak denoise MLX allocation +was 17.87 GiB. System swap use rose from 1231.19 to 2753.75 MiB. +The output has 124 H.264 frames at 832x480 and stereo AAC at 32 kHz. + +This is one combined-mode measurement. It preserves spatial resolution but +still combines frame interpolation with approximate decoding. Frame samples +retain more fine detail than the spatial fast experiment below. Motion and +speech need human review; do not infer native dense generation parity. + +### Spatial fast experiment, not the preferred quality path + +A separate uncached generation with `--fast-spatial --video-decode-backend +taeh3` produced a 124-frame, 832x480 MP4 with stereo audio in **95.75 s wall +time**. Seed 2028, four denoising calls, INT6 weights, dense attention, and no +temporal fast mode were used. The internal canvas was 416x256, then cropped +and upscaled to the requested output size. + +That run spent 16.66 s encoding the prompt, 77.05 s denoising, 0.34 s decoding +video, 0.30 s upscaling, 0.73 s decoding audio, and 0.40 s muxing. Peak denoise +MLX allocation was 16.85 GiB. System swap use rose from 1041.94 to 1247.19 MiB; +this measurement does not attribute that increase to a particular process. +It is one end-to-end result, not a repeated benchmark or native-resolution +quality comparison. The reduced canvas visibly loses detail, especially in +the opening frames. Speech intelligibility and motion quality need human +review before treating this combination as a final-output preset. diff --git a/examples/inference/basic/mlx_fasth3.py b/examples/inference/basic/mlx_fasth3.py index d13bdcd03..c85d72474 100644 --- a/examples/inference/basic/mlx_fasth3.py +++ b/examples/inference/basic/mlx_fasth3.py @@ -98,6 +98,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--rife-weights-dir", type=Path, default=None, help="optional local mlx-community/RIFE-4.25 snapshot") parser.add_argument("--vae-dtype", choices=("fp32", "fp16", "bf16"), default="fp32") + parser.add_argument("--video-decode-backend", choices=("h3-vae", "taeh3"), default="h3-vae", + help="full H3 VAE or approximate TAEH3 preview decoder; audio is unchanged") + parser.add_argument("--taeh3-checkpoint", type=Path, default=None, + help="local TAEH3 safetensors; otherwise download hash-verified upstream weights") + parser.add_argument("--taeh3-chunk-size", type=int, default=5, + help="latent frames per TAEH3 chunk; causal memory persists across chunks") parser.add_argument("--prompt-cache-dir", type=Path, default=None, help="directory for reusable prompt embedding caches") parser.add_argument( @@ -169,6 +175,9 @@ def main() -> None: model_root=args.model_root, mlx_dit_checkpoint=args.mlx_checkpoint, vae_dtype=args.vae_dtype, + video_decode_backend=args.video_decode_backend, + taeh3_checkpoint=args.taeh3_checkpoint, + taeh3_chunk_size=args.taeh3_chunk_size, prompt_cache_dir=args.prompt_cache_dir, ) result = pipeline.generate( @@ -201,6 +210,7 @@ def main() -> None: "timings_s": {k: round(v, 2) for k, v in result.timings.items()}, "peak_memory_gib": {k: round(v, 2) for k, v in result.peak_memory_gib.items()}, "vsa": result.vsa, + "video_decode_backend": result.video_decode_backend, "audio_samples": int(result.waveform.shape[-1]), }, indent=2)) diff --git a/examples/inference/basic/mlx_h3_decode_benchmark.py b/examples/inference/basic/mlx_h3_decode_benchmark.py new file mode 100644 index 000000000..76d30c74b --- /dev/null +++ b/examples/inference/basic/mlx_h3_decode_benchmark.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Compare H3 decoders on saved normalized video rows without rerunning denoise. + +The input NPZ must contain a ``video`` array of packed diffusion rows. Geometry +and the DiT manifest must match the generation that produced those rows. +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import platform +import resource +import subprocess +import time +from pathlib import Path + +import numpy as np + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--latents", type=Path, required=True) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--mlx-checkpoint", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--num-frames", type=int, default=124) + parser.add_argument("--backends", nargs="+", choices=("h3-vae", "taeh3"), default=["h3-vae", "taeh3"]) + parser.add_argument("--taeh3-checkpoint", type=Path) + parser.add_argument("--taeh3-chunk-size", type=int, default=5) + parser.add_argument("--vae-dtype", choices=("fp32", "fp16", "bf16"), default="fp32") + parser.add_argument("--repeats", type=int, default=1) + args = parser.parse_args() + if args.repeats < 1: + parser.error("--repeats must be positive") + if args.output_dir.exists(): + parser.error("--output-dir must be a new directory to preserve previous results") + + import mlx.core as mx + from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline, _cleanup_mlx + from fastvideo.mlx_runtime.minimax_h3_taeh3 import ensure_taeh3_checkpoint + + with np.load(args.latents) as archive: + rows = archive["video"] + if not np.isfinite(rows).all(): + raise ValueError("The saved video rows contain non-finite values") + args.output_dir.mkdir(parents=True) + checksum = hashlib.sha256() + with args.latents.open("rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + checksum.update(chunk) + checkpoint = ensure_taeh3_checkpoint(args.taeh3_checkpoint) if "taeh3" in args.backends else None + report = { + "mlx": mx.__version__, + "platform": platform.platform(), + "device": mx.device_info(), + "latents": str(args.latents.resolve()), + "latents_sha256": checksum.hexdigest(), + "geometry": [args.height, args.width, args.num_frames], + "dtype": args.vae_dtype, + "chunk_size": args.taeh3_chunk_size, + "checkpoint_download_excluded": True, + "decoder_loading_included": True, + "trials": [], + } + for repeat in range(args.repeats): + # Reverse each paired trial's order to expose warmup/order effects. + order = args.backends if repeat % 2 == 0 else args.backends[::-1] + for backend in order: + pipeline = MiniMaxH3MLXPipeline(model_root=args.model_root, + mlx_dit_checkpoint=args.mlx_checkpoint, + video_decode_backend=backend, + taeh3_checkpoint=checkpoint if backend == "taeh3" else None, + taeh3_chunk_size=args.taeh3_chunk_size, + vae_dtype=args.vae_dtype) + _cleanup_mlx() + mx.reset_peak_memory() + before = subprocess.check_output(["sysctl", "-n", "vm.swapusage"], text=True).strip() + started = time.perf_counter() + frames = pipeline.decode_video(rows, height=args.height, width=args.width, num_frames=args.num_frames) + elapsed = time.perf_counter() - started + trial = { + "repeat": repeat, + "backend": backend, + "decode_s": elapsed, + "shape": list(frames.shape), + "peak_active_gib": mx.get_peak_memory() / 2**30, + "process_lifetime_rss_peak_gib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**30, + "swap_before": before, + "swap_after": subprocess.check_output(["sysctl", "-n", "vm.swapusage"], text=True).strip(), + } + if repeat == 0: + np.save(args.output_dir / f"{backend}_frames.npy", frames) + report["trials"].append(trial) + (args.output_dir / "report.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(trial), flush=True) + del frames, pipeline + _cleanup_mlx() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/mlx_runtime/minimax_h3_pipeline.py b/fastvideo/mlx_runtime/minimax_h3_pipeline.py index 4310ee8c3..5f517855f 100644 --- a/fastvideo/mlx_runtime/minimax_h3_pipeline.py +++ b/fastvideo/mlx_runtime/minimax_h3_pipeline.py @@ -71,6 +71,7 @@ class GenerationResult: timings: dict[str, float] = field(default_factory=dict) peak_memory_gib: dict[str, float] = field(default_factory=dict) vsa: dict[str, Any] = field(default_factory=dict) + video_decode_backend: str = "h3-vae" @dataclass(frozen=True) @@ -305,6 +306,9 @@ class MiniMaxH3MLXPipeline: conditioner_dir: str | Path | None = None, tokenizer_dir: str | Path | None = None, metal_wired_limit_gib: float | None = None, + video_decode_backend: str = "h3-vae", + taeh3_checkpoint: str | Path | None = None, + taeh3_chunk_size: int = 5, ) -> None: import mlx.core as mx @@ -322,6 +326,15 @@ class MiniMaxH3MLXPipeline: self.model_root = Path(model_root) self.dit_checkpoint = Path(mlx_dit_checkpoint) self.vae_dtype = vae_dtype + if video_decode_backend not in ("h3-vae", "taeh3"): + raise ValueError(f"Unknown H3 video decoder: {video_decode_backend}") + if taeh3_chunk_size < 1: + raise ValueError("taeh3_chunk_size must be positive.") + if taeh3_checkpoint is not None and video_decode_backend != "taeh3": + raise ValueError("taeh3_checkpoint requires video_decode_backend='taeh3'.") + self.video_decode_backend = video_decode_backend + self.taeh3_checkpoint = taeh3_checkpoint + self.taeh3_chunk_size = taeh3_chunk_size self.prompt_cache_dir = Path(prompt_cache_dir) if prompt_cache_dir else None self.conditioner_dir = Path(conditioner_dir) if conditioner_dir else self.model_root / "text_encoder" self.tokenizer_dir = Path(tokenizer_dir) if tokenizer_dir else self.model_root / "tokenizer" @@ -344,7 +357,7 @@ class MiniMaxH3MLXPipeline: missing.append(str(self.dit_checkpoint)) vae_dir = self.model_root / "vae" audio_dir = self.model_root / "audio_vae" - if not (vae_dir.exists() and any(vae_dir.glob("*.safetensors"))): + if self.video_decode_backend == "h3-vae" and not (vae_dir.exists() and any(vae_dir.glob("*.safetensors"))): missing.append(str(vae_dir)) if not (audio_dir.exists() and any(audio_dir.glob("*.safetensors"))): missing.append(str(audio_dir)) @@ -562,6 +575,22 @@ class MiniMaxH3MLXPipeline: from fastvideo.mlx_runtime.minimax_h3_video_vae import mlx_h3_video_vae_from_dir geometry = self.resolve_geometry(height, width, num_frames, enforce_duration=False) + if self.video_decode_backend == "taeh3": + from fastvideo.mlx_runtime.minimax_h3_taeh3 import decode_latents_taeh3_mlx + + if self._dit_in_channels != 24: + raise ValueError("TAEH3 requires a 24-channel H3 checkpoint.") + latents = unpatchify_video_tokens(video_rows, geometry["latent_frame_count"], geometry["latent_height"], + geometry["latent_width"], self._dit_in_channels, self._dit_patch_size) + pixels = decode_latents_taeh3_mlx(latents, + checkpoint_path=self.taeh3_checkpoint, + dtype=self.vae_dtype, + chunk_size=self.taeh3_chunk_size) + frames = (pixels[0] * 255.0).astype(np.uint8) + if frames.shape != (geometry["num_frames"], height, width, 3): + raise RuntimeError(f"TAEH3 produced unexpected frame shape: {frames.shape}") + _cleanup_mlx() + return frames vae = mlx_h3_video_vae_from_dir(self.model_root / "vae", include_encoder=False, storage_dtype=self.vae_dtype) expected_height = height // vae.spatial_compression_ratio expected_width = width // vae.spatial_compression_ratio @@ -775,6 +804,14 @@ class MiniMaxH3MLXPipeline: spatial_plan, ) + if self.video_decode_backend == "taeh3": + from fastvideo.mlx_runtime.minimax_h3_taeh3 import ensure_taeh3_checkpoint + + started = time.perf_counter() + ensure_taeh3_checkpoint(self.taeh3_checkpoint) + timings["decoder_prepare_s"] = time.perf_counter() - started + logger.warning("TAEH3 is an approximate preview decoder; reconstruction differs from the full H3 VAE.") + _reset_peak_memory() started = time.perf_counter() text_rows, token_tags = self.encode_prompt(prompt) @@ -872,8 +909,8 @@ class MiniMaxH3MLXPipeline: video_path = self.mux(frames, waveform, output_path) timings["mux_s"] = time.perf_counter() - started timings["generate_s"] = sum( - timings.get(key, 0.0) for key in ("condition_s", "denoise_s", "video_decode_s", "rife_s", - "spatial_upsample_s", "audio_decode_s", "mux_s")) + timings.get(key, 0.0) for key in ("decoder_prepare_s", "condition_s", "denoise_s", "video_decode_s", + "rife_s", "spatial_upsample_s", "audio_decode_s", "mux_s")) result = GenerationResult( video_path=str(video_path), @@ -882,6 +919,7 @@ class MiniMaxH3MLXPipeline: sample_rate=32000, timings=timings, peak_memory_gib=peaks, + video_decode_backend=self.video_decode_backend, vsa=getattr(self, "last_vsa_stats", None) or { "enabled": vsa_config.enabled, "checkpoint_vsa_capable": mlx_h3_checkpoint_vsa_capable(self.dit_checkpoint), diff --git a/fastvideo/mlx_runtime/minimax_h3_taeh3.py b/fastvideo/mlx_runtime/minimax_h3_taeh3.py new file mode 100644 index 000000000..f4138037d --- /dev/null +++ b/fastvideo/mlx_runtime/minimax_h3_taeh3.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Optional, approximate MiniMax H3 tiny decoder, executed entirely in MLX. + +Architecture and temporal mapping adapted from madebyollin/taehv at +62f7591f59dfbb4c3c02b7a621d180a9eeaba26c (MIT, Ollin Boer Bohan). +See fastvideo/third_party/taehv/LICENSE. Weights remain in the upstream format. +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 numpy as np + +TAEH3_REVISION = "62f7591f59dfbb4c3c02b7a621d180a9eeaba26c" +TAEH3_URL = f"https://raw.githubusercontent.com/madebyollin/taehv/{TAEH3_REVISION}/safetensors/taeh3.safetensors" +TAEH3_SHA256 = "4fd022bfcab08772fe0536b17ea1a3bbb5625be11e397868d1c5d891863d4c13" + + +def ensure_taeh3_checkpoint(checkpoint_path: str | Path | None = None) -> Path: + """Fetch only pinned weights, atomically; never download executable code. + + Explicit local safetensors paths may contain custom trained weights. + Managed cache entries must match the upstream SHA-256 digest. + """ + 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 MLX TAEH3 decoder requires a .safetensors checkpoint.") + return path + path = Path.home() / ".cache/fastvideo/taehv/taeh3.safetensors" + + def verify(candidate: Path) -> None: + with candidate.open("rb") as handle: + hasher = hashlib.sha256() + 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) + return path + + +class MLXTAEH3Decoder: + """Decode H3 NTCHW latents with bounded temporal feature memory. + + Chunks carry the last input of every MemBlock into the next chunk. + Chunking limits activation memory without resetting the causal state. + The full H3 VAE remains the quality reference, not a numerical oracle for + this independently trained tiny decoder. + """ + + def __init__(self, checkpoint_path: str | Path, *, dtype: str = "fp32") -> None: + import mlx.core as mx + + if dtype not in ("fp32", "fp16", "bf16"): + raise ValueError(f"Unsupported TAEH3 dtype: {dtype}") + self.dtype = {"fp32": mx.float32, "fp16": mx.float16, "bf16": mx.bfloat16}[dtype] + raw = mx.load(str(checkpoint_path)) + expected: 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[f"{prefix}.weight"] = (channels, channels * (2 if layer == 0 else 1), 3, 3) + expected[f"{prefix}.bias"] = (channels, ) + actual = {key for key in raw if key.startswith("decoder.")} + if actual != set(expected): + raise ValueError(f"TAEH3 decoder keys mismatch: missing={set(expected) - actual}, " + f"unexpected={actual - set(expected)}") + self.weights = {} + for key, shape in expected.items(): + value = raw[key] + if tuple(value.shape) != shape: + raise ValueError(f"TAEH3 weight {key} has shape {value.shape}, expected {shape}") + # Keep convolutions and features channels-last throughout execution. + value = value.transpose(0, 2, 3, 1) if value.ndim == 4 else value + self.weights[key] = mx.contiguous(value.astype(self.dtype)) + mx.eval(*self.weights.values()) + + def _conv(self, x: Any, name: str) -> Any: + import mlx.core as mx + + weight = self.weights[f"{name}.weight"] + y = mx.conv2d(x, weight, padding=weight.shape[1] // 2) + bias = self.weights.get(f"{name}.bias") + return y if bias is None else y + bias + + def _chunk(self, x: Any, memory: dict[int, Any]) -> Any: + import mlx.core as mx + + n, t, h, w, c = x.shape + x = mx.maximum(self._conv(mx.tanh(x.reshape(n * t, h, w, c) / 3.0) * 3.0, "decoder.1"), 0) + 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: + _, h, w, c = x.shape + sequence = x.reshape(n, -1, h, w, c) + previous = memory.get(index) + if previous is None: + previous = mx.zeros_like(sequence[:, :1]) + past = mx.concatenate([previous, sequence[:, :-1]], axis=1).reshape(x.shape) + # Preserve the final feature frame for the next execution chunk. + memory[index] = mx.contiguous(sequence[:, -1:]) + y = mx.concatenate([x, past], axis=-1) + for layer in (0, 2, 4): + y = self._conv(y, f"decoder.{index}.conv.{layer}") + if layer != 4: + y = mx.maximum(y, 0) + x = mx.maximum(x + y, 0) + nt, h, w, c = x.shape + x = mx.broadcast_to(x.reshape(nt, h, 1, w, 1, c), (nt, h, 2, w, 2, c)) + x = x.reshape(nt, h * 2, w * 2, c) + x = self._conv(x, f"decoder.{grow}.conv") + # Torch TGrow splits its channel-major output into temporal frames. + x = x.reshape(nt, h * 2, w * 2, stride, c).transpose(0, 3, 1, 2, 4) + x = self._conv(x.reshape(nt * stride, h * 2, w * 2, c), f"decoder.{projection}") + x = self._conv(mx.maximum(x, 0), "decoder.22") + nt, h, w, _ = x.shape + # Torch pixel_shuffle stores channels as (RGB, patch_y, patch_x). + x = x.reshape(nt, h, w, 3, 2, 2).transpose(0, 1, 4, 2, 5, 3) + return mx.clip(x.reshape(n, -1, h * 2, w * 2, 3), 0, 1) + + def decode_ntchw(self, latents: Any, *, chunk_size: int = 5) -> Any: + """Return NTCHW RGB in [0, 1] for H3's valid 5*k-3 latent lengths.""" + import mlx.core as mx + + 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 {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.astype(self.dtype).transpose(0, 1, 3, 4, 2) + memory: dict[int, Any] = {} + frames = [] + for start in range(0, x.shape[1], chunk_size): + decoded = self._chunk(x[:, start:start + chunk_size], memory) + # Remove three raw frames from each five-latent (20-frame) group. + keep = [i for i in range(decoded.shape[1]) if (start * 4 + i) % 20 >= 3] + decoded = decoded[:, mx.array(keep, dtype=mx.int32)] + mx.eval(decoded, *memory.values()) + frames.append(decoded) + return mx.concatenate(frames, axis=1).transpose(0, 1, 4, 2, 3) + + +def decode_latents_taeh3_mlx(latents: np.ndarray, + *, + checkpoint_path: str | Path | None = None, + dtype: str = "fp32", + chunk_size: int = 5) -> np.ndarray: + """Decode normalized NCTHW diffusion latents into NTHWC float RGB.""" + import mlx.core as mx + + if latents.ndim != 5: + raise ValueError(f"Expected NCTHW latents, got {latents.shape}") + checkpoint = ensure_taeh3_checkpoint(checkpoint_path) + decoder = MLXTAEH3Decoder(checkpoint, dtype=dtype) + output = decoder.decode_ntchw(mx.array(latents.transpose(0, 2, 1, 3, 4)), chunk_size=chunk_size) + mx.eval(output) + return np.asarray(output.astype(mx.float32)).transpose(0, 1, 3, 4, 2) diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py index f3dbe1212..5cc04cc29 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py @@ -119,6 +119,7 @@ def test_default_wired_limit_scales_down_and_keeps_tested_cap() -> None: def test_mux_cleans_temporary_files_after_ffmpeg_failure(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) + pipeline.video_decode_backend = "h3-vae" output = tmp_path / "output.mp4" frames = np.zeros((1, 2, 2, 3), dtype=np.uint8) waveform = np.zeros((2, 32), dtype=np.float32) diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py index 66964da53..370aa1c54 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py @@ -91,6 +91,7 @@ def _generate_with_mocked_phases(monkeypatch, tmp_path, **generate_kwargs): calls: dict = {} pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) + pipeline.video_decode_backend = "h3-vae" pipeline.dit_checkpoint = tmp_path monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline._validate_checkpoint_step_ladder", diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py index 4e4b44f75..e2a4e3778 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py @@ -163,6 +163,7 @@ def test_dense_generate_ignores_unused_vsa_parameters(tmp_path, monkeypatch): from fastvideo.mlx_runtime import minimax_h3_pipeline as pipeline_mod pipeline = object.__new__(pipeline_mod.MiniMaxH3MLXPipeline) + pipeline.video_decode_backend = "h3-vae" pipeline.dit_checkpoint = tmp_path monkeypatch.setattr(pipeline_mod, "_validate_checkpoint_step_ladder", lambda *a: None) def stop_at_media_preflight(**kwargs): diff --git a/fastvideo/tests/mlx/test_mlx_taeh3.py b/fastvideo/tests/mlx/test_mlx_taeh3.py new file mode 100644 index 000000000..788379997 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_taeh3.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +"""TAEH3 port parity. Requires Apple MLX and an explicit upstream checkout. + +TAEH3_REFERENCE_DIR=/path/to/taehv python -m pytest fastvideo/tests/mlx/test_mlx_taeh3.py +No weights or source code are downloaded by this test. +""" +import importlib.util +import os +from pathlib import Path + +import numpy as np +import pytest + +mx = pytest.importorskip('mlx.core') +from fastvideo.mlx_runtime.minimax_h3_taeh3 import MLXTAEH3Decoder, ensure_taeh3_checkpoint + + +@pytest.fixture(scope='module') +def models(): + torch = pytest.importorskip('torch') + location = os.environ.get('TAEH3_REFERENCE_DIR') + if not location: + pytest.skip('Set TAEH3_REFERENCE_DIR to an upstream taehv checkout') + root = Path(location) + spec = importlib.util.spec_from_file_location('taeh3_reference', root / 'taehv.py') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + # FP64 avoids the CPU FP32 convolution's batch-size-dependent rounding. + # The original FP32 comparison is recorded separately in the validation report. + reference = module.TAEHV(str(root / 'taeh3.pth')).double().eval() + actual = MLXTAEH3Decoder(root / 'safetensors/taeh3.safetensors') + return torch, reference, actual + + +@pytest.mark.parametrize('reference_mode', ['serial_fp32', 'parallel_fp64']) +@pytest.mark.parametrize('batch,time,chunk', [(1, 2, 5), (1, 7, 1), (2, 7, 3), (1, 12, 5), (1, 37, 5)]) +def test_mlx_taeh3_matches_upstream(models, batch, time, chunk, reference_mode): + torch, reference, actual = models + latent = np.random.default_rng(123).standard_normal((batch, time, 24, 2, 3)).astype(np.float32) + with torch.no_grad(): + dtype = torch.float64 if reference_mode == "parallel_fp64" else torch.float32 + reference = reference.to(dtype=dtype) + expected = reference.decode_video(torch.from_numpy(latent).to(dtype=dtype), + parallel=reference_mode == "parallel_fp64", + show_progress_bar=False).float().numpy() + output = np.asarray(actual.decode_ntchw(mx.array(latent), chunk_size=chunk)) + assert output.shape == expected.shape + assert np.isfinite(output).all() + np.testing.assert_allclose(output, expected, atol=1e-5, rtol=1e-5) + + +def test_invalid_latent_length_rejected(models): + _, _, actual = models + with pytest.raises(ValueError, match='5\\*k-3'): + actual.decode_ntchw(mx.zeros((1, 5, 24, 2, 2))) + + +def test_cache_tampering_rejected(monkeypatch, tmp_path): + monkeypatch.setattr(Path, 'home', lambda: tmp_path) + cache = tmp_path / '.cache/fastvideo/taehv/taeh3.safetensors' + cache.parent.mkdir(parents=True) + cache.write_bytes(b'corrupt') + with pytest.raises(RuntimeError, match='SHA-256'): + ensure_taeh3_checkpoint() + + +def test_local_checkpoint_requires_safetensors(tmp_path): + path = tmp_path / 'weights.pth' + path.touch() + with pytest.raises(ValueError, match='safetensors'): + ensure_taeh3_checkpoint(path) + + +def test_chunk_size_does_not_reset_temporal_state(models): + _, _, actual = models + latent = mx.array(np.random.default_rng(11).standard_normal((1, 12, 24, 2, 3)).astype(np.float32)) + expected = np.asarray(actual.decode_ntchw(latent, chunk_size=12)) + for size in (1, 3, 5): + np.testing.assert_array_equal(np.asarray(actual.decode_ntchw(latent, chunk_size=size)), expected) + + +def test_strided_latents_match_contiguous(models): + _, _, actual = models + latent = mx.array(np.random.default_rng(9).standard_normal((1, 7, 24, 2, 6)).astype(np.float32))[:, :, :, :, ::2] + expected = actual.decode_ntchw(mx.contiguous(latent)) + np.testing.assert_array_equal(np.asarray(actual.decode_ntchw(latent)), np.asarray(expected)) + + +def test_incompatible_checkpoint_fails_before_decode(tmp_path): + path = tmp_path / 'bad.safetensors' + mx.save_safetensors(str(path), {'decoder.1.weight': mx.zeros((1, 1, 1, 1))}) + with pytest.raises(ValueError, match='keys mismatch'): + MLXTAEH3Decoder(path) + + +def test_failed_download_does_not_leave_cache(monkeypatch, tmp_path): + import io + from fastvideo.mlx_runtime import minimax_h3_taeh3 as module + monkeypatch.setattr(Path, 'home', lambda: tmp_path) + monkeypatch.setattr(module.urllib.request, 'urlopen', lambda *args, **kwargs: io.BytesIO(b'corrupt')) + with pytest.raises(RuntimeError, match='SHA-256'): + ensure_taeh3_checkpoint() + assert list((tmp_path / '.cache/fastvideo/taehv').iterdir()) == [] + + +def test_pipeline_taeh3_skips_full_vae_and_denormalization(monkeypatch): + from fastvideo.mlx_runtime import minimax_h3_taeh3 as module + from fastvideo.mlx_runtime import minimax_h3_video_vae as full_vae + from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline + + pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) + pipeline.video_decode_backend = 'taeh3' + pipeline.taeh3_checkpoint = None + pipeline.taeh3_chunk_size = 5 + pipeline.vae_dtype = 'fp32' + pipeline._dit_in_channels = 24 + pipeline._dit_patch_size = (1, 2, 2) + + def fail(*args, **kwargs): + raise AssertionError('Full VAE must not load for TAEH3') + + def fake_decode(latents, **kwargs): + assert latents.shape == (1, 24, 2, 2, 2) + np.testing.assert_array_equal(latents, np.full_like(latents, 2.0)) + return np.full((1, 5, 32, 32, 3), 0.5, dtype=np.float32) + + monkeypatch.setattr(full_vae, 'mlx_h3_video_vae_from_dir', fail) + monkeypatch.setattr(module, 'decode_latents_taeh3_mlx', fake_decode) + frames = pipeline.decode_video(np.full((2, 96), 2.0, dtype=np.float32), height=32, width=32, num_frames=5) + assert frames.shape == (5, 32, 32, 3) + assert frames.dtype == np.uint8 + assert np.all(frames == 127)