File size: 579 Bytes
ffc54ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 | 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())
|