From 6bf533e808041326f08a17bdfc07fe3bcb7fec8b Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 23 Aug 2026 05:00:29 +0300 Subject: [PATCH] perf(loaders): stop re-reading the checkpoint for inference settings Attention and SLA are applied to a clone of a one-entry weight cache keyed on file identity + dtypes, so a sparsity change skips the disk. GGUF dequantization now runs across threads instead of serially. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 9 ++++++ gguf_support.py | 44 +++++++++++++++++++-------- loaders.py | 52 ++++++++++++++++++++++++++++++++ tests/test_loaders.py | 70 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 162 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dc151129..b23780a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,15 @@ ## [Unreleased] ### 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. +- **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, diff --git a/gguf_support.py b/gguf_support.py index 30aee2af..8eb13ae4 100644 --- a/gguf_support.py +++ b/gguf_support.py @@ -262,18 +262,35 @@ def _load_native(path: str) -> tuple[dict, dict, str]: import torch reader = gguf.GGUFReader(path) - sd: dict = {} - quantized = 0 - for tensor in reader.tensors: - name = str(tensor.name) + tensors = list(reader.tensors) + # Dequantizing serially is where the minute goes: gguf.quants.dequantize is numpy, and a + # video checkpoint is thousands of tensors. numpy releases the GIL for the array work, so + # threads are real parallelism here. Bounded, not one per core: each worker holds a + # dequantized tensor AND its copy, so an unbounded pool trades load time for a memory + # spike on the largest weights. + workers = max(1, min(8, (os.cpu_count() or 2) - 1)) + quantized = sum( + 1 for t in tensors + if t.tensor_type not in (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16)) + + def _dequant(tensor): 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 + return tensor.data + # gguf.quants.dequantize is the reference implementation for every type the + # package knows — do not reimplement it per quant type. + return gguf.quants.dequantize(tensor.data, qtype) + + if workers > 1 and len(tensors) > 1: + from concurrent.futures import ThreadPoolExecutor + with ThreadPoolExecutor(max_workers=workers) as pool: + arrays = list(pool.map(_dequant, tensors)) + else: + arrays = [_dequant(t) for t in tensors] + + sd: dict = {} + for tensor, arr in zip(tensors, arrays): + name = str(tensor.name) 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)) @@ -290,7 +307,8 @@ def _load_native(path: str) -> tuple[dict, dict, str]: 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.") + 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.") return sd, {}, note diff --git a/loaders.py b/loaders.py index a0027eba..df335bef 100644 --- a/loaders.py +++ b/loaders.py @@ -11,6 +11,7 @@ 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 @@ -128,6 +129,31 @@ 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 same +# lifetime ComfyUI's node-output cache already gave 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.""" @@ -240,6 +266,20 @@ 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. + 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) + path = None misnamed = False if gguf_support.is_gguf(model_name): @@ -280,6 +320,18 @@ 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 + 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 diff --git a/tests/test_loaders.py b/tests/test_loaders.py index 191667c1..e8f3b249 100644 --- a/tests/test_loaders.py +++ b/tests/test_loaders.py @@ -231,3 +231,73 @@ 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