Add H3 VAE pixel comparator

tools/compare_vae_pixels.py

@@ -0,0 +1,79 @@
+"""Decode the captured final H3 video latent and compare against Comfy frames."""
+
+import argparse
+from pathlib import Path
+import subprocess
+
+import numpy as np
+from PIL import Image
+import torch
+
+from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
+
+
+def load_reference_frames(path: Path, frames: int) -> torch.Tensor:
+    files = sorted(path.glob("fl2va-reference_*.png"))[-frames:]
+    if len(files) != frames:
+        raise FileNotFoundError(f"expected {frames} reference PNGs in {path}, found {len(files)}")
+    arrays = [np.asarray(Image.open(file).convert("RGB"), dtype=np.uint8) for file in files]
+    return torch.from_numpy(np.stack(arrays, axis=0))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--latent", default="/artifacts/capture/output_11.pt")
+    parser.add_argument("--frames", type=int, default=22)
+    parser.add_argument("--reference-frames", type=Path, default=Path("/artifacts/fl2va-sampler-reference-e679178/reference_frames"))
+    parser.add_argument("--output", type=Path, default=Path("/output/h3-blackwell-runtime/direct-vae-same-latent.mp4"))
+    args = parser.parse_args()
+
+    state = torch.load(args.latent, map_location="cuda", weights_only=False)
+    latent = state["video"].to("cuda")
+    vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
+    with torch.inference_mode():
+        decoded = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, : args.frames]
+    pixels = ((decoded[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()
+    reference = load_reference_frames(args.reference_frames, args.frames)
+    delta = (pixels.to(torch.int16) - reference.to(torch.int16)).abs()
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    raw = args.output.with_suffix(".rgb")
+    pixels.numpy().tofile(raw)
+    subprocess.run(
+        [
+            "ffmpeg",
+            "-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,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    raw.unlink()
+    print(
+        {
+            "output": str(args.output),
+            "shape": tuple(pixels.shape),
+            "mean_abs_pixel_delta": float(delta.float().mean()),
+            "max_abs_pixel_delta": int(delta.max()),
+        }
+    )
+
+
+if __name__ == "__main__":
+    main()