daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add standalone H3 latent decode tooling

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 00:11:17 +07:00
parent c3d72d9f1e
commit 6b0050adfa
2 changed files with 56 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
tools/decode_video_latent.py Normal file
View file

@ -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})

8
tools/direct_t2v_preview.py
View file

@ -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():