105 lines
4.6 KiB
Python
105 lines
4.6 KiB
Python
# 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()
|