"""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()