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

invokeai/backend/quantization/gguf/loaders.py

@@ -36,16 +36,53 @@ 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.<tensor name>`` 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
             # while tensors still hold references to the mapped memory.
             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(

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)