nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

revert(gguf): back to the load path from 6f04642

Browse Source 
Reverts 6bf533e, 78c9625, 4bacc9c, 993e5e4, ae337b5 and c4d8018 —
gguf_support.py, loaders.py and their tests are byte-identical to
6f04642. Nothing else is touched.

The quantized path loaded fast and gave the sampler a model it could not
run: shape [21504, 5376] against an input of 22848x5376, one 1344-token
unit apart. Expanding at load is slow and works.
This commit is contained in:
digital-garbage
2026-08-23 08:12:43 +03:00
parent c4d8018c09
commit a490b8a6ed
5 changed files with 36 additions and 575 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+6 -23
View File
@@ -2,30 +2,13 @@
## [Unreleased]
### Reverted
- **GGUF loading is back to expanding the checkpoint at load.** The quantized path through
ComfyUI-GGUF loaded in seconds instead of a minute, and produced a model the sampler could
not run: a packed sequence one 1344-token unit short of what the model built. Reverted
whole rather than patched further.
### Fixed
- **Changing an attention or SLA setting no longer re-reads the checkpoint.** They live on
the loader node, so ComfyUI invalidated the whole node and the file was loaded again — a
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 expanded the whole checkpoint
instead. The name is added to its IMAGE list (never the text one) for a file whose header
declares it. Since that steps around the pack's own guard, what the backend hands back is
checked here instead: anything that is not a state dict of named tensors raises, and the
existing fallback turns that into the slow-but-working path. The status line says the
override is in effect, because a wrong generation should be traceable to it.
- **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
than one per core, because each worker holds a dequantized tensor and its copy.
- **The updater installs only what is MISSING, and never upgrades anything.** It ran
`pip install -r requirements.txt`, which upgrades any package below its version floor —
and on an install full of compiled extensions (torch, comfy-kitchen, comfy-aimdo,
gguf_support.py
+21 -175
View File
@@ -27,7 +27,6 @@ from __future__ import annotations
import importlib
import importlib.util
import itertools
import os
import sys
import types
@@ -230,138 +229,27 @@ 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:
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
sd = read(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"; {leftover}" if leftover else ""))
return sd, options, f"gguf: quantized (ComfyUI-GGUF at {os.path.basename(d)})"
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)
return sd, options, (f"gguf: ComfyUI-GGUF could not read this file ({e}) — "
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:
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:
"""Raise unless `sd` is a usable state dict. Cheap, and only shape — nothing here can
tell whether the NUMBERS are right, which is why the status line says to suspect this
first if a generation comes out wrong."""
if not isinstance(sd, dict) or not sd:
raise RuntimeError(
f"{os.path.basename(path)}: the GGUF backend returned "
f"{type(sd).__name__} instead of a state dict")
for key, value in sd.items():
if not isinstance(key, str):
raise RuntimeError(
f"{os.path.basename(path)}: the GGUF backend returned a state dict keyed by "
f"{type(key).__name__}, not by tensor name")
if not hasattr(value, "shape"):
raise RuntimeError(
f"{os.path.basename(path)}: {key!r} came back as {type(value).__name__}, "
f"which is not a tensor")
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 "
f"it, so the file loads quantized instead of being expanded — if the weights come "
f"out wrong, this is the first thing to suspect")
f"fell back. {note}")
return _load_native(path)
def _load_native(path: str) -> tuple[dict, dict, str]:
@@ -369,44 +257,24 @@ def _load_native(path: str) -> tuple[dict, dict, str]:
Correct but not thrifty: the model ends up at its dequantized size, so this buys the
ability to LOAD a .gguf, not GGUF's memory saving. The caller states that.
Expansion is STREAMED. Dequantizing produces float32 — four times the size the weight
will occupy once it is cast — so collecting every tensor before converting any of them
peaks at several times the finished model in host RAM. On a video checkpoint that is
tens of gigabytes of transient allocation, which is enough to take the machine down
rather than merely the load. Each tensor is therefore cast and released as it is
produced, and only a small window is ever in flight.
"""
import gguf
import numpy as np
import torch
reader = gguf.GGUFReader(path)
tensors = list(reader.tensors)
raw = (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16)
quantized = sum(1 for t in tensors if t.tensor_type not in raw)
# 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
# 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. Install ComfyUI-GGUF to skip it entirely.",
flush=True)
def _expand(tensor):
"""One tensor, fully finished: nothing float32 outlives this call."""
qtype = tensor.tensor_type
arr = tensor.data if qtype in raw else gguf.quants.dequantize(tensor.data, qtype)
if arr.dtype == np.float32:
# `.to` allocates the half-width copy and the float32 array is freed on return,
# so the wide form never accumulates.
t = torch.from_numpy(arr).to(torch.float16)
else:
# Detach from the memory-mapped file; the reader's mapping does not outlive us.
t = torch.from_numpy(arr.copy())
sd: dict = {}
quantized = 0
for tensor in reader.tensors:
name = str(tensor.name)
qtype = tensor.tensor_type
if qtype in (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
arr = tensor.data
else:
# gguf.quants.dequantize is the reference implementation for every type the
# package knows — do not reimplement it per quant type.
arr = gguf.quants.dequantize(tensor.data, qtype)
quantized += 1
t = torch.from_numpy(arr.copy())
# GGUF stores dimensions in the opposite order to torch.
shape = tuple(int(d) for d in reversed(tensor.shape))
expected = 1
@@ -419,32 +287,10 @@ def _load_native(path: str) -> tuple[dict, dict, str]:
f"{os.path.basename(path)}: tensor {name!r} dequantized to {t.numel()} "
f"elements but its header declares {shape} ({expected}). This container is "
f"not laid out the way the gguf package describes it.")
return name, t.reshape(shape) if shape else t
sd: dict = {}
if workers > 1 and len(tensors) > 1:
from collections import deque
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor(max_workers=workers) as pool:
remaining = iter(tensors)
# A fixed window rather than `pool.map`, which submits every tensor at once and
# buffers each result until the one before it is consumed — the whole model in
# float32, which is the allocation this function exists to avoid.
flight = deque(pool.submit(_expand, t)
for t in itertools.islice(remaining, workers * 2))
while flight:
name, t = flight.popleft().result()
sd[name] = t
nxt = next(remaining, None)
if nxt is not None:
flight.append(pool.submit(_expand, nxt))
else:
for tensor in tensors:
name, t = _expand(tensor)
sd[name] = t
note = (f"gguf: dequantized at load ({quantized} quantized tensors expanded across "
f"{workers} thread(s)) — the file loads, but it occupies its full size in VRAM, "
f"and expanding it is why this is slower than a .safetensors. Install "
f"ComfyUI-GGUF for a lazy quantized load.")
if shape:
t = t.reshape(shape)
sd[name] = t.to(torch.float16) if t.dtype == torch.float32 else t
note = (f"gguf: dequantized at load ({quantized} quantized tensors expanded) — the file "
f"loads, but it occupies its full size in VRAM. Install ComfyUI-GGUF to keep it "
f"quantized.")
return sd, {}, note
loaders.py
-80
View File
@@ -11,7 +11,6 @@ SageAttention is installed, flash when flash-attn is, xformers when enabled —
lies about the rest.
"""
import logging
import os
import comfy.sd
import comfy.utils
@@ -129,35 +128,6 @@ def set_fp16_accumulation(enabled):
return bool(enabled)
# One loaded checkpoint, so changing an inference setting on the loader does not re-read the
# file. Replaced whenever the key changes, so this never holds more than one model.
#
# The reference is strong on purpose: ComfyUI drops this node's cached output BEFORE it
# re-runs it, which is exactly the moment a weak one would be dead. The cost is that this
# entry outlives a "free model and node cache", so one checkpoint stays resident until a
# different file or dtype replaces it. Loading anything else releases it.
_WEIGHT_CACHE: dict = {"key": None, "model": None}
def _weight_cache_key(model_name, weight_dtype, compute_dtype):
"""Everything that changes the WEIGHTS: the file (by identity, size and mtime) and the
dtypes. Not the attention backend, not SLA — those are applied to a clone.
Size and mtime are in the key so a file replaced under the same name is re-read. An
unreadable stat returns None, which never matches and so always reloads: failing to
notice a changed file is worse than an occasional extra read.
"""
try:
path = gguf_support.gguf_path("diffusion_models", model_name) if gguf_support.is_gguf(model_name) \
else folder_paths.get_full_path("diffusion_models", model_name)
if not path:
return None
st = os.stat(path)
return (path, st.st_size, st.st_mtime_ns, str(weight_dtype), str(compute_dtype))
except Exception: # noqa: BLE001
return None
class FunPackDiffusionModelLoader:
"""Loads a diffusion model with quantization and an attention backend chosen per node."""
@@ -259,24 +229,6 @@ class FunPackDiffusionModelLoader:
fp16_accumulation=False, sla_sparsity=None, sla_block_size=None,
sla_protect_audio=None, sla_min_seq_len=None, sla_dense_last_steps=None,
sla=False):
# ComfyUI reports a node failure as "<node>: <message>" with the traceback going to
# the console. A bare "'dict' object has no attribute 'startswith'" names neither the
# step it came from nor the file it was reading, which is most of what you need. Each
# stage sets `stage`, and the re-raise says which one.
stage = {"at": "starting"}
try:
return self._load_model(stage, model_name, weight_dtype, compute_dtype, attention,
fp16_accumulation, sla_sparsity, sla_block_size,
sla_protect_audio, sla_min_seq_len, sla_dense_last_steps,
sla)
except Exception as e:
raise type(e)(f"{e} [while {stage['at']}: {model_name}]").with_traceback(
e.__traceback__) from e
def _load_model(self, stage, model_name, weight_dtype, compute_dtype, attention,
fp16_accumulation=False, sla_sparsity=None, sla_block_size=None,
sla_protect_audio=None, sla_min_seq_len=None, sla_dense_last_steps=None,
sla=False):
notes = [f"FunPack Diffusion Model Loader | {model_name}"]
accum = set_fp16_accumulation(fp16_accumulation)
@@ -288,22 +240,6 @@ class FunPackDiffusionModelLoader:
model_options = weight_model_options(weight_dtype)
notes.append(f"weight dtype: {weight_dtype}")
# Attention and SLA are INFERENCE settings, but they live on this node, so ComfyUI
# invalidates the whole node when one changes and the checkpoint is read from disk
# again — a minute of dequantizing a GGUF to alter a sparsity ratio. The weights only
# depend on the file and the dtypes, so those are the cache key and everything else
# is applied to a clone. One entry, replaced whenever the key changes: this holds no
# more than ComfyUI's own node-output cache already did.
stage["at"] = "checking the weight cache"
cache_key = _weight_cache_key(model_name, weight_dtype, compute_dtype)
cached = _WEIGHT_CACHE.get("model") if _WEIGHT_CACHE.get("key") == cache_key else None
if cached is not None:
notes.append("weights: reused (unchanged file and dtypes)")
model = cached.clone()
return self._finish(model, notes, attention, sla, sla_sparsity, sla_block_size,
sla_min_seq_len, sla_dense_last_steps, sla_protect_audio)
stage["at"] = "locating the file"
path = None
misnamed = False
if gguf_support.is_gguf(model_name):
@@ -320,7 +256,6 @@ class FunPackDiffusionModelLoader:
if misnamed:
notes.append(f"{model_name} is named .safetensors but is a GGUF container — "
f"loaded as GGUF")
stage["at"] = "reading the GGUF container"
state_dict, gguf_options, gguf_note = gguf_support.load_state_dict(path)
# The quantized path needs its own torch operations, and they must not be lost to
# the dtype options merged above — a GGUF loaded with stock ops would try to matmul
@@ -329,9 +264,7 @@ class FunPackDiffusionModelLoader:
metadata = None
notes.append(gguf_note)
else:
stage["at"] = "reading the safetensors file"
state_dict, metadata = comfy.utils.load_torch_file(path, return_metadata=True)
stage["at"] = "handing the state dict to ComfyUI"
model = comfy.sd.load_diffusion_model_state_dict(
state_dict, model_options=model_options, metadata=metadata)
if model is None:
@@ -347,19 +280,6 @@ class FunPackDiffusionModelLoader:
elif dtype is not None:
notes.append("compute dtype: unsupported by this ComfyUI, ignored")
# Kept BEFORE the attention work, so what is cached is the plain weights. Everything
# below is applied to a clone, which is what lets a sparsity change skip the disk.
_WEIGHT_CACHE["key"] = cache_key
_WEIGHT_CACHE["model"] = model
stage["at"] = "applying attention / SLA"
return self._finish(model.clone(), notes, attention, sla, sla_sparsity,
sla_block_size, sla_min_seq_len, sla_dense_last_steps,
sla_protect_audio)
def _finish(self, model, notes, attention, sla, sla_sparsity, sla_block_size,
sla_min_seq_len, sla_dense_last_steps, sla_protect_audio):
"""Attention + SLA onto an already-loaded model. Separate from the read so it can run
against a cached one."""
override = attention_override(attention)
# SLA wraps the chosen backend rather than replacing it: there is one override
tests/test_gguf_support.py
+9 -227
View File
@@ -5,14 +5,10 @@ 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 gc
import os
import sys
import types
import weakref
import pytest
import torch
sys.path.insert(0, ".")
import gguf_support # noqa: E402
@@ -124,13 +120,12 @@ def test_pack_load_passes_custom_operations_through(tmp_path, monkeypatch):
d = tmp_path / "pack"
d.mkdir()
(d / "loader.py").write_text(
"class _T:\n shape = (1,)\n\n"
"def gguf_sd_loader(path):\n return {'w': _T()}\n")
"def gguf_sd_loader(path):\n return {'w': 'quantized'}\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 list(sd) == ["w"]
assert sd == {"w": "quantized"}
assert "custom_operations" in options
assert type(options["custom_operations"]).__name__ == "GGMLOps"
assert "quantized" in note
@@ -164,19 +159,18 @@ 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(
"class Marker:\n shape = (1,)\n name = 'dequant'\n")
(d / "dequant.py").write_text("MARKER = '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"].name == "dequant"
assert sd == {"w": "dequant"}
assert type(options["custom_operations"]).__name__ == "GGMLOps"
@@ -188,16 +182,15 @@ 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"class _T:\n shape = (1,)\n name = '{marker}'\n"
f"def gguf_sd_loader(path):\n return {{'w': _T()}}\n")
f"def gguf_sd_loader(path):\n return {{'w': '{marker}'}}\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"].name == "first"
assert gguf_support.load_state_dict("/x.gguf")[0] == {"w": "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"].name == "second"
assert gguf_support.load_state_dict("/x.gguf")[0] == {"w": "second"}
def test_pack_refusing_an_architecture_falls_back(tmp_path, monkeypatch):
@@ -247,214 +240,3 @@ 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.
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" 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"}
def test_a_backend_that_returns_something_else_is_refused(tmp_path):
"""The pack's own guard is what we stepped around, so the check it was doing happens
here: handing a non-state-dict to ComfyUI fails deep in core with a message that makes no
sense out here. Raising instead drops to the slow path that works."""
for bad in (None, [], "sd", {}, {1: object()}):
with pytest.raises(RuntimeError):
gguf_support._assert_state_dict(bad, "m.gguf")
def test_a_non_tensor_value_is_named(tmp_path):
with pytest.raises(RuntimeError, match="not a tensor"):
gguf_support._assert_state_dict({"w": {"nested": 1}}, "m.gguf")
def test_a_real_state_dict_passes():
import torch
gguf_support._assert_state_dict({"w": torch.zeros(2)}, "m.gguf")
def _fake_gguf(monkeypatch, tensors):
"""A stand-in for the `gguf` package over `tensors` — (name, shape, values) triples.
Every tensor is reported as quantized, so `_load_native` takes its dequantize path for
all of them, and each dequantize hands back a fresh float32 array whose lifetime the
caller can track.
"""
import numpy as np
fake = types.ModuleType("gguf")
class _QType:
F32 = "F32"
F16 = "F16"
Q8 = "Q8"
class _Tensor:
def __init__(self, name, shape, values):
self.name = name
# GGUF reports dimensions in the opposite order to torch.
self.shape = tuple(reversed(shape))
self.tensor_type = _QType.Q8
self.data = np.asarray(values, dtype=np.float32)
class _Reader:
def __init__(self, path):
self.tensors = [_Tensor(n, s, v) for n, s, v in tensors]
produced = []
live = []
def _dequantize(data, qtype):
# Sampled BEFORE this call's own array exists, so it counts what the loader is still
# holding on to. Peak is the whole point: every array is freed by the time the load
# returns no matter how it was written.
live.append(sum(1 for ref in produced if ref() is not None))
arr = np.asarray(data, dtype=np.float32).copy()
produced.append(weakref.ref(arr))
return arr
quants = types.ModuleType("gguf.quants")
quants.dequantize = _dequantize
fake.GGMLQuantizationType = _QType
fake.GGUFReader = _Reader
fake.quants = quants
monkeypatch.setitem(sys.modules, "gguf", fake)
monkeypatch.setitem(sys.modules, "gguf.quants", quants)
return produced, live
def test_load_native_returns_half_width_tensors_in_torch_order(monkeypatch):
monkeypatch.setattr(gguf_support.os, "cpu_count", lambda: 2) # forces the serial path
_fake_gguf(monkeypatch, [("a", (2, 3), [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])])
sd, options, note = gguf_support._load_native("/models/m.gguf")
assert options == {}
assert list(sd) == ["a"]
assert tuple(sd["a"].shape) == (2, 3)
assert sd["a"].dtype is torch.float16
assert sd["a"][1][2].item() == 6.0
assert "1 quantized tensors" in note
def test_load_native_does_not_keep_every_expanded_tensor_alive(monkeypatch):
"""The float32 form is four times the finished weight, so holding all of them at once
peaks at several times the model in host RAM — enough to take a machine down on a video
checkpoint. Each one must be cast and released as it is produced."""
monkeypatch.setattr(gguf_support.os, "cpu_count", lambda: 2) # forces the serial path
produced, live = _fake_gguf(
monkeypatch, [(f"t{i}", (4, 4), [[float(i)] * 4] * 4) for i in range(24)])
gguf_support._load_native("/models/m.gguf")
assert len(produced) == 24
# Collecting them all before converting any would climb 0, 1, 2, ... 23.
assert max(live) <= 2, f"{max(live)} expanded tensors were held at once"
gc.collect()
assert [ref for ref in produced if ref() is not None] == []
def test_load_native_rejects_a_tensor_that_contradicts_its_header(monkeypatch):
monkeypatch.setattr(gguf_support.os, "cpu_count", lambda: 2)
_fake_gguf(monkeypatch, [("wrong", (5, 5), [1.0, 2.0, 3.0])])
with pytest.raises(RuntimeError) as excinfo:
gguf_support._load_native("/models/m.gguf")
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)) == ""
tests/test_loaders.py
-70
View File
@@ -231,73 +231,3 @@ def test_sla_is_a_toggle_not_a_rival_backend():
assert sla_attention.SLA_NAME not in loaders.attention_choices()
spec = loaders.FunPackDiffusionModelLoader.INPUT_TYPES()
assert spec["optional"]["sla"][1]["default"] is False
# ── changing an inference setting must not re-read the checkpoint ─────────────
# Attention and SLA live on the loader node, so ComfyUI invalidates the whole node when one
# changes and the file is read again — a minute of dequantizing a GGUF to alter a sparsity
# ratio. Only the file and the dtypes decide the weights.
class _CacheModel:
def __init__(self):
self.model_options = {}
self.clones = 0
def clone(self):
c = _CacheModel()
c.clones = self.clones + 1
return c
def _key_env(tmp_path, monkeypatch):
import loaders
f = tmp_path / "m.safetensors"
f.write_bytes(b"x" * 16)
monkeypatch.setattr(loaders.folder_paths, "get_full_path", lambda folder, n: str(f),
raising=False)
loaders._WEIGHT_CACHE.update({"key": None, "model": None})
return loaders, f
def test_the_key_ignores_attention_and_sla(tmp_path, monkeypatch):
loaders, _ = _key_env(tmp_path, monkeypatch)
a = loaders._weight_cache_key("m.safetensors", "default", "bf16")
b = loaders._weight_cache_key("m.safetensors", "default", "bf16")
assert a == b and a is not None
def test_a_different_dtype_is_a_different_key(tmp_path, monkeypatch):
loaders, _ = _key_env(tmp_path, monkeypatch)
assert (loaders._weight_cache_key("m.safetensors", "default", "bf16")
!= loaders._weight_cache_key("m.safetensors", "fp8_e4m3fn", "bf16"))
def test_a_replaced_file_is_re_read(tmp_path, monkeypatch):
"""Same name, different bytes: size and mtime are in the key so the cache misses."""
loaders, f = _key_env(tmp_path, monkeypatch)
before = loaders._weight_cache_key("m.safetensors", "default", "bf16")
f.write_bytes(b"y" * 32)
assert loaders._weight_cache_key("m.safetensors", "default", "bf16") != before
def test_an_unreadable_file_never_matches(tmp_path, monkeypatch):
"""Failing to notice a changed file is worse than an occasional extra read."""
import loaders
monkeypatch.setattr(loaders.folder_paths, "get_full_path", lambda folder, n: None,
raising=False)
assert loaders._weight_cache_key("gone.safetensors", "default", "bf16") is None
def test_the_cached_model_is_cloned_never_handed_out(tmp_path, monkeypatch):
"""The cached object must not collect the model_options of every run that reused it."""
loaders, _ = _key_env(tmp_path, monkeypatch)
base = _CacheModel()
key = loaders._weight_cache_key("m.safetensors", "default", "bf16")
loaders._WEIGHT_CACHE.update({"key": key, "model": base})
node = loaders.FunPackDiffusionModelLoader()
out, _status = node._finish(loaders._WEIGHT_CACHE["model"].clone(), [], "default",
False, 0.9, "64", 8192, 0, True)
out.model_options.setdefault("transformer_options", {})["marker"] = 1
assert base.model_options == {} # untouched
assert out.clones == 1