fix(gguf): unwrap a (state dict, architecture) return; drop the disk cache

ComfyUI-GGUF hands back either the state dict or a (state_dict, arch)
pair depending on version. Passed on whole, ComfyUI iterates the tuple,
gets the dict as its first element and raises "'dict' object has no
attribute 'startswith'" — an error naming nothing that led back here.
Unwrapped, and the shape check now runs on every load rather than only
when FunPack had overridden the architecture list, which is why a pack
that already knew the architecture went through unvalidated.

The expansion cache is gone. It wrote the model's full dequantized size
beside it, and rented disk is billed. A leftover from it is named with
its size on the next load; deleting it is left to the user.
This commit is contained in:
digital-garbage
2026-08-23 05:53:26 +03:00
parent ae337b5196
commit c4d8018c09
2 changed files with 111 additions and 150 deletions
+39 -71
View File
@@ -230,6 +230,7 @@ def _load(path: str, clip: bool) -> tuple[dict, dict, str]:
which = backend()
if which is None:
raise RuntimeError(UNAVAILABLE)
leftover = note_leftover_cache(path)
if which == "pack":
d = _pack_dir()
try:
@@ -240,86 +241,53 @@ def _load(path: str, clip: bool) -> tuple[dict, dict, str]:
if clip:
read = getattr(loader, "gguf_clip_loader", None) or read
sd = read(path)
# The pack's own guard is what we just stepped around, so the check it was doing
# has to happen here instead: an untested architecture can return something that
# is not a state dict at all, and handing that to ComfyUI fails deep inside core
# with a message that makes no sense out here. Anything unexpected raises, which
# the fallback below turns into the slow-but-working path.
if allowed:
_assert_state_dict(sd, path)
# Some versions hand back (state_dict, architecture) rather than the dict alone.
# Passed on as-is, ComfyUI iterates the TUPLE, gets the dict as its first element
# and fails with "'dict' object has no attribute 'startswith'" — an error with
# nothing in it that points back here.
if isinstance(sd, (tuple, list)) and len(sd) == 2 and isinstance(sd[0], dict):
sd = sd[0]
# Checked on every path, not just when the architecture list was overridden: the
# shape of what the pack returns is a moving target across its versions, and
# anything unexpected fails deep inside core where the message makes no sense out
# here. Raising turns it into the fallback below instead.
_assert_state_dict(sd, path)
options = {}
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)})"
+ (f"; {allowed}" if allowed else ""))
+ (f"; {allowed}" if allowed else "")
+ (f"; {leftover}" if leftover 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_cached(path)
sd, options, note = _load_native(path)
return sd, options, (f"gguf: ComfyUI-GGUF could not read this file ({e}) — "
f"fell back. {note}")
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")
f"fell back. {note}"
+ (f" {leftover}" if leftover else ""))
sd, options, note = _load_native(path)
return sd, options, note + (f" {leftover}" if leftover else "")
LEFTOVER_SUFFIX = ".dequantized.safetensors"
def note_leftover_cache(path: str) -> str:
"""FunPack used to keep an expanded copy of a .gguf beside it. It no longer does — on a
rented box disk is billed, and the copy is the size of the whole model. Any file left
over from that is dead weight the user has no reason to guess at, so name it and its
size once. Never deleted here: it is a big file in the user's model folder."""
leftover = path + LEFTOVER_SUFFIX
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
if not os.path.isfile(leftover):
return ""
gb = os.path.getsize(leftover) / 1024 ** 3
except OSError:
return ""
return (f"an old FunPack expansion cache is still beside this model "
f"({os.path.basename(leftover)}, {gb:.1f} GB) — nothing reads it any more, "
f"delete it to reclaim the space")
def _assert_state_dict(sd, path: str) -> None:
@@ -420,11 +388,11 @@ def _load_native(path: str) -> tuple[dict, dict, str]:
# Bounded, and deliberately small: the window below keeps `workers * 2` float32 tensors
# alive at once, and the largest weight in a video model is around a gigabyte at that
# width. More threads would expand faster and spike higher; the durable speed win is the
# dequantized cache in `_load_native_cached`, not this pool.
# ComfyUI-GGUF pack's quantized path, not this pool.
workers = max(1, min(4, (os.cpu_count() or 2) - 1))
print(f"[FunPack] gguf: expanding {quantized} of {len(tensors)} tensors from "
f"{os.path.basename(path)} across {workers} thread(s). This is the slow path — "
f"minutes on a video checkpoint, and the result is cached beside the file.",
f"minutes on a video checkpoint. Install ComfyUI-GGUF to skip it entirely.",
flush=True)
def _expand(tensor):
+72 -79
View File
@@ -124,12 +124,13 @@ def test_pack_load_passes_custom_operations_through(tmp_path, monkeypatch):
d = tmp_path / "pack"
d.mkdir()
(d / "loader.py").write_text(
"def gguf_sd_loader(path):\n return {'w': 'quantized'}\n")
"class _T:\n shape = (1,)\n\n"
"def gguf_sd_loader(path):\n return {'w': _T()}\n")
(d / "ops.py").write_text("class GGMLOps:\n pass\n")
monkeypatch.setattr(gguf_support, "_pack_dir", lambda: str(d))
sd, options, note = gguf_support.load_state_dict("/whatever.gguf")
assert sd == {"w": "quantized"}
assert list(sd) == ["w"]
assert "custom_operations" in options
assert type(options["custom_operations"]).__name__ == "GGMLOps"
assert "quantized" in note
@@ -163,18 +164,19 @@ def test_pack_modules_can_use_relative_imports(tmp_path, monkeypatch):
d.mkdir()
(d / "__init__.py").write_text("raise AssertionError('the pack __init__ must not run')")
(d / "ops.py").write_text("class GGMLOps:\n pass\n")
(d / "dequant.py").write_text("MARKER = 'dequant'\n")
(d / "dequant.py").write_text(
"class Marker:\n shape = (1,)\n name = 'dequant'\n")
(d / "loader.py").write_text(
"from .ops import GGMLOps\n"
"from .dequant import MARKER\n"
"from .dequant import Marker\n"
"def gguf_sd_loader(path):\n"
" return {'w': MARKER}\n"
" return {'w': Marker()}\n"
)
monkeypatch.setattr(gguf_support, "_pack_dir", lambda: str(d))
monkeypatch.delitem(sys.modules, gguf_support._PKG_NAME, raising=False)
sd, options, note = gguf_support.load_state_dict("/whatever.gguf")
assert sd == {"w": "dequant"}
assert sd["w"].name == "dequant"
assert type(options["custom_operations"]).__name__ == "GGMLOps"
@@ -186,15 +188,16 @@ def test_switching_pack_directories_remounts(tmp_path, monkeypatch):
d.mkdir()
(d / "ops.py").write_text("class GGMLOps:\n pass\n")
(d / "loader.py").write_text(
f"def gguf_sd_loader(path):\n return {{'w': '{marker}'}}\n")
f"class _T:\n shape = (1,)\n name = '{marker}'\n"
f"def gguf_sd_loader(path):\n return {{'w': _T()}}\n")
monkeypatch.delitem(sys.modules, gguf_support._PKG_NAME, raising=False)
monkeypatch.setattr(gguf_support, "_pack_dir", lambda: str(tmp_path / "a"))
assert gguf_support.load_state_dict("/x.gguf")[0] == {"w": "first"}
assert gguf_support.load_state_dict("/x.gguf")[0]["w"].name == "first"
# No manual cache clearing: remounting a different directory must purge it by itself.
monkeypatch.setattr(gguf_support, "_pack_dir", lambda: str(tmp_path / "b"))
assert gguf_support.load_state_dict("/x.gguf")[0] == {"w": "second"}
assert gguf_support.load_state_dict("/x.gguf")[0]["w"].name == "second"
def test_pack_refusing_an_architecture_falls_back(tmp_path, monkeypatch):
@@ -250,76 +253,6 @@ def test_gguf_is_a_declared_requirement():
# 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)
@@ -465,3 +398,63 @@ def test_load_native_rejects_a_tensor_that_contradicts_its_header(monkeypatch):
assert "'wrong'" in str(excinfo.value)
assert "(25)" in str(excinfo.value)
def test_a_state_dict_returned_beside_its_architecture_is_unwrapped(tmp_path, monkeypatch):
"""Some pack versions return (state_dict, architecture). Passed on whole, ComfyUI iterates
the TUPLE and fails with "'dict' object has no attribute 'startswith'" — an error with
nothing in it that points back at the loader."""
d = tmp_path / "pack"
d.mkdir()
(d / "ops.py").write_text("class GGMLOps:\n pass\n")
(d / "loader.py").write_text(
"class _T:\n shape = (1,)\n\n"
"def gguf_sd_loader(path):\n return ({'w': _T()}, 'minimax_h3')\n")
monkeypatch.setattr(gguf_support, "_pack_dir", lambda: str(d))
monkeypatch.delitem(sys.modules, gguf_support._PKG_NAME, raising=False)
sd, options, note = gguf_support.load_state_dict("/whatever.gguf")
assert isinstance(sd, dict) and list(sd) == ["w"]
assert "custom_operations" in options
def test_a_known_architecture_is_still_validated(tmp_path, monkeypatch):
"""Validation used to run only when FunPack had overridden the architecture list, so a
pack that already knew the architecture returned straight through unchecked."""
d = tmp_path / "pack"
d.mkdir()
(d / "ops.py").write_text("class GGMLOps:\n pass\n")
(d / "loader.py").write_text("def gguf_sd_loader(path):\n return ['not', 'a', 'dict']\n")
monkeypatch.setattr(gguf_support, "_pack_dir", lambda: str(d))
monkeypatch.setattr(gguf_support, "_allow_architectures", lambda loader, path: "")
monkeypatch.setattr(gguf_support, "_load_native",
lambda path: ({"w": 1}, {}, "gguf: dequantized at load"))
monkeypatch.delitem(sys.modules, gguf_support._PKG_NAME, raising=False)
sd, options, note = gguf_support.load_state_dict("/whatever.gguf")
assert sd == {"w": 1} # refused, and fell back rather than reaching ComfyUI
assert "could not read this file" in note
def test_an_old_expansion_cache_is_named_but_never_deleted(tmp_path):
"""FunPack no longer writes one. Disk is billed on a rental, so a leftover has to be
findable — and it is the user's file, in the user's model folder, so it is not removed."""
model = tmp_path / "m.gguf"
model.write_bytes(b"GGUF")
leftover = tmp_path / ("m.gguf" + gguf_support.LEFTOVER_SUFFIX)
leftover.write_bytes(b"x" * 2048)
note = gguf_support.note_leftover_cache(str(model))
assert "m.gguf.dequantized.safetensors" in note
assert "delete it" in note
assert leftover.exists()
def test_no_leftover_cache_says_nothing(tmp_path):
model = tmp_path / "m.gguf"
model.write_bytes(b"GGUF")
assert gguf_support.note_leftover_cache(str(model)) == ""