[feat] Add an optional MLX TAEH3 preview decoder (#1794)

Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>

Author: aryan5v

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.

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))
 

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()

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 @@ def __init__(
         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 @@ def __init__(
         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 @@ def _validate_inputs_before_loading(self) -> bool:
             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 @@ def decode_video(self,
         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 @@ def generate(
             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 @@ def generate(
         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 @@ def generate(
             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),