FLUX.2-klein-4B: model metadata validation fails — tensor shape/naming mismatch (QKV [3072,3072] vs [3072,9216], VAE 32ch vs 16ch)

[PBCN216]

Environment

• sd.cpp: master-788-8a51eb9 (2026-07-22), compiled from source, Vulkan backend
• OS: Windows 11, Intel Iris Xe iGPU (Vulkan), 31GB RAM, i5-13500H
• Model: BFL FLUX.2-klein-4B from ModelScope/HuggingFace (diffusers format)

Repro steps

1. Download FLUX.2-klein-4B safetensors (transformer/ + text_encoder/ + vae/ layout)
2. Compile sd.cpp latest master with Vulkan
3. Run (both direct safetensors and GGUF convert path fail identically):

Direct:

sd-cli --diffusion-model transformer/diffusion_pytorch_model.safetensors \
  --llm text_encoder/ --vae vae/diffusion_pytorch_model.safetensors --vae-format flux2 \
  --backend diffusion=vulkan0 -p "cat" --cfg-scale 1 --steps 4 -H 512 -W 512

GGUF convert then load:

sd-cli --mode convert --diffusion-model transformer/..safetensors --llm text_encoder/ \
  --vae vae/..safetensors -o model.gguf --type q4_0
sd-cli -m model.gguf --backend diffusion=vulkan0 -p "cat" --cfg-scale 1 --steps 4

Expected behavior

FLUX.2 Klein loads and generates images (README says supported since 2026-01-18, PR #1193).

Actual behavior

Both paths fail identically at model metadata validation:

Version: Flux Control
ERROR: Diffusion model tensor 'model.diffusion_model.double_blocks.0.img_attn.qkv.weight'
  has wrong shape in model metadata: got [3072, 3072, 1, 1], expected [3072, 9216, 1, 1]
ERROR: Diffusion model tensor 'model.diffusion_model.double_blocks.0.img_attn.proj.bias'
  not in model metadata
ERROR: Diffusion model tensor 'model.diffusion_model.single_blocks.0.linear1.bias'
  not in model metadata
ERROR: Diffusion model tensor 'model.diffusion_model.single_blocks.0.linear1.weight'
  not in model metadata
... (all single_blocks.* tensors missing, repeats for double_blocks 0-4)
ERROR: VAE tensor 'first_stage_model.decoder.conv_in.weight'
  has wrong shape: got [3, 3, 32, 512], expected [3, 3, 16, 512]
ERROR: VAE tensor 'first_stage_model.encoder.conv_out.bias'
  has wrong shape: got [64, 1, 1, 1], expected [32, 1, 1, 1]
ERROR: VAE tensor 'first_stage_model.encoder.conv_out.weight'
  has wrong shape: got [3, 3, 512, 64], expected [3, 3, 512, 32]
stable-diffusion.cpp:1573 - model metadata validation failed

Root cause analysis

FLUX.2 Klein architecture differs from FLUX.1 in 3 ways not handled by current model_manager.cpp:

1. Separate Q/K/V projections: 3072→3072 each, instead of fused 3072→9216
2. single_blocks topology: Klein uses both single_blocks (20 layers) and double_blocks (5 layers), but model_manager expects only double_blocks
3. FLUX.2 VAE: 32 latent channels + 64 encoder output channels, instead of 16+32

PR #1193's example command uses single-file safetensors from ComfyOrg layout. May have worked at merge time but appears broken on current master with BFL's official diffusers-format safetensors.

Additional context

• Compiled with clang-cl 22.1.8 + MSVC 14.44, Vulkan 1.4.350 via Ninja
• Same error on both Vulkan and CPU backends
• leejet's own GGUF (FLUX.2-klein-4B-GGUF on HF) also fails — was converted with older converter missing KV metadata
• Manually converting with current master's convert mode produces GGUF with same tensor shape errors

[Green-Sky]

Interesting, cant compile myself right now, but can you check https://huggingface.co/Green-Sky/bonsai-image-binary-4B-GGUF/blob/main/bonsai_image_4b-mod_q8_0-q1_0.gguf ? I converted that recently and still works with the version I have locally. (+ its a small file)

[PBCN216]

Thanks for the reply @Green-Sky!

Quick update: we got FLUX.2 Klein 4B working (Intel Iris Xe Vulkan, Win11) by
bypassing GGUF entirely. The working approach matches #1216:

1. Diffusion: BFL single-file flux-2-klein-4b.safetensors (not diffusers subdirs)
2. VAE: Comfy-Org flux2-vae.safetensors (32ch, not BFL bundled VAE)
3. Text encoder: merged BFL's 2 shards into one qwen_3_4b_merged.safetensors
   (Comfy-Org's single-file TE only has 10 layers, incomplete)

Key flags for low-VRAM iGPU (Iris Xe 1GB):
--auto-fit --max-vram vulkan0=1 --offload-to-cpu --vae-conv-direct --vae-tiling
(no --diffusion-fa, flash attention OOMs on 1GB)

Result: 512x512 in ~80s, quality as expected for Klein 4B.

Root cause: sd.cpp model_manager doesn't recognize BFL's diffusers-format
subdirectory layout for Klein's separate QKV + single_blocks + 32ch VAE,
but the single-file safetensors path works on current master-788.