From 9ce5e51dce6c7b47feddb4da58c0a199665a9f13 Mon Sep 17 00:00:00 2001 From: ayanoby Date: Mon, 3 Aug 2026 00:59:13 +0300 Subject: [PATCH] Fix QK-norm .scale/.weight key mismatch on Flux-compat GGUF loads Some third-party GGUF quantizations remap a non-Flux architecture with QK-normalization (e.g. LongCat-Image-Edit, converted to Flux-compatible 'bfl_format' naming for loader reuse) onto double_blocks/single_blocks naming, but store the norm weights under a .scale suffix instead of the .weight suffix that comfy's RMSNorm module (torch.nn.RMSNorm) actually registers its parameter as. Since load_diffusion_model_state_dict uses strict=False, this mismatch never raises: it's logged as 'unet missing'/'unet unexpected' warnings, and every affected RMSNorm ends up with an uninitialized weight tensor. The model still runs, but NaNs propagate through the rest of the forward pass and the final image comes out solid black. Architecture detection in comfy/model_detection.py already treats .weight and .scale as interchangeable for these keys via any_suffix_in(), so this makes the actual state dict loading consistent with that. Verified against stduhpf/LongCat-Image-Edit-gguf (Q4_K_M): before this patch, loading produces the missing/unexpected warnings above and a black image; after, no warnings and normal output, using the same GGUF file (no need to fall back to the full bf16 checkpoint). Full writeup: https://github.com/ayanoby/comfyui-gguf-qknorm-fix --- loader.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/loader.py b/loader.py index 7cefb11..eeaffbf 100644 --- a/loader.py +++ b/loader.py @@ -87,6 +87,12 @@ def gguf_sd_loader(path, handle_prefix="model.diffusion_model.", is_text_model=F if not tensor_name.startswith(handle_prefix): continue sd_key = tensor_name[prefix_len:] + # Some GGUF conversions (e.g. sd.cpp-style "bfl_format" repacks of + # non-Flux archs like LongCat) name QK-norm params ".scale" instead + # of the ".weight" that comfy's RMSNorm module actually registers, + # so they'd otherwise silently fail to load (-> NaNs downstream). + if sd_key.endswith(".query_norm.scale") or sd_key.endswith(".key_norm.scale"): + sd_key = sd_key[:-len(".scale")] + ".weight" tensors.append((sd_key, tensor)) # detect and verify architecture