From 6a259a14fa412f3731da1ab19b94837acf8f3e47 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Mon, 10 Aug 2026 12:31:36 +0800 Subject: [PATCH] feat(nodes): flag missing local models at queue and load time (#1057) --- py/nodes/checkpoint_loader.py | 6 +- py/nodes/create_hook_lora.py | 6 + py/nodes/lora_loader.py | 6 + py/nodes/lora_randomizer.py | 6 + py/nodes/lora_stacker.py | 7 +- py/nodes/unet_loader.py | 5 +- py/nodes/utils.py | 152 +++++++++++ py/nodes/wanvideo_lora_select.py | 7 +- py/routes/handlers/model_handlers.py | 33 +++ .../lorasWidgetAvailability.test.js | 228 ++++++++++++++++ tests/nodes/test_checkpoint_name_filtering.py | 85 ++++++ tests/nodes/test_lora_validation.py | 252 ++++++++++++++++++ tests/routes/test_models_changed_broadcast.py | 40 +++ web/comfyui/lm_styles.css | 19 ++ web/comfyui/loras_widget.js | 112 +++++++- web/comfyui/loras_widget_utils.js | 242 +++++++++++++++++ 16 files changed, 1200 insertions(+), 6 deletions(-) create mode 100644 tests/frontend/components/lorasWidgetAvailability.test.js create mode 100644 tests/nodes/test_checkpoint_name_filtering.py create mode 100644 tests/nodes/test_lora_validation.py create mode 100644 tests/routes/test_models_changed_broadcast.py diff --git a/py/nodes/checkpoint_loader.py b/py/nodes/checkpoint_loader.py index 87dc497a..86ff4836 100644 --- a/py/nodes/checkpoint_loader.py +++ b/py/nodes/checkpoint_loader.py @@ -1,4 +1,5 @@ import logging +import os from typing import Any, List, Tuple import comfy.sd # pyright: ignore[reportMissingImports] import folder_paths # pyright: ignore[reportMissingImports] @@ -58,7 +59,10 @@ class CheckpointLoaderLM: for item in cache.raw_data: if item.get("sub_type") == "checkpoint": file_path = item.get("file_path", "") - if file_path: + # Only offer models that still exist on disk so ComfyUI + # flags missing checkpoints at queue time via + # "value not in list" (the scanner cache can be stale). + if file_path and os.path.exists(file_path): # Format using relative path with OS-native separator formatted_name = _format_model_name_for_comfyui( file_path, model_roots diff --git a/py/nodes/create_hook_lora.py b/py/nodes/create_hook_lora.py index d192cde5..385e6e2d 100644 --- a/py/nodes/create_hook_lora.py +++ b/py/nodes/create_hook_lora.py @@ -15,6 +15,7 @@ from .utils import ( any_type, apply_lora_syntax_format, get_loras_list, + validate_lora_entries, ) logger = logging.getLogger(__name__) @@ -42,6 +43,11 @@ class CreateHookLoraLM: "optional": FlexibleOptionalInputType(any_type), } + @classmethod + def VALIDATE_INPUTS(cls, loras=None): + """Queue-time validation: reject missing local LoRAs before execution.""" + return validate_lora_entries({"loras": loras}) or True + RETURN_TYPES = ("HOOKS", "STRING", "STRING") RETURN_NAMES = ("HOOKS", "trigger_words", "active_loras") FUNCTION = "create_hook" diff --git a/py/nodes/lora_loader.py b/py/nodes/lora_loader.py index 74da5804..7683293a 100644 --- a/py/nodes/lora_loader.py +++ b/py/nodes/lora_loader.py @@ -14,6 +14,7 @@ from .utils import ( get_loras_list, nunchaku_load_lora, parse_lora_syntax, + validate_lora_entries, ) logger = logging.getLogger(__name__) @@ -142,6 +143,11 @@ class LoraLoaderLM: "optional": FlexibleOptionalInputType(any_type), } + @classmethod + def VALIDATE_INPUTS(cls, loras=None): + """Queue-time validation: reject missing local LoRAs before execution.""" + return validate_lora_entries({"loras": loras}) or True + RETURN_TYPES = ("MODEL", "CLIP", "STRING", "STRING") RETURN_NAMES = ("MODEL", "CLIP", "trigger_words", "loaded_loras") FUNCTION = "load_loras" diff --git a/py/nodes/lora_randomizer.py b/py/nodes/lora_randomizer.py index 7280032a..23b410af 100644 --- a/py/nodes/lora_randomizer.py +++ b/py/nodes/lora_randomizer.py @@ -9,6 +9,7 @@ and tracks the last used combination for reuse. import logging import os from ..utils.utils import get_lora_info +from .utils import validate_lora_entries logger = logging.getLogger(__name__) @@ -31,6 +32,11 @@ class LoraRandomizerLM: }, } + @classmethod + def VALIDATE_INPUTS(cls, loras=None): + """Queue-time validation: reject missing local LoRAs before execution.""" + return validate_lora_entries({"loras": loras}) or True + RETURN_TYPES = ("LORA_STACK",) RETURN_NAMES = ("LORA_STACK",) diff --git a/py/nodes/lora_stacker.py b/py/nodes/lora_stacker.py index 2b3549a1..aa014d98 100644 --- a/py/nodes/lora_stacker.py +++ b/py/nodes/lora_stacker.py @@ -1,6 +1,6 @@ import os from ..utils.utils import get_lora_info -from .utils import FlexibleOptionalInputType, any_type, apply_lora_syntax_format, extract_lora_name, get_loras_list +from .utils import FlexibleOptionalInputType, any_type, apply_lora_syntax_format, extract_lora_name, get_loras_list, validate_lora_entries import logging @@ -22,6 +22,11 @@ class LoraStackerLM: "optional": FlexibleOptionalInputType(any_type), } + @classmethod + def VALIDATE_INPUTS(cls, loras=None): + """Queue-time validation: reject missing local LoRAs before execution.""" + return validate_lora_entries({"loras": loras}) or True + RETURN_TYPES = ("LORA_STACK", "STRING", "STRING") RETURN_NAMES = ("LORA_STACK", "trigger_words", "active_loras") FUNCTION = "stack_loras" diff --git a/py/nodes/unet_loader.py b/py/nodes/unet_loader.py index 0ab324b5..355403fa 100644 --- a/py/nodes/unet_loader.py +++ b/py/nodes/unet_loader.py @@ -74,7 +74,10 @@ class UNETLoaderLM: for item in cache.raw_data: if item.get("sub_type") == "diffusion_model": file_path = item.get("file_path", "") - if file_path: + # Only offer models that still exist on disk so ComfyUI + # flags missing diffusion models at queue time via + # "value not in list" (the scanner cache can be stale). + if file_path and os.path.exists(file_path): # Format using relative path with OS-native separator formatted_name = _format_model_name_for_comfyui( file_path, model_roots diff --git a/py/nodes/utils.py b/py/nodes/utils.py index 44a33bc7..8d6b1c2c 100644 --- a/py/nodes/utils.py +++ b/py/nodes/utils.py @@ -44,6 +44,7 @@ import re import logging import copy import sys +import asyncio import folder_paths # pyright: ignore[reportMissingImports] logger = logging.getLogger(__name__) @@ -111,6 +112,157 @@ def get_loras_list(kwargs): return [] +_LORA_EXTENSIONS = (".safetensors", ".ckpt", ".pt", ".bin") + + +def _strip_lora_extension(name: str) -> str: + """Strip a known LoRA model extension from a name (case-insensitive).""" + lowered = name.lower() + for ext in _LORA_EXTENSIONS: + if lowered.endswith(ext): + return name[: -len(ext)] + return name + + +def _find_missing_loras(names: list[str]) -> list[str]: + """Return the names that cannot be resolved to an existing local LoRA file. + + Mirrors the matching semantics of ``get_lora_info_absolute`` + (py/utils/utils.py): after stripping the extension, a name matches a cached + LoRA when it equals the cached file name or the ``folder/file`` path. As a + fallback, a name containing a folder that only matches by basename resolves + to the first basename match (same behavior as the runtime resolver). Raw + absolute paths that exist on disk are always considered available. + + The scanner cache is fetched once for all names; the cache may be stale, so + resolved paths are additionally verified with ``os.path.isfile``. + """ + if not names: + return [] + + async def _check() -> list[str]: + from ..services.service_registry import ServiceRegistry + + scanner = await ServiceRegistry.get_lora_scanner() + # The scanner cache may not be hydrated yet (startup, library path + # change). An empty cache is not authoritative — treat it as "cannot + # verify" and skip validation instead of flagging every active LoRA + # as missing. + if getattr(scanner, "_cache", None) is None or getattr( + scanner, "_is_initializing", False + ): + return [] + cache = await scanner.get_cached_data() + + lookup = {} + basename_candidates = {} + for item in cache.raw_data: + file_path = item.get("file_path") + if not file_path: + continue + file_name = item.get("file_name", "") + folder = item.get("folder", "") + file_name_no_ext = _strip_lora_extension(file_name) + path_name_no_ext = ( + f"{folder}/{file_name_no_ext}".replace("\\", "/") + if folder + else file_name_no_ext + ) + lookup.setdefault(file_name_no_ext, file_path) + lookup.setdefault(path_name_no_ext, file_path) + basename_candidates.setdefault(file_name_no_ext, []).append( + (folder, file_path) + ) + + missing = [] + for name in names: + if not name: + continue + normalized = name.replace("\\", "/") + # Raw absolute paths (outside the library) are usable as-is. + if os.path.isfile(normalized): + continue + no_ext = _strip_lora_extension(normalized) + file_path = lookup.get(no_ext) + if file_path is None and "/" in no_ext: + # A name with a folder that matches only by basename resolves + # at runtime like get_lora_info_absolute's fallback does: + # prefer a candidate whose folder prefixes the name, else the + # first basename match. + folder, basename = no_ext.rsplit("/", 1) + candidates = basename_candidates.get(basename, []) + file_path = next( + ( + fp + for fld, fp in candidates + if fld and no_ext.startswith(fld + "/") + ), + None, + ) + if file_path is None and candidates: + file_path = candidates[0][1] + if file_path is None or not os.path.isfile(file_path): + missing.append(name) + return missing + + try: + # Check if we're already in an event loop + loop = asyncio.get_running_loop() + # If we're in a running loop, run the async check in a separate thread + import concurrent.futures + + def run_in_thread(): + new_loop = asyncio.new_event_loop() + asyncio.set_event_loop(new_loop) + try: + return new_loop.run_until_complete(_check()) + finally: + new_loop.close() + + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit(run_in_thread) + return future.result() + except RuntimeError: + # No event loop is running, we can use asyncio.run() + return asyncio.run(_check()) + + +def validate_lora_entries(kwargs): + """Validate active LoRA widget entries against the local library. + + Used by node ``VALIDATE_INPUTS`` implementations so ComfyUI rejects the + prompt at queue time (``custom_validation_failed``) when an active entry + references a LoRA that is not available locally — mirroring how built-in + loader nodes flag missing models before execution starts. + + Returns: + None when every active entry resolves to an existing local file, + otherwise a descriptive error string listing the missing LoRAs. + Verification failures (e.g. scanner not ready) are treated as valid + so queueing is never blocked by validation machinery itself. + """ + # Missing/empty loras input is always valid; skip get_loras_list so it + # does not log a warning for the None case on every queue. + if not kwargs.get("loras"): + return None + loras = get_loras_list(kwargs) + active_names = [] + for lora in loras: + if not isinstance(lora, dict): + continue + if not lora.get("active", False): + continue + active_names.append(apply_lora_syntax_format(str(lora.get("name") or ""))) + try: + missing = _find_missing_loras(active_names) + except Exception: + logger.exception("Failed to validate LoRA entries against the local library") + return None + if not missing: + return None + return "Missing LoRA(s) in local library: " + ", ".join(missing) + + def load_state_dict_in_safetensors(path, device="cpu", filter_prefix=""): """Simplified version of load_state_dict_in_safetensors that just loads from a local path""" import safetensors.torch diff --git a/py/nodes/wanvideo_lora_select.py b/py/nodes/wanvideo_lora_select.py index 8c49eb22..576db0ce 100644 --- a/py/nodes/wanvideo_lora_select.py +++ b/py/nodes/wanvideo_lora_select.py @@ -1,7 +1,7 @@ import os from ..utils.utils import get_lora_info_absolute from ..config import config -from .utils import FlexibleOptionalInputType, any_type, get_loras_list +from .utils import FlexibleOptionalInputType, any_type, get_loras_list, validate_lora_entries import logging logger = logging.getLogger(__name__) @@ -35,6 +35,11 @@ class WanVideoLoraSelectLM: "optional": FlexibleOptionalInputType(any_type), } + @classmethod + def VALIDATE_INPUTS(cls, loras=None): + """Queue-time validation: reject missing local LoRAs before execution.""" + return validate_lora_entries({"loras": loras}) or True + RETURN_TYPES = ("WANVIDLORA", "STRING", "STRING") RETURN_NAMES = ("lora", "trigger_words", "active_loras") FUNCTION = "process_loras" diff --git a/py/routes/handlers/model_handlers.py b/py/routes/handlers/model_handlers.py index 385ba993..2ba5998b 100644 --- a/py/routes/handlers/model_handlers.py +++ b/py/routes/handlers/model_handlers.py @@ -51,6 +51,29 @@ LICENSE_FIELDS = ( ) +_broadcast_models_changed_tasks: set = set() + + +def _broadcast_models_changed() -> None: + """Notify connected clients that the local model library changed. + + The ComfyUI graph page listens for this event to invalidate its cached + model availability data (loras widget missing-model cues / error flags) + without waiting for the cache TTL to expire. + """ + try: + from ...services.websocket_manager import ws_manager + + task = asyncio.create_task(ws_manager.broadcast({"type": "models_changed"})) + # Keep a reference so the task is not garbage-collected mid-await. + _broadcast_models_changed_tasks.add(task) + task.add_done_callback(_broadcast_models_changed_tasks.discard) + except Exception: + logging.getLogger(__name__).debug( + "Failed to broadcast models_changed", exc_info=True + ) + + class ModelPageView: """Render the HTML view for model listings.""" @@ -460,6 +483,7 @@ class ModelManagementHandler: return web.Response(text="Model path is required", status=400) result = await self._lifecycle_service.delete_model(file_path) + _broadcast_models_changed() return web.json_response(result) except ValueError as exc: return web.json_response({"success": False, "error": str(exc)}, status=400) @@ -931,6 +955,8 @@ class ModelManagementHandler: file_path=file_path, new_file_name=new_file_name ) + _broadcast_models_changed() + return web.json_response( { **result, @@ -959,6 +985,7 @@ class ModelManagementHandler: ) result = await self._lifecycle_service.bulk_delete_models(file_paths) + _broadcast_models_changed() return web.json_response(result) except ValueError as exc: return web.json_response({"success": False, "error": str(exc)}, status=400) @@ -1061,6 +1088,7 @@ class ModelQueryHandler: await self._service.scan_models( force_refresh=True, rebuild_cache=full_rebuild ) + _broadcast_models_changed() if self._service.scanner.is_cancelled(): return web.json_response( { @@ -2235,6 +2263,8 @@ class ModelMoveHandler: result = await self._move_service.move_model( file_path, target_path, use_default_paths=use_default_paths ) + if result.get("success"): + _broadcast_models_changed() status = 200 if result.get("success") else 500 return web.json_response(result, status=status) except Exception as exc: @@ -2254,6 +2284,8 @@ class ModelMoveHandler: result = await self._move_service.move_models_bulk( file_paths, target_path, use_default_paths=use_default_paths ) + if result.get("success"): + _broadcast_models_changed() return web.json_response(result) except Exception as exc: self._logger.error("Error moving models in bulk: %s", exc, exc_info=True) @@ -2299,6 +2331,7 @@ class ModelAutoOrganizeHandler: progress_callback=self._progress_callback, exclusion_patterns=exclusion_patterns, ) + _broadcast_models_changed() return web.json_response(result.to_dict()) except AutoOrganizeInProgressError: return web.json_response( diff --git a/tests/frontend/components/lorasWidgetAvailability.test.js b/tests/frontend/components/lorasWidgetAvailability.test.js new file mode 100644 index 00000000..d4acff37 --- /dev/null +++ b/tests/frontend/components/lorasWidgetAvailability.test.js @@ -0,0 +1,228 @@ +import { describe, it, expect, beforeEach, vi } from 'vitest'; + +const { + APP_MODULE, + API_MODULE, + UTILS_MODULE, +} = vi.hoisted(() => ({ + APP_MODULE: new URL('../../../scripts/app.js', import.meta.url).pathname, + API_MODULE: new URL('../../../scripts/api.js', import.meta.url).pathname, + UTILS_MODULE: new URL('../../../web/comfyui/loras_widget_utils.js', import.meta.url).pathname, +})); + +vi.mock(APP_MODULE, () => ({ + app: { graph: {} }, +})); + +const { fetchApiMock } = vi.hoisted(() => ({ fetchApiMock: vi.fn() })); +vi.mock(API_MODULE, () => ({ + api: { fetchApi: fetchApiMock }, +})); + +import { + normalizeLoraNameKey, + buildAvailableLoraSet, + isLoraNameAvailable, + getAvailableLoras, + getAvailableLorasSync, + resetAvailableLorasCache, + onLibraryChanged, + handleLibraryChangeMessage, +} from '../../../web/comfyui/loras_widget_utils.js'; + +describe('normalizeLoraNameKey', () => { + it('normalizes backslashes to forward slashes', () => { + expect(normalizeLoraNameKey('sub\\folder\\lora.safetensors')).toBe( + 'sub/folder/lora' + ); + }); + + it('strips known model extensions case-insensitively', () => { + expect(normalizeLoraNameKey('lora.safetensors')).toBe('lora'); + expect(normalizeLoraNameKey('lora.CKPT')).toBe('lora'); + expect(normalizeLoraNameKey('lora.pt')).toBe('lora'); + expect(normalizeLoraNameKey('lora.bin')).toBe('lora'); + }); + + it('keeps extensions the backend does not strip', () => { + // Matches backend _strip_lora_extension: .gguf is not a LoRA extension. + expect(normalizeLoraNameKey('lora.gguf')).toBe('lora.gguf'); + }); + + it('leaves extension-less names unchanged', () => { + expect(normalizeLoraNameKey('lora')).toBe('lora'); + }); +}); + +describe('buildAvailableLoraSet', () => { + it('registers both path and basename forms without extension', () => { + const set = buildAvailableLoraSet(['sub/a.safetensors', 'b.ckpt']); + expect(set.has('sub/a')).toBe(true); + expect(set.has('a')).toBe(true); + expect(set.has('b')).toBe(true); + }); + + it('ignores empty entries', () => { + const set = buildAvailableLoraSet([null, '', 'sub/c.safetensors']); + expect(set.has('sub/c')).toBe(true); + expect(set.has('')).toBe(false); + }); +}); + +describe('isLoraNameAvailable', () => { + const set = buildAvailableLoraSet(['sub/a.safetensors', 'b.ckpt']); + + it('treats everything as available while the set is not loaded', () => { + expect(isLoraNameAvailable('anything.safetensors', null)).toBe(true); + }); + + it('matches by basename with or without extension', () => { + expect(isLoraNameAvailable('a', set)).toBe(true); + expect(isLoraNameAvailable('a.safetensors', set)).toBe(true); + expect(isLoraNameAvailable('b.ckpt', set)).toBe(true); + }); + + it('matches by full folder path', () => { + expect(isLoraNameAvailable('sub/a.safetensors', set)).toBe(true); + }); + + it('reports names not in the library', () => { + expect(isLoraNameAvailable('missing.safetensors', set)).toBe(false); + }); + + it('falls back to the basename for folder-qualified names', () => { + // Mirror of the backend basename fallback: a folder prefix that does not + // match a stored path still resolves when the basename exists. + expect(isLoraNameAvailable('sub/b.ckpt', set)).toBe(true); + expect(isLoraNameAvailable('any/folder/a.safetensors', set)).toBe(true); + expect(isLoraNameAvailable('sub/missing.safetensors', set)).toBe(false); + }); + + it('treats absolute paths as available without verification', () => { + expect(isLoraNameAvailable('/abs/path/x.safetensors', set)).toBe(true); + expect(isLoraNameAvailable('C:/abs/path/x.safetensors', set)).toBe(true); + }); +}); + +describe('getAvailableLoras caching', () => { + beforeEach(() => { + fetchApiMock.mockReset(); + resetAvailableLorasCache(); + }); + + it('fetches the cycler list and builds the availability set', async () => { + fetchApiMock.mockResolvedValue({ + ok: true, + json: async () => ({ + success: true, + loras: [{ file_name: 'sub/a.safetensors' }, { file_name: 'b.ckpt' }], + }), + }); + + const set = await getAvailableLoras(); + expect(fetchApiMock).toHaveBeenCalledWith('/lm/loras/cycler-list', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: '{}', + }); + expect(set.has('sub/a')).toBe(true); + expect(set.has('b')).toBe(true); + expect(getAvailableLorasSync().has('sub/a')).toBe(true); + }); + + it('shares a single in-flight request between concurrent callers', async () => { + fetchApiMock.mockResolvedValue({ + ok: true, + json: async () => ({ success: true, loras: [] }), + }); + + await Promise.all([getAvailableLoras(), getAvailableLoras()]); + expect(fetchApiMock).toHaveBeenCalledTimes(1); + }); + + it('does not refetch while the cache is fresh', async () => { + fetchApiMock.mockResolvedValue({ + ok: true, + json: async () => ({ success: true, loras: [] }), + }); + + await getAvailableLoras(); + await getAvailableLoras(); + expect(fetchApiMock).toHaveBeenCalledTimes(1); + }); + + it('resolves to null and does not cache on fetch failure', async () => { + fetchApiMock.mockRejectedValue(new Error('network down')); + + const result = await getAvailableLoras(); + expect(result).toBeNull(); + expect(getAvailableLorasSync()).toBeNull(); + }); + + it('resolves to null on non-ok response', async () => { + fetchApiMock.mockResolvedValue({ ok: false }); + + const result = await getAvailableLoras(); + expect(result).toBeNull(); + }); +}); + +describe('library change invalidation', () => { + beforeEach(() => { + fetchApiMock.mockReset(); + resetAvailableLorasCache(); + }); + + it('invalidates the availability cache and notifies listeners on models_changed', async () => { + fetchApiMock.mockResolvedValue({ + ok: true, + json: async () => ({ + success: true, + loras: [{ file_name: 'old.safetensors' }], + }), + }); + + const listener = vi.fn(); + const unsubscribe = onLibraryChanged(listener); + try { + await getAvailableLoras(); + expect(getAvailableLorasSync().has('old')).toBe(true); + + // Simulate a deletion in the Lora Manager UI: the cache is dropped and + // listeners are notified so widgets re-render with fresh data. + handleLibraryChangeMessage({ type: 'models_changed' }); + expect(listener).toHaveBeenCalledTimes(1); + expect(getAvailableLorasSync()).toBeNull(); + + fetchApiMock.mockResolvedValue({ + ok: true, + json: async () => ({ success: true, loras: [] }), + }); + await getAvailableLoras(); + expect(getAvailableLorasSync().has('old')).toBe(false); + } finally { + unsubscribe(); + } + }); + + it('ignores unrelated messages', () => { + const listener = vi.fn(); + const unsubscribe = onLibraryChanged(listener); + try { + handleLibraryChangeMessage({ type: 'download_progress' }); + handleLibraryChangeMessage({ type: 'init_progress' }); + handleLibraryChangeMessage(null); + expect(listener).not.toHaveBeenCalled(); + } finally { + unsubscribe(); + } + }); + + it('unsubscribes listeners', () => { + const listener = vi.fn(); + const unsubscribe = onLibraryChanged(listener); + unsubscribe(); + handleLibraryChangeMessage({ type: 'models_changed' }); + expect(listener).not.toHaveBeenCalled(); + }); +}); diff --git a/tests/nodes/test_checkpoint_name_filtering.py b/tests/nodes/test_checkpoint_name_filtering.py new file mode 100644 index 00000000..e17a5b04 --- /dev/null +++ b/tests/nodes/test_checkpoint_name_filtering.py @@ -0,0 +1,85 @@ +"""Tests for the checkpoint/unet combo-name existence filtering. + +Deleted files must drop out of the combo list so ComfyUI flags the node at +queue time via "value not in list" instead of failing at execution time. +""" + +import pytest + +from py.nodes.checkpoint_loader import CheckpointLoaderLM +from py.nodes.unet_loader import UNETLoaderLM + + +class _FakeCache: + def __init__(self, raw_data): + self.raw_data = raw_data + + +class _FakeScanner: + def __init__(self, raw_data, model_roots): + self._raw_data = raw_data + self._model_roots = model_roots + + async def get_cached_data(self, force_refresh=False): + return _FakeCache(self._raw_data) + + def get_model_roots(self): + return self._model_roots + + +@pytest.fixture +def checkpoint_library(tmp_path, monkeypatch): + from py.services.service_registry import ServiceRegistry + + existing = tmp_path / "keep.safetensors" + existing.write_bytes(b"x") + deleted = tmp_path / "deleted.safetensors" # referenced but never created + + raw_data = [ + {"sub_type": "checkpoint", "file_path": str(existing)}, + {"sub_type": "checkpoint", "file_path": str(deleted)}, + # Wrong type must stay excluded by the sub_type filter. + {"sub_type": "diffusion_model", "file_path": str(existing)}, + ] + async def _fake_scanner(): + return _FakeScanner(raw_data, [str(tmp_path)]) + + monkeypatch.setattr( + ServiceRegistry, "get_checkpoint_scanner", _fake_scanner + ) + return tmp_path + + +def test_checkpoint_names_drop_deleted_files(checkpoint_library): + assert CheckpointLoaderLM._get_checkpoint_names() == ["keep.safetensors"] + + +def test_unet_names_drop_deleted_files(tmp_path, monkeypatch): + from py.services.service_registry import ServiceRegistry + + existing = tmp_path / "keep.safetensors" + existing.write_bytes(b"x") + deleted = tmp_path / "deleted.safetensors" + + raw_data = [ + {"sub_type": "diffusion_model", "file_path": str(existing)}, + {"sub_type": "diffusion_model", "file_path": str(deleted)}, + {"sub_type": "checkpoint", "file_path": str(existing)}, + ] + async def _fake_scanner(): + return _FakeScanner(raw_data, [str(tmp_path)]) + + monkeypatch.setattr( + ServiceRegistry, "get_checkpoint_scanner", _fake_scanner + ) + assert UNETLoaderLM._get_unet_names() == ["keep.safetensors"] + + +def test_checkpoint_names_empty_when_scanner_fails(tmp_path, monkeypatch): + from py.services.service_registry import ServiceRegistry + + def _boom(): + raise RuntimeError("scanner not available") + + monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _boom) + assert CheckpointLoaderLM._get_checkpoint_names() == [] diff --git a/tests/nodes/test_lora_validation.py b/tests/nodes/test_lora_validation.py new file mode 100644 index 00000000..d113da39 --- /dev/null +++ b/tests/nodes/test_lora_validation.py @@ -0,0 +1,252 @@ +"""Tests for queue-time LoRA validation helpers (validate_lora_entries).""" + +import pytest + +from py.nodes.utils import _find_missing_loras, validate_lora_entries + + +class _FakeCache: + def __init__(self, raw_data): + self.raw_data = raw_data + + +class _FakeScanner: + def __init__(self, raw_data): + self._raw_data = raw_data + # Non-None cache marks the scanner as initialized; None means the + # real scanner has not hydrated yet (validation must stay lenient). + self._cache = object() + self._is_initializing = False + + async def get_cached_data(self, force_refresh=False): + return _FakeCache(self._raw_data) + + +@pytest.fixture +def lora_library(tmp_path): + """Create a fake on-disk LoRA library plus a matching scanner cache.""" + existing_a = tmp_path / "a.safetensors" + existing_b = tmp_path / "sub" / "b.safetensors" + existing_a.write_bytes(b"a") + existing_b.parent.mkdir() + existing_b.write_bytes(b"b") + + # "gone" is referenced by the cache but the file was deleted afterwards, + # simulating a stale scanner cache. + gone = tmp_path / "gone.safetensors" + + raw_data = [ + {"file_name": "a.safetensors", "folder": "", "file_path": str(existing_a)}, + { + "file_name": "b.safetensors", + "folder": "sub", + "file_path": str(existing_b), + }, + {"file_name": "gone.safetensors", "folder": "", "file_path": str(gone)}, + ] + return raw_data + + +@pytest.fixture +def mock_lora_scanner(lora_library, monkeypatch): + from py.services.service_registry import ServiceRegistry + + async def _fake_scanner(): + return _FakeScanner(lora_library) + + monkeypatch.setattr( + ServiceRegistry, "get_lora_scanner", _fake_scanner + ) + + +def _entries(*names): + return [ + {"active": True, "name": name, "strength": 1.0, "clipStrength": 1.0} + for name in names + ] + + +def test_validate_flat_name_with_extension(mock_lora_scanner): + assert validate_lora_entries({"loras": _entries("a.safetensors")}) is None + + +def test_validate_flat_name_without_extension(mock_lora_scanner): + assert validate_lora_entries({"loras": _entries("a")}) is None + + +def test_validate_subfolder_name(mock_lora_scanner): + assert validate_lora_entries({"loras": _entries("sub/b.safetensors")}) is None + assert validate_lora_entries({"loras": _entries("sub/b")}) is None + + +def test_validate_stale_cache_entry_reported(mock_lora_scanner): + result = validate_lora_entries({"loras": _entries("gone.safetensors")}) + assert result is not None + assert "gone.safetensors" in result + + +def test_validate_unknown_name_reported(mock_lora_scanner): + result = validate_lora_entries({"loras": _entries("nope.safetensors")}) + assert result is not None + assert "nope.safetensors" in result + + +def test_validate_multiple_missing_all_listed(mock_lora_scanner): + result = validate_lora_entries( + {"loras": _entries("gone.safetensors", "nope.safetensors")} + ) + assert result is not None + assert "gone.safetensors" in result + assert "nope.safetensors" in result + + +def test_validate_inactive_entries_ignored(mock_lora_scanner): + kwargs = {"loras": [{"active": False, "name": "gone.safetensors", "strength": 1.0}]} + assert validate_lora_entries(kwargs) is None + + +def test_validate_value_wrapper_format(mock_lora_scanner): + kwargs = {"loras": {"__value__": _entries("a.safetensors")}} + assert validate_lora_entries(kwargs) is None + kwargs_missing = {"loras": {"__value__": _entries("gone.safetensors")}} + assert validate_lora_entries(kwargs_missing) is not None + + +def test_validate_legacy_basename_fallback(mock_lora_scanner): + # A name with a folder that only matches by basename resolves at runtime + # via get_lora_info_absolute's fallback, so it must not be flagged. + assert validate_lora_entries({"loras": _entries("other/b.safetensors")}) is None + + +def test_validate_existing_absolute_path_ok(mock_lora_scanner, lora_library): + existing_a = lora_library[0]["file_path"] + assert validate_lora_entries({"loras": _entries(existing_a)}) is None + + +def test_validate_missing_absolute_path_reported(mock_lora_scanner): + result = validate_lora_entries( + {"loras": _entries("/nonexistent/path/x.safetensors")} + ) + # Legacy syntax format normalizes to the basename, matching execution. + assert result is not None + assert "x.safetensors" in result + + +def test_validate_empty_or_missing_loras_is_valid(mock_lora_scanner): + assert validate_lora_entries({}) is None + assert validate_lora_entries({"loras": []}) is None + + +def test_validate_scanner_failure_is_lenient(monkeypatch): + from py.services.service_registry import ServiceRegistry + + def _boom(): + raise RuntimeError("scanner not initialized") + + monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _boom) + assert validate_lora_entries({"loras": _entries("a.safetensors")}) is None + + +def test_find_missing_loras_empty(mock_lora_scanner): + assert _find_missing_loras([]) == [] + + +def test_validate_inputs_only_reports_loras_input(mock_lora_scanner): + """VALIDATE_INPUTS must fail for the single 'loras' input only — ComfyUI + creates one custom_validation_failed error per declared input, so an + explicit loras parameter keeps the report to a single input error.""" + from py.nodes.lora_loader import LoraLoaderLM + from py.nodes.lora_stacker import LoraStackerLM + from py.nodes.create_hook_lora import CreateHookLoraLM + from py.nodes.wanvideo_lora_select import WanVideoLoraSelectLM + from py.nodes.lora_randomizer import LoraRandomizerLM + + import inspect + + for cls in ( + LoraLoaderLM, + LoraStackerLM, + CreateHookLoraLM, + WanVideoLoraSelectLM, + LoraRandomizerLM, + ): + spec = inspect.getfullargspec(getattr(cls, "VALIDATE_INPUTS")) + assert spec.varkw is None, f"{cls.__name__} VALIDATE_INPUTS must not accept **kwargs" + assert "loras" in spec.args, f"{cls.__name__} VALIDATE_INPUTS must declare loras" + + result = LoraLoaderLM.VALIDATE_INPUTS( + loras=[ + {"active": True, "name": "gone.safetensors", "strength": 1.0}, + ] + ) + assert result is not None + assert "gone.safetensors" in result + + assert LoraLoaderLM.VALIDATE_INPUTS(loras=None) is True + assert LoraLoaderLM.VALIDATE_INPUTS() is True + assert LoraLoaderLM.VALIDATE_INPUTS(loras=[]) is True + + +def test_validate_lenient_when_scanner_not_initialized(monkeypatch): + """An unhydrated scanner (empty cache) must not reject every LoRA.""" + from py.services.service_registry import ServiceRegistry + + class _UninitializedScanner(_FakeScanner): + def __init__(self, raw_data): + super().__init__(raw_data) + self._cache = None + + async def _fake_scanner(): + return _UninitializedScanner([]) + + monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _fake_scanner) + assert validate_lora_entries({"loras": _entries("a.safetensors")}) is None + + +def test_validate_lenient_while_scanner_initializing(monkeypatch): + from py.services.service_registry import ServiceRegistry + + class _InitializingScanner(_FakeScanner): + def __init__(self, raw_data): + super().__init__(raw_data) + self._is_initializing = True + + async def _fake_scanner(): + return _InitializingScanner([]) + + monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _fake_scanner) + assert validate_lora_entries({"loras": _entries("a.safetensors")}) is None + + +def test_find_missing_basename_fallback_prefers_folder_prefix( + tmp_path, monkeypatch +): + """Folder-prefix candidates win over the first basename match, mirroring + get_lora_info_absolute's fallback ordering.""" + from py.services.service_registry import ServiceRegistry + + root_a = tmp_path / "root_a" + root_b = tmp_path / "root_b" + root_a.mkdir() + root_b.mkdir() + live_a = root_a / "x.safetensors" + live_a.write_bytes(b"x") + stale_b = root_b / "x.safetensors" # same basename, file deleted + + raw_data = [ + {"file_name": "x.safetensors", "folder": "root_a", "file_path": str(live_a)}, + {"file_name": "x.safetensors", "folder": "root_b", "file_path": str(stale_b)}, + ] + + async def _fake_scanner(): + return _FakeScanner(raw_data) + + monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _fake_scanner) + # Exact path matches are authoritative. + assert _find_missing_loras(["root_b/x"]) == ["root_b/x"] + # Folder-prefix fallback beats the first basename match: root_a is live. + assert _find_missing_loras(["root_a/deep/x"]) == [] + # Prefix match wins over the first candidate: root_b is stale. + assert _find_missing_loras(["root_b/deep/x"]) == ["root_b/deep/x"] + # No prefix match: falls back to the first basename candidate (root_a). + assert _find_missing_loras(["other/x"]) == [] diff --git a/tests/routes/test_models_changed_broadcast.py b/tests/routes/test_models_changed_broadcast.py new file mode 100644 index 00000000..fa242f61 --- /dev/null +++ b/tests/routes/test_models_changed_broadcast.py @@ -0,0 +1,40 @@ +"""Tests for the models_changed WebSocket broadcast helper.""" + +import asyncio + +import pytest + +from py.routes.handlers import model_handlers + + +@pytest.mark.asyncio +async def test_broadcast_models_changed_payload(monkeypatch): + """The helper must broadcast the models_changed event to all clients.""" + sent = [] + + async def fake_broadcast(data): + sent.append(data) + + monkeypatch.setattr( + "py.services.websocket_manager.ws_manager.broadcast", fake_broadcast + ) + # Give the create_task a chance to run. + for _ in range(10): + await asyncio.sleep(0) + + model_handlers._broadcast_models_changed() + await asyncio.sleep(0) + assert sent == [{"type": "models_changed"}] + + +@pytest.mark.asyncio +async def test_broadcast_models_changed_survives_broadcast_failure(monkeypatch): + """A failing broadcast must not raise out of the mutation handler.""" + async def fake_broadcast(data): + raise RuntimeError("socket closed") + + monkeypatch.setattr( + "py.services.websocket_manager.ws_manager.broadcast", fake_broadcast + ) + model_handlers._broadcast_models_changed() + await asyncio.sleep(0) diff --git a/web/comfyui/lm_styles.css b/web/comfyui/lm_styles.css index ffb6b249..235bba83 100644 --- a/web/comfyui/lm_styles.css +++ b/web/comfyui/lm_styles.css @@ -247,6 +247,25 @@ border-left: 3px solid rgba(245, 158, 11, 0.6) !important; } +.lm-lora-entry[data-missing="true"] { + background-color: rgba(220, 38, 38, 0.16); + border: 1px solid rgba(220, 38, 38, 0.5); +} + +.lm-lora-entry[data-missing="true"] .lm-lora-name { + color: rgba(252, 165, 165, 0.95); +} + +.lm-lora-clip-entry[data-missing="true"] { + background-color: rgba(220, 38, 38, 0.12); + border: 1px solid rgba(220, 38, 38, 0.4); + border-left: 2px solid rgba(220, 38, 38, 0.6); +} + +.lm-lora-clip-entry[data-missing="true"] .lm-lora-name { + color: rgba(252, 165, 165, 0.95); +} + .lm-lora-name { margin-left: 4px; flex: 1; diff --git a/web/comfyui/loras_widget.js b/web/comfyui/loras_widget.js index 8a9e1c16..8599366c 100644 --- a/web/comfyui/loras_widget.js +++ b/web/comfyui/loras_widget.js @@ -3,7 +3,11 @@ import { parseLoraValue, formatLoraValue, shouldShowClipEntry, - syncClipStrengthIfCollapsed + syncClipStrengthIfCollapsed, + getAvailableLoras, + getAvailableLorasSync, + isLoraNameAvailable, + onLibraryChanged } from "./loras_widget_utils.js"; import { initDrag, createContextMenu, initHeaderDrag, initReorderDrag, handleKeyboardNavigation } from "./loras_widget_events.js"; import { forwardMiddleMouseToCanvas, forwardWheelToCanvas, enableListWheelScroll } from "./utils.js"; @@ -140,6 +144,48 @@ export function addLorasWidget(node, name, opts, callback) { emitSelectionChange(buildSelectionPayload(loraName)); } }; + + // Mirror ComfyUI's setNodeHasErrors: has_errors is not an auto-tracked + // litegraph property, so the node:property:changed event must be fired + // manually for the Vue renderer to pick up the error state. + // + // The flag is applied asynchronously (setTimeout 0): applying it during + // LGraphNode.configure makes ComfyUI's errorNodeWidgets.onConfigure create + // a fallback UNKNOWN widget for every widgets_values entry, because it + // treats has_errors as "node definition missing". + let pendingErrorFlag = null; + let errorFlagTimer = null; + + const flushErrorFlag = () => { + errorFlagTimer = null; + const hasMissing = pendingErrorFlag; + pendingErrorFlag = null; + if (typeof hasMissing !== 'boolean') { + return; + } + const oldValue = node.has_errors === true; + if (oldValue === hasMissing) { + return; + } + node.has_errors = hasMissing; + if (node.graph) { + node.graph.trigger('node:property:changed', { + type: 'node:property:changed', + nodeId: node.id, + property: 'has_errors', + oldValue, + newValue: hasMissing + }); + node.graph.setDirtyCanvas(true, true); + } + }; + + const updateNodeErrorFlag = (hasMissing) => { + pendingErrorFlag = hasMissing; + if (errorFlagTimer === null) { + errorFlagTimer = setTimeout(flushErrorFlag, 0); + } + }; // Add keyboard event listener to container container.addEventListener('keydown', (e) => { @@ -220,6 +266,7 @@ export function addLorasWidget(node, name, opts, callback) { emptyMessage.textContent = "No LoRAs added"; emptyMessage.className = "lm-lora-empty-state"; container.appendChild(emptyMessage); + updateNodeErrorFlag(false); return; } @@ -267,8 +314,21 @@ export function addLorasWidget(node, name, opts, callback) { initHeaderDrag(header, widget, renderLoras); // Render each lora entry + const availableSet = getAvailableLorasSync(); + if (!availableSet) { + // Availability data missing (workflow switch without node recreation, + // cache expiry): fetch it and re-render once it lands so missing cues + // and the node flag always resolve. Only re-render on success to avoid + // retry loops on failure. + getAvailableLoras().then((set) => { + if (set && !widget.__dragActive && container.isConnected) { + renderLoras(widget.value, widget); + } + }); + } lorasData.forEach((loraData) => { const { name, strength, clipStrength, active } = loraData; + const missing = !isLoraNameAvailable(name, availableSet); // Determine expansion state using our helper function const isExpanded = shouldShowClipEntry(loraData); @@ -283,6 +343,10 @@ export function addLorasWidget(node, name, opts, callback) { loraEl.dataset.active = active ? "true" : "false"; loraEl.dataset.locked = (loraData.locked || false) ? "true" : "false"; + if (missing) { + loraEl.setAttribute("data-missing", "true"); + } + // Add click handler for selection loraEl.addEventListener('click', (e) => { // Skip if clicking on interactive elements @@ -374,6 +438,9 @@ export function addLorasWidget(node, name, opts, callback) { const nameEl = document.createElement("div"); nameEl.textContent = name; nameEl.className = "lm-lora-name"; + if (missing) { + nameEl.title = "LoRA not found in local library"; + } // Move preview tooltip events to nameEl instead of loraEl let previewTimer = null; // Timer for delayed preview @@ -387,7 +454,8 @@ export function addLorasWidget(node, name, opts, callback) { nameEl.addEventListener('mouseenter', (e) => { e.stopPropagation(); - if (shouldSuppressPreview()) { + // Missing LoRAs have no preview data — skip the placeholder tooltip. + if (missing || shouldSuppressPreview()) { return; } previewTimer = setTimeout(async () => { @@ -544,10 +612,17 @@ export function addLorasWidget(node, name, opts, callback) { clipEl.dataset.loraName = name; clipEl.dataset.active = active ? "true" : "false"; + if (missing) { + clipEl.setAttribute("data-missing", "true"); + } + // Create clip name display const clipNameEl = document.createElement("div"); clipNameEl.textContent = "[clip] " + name; clipNameEl.className = "lm-lora-name"; + if (missing) { + clipNameEl.title = "LoRA not found in local library"; + } // Create clip strength control const clipStrengthControl = document.createElement("div"); @@ -667,6 +742,16 @@ export function addLorasWidget(node, name, opts, callback) { updateEntrySelection(entry, entryLoraName === selectedLora); }); + // Flag the node when any active entry references a LoRA missing locally. + // Skipped while the availability set is not loaded (null) to avoid + // clearing or setting the flag based on incomplete information. + const hasMissingActive = availableSet + ? lorasData.some( + (lora) => lora.active && !isLoraNameAvailable(lora.name, availableSet) + ) + : null; + updateNodeErrorFlag(hasMissingActive); + const selectionExists = selectedLora ? currentLorasData.some((lora) => lora.name === selectedLora) : false; @@ -767,7 +852,30 @@ export function addLorasWidget(node, name, opts, callback) { widget.callback = callback; + // Invalidate the availability cache and re-render when the local library + // changes (e.g. a LoRA is deleted from the Lora Manager UI) so missing + // cues and the node error flag update without waiting for the TTL. + const unsubscribeLibraryChange = onLibraryChanged(() => { + if (!widget.__dragActive && container.isConnected) { + renderLoras(widget.value, widget); + } + }); + + // Fetch the local library and re-render once available so missing entries + // get their visual cue and the node error flag as soon as the data lands. + getAvailableLoras().then(() => { + if (!widget.__dragActive && container.isConnected) { + renderLoras(widget.value, widget); + } + }); + widget.onRemove = () => { + unsubscribeLibraryChange(); + if (errorFlagTimer !== null) { + clearTimeout(errorFlagTimer); + errorFlagTimer = null; + pendingErrorFlag = null; + } while (container.firstChild) { container.removeChild(container.firstChild); } diff --git a/web/comfyui/loras_widget_utils.js b/web/comfyui/loras_widget_utils.js index bcc8eeb9..4b23bfad 100644 --- a/web/comfyui/loras_widget_utils.js +++ b/web/comfyui/loras_widget_utils.js @@ -1,4 +1,246 @@ import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +// Mirrors the backend resolver (get_lora_info_absolute): a ".ckpt"/".pt" +// reference resolves to the same-named .safetensors file. The scanner only +// indexes .safetensors, but keeping these here lets legacy references match. +const LORA_FILE_EXTENSIONS = [".safetensors", ".ckpt", ".pt", ".bin"]; + +/** + * Strip a known LoRA model extension from a name (case-insensitive). + * + * The two sides of the availability check differ: + * - The collection side (cycler-list `file_name`) is stored extension-free + * (scanner convention), so stripping is a no-op there. + * - The widget entry side comes from the autocomplete path, which returns + * on-disk relative paths WITH the extension (e.g. + * "Illustrious/lazyhand.safetensors"), so stripping is required to match. + */ +export function stripLoraExtension(name) { + const lowered = String(name || "").toLowerCase(); + for (const ext of LORA_FILE_EXTENSIONS) { + if (lowered.endsWith(ext)) { + return name.slice(0, -ext.length); + } + } + return name; +} + +/** + * Normalize a LoRA name for availability lookup: forward slashes and no + * extension, mirroring the backend matching in get_lora_info_absolute. + */ +export function normalizeLoraNameKey(name) { + return stripLoraExtension(String(name || "").replace(/\\/g, "/")); +} + +/** + * Build the lookup set of available LoRA names from relative paths like + * "folder/lora.safetensors". Both the full path and the bare basename are + * registered (extension stripped), matching how users can reference LoRAs. + */ +export function buildAvailableLoraSet(relativePaths) { + const set = new Set(); + for (const p of relativePaths || []) { + const normalized = normalizeLoraNameKey(p); + if (!normalized) continue; + set.add(normalized); + const slash = normalized.lastIndexOf("/"); + if (slash >= 0) { + set.add(normalized.slice(slash + 1)); + } + } + return set; +} + +/** + * Check whether a widget entry name is available locally. + * + * When the availability set is not loaded yet (null), every name is treated + * as available so entries are never falsely flagged while the fetch is + * pending. Absolute paths outside the library cannot be verified + * client-side and are treated as available. A folder-qualified name that + * does not match a stored path falls back to its basename, mirroring the + * backend resolver (get_lora_info_absolute's basename fallback and the + * legacy syntax format). + */ +export function isLoraNameAvailable(name, availableSet) { + if (!availableSet) { + return true; + } + const normalized = String(name || "").replace(/\\/g, "/"); + if (normalized.startsWith("/") || /^[A-Za-z]:\//.test(normalized)) { + return true; + } + const key = normalizeLoraNameKey(name); + if (availableSet.has(key)) { + return true; + } + const slash = key.lastIndexOf("/"); + if (slash >= 0) { + return availableSet.has(key.slice(slash + 1)); + } + return false; +} + +const AVAILABLE_LORAS_TTL_MS = 60000; +let availableLorasCache = null; +let availableLorasPromise = null; +let availabilityGeneration = 0; + +async function refreshAvailableLoras() { + const generation = availabilityGeneration; + try { + const response = await api.fetchApi("/lm/loras/cycler-list", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: "{}", + }); + if (!response || !response.ok) { + return null; + } + const data = await response.json(); + const paths = (data?.loras || []) + .map((lora) => lora?.file_name) + .filter(Boolean); + const set = buildAvailableLoraSet(paths); + if (generation !== availabilityGeneration) { + // Stale response: the cache was invalidated while this fetch was in + // flight, do not repopulate it with pre-change data. + return null; + } + availableLorasCache = { set, at: Date.now() }; + return set; + } catch (error) { + console.warn("Failed to fetch available LoRAs:", error); + return null; + } +} + +/** + * Fetch the set of available LoRA names, cached with a TTL. Concurrent + * callers share a single in-flight request. Resolves to null on failure. + */ +export function getAvailableLoras() { + connectLibraryChangeSocket(); + if ( + availableLorasCache && + Date.now() - availableLorasCache.at < AVAILABLE_LORAS_TTL_MS + ) { + return Promise.resolve(availableLorasCache.set); + } + if (!availableLorasPromise) { + availableLorasPromise = refreshAvailableLoras().finally(() => { + availableLorasPromise = null; + }); + } + return availableLorasPromise; +} + +/** + * Synchronous snapshot of the cached availability set, or null when the + * cache is not loaded (or expired). + */ +export function getAvailableLorasSync() { + if ( + availableLorasCache && + Date.now() - availableLorasCache.at < AVAILABLE_LORAS_TTL_MS + ) { + return availableLorasCache.set; + } + return null; +} + +/** + * Drop the cached availability data (used by tests and by callers that need + * a forced refresh of the local library state). In-flight fetches started + * before the reset are invalidated via the generation counter. + */ +export function resetAvailableLorasCache() { + availabilityGeneration += 1; + availableLorasCache = null; + availableLorasPromise = null; +} + +// The Lora Manager UI and the ComfyUI graph page are separate pages; the +// backend broadcasts "models_changed" over its WebSocket when the local +// library changes (delete/rename/move/scan), so the graph page can +// invalidate its availability cache immediately instead of waiting for the +// TTL to expire. +const libraryChangeListeners = new Set(); + +/** + * Register a callback fired whenever the local model library changes. + * Returns an unsubscribe function. + */ +export function onLibraryChanged(callback) { + libraryChangeListeners.add(callback); + return () => { + libraryChangeListeners.delete(callback); + }; +} + +/** + * Process a library-change WebSocket message. Exported for testability. + */ +export function handleLibraryChangeMessage(data) { + if (!data || data.type !== "models_changed") { + return; + } + resetAvailableLorasCache(); + for (const listener of libraryChangeListeners) { + try { + listener(); + } catch (error) { + console.warn("Library change listener failed:", error); + } + } +} + +const LIBRARY_WS_RECONNECT_MS = 30000; +let libraryWs = null; +let libraryWsRetryTimer = null; + +function connectLibraryChangeSocket() { + if (libraryWs || typeof WebSocket === "undefined") { + return; + } + const protocol = window.location.protocol === "https:" ? "wss://" : "ws://"; + let ws; + try { + ws = new WebSocket(`${protocol}${window.location.host}/ws/fetch-progress`); + } catch (error) { + return; + } + libraryWs = ws; + ws.onmessage = (event) => { + try { + handleLibraryChangeMessage(JSON.parse(event.data)); + } catch (error) { + // Non-JSON messages from other broadcasters are ignored. + } + }; + ws.onclose = () => { + libraryWs = null; + if (libraryWsRetryTimer === null) { + libraryWsRetryTimer = setTimeout(() => { + libraryWsRetryTimer = null; + connectLibraryChangeSocket(); + }, LIBRARY_WS_RECONNECT_MS); + } + }; + ws.onerror = () => { + ws.close(); + }; +} + +/** + * Ensure the library-change WebSocket is connected (idempotent). Called on + * the first availability fetch; safe in environments without WebSocket. + */ +export function ensureLibraryChangeSocket() { + connectLibraryChangeSocket(); +} // Parse LoRA entries from value export function parseLoraValue(value) {