nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
This commit is contained in:
Aryan Kumar
2026-08-31 05:08:06 -07:00
committed by GitHub
co-authored by Aryan Kumar
parent f82d8be4bf
commit a28f2bab4b
9 changed files with 642 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/inference/basic/mlx_fasth3.py
+10
View File
@@ -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
+104
View File
@@ -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()