diff --git a/tools/compare_vae_pixels.py b/tools/compare_vae_pixels.py new file mode 100644 index 0000000..d27c42d --- /dev/null +++ b/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()