| """Reproduce the README's fp16 comparison: decode a Qwen-Image-2.1 RGBA latent with the official VAE (bf16 and fp16) |
| and with this VAE (fp16), composited over a checkerboard. Pixels where the decoder output is NaN / inf are painted magenta. |
| |
| python fp16_demo.py |
| """ |
| import torch |
| from diffusers import AutoencoderKLQwenImage21 |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
|
|
| REPO = "madebyollin/texture-fix-vae-for-qwen-image-2.1" |
| |
| z = torch.load(hf_hub_download(REPO, "images/rgba_dragon_sticker_1024.latent.pt"), weights_only=True) |
|
|
|
|
| @torch.no_grad() |
| def decode(repo, subfolder, dtype): |
| vae = AutoencoderKLQwenImage21.from_pretrained(repo, subfolder=subfolder, torch_dtype=dtype).to("cuda") |
| mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1) |
| std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1) |
| return vae.decode((z.unsqueeze(2) * std + mean).to("cuda", dtype)).sample[:, :, 0].float().cpu() |
|
|
|
|
| def save(y, name): |
| bad = ~y.isfinite().all(1, keepdim=True) |
| y = y.nan_to_num(0).clamp(-1, 1).mul(0.5).add(0.5) |
| h, w = y.shape[-2:] |
| checker = ((torch.arange(h)[:, None] // 32 + torch.arange(w)[None] // 32) % 2).float() * 0.2 + 0.6 |
| rgb = y[:, :3] * y[:, 3:] + checker * (1 - y[:, 3:]) |
| rgb = torch.where(bad, torch.tensor([1.0, 0.0, 1.0]).view(1, 3, 1, 1), rgb) |
| Image.fromarray(rgb[0].permute(1, 2, 0).mul(255).round().byte().numpy()).save(name) |
| print(f"{name}: {bad.float().mean().item():.1%} of pixels are NaN / inf") |
|
|
|
|
| save(decode("Qwen/Qwen-Image-2.1", "vae", torch.bfloat16), "fp16-demo-official-bf16.png") |
| save(decode("Qwen/Qwen-Image-2.1", "vae", torch.float16), "fp16-demo-official-fp16.png") |
| save(decode(REPO, None, torch.float16), "fp16-demo-texture-fix-fp16.png") |
|
|
