From 78c9625a1aefc27df913cc75e4b9c83995209627 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 23 Aug 2026 05:04:07 +0300 Subject: [PATCH] perf(gguf): load quantized where possible, expand only once Add an unlisted architecture to ComfyUI-GGUF's image list so the quantized path works, and cache the dequantized weights beside the file when it does not, so the expansion is paid once. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 12 ++++ gguf_support.py | 123 ++++++++++++++++++++++++++++++++++++- tests/test_gguf_support.py | 111 +++++++++++++++++++++++++++++++++ 3 files changed, 243 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b23780a1..2ed00b5d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,18 @@ minute of dequantizing a GGUF to alter a sparsity ratio. The weights depend only on the file (identity, size, mtime) and the dtypes; everything else is applied to a clone. One entry, replaced whenever the key changes. +- **A GGUF architecture ComfyUI-GGUF has not been told about now loads quantized anyway.** + The pack refuses anything off its tested list — `Unexpected architecture type in GGUF file: + 'minimax_h3'` — which cost the quantized path entirely and fell back to expanding the whole + checkpoint. The name is now added to its IMAGE list (never the text one) for a file whose + header actually declares it, and the status line says so, because it is an override of + another project's own guard. +- **When it does have to expand a GGUF, it expands ONCE.** The dequantized weights are + written beside the file as an ordinary safetensors and reused on every later load, which + then runs at safetensors speed. It costs the model's dequantized size in disk — the same + size that path was already going to occupy in VRAM — and the status line names the file to + delete. A cache older than the model is ignored; one that cannot be written is a slow + load, not a failed one. - **GGUF dequantizes across threads instead of one tensor at a time.** The native path expands every quantized tensor at load with numpy, serially, which is where the wait went. numpy releases the GIL for that work, so it now runs on up to 8 threads — bounded rather diff --git a/gguf_support.py b/gguf_support.py index 8eb13ae4..0d6d244d 100644 --- a/gguf_support.py +++ b/gguf_support.py @@ -234,6 +234,7 @@ def _load(path: str, clip: bool) -> tuple[dict, dict, str]: try: loader = _load_pack_module(d, "loader") ops = _load_pack_module(d, "ops") + allowed = _allow_architectures(loader, path) read = loader.gguf_sd_loader if clip: read = getattr(loader, "gguf_clip_loader", None) or read @@ -242,14 +243,130 @@ def _load(path: str, clip: bool) -> tuple[dict, dict, str]: ggml_ops = getattr(ops, "GGMLOps", None) if ggml_ops is not None: options["custom_operations"] = ggml_ops() - return sd, options, f"gguf: quantized (ComfyUI-GGUF at {os.path.basename(d)})" + return sd, options, (f"gguf: quantized (ComfyUI-GGUF at {os.path.basename(d)})" + + (f"; {allowed}" if allowed else "")) except Exception as e: # noqa: BLE001 — any refusal is a reason to fall back if importlib.util.find_spec("gguf") is None: raise - sd, options, note = _load_native(path) + sd, options, note = _load_native_cached(path) return sd, options, (f"gguf: ComfyUI-GGUF could not read this file ({e}) — " f"fell back. {note}") - return _load_native(path) + return _load_native_cached(path) + + +CACHE_SUFFIX = ".dequantized.safetensors" + + +def cache_path(path: str) -> str: + """Where a dequantized copy of `path` lives. Beside the original, so it is obvious what + it belongs to and deleting the model takes it with you.""" + return path + CACHE_SUFFIX + + +def _cache_is_current(src: str, dst: str) -> bool: + """Only trust a cache written AFTER the file it came from.""" + try: + return os.path.isfile(dst) and os.path.getmtime(dst) >= os.path.getmtime(src) + except OSError: + return False + + +def _load_native_cached(path: str) -> tuple[dict, dict, str]: + """`_load_native`, but the expansion happens once and is kept on disk. + + The native path has to expand every quantized tensor before the model can be used, and + that is minutes on a video checkpoint — paid again on every launch, for a result that is + identical every time. Written next to the file as an ordinary safetensors, which then + loads at safetensors speed (memory-mapped, no dequantization at all). + + It costs the model's dequantized size in DISK. That is the same size it was going to + occupy in VRAM on this path anyway, so nothing is being traded that was not already + spent — but it is real disk, and it is stated rather than left to be discovered. + """ + import comfy.utils + + dst = cache_path(path) + if _cache_is_current(path, dst): + try: + sd = comfy.utils.load_torch_file(dst) + return sd, {}, (f"gguf: loaded the dequantized cache beside it " + f"({os.path.basename(dst)}) — delete that file to rebuild it") + except Exception as e: # noqa: BLE001 + print(f"[FunPack] gguf: the dequantized cache could not be read ({e}) — " + f"expanding the .gguf again") + + sd, options, note = _load_native(path) + try: + import safetensors.torch + tmp = dst + ".part" + safetensors.torch.save_file({k: v.contiguous() for k, v in sd.items()}, tmp) + os.replace(tmp, dst) # atomic: a half-written cache is never left behind + note += (f" Written to {os.path.basename(dst)} so the next load skips this " + f"(costs its dequantized size in disk; delete it to reclaim).") + except Exception as e: # noqa: BLE001 + # A cache that cannot be written is a slow load, not a failed one. + note += f" (could not write the dequantized cache: {e})" + try: + if os.path.exists(dst + ".part"): + os.remove(dst + ".part") + except OSError: + pass + return sd, options, note + + +def read_architecture(path: str) -> Optional[str]: + """`general.architecture` out of a GGUF header, or None. Header only — no tensor data.""" + try: + import gguf + reader = gguf.GGUFReader(path) + field = reader.fields.get("general.architecture") + if field is None: + return None + return str(bytes(field.parts[field.data[0]]), encoding="utf-8") + except Exception: # noqa: BLE001 + return None + + +def _allow_architectures(loader, path: str) -> str: + """Let ComfyUI-GGUF read an architecture its allowlist does not name. Returns a note. + + The pack keeps a list of architectures it has been tested against and refuses anything + else. That is the right default for the pack and a dead end here: the refusal costs the + QUANTIZED path, and the fallback expands the whole checkpoint at load — minutes of + dequantizing, and the file's full size in VRAM, which is the entire reason someone chose + a GGUF. + + The reading itself is generic: the pack walks the tensors and wraps them, and the + architecture only steers the key renaming that image models do not use. So the name is + added to the IMAGE list, never the text one, and only for a file whose header actually + declares it. Announced, because it is an override of another project's own guard and a + wrong result here should be traceable to this line. + """ + arch = read_architecture(path) + if not arch: + return "" + added = [] + for attr in ("IMG_ARCH_LIST", "ARCH_LIST"): + lst = getattr(loader, attr, None) + if lst is None: + continue + try: + if arch in lst: + return "" # the pack already knows it; nothing to override + if isinstance(lst, set): + lst.add(arch) + elif isinstance(lst, list): + lst.append(arch) + else: + continue + added.append(attr) + except Exception: # noqa: BLE001 + continue + if not added: + return "" + return (f"architecture {arch!r} is not on ComfyUI-GGUF's tested list and FunPack added it " + f"so the file loads quantized instead of being expanded — if the weights come out " + f"wrong, this is why") def _load_native(path: str) -> tuple[dict, dict, str]: diff --git a/tests/test_gguf_support.py b/tests/test_gguf_support.py index 8d74b174..a2bf2e6e 100644 --- a/tests/test_gguf_support.py +++ b/tests/test_gguf_support.py @@ -5,6 +5,7 @@ which backend gets picked, what happens when none is available, and whether .ggu are found at all (core's extension set excludes them, which is the whole reason this module exists). """ +import os import sys import types @@ -240,3 +241,113 @@ def test_gguf_is_a_declared_requirement(): reqs = pathlib.Path("requirements.txt").read_text(encoding="utf-8") lines = [ln.strip() for ln in reqs.splitlines() if ln.strip() and not ln.startswith("#")] assert any(ln == "gguf" or ln.startswith("gguf") for ln in lines) + + +# ── expand once, not every launch ───────────────────────────────────────────── +# The native path has to dequantize every tensor before the model is usable — minutes on a +# video checkpoint, paid again on every launch for a result that is identical every time. + + +def test_the_cache_sits_beside_the_model(tmp_path): + src = str(tmp_path / "m.gguf") + assert gguf_support.cache_path(src) == src + gguf_support.CACHE_SUFFIX + + +def test_a_cache_older_than_the_model_is_not_trusted(tmp_path): + src = tmp_path / "m.gguf" + src.write_bytes(b"GGUF") + dst = tmp_path / ("m.gguf" + gguf_support.CACHE_SUFFIX) + dst.write_bytes(b"old") + os.utime(dst, (1, 1)) # written before the model + assert not gguf_support._cache_is_current(str(src), str(dst)) + + +def test_a_cache_newer_than_the_model_is_used(tmp_path, monkeypatch): + src = tmp_path / "m.gguf" + src.write_bytes(b"GGUF") + dst = tmp_path / ("m.gguf" + gguf_support.CACHE_SUFFIX) + dst.write_bytes(b"new") + os.utime(src, (1, 1)) + assert gguf_support._cache_is_current(str(src), str(dst)) + + +def test_a_missing_cache_is_simply_not_current(tmp_path): + src = tmp_path / "m.gguf" + src.write_bytes(b"GGUF") + assert not gguf_support._cache_is_current(str(src), str(tmp_path / "nope")) + + +def test_the_expansion_is_written_and_reused(tmp_path, monkeypatch): + """The whole point: the second load must not dequantize anything.""" + import torch + src = tmp_path / "m.gguf" + src.write_bytes(b"GGUF") + calls = [] + + def fake_native(path): + calls.append(path) + return {"w": torch.zeros(2, 2)}, {}, "gguf: dequantized at load" + + monkeypatch.setattr(gguf_support, "_load_native", fake_native) + # The stub comfy tree has no real loader; the cache is an ordinary safetensors file. + import safetensors.torch + monkeypatch.setattr("comfy.utils.load_torch_file", + lambda p, **kw: safetensors.torch.load_file(p), raising=False) + sd1, _, note1 = gguf_support._load_native_cached(str(src)) + assert calls == [str(src)] and "Written to" in note1 + sd2, _, note2 = gguf_support._load_native_cached(str(src)) + assert calls == [str(src)] # never expanded a second time + assert "dequantized cache" in note2 + assert torch.equal(sd1["w"], sd2["w"]) + + +def test_an_unwritable_cache_is_a_slow_load_not_a_failed_one(tmp_path, monkeypatch): + import torch + src = tmp_path / "m.gguf" + src.write_bytes(b"GGUF") + monkeypatch.setattr(gguf_support, "_load_native", + lambda p: ({"w": torch.zeros(2, 2)}, {}, "expanded")) + monkeypatch.setattr("safetensors.torch.save_file", + lambda *a, **kw: (_ for _ in ()).throw(OSError("disk full"))) + sd, _, note = gguf_support._load_native_cached(str(src)) + assert "w" in sd and "could not write" in note + assert not os.path.exists(gguf_support.cache_path(str(src)) + ".part") # no debris + + +# ── letting the pack read an architecture it has not been told about ────────── + + +class _Loader: + def __init__(self, arches): + self.IMG_ARCH_LIST = set(arches) + + +def test_an_unknown_architecture_is_added_so_the_file_loads_quantized(monkeypatch): + monkeypatch.setattr(gguf_support, "read_architecture", lambda p: "minimax_h3") + loader = _Loader({"flux", "sd3"}) + note = gguf_support._allow_architectures(loader, "x.gguf") + assert "minimax_h3" in loader.IMG_ARCH_LIST + assert "minimax_h3" in note and "FunPack added it" in note + + +def test_a_known_architecture_is_left_alone_and_says_nothing(monkeypatch): + monkeypatch.setattr(gguf_support, "read_architecture", lambda p: "flux") + loader = _Loader({"flux"}) + assert gguf_support._allow_architectures(loader, "x.gguf") == "" + + +def test_an_unreadable_architecture_changes_nothing(monkeypatch): + monkeypatch.setattr(gguf_support, "read_architecture", lambda p: None) + loader = _Loader({"flux"}) + assert gguf_support._allow_architectures(loader, "x.gguf") == "" + assert loader.IMG_ARCH_LIST == {"flux"} + + +def test_the_text_encoder_list_is_never_touched(monkeypatch): + """Image and text architectures steer different key renaming in the pack; guessing a + text arch into the image list is one thing, the reverse is another.""" + monkeypatch.setattr(gguf_support, "read_architecture", lambda p: "minimax_h3") + loader = _Loader({"flux"}) + loader.TXT_ARCH_LIST = {"t5"} + gguf_support._allow_architectures(loader, "x.gguf") + assert loader.TXT_ARCH_LIST == {"t5"}