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

Add upstream VAE latent decode mode

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

23
tools/decode_video_latent.py
View file

@ -13,6 +13,7 @@ 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("--implementation", choices=("direct", "upstream"), default="direct")
parser.add_argument("--frames-dir", type=Path)
args = parser.parse_args()
@ -20,11 +21,21 @@ 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()
if args.implementation == "direct":
vae = MiniMaxH3VideoVAE.from_safetensors(
"/vae/minimax_h3_video_vae_fp16.safetensors",
device="cuda",
tiling=not args.no_tiling,
).eval()
else:
from safetensors.torch import load_file
from h3_blackwell_runtime.upstream_vae import MiniMaxH3VideoVAE as UpstreamMiniMaxH3VideoVAE
vae = UpstreamMiniMaxH3VideoVAE(tiling=not args.no_tiling).to("cuda").eval()
checkpoint = load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda")
missing, unexpected = vae.load_state_dict(checkpoint, strict=False)
print({"upstream_missing": len(missing), "upstream_unexpected": len(unexpected)}, flush=True)
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()
@ -45,4 +56,4 @@ subprocess.run([
"-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})
print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape), "tiling": not args.no_tiling, "implementation": args.implementation})