import json, time, sys, torch
from diffusers import AutoencoderKLQwenImage21
from common import BASE
from vae_fast import enable_fast_decode
mode = sys.argv[1]
if "bench" in mode: torch.backends.cudnn.benchmark = True
vae = AutoencoderKLQwenImage21.from_pretrained(BASE, subfolder="vae", torch_dtype=torch.bfloat16).cuda().eval()
if "cl" in mode: vae = vae.to(memory_format=torch.channels_last)
if mode != "eager": enable_fast_decode(vae, compile="compile" in mode)
with torch.inference_mode():
    for s in (64, 128):
        z = torch.randn(1, 64, 1, s, s, device="cuda", dtype=torch.bfloat16)
        for _ in range(3): vae.decode(z, return_dict=False)
        torch.cuda.synchronize(); t = time.perf_counter()
        for _ in range(5): vae.decode(z, return_dict=False)
        torch.cuda.synchronize(); print(json.dumps({"mode": mode, "px": s * 16, "ms": (time.perf_counter() - t) / 5 * 1e3}))