From 280181f92ec5a5b3985f7c5496fd35f493ca247d Mon Sep 17 00:00:00 2001 From: Will Miao Date: Fri, 14 Aug 2026 15:21:28 +0800 Subject: [PATCH] feat(metadata-overwrite): support wired SAMPLER input on sampler field The sampler field now accepts either a manual string or a SAMPLER connection. When wired, the sampler name is extracted from the KSAMPLER object's sampler_function __name__ (sample_euler -> euler), with special-casing for dpm_fast/dpm_adaptive local closures and uni_pc/uni_pc_bh2 function names. - sampler input declared as "STRING,SAMPLER" with widgetType STRING, mirroring the existing model field union pattern - shared collect_overwrite_params() handles the non-str branch so the node and the metadata extractor conversion logic stay in sync; unrecognized sampler functions are logged and skipped - note: ddim is constructed by ComfyUI as euler with random inpaint, so the ddim name is unrecoverable and extracts as euler --- py/metadata_collector/overwrite_utils.py | 13 ++- py/nodes/metadata_overwrite.py | 16 +++- py/utils/utils.py | 36 ++++++++ .../test_metadata_collector.py | 91 +++++++++++++++++++ 4 files changed, 151 insertions(+), 5 deletions(-) diff --git a/py/metadata_collector/overwrite_utils.py b/py/metadata_collector/overwrite_utils.py index 54d58ef3..5aee305f 100644 --- a/py/metadata_collector/overwrite_utils.py +++ b/py/metadata_collector/overwrite_utils.py @@ -8,7 +8,7 @@ cannot drift between the two paths. import logging from typing import Any, Dict -from ..utils.utils import model_patcher_to_name +from ..utils.utils import model_patcher_to_name, sampler_object_to_name from .constants import CLIP_SKIP_SENTINEL, METADATA_OVERWRITE_FIELDS logger = logging.getLogger(__name__) @@ -22,7 +22,9 @@ def collect_overwrite_params(values: Dict[str, Any]) -> Dict[str, Any]: 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. + stored as a ComfyUI-style relative path. The ``sampler`` field likewise + accepts a manual string or a wired SAMPLER (KSAMPLER) connection, from + which the sampler name is extracted via the sampler function's name. """ result: Dict[str, Any] = {} for key in METADATA_OVERWRITE_FIELDS: @@ -34,6 +36,13 @@ def collect_overwrite_params(values: Dict[str, Any]) -> Dict[str, Any]: "Could not extract model name from wired MODEL input " "(no cached_patcher_init); model metadata overwrite skipped" ) + elif key == "sampler" and not isinstance(value, str): + value = sampler_object_to_name(value) + if value is None: + logger.warning( + "Could not extract sampler name from wired SAMPLER input " + "(unrecognized sampler function); sampler metadata overwrite skipped" + ) if key == "clip_skip": if value != CLIP_SKIP_SENTINEL: result[key] = value diff --git a/py/nodes/metadata_overwrite.py b/py/nodes/metadata_overwrite.py index 337d8e92..58956852 100644 --- a/py/nodes/metadata_overwrite.py +++ b/py/nodes/metadata_overwrite.py @@ -71,10 +71,18 @@ class MetadataOverwriteLM: }, ), "sampler": ( - "STRING", + "STRING,SAMPLER", { "default": "", - "tooltip": "Sampler name. Only overwrites when non-empty.", + "widgetType": "STRING", + "tooltip": ( + "Sampler name. Fill in the name manually or " + "connect a SAMPLER output (e.g. KSamplerSelect) " + "— the sampler name is then extracted " + "automatically. Note: ddim is recorded as " + "euler (ComfyUI internal representation). " + "Only overwrites when non-empty." + ), }, ), "scheduler": ( @@ -164,6 +172,8 @@ class MetadataOverwriteLM: 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. + stored as a ComfyUI-style relative path. The ``sampler`` field + likewise accepts a manual string or a wired SAMPLER (KSAMPLER) + connection, from which the sampler name is extracted automatically. """ return (collect_overwrite_params(kwargs),) diff --git a/py/utils/utils.py b/py/utils/utils.py index 8670d497..bcfe4d7d 100644 --- a/py/utils/utils.py +++ b/py/utils/utils.py @@ -323,6 +323,42 @@ def model_patcher_to_name(model_patcher: Any) -> Optional[str]: return _abs_model_path_to_name(abs_path) +def sampler_object_to_name(sampler: Any) -> Optional[str]: + """Extract a ComfyUI-style sampler name from a SAMPLER (KSAMPLER) object. + + Standard outputs (KSamplerSelect, most built-in sampler nodes) round-trip + losslessly via the underlying sampler function's ``__name__`` + (``sample_euler`` -> ``euler``). A few edge cases need special-casing + because the function name diverges from the ``SAMPLER_NAMES`` entry: + + - ``dpm_fast`` / ``dpm_adaptive`` are local closures inside + ``comfy.samplers.ksampler`` (``dpm_fast_function`` / ``dpm_adaptive_function``) + - ``uni_pc`` / ``uni_pc_bh2`` use ``sample_unipc`` / ``sample_unipc_bh2`` + + ``ddim`` is constructed by ComfyUI as ``euler`` with random inpaint, so + the original ``ddim`` name is unrecoverable (extracts as ``euler``). + Custom sampler nodes that pass non-``sample_*`` functions return None. + + Returns None when the name cannot be recovered. + """ + sampler_function = getattr(sampler, "sampler_function", None) + func_name = getattr(sampler_function, "__name__", None) + if not isinstance(func_name, str) or not func_name: + return None + if func_name == "dpm_fast_function": + return "dpm_fast" + if func_name == "dpm_adaptive_function": + return "dpm_adaptive" + if func_name.startswith("sample_"): + name = func_name[len("sample_"):] + if name == "unipc": + return "uni_pc" + if name == "unipc_bh2": + return "uni_pc_bh2" + return name or None + return None + + def _abs_model_path_to_name(abs_path: str) -> str: """Convert an absolute model path to a ComfyUI-style relative name. diff --git a/tests/metadata_collector/test_metadata_collector.py b/tests/metadata_collector/test_metadata_collector.py index 8f22540e..f9efd403 100644 --- a/tests/metadata_collector/test_metadata_collector.py +++ b/tests/metadata_collector/test_metadata_collector.py @@ -1373,6 +1373,97 @@ def test_metadata_overwrite_extractor_empty_inputs(metadata_registry): metadata_registry.clear_metadata() +def _make_ksampler(func_name: str): + """Build a duck-typed comfy.samplers.KSAMPLER stub with a named function.""" + def _sampler_function(*args, **kwargs): + pass + + _sampler_function.__name__ = func_name + return SimpleNamespace(sampler_function=_sampler_function) + + +def test_metadata_overwrite_extractor_sampler_union(metadata_registry): + """Wired SAMPLER objects should be converted to sampler names.""" + from py.metadata_collector.constants import CLIP_SKIP_SENTINEL + + metadata_registry.start_collection("prompt-ow-sampler") + metadata = metadata_registry.prompt_metadata["prompt-ow-sampler"] + + inputs: Dict[str, Any] = {key: "" for key in METADATA_OVERWRITE_FIELDS} + inputs.update({"seed": 0, "steps": 0, "cfg_scale": 0.0, "clip_skip": CLIP_SKIP_SENTINEL}) + inputs["sampler"] = _make_ksampler("sample_euler") + + MetadataOverwriteExtractor.extract("ow-sampler-1", inputs, None, metadata) + + params = metadata[OVERWRITE]["ow-sampler-1"]["parameters"] + assert params["sampler"] == "euler" + + metadata_registry.clear_metadata() + + +def test_metadata_overwrite_extractor_sampler_union_special_cases(metadata_registry): + """Sampler functions whose names diverge from SAMPLER_NAMES entries.""" + from py.metadata_collector.constants import CLIP_SKIP_SENTINEL + + metadata_registry.start_collection("prompt-ow-sampler2") + metadata = metadata_registry.prompt_metadata["prompt-ow-sampler2"] + + cases = [ + ("dpm_fast_function", "dpm_fast"), + ("dpm_adaptive_function", "dpm_adaptive"), + ("sample_unipc", "uni_pc"), + ("sample_unipc_bh2", "uni_pc_bh2"), + ("sample_dpmpp_2m_sde", "dpmpp_2m_sde"), + ] + for i, (func_name, expected) in enumerate(cases): + inputs: Dict[str, Any] = {key: "" for key in METADATA_OVERWRITE_FIELDS} + inputs.update({"seed": 0, "steps": 0, "cfg_scale": 0.0, "clip_skip": CLIP_SKIP_SENTINEL}) + inputs["sampler"] = _make_ksampler(func_name) + MetadataOverwriteExtractor.extract(f"ow-sampler-{i}", inputs, None, metadata) + + params_by_node = {node_id: entry["parameters"] for node_id, entry in metadata[OVERWRITE].items()} + for i, (func_name, expected) in enumerate(cases): + assert params_by_node[f"ow-sampler-{i}"]["sampler"] == expected, func_name + + metadata_registry.clear_metadata() + + +def test_metadata_overwrite_extractor_sampler_union_unrecognized_skipped(metadata_registry): + """Unrecognized sampler functions should skip the field, not crash.""" + from py.metadata_collector.constants import CLIP_SKIP_SENTINEL + + metadata_registry.start_collection("prompt-ow-sampler3") + metadata = metadata_registry.prompt_metadata["prompt-ow-sampler3"] + + inputs: Dict[str, Any] = {key: "" for key in METADATA_OVERWRITE_FIELDS} + inputs.update({"seed": 0, "steps": 0, "cfg_scale": 0.0, "clip_skip": CLIP_SKIP_SENTINEL}) + inputs["sampler"] = _make_ksampler("my_custom_sampler_function") + + MetadataOverwriteExtractor.extract("ow-sampler-unrec", inputs, None, metadata) + + assert not metadata[OVERWRITE] + + metadata_registry.clear_metadata() + + +def test_metadata_overwrite_extractor_sampler_union_no_sampler_function(metadata_registry): + """Objects without a sampler_function (e.g. old KUNASampler classes) are skipped.""" + from py.metadata_collector.constants import CLIP_SKIP_SENTINEL + + metadata_registry.start_collection("prompt-ow-sampler4") + metadata = metadata_registry.prompt_metadata["prompt-ow-sampler4"] + + inputs: Dict[str, Any] = {key: "" for key in METADATA_OVERWRITE_FIELDS} + inputs.update({"seed": 0, "steps": 0, "cfg_scale": 0.0, "clip_skip": CLIP_SKIP_SENTINEL}) + inputs["sampler"] = SimpleNamespace() + + MetadataOverwriteExtractor.extract("ow-sampler-nofn", inputs, None, metadata) + + assert not metadata[OVERWRITE] + + metadata_registry.clear_metadata() + + def test_extract_generation_params_applies_overwrite(metadata_registry, populated_registry, monkeypatch): """overwrite values should replace inferred params in extract_generation_params.""" import py.metadata_collector.metadata_processor as mp