diff --git a/nodes/consensus_merger.py b/nodes/consensus_merger.py index f48e91e..c3a2921 100644 --- a/nodes/consensus_merger.py +++ b/nodes/consensus_merger.py @@ -7,13 +7,14 @@ import logging import os from contextlib import closing from dataclasses import asdict, dataclass, field -from typing import Iterable +from typing import Any, Iterable import comfy.utils import folder_paths import torch import torch.nn.functional as F from comfy_api.latest import io +from comfy_api.latest._io import _DynamicGroup from tqdm import tqdm from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned @@ -3088,8 +3089,273 @@ class CWBEmbeddingMultiMerger(io.ComfyNode): return io.NodeOutput(filename, documentation, report) +def _embedding_template() -> list[io.Input]: + return [ + io.Combo.Input( + "embedding_name", + options=folder_paths.get_filename_list("embeddings"), + tooltip="The name of the Embedding file to apply.", + ), + io.Boolean.Input( + "enabled", + default=True, + tooltip="Turn off to skip this Embedding without changing its file.", + ), + ] + + +def build_embedding_stack_dynamic_combo( + depth: int = 1, max_depth: int = 30 +) -> io.DynamicCombo.Input: + sub_inputs: list[io.Input] = [ + io.Combo.Input( + "embedding_name", + options=folder_paths.get_filename_list("embeddings"), + tooltip="The name of the Embedding file to apply.", + ), + io.Boolean.Input( + "enabled", + default=True, + tooltip="Turn off to skip this Embedding without removing its settings.", + ), + ] + if depth < max_depth: + sub_inputs.append( + build_embedding_stack_dynamic_combo( + depth + 1, max_depth=max_depth + ) + ) + + input_id = "add_embedding" if depth == 1 else f"add_embedding_{depth}" + return io.DynamicCombo.Input( + input_id, + options=[ + io.DynamicCombo.Option("False", inputs=[]), + io.DynamicCombo.Option("True", inputs=sub_inputs), + ], + display_name="add_embedding", + tooltip="Enable to configure and add an Embedding to the stack.", + ) + + +def _parse_bool(value: Any, default: bool = True) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, str): + return value.strip().lower() in ("true", "1", "yes") + return bool(value) + + +def extract_embedding_stack_entries( + add_embedding_data: Any = None, embedding_list: str = "", **kwargs: Any +) -> list[dict[str, Any]]: + entries: list[dict[str, Any]] = [] + stack_active = False + + # 1. Walk nested dictionary structure + curr = add_embedding_data + depth = 1 + while isinstance(curr, dict): + key = f"add_embedding_{depth}" if depth > 1 else "add_embedding" + combo_val = curr.get(key, curr.get("add_embedding")) + is_active = _parse_bool(combo_val, default=False) + if not is_active: + break + + stack_active = True + embedding_name = curr.get("embedding_name") + enabled = _parse_bool(curr.get("enabled", True), default=True) + if embedding_name: + entries.append( + { + "embedding_name": str(embedding_name), + "enabled": bool(enabled), + } + ) + + depth += 1 + next_key = f"add_embedding_{depth}" + next_child = curr.get(next_key) + if not isinstance(next_child, dict): + for k, v in curr.items(): + if k.startswith("add_embedding") and k != key and isinstance(v, dict): + next_child = v + break + curr = next_child + + # 2. Fallback: inspect kwargs if add_embedding was not a dict + if not stack_active and kwargs: + raw_add = kwargs.get("add_embedding", add_embedding_data) + single_active = _parse_bool(raw_add, default=False) + if single_active: + stack_active = True + embedding_name = kwargs.get("embedding_name") + if embedding_name: + enabled = _parse_bool(kwargs.get("enabled", True), default=True) + entries.append( + { + "embedding_name": str(embedding_name), + "enabled": bool(enabled), + } + ) + + # 3. String list fallback when the stack is set to False + if not stack_active: + raw_list = embedding_list if embedding_list else kwargs.get("embedding_list", "") + if not raw_list and isinstance(add_embedding_data, dict): + raw_list = add_embedding_data.get("embedding_list", "") + if isinstance(raw_list, str) and raw_list.strip(): + for line in raw_list.splitlines(): + for part in line.split(","): + cleaned = part.strip() + if cleaned: + entries.append( + { + "embedding_name": cleaned, + "enabled": True, + } + ) + + return entries + + +def enable_dynamic_group_link_support() -> bool: + """Allow _DynamicGroup inputs to accept connected COMFY_DYNAMICGROUP_V3 socket links.""" + try: + from comfy_api.latest._io import ( + DYNAMIC_INPUT_LOOKUP, + finalize_prefix, + register_dynamic_input_func, + ) + from execution import is_link + + orig_func = DYNAMIC_INPUT_LOOKUP.get(_DynamicGroup.io_type) + if orig_func is None or getattr(orig_func, "_supports_links", False): + return True + + def _expand_dynamic_group_with_links( + out_dict, live_inputs, value, input_type, curr_prefix + ): + finalized_prefix = finalize_prefix(curr_prefix) + val = live_inputs.get(finalized_prefix) + if is_link(val): + out_dict[input_type][finalized_prefix] = value + return + return orig_func(out_dict, live_inputs, value, input_type, curr_prefix) + + _expand_dynamic_group_with_links._supports_links = True + register_dynamic_input_func( + _DynamicGroup.io_type, _expand_dynamic_group_with_links + ) + return True + except Exception: + return False + + +enable_dynamic_group_link_support() + + +class CWBEmbeddingStack(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="CWBEmbeddingStack", + display_name="CWB Embedding Stack", + category="ModelUtils/Configuration", + description="Specifies a stack of Embeddings using a recursive DynamicCombo template or a newline/comma-separated list when add_embedding is False.", + search_aliases=[ + "embedding", + "embedding stack", + "stack", + "dynamic group", + "embedding template", + ], + inputs=[ + build_embedding_stack_dynamic_combo(depth=1, max_depth=30), + io.String.Input( + "embedding_list", + default="", + multiline=True, + tooltip="Newline or comma-separated list of embedding path names used when add_embedding is False.", + ), + ], + outputs=[ + io.Custom("COMFY_DYNAMICGROUP_V3").Output( + "embeddings", + display_name="embeddings", + tooltip="The Embedding stack specification using the DynamicGroup format.", + ), + ], + ) + + @classmethod + def execute(cls, add_embedding=None, embedding_list: str = "", **kwargs) -> io.NodeOutput: + entries = extract_embedding_stack_entries( + add_embedding, embedding_list=embedding_list, **kwargs + ) + return io.NodeOutput(entries) + + +class CWBEmbeddingStackMerger(CWBEmbeddingMultiMerger): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="CWBEmbeddingStackMerger", + display_name="CWB Embedding Stack Multi-Merge", + category="ModelUtils/Embeddings/Merge", + description="CWB-merge an arbitrary stack or list of embeddings, then coalesce locally similar output vectors.", + inputs=[ + _DynamicGroup.Input( + "embeddings", + template=_embedding_template(), + min=1, + max=20, + tooltip="Each row includes one embedding in the merge, or connect an Embedding Stack.", + group_name="Embedding", + ), + *_embedding_coalesce_controls("cwb_merged_coalesced_embedding"), + ], + outputs=[ + io.AnyType.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + io.String.Output(display_name="cwb_report"), + ], + is_experimental=True, + ) + + @classmethod + def execute(cls, embeddings: list[dict] | None = None, **kwargs): + documentation = load_documentation_from_file("consensus_mergers.md") + if kwargs["execution_mode"] == "DOCUMENTATION ONLY": + return io.NodeOutput( + "Documentation mode active. No merge performed.", + documentation, + "CWB report unavailable because no merge was performed.", + ) + if kwargs["legacy_boundary_search"] and not kwargs["vision_boundary_embeddings"]: + raise ValueError("Legacy boundary search requires Vision Boundary Embeddings.") + names = [ + str(row["embedding_name"]).strip() + for row in (embeddings or []) + if isinstance(row, dict) + and _parse_bool(row.get("enabled", True), default=True) + and row.get("embedding_name") + and str(row["embedding_name"]).strip() + and str(row["embedding_name"]).strip() != "None" + ] + if len(names) < 2: + raise ValueError("At least 2 enabled embeddings are required for CWB Embedding Multi-Merge.") + filename, report = ConsensusMergerLogic.execute( + names, "embeddings", _embedding_coalesce_params(kwargs), embedding_union=True + ) + return io.NodeOutput(filename, documentation, report) + + CWB_MERGER_NODES = [ CWBCustomConfiguration, + CWBEmbeddingStack, CWBCheckpointTwoMerger, CWBCheckpointThreeMerger, CWBModelTwoMerger, @@ -3103,4 +3369,5 @@ CWB_MERGER_NODES = [ CWBEmbeddingThreeMerger, CWBEmbeddingSelfCoalesce, CWBEmbeddingMultiMerger, + CWBEmbeddingStackMerger, ] diff --git a/tests/test_consensus_merger.py b/tests/test_consensus_merger.py index e52aa4e..836752d 100644 --- a/tests/test_consensus_merger.py +++ b/tests/test_consensus_merger.py @@ -622,9 +622,9 @@ def test_norm_rescale_and_dsc_bandpass_are_finite(cwb): def test_all_cwb_schemas_and_compact_control_contract(cwb): - assert len(cwb.CWB_MERGER_NODES) == 14 + assert len(cwb.CWB_MERGER_NODES) == 16 schemas = [node.define_schema() for node in cwb.CWB_MERGER_NODES] - assert len({schema.node_id for schema in schemas}) == 14 + assert len({schema.node_id for schema in schemas}) == 16 advanced = { "consensus_type", "alignment_method", "alignment_threshold", "similarity_threshold", "power_alpha", "diversity_beta", @@ -639,7 +639,7 @@ def test_all_cwb_schemas_and_compact_control_contract(cwb): ids = [value.id for value in schema.inputs] assert schema.description assert all(value.tooltip for value in schema.inputs) - if node is cwb.CWBCustomConfiguration: + if node in (cwb.CWBCustomConfiguration, cwb.CWBEmbeddingStack): continue assert [output.display_name for output in schema.outputs] == [ "output_filename", "documentation", "cwb_report" @@ -658,7 +658,11 @@ def test_all_cwb_schemas_and_compact_control_contract(cwb): "include_mode", ] continue - if node in (cwb.CWBEmbeddingSelfCoalesce, cwb.CWBEmbeddingMultiMerger): + if node in ( + cwb.CWBEmbeddingSelfCoalesce, + cwb.CWBEmbeddingMultiMerger, + cwb.CWBEmbeddingStackMerger, + ): assert "vision_boundary_embeddings" in ids assert "legacy_boundary_search" in ids assert "boundary_reference_embedding" in ids @@ -1070,6 +1074,91 @@ def test_embedding_multi_wrapper_forwards_selected_inputs(monkeypatch, cwb): assert result.result == ("merged.safetensors", "DOCS", "REPORT") +def test_embedding_stack_and_stack_merger(monkeypatch, cwb): + stack_schema = cwb.CWBEmbeddingStack.define_schema() + assert stack_schema.node_id == "CWBEmbeddingStack" + assert [inp.id for inp in stack_schema.inputs] == ["add_embedding", "embedding_list"] + assert stack_schema.outputs[0].get_io_type() == "COMFY_DYNAMICGROUP_V3" + + # 1. Active recursive DynamicCombo stack + stack_data = { + "add_embedding": "True", + "embedding_name": "emb1.safetensors", + "enabled": True, + "add_embedding_2": { + "add_embedding_2": "True", + "embedding_name": "emb2.safetensors", + "enabled": False, + "add_embedding_3": { + "add_embedding_3": "True", + "embedding_name": "emb3.safetensors", + "enabled": True, + "add_embedding_4": {"add_embedding_4": "False"}, + }, + }, + } + stack_out = cwb.CWBEmbeddingStack.execute( + add_embedding=stack_data, embedding_list="ignored1.safetensors, ignored2.safetensors" + ) + assert stack_out.result[0] == [ + {"embedding_name": "emb1.safetensors", "enabled": True}, + {"embedding_name": "emb2.safetensors", "enabled": False}, + {"embedding_name": "emb3.safetensors", "enabled": True}, + ] + + # 2. Stack set to False -> uses newline/comma-separated string list + list_out = cwb.CWBEmbeddingStack.execute( + add_embedding="False", + embedding_list="first.safetensors, second.safetensors\nthird.safetensors, \nfourth.safetensors", + ) + assert list_out.result[0] == [ + {"embedding_name": "first.safetensors", "enabled": True}, + {"embedding_name": "second.safetensors", "enabled": True}, + {"embedding_name": "third.safetensors", "enabled": True}, + {"embedding_name": "fourth.safetensors", "enabled": True}, + ] + + # 3. CWBEmbeddingStackMerger is a child of CWBEmbeddingMultiMerger and executes enabled entries + assert issubclass(cwb.CWBEmbeddingStackMerger, cwb.CWBEmbeddingMultiMerger) + monkeypatch.setattr(cwb, "load_documentation_from_file", lambda _: "DOCS") + captured = {} + + def execute(names, model_type, params, **options): + captured.update(names=names, model_type=model_type, params=params, options=options) + return "merged.safetensors", "REPORT" + + monkeypatch.setattr(cwb.ConsensusMergerLogic, "execute", execute) + base_kwargs = { + "execution_mode": "MERGE", + "cwb_preset": "balanced_sim_mean", + "cwb_config": None, + "target_vector_count": 0, + "coalesce_similarity_threshold": 0.95, + "coalesce_position_window": 0.1, + "vision_boundary_embeddings": False, + "legacy_boundary_search": False, + "boundary_reference_embedding": "None", + "boundary_similarity_threshold": 0.95, + "output_filename": "merged", + "save_dtype": "fp32", + "process_device": "cpu", + "lazy_load": True, + "force_clear_cache": False, + "override_dtype": False, + } + res = cwb.CWBEmbeddingStackMerger.execute( + embeddings=stack_out.result[0], **base_kwargs + ) + assert captured["names"] == ["emb1.safetensors", "emb3.safetensors"] + assert res.result == ("merged.safetensors", "DOCS", "REPORT") + + with pytest.raises(ValueError, match="At least 2 enabled embeddings"): + cwb.CWBEmbeddingStackMerger.execute( + embeddings=[{"embedding_name": "only_one.safetensors", "enabled": True}], + **base_kwargs, + ) + + def test_streaming_dtype_nonfloat_and_secondary_preservation( monkeypatch, tmp_path, cwb ):