Skip to content

fix(gguf): honor ComfyUI's comfy.gguf.orig_shape metadata - #9564

Open
Pfannkuchensack wants to merge 2 commits into
invoke-ai:mainfrom
Pfannkuchensack:fix/gguf_comfy_orig_shape
Open

fix(gguf): honor ComfyUI's comfy.gguf.orig_shape metadata#9564
Pfannkuchensack wants to merge 2 commits into
invoke-ai:mainfrom
Pfannkuchensack:fix/gguf_comfy_orig_shape

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Aug 31, 2026

Copy link
Copy Markdown
Member

Summary

Kind of change: fix (backend, GGUF model loading).

Why: Certain ComfyUI-produced GGUF checkpoints install fine but fail at generation with a weight shape mismatch, e.g. for a Krea-2 Q4_K_M checkpoint:

size mismatch for img_in.weight: copying a param with shape torch.Size([1536, 256])
from checkpoint, the shape in current model is torch.Size([6144, 64])

ComfyUI's GGUF converter can only quantize 2-D tensors, so it reshapes any tensor whose native shape the quantizer rejects and records the native shape in a comfy.gguf.orig_shape.<tensor name> KV entry. gguf_sd_loader ignored those entries and used the stored shape instead, so load_state_dict saw a wrongly-shaped tensor.

The affected file's header shows exactly this — note that both shapes have the same element count, so it is purely a reshape:

comfy.gguf.orig_shape.first.weight  [6144, 64]
first.weight                        [256, 1536]  F32

This is not Krea-2 specific; other Krea-2 GGUFs load only because their conversion did not have to reshape this tensor.

How: gguf_sd_loader now reads the comfy.gguf.orig_shape.* metadata and uses the declared native shape for the GGMLTensor. An entry whose element count does not match the stored tensor raises with an explicit message rather than producing a silently wrong view; a malformed entry is logged and ignored.

Related Issues / Discussions

Closes #9537

QA Instructions

Unit tests:

pytest tests/backend/quantization/gguf/test_loaders.py

They cover all three paths: metadata present (shape is applied), metadata absent (unchanged behaviour), metadata inconsistent (raises).

Manual, with the checkpoint from #9537 (Krea-2 SAT-IOR v2 Q4_K_M GGUF from CivitAI):

  1. Install the GGUF via the model manager — it is probed as main / krea-2 / gguf_quantized.
  2. Run a text-to-image generation with it (a standalone Qwen3-VL encoder and the Qwen-Image VAE, since a single-file checkpoint bundles neither).
  3. Before this change: generation fails with the size mismatch above. After: load_state_dict reports no missing or unexpected keys, and generation completes with a coherent image (verified at 512x512, 4 steps).

Also worth a quick regression check that an unaffected GGUF (e.g. an existing FLUX or Krea-2 GGUF without comfy.gguf.orig_shape keys) still loads — those files take the unchanged code path.

Merge Plan

Nothing special — self-contained backend change, no schema or frontend impact.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Tests added / updated (if applicable)
  • ❗Changes to a redux slice have a corresponding migration
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

ComfyUI's GGUF converter can only quantize 2-D tensors, so it reshapes any
tensor the quantizer rejects and records the native shape under a
`comfy.gguf.orig_shape.<tensor name>` KV entry. `gguf_sd_loader` ignored those
entries and used the stored shape, so such a checkpoint failed at load with a
size mismatch.

Concretely, Krea-2's `first.weight` is (6144, 64) but is stored as (1536, 256),
which produced:

    size mismatch for img_in.weight: copying a param with shape
    torch.Size([1536, 256]) from checkpoint, the shape in current model is
    torch.Size([6144, 64])

The loader now reads the metadata and uses the declared native shape, rejecting
an entry whose element count doesn't match the stored tensor and warning on a
malformed one. This is architecture-agnostic, not a Krea-2 special case.

Verified end-to-end: the affected checkpoint from the issue installs, loads and
generates a coherent image.

Closes invoke-ai#9537

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@github-actions github-actions Bot added python PRs that change python files backend PRs that change backend files python-tests PRs that change python tests labels Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backend PRs that change backend files python PRs that change python files python-tests PRs that change python tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[bug]: At least some Krea 2 Q4_K_M GGUFs fail on generation

1 participant