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

invokeai/backend/quantization/gguf/ggml_tensor.py

@@ -184,8 +184,14 @@ def get_dequantized_tensor(self):
             ).to(self.compute_dtype)
         else:
             # There is no GPU implementation for this quantization type, so fallback to the numpy implementation.
+            # The numpy implementation infers the output shape from the stored (possibly reshaped) data, so reshape
+            # to the logical shape here - see GGMLTensor.tensor_shape.
             new = gguf.quants.dequantize(self.quantized_data.cpu().numpy(), self._ggml_quantization_type)
-            return torch.from_numpy(new).to(self.quantized_data.device, dtype=self.compute_dtype)
+            return (
+                torch.from_numpy(new)
+                .reshape(self.tensor_shape)
+                .to(self.quantized_data.device, dtype=self.compute_dtype)
+            )
 
     @classmethod
     def __torch_dispatch__(cls, func, types, args, kwargs):

invokeai/backend/quantization/gguf/loaders.py

@@ -1,7 +1,10 @@
 import gc
+import math
 from pathlib import Path
+from typing import Any
 
 import gguf
+import numpy as np
 import torch
 
 from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
@@ -36,16 +39,77 @@ def close(self):
         gc.collect()
 
 
+ORIG_SHAPE_KEY_PREFIX = "comfy.gguf.orig_shape."
+
+
+def _coerce_dim(value: Any) -> int | None:
+    """Coerce a single value from a GGUF metadata array to a positive dimension, or None if it isn't one.
+
+    Metadata is untrusted input, so anything that is not a finite, integral, positive number is rejected rather
+    than silently truncated (``int(2.5) == 2``) or allowed to raise (``int(float("inf"))``).
+    """
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, (int, np.integer)):
+        dim = int(value)
+    elif isinstance(value, (float, np.floating)):
+        value = float(value)
+        if not math.isfinite(value) or not value.is_integer():
+            return None
+        dim = int(value)
+    else:
+        return None
+    return dim if dim > 0 else None
+
+
+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:
+            contents = field.contents()
+        except Exception as e:
+            logger.warning(f"Ignoring malformed GGUF metadata key {key!r}: {e}")
+            continue
+        if not isinstance(contents, (list, tuple)):
+            logger.warning(f"Ignoring malformed GGUF metadata key {key!r}: expected an array, got {contents!r}")
+            continue
+        dims = tuple(_coerce_dim(v) for v in contents)
+        if not dims or any(d is None for d in dims):
+            logger.warning(f"Ignoring malformed GGUF metadata key {key!r}: {tuple(contents)!r}")
+            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_ggml_tensor.py

@@ -3,6 +3,7 @@
 import torch
 
 from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+from invokeai.backend.quantization.gguf.utils import DEQUANTIZE_FUNCTIONS, TORCH_COMPATIBLE_QTYPES
 from invokeai.backend.util.calc_tensor_size import calc_tensor_size
 
 
@@ -126,3 +127,28 @@ def test_ggml_tensor_calc_size():
     compression_ratio = calc_tensor_size(x) / calc_tensor_size(x_quantized)
     # Assert that the compression ratio is approximately 4x.
     assert abs(compression_ratio - 4) < 0.5
+
+
+def test_ggml_tensor_numpy_fallback_dequantize_uses_logical_shape():
+    """Qtypes without a torch dequantize kernel fall back to numpy, which infers the shape from the stored data.
+
+    ComfyUI reshapes tensors before quantizing them, so the stored shape is not necessarily the logical one - the
+    fallback must still return the logical shape (see GGMLTensor.tensor_shape).
+    """
+    qtype = gguf.GGMLQuantizationType.IQ4_NL
+    assert qtype not in DEQUANTIZE_FUNCTIONS and qtype not in TORCH_COMPATIBLE_QTYPES
+
+    # One IQ4_NL block is 32 elements in 18 bytes, so this holds 8 rows of 32 elements = 256 elements, which the
+    # numpy fallback dequantizes to (8, 32). The logical shape is a different view of the same 256 elements.
+    generator = torch.Generator().manual_seed(123)
+    quantized = torch.randint(0, 256, (8, 18), dtype=torch.uint8, generator=generator)
+    logical_shape = torch.Size((4, 64))
+
+    t = GGMLTensor(
+        data=quantized,
+        ggml_quantization_type=qtype,
+        tensor_shape=logical_shape,
+        compute_dtype=torch.float32,
+    )
+
+    assert t.get_dequantized_tensor().shape == logical_shape

tests/backend/quantization/gguf/test_loaders.py

@@ -0,0 +1,68 @@
+from collections.abc import Sequence
+
+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: Sequence[float] | 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", list(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)
+
+
+@pytest.mark.parametrize(
+    "orig_shape",
+    [
+        pytest.param([2.5, 157286.4], id="non-integral"),
+        pytest.param([float("inf")], id="infinite"),
+        pytest.param([float("nan"), 64.0], id="nan"),
+        pytest.param([-6144.0, -64.0], id="negative"),
+    ],
+)
+def test_gguf_sd_loader_ignores_malformed_orig_shape(tmp_path, orig_shape):
+    """Malformed metadata must be ignored with a warning - never truncated to a wrong shape, never raised."""
+    path = tmp_path / "model.gguf"
+    _write_gguf(path, orig_shape=orig_shape)
+
+    sd = gguf_sd_loader(path, compute_dtype=torch.bfloat16)
+
+    assert tuple(sd["first.weight"].shape) == (1536, 256)