Fix QK-norm .scale/.weight key mismatch on Flux-compat GGUF loads (silent NaN/black image)

[ayanoby]

Problem

Some third-party GGUF quantizations that remap a non-Flux architecture with QK-normalization onto Flux's double_blocks/single_blocks tensor naming (for free compatibility with this loader's existing Flux-family support) store the QK-norm weights under a .scale suffix instead of the .weight suffix that comfy's RMSNorm (comfy/ops.py, a thin wrapper around torch.nn.RMSNorm) actually registers its parameter as.

Because comfy.sd.load_diffusion_model_state_dict loads with strict=False, this mismatch never raises — it just logs unet missing: [...norm.query_norm.weight, ...] / unet unexpected: [...norm.query_norm.scale, ...] warnings for every affected layer and silently drops ~140 tensors. Every affected RMSNorm.weight is then left uninitialized. The model still runs to completion, but NaNs propagate through the rest of the forward pass and the decoded image comes out solid black — with no hard error anywhere to point at the cause.

Worth noting: comfy/model_detection.py's architecture detection already treats .weight and .scale as interchangeable for exactly these keys via any_suffix_in(state_dict_keys, key_prefix, 'double_blocks.0.img_attn.norm.key_norm.', ["weight", "scale"]) — so the model gets correctly identified as Flux-family, it just doesn't get correctly loaded.

Root cause

Confirmed with stduhpf/LongCat-Image-Edit-gguf (Q4_K_M), a GGUF quantization of meituan-longcat/LongCat-Image-Edit repacked into Flux-compatible ("bfl_format") naming. LongCat's DiT has an extra QK-normalization step that Flux.1 doesn't have; everything else in the architecture lined up 1:1 with Flux's naming under this repack, but this one component used the scale naming convention (following llama.cpp/sd.cpp's own norm-tensor naming) instead of weight.

This is a different issue from #383 (1D tensors incorrectly GGUF-quantized instead of stored as raw fp32/bf16, fixed in a57094a) — that one is about quantization/dtype, this one is purely a key-naming mismatch, and reproduces on current main.

Fix

Rename .query_norm.scale / .key_norm.scale → .weight at load time, right after sd_key is resolved in gguf_sd_loader(), before the tensor is added to the state dict. Minimal, only touches keys with these two exact suffixes, shouldn't affect any GGUF file that doesn't have this specific naming quirk.

Verification

Before: loading stduhpf/LongCat-Image-Edit-gguf (Q4_K_M) produces the missing/unexpected warnings above during model load, then RuntimeWarning: invalid value encountered in cast during sampling/decode, and a fully black output image.

After: no missing/unexpected warnings, no cast warnings, normal (non-black) output — using the same 3.7GB GGUF file, no need to fall back to the full ~12.5GB bf16 checkpoint.

(Ruled out other causes first: an isolated VAEEncode→VAEDecode round-trip on the same VAE reproduced the input correctly, and the workflow graph validated cleanly via graphToPrompt() — so the VAE and graph shape weren't the problem before landing on the UNET weights themselves.)

Full write-up with logs: https://github.com/ayanoby/comfyui-gguf-qknorm-fix

Test plan

• Confirmed unet missing/unet unexpected warnings for QK-norm tensors disappear after the patch
• Confirmed output image is no longer solid black, using the same GGUF file
• Not tested against other GGUF files/architectures beyond the one above — the rename is scoped tightly enough that it should be a no-op for anything not hitting this exact naming pattern, but flagging for maintainer awareness