28 lines
1 KiB
Python
28 lines
1 KiB
Python
"""Validate and optionally execute the direct H3 video VAE decoder."""
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("checkpoint", nargs="?", default="/vae/minimax_h3_video_vae_fp16.safetensors")
|
|
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
|
parser.add_argument("--validate-only", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
vae = MiniMaxH3VideoVAE.from_safetensors(args.checkpoint, device=args.device).eval()
|
|
print(f"validated checkpoint keys and shapes: {args.checkpoint}")
|
|
if args.validate_only:
|
|
return
|
|
with torch.inference_mode():
|
|
latent = torch.zeros(1, 24, 1, 1, 1, device=args.device, dtype=next(vae.parameters()).dtype)
|
|
decoded = vae.decode(latent)
|
|
print(f"decoded latent {tuple(latent.shape)} -> {tuple(decoded.shape)} ({decoded.dtype}, {decoded.device})")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|