perf(model-manager): identify GGUF models from the header alone

invokeai/backend/model_manager/model_on_disk.py

@@ -12,7 +12,7 @@
 from invokeai.app.services.config.config_default import get_config
 from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash
 from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
-from invokeai.backend.quantization.gguf.loaders import WrappedGGUFReader, gguf_sd_loader
+from invokeai.backend.quantization.gguf.loaders import WrappedGGUFReader, gguf_header_state_dict
 from invokeai.backend.quantization.sdnq.loaders import sdnq_sd_loader
 from invokeai.backend.util.logging import InvokeAILogger
 from invokeai.backend.util.silence_warnings import SilenceWarnings
@@ -201,9 +201,10 @@ def load_state_dict(self, path: Optional[Path] = None) -> StateDict:
                 checkpoint = torch.load(path, map_location="meta")
                 assert isinstance(checkpoint, dict)
             elif path.suffix.endswith(".gguf"):
-                # Identification reads shapes and never builds a module. Whether a Q8_CR file's quantization
-                # is supported is decided from its metadata by `ModelConfigFactory`, once, not per candidate.
-                checkpoint = gguf_sd_loader(path, compute_dtype=torch.float32, q8_cr="ignore")
+                # Shapes and quantization types from the header, like the safetensors branch: the tensor data
+                # is never read. Whether a Q8_CR file's quantization is supported is decided from its metadata
+                # by `ModelConfigFactory`, once, not per candidate.
+                checkpoint = gguf_header_state_dict(path, compute_dtype=torch.float32)
             elif path.suffix.endswith(".safetensors"):
                 if _is_sdnq_safetensors(path):
                     checkpoint = sdnq_sd_loader(path, compute_dtype=torch.float32)

invokeai/backend/quantization/gguf/loaders.py

@@ -97,7 +97,7 @@ def _read_comfy_orig_shapes(reader: gguf.GGUFReader) -> dict[str, torch.Size]:
 
 COMFY_QUANT_KEY_PREFIX = "comfy.gguf.quant."
 
-Q8CRMode = Literal["refuse", "decode", "ignore"]
+Q8CRMode = Literal["refuse", "decode"]
 
 
 def parse_q8_cr_markers(fields: Mapping[str, Any], source: str) -> dict[str, dict[str, Any]]:
@@ -144,6 +144,46 @@ def parse_q8_cr_markers(fields: Mapping[str, Any], source: str) -> dict[str, dic
     return layers
 
 
+def _logical_shape(tensor: gguf.ReaderTensor, orig_shapes: Mapping[str, torch.Size]) -> torch.Size:
+    """The shape a tensor has in the model: GGUF's reversed dims, or the native shape ComfyUI recorded."""
+    shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
+    orig_shape = orig_shapes.get(tensor.name)
+    if orig_shape is None:
+        return shape
+    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)}."
+        )
+    return orig_shape
+
+
+def gguf_header_state_dict(path: Path, compute_dtype: torch.dtype) -> dict[str, torch.Tensor]:
+    """A GGUF file's tensors as `gguf_sd_loader` returns them, but on the `meta` device: read from the header.
+
+    Model identification needs names, shapes and quantization types, never values, and the header holds them
+    all. `gguf_sd_loader` copies every tensor out of the file -- for a 23.6 GB transformer, 22 GiB of RAM per
+    probe -- to answer questions the first few MB answer. Each tensor comes back as the `GGMLTensor` the full
+    read gives, with the same packed and logical shapes; a ComfyUI-GGUF ``Q8_CR`` layer comes back as its raw
+    ``I8`` codes, without the markers a decoding read adds, since whether it is supported is decided from the
+    metadata. The reader maps the file but touches only the header.
+    """
+    with WrappedGGUFReader(path) as reader:
+        orig_shapes = _read_comfy_orig_shapes(reader)
+        sd: dict[str, torch.Tensor] = {}
+        for tensor in reader.tensors:
+            shape = _logical_shape(tensor, orig_shapes)
+            stored_dtype = torch.from_numpy(np.empty(0, dtype=tensor.data.dtype)).dtype
+            stored_shape = shape if tensor.tensor_type in TORCH_COMPATIBLE_QTYPES else tensor.data.shape
+            sd[tensor.name] = GGMLTensor(
+                torch.empty(stored_shape, dtype=stored_dtype, device="meta"),
+                ggml_quantization_type=tensor.tensor_type,
+                tensor_shape=shape,
+                compute_dtype=compute_dtype,
+            )
+        return sd
+
+
 def gguf_sd_loader(path: Path, compute_dtype: torch.dtype, *, q8_cr: Q8CRMode = "refuse") -> dict[str, torch.Tensor]:
     """Read a GGUF file into a state dict of ``GGMLTensor``s, dequantized on the fly at use.
 
@@ -155,17 +195,14 @@ def gguf_sd_loader(path: Path, compute_dtype: torch.dtype, *, q8_cr: Q8CRMode =
     - ``"refuse"``: any other loader. Raises before the tensors are read; without it the int8 codes reach
       a plain ``nn.Linear`` and the load fails on a dtype error that names neither the file nor the format.
       Identification refuses such files at install already, so this guards records made before it did.
-    - ``"ignore"``: model identification, which needs shapes and never builds a module. The file reads as
-      it always has; whether its quantization is supported is the installer's question, answered from
-      the metadata (see ``ModelConfigFactory``), so it is not raised here once per candidate config.
+
+    Model identification does not come here: it reads the header (`gguf_header_state_dict`).
     """
     with WrappedGGUFReader(path) as reader:
         sd: dict[str, torch.Tensor] = {}
         orig_shapes = _read_comfy_orig_shapes(reader)
-        int8_layers: dict[str, dict[str, Any]] = {}
-        if q8_cr != "ignore":
-            quant_fields = {k: f.contents() for k, f in reader.fields.items() if k.startswith(COMFY_QUANT_KEY_PREFIX)}
-            int8_layers = parse_q8_cr_markers(quant_fields, path.name)
+        quant_fields = {k: f.contents() for k, f in reader.fields.items() if k.startswith(COMFY_QUANT_KEY_PREFIX)}
+        int8_layers = parse_q8_cr_markers(quant_fields, path.name)
         if int8_layers and q8_cr == "refuse":
             raise ValueError(
                 f"{path.name} holds {len(int8_layers)} ComfyUI-GGUF Q8_CR (int8 convrot) layer(s), e.g. "
@@ -184,15 +221,7 @@ def gguf_sd_loader(path: Path, compute_dtype: torch.dtype, *, q8_cr: Q8CRMode =
             # 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
+            shape = _logical_shape(tensor, orig_shapes)
 
             if tensor.name in int8_layers or tensor.name in int8_scales:
                 expected = gguf.GGMLQuantizationType.I8 if tensor.name in int8_layers else gguf.GGMLQuantizationType.F32

tests/backend/quantization/gguf/test_loaders.py

@@ -5,8 +5,9 @@
 import pytest
 import torch
 
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
 from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
-from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader
+from invokeai.backend.quantization.gguf.loaders import gguf_header_state_dict, gguf_sd_loader
 from tests.fixtures.quantized_payloads import q8_cr_marker, quantize_convrot, write_gguf
 
 
@@ -114,16 +115,17 @@ def test_a_q8_cr_file_is_refused_unless_the_caller_decodes_int8(tmp_path):
         gguf_sd_loader(path, compute_dtype=torch.float32)
 
 
-def test_identification_reads_a_q8_cr_file_as_it_always_has(tmp_path):
+def test_the_header_read_reports_q8_cr_codes_as_i8_without_raising(tmp_path):
     """Identification asks only for shapes, once per candidate config. A format it cannot decode must not
     raise there -- the factory refuses it from the metadata, with its reason -- and nothing is translated."""
 
     path = tmp_path / "model.gguf"
     write_gguf(path, _q8_cr_layer("proj"), quant={"proj.weight": {"format": "int4_cr", "backing": "w4a4"}})
 
-    sd = gguf_sd_loader(path, compute_dtype=torch.float32, q8_cr="ignore")
+    sd = gguf_header_state_dict(path, compute_dtype=torch.float32)
 
     assert isinstance(sd["proj.weight"], GGMLTensor)
+    assert sd["proj.weight"]._ggml_quantization_type is gguf.GGMLQuantizationType.I8
     assert "proj.comfy_quant" not in sd
 
 
@@ -190,3 +192,57 @@ def test_a_q8_cr_file_the_decode_cannot_read_is_refused(tmp_path, quant, tensors
 
     with pytest.raises(ValueError, match=message):
         gguf_sd_loader(path, compute_dtype=torch.float32, q8_cr="decode")
+
+
+def _write_mixed_gguf(path) -> None:
+    """One tensor of each storage the releases use: F32, F16, BF16, Q8_0, a K-quant, and a tensor stored flat
+    with its native shape in ``comfy.gguf.orig_shape``."""
+    rng = np.random.default_rng(0)
+    writer = gguf.GGUFWriter(str(path), "flux")
+    writer.add_tensor("norm.weight", rng.standard_normal(64).astype(np.float32))
+    writer.add_tensor("gate.weight", rng.standard_normal((8, 64)).astype(np.float16))
+    bf16 = gguf.quants.quantize(rng.standard_normal((8, 64)).astype(np.float32), gguf.GGMLQuantizationType.BF16)
+    writer.add_tensor("proj.bias", bf16, raw_shape=bf16.shape, raw_dtype=gguf.GGMLQuantizationType.BF16)
+    q8 = gguf.quants.quantize(rng.standard_normal((16, 64)).astype(np.float32), gguf.GGMLQuantizationType.Q8_0)
+    writer.add_tensor("proj.weight", q8, raw_shape=q8.shape, raw_dtype=gguf.GGMLQuantizationType.Q8_0)
+    q4k = rng.integers(0, 256, size=(4, 144), dtype=np.uint8)
+    writer.add_tensor("attn.weight", q4k, raw_shape=q4k.shape, raw_dtype=gguf.GGMLQuantizationType.Q4_K)
+    writer.add_tensor("embedding", rng.standard_normal(96).astype(np.float32))
+    writer.add_array("comfy.gguf.orig_shape.embedding", [1, 96])
+    writer.write_header_to_file()
+    writer.write_kv_data_to_file()
+    writer.write_tensors_to_file()
+    writer.close()
+
+
+def test_the_header_state_dict_is_the_full_read_without_the_data(tmp_path):
+    """Identification probes the header version. Every candidate config reads names, shapes and quantization
+    types off it, so each tensor has to look exactly as the full read's does -- packed and logical shape, storage
+    dtype, type -- while holding nothing."""
+    path = tmp_path / "model.gguf"
+    _write_mixed_gguf(path)
+
+    full = gguf_sd_loader(path, compute_dtype=torch.float32)
+    header = gguf_header_state_dict(path, compute_dtype=torch.float32)
+
+    assert header.keys() == full.keys()
+    for name, read in full.items():
+        probed = header[name]
+        assert isinstance(probed, GGMLTensor) and isinstance(read, GGMLTensor), name
+        assert probed.quantized_data.is_meta, name
+        assert probed._ggml_quantization_type == read._ggml_quantization_type, name
+        assert probed.tensor_shape == read.tensor_shape, name
+        assert probed.quantized_data.shape == read.quantized_data.shape, name
+        assert probed.quantized_data.dtype == read.quantized_data.dtype, name
+        assert probed.compute_dtype == read.compute_dtype, name
+    assert header["embedding"].shape == (1, 96)
+
+
+def test_identification_reads_only_the_header(tmp_path):
+    """The install probes every file once per candidate config; for a GGUF that must not copy its tensors."""
+    path = tmp_path / "model.gguf"
+    _write_mixed_gguf(path)
+
+    sd = ModelOnDisk(path).load_state_dict()
+
+    assert sd and all(isinstance(value, GGMLTensor) and value.quantized_data.is_meta for value in sd.values())