From f6166ede682c721de2e2c7ef6fa1f5fbe2d73d21 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Fri, 14 Aug 2026 14:36:42 +0700 Subject: [PATCH] Add VAE decode benchmark tool --- tools/benchmark_video_vae_decode.py | 173 ++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 tools/benchmark_video_vae_decode.py diff --git a/tools/benchmark_video_vae_decode.py b/tools/benchmark_video_vae_decode.py new file mode 100644 index 0000000..f8a05ac --- /dev/null +++ b/tools/benchmark_video_vae_decode.py @@ -0,0 +1,173 @@ +"""Benchmark MiniMax H3 video VAE decode variants on a saved latent.""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import time +from pathlib import Path + +import torch + + +def sync() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def timed(stage: str, rows: list[dict], fn): + sync() + start = time.perf_counter() + value = fn() + sync() + elapsed = time.perf_counter() - start + rows.append({"stage": stage, "seconds": elapsed}) + return value + + +def dtype_from_name(name: str) -> torch.dtype: + return { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + }[name] + + +def ffmpeg_command(loglevel: str, *parts: str) -> list[str]: + return ["ffmpeg", "-hide_banner", "-loglevel", loglevel, *parts] + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--latent", type=Path, required=True) + parser.add_argument("--output", type=Path) + parser.add_argument("--metrics", type=Path, required=True) + parser.add_argument("--implementation", choices=("direct", "upstream"), default="direct") + parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float32") + parser.add_argument("--no-tiling", action="store_true") + parser.add_argument("--tile-size", type=int) + parser.add_argument("--tile-overlap", type=int) + parser.add_argument("--trace-calls", action="store_true") + parser.add_argument("--skip-video", action="store_true") + parser.add_argument("--ffmpeg-loglevel", default="error") + args = parser.parse_args() + + rows: list[dict] = [] + call_rows: list[dict] = [] + dtype = dtype_from_name(args.dtype) + + state = timed("latent_load", rows, lambda: torch.load(args.latent, map_location="cuda", weights_only=False)) + latent = state["latent"].to("cuda") if isinstance(state, dict) else state.to("cuda") + frames = int(state.get("frames", latent.shape[2] * 4)) if isinstance(state, dict) else latent.shape[2] * 4 + + if args.implementation == "direct": + from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE + + vae = timed( + "vae_load", + rows, + lambda: MiniMaxH3VideoVAE.from_safetensors( + "/vae/minimax_h3_video_vae_fp16.safetensors", + device="cuda", + tiling=not args.no_tiling, + dtype=dtype, + ).eval(), + ) + else: + from safetensors.torch import load_file + + from h3_blackwell_runtime.upstream_vae import MiniMaxH3VideoVAE + + def load_upstream(): + model = MiniMaxH3VideoVAE(tiling=not args.no_tiling).to("cuda").eval() + checkpoint = load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda") + checkpoint = {key: value.to(dtype=dtype) for key, value in checkpoint.items()} + model.load_state_dict(checkpoint, strict=True) + return model + + vae = timed("vae_load", rows, load_upstream) + + if args.tile_size is not None: + vae.tile_size = args.tile_size + if args.tile_overlap is not None: + vae.tile_overlap_min = args.tile_overlap + + if args.trace_calls: + original_decode_pixels = vae._decode_pixels + + def traced_decode_pixels(z): + sync() + start = time.perf_counter() + result = original_decode_pixels(z) + sync() + call_rows.append({ + "index": len(call_rows), + "latent_shape": tuple(z.shape), + "seconds": time.perf_counter() - start, + }) + return result + + vae._decode_pixels = traced_decode_pixels + + with torch.inference_mode(): + decoded = timed("vae_decode", rows, lambda: vae.decode(latent.to(dtype))[:, :, :frames]) + pixels = timed( + "pixelize_cpu", + rows, + lambda: ((decoded[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu(), + ) + + if args.output is not None and not args.skip_video: + args.output.parent.mkdir(parents=True, exist_ok=True) + raw = args.output.with_suffix(".rgb") + timed("raw_write", rows, lambda: pixels.numpy().tofile(raw)) + timed( + "ffmpeg_encode", + rows, + lambda: subprocess.run( + ffmpeg_command( + args.ffmpeg_loglevel, + "-y", + "-f", + "rawvideo", + "-pixel_format", + "rgb24", + "-video_size", + f"{pixels.shape[2]}x{pixels.shape[1]}", + "-framerate", + "24", + "-i", + str(raw), + "-an", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + str(args.output), + ), + check=True, + ), + ) + raw.unlink() + + metrics = { + "implementation": args.implementation, + "dtype": args.dtype, + "tiling": not args.no_tiling, + "tile_size": vae.tile_size, + "tile_overlap_min": vae.tile_overlap_min, + "frames": frames, + "latent_shape": tuple(latent.shape), + "output": str(args.output) if args.output is not None else None, + "stages": rows, + "decode_pixel_calls": call_rows, + "total_seconds": sum(row["seconds"] for row in rows), + } + args.metrics.parent.mkdir(parents=True, exist_ok=True) + args.metrics.write_text(json.dumps(metrics, indent=2), encoding="utf-8") + print(json.dumps(metrics, indent=2), flush=True) + + +if __name__ == "__main__": + main()