diff --git a/README.md b/README.md index b33fe44..1972def 100644 --- a/README.md +++ b/README.md @@ -35,15 +35,28 @@ The only time a build runs is the first load after install, or when you explicitly refresh (below). ## Refreshing after you add / remove models or LoRAs -The cache holds the old model lists until you refresh, so new files won't appear -until you do one of: +The cache holds the old model lists until you refresh. Three modes are available +from the **`Extensions`** menu (and the command palette): -- **Menu:** `Extensions ▸ 🔄 Refresh Models / LoRAs` (also in the command palette). -- **Graph node:** `🔄 Refresh Models/LoRAs (Tenaciousload)` (for automated workflows). -- **HTTP:** `POST /tenaciousload/refresh`, then `GET /object_info?nocache=1`. +| Mode | What it does | Speed | +|------|--------------|-------| +| ⚡ **Quick refresh** | Re-walks only the folders whose timestamp **changed** since the last scan; reuses the cache for the rest. Catches new / removed / renamed files. | Fast on local disks; **~2× faster** on a slow network mount (it still has to stat every folder to find which changed). | +| 🔄 **Full refresh** | Clears ComfyUI's folder cache and re-walks **everything**. Catches moves/deletes anywhere. | Slowest (the original behaviour). | +| ➕ **Register new file…** | You give it the path(s) of the file(s) you just added; it appends them to the cache with **no folder walk**. | Instant disk-wise — only the `object_info` rebuild remains. | -A refresh re-walks your model folders (slow over a network mount, ~minutes) — the -button shows a "refreshing…" toast meanwhile. Normal loads stay instant. +Also available: +- **Graph node** `🔄 Refresh Models/LoRAs (Tenaciousload)` with a `mode` widget + (`quick` / `full`), for automated workflows. +- **HTTP:** `POST /tenaciousload/refresh` with + `{"mode": "quick" | "full" | "register", "folder": "loras", "files": ["pack/new.safetensors"]}`, + then `GET /object_info?nocache=1`. + +> The **first** Quick refresh after install builds a folder index (one full walk), +> so it's as slow as a Full refresh that one time; every Quick refresh after that +> is incremental. The index is saved to `./cache/scan_snapshot.json`. + +Whichever mode you pick, the button shows a "refreshing…" toast and normal loads +stay instant. ## Requirements **None to install.** Only ComfyUI itself (tested on 0.23.0) and Python ≥ 3.8. diff --git a/__init__.py b/__init__.py index cbf7a88..f6cbb0a 100644 --- a/__init__.py +++ b/__init__.py @@ -4,21 +4,28 @@ ComfyUI-Tenaciousload Self-contained fix for slow / black-screen ComfyUI loading when you have a huge model/LoRA collection (especially on a network mount). -The cause: /api/object_info is tens of MB and takes minutes to build on EVERY -page load (it re-walks the model folders), and the build blocks ComfyUI's event -loop the whole time. This pack installs an in-process caching layer that: +It injects an aiohttp middleware that caches the huge /api/object_info response +in memory and on disk (survives restarts) and serves it gzipped, so the slow +build (which freezes ComfyUI's event loop) runs only on the first load or an +explicit refresh — not on every page load. - * caches the built object_info in memory AND on disk (so it survives restarts), - * serves it gzipped in ~milliseconds instead of rebuilding every load, - * exposes a one-click "Refresh Models / LoRAs" button + graph node to rebuild - the cache after you add/remove models. +Three refresh modes are exposed (menu buttons, a graph node, and HTTP): + * full - clear ComfyUI's folder cache -> full re-walk of every model + folder. Most thorough (catches moves/deletes anywhere). Slowest. + * quick - incremental: re-walk only the folders whose timestamp changed + since the last scan, reuse the cache for the rest. Much faster on + local disks; ~2x on a slow network mount (it still has to stat + every folder to find which changed). + * register - append specific file path(s) to the cache with NO folder walk. + Instant disk-wise; use right after downloading a known file. -No nginx, no docker, no extra port. Just install the pack and restart ComfyUI. +All modes then rebuild the object_info cache so new files show up. """ import os import gzip import json +import time import logging import threading @@ -31,20 +38,22 @@ log = logging.getLogger("Tenaciousload") WEB_DIRECTORY = "./web" -# --- cache storage -------------------------------------------------------------- +# --------------------------------------------------------------------------- # +# object_info response cache (memory + disk) +# --------------------------------------------------------------------------- # _CACHE_DIR = os.path.join(os.path.dirname(__file__), "cache") _RAW_PATH = os.path.join(_CACHE_DIR, "object_info.json") _GZ_PATH = os.path.join(_CACHE_DIR, "object_info.json.gz") +_SNAP_PATH = os.path.join(_CACHE_DIR, "scan_snapshot.json") _OBJECT_INFO_PATHS = ("/object_info", "/api/object_info") _GZIP_LEVEL = 5 _lock = threading.Lock() -_mem = {"raw": None, "gz": None} # bytes / bytes +_mem = {"raw": None, "gz": None} _disk_loaded = False def _load_from_disk(): - """Populate the in-memory cache from disk once (cheap, makes restarts instant).""" global _disk_loaded _disk_loaded = True try: @@ -63,7 +72,6 @@ def _load_from_disk(): def _store(raw_bytes): - """Cache a freshly built object_info body (memory + disk).""" with _lock: gz = gzip.compress(raw_bytes, _GZIP_LEVEL) _mem["raw"] = raw_bytes @@ -80,7 +88,6 @@ def _store(raw_bytes): def invalidate_object_info_cache(): - """Drop the cached object_info so the next request rebuilds it.""" with _lock: _mem["raw"] = None _mem["gz"] = None @@ -94,7 +101,7 @@ def invalidate_object_info_cache(): def clear_comfy_model_cache(): - """Force ComfyUI to re-scan every model/LoRA folder on the next build.""" + """Full reset: force ComfyUI to re-walk every model/LoRA folder next build.""" try: folder_paths.filename_list_cache.clear() except Exception as e: # pragma: no cover @@ -105,11 +112,197 @@ def clear_comfy_model_cache(): log.warning("Tenaciousload: could not clear cache_helper: %s", e) +# --------------------------------------------------------------------------- # +# Incremental folder scanning (quick mode) +# --------------------------------------------------------------------------- # +# Snapshot layout: { folder_name: { root: { dirpath: {"m": mtime, "f": [names], "d": [subdir names]} } } } +_scan_lock = threading.Lock() +_snapshot = None +_SKIP_FOLDERS = {"custom_nodes"} + + +def _load_snapshot(): + global _snapshot + if _snapshot is not None: + return + try: + with open(_SNAP_PATH) as f: + _snapshot = json.load(f) + except Exception: + _snapshot = {} + + +def _save_snapshot(): + try: + os.makedirs(_CACHE_DIR, exist_ok=True) + tmp = _SNAP_PATH + ".tmp" + with open(tmp, "w") as f: + json.dump(_snapshot, f) + os.replace(tmp, _SNAP_PATH) + except Exception as e: # pragma: no cover + log.warning("Tenaciousload: could not save scan snapshot: %s", e) + + +def _scandir_immediate(d): + """One scandir of a single directory -> {m, f, d} or None if inaccessible.""" + files, subdirs = [], [] + try: + with os.scandir(d) as it: + for e in it: + if e.name == ".git": + continue + try: + if e.is_dir(follow_symlinks=True): + subdirs.append(e.name) + elif e.is_file(follow_symlinks=True): + files.append(e.name) + except OSError: + continue + except (FileNotFoundError, NotADirectoryError, PermissionError): + return None + try: + m = os.path.getmtime(d) + except OSError: + m = 0.0 + return {"m": m, "f": files, "d": subdirs} + + +def _scan_root_incremental(root, old): + """Walk a root, scandir-ing only dirs whose mtime changed; reuse the rest.""" + new = {} + scanned = reused = 0 + stack = [root] + while stack: + d = stack.pop() + try: + m = os.path.getmtime(d) + except OSError: + continue # directory disappeared -> drop it + rec = old.get(d) + if rec is not None and rec["m"] == m: + new[d] = rec + reused += 1 + for sub in rec["d"]: + stack.append(os.path.join(d, sub)) + else: + fresh = _scandir_immediate(d) + if fresh is None: + continue + new[d] = fresh + scanned += 1 + for sub in fresh["d"]: + stack.append(os.path.join(d, sub)) + return new, scanned, reused + + +def incremental_scan_folder(folder_name): + """Update folder_paths' cached file list for one folder type, incrementally.""" + folder_name = folder_paths.map_legacy(folder_name) + fnp = folder_paths.folder_names_and_paths.get(folder_name) + if not fnp: + return None + roots, exts = fnp[0], fnp[1] + if _snapshot is None: + _load_snapshot() + folder_snap = _snapshot.setdefault(folder_name, {}) + all_rel, dirs_mtime = set(), {} + scanned = reused = 0 + for root in roots: + if not os.path.isdir(root): + folder_snap.pop(root, None) + continue + new, s, r = _scan_root_incremental(root, folder_snap.get(root, {})) + folder_snap[root] = new + scanned += s + reused += r + for d, rec in new.items(): + dirs_mtime[d] = rec["m"] + for fname in rec["f"]: + all_rel.add(os.path.relpath(os.path.join(d, fname), root)) + filtered = folder_paths.filter_files_extensions(all_rel, exts) + folder_paths.filename_list_cache[folder_name] = (filtered, dirs_mtime, time.perf_counter()) + return {"folder": folder_name, "files": len(filtered), "scanned": scanned, "reused": reused} + + +def quick_rescan_all(): + """Incrementally refresh every model folder type. Returns a per-folder summary.""" + with _scan_lock: + if _snapshot is None: + _load_snapshot() + results = [] + for folder_name in list(folder_paths.folder_names_and_paths.keys()): + if folder_name in _SKIP_FOLDERS: + continue + try: + r = incremental_scan_folder(folder_name) + if r and (r["scanned"] or r["files"]): + results.append(r) + except Exception as e: # pragma: no cover + log.warning("Tenaciousload: quick scan of '%s' failed: %s", folder_name, e) + _save_snapshot() + # also drop ComfyUI's strong request-cache so the new lists are picked up + try: + folder_paths.cache_helper.clear() + except Exception: + pass + return results + + +def register_files(folder_name, rel_paths): + """Append specific files to a folder's cache with no disk walk. Returns counts.""" + folder_name = folder_paths.map_legacy(folder_name) + fnp = folder_paths.folder_names_and_paths.get(folder_name) + if not fnp: + return {"added": 0, "skipped": len(rel_paths), "folder": folder_name} + roots, exts = fnp[0], fnp[1] + cache = folder_paths.filename_list_cache.get(folder_name) + if cache is None: + cache = folder_paths.get_filename_list_(folder_name) + files, dirs = set(cache[0]), dict(cache[1]) + added = skipped = 0 + for rp in rel_paths: + rp = (rp or "").strip().strip("/\\") + if not rp: + continue + placed = False + for root in roots: + full = os.path.join(root, rp) + if os.path.exists(full): + d, rootn = os.path.dirname(full), os.path.normpath(root) + while True: # bump mtimes from the file's dir up to the root + try: + if os.path.isdir(d): + dirs[d] = os.path.getmtime(d) + except OSError: + pass + if os.path.normpath(d) == rootn: + break + parent = os.path.dirname(d) + if parent == d: + break + d = parent + files.add(rp) + placed = True + break + if placed: + added += 1 + else: + skipped += 1 + filtered = folder_paths.filter_files_extensions(files, exts) + folder_paths.filename_list_cache[folder_name] = (filtered, dirs, time.perf_counter()) + try: + folder_paths.cache_helper.clear() + except Exception: + pass + return {"added": added, "skipped": skipped, "folder": folder_name, "files": len(filtered)} + + +# --------------------------------------------------------------------------- # +# object_info caching middleware +# --------------------------------------------------------------------------- # def _serve_cached(request): - raw = _mem["raw"] - gz = _mem["gz"] - accepts_gzip = "gzip" in request.headers.get("Accept-Encoding", "") - if accepts_gzip and gz is not None: + raw, gz = _mem["raw"], _mem["gz"] + if "gzip" in request.headers.get("Accept-Encoding", "") and gz is not None: return web.Response( body=gz, status=200, content_type="application/json", headers={"Content-Encoding": "gzip", "X-Tenaciousload-Cache": "HIT", @@ -125,7 +318,6 @@ def _serve_cached(request): @web.middleware async def _object_info_cache_mw(request, handler): - # Only touch the full object_info list (not /object_info/{node_class}). if request.method != "GET" or request.path not in _OBJECT_INFO_PATHS: return await handler(request) @@ -133,11 +325,9 @@ async def _object_info_cache_mw(request, handler): if not _disk_loaded: _load_from_disk() - nocache = "nocache" in request.query - if not nocache and _mem["raw"] is not None: + if "nocache" not in request.query and _mem["raw"] is not None: return _serve_cached(request) - # MISS (or forced refresh): build via the real handler, then cache the body. resp = await handler(request) try: body = getattr(resp, "body", None) @@ -151,8 +341,7 @@ async def _object_info_cache_mw(request, handler): def _install_middleware(): try: - app = PromptServer.instance.app - app.middlewares.insert(0, _object_info_cache_mw) + PromptServer.instance.app.middlewares.insert(0, _object_info_cache_mw) log.info("Tenaciousload: object_info cache middleware installed") except Exception as e: log.error("Tenaciousload: could not install cache middleware (loads will be slow): %s", e) @@ -161,16 +350,42 @@ def _install_middleware(): _install_middleware() -# --- refresh API ---------------------------------------------------------------- +# --------------------------------------------------------------------------- # +# Refresh API +# --------------------------------------------------------------------------- # @PromptServer.instance.routes.post("/tenaciousload/refresh") async def _refresh(request): + try: + data = await request.json() + except Exception: + data = {} + mode = (data.get("mode") or "full").lower() + + if mode == "quick": + summary = quick_rescan_all() + invalidate_object_info_cache() + rescanned = sum(s["scanned"] for s in summary) + log.info("Tenaciousload: quick refresh — %d folders touched, %d dirs rescanned", len(summary), rescanned) + return web.json_response({"status": "ok", "mode": "quick", "folders": summary}) + + if mode == "register": + folder = data.get("folder") or "loras" + files = data.get("files") or [] + result = register_files(folder, files) + invalidate_object_info_cache() + log.info("Tenaciousload: register — %s", result) + return web.json_response({"status": "ok", "mode": "register", **result}) + + # default: full clear_comfy_model_cache() invalidate_object_info_cache() - log.info("Tenaciousload: caches cleared via refresh endpoint") - return web.json_response({"status": "ok"}) + log.info("Tenaciousload: full refresh — folder cache cleared") + return web.json_response({"status": "ok", "mode": "full"}) -# --- optional graph node (for workflow automation) ------------------------------ +# --------------------------------------------------------------------------- # +# Optional graph node (workflow automation) +# --------------------------------------------------------------------------- # class _AnyType(str): def __ne__(self, other): return False @@ -182,17 +397,23 @@ ANY = _AnyType("*") class TenaciousloadRefresh: @classmethod def INPUT_TYPES(cls): - return {"optional": {"trigger": (ANY, {})}} + return { + "required": {"mode": (["quick", "full"], {"default": "quick"})}, + "optional": {"trigger": (ANY, {})}, + } RETURN_TYPES = (ANY,) RETURN_NAMES = ("trigger",) FUNCTION = "refresh" CATEGORY = "Tenaciousload" OUTPUT_NODE = True - DESCRIPTION = "Clear ComfyUI's model/LoRA cache and invalidate the object_info cache." + DESCRIPTION = "Refresh the model/LoRA cache (quick = changed folders only, full = rescan everything)." - def refresh(self, trigger=None): - clear_comfy_model_cache() + def refresh(self, mode="quick", trigger=None): + if mode == "full": + clear_comfy_model_cache() + else: + quick_rescan_all() invalidate_object_info_cache() return (trigger,) diff --git a/web/tenaciousload.js b/web/tenaciousload.js index 4c98970..ec9814a 100644 --- a/web/tenaciousload.js +++ b/web/tenaciousload.js @@ -13,21 +13,38 @@ function notify(severity, detail, life = 6000) { console[severity === "error" ? "error" : "log"]("[Tenaciousload]", detail); } -async function doRefresh() { +async function runRefresh(mode, extra) { if (busy) { notify("warn", "A refresh is already running…"); return; } busy = true; - notify("info", "Refreshing models/LoRAs — scanning over the network, this can take a few minutes…", 10000); + const label = + mode === "quick" ? "Quick refresh (changed folders)" : + mode === "register" ? "Registering file(s)" : "Full refresh (rescan all)"; + notify("info", `${label} — rebuilding model lists… this can take a moment.`, 10000); try { - // 1) clear ComfyUI's folder cache + drop the object_info cache. - await api.fetchApi("/tenaciousload/refresh", { method: "POST" }); - // 2) force a fresh object_info build and re-cache it (this is the slow step). + // 1) run the chosen refresh mode on the server + const res = await api.fetchApi("/tenaciousload/refresh", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ mode, ...(extra || {}) }), + }); + const data = await res.json().catch(() => ({})); + // 2) rebuild + re-cache object_info (the slow build, if any) await api.fetchApi("/object_info?nocache=1").then((r) => r.arrayBuffer()); - // 3) update open node dropdowns live if the frontend supports it. + // 3) update open dropdowns live if supported try { await app.refreshComboInNodes?.(); } catch (e) { /* non-fatal */ } - notify("success", "Done — new models are now available (reload the page if a dropdown still looks stale).", 8000); + + let detail = "Done — new models are available (reload the page if a dropdown still looks stale)."; + if (mode === "register") { + detail = `Registered ${data.added || 0} file(s)` + + (data.skipped ? `, ${data.skipped} not found on disk` : "") + ". " + detail; + } else if (mode === "quick") { + const dirs = (data.folders || []).reduce((a, f) => a + (f.scanned || 0), 0); + detail = `Quick scan: ${dirs} folder(s) rescanned. ` + detail; + } + notify("success", detail, 8000); } catch (e) { notify("error", "Refresh failed: " + (e?.message || e), 10000); } finally { @@ -35,21 +52,45 @@ async function doRefresh() { } } +async function doRegister() { + const folder = (prompt("Model folder type (loras, checkpoints, vae, …):", "loras") || "").trim(); + if (!folder) return; + const raw = prompt( + `New file path(s) relative to the '${folder}' folder\n(comma- or newline-separated, e.g. mypack/newlora.safetensors):`, + "", + ); + if (!raw) return; + const files = raw.split(/[\n,]+/).map((s) => s.trim()).filter(Boolean); + if (!files.length) return; + await runRefresh("register", { folder, files }); +} + app.registerExtension({ name: "Tenaciousload.Refresh", commands: [ { - id: "Tenaciousload.refresh", - label: "🔄 Refresh Models / LoRAs", + id: "Tenaciousload.quick", + label: "⚡ Quick refresh (changed folders)", + icon: "pi pi-bolt", + function: () => runRefresh("quick"), + }, + { + id: "Tenaciousload.full", + label: "🔄 Full refresh (rescan all models/LoRAs)", icon: "pi pi-refresh", - function: doRefresh, + function: () => runRefresh("full"), + }, + { + id: "Tenaciousload.register", + label: "➕ Register new model file…", + icon: "pi pi-plus", + function: doRegister, }, ], - // Adds the command under the top "Extensions" menu (and the command palette). menuCommands: [ { path: ["Extensions"], - commands: ["Tenaciousload.refresh"], + commands: ["Tenaciousload.quick", "Tenaciousload.full", "Tenaciousload.register"], }, ], });