From ee29a281121002d7629383b91e01440d713fe457 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 31 Aug 2026 23:35:05 +0200 Subject: [PATCH] fix(gguf): honor ComfyUI's comfy.gguf.orig_shape metadata 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.` 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 #9537 Co-Authored-By: Claude Opus 5 (1M context) --- invokeai/backend/quantization/gguf/loaders.py | 37 +++++++++++++++ .../backend/quantization/gguf/test_loaders.py | 47 +++++++++++++++++++ 2 files changed, 84 insertions(+) create mode 100644 tests/backend/quantization/gguf/test_loaders.py diff --git a/invokeai/backend/quantization/gguf/loaders.py b/invokeai/backend/quantization/gguf/loaders.py index cb8ac2dbeb6..ce8e1ed5164 100644 --- a/invokeai/backend/quantization/gguf/loaders.py +++ b/invokeai/backend/quantization/gguf/loaders.py @@ -36,9 +36,38 @@ def close(self): gc.collect() +ORIG_SHAPE_KEY_PREFIX = "comfy.gguf.orig_shape." + + +def _read_comfy_orig_shapes(reader: gguf.GGUFReader) -> dict[str, torch.Size]: + """Read ComfyUI's ``comfy.gguf.orig_shape.`` metadata. + + ComfyUI's GGUF converter can only quantize 2-D tensors, so it reshapes any tensor whose native + rank/shape the quantizer rejects (e.g. Krea-2's ``first.weight`` of (6144, 64)) into a workable + 2-D shape and records the native shape under this key. Without honoring it, the tensor loads with + the reshaped shape and ``load_state_dict`` fails with a size mismatch. + """ + orig_shapes: dict[str, torch.Size] = {} + for key, field in reader.fields.items(): + if not key.startswith(ORIG_SHAPE_KEY_PREFIX): + continue + tensor_name = key[len(ORIG_SHAPE_KEY_PREFIX) :] + try: + dims = tuple(int(v) for v in field.contents()) + except (TypeError, ValueError) as e: + logger.warning(f"Ignoring malformed GGUF metadata key {key!r}: {e}") + continue + if not dims or any(d <= 0 for d in dims): + logger.warning(f"Ignoring malformed GGUF metadata key {key!r}: {dims}") + continue + orig_shapes[tensor_name] = torch.Size(dims) + return orig_shapes + + def gguf_sd_loader(path: Path, compute_dtype: torch.dtype) -> dict[str, GGMLTensor]: with WrappedGGUFReader(path) as reader: sd: dict[str, GGMLTensor] = {} + orig_shapes = _read_comfy_orig_shapes(reader) for tensor in reader.tensors: # Use .copy() to create a true copy of the data, not a view. # This is critical on Windows where the memory-mapped file cannot be deleted @@ -46,6 +75,14 @@ def gguf_sd_loader(path: Path, compute_dtype: torch.dtype) -> dict[str, GGMLTens torch_tensor = torch.from_numpy(tensor.data.copy()) shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) + orig_shape = orig_shapes.get(tensor.name) + if orig_shape is not None: + if orig_shape.numel() != shape.numel(): + raise ValueError( + f"GGUF tensor {tensor.name!r} declares original shape {tuple(orig_shape)}, which has a " + f"different element count than its stored shape {tuple(shape)}." + ) + shape = orig_shape if tensor.tensor_type in TORCH_COMPATIBLE_QTYPES: torch_tensor = torch_tensor.view(*shape) sd[tensor.name] = GGMLTensor( diff --git a/tests/backend/quantization/gguf/test_loaders.py b/tests/backend/quantization/gguf/test_loaders.py new file mode 100644 index 00000000000..c087f45ef5d --- /dev/null +++ b/tests/backend/quantization/gguf/test_loaders.py @@ -0,0 +1,47 @@ +import gguf +import numpy as np +import pytest +import torch + +from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader + + +def _write_gguf(path, *, orig_shape: tuple[int, ...] | None) -> None: + """Write a tiny GGUF holding one F32 tensor stored 2-D as (256, 1536) i.e. torch (1536, 256).""" + writer = gguf.GGUFWriter(str(path), "krea2") + stored = np.arange(1536 * 256, dtype=np.float32).reshape(1536, 256) + writer.add_tensor("first.weight", stored, raw_dtype=gguf.GGMLQuantizationType.F32) + if orig_shape is not None: + writer.add_array("comfy.gguf.orig_shape.first.weight", [int(d) for d in orig_shape]) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + +def test_gguf_sd_loader_honors_comfy_orig_shape(tmp_path): + """ComfyUI reshapes non-2-D tensors before quantizing; the recorded native shape must win.""" + path = tmp_path / "model.gguf" + _write_gguf(path, orig_shape=(6144, 64)) + + sd = gguf_sd_loader(path, compute_dtype=torch.bfloat16) + + assert tuple(sd["first.weight"].shape) == (6144, 64) + assert tuple(sd["first.weight"].get_dequantized_tensor().shape) == (6144, 64) + + +def test_gguf_sd_loader_without_orig_shape(tmp_path): + path = tmp_path / "model.gguf" + _write_gguf(path, orig_shape=None) + + sd = gguf_sd_loader(path, compute_dtype=torch.bfloat16) + + assert tuple(sd["first.weight"].shape) == (1536, 256) + + +def test_gguf_sd_loader_rejects_orig_shape_with_wrong_element_count(tmp_path): + path = tmp_path / "model.gguf" + _write_gguf(path, orig_shape=(6144, 65)) + + with pytest.raises(ValueError, match="different element count"): + gguf_sd_loader(path, compute_dtype=torch.bfloat16)