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:
+39
-71
@@ -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
@@ -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)) == ""
|
||||
|
||||
Reference in New Issue
Block a user