Fix Krea2/Qwen-Image flat samples: use a VAE with non-zero quant convs

[skydam]

I've never submitted something like this before and i want to be upfront: Claude Code came up with the solution, but i was getting these brown sample images. Claude Code says it's the VAE.

Problem

Training Krea2 (raw or turbo) produces sample previews that are a uniform
muddy-brown field — identical for every prompt, step, and run (the output JPEGs
are byte-for-byte identical). FLUX / z-image are unaffected.

Root cause

Krea2 defaults its VAE to Qwen/Qwen-Image, loaded via AutoencoderKLQwenImage.
In the current main revision of that repo, the VAE's quant_conv and
post_quant_conv weights are all zero. The pinned diffusers commit applies
these convs on both encode and decode, which zeros the latent path:

• decode → constant output regardless of latent → flat brown image
• encode → returns N(0,1) noise instead of the image → so training also
  silently learns against noise targets, not just broken previews.

Qwen/Qwen-Image-2512 is the same VAE (byte-identical encoder/decoder weights)
with the quant convs populated. extensions_built_in/diffusion_models/qwen_image/ qwen_image.py loads the same Qwen/Qwen-Image VAE, so that arch is likely
affected too.

Fix

• Default Krea2 to Qwen/Qwen-Image-2512.
• Guard _load_vae to raise on an all-zero quant_conv so this fails loudly
  instead of silently wasting training hours.

Minimal repro (no DiT needed)

import torch
from diffusers import AutoencoderKLQwenImage

for repo in ["Qwen/Qwen-Image", "Qwen/Qwen-Image-2512"]:
    vae = AutoencoderKLQwenImage.from_pretrained(
        repo, subfolder="vae", torch_dtype=torch.float32).to("cuda").eval()
    a, b = torch.randn(1,16,1,64,64,device="cuda"), torch.randn(1,16,1,64,64,device="cuda")
    with torch.no_grad():
        da, db = vae.decode(a).sample, vae.decode(b).sample
    print(repo,
          "| post_quant_conv std:", round(vae.post_quant_conv.weight.std().item(), 4),
          "| decode response to latent:", round(((da-db).norm()/da.norm()).item(), 4))
# Qwen/Qwen-Image      -> post_quant_conv std: 0.0    | decode response: 0.0   (dead)
# Qwen/Qwen-Image-2512 -> post_quant_conv std: 0.104  | decode response: ~1.0  (works)

Verification

With the fix, the VAE round-trip reconstructs the input (recon pixel std matches
the source image) and Krea2 generation produces correct, prompt-distinct images
(verified end-to-end on Krea-2-Raw: "a red sports car on a city street" and
"an oil painting of snowy mountains at sunset" both render correctly).

Alternative the maintainer may prefer: pin a Qwen/Qwen-Image revision whose VAE
has populated quant convs. The same fix likely applies to qwen_image.py.

🤖 Generated with Claude Code (https://claude.com/claude-code)