84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
"""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/fl2va-sampler-reference/step_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)
|
|
if "denoised" in state:
|
|
video_shape = (1, 24, 7, 12, 20)
|
|
video_count = torch.tensor(video_shape).prod().item()
|
|
latent = state["denoised"].to("cuda").reshape(-1)[:video_count].reshape(video_shape)
|
|
else:
|
|
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()
|