diff --git a/invokeai/backend/quantization/gguf/ggml_tensor.py b/invokeai/backend/quantization/gguf/ggml_tensor.py index af895fb3eee..9d7219fe699 100644 --- a/invokeai/backend/quantization/gguf/ggml_tensor.py +++ b/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): diff --git a/invokeai/backend/quantization/gguf/loaders.py b/invokeai/backend/quantization/gguf/loaders.py index cb8ac2dbeb6..0ef98433f79 100644 --- a/invokeai/backend/quantization/gguf/loaders.py +++ b/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,9 +39,62 @@ 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.`` 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 @@ -46,6 +102,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_ggml_tensor.py b/tests/backend/quantization/gguf/test_ggml_tensor.py index 33506c0bddc..01c3aae6466 100644 --- a/tests/backend/quantization/gguf/test_ggml_tensor.py +++ b/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 diff --git a/tests/backend/quantization/gguf/test_loaders.py b/tests/backend/quantization/gguf/test_loaders.py new file mode 100644 index 00000000000..4539af5a72c --- /dev/null +++ b/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)