import time, torch
from diffusers import AutoencoderKLQwenImage21
from common import BASE
from vae_fast import enable_fast_decode
vae = AutoencoderKLQwenImage21.from_pretrained(BASE, subfolder='vae', torch_dtype=torch.bfloat16).cuda().eval()
torch.manual_seed(0)
z = torch.randn(1, 64, 1, 64, 64, device='cuda', dtype=torch.bfloat16)
with torch.inference_mode():
    ref = vae.decode(z, return_dict=False)[0]
    enable_fast_decode(vae)
    out = vae.decode(z, return_dict=False)[0]
    print('exact equal', torch.equal(ref, out), (ref.float() - out.float()).abs().max().item())