From 1e7cad595c9ddd77eb8a7bc4ba9869e33024e9ca Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Tue, 28 Jul 2026 14:43:37 +0200 Subject: [PATCH] refactor: extract custom-node model adapters from inline handlers --- __init__.py | 31 ++- core/custom_nodes/README.md | 16 ++ core/custom_nodes/__init__.py | 29 +++ core/custom_nodes/base.py | 19 ++ core/custom_nodes/lora_manager.py | 234 ++++++++++++++++++ core/custom_nodes/registry.py | 158 ++++++++++++ .../custom_nodes/rgthree_power_lora_loader.py | 12 + core/resolver.py | 38 +-- core/workflow_analyzer.py | 124 ++-------- core/workflow_updater.py | 96 +------ tests/test_downloads_tab_workflow_route.mjs | 134 +++++----- tests/test_resolver_robustness.py | 4 +- tests/test_workflow_analyzer.py | 35 +++ web/resolver.js | 4 +- web/resolver/actions/queue_methods.js | 31 ++- web/resolver/custom_nodes/README.md | 22 ++ web/resolver/custom_nodes/lora_manager.js | 51 ++++ web/resolver/custom_nodes/registry.js | 96 +++++++ .../custom_nodes/rgthree_power_lora_loader.js | 50 ++++ web/resolver/model_resolver.js | 191 +++----------- .../search/download_target_methods.js | 6 +- web/resolver/shell/lifecycle_graph_methods.js | 3 +- web/resolver/views/tabs_loaded_methods.js | 40 +-- 23 files changed, 934 insertions(+), 490 deletions(-) create mode 100644 core/custom_nodes/README.md create mode 100644 core/custom_nodes/__init__.py create mode 100644 core/custom_nodes/base.py create mode 100644 core/custom_nodes/lora_manager.py create mode 100644 core/custom_nodes/registry.py create mode 100644 core/custom_nodes/rgthree_power_lora_loader.py create mode 100644 web/resolver/custom_nodes/README.md create mode 100644 web/resolver/custom_nodes/lora_manager.js create mode 100644 web/resolver/custom_nodes/registry.js create mode 100644 web/resolver/custom_nodes/rgthree_power_lora_loader.js diff --git a/__init__.py b/__init__.py index 1903271..624613d 100644 --- a/__init__.py +++ b/__init__.py @@ -355,6 +355,10 @@ class ModelResolverExtension: # Import resolver modules try: + from .core.custom_nodes import ( + adapt_custom_node_loaded_model, + should_skip_existing_custom_node_reference, + ) from .core.metadata_audit import audit_metadata_sizes from .core.metadata_builder import ( build_missing_local_metadata, @@ -712,22 +716,14 @@ class ModelResolverExtension: force_rescan=force_rescan, ) - # Filter out LoraManager lorAs that already exist locally (exists=True) - # These should not appear in missing models at all missing_models = result.get("missing_models", []) filtered_missing = [] for missing in missing_models: - is_lora = missing.get("is_lora_v2") - exists = missing.get("exists") name = missing.get("name") or missing.get("original_path", "") - self.logger.debug( - f"Filtering: {name} is_lora_v2={is_lora} exists={exists}" - ) - - # Skip LoraManager lorAs that already exist locally - if is_lora and exists: + if should_skip_existing_custom_node_reference(missing): self.logger.info( - f"Filtered out LoraManager lora: {name}" + "Filtered existing custom-node model " + f"reference: {name}" ) continue filtered_missing.append(missing) @@ -2001,10 +1997,11 @@ class ModelResolverExtension: if ref.get("strength") is not None: strength = ref.get("strength") - # For text-based lora loaders (LoraLoaderV2, LoraManager), get strength from ref - if ref.get("is_lora_v2"): - strength = ref.get("strength") - model_name = ref.get("name", model_name) + model_name, strength = adapt_custom_node_loaded_model( + ref, + model_name, + strength, + ) # Check if model exists locally exists = ref.get("exists", False) @@ -2050,7 +2047,9 @@ class ModelResolverExtension: "strength": strength, "original_path": original_path, "is_urn": ref.get("is_urn", False), - "is_lora_v2": ref.get("is_lora_v2", False), + "custom_node_adapter": ref.get( + "custom_node_adapter" + ), "active": ref.get("active"), "connected": ref.get("connected", True), "resolved_path": ( diff --git a/core/custom_nodes/README.md b/core/custom_nodes/README.md new file mode 100644 index 0000000..8d9bbc7 --- /dev/null +++ b/core/custom_nodes/README.md @@ -0,0 +1,16 @@ +# Backend custom-node model adapters + +Backend integrations for non-standard workflow node formats live here. + +Each adapter declares its node types, category hints, and fixed widget +categories. Optional hooks can provide: + +- custom workflow reference extraction; +- lightweight potential-reference detection; +- custom workflow path updates; +- existing-reference filtering; +- Loaded Models display adjustments. + +Register new adapters in `registry.py`. Generic dictionary-backed model fields, +including `nested_key`, stay in the shared analyzer and updater because they are +not specific to one node package. diff --git a/core/custom_nodes/__init__.py b/core/custom_nodes/__init__.py new file mode 100644 index 0000000..112bc78 --- /dev/null +++ b/core/custom_nodes/__init__.py @@ -0,0 +1,29 @@ +"""Backend custom-node model adapter registry.""" + +from .registry import ( + CUSTOM_NODE_MODEL_ADAPTERS, + adapt_custom_node_loaded_model, + analyze_custom_node_references, + custom_node_has_potential_model_reference, + get_custom_node_adapter_for_reference, + get_custom_node_category_hints, + get_custom_node_model_adapter, + get_custom_node_resolution_metadata, + get_custom_node_widget_categories, + should_skip_existing_custom_node_reference, + update_custom_node_model_path, +) + +__all__ = [ + "CUSTOM_NODE_MODEL_ADAPTERS", + "adapt_custom_node_loaded_model", + "analyze_custom_node_references", + "custom_node_has_potential_model_reference", + "get_custom_node_adapter_for_reference", + "get_custom_node_category_hints", + "get_custom_node_model_adapter", + "get_custom_node_resolution_metadata", + "get_custom_node_widget_categories", + "should_skip_existing_custom_node_reference", + "update_custom_node_model_path", +] diff --git a/core/custom_nodes/base.py b/core/custom_nodes/base.py new file mode 100644 index 0000000..c39144e --- /dev/null +++ b/core/custom_nodes/base.py @@ -0,0 +1,19 @@ +"""Shared contract for backend custom-node model adapters.""" + +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, Optional, Tuple + + +@dataclass(frozen=True) +class CustomNodeModelAdapter: + """Describe backend behavior supplied by one custom-node integration.""" + + adapter_id: str + node_types: Tuple[str, ...] + category_hint: Optional[str] = None + widget_categories: Dict[int, str] = field(default_factory=dict) + analyze_references: Optional[Callable[..., Any]] = None + has_potential_reference: Optional[Callable[..., bool]] = None + update_model_path: Optional[Callable[..., Optional[bool]]] = None + should_skip_existing: Optional[Callable[..., bool]] = None + adapt_loaded_model: Optional[Callable[..., Any]] = None diff --git a/core/custom_nodes/lora_manager.py b/core/custom_nodes/lora_manager.py new file mode 100644 index 0000000..5b69de7 --- /dev/null +++ b/core/custom_nodes/lora_manager.py @@ -0,0 +1,234 @@ +"""Backend adapter for ComfyUI-Lora-Manager loader nodes.""" + +import os +from typing import Any, Dict, List, Optional + +from ..log_system import create_module_logger +from .base import CustomNodeModelAdapter + +log = create_module_logger(__name__) + +ADAPTER_ID = "lora-manager" +NODE_TYPES = ( + "LoraLoaderV2", + "Lora Loader (LoraManager)", + "Lora Stacker (LoraManager)", +) +LORA_LIST_WIDGET_INDEX = 2 +TEXT_WIDGET_INDEX = 1 + + +def analyze_references( + node: Dict[str, Any], + available_models: Optional[List[Dict[str, Any]]] = None, + *, + is_active: bool, + get_widget_name_hint: Any, +) -> Optional[List[Dict[str, Any]]]: + """Extract LoRA references stored in Lora Manager's list widget.""" + widgets_values = node.get("widgets_values", []) + if len(widgets_values) < 3: + return None + + from ..scanner import get_model_files + + all_loras = ( + available_models if available_models is not None else get_model_files() + ) + lora_files = [ + model for model in all_loras if model.get("category") == "loras" + ] + lora_lookup: Dict[str, List[Dict[str, Any]]] = {} + for lora_file in lora_files: + filename = lora_file.get("filename", "") + if not filename: + continue + base_name = os.path.splitext(filename)[0] + lora_lookup.setdefault(base_name, []).append(lora_file) + + lora_list = widgets_values[LORA_LIST_WIDGET_INDEX] + if not isinstance(lora_list, list): + return [] + + node_id = node.get("id") + node_type = node.get("type", "") + node_title = str(node.get("title", "") or "").strip() + model_refs: List[Dict[str, Any]] = [] + for lora_item in lora_list: + if not isinstance(lora_item, dict): + continue + + name = lora_item.get("name", "") + if not name: + continue + + lora_exists = False + lora_full_path = None + if name in lora_lookup: + lora_full_path = lora_lookup[name][0].get("path") + lora_exists = ( + os.path.exists(lora_full_path) if lora_full_path else False + ) + else: + for extension in [".safetensors", ".ckpt", ".pt", ".pth"]: + test_name = name + extension + if test_name not in lora_lookup: + continue + lora_full_path = lora_lookup[test_name][0].get("path") + lora_exists = ( + os.path.exists(lora_full_path) if lora_full_path else False + ) + if lora_exists: + break + + log.debug(f"Lora {name}: exists={lora_exists}, path={lora_full_path}") + model_refs.append( + { + "node_id": node_id, + "node_type": node_type, + "widget_index": LORA_LIST_WIDGET_INDEX, + "widget_name": get_widget_name_hint( + node, LORA_LIST_WIDGET_INDEX + ), + "original_path": name, + "name": name, + "strength": float(lora_item.get("strength", 1.0)), + "active": lora_item.get("active", True), + "node_title": node_title, + "category": "loras", + "category_hints": ["loras"], + "folder_key_hints": ["loras"], + "full_path": lora_full_path, + "exists": lora_exists, + "is_urn": False, + "custom_node_adapter": ADAPTER_ID, + "connected": is_active, + } + ) + return model_refs + + +def has_potential_reference(node: Dict[str, Any]) -> bool: + """Return whether the serialized LoRA list contains a named entry.""" + widgets_values = node.get("widgets_values") + if not isinstance(widgets_values, list) or len(widgets_values) < 3: + return False + lora_list = widgets_values[LORA_LIST_WIDGET_INDEX] + if not isinstance(lora_list, list): + return False + return any( + isinstance(item, dict) and str(item.get("name") or "").strip() + for item in lora_list + ) + + +def update_model_path( + node: Dict[str, Any], + widget_index: int, + resolved_model: Optional[Dict[str, Any]], + mapping: Optional[Dict[str, Any]], +) -> Optional[bool]: + """Update one LoRA name in both the list and formatted text widgets.""" + mapping = mapping or {} + adapter_id = mapping.get("custom_node_adapter") + is_legacy_mapping = mapping.get("is_lora_v2") is True + if adapter_id != ADAPTER_ID and not is_legacy_mapping: + return None + + original_name = ( + mapping.get("custom_node_original_identity") + or mapping.get("original_lora_name") + ) + if not original_name or widget_index != LORA_LIST_WIDGET_INDEX: + return None + + widgets_values = node.get("widgets_values", []) + lora_list = widgets_values[LORA_LIST_WIDGET_INDEX] + if not isinstance(lora_list, list): + log.warning( + "Lora Manager list widget is not a list: " + f"{type(lora_list)}" + ) + return False + + new_name = None + if resolved_model: + new_name = resolved_model.get("filename") or resolved_model.get( + "name", "" + ) + if new_name and "." in new_name: + new_name = new_name.rsplit(".", 1)[0] + if not new_name: + return False + + original_stripped = str(original_name).strip() + updated = False + for lora_item in lora_list: + if not isinstance(lora_item, dict): + continue + current_name = str(lora_item.get("name", "")).strip() + if ( + current_name == original_stripped + or current_name.lower() == original_stripped.lower() + ): + lora_item["name"] = new_name + updated = True + break + + if not updated: + available = [ + item.get("name") for item in lora_list if isinstance(item, dict) + ] + log.warning( + f"Lora '{original_name}' not found in Lora Manager list. " + f"Available: {available}" + ) + return False + + if ( + len(widgets_values) > TEXT_WIDGET_INDEX + and isinstance(widgets_values[TEXT_WIDGET_INDEX], str) + ): + old_text = widgets_values[TEXT_WIDGET_INDEX] + new_text = old_text.replace( + f" {new_name}" + ) + return True + + +def should_skip_existing(reference: Dict[str, Any]) -> bool: + """Existing list entries do not require matching or relinking.""" + return reference.get("exists") is True + + +def adapt_loaded_model( + reference: Dict[str, Any], + model_name: str, + strength: Any, +) -> tuple[str, Any]: + """Use the list entry's display name and strength.""" + return ( + reference.get("name", model_name), + reference.get("strength", strength), + ) + + +ADAPTER = CustomNodeModelAdapter( + adapter_id=ADAPTER_ID, + node_types=NODE_TYPES, + category_hint="loras", + widget_categories={LORA_LIST_WIDGET_INDEX: "loras"}, + analyze_references=analyze_references, + has_potential_reference=has_potential_reference, + update_model_path=update_model_path, + should_skip_existing=should_skip_existing, + adapt_loaded_model=adapt_loaded_model, +) diff --git a/core/custom_nodes/registry.py b/core/custom_nodes/registry.py new file mode 100644 index 0000000..3421142 --- /dev/null +++ b/core/custom_nodes/registry.py @@ -0,0 +1,158 @@ +"""Registry and dispatch helpers for backend custom-node model adapters.""" + +from typing import Any, Dict, List, Optional + +from .base import CustomNodeModelAdapter +from .lora_manager import ADAPTER as LORA_MANAGER_ADAPTER +from .rgthree_power_lora_loader import ( + ADAPTER as RGTHREE_POWER_LORA_LOADER_ADAPTER, +) + +CUSTOM_NODE_MODEL_ADAPTERS = ( + LORA_MANAGER_ADAPTER, + RGTHREE_POWER_LORA_LOADER_ADAPTER, +) +_ADAPTERS_BY_ID = { + adapter.adapter_id: adapter for adapter in CUSTOM_NODE_MODEL_ADAPTERS +} +_ADAPTERS_BY_NODE_TYPE = { + node_type: adapter + for adapter in CUSTOM_NODE_MODEL_ADAPTERS + for node_type in adapter.node_types +} + + +def get_custom_node_model_adapter( + node_or_type: Any, +) -> Optional[CustomNodeModelAdapter]: + """Return the adapter registered for a node dictionary or node type.""" + node_type = ( + node_or_type.get("type", "") + if isinstance(node_or_type, dict) + else str(node_or_type or "") + ) + return _ADAPTERS_BY_NODE_TYPE.get(node_type) + + +def get_custom_node_adapter_for_reference( + reference: Dict[str, Any], +) -> Optional[CustomNodeModelAdapter]: + """Resolve an adapter from normalized or legacy reference metadata.""" + adapter_id = reference.get("custom_node_adapter") + if adapter_id in _ADAPTERS_BY_ID: + return _ADAPTERS_BY_ID[adapter_id] + adapter = get_custom_node_model_adapter(reference.get("node_type")) + if adapter: + return adapter + if reference.get("is_lora_v2"): + return LORA_MANAGER_ADAPTER + return None + + +def get_custom_node_category_hints() -> Dict[str, str]: + """Return node-type category hints contributed by adapters.""" + return { + node_type: adapter.category_hint + for adapter in CUSTOM_NODE_MODEL_ADAPTERS + if adapter.category_hint + for node_type in adapter.node_types + } + + +def get_custom_node_widget_categories() -> Dict[str, Dict[int, str]]: + """Return fixed serialized widget categories contributed by adapters.""" + return { + node_type: dict(adapter.widget_categories) + for adapter in CUSTOM_NODE_MODEL_ADAPTERS + if adapter.widget_categories + for node_type in adapter.node_types + } + + +def get_custom_node_resolution_metadata( + reference: Dict[str, Any], +) -> Dict[str, Any]: + """Normalize adapter metadata carried into workflow update mappings.""" + adapter = get_custom_node_adapter_for_reference(reference) + if not adapter: + return {} + original_identity = ( + reference.get("custom_node_original_identity") + or reference.get("original_lora_name") + or reference.get("name") + or reference.get("original_path") + ) + return { + "custom_node_adapter": adapter.adapter_id, + "custom_node_original_identity": original_identity, + } + + +def analyze_custom_node_references( + node: Dict[str, Any], + available_models: Optional[List[Dict[str, Any]]] = None, + **context: Any, +) -> Optional[List[Dict[str, Any]]]: + """Dispatch custom serialized model extraction when an adapter handles it.""" + adapter = get_custom_node_model_adapter(node) + if not adapter or not adapter.analyze_references: + return None + return adapter.analyze_references( + node, + available_models, + **context, + ) + + +def custom_node_has_potential_model_reference( + node: Dict[str, Any], +) -> bool: + """Return whether a registered adapter finds a model reference.""" + adapter = get_custom_node_model_adapter(node) + return bool( + adapter + and adapter.has_potential_reference + and adapter.has_potential_reference(node) + ) + + +def update_custom_node_model_path( + node: Dict[str, Any], + widget_index: int, + resolved_model: Optional[Dict[str, Any]], + mapping: Optional[Dict[str, Any]], +) -> Optional[bool]: + """Dispatch a custom workflow update, returning None when not handled.""" + adapter = get_custom_node_model_adapter(node) + if not adapter or not adapter.update_model_path: + return None + return adapter.update_model_path( + node, + widget_index, + resolved_model, + mapping, + ) + + +def should_skip_existing_custom_node_reference( + reference: Dict[str, Any], +) -> bool: + """Return whether an existing adapter reference needs no resolution.""" + adapter = get_custom_node_adapter_for_reference(reference) + return bool( + adapter + and adapter.should_skip_existing + and adapter.should_skip_existing(reference) + ) + + +def adapt_custom_node_loaded_model( + reference: Dict[str, Any], + model_name: str, + strength: Any, +) -> tuple[str, Any]: + """Apply adapter-specific display data for Loaded Models.""" + adapter = get_custom_node_adapter_for_reference(reference) + if not adapter or not adapter.adapt_loaded_model: + return model_name, strength + return adapter.adapt_loaded_model(reference, model_name, strength) diff --git a/core/custom_nodes/rgthree_power_lora_loader.py b/core/custom_nodes/rgthree_power_lora_loader.py new file mode 100644 index 0000000..b446d63 --- /dev/null +++ b/core/custom_nodes/rgthree_power_lora_loader.py @@ -0,0 +1,12 @@ +"""Backend adapter metadata for rgthree Power Lora Loader.""" + +from .base import CustomNodeModelAdapter + +ADAPTER_ID = "rgthree-power-lora-loader" +NODE_TYPES = ("Power Lora Loader (rgthree)",) + +ADAPTER = CustomNodeModelAdapter( + adapter_id=ADAPTER_ID, + node_types=NODE_TYPES, + category_hint="loras", +) diff --git a/core/resolver.py b/core/resolver.py index ecd8c13..041f1f2 100644 --- a/core/resolver.py +++ b/core/resolver.py @@ -11,6 +11,11 @@ import threading from typing import Any, Callable, Dict, List, Optional, Tuple from urllib.parse import unquote +from .custom_nodes import ( + custom_node_has_potential_model_reference, + get_custom_node_resolution_metadata, + should_skip_existing_custom_node_reference, +) from .log_system import create_module_logger log = create_module_logger(__name__) @@ -456,19 +461,8 @@ def node_has_potential_model_reference(node: Dict[str, Any]) -> bool: if not isinstance(widgets_values, list) or not widgets_values: return False - node_type = node.get("type", "") - if ( - node_type in { - "LoraLoaderV2", - "Lora Loader (LoraManager)", - "Lora Stacker (LoraManager)", - } - and len(widgets_values) >= 3 - and isinstance(widgets_values[2], list) - ): - for lora_item in widgets_values[2]: - if isinstance(lora_item, dict) and str(lora_item.get("name") or "").strip(): - return True + if custom_node_has_potential_model_reference(node): + return True for idx, value in enumerate(widgets_values): model_widget_category_hint = get_model_widget_category_hint(node, idx) @@ -1177,14 +1171,11 @@ def analyze_and_find_matches( } ) - # Skip LoraManager lorAs that already exist locally (exists=True means no linking needed) - is_lora_v2 = missing.get("is_lora_v2") - exists = missing.get("exists") name = missing.get("name") or missing.get("original_path", "") - log.debug(f"Checking {name}: is_lora_v2={is_lora_v2}, exists={exists}") - - if is_lora_v2 and exists: - log.info(f"Skipping LoraManager lora {name} - already exists locally") + if should_skip_existing_custom_node_reference(missing): + log.info( + f"Skipping existing custom-node model reference: {name}" + ) continue missing_with_matches.append({ @@ -1260,13 +1251,10 @@ def apply_resolution( "is_top_level": resolution.get( "is_top_level" ), # True for top-level nodes, False for nodes in subgraph definitions - "is_lora_v2": resolution.get("is_lora_v2"), # Flag for LoraManager nodes - "original_lora_name": resolution.get( - "original_lora_name" - ), # Original lora name for LoraManager replacement "nested_key": resolution.get( "nested_key" - ), # For dict-type widgets (e.g. Power Lora Loader) + ), # For dict-type widget values + **get_custom_node_resolution_metadata(resolution), } # If resolved_model provided, extract path if needed diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 731437a..eb3c683 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -40,6 +40,12 @@ except ImportError: log.warning("Model Resolver: folder_paths not available yet - will retry later") +from .custom_nodes import ( + analyze_custom_node_references, + get_custom_node_category_hints, + get_custom_node_model_adapter, + get_custom_node_widget_categories, +) from .type_utils import ( MODEL_EXTENSIONS, URN_TYPE_MAP, @@ -155,6 +161,7 @@ NODE_TYPE_MODEL_WIDGET_CATEGORIES = { "LTXVGemmaCLIPModelLoader": {0: "text_encoders"}, "LTXAVTextEncoderLoader": {0: "text_encoders", 1: "checkpoints"}, } +NODE_TYPE_MODEL_WIDGET_CATEGORIES.update(get_custom_node_widget_categories()) # Mapping of common node types to their expected model category. # Derived dynamically from NODE_TYPE_MODEL_WIDGET_CATEGORIES. @@ -163,14 +170,7 @@ NODE_TYPE_TO_CATEGORY_HINTS = { for node_type, widget_map in NODE_TYPE_MODEL_WIDGET_CATEGORIES.items() if widget_map } - -# Fallbacks and overrides for custom nodes that do not use standard model widget indices -NODE_TYPE_TO_CATEGORY_HINTS.update({ - "LoraLoaderV2": "loras", - "Lora Loader (LoraManager)": "loras", - "Lora Stacker (LoraManager)": "loras", - "Power Lora Loader (rgthree)": "loras", -}) +NODE_TYPE_TO_CATEGORY_HINTS.update(get_custom_node_category_hints()) # Model category hints by widget/input name. Workflow JSON does not always preserve # widget names, but when it does this catches custom loaders without a node-type entry. @@ -223,8 +223,6 @@ MODEL_OUTPUT_TYPE_TO_CATEGORY = { } # Keys within dict-type widget values that contain model file references. -# Some nodes (e.g. rgthree Power Lora Loader) store model info as objects like -# {"on": true, "lora": "name.safetensors", "strength": 1.0} inside widgets_values. # Maps nested key name -> category hint. NESTED_MODEL_KEYS = { "lora": "loras", @@ -1617,98 +1615,14 @@ def get_node_model_info( if not widgets_values: return model_refs - # Special handling for text-based lora loaders (LoraLoaderV2, LoraManager, etc.) - lora_text_types = [ - "LoraLoaderV2", - "Lora Loader (LoraManager)", - "Lora Stacker (LoraManager)", - ] - is_lora_text = node_type in lora_text_types - if is_lora_text and len(widgets_values) >= 3: - # widgets_values[0] = {"version": 1, "textWidgetName": "text"} - # widgets_values[1] = " " - # widgets_values[2] = [{"name": "...", "strength": 1, "active": true}, ...] - - # Get all lora files using scanner for recursive search - from .scanner import get_model_files - - all_loras = available_models if available_models is not None else get_model_files() - lora_files = [m for m in all_loras if m.get("category") == "loras"] - - # Build a lookup by filename (without extension) - lora_lookup = {} - for lf in lora_files: - fname = lf.get("filename", "") - if fname: - # Get name without extension for matching - base_name = os.path.splitext(fname)[0] - if base_name not in lora_lookup: - lora_lookup[base_name] = [] - lora_lookup[base_name].append(lf) - - lora_list = widgets_values[2] - if isinstance(lora_list, list): - for lora_item in lora_list: - if isinstance(lora_item, dict): - name = lora_item.get("name", "") - strength = lora_item.get("strength", 1.0) - active = lora_item.get("active", True) - - if name: - # Check if lora exists locally using scanner data (recursive search) - lora_exists = False - lora_full_path = None - - # Try exact name first (without extension) - if name in lora_lookup: - lora_full_path = lora_lookup[name][0].get("path") - lora_exists = ( - os.path.exists(lora_full_path) - if lora_full_path - else False - ) - else: - # Try with common extensions - for ext in [".safetensors", ".ckpt", ".pt", ".pth"]: - test_name = name + ext - if test_name in lora_lookup: - lora_full_path = lora_lookup[test_name][0].get( - "path" - ) - lora_exists = ( - os.path.exists(lora_full_path) - if lora_full_path - else False - ) - if lora_exists: - break - - log.debug( - f"Lora {name}: exists={lora_exists}, path={lora_full_path}" - ) - - model_refs.append( - { - "node_id": node_id, - "node_type": node_type, - "widget_index": 2, # Index in lora list - "widget_name": get_widget_name_hint(node, 2), - "original_path": name, - "name": name, - "strength": float(strength), - "active": active, - "node_title": node_title, - "category": "loras", - "category_hints": ["loras"], - "folder_key_hints": ["loras"], - "full_path": lora_full_path, - "exists": lora_exists, - "is_urn": False, - "is_lora_v2": is_lora_text, - "connected": is_active, - } - ) - return model_refs + custom_model_refs = analyze_custom_node_references( + node, + available_models, + is_active=is_active, + get_widget_name_hint=get_widget_name_hint, + ) + if custom_model_refs is not None: + return custom_model_refs # For each widget value, check if it looks like a model file or URN for idx, value in enumerate(widgets_values): @@ -1802,8 +1716,7 @@ def get_node_model_info( category_backed_model_widget or named_model_file_widget ), ): - # Check for dict-type widget values containing model references (e.g. Power Lora Loader) - # Some nodes store model info as objects like {"on": true, "lora": "name.safetensors", "strength": 1.0} + # Check for dict-type widget values containing model references. if isinstance(value, dict): for nested_key, nested_category_hint in NESTED_MODEL_KEYS.items(): nested_value = value.get(nested_key) @@ -1850,6 +1763,9 @@ def get_node_model_info( "connected": is_active, "nested_key": nested_key, # Track nested key for updates } + custom_adapter = get_custom_node_model_adapter(node) + if custom_adapter: + ref["custom_node_adapter"] = custom_adapter.adapter_id if nested_key == "lora": strength = value.get("strength") diff --git a/core/workflow_updater.py b/core/workflow_updater.py index 7a76851..3252d90 100644 --- a/core/workflow_updater.py +++ b/core/workflow_updater.py @@ -7,6 +7,7 @@ Updates workflow JSON by replacing model paths in nodes. import os from typing import Any, Dict, List, Optional +from .custom_nodes import update_custom_node_model_path from .log_system import create_module_logger log = create_module_logger(__name__) @@ -140,7 +141,7 @@ def update_model_path( resolved_model: Model dict from scanner subgraph_id: Subgraph ID (if node is in subgraph) is_top_level: Whether node is in top-level (not subgraph definition) - mapping: Full mapping dict with is_lora_v2 and original_lora_name for LoraManager + mapping: Full mapping dictionary for standard or custom node updates Returns: True if update was successful, False otherwise @@ -212,89 +213,14 @@ def update_model_path( if not category and resolved_model: category = resolved_model.get("category") - # Check if this is a LoraManager (LoraLoaderV2) node with is_lora_v2 flag - is_lora_v2 = mapping.get("is_lora_v2") if mapping else False - original_lora_name = mapping.get("original_lora_name") if mapping else None - log.info( - f"@@ update_model_path: node={node_id}, idx={widget_index}, is_lora_v2={is_lora_v2}, original={original_lora_name}" + custom_update_result = update_custom_node_model_path( + node, + widget_index, + resolved_model, + mapping, ) - - # Special handling for LoraManager nodes: update lora name in lora list - if is_lora_v2 and original_lora_name and widget_index == 2: - # For LoraManager, we need to update the name in the lora list (widgets_values[2]) - lora_list = widgets_values[2] - log.info( - f">>> LoraManager: original_lora={original_lora_name}, widget_index={widget_index}, list_type={type(lora_list)}" - ) - - if isinstance(lora_list, list): - # Get the new lora name from resolved_model - new_lora_name = None - if resolved_model: - # Use filename without extension - new_lora_name = resolved_model.get("filename") or resolved_model.get( - "name", "" - ) - if new_lora_name and "." in new_lora_name: - # Remove extension if present - new_lora_name = new_lora_name.rsplit(".", 1)[0] - - log.debug( - f"LoraManager update: original={original_lora_name}, new={new_lora_name}" - ) - - if new_lora_name: - updated = False - original_stripped = original_lora_name.strip() - - for lora_item in lora_list: - if isinstance(lora_item, dict): - lora_name_in_list = lora_item.get("name", "").strip() - # Try exact match first, then case-insensitive - if lora_name_in_list == original_stripped: - lora_item["name"] = new_lora_name - updated = True - log.info( - f"Updated LoraManager lora (exact): {original_lora_name} -> {new_lora_name}" - ) - break - elif lora_name_in_list.lower() == original_stripped.lower(): - lora_item["name"] = new_lora_name - updated = True - log.info( - f"Updated LoraManager lora (case-insensitive): {original_lora_name} -> {new_lora_name}" - ) - break - - if updated: - # Also update the text widget (widgets_values[1]) which contains formatted lora string - if len(widgets_values) > 1 and isinstance(widgets_values[1], str): - old_text = widgets_values[1] - # Replace the lora name in the text format - new_text = old_text.replace( - f" { - const CALLBACK_DRIVEN_MODEL_WIDGETS = { - LoraLoaderV2: new Set(['text', 'loras']), - 'Lora Loader (LoraManager)': new Set(['text', 'loras']), - 'Lora Stacker (LoraManager)': new Set(['text', 'loras']), - 'Power Lora Loader (rgthree)': new Set(), +test('custom node registry exposes normalized LoRA Manager entries', () => { + const node = { + comfyClass: 'Lora Loader (LoraManager)', + widgets: [{ + name: 'loras', + value: [ + 'plain.safetensors', + { name: 'weighted.safetensors', strength: '0.65', active: false }, + ], + }], }; - const configureCallbackDrivenModelWidgets = eval( - `(${extractMethod(modelResolverSource, 'configureCallbackDrivenModelWidgets')})` + + assert.equal(getCustomNodeModelAdapter(node)?.id, 'lora-manager'); + assert.deepEqual(getCustomNodeModelEntries(node), [ + { + identity: 'plain.safetensors', + active: true, + strength: null, + }, + { + identity: 'weighted.safetensors', + active: false, + strength: 0.65, + }, + ]); +}); + +test('LoRA Manager adapter refreshes analysis only after its model list changes', () => { + const configureCustomNodeModelAdapter = eval( + `(${extractMethod(modelResolverSource, 'configureCustomNodeModelAdapter')})` ); - const getCallbackDrivenModelListSignature = eval( - `(${extractMethod(modelResolverSource, 'getCallbackDrivenModelListSignature')})` + const getCustomNodeModelListSignature = eval( + `(${extractMethod(modelResolverSource, 'getCustomNodeModelListSignature')})` ); - const getCallbackDrivenModelStrengthSignature = eval( - `(${extractMethod(modelResolverSource, 'getCallbackDrivenModelStrengthSignature')})` + const getCustomNodeModelStrengthSignature = eval( + `(${extractMethod(modelResolverSource, 'getCustomNodeModelStrengthSignature')})` ); const calls = []; const lorasWidget = { @@ -1210,9 +1240,9 @@ test('LoRA Manager callback-driven widgets refresh the open workflow analysis', }, }; const resolver = { - callbackDrivenModelWidgetStates: new WeakMap(), - getCallbackDrivenModelListSignature, - getCallbackDrivenModelStrengthSignature, + customNodeModelAdapterStates: new WeakMap(), + getCustomNodeModelListSignature, + getCustomNodeModelStrengthSignature, scheduleNodeContextMenuAnalysis: () => calls.push(['context-analysis']), dialog: { isWorkflowRefreshSuppressed: () => false, @@ -1224,9 +1254,9 @@ test('LoRA Manager callback-driven widgets refresh the open workflow analysis', widgets: [textWidget, lorasWidget, unrelatedWidget], }; - configureCallbackDrivenModelWidgets.call(resolver, node); + configureCustomNodeModelAdapter.call(resolver, node); lorasWidget.value = [{ name: 'existing', strength: 1, active: true }]; - configureCallbackDrivenModelWidgets.call(resolver, node); + configureCustomNodeModelAdapter.call(resolver, node); textWidget.callback(' { - const CALLBACK_DRIVEN_MODEL_WIDGETS = { - LoraLoaderV2: new Set(['text', 'loras']), - 'Lora Loader (LoraManager)': new Set(['text', 'loras']), - 'Lora Stacker (LoraManager)': new Set(['text', 'loras']), - 'Power Lora Loader (rgthree)': new Set(), - }; - const configureCallbackDrivenModelWidgets = eval( - `(${extractMethod(modelResolverSource, 'configureCallbackDrivenModelWidgets')})` +test('custom node adapter refresh respects internal workflow update suppression', () => { + const configureCustomNodeModelAdapter = eval( + `(${extractMethod(modelResolverSource, 'configureCustomNodeModelAdapter')})` ); - const getCallbackDrivenModelListSignature = eval( - `(${extractMethod(modelResolverSource, 'getCallbackDrivenModelListSignature')})` + const getCustomNodeModelListSignature = eval( + `(${extractMethod(modelResolverSource, 'getCustomNodeModelListSignature')})` ); - const getCallbackDrivenModelStrengthSignature = eval( - `(${extractMethod(modelResolverSource, 'getCallbackDrivenModelStrengthSignature')})` + const getCustomNodeModelStrengthSignature = eval( + `(${extractMethod(modelResolverSource, 'getCustomNodeModelStrengthSignature')})` ); const calls = []; const widget = { @@ -1266,9 +1290,9 @@ test('callback-driven widget refresh respects internal workflow update suppressi callback: () => calls.push('original'), }; const resolver = { - callbackDrivenModelWidgetStates: new WeakMap(), - getCallbackDrivenModelListSignature, - getCallbackDrivenModelStrengthSignature, + customNodeModelAdapterStates: new WeakMap(), + getCustomNodeModelListSignature, + getCustomNodeModelStrengthSignature, scheduleNodeContextMenuAnalysis: () => calls.push('context-analysis'), dialog: { isWorkflowRefreshSuppressed: () => true, @@ -1276,7 +1300,7 @@ test('callback-driven widget refresh respects internal workflow update suppressi }, }; - configureCallbackDrivenModelWidgets.call(resolver, { + configureCustomNodeModelAdapter.call(resolver, { comfyClass: 'Lora Stacker (LoraManager)', widgets: [widget], }); @@ -1287,20 +1311,14 @@ test('callback-driven widget refresh respects internal workflow update suppressi }); test('rgthree Power Lora Loader updates strength without reloading workflow analysis', async () => { - const CALLBACK_DRIVEN_MODEL_WIDGETS = { - LoraLoaderV2: new Set(['text', 'loras']), - 'Lora Loader (LoraManager)': new Set(['text', 'loras']), - 'Lora Stacker (LoraManager)': new Set(['text', 'loras']), - 'Power Lora Loader (rgthree)': new Set(), - }; - const configureCallbackDrivenModelWidgets = eval( - `(${extractMethod(modelResolverSource, 'configureCallbackDrivenModelWidgets')})` + const configureCustomNodeModelAdapter = eval( + `(${extractMethod(modelResolverSource, 'configureCustomNodeModelAdapter')})` ); - const getCallbackDrivenModelListSignature = eval( - `(${extractMethod(modelResolverSource, 'getCallbackDrivenModelListSignature')})` + const getCustomNodeModelListSignature = eval( + `(${extractMethod(modelResolverSource, 'getCustomNodeModelListSignature')})` ); - const getCallbackDrivenModelStrengthSignature = eval( - `(${extractMethod(modelResolverSource, 'getCallbackDrivenModelStrengthSignature')})` + const getCustomNodeModelStrengthSignature = eval( + `(${extractMethod(modelResolverSource, 'getCustomNodeModelStrengthSignature')})` ); const calls = []; const node = { @@ -1318,9 +1336,9 @@ test('rgthree Power Lora Loader updates strength without reloading workflow anal }, }; const resolver = { - callbackDrivenModelWidgetStates: new WeakMap(), - getCallbackDrivenModelListSignature, - getCallbackDrivenModelStrengthSignature, + customNodeModelAdapterStates: new WeakMap(), + getCustomNodeModelListSignature, + getCustomNodeModelStrengthSignature, scheduleNodeContextMenuAnalysis: () => calls.push('context-analysis'), dialog: { isWorkflowRefreshSuppressed: () => false, @@ -1329,7 +1347,7 @@ test('rgthree Power Lora Loader updates strength without reloading workflow anal }, }; - configureCallbackDrivenModelWidgets.call(resolver, node); + configureCustomNodeModelAdapter.call(resolver, node); node.widgets[0].value.strength = 0.5; node.setDirtyCanvas(true, true); @@ -1356,17 +1374,11 @@ test('rgthree Power Lora Loader updates strength without reloading workflow anal }); test('rgthree dynamic lora widgets bypass the generic full workflow refresh', () => { - const CALLBACK_DRIVEN_MODEL_WIDGETS = { - LoraLoaderV2: new Set(['text', 'loras']), - 'Lora Loader (LoraManager)': new Set(['text', 'loras']), - 'Lora Stacker (LoraManager)': new Set(['text', 'loras']), - 'Power Lora Loader (rgthree)': new Set(), - }; const configureNodeContextMenu = eval( `(${extractMethod(modelResolverSource, 'configureNodeContextMenu')})` ); - const isCallbackDrivenModelWidget = eval( - `(${extractMethod(modelResolverSource, 'isCallbackDrivenModelWidget')})` + const isCustomNodeModelWidget = eval( + `(${extractMethod(modelResolverSource, 'isCustomNodeModelWidget')})` ); const calls = []; class NodeType {} @@ -1376,11 +1388,11 @@ test('rgthree dynamic lora widgets bypass the generic full workflow refresh', () }; const node = new NodeType(); const resolver = { - callbackDrivenModelWidgetStates: new WeakMap([ + customNodeModelAdapterStates: new WeakMap([ [node, { notify: () => calls.push('selective-update') }], ]), getResolvedModelsForNodeContextMenu: () => [], - isCallbackDrivenModelWidget, + isCustomNodeModelWidget, scheduleNodeContextMenuAnalysis: () => calls.push('context-analysis'), dialog: { isWorkflowRefreshSuppressed: () => false, @@ -1407,7 +1419,7 @@ test('LoRA Manager text onWidgetChanged waits for a confirmed model-list change' }; const resolver = { getResolvedModelsForNodeContextMenu: () => [], - isCallbackDrivenModelWidget: (_node, widgetName) => widgetName === 'text', + isCustomNodeModelWidget: (_node, widgetName) => widgetName === 'text', scheduleNodeContextMenuAnalysis: () => calls.push('context-analysis'), dialog: { isWorkflowRefreshSuppressed: () => false, diff --git a/tests/test_resolver_robustness.py b/tests/test_resolver_robustness.py index 11539a1..500aa7c 100644 --- a/tests/test_resolver_robustness.py +++ b/tests/test_resolver_robustness.py @@ -143,8 +143,8 @@ class ModelResolverRobustnessTests(unittest.TestCase): ] } mapping = { - "is_lora_v2": True, - "original_lora_name": "OLD_LORA_NAME" + "custom_node_adapter": "lora-manager", + "custom_node_original_identity": "OLD_LORA_NAME", } resolved_model = { "filename": "new_lora_name.safetensors", diff --git a/tests/test_workflow_analyzer.py b/tests/test_workflow_analyzer.py index 6794775..b3dba12 100644 --- a/tests/test_workflow_analyzer.py +++ b/tests/test_workflow_analyzer.py @@ -1141,6 +1141,37 @@ class WorkflowLoraStrengthTests(unittest.TestCase): self.assertEqual(0.45, strength) +class WorkflowCustomNodeAdapterTests(unittest.TestCase): + def test_lora_manager_list_is_extracted_by_adapter(self): + workflow = { + "nodes": [ + { + "id": 390, + "type": "Lora Loader (LoraManager)", + "widgets_values": [ + {"version": 1, "textWidgetName": "text"}, + "", + [ + { + "name": "missing_style", + "strength": 0.65, + "active": True, + } + ], + ], + } + ] + } + + refs = analyze_workflow_models(workflow, available_models=[]) + + self.assertEqual(1, len(refs)) + self.assertEqual("lora-manager", refs[0]["custom_node_adapter"]) + self.assertEqual("missing_style", refs[0]["original_path"]) + self.assertEqual(0.65, refs[0]["strength"]) + self.assertTrue(refs[0]["active"]) + + class WorkflowMissingReferenceGroupingTests(unittest.TestCase): def test_power_lora_duplicate_missing_tracks_all_node_refs(self): workflow = { @@ -1166,6 +1197,10 @@ class WorkflowMissingReferenceGroupingTests(unittest.TestCase): missing = identify_missing_models(refs, []) self.assertEqual(2, len(refs)) + self.assertEqual( + ["rgthree-power-lora-loader", "rgthree-power-lora-loader"], + [ref["custom_node_adapter"] for ref in refs], + ) self.assertEqual(1, len(missing)) self.assertEqual(2, missing[0]["reference_count"]) self.assertEqual( diff --git a/web/resolver.js b/web/resolver.js index c5d0826..62647f9 100644 --- a/web/resolver.js +++ b/web/resolver.js @@ -35,12 +35,12 @@ app.registerExtension({ }, nodeCreated(node) { modelResolver.configureWorkflowDependencyMarkerNode(node); - modelResolver.configureCallbackDrivenModelWidgets(node); + modelResolver.configureCustomNodeModelAdapter(node); modelResolver.scheduleNodeContextMenuAnalysis(); }, loadedGraphNode(node) { modelResolver.configureWorkflowDependencyMarkerNode(node); - modelResolver.configureCallbackDrivenModelWidgets(node); + modelResolver.configureCustomNodeModelAdapter(node); modelResolver.scheduleNodeContextMenuAnalysis(); }, afterConfigureGraph() { diff --git a/web/resolver/actions/queue_methods.js b/web/resolver/actions/queue_methods.js index d6c9157..9b92e00 100644 --- a/web/resolver/actions/queue_methods.js +++ b/web/resolver/actions/queue_methods.js @@ -1672,7 +1672,12 @@ export const queueMethods = { subgraph_id: ref.subgraph_id, subgraph_name: ref.subgraph_name, is_top_level: ref.is_top_level, - is_lora_v2: ref.is_lora_v2, + custom_node_adapter: ref.custom_node_adapter, + custom_node_original_identity: ( + ref.custom_node_original_identity + || ref.name + || ref.original_path + ), nested_key: ref.nested_key })); }, @@ -1695,8 +1700,18 @@ export const queueMethods = { category: ref.category || baseResolution.category, subgraph_id: ref.subgraph_id ?? baseResolution.subgraph_id, is_top_level: ref.is_top_level ?? baseResolution.is_top_level, - is_lora_v2: ref.is_lora_v2 ?? baseResolution.is_lora_v2, - original_lora_name: ref.name || ref.original_path || baseResolution.original_lora_name || baseResolution.name || baseResolution.original_path, + custom_node_adapter: ( + ref.custom_node_adapter + || baseResolution.custom_node_adapter + ), + custom_node_original_identity: ( + ref.custom_node_original_identity + || ref.name + || ref.original_path + || baseResolution.custom_node_original_identity + || baseResolution.name + || baseResolution.original_path + ), nested_key: ref.nested_key ?? baseResolution.nested_key, original_path: ref.original_path || baseResolution.original_path, node_type: ref.node_type || baseResolution.node_type, @@ -1708,7 +1723,7 @@ export const queueMethods = { item.subgraph_id || '', item.is_top_level !== false ? 'T' : 'F', item.nested_key || '', - item.original_lora_name || item.original_path || '' + item.custom_node_original_identity || item.original_path || '' ].join(':'); if (seen.has(itemKey)) continue; seen.add(itemKey); @@ -2155,8 +2170,12 @@ export const queueMethods = { urn_string: missing.urn_string, subgraph_id: missing.subgraph_id, is_top_level: missing.is_top_level, - is_lora_v2: missing.is_lora_v2, - original_lora_name: missing.original_lora_name || missing.name || missing.original_path, + custom_node_adapter: missing.custom_node_adapter, + custom_node_original_identity: ( + missing.custom_node_original_identity + || missing.name + || missing.original_path + ), nested_key: missing.nested_key, node_refs: this.getResolutionNodeRefs(missing), node_type: missing.node_type, diff --git a/web/resolver/custom_nodes/README.md b/web/resolver/custom_nodes/README.md new file mode 100644 index 0000000..20a53a4 --- /dev/null +++ b/web/resolver/custom_nodes/README.md @@ -0,0 +1,22 @@ +# Custom node model adapters + +This directory contains frontend integrations for custom nodes whose model +widgets do not follow ComfyUI's standard widget behavior. + +Each adapter exports: + +- `id`: stable integration identifier. +- `nodeTypes`: supported ComfyUI node type names. +- `category`: default model category for download targeting. +- `isModelWidget(widgetName)`: identifies widget changes handled by the adapter. +- `getModelEntries(node)`: returns normalized source data with `identity`, + `active`, and `strength`. +- `observe(node, notify)`: attaches any node-specific change hooks and calls + `notify` after a relevant mutation. + +Register new adapters in `registry.js`. The registry owns entry normalization +and model-list/strength signatures. Generic workflow refresh, cache, and Loaded +Models rendering stay outside adapters. + +Legacy custom-node metadata normalization also stays in the registry so generic +views do not need package-specific field knowledge. diff --git a/web/resolver/custom_nodes/lora_manager.js b/web/resolver/custom_nodes/lora_manager.js new file mode 100644 index 0000000..5b9cf11 --- /dev/null +++ b/web/resolver/custom_nodes/lora_manager.js @@ -0,0 +1,51 @@ +const LORA_MANAGER_NODE_TYPES = Object.freeze([ + 'LoraLoaderV2', + 'Lora Loader (LoraManager)', + 'Lora Stacker (LoraManager)', +]); +const LORA_MANAGER_WIDGET_NAMES = new Set(['text', 'loras']); + +function getLoraManagerEntries(node) { + const widgets = Array.isArray(node?.widgets) ? node.widgets : []; + const lorasWidget = widgets.find(widget => widget?.name === 'loras'); + const loras = Array.isArray(lorasWidget?.value) ? lorasWidget.value : []; + return loras.map((lora) => { + if (lora && typeof lora === 'object') { + return { + identity: lora.name || lora.filename || lora.path || '', + active: lora.active !== false, + strength: lora.strength, + }; + } + return { + identity: lora, + active: true, + strength: null, + }; + }); +} + +function observeLoraManager(node, notify) { + for (const widget of node.widgets || []) { + if (!LORA_MANAGER_WIDGET_NAMES.has(widget?.name)) continue; + if (widget.callback?.__modelResolverCustomNodeAdapter) continue; + + const originalCallback = widget.callback; + const wrappedCallback = function() { + const result = originalCallback?.apply(this, arguments); + notify(); + return result; + }; + wrappedCallback.__modelResolverCustomNodeAdapter = true; + widget.callback = wrappedCallback; + } +} + +export const loraManagerAdapter = Object.freeze({ + id: 'lora-manager', + nodeTypes: LORA_MANAGER_NODE_TYPES, + category: 'loras', + isModelWidget: widgetName => LORA_MANAGER_WIDGET_NAMES.has(widgetName), + getModelEntries: getLoraManagerEntries, + observe: observeLoraManager, +}); diff --git a/web/resolver/custom_nodes/registry.js b/web/resolver/custom_nodes/registry.js new file mode 100644 index 0000000..8c2b65e --- /dev/null +++ b/web/resolver/custom_nodes/registry.js @@ -0,0 +1,96 @@ +import { loraManagerAdapter } from './lora_manager.js'; +import { rgthreePowerLoraLoaderAdapter } from './rgthree_power_lora_loader.js'; + +export const CUSTOM_NODE_MODEL_ADAPTERS = Object.freeze([ + loraManagerAdapter, + rgthreePowerLoraLoaderAdapter, +]); + +function getNodeTypeCandidates(node) { + return [ + node?.comfyClass, + node?.type, + node?.constructor?.comfyClass, + node?.constructor?.ComfyClass, + ].filter(Boolean); +} + +function normalizeModelEntry(entry) { + if (!entry || typeof entry !== 'object') return null; + const identity = String(entry.identity || '').trim(); + if (!identity) return null; + const hasStrength = ( + entry.strength !== null + && entry.strength !== undefined + && entry.strength !== '' + ); + const numericStrength = hasStrength ? Number(entry.strength) : Number.NaN; + return { + identity, + active: entry.active !== false, + strength: Number.isFinite(numericStrength) ? numericStrength : null, + }; +} + +export function getCustomNodeModelAdapter(node) { + const nodeTypes = getNodeTypeCandidates(node); + return CUSTOM_NODE_MODEL_ADAPTERS.find( + adapter => adapter.nodeTypes.some(nodeType => nodeTypes.includes(nodeType)) + ) || null; +} + +export function getCustomNodeModelCategory(nodeType = '') { + const normalizedNodeType = String(nodeType || '').trim(); + const adapter = CUSTOM_NODE_MODEL_ADAPTERS.find( + candidate => candidate.nodeTypes.includes(normalizedNodeType) + ); + return adapter?.category || ''; +} + +export function getCustomNodeOriginalIdentity(model = {}) { + return ( + model.custom_node_original_identity + || model.original_lora_name + || '' + ); +} + +export function getCustomNodeModelEntries(node) { + const adapter = getCustomNodeModelAdapter(node); + if (!adapter) return []; + return (adapter.getModelEntries(node) || []) + .map(normalizeModelEntry) + .filter(Boolean); +} + +export function isCustomNodeModelWidget(node, widgetName = '') { + return Boolean( + getCustomNodeModelAdapter(node)?.isModelWidget?.(widgetName) + ); +} + +export function getCustomNodeModelListSignature(node) { + const entries = getCustomNodeModelEntries(node) + .map(entry => ({ + identity: entry.identity, + active: entry.active, + })) + .sort((left, right) => ( + left.identity.localeCompare(right.identity) + || Number(left.active) - Number(right.active) + )); + return JSON.stringify(entries); +} + +export function getCustomNodeModelStrengthSignature(node) { + const entries = getCustomNodeModelEntries(node) + .map(entry => ({ + identity: entry.identity, + strength: entry.strength, + })) + .sort((left, right) => ( + left.identity.localeCompare(right.identity) + || (left.strength ?? 0) - (right.strength ?? 0) + )); + return JSON.stringify(entries); +} diff --git a/web/resolver/custom_nodes/rgthree_power_lora_loader.js b/web/resolver/custom_nodes/rgthree_power_lora_loader.js new file mode 100644 index 0000000..764c2c4 --- /dev/null +++ b/web/resolver/custom_nodes/rgthree_power_lora_loader.js @@ -0,0 +1,50 @@ +const POWER_LORA_LOADER_NODE_TYPE = 'Power Lora Loader (rgthree)'; +const POWER_LORA_WIDGET_PATTERN = /^lora_\d+$/; + +function getPowerLoraEntries(node) { + const widgets = Array.isArray(node?.widgets) ? node.widgets : []; + return widgets + .filter(widget => ( + POWER_LORA_WIDGET_PATTERN.test(widget?.name) + && widget?.value?.lora + )) + .map(widget => ({ + identity: widget.value.lora, + active: widget.value.on !== false, + strength: widget.value.strength, + })); +} + +function observePowerLoraLoader(node, notify) { + if ( + typeof node.setDirtyCanvas !== 'function' + || node.setDirtyCanvas.__modelResolverCustomNodeAdapter + ) { + return; + } + + const originalSetDirtyCanvas = node.setDirtyCanvas; + let inspectionQueued = false; + const wrappedSetDirtyCanvas = function() { + const result = originalSetDirtyCanvas.apply(this, arguments); + if (!inspectionQueued) { + inspectionQueued = true; + queueMicrotask(() => { + inspectionQueued = false; + notify(); + }); + } + return result; + }; + wrappedSetDirtyCanvas.__modelResolverCustomNodeAdapter = true; + node.setDirtyCanvas = wrappedSetDirtyCanvas; +} + +export const rgthreePowerLoraLoaderAdapter = Object.freeze({ + id: 'rgthree-power-lora-loader', + nodeTypes: Object.freeze([POWER_LORA_LOADER_NODE_TYPE]), + category: 'loras', + isModelWidget: widgetName => POWER_LORA_WIDGET_PATTERN.test(widgetName), + getModelEntries: getPowerLoraEntries, + observe: observePowerLoraLoader, +}); diff --git a/web/resolver/model_resolver.js b/web/resolver/model_resolver.js index e200341..9c20b67 100644 --- a/web/resolver/model_resolver.js +++ b/web/resolver/model_resolver.js @@ -14,6 +14,12 @@ import { matchesWorkflowModelReference, toResolverContextModel, } from "./node_context_menu.js"; +import { + getCustomNodeModelAdapter, + getCustomNodeModelListSignature as createCustomNodeModelListSignature, + getCustomNodeModelStrengthSignature as createCustomNodeModelStrengthSignature, + isCustomNodeModelWidget as matchesCustomNodeModelWidget, +} from "./custom_nodes/registry.js"; const log = createModuleLogger('model_resolver'); export const MODEL_RESOLVER_OPEN_COMMAND_ID = "ModelResolver.OpenModelResolver"; @@ -36,12 +42,6 @@ const WORKFLOW_DEPENDENCY_MARKER_BUTTON_NAME = "Open Model Resolver"; const WORKFLOW_DEPENDENCY_MARKER_AUTO_ID = -918273646; const WORKFLOW_DEPENDENCY_MARKER_DEFAULT_SIZE = Object.freeze([160, 40]); const WORKFLOW_DEPENDENCY_MARKER_MIN_SIZE = Object.freeze([160, 40]); -const CALLBACK_DRIVEN_MODEL_WIDGETS = Object.freeze({ - LoraLoaderV2: new Set(["text", "loras"]), - "Lora Loader (LoraManager)": new Set(["text", "loras"]), - "Lora Stacker (LoraManager)": new Set(["text", "loras"]), - "Power Lora Loader (rgthree)": new Set(), -}); function getKeybindingSetting(id) { try { @@ -135,7 +135,7 @@ export class ModelResolver { this.nodeContextAnalysisSignature = null; this.nodeContextAnalysisPromise = null; this.nodeContextAnalysisTimer = null; - this.callbackDrivenModelWidgetStates = new WeakMap(); + this.customNodeModelAdapterStates = new WeakMap(); } setup = async () => { @@ -530,12 +530,12 @@ export class ModelResolver { nodeType.prototype.onWidgetChanged = function() { const result = originalOnWidgetChanged?.apply(this, arguments); const widgetName = arguments[0]; - const handledByCallbackDrivenWidget = owner.isCallbackDrivenModelWidget?.( + const handledByCustomNodeAdapter = owner.isCustomNodeModelWidget?.( this, widgetName ); - if (handledByCallbackDrivenWidget) { - owner.callbackDrivenModelWidgetStates?.get(this)?.notify?.(); + if (handledByCustomNodeAdapter) { + owner.customNodeModelAdapterStates?.get(this)?.notify?.(); } else if (!owner.dialog?.isWorkflowRefreshSuppressed?.()) { if (!owner.dialog?.isWorkflowStrengthWidgetName?.(widgetName)) { owner.scheduleNodeContextMenuAnalysis(); @@ -547,133 +547,41 @@ export class ModelResolver { nodeType.prototype.__modelResolverContextMenuPatched = true; } - isCallbackDrivenModelWidget(node, widgetName = '') { - const nodeTypes = [ - node?.comfyClass, - node?.type, - node?.constructor?.comfyClass, - node?.constructor?.ComfyClass, - ]; - const nodeType = nodeTypes.find(candidate => CALLBACK_DRIVEN_MODEL_WIDGETS[candidate]); - if ( - nodeType === 'Power Lora Loader (rgthree)' - && /^lora_\d+$/.test(widgetName) - ) { - return true; - } - return Boolean(CALLBACK_DRIVEN_MODEL_WIDGETS[nodeType]?.has(widgetName)); + isCustomNodeModelWidget(node, widgetName = '') { + return matchesCustomNodeModelWidget(node, widgetName); } - getCallbackDrivenModelListSignature(node) { - const widgets = Array.isArray(node?.widgets) ? node.widgets : []; - const isRgthreePowerLoraLoader = [ - node?.comfyClass, - node?.type, - node?.constructor?.comfyClass, - node?.constructor?.ComfyClass, - ].includes('Power Lora Loader (rgthree)'); - const lorasWidget = widgets.find(widget => widget?.name === 'loras'); - const loras = isRgthreePowerLoraLoader - ? widgets - .filter(widget => /^lora_\d+$/.test(widget?.name) && widget?.value?.lora) - .map(widget => ({ - name: widget.value.lora, - active: widget.value.on !== false, - })) - : (Array.isArray(lorasWidget?.value) ? lorasWidget.value : []); - const normalized = loras - .map((lora) => { - if (lora && typeof lora === 'object') { - const name = String( - lora.name - || lora.filename - || lora.path - || '' - ).trim(); - return name ? { name, active: lora.active !== false } : null; - } - const name = String(lora || '').trim(); - return name ? { name, active: true } : null; - }) - .filter(Boolean) - .sort((left, right) => ( - left.name.localeCompare(right.name) - || Number(left.active) - Number(right.active) - )); - return JSON.stringify(normalized); + getCustomNodeModelListSignature(node) { + return createCustomNodeModelListSignature(node); } - getCallbackDrivenModelStrengthSignature(node) { - const widgets = Array.isArray(node?.widgets) ? node.widgets : []; - const isRgthreePowerLoraLoader = [ - node?.comfyClass, - node?.type, - node?.constructor?.comfyClass, - node?.constructor?.ComfyClass, - ].includes('Power Lora Loader (rgthree)'); - const lorasWidget = widgets.find(widget => widget?.name === 'loras'); - const loras = isRgthreePowerLoraLoader - ? widgets - .filter(widget => /^lora_\d+$/.test(widget?.name) && widget?.value?.lora) - .map(widget => ({ - name: widget.value.lora, - strength: widget.value.strength, - })) - : (Array.isArray(lorasWidget?.value) ? lorasWidget.value : []); - const normalized = loras - .map((lora) => { - if (!lora || typeof lora !== 'object') return null; - const name = String( - lora.name - || lora.filename - || lora.path - || '' - ).trim(); - if (!name) return null; - const numericStrength = Number(lora.strength); - return { - name, - strength: Number.isFinite(numericStrength) ? numericStrength : null, - }; - }) - .filter(Boolean) - .sort((left, right) => ( - left.name.localeCompare(right.name) - || (left.strength ?? 0) - (right.strength ?? 0) - )); - return JSON.stringify(normalized); + getCustomNodeModelStrengthSignature(node) { + return createCustomNodeModelStrengthSignature(node); } - configureCallbackDrivenModelWidgets(node) { - const nodeTypes = [ - node?.comfyClass, - node?.type, - node?.constructor?.comfyClass, - node?.constructor?.ComfyClass, - ]; - const nodeType = nodeTypes.find(candidate => CALLBACK_DRIVEN_MODEL_WIDGETS[candidate]); - const watchedWidgetNames = CALLBACK_DRIVEN_MODEL_WIDGETS[nodeType]; - if (!watchedWidgetNames || !Array.isArray(node?.widgets)) return; + configureCustomNodeModelAdapter(node) { + const adapter = getCustomNodeModelAdapter(node); + if (!adapter || !Array.isArray(node?.widgets)) return; const owner = this; - if (!(this.callbackDrivenModelWidgetStates instanceof WeakMap)) { - this.callbackDrivenModelWidgetStates = new WeakMap(); + if (!(this.customNodeModelAdapterStates instanceof WeakMap)) { + this.customNodeModelAdapterStates = new WeakMap(); } - let state = this.callbackDrivenModelWidgetStates.get(node); + let state = this.customNodeModelAdapterStates.get(node); if (!state) { state = { - signature: this.getCallbackDrivenModelListSignature(node), - strengthSignature: this.getCallbackDrivenModelStrengthSignature(node), + signature: this.getCustomNodeModelListSignature(node), + strengthSignature: this.getCustomNodeModelStrengthSignature(node), }; - this.callbackDrivenModelWidgetStates.set(node, state); + this.customNodeModelAdapterStates.set(node, state); } else { - state.signature = this.getCallbackDrivenModelListSignature(node); - state.strengthSignature = this.getCallbackDrivenModelStrengthSignature(node); + state.signature = this.getCustomNodeModelListSignature(node); + state.strengthSignature = this.getCustomNodeModelStrengthSignature(node); } - const notifyIfModelListChanged = () => { - const nextSignature = owner.getCallbackDrivenModelListSignature(node); - const nextStrengthSignature = owner.getCallbackDrivenModelStrengthSignature(node); + const notifyIfModelsChanged = () => { + const nextSignature = owner.getCustomNodeModelListSignature(node); + const nextStrengthSignature = owner.getCustomNodeModelStrengthSignature(node); const modelListChanged = nextSignature !== state.signature; const strengthChanged = nextStrengthSignature !== state.strengthSignature; state.signature = nextSignature; @@ -689,43 +597,8 @@ export class ModelResolver { owner.dialog?.updateLoadedModelStrengthsFromNode?.(node); } }; - state.notify = notifyIfModelListChanged; - - for (const widget of node.widgets) { - if (!watchedWidgetNames.has(widget?.name)) continue; - if (widget.callback?.__modelResolverCallbackDrivenWidget) continue; - - const originalCallback = widget.callback; - const wrappedCallback = function() { - const result = originalCallback?.apply(this, arguments); - notifyIfModelListChanged(); - return result; - }; - wrappedCallback.__modelResolverCallbackDrivenWidget = true; - widget.callback = wrappedCallback; - } - - if ( - nodeType === 'Power Lora Loader (rgthree)' - && typeof node.setDirtyCanvas === 'function' - && !node.setDirtyCanvas.__modelResolverCallbackDrivenNode - ) { - const originalSetDirtyCanvas = node.setDirtyCanvas; - let inspectionQueued = false; - const wrappedSetDirtyCanvas = function() { - const result = originalSetDirtyCanvas.apply(this, arguments); - if (!inspectionQueued) { - inspectionQueued = true; - queueMicrotask(() => { - inspectionQueued = false; - notifyIfModelListChanged(); - }); - } - return result; - }; - wrappedSetDirtyCanvas.__modelResolverCallbackDrivenNode = true; - node.setDirtyCanvas = wrappedSetDirtyCanvas; - } + state.notify = notifyIfModelsChanged; + adapter.observe(node, notifyIfModelsChanged); } waitForResolverDialogReady(timeoutMs = 2500) { diff --git a/web/resolver/search/download_target_methods.js b/web/resolver/search/download_target_methods.js index ad4b71b..c9d63e4 100644 --- a/web/resolver/search/download_target_methods.js +++ b/web/resolver/search/download_target_methods.js @@ -2,6 +2,7 @@ import { app } from "../../../../../scripts/app.js"; import { api } from "../../../../../scripts/api.js"; import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; +import { getCustomNodeModelCategory } from "../custom_nodes/registry.js"; import { createFloatingTreePicker } from "../utils/tree_picker.js"; import { safeStorage } from "../utils/html_utils.js"; const localStorage = safeStorage; @@ -293,6 +294,10 @@ export const downloadTargetMethods = { } } } + const customNodeCategory = getCustomNodeModelCategory(normalizedNodeType); + if (customNodeCategory) { + return customNodeCategory; + } const exactMap = { CheckpointLoaderSimple: 'checkpoints', CheckpointLoader: 'checkpoints', @@ -308,7 +313,6 @@ export const downloadTargetMethods = { LoraLoaderModelOnly: 'loras', LoraLoaderBypass: 'loras', LoraLoaderBypassModelOnly: 'loras', - LoraLoaderV2: 'loras', ControlNetLoader: 'controlnet', CLIPLoader: 'text_encoders', DualCLIPLoader: 'text_encoders', diff --git a/web/resolver/shell/lifecycle_graph_methods.js b/web/resolver/shell/lifecycle_graph_methods.js index a683a2b..1775f5c 100644 --- a/web/resolver/shell/lifecycle_graph_methods.js +++ b/web/resolver/shell/lifecycle_graph_methods.js @@ -2,6 +2,7 @@ import { app } from "../../../../../scripts/app.js"; import { api } from "../../../../../scripts/api.js"; import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; +import { getCustomNodeOriginalIdentity } from "../custom_nodes/registry.js"; import { safeStorage } from "../utils/html_utils.js"; export const lifecycleGraphMethods = { async show(workflow = null) { @@ -798,7 +799,7 @@ export const lifecycleGraphMethods = { }, getMissingModelIdentityPart(missing = {}) { - return missing.original_lora_name + return getCustomNodeOriginalIdentity(missing) || missing.original_path || missing.expected_filename || missing.name diff --git a/web/resolver/views/tabs_loaded_methods.js b/web/resolver/views/tabs_loaded_methods.js index 038a8d0..f365700 100644 --- a/web/resolver/views/tabs_loaded_methods.js +++ b/web/resolver/views/tabs_loaded_methods.js @@ -3,6 +3,10 @@ import { api } from "../../../../../scripts/api.js"; import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; import { escapeHtml } from "../utils/html_utils.js"; +import { + getCustomNodeModelEntries, + getCustomNodeOriginalIdentity, +} from "../custom_nodes/registry.js"; export const tabsLoadedMethods = { escapeHtml(value) { return escapeHtml(value); @@ -288,31 +292,12 @@ export const tabsLoadedMethods = { return false; } - const widgets = Array.isArray(node?.widgets) ? node.widgets : []; - const dynamicLoras = widgets - .filter(widget => /^lora_\d+$/.test(widget?.name) && widget?.value?.lora) - .map(widget => widget.value); - const lorasWidget = widgets.find(widget => widget?.name === 'loras'); - const loras = dynamicLoras.length - ? dynamicLoras - : (Array.isArray(lorasWidget?.value) ? lorasWidget.value : []); - const entries = loras - .map((lora) => { - if (!lora || typeof lora !== 'object') return null; - const identity = lora.lora - || lora.name - || lora.filename - || lora.path - || ''; - const strength = Number(lora.strength); - return identity && Number.isFinite(strength) - ? { - identity: this.normalizeLoadedModelIdentity(identity), - strength, - } - : null; - }) - .filter(Boolean); + const entries = getCustomNodeModelEntries(node) + .filter(entry => Number.isFinite(entry.strength)) + .map(entry => ({ + identity: this.normalizeLoadedModelIdentity(entry.identity), + strength: entry.strength, + })); if (!entries.length) return false; const nodeId = String(node?.id ?? ''); @@ -320,7 +305,7 @@ export const tabsLoadedMethods = { for (const model of this.cachedLoadedModelsData.loaded_models) { if (String(model.node_id ?? '') !== nodeId) continue; const identity = this.normalizeLoadedModelIdentity( - model.original_lora_name + getCustomNodeOriginalIdentity(model) || model.original_path || model.name || model.filename @@ -386,8 +371,7 @@ export const tabsLoadedMethods = { byCategory[cat] = { active: [], inactive: [] }; } - // Some loaders expose per-model enablement, e.g. LoraManager's - // active flag and rgthree Power Lora Loader's on flag. + // Some loaders expose per-model enablement separately from links. const isActive = model.active !== false && model.connected !== false; if (isActive) {