48 lines
1.8 KiB
Python
48 lines
1.8 KiB
Python
"""Decode a saved H3 video latent with direct VAE settings."""
|
|
|
|
import argparse
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--latent", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--no-tiling", action="store_true")
|
|
parser.add_argument("--frames-dir", type=Path)
|
|
args = parser.parse_args()
|
|
|
|
state = 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
|
|
|
|
vae = MiniMaxH3VideoVAE.from_safetensors(
|
|
"/vae/minimax_h3_video_vae_fp16.safetensors",
|
|
device="cuda",
|
|
tiling=not args.no_tiling,
|
|
).eval()
|
|
with torch.inference_mode():
|
|
pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]
|
|
pixels = ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()
|
|
|
|
if args.frames_dir is not None:
|
|
from PIL import Image
|
|
|
|
args.frames_dir.mkdir(parents=True, exist_ok=True)
|
|
for index, frame in enumerate(pixels.numpy()):
|
|
Image.fromarray(frame).save(args.frames_dir / f"frame_{index:04d}.png")
|
|
|
|
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)
|
|
raw.unlink()
|
|
print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape), "tiling": not args.no_tiling})
|