Add standalone H3 latent decode tooling

tools/decode_video_latent.py

@@ -0,0 +1,48 @@
+"""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})

tools/direct_t2v_preview.py

@@ -33,6 +33,8 @@ parser.add_argument("--seed", type=int, default=440204)
 parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
 parser.add_argument("--model-timesteps-capture", type=Path, help="Directory containing captured input_XX.pt H3 timesteps for strict parity checks.")
 parser.add_argument("--profile-memory", action="store_true")
+parser.add_argument("--save-latent", type=Path)
+parser.add_argument("--skip-decode", action="store_true")
 args = parser.parse_args()
 started = time.perf_counter()
 last_report = started
@@ -77,6 +79,12 @@ if args.model_timesteps_capture is not None:
     ]
 latent = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps)
 report_memory("sampled")
+if args.save_latent is not None:
+    args.save_latent.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"latent": latent.detach().cpu(), "frames": frames, "width": args.width, "height": args.height, "prompt": args.prompt, "seed": args.seed}, args.save_latent)
+    print({"latent": str(args.save_latent)}, flush=True)
+if args.skip_decode:
+    raise SystemExit(0)
 vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
 report_memory("vae_loaded")
 with torch.inference_mode():