diff --git a/py/metadata_collector/node_extractors.py b/py/metadata_collector/node_extractors.py index 5421c306..ab203132 100644 --- a/py/metadata_collector/node_extractors.py +++ b/py/metadata_collector/node_extractors.py @@ -2,7 +2,8 @@ import json import os import re -from .constants import CLIP_SKIP_SENTINEL, MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER, OVERWRITE, METADATA_OVERWRITE_FIELDS +from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER, OVERWRITE +from .overwrite_utils import collect_overwrite_params def _store_checkpoint_metadata(metadata, node_id, model_name): @@ -1233,14 +1234,7 @@ class MetadataOverwriteExtractor(NodeMetadataExtractor): if not inputs: return - overwrite_params = {} - for key in METADATA_OVERWRITE_FIELDS: - value = inputs.get(key) - if key == "clip_skip": - if value != CLIP_SKIP_SENTINEL: - overwrite_params[key] = value - elif value: # truthy — only overwrite when user provided a real value - overwrite_params[key] = value + overwrite_params = collect_overwrite_params(inputs) if overwrite_params: metadata.setdefault(OVERWRITE, {}) diff --git a/py/metadata_collector/overwrite_utils.py b/py/metadata_collector/overwrite_utils.py new file mode 100644 index 00000000..54d58ef3 --- /dev/null +++ b/py/metadata_collector/overwrite_utils.py @@ -0,0 +1,42 @@ +"""Shared helpers for Metadata Overwrite node metadata collection. + +Used by both the MetadataOverwriteLM node (execution time) and the +MetadataOverwriteExtractor (hook time) so the conversion/filtering logic +cannot drift between the two paths. +""" + +import logging +from typing import Any, Dict + +from ..utils.utils import model_patcher_to_name +from .constants import CLIP_SKIP_SENTINEL, METADATA_OVERWRITE_FIELDS + +logger = logging.getLogger(__name__) + + +def collect_overwrite_params(values: Dict[str, Any]) -> Dict[str, Any]: + """Convert node input values into non-default overwrite parameters. + + For most fields, a falsy value (empty string, 0) means "not set" and is + skipped. clip_skip uses a dedicated sentinel (-25) so that a wired value + of 0 is preserved. The ``model`` field accepts either a manual string or + a wired MODEL (ModelPatcher) connection; in the latter case the source + model name is extracted from the patcher's ``cached_patcher_init`` and + stored as a ComfyUI-style relative path. + """ + result: Dict[str, Any] = {} + for key in METADATA_OVERWRITE_FIELDS: + value = values.get(key) + if key == "model" and not isinstance(value, str): + value = model_patcher_to_name(value) + if value is None: + logger.warning( + "Could not extract model name from wired MODEL input " + "(no cached_patcher_init); model metadata overwrite skipped" + ) + if key == "clip_skip": + if value != CLIP_SKIP_SENTINEL: + result[key] = value + elif value: + result[key] = value + return result diff --git a/py/nodes/metadata_overwrite.py b/py/nodes/metadata_overwrite.py index c36613a5..337d8e92 100644 --- a/py/nodes/metadata_overwrite.py +++ b/py/nodes/metadata_overwrite.py @@ -9,10 +9,8 @@ but users may wire 0 to express "no clip skip / default". from typing import Any -from ..metadata_collector.constants import ( - CLIP_SKIP_SENTINEL as _CLIP_SKIP_SENTINEL, - METADATA_OVERWRITE_FIELDS, -) +from ..metadata_collector.constants import CLIP_SKIP_SENTINEL as _CLIP_SKIP_SENTINEL +from ..metadata_collector.overwrite_utils import collect_overwrite_params class MetadataOverwriteLM: @@ -87,12 +85,16 @@ class MetadataOverwriteLM: }, ), "model": ( - "STRING", + "STRING,MODEL", { "default": "", + "widgetType": "STRING", "tooltip": ( "The checkpoint or diffusion model (UNet) used " - "for generation. Only overwrites when non-empty." + "for generation. Fill in the name manually or " + "connect a MODEL output — the model name is then " + "extracted automatically. Only overwrites when " + "non-empty." ), }, ), @@ -158,13 +160,10 @@ class MetadataOverwriteLM: For most fields, a falsy value (empty string, 0) means "not set" and is skipped. clip_skip uses a dedicated sentinel (-25) so that a wired value of 0 is preserved and reaches the metadata pipeline. + + The ``model`` field accepts either a manual string or a wired MODEL + (ModelPatcher) connection; in the latter case the underlying model + name is extracted from the patcher's ``cached_patcher_init`` and + stored as a ComfyUI-style relative path. """ - result: dict[str, Any] = {} - for key in METADATA_OVERWRITE_FIELDS: - value = kwargs.get(key) - if key == "clip_skip": - if value != _CLIP_SKIP_SENTINEL: - result[key] = value - elif value: - result[key] = value - return (result,) + return (collect_overwrite_params(kwargs),) diff --git a/py/nodes/unet_loader.py b/py/nodes/unet_loader.py index a17bfc6c..d8d909c2 100644 --- a/py/nodes/unet_loader.py +++ b/py/nodes/unet_loader.py @@ -7,6 +7,21 @@ from ..utils.utils import get_checkpoint_info_absolute, _format_model_name_for_c logger = logging.getLogger(__name__) +def _reload_gguf_unet( + unet_path: str, weight_dtype: str, disable_dynamic: bool = False +) -> object: + """Reload a GGUF diffusion model from disk (cached_patcher_init factory). + + Mirrors the GGUF branch of UNETLoaderLM.load_unet so ModelPatcher + deepclone/dynamic machinery can rebuild GGUF models with the correct + GGMLOps. ``disable_dynamic`` is accepted for signature compatibility + with core ComfyUI loaders. + """ + loader = UNETLoaderLM() + model, = loader._load_gguf_unet(unet_path, unet_path, weight_dtype) + return model + + class UNETLoaderLM: """UNET Loader with support for extra folder paths @@ -196,6 +211,12 @@ class UNETLoaderLM: # Wrap with GGUFModelPatcher model = GGUFModelPatcher.clone(model) + # Register a reload factory so the MODEL carries its source path + # (cached_patcher_init) like core ComfyUI loaders do — required + # for model-name extraction downstream and for ModelPatcher + # deepclone/dynamic machinery. + model.cached_patcher_init = (_reload_gguf_unet, (unet_path, weight_dtype)) + return (model,) except Exception as e: diff --git a/py/utils/utils.py b/py/utils/utils.py index e61be515..4ea8c8bf 100644 --- a/py/utils/utils.py +++ b/py/utils/utils.py @@ -1,7 +1,7 @@ from difflib import SequenceMatcher import os import re -from typing import Dict +from typing import Any, Dict, List, Optional from ..services.service_registry import ServiceRegistry from ..config import config from ..services.settings_manager import get_settings_manager @@ -294,6 +294,53 @@ def _format_model_name_for_comfyui(file_path: str, model_roots: list) -> str: return os.path.basename(file_path) +def model_patcher_to_name(model_patcher: Any) -> Optional[str]: + """Extract a ComfyUI-style model name from a MODEL (ModelPatcher) object. + + Core ComfyUI loaders record the absolute weight file path on the patcher's + ``cached_patcher_init`` attribute: + - load_checkpoint_guess_config -> (fn, (ckpt_path, ...), index) + - load_diffusion_model -> (fn, (unet_path, model_options)) + Patcher clones (LoRA loaders, model merges, ...) preserve the attribute, + so the name is recoverable anywhere downstream of a core loader — including + from LoRA Manager's own loaders (CheckpointLoaderLM / UNETLoaderLM), which + call the same core load functions. + + The absolute path is converted to the ComfyUI-style relative name used by + the metadata pipeline (covering standard ComfyUI roots and LoRA Manager + extra folder paths). + + Returns None when the path cannot be recovered (e.g. third-party loaders + that never set ``cached_patcher_init``). + """ + init = getattr(model_patcher, "cached_patcher_init", None) + if not isinstance(init, (tuple, list)) or len(init) < 2: + return None + args = init[1] + abs_path = args[0] if args else None + if not isinstance(abs_path, str) or not abs_path: + return None + return _abs_model_path_to_name(abs_path) + + +def _abs_model_path_to_name(abs_path: str) -> str: + """Convert an absolute model path to a ComfyUI-style relative name. + + Tries standard ComfyUI model roots plus LoRA Manager extra folder paths; + falls back to the bare filename. + """ + try: + roots: List[str] = list(config.base_models_roots or []) + roots.extend(config.extra_checkpoints_roots or []) + roots.extend(config.extra_unet_roots or []) + formatted = _format_model_name_for_comfyui(abs_path, roots) + if formatted: + return formatted + except Exception: + pass + return os.path.basename(abs_path) + + def fuzzy_match(text: str, pattern: str, threshold: float = 0.85) -> bool: """ Check if text matches pattern using fuzzy matching.