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

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

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

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)

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)

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",

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

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)