diff --git a/__init__.py b/__init__.py index b259dcd..26c84d0 100644 --- a/__init__.py +++ b/__init__.py @@ -1,38 +1,38 @@ import random import re - +# 真の動的入力を実現するための辞書クラス class StringToolsOptionalDict(dict): - getitem_default_callback = None - - def __init__(self, *args, get_default_callback=None, **kwargs): - super().__init__(*args, **kwargs) - self.getitem_default_callback = get_default_callback - def __contains__(self, key: object) -> bool: - return True + # text_NN 形式のキーなら常に存在するとみなす + if isinstance(key, str) and re.match(r"text_\d+", key): + return True + return super().__contains__(key) def __getitem__(self, key): - if callable(self.getitem_default_callback): - return super().get(key, self.getitem_default_callback(key)) - else: - return super().get(key, None) + if isinstance(key, str) and re.match(r"text_\d+", key): + return super().get(key, ("STRING", {"forceInput": True})) + return super().__getitem__(key) + + def get(self, key, default=None): + if isinstance(key, str) and re.match(r"text_\d+", key): + return super().get(key, ("STRING", {"forceInput": True})) + return super().get(key, default) -def sort_kwargs_value(basename, kwargs, separator="_"): - sorted_args = {} +def sort_kwargs_value(basename, kwargs): basename_dict = {} - for key, value in kwargs.items(): - if separator.join(key.split(separator)[:-1]) != basename: - sorted_args[key] = value - else: - basename_dict[key] = value - - sorted_basename_dict = sorted( - basename_dict.items(), key=lambda kv: int(kv[0].split(separator)[-1]) - ) + if key.startswith(basename) and value is not None: + parts = key.split("_") + if len(parts) > 1 and parts[1].isdigit(): + # value がリストの場合は最初の要素を取り出し、そうでなければそのまま文字列化 + if isinstance(value, list): + basename_dict[int(parts[1])] = str(value[0]) if value else "" + else: + basename_dict[int(parts[1])] = str(value) + sorted_basename_dict = sorted(basename_dict.items()) return list(map(lambda kv: kv[1], sorted_basename_dict)) @@ -40,148 +40,119 @@ def get_node(kwargs): extra_pnginfo = kwargs.get("extra_pnginfo", None) unique_id = kwargs.get("unique_id", None) - # INPUT_IS_LIST = True のノードではhidden入力もリストにラップされるためアンラップ if isinstance(extra_pnginfo, list): - extra_pnginfo = extra_pnginfo[0] + extra_pnginfo = extra_pnginfo[0] if extra_pnginfo else None if isinstance(unique_id, list): - unique_id = unique_id[0] + unique_id = unique_id[0] if unique_id else None if extra_pnginfo is None or unique_id is None: - print( - re.sub( - pattern="^\s*> ", - repl="", - string=""" - > "extra_pnginfo" or "unique_id" not found in kwargs. - > please add to INPUT_TYPES: - > "hidden": { - > "extra_pnginfo": "EXTRA_PNGINFO", - > "unique_id": "UNIQUE_ID", - > } - """, - flags=re.MULTILINE, - ) - ) return None - workflow = extra_pnginfo["workflow"] - node_path = [int(p) for p in unique_id.split(":")] + workflow = extra_pnginfo.get("workflow", None) + if workflow is None: + return None - subgraphs = workflow["definitions"]["subgraphs"] if "definitions" in workflow and "subgraphs" in workflow["definitions"] else [] - subgraphs_by_id = {} - for subgraph in subgraphs: - subgraphs_by_id[subgraph["id"]] = subgraph + node_path = [unique_id] + + def walkdown_node_path(current_path, workflow): + if not current_path: return None + node_id = current_path[0] + node = None + for n in workflow.get("nodes", []): + if str(n["id"]) == str(node_id): + node = n + break - def walkdown_node_path(current_path, graph): - nodes_by_id = {} - for n in graph["nodes"]: - nodes_by_id[n["id"]] = n - node = nodes_by_id.get(current_path[0]) if node is None: return None - if node["type"] in subgraphs_by_id: - return walkdown_node_path(current_path[1:], subgraphs_by_id[node["type"]]) - else: - return node + if node["type"] == "Subgraph" or "subgraph" in node: + subgraphs_by_id = workflow.get("definitions", {}) + if node["type"] in subgraphs_by_id: + return walkdown_node_path(current_path[1:], subgraphs_by_id[node["type"]]) + + return node node = walkdown_node_path(node_path, workflow) - return node -def get_input_basename(input_name): - return input_name.split('_')[0] - -def get_input_extraname(input_name): - parts = input_name.split('_', 1) - return parts[1] if len(parts) > 1 else "" - -def calculate_weights_from_prompt(prompt, target_node_id, with_weights_classes, basename="text"): +def calculate_weights_from_prompt(prompt, target_node_id, with_weights_classes, weight_prefix="text"): + # JS 実装 (index.ts) と同様に ID をフラット化してマッチングするための前処理 + # 末尾の数値IDをキーにしたマップを作成する flatten_prompt = {} - for key, value in prompt.items(): - flatten_id = key.split(':')[-1] - flatten_prompt[flatten_id] = value + for k, v in prompt.items(): + flatten_id = str(k).split(':')[-1] + flatten_prompt[flatten_id] = v - if str(target_node_id) not in flatten_prompt: - return {} + # ターゲットIDも同様に正規化 + target_id_str = str(target_node_id).split(':')[-1] - node = flatten_prompt[str(target_node_id)] - weights_result = {} - - def walkdown(type_name, id_str, current_sum): - target_node = flatten_prompt.get(str(id_str.split(':')[-1])) - if not target_node: - return current_sum + # 再帰的にリーフノードの合計数を数える内部関数 + def count_leaves(current_id): + curr_id_str = str(current_id).split(':')[-1] + if not curr_id_str or curr_id_str not in flatten_prompt: + return 1 - for input_name, value in target_node.get("inputs", {}).items(): - if get_input_basename(input_name) != basename: - continue - - start = 0 - if target_node.get("class_type") in with_weights_classes: - start = 1 - if isinstance(value, list) and len(value) > 0: - current_sum += walkdown(target_node.get("class_type"), value[0], start) - return current_sum - - for input_name, value in node.get("inputs", {}).items(): - if get_input_basename(input_name) != basename: - continue + node_info = flatten_prompt[curr_id_str] + class_type = node_info.get("class_type") - extraname = get_input_extraname(input_name) - if isinstance(value, list) and len(value) > 0: - total_sum = walkdown(node.get("class_type"), value[0], 0) - if total_sum == 0: - total_sum = 1 - weights_result[f"weight_{extraname}"] = total_sum + # ターゲットクラスの場合 + if class_type in with_weights_classes: + node_inputs = node_info.get("inputs", {}) + total = 0 + for key, value in node_inputs.items(): + if key.startswith(weight_prefix): + if isinstance(value, list) and len(value) >= 2: + total += count_leaves(value[0]) + else: + total += 1 + elif key == "text_list": + if isinstance(value, list) and len(value) >= 2: + total += count_leaves(value[0]) + else: + total += 1 + return max(1, total) + + # 中間ノードの場合は全入力を再帰探索 + node_inputs = node_info.get("inputs", {}) + total = 0 + has_connections = False + for key, value in node_inputs.items(): + if isinstance(value, list) and len(value) >= 2: + has_connections = True + total += count_leaves(value[0]) + + if not has_connections: + return 1 + return max(1, total) - return weights_result + weights = {} + main_node = flatten_prompt.get(target_id_str) + if not main_node: + return weights - -def format(num): - return "{:6.2f}".format(num) - - -class StringToolsSeed: - RETURN_TYPES = ("INT",) - FUNCTION = "process" - CATEGORY = "string-tools" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "seed": ( - "INT", - { - "default": 0, - "min": 0, - "max": 0xffffffffffffffff, - } - ), - } - } - - def process(self, seed): - return (seed,) + main_inputs = main_node.get("inputs", {}) + for key, value in main_inputs.items(): + if key.startswith(weight_prefix) or key == "text_list": + if isinstance(value, list) and len(value) >= 2: + weights[key] = count_leaves(value[0]) + else: + weights[key] = 1 + + return weights class StringToolsString: - def __init__(self): - pass - @classmethod def INPUT_TYPES(s): return { "required": { "text": ( "STRING", - { - "dynamicPrompts": False, - }, + {"default": "", "multiline": False, "dynamicPrompts": False}, ), - }, + } } RETURN_TYPES = ("STRING",) @@ -193,21 +164,15 @@ class StringToolsString: class StringToolsText: - def __init__(self): - pass - @classmethod def INPUT_TYPES(s): return { "required": { "text": ( "STRING", - { - "multiline": True, - "dynamicPrompts": False, - }, + {"default": "", "multiline": True, "dynamicPrompts": False}, ), - }, + } } RETURN_TYPES = ("STRING",) @@ -219,93 +184,15 @@ class StringToolsText: class StringToolsConcat: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": {}, - "optional": { - "separator": ( - "STRING", - {"forceInput": True}, - ), - }, - "hidden": { - "text": ( - "STRING", - {"forceInput": True}, - ), - }, - } - - RETURN_TYPES = ("STRING",) - FUNCTION = "process" - CATEGORY = "string-tools" - - def process(self, *args, **kwargs): - if "separator" in kwargs: - separator = kwargs["separator"] - del kwargs["separator"] - else: - separator = "\n" - return (separator.join(sort_kwargs_value("text", kwargs)),) - -class StringToolsConcatList: - def __init__(self): - pass - @classmethod def INPUT_TYPES(s): return { "required": { - "text_list": ( - "STRING", - {"forceInput": True}, - ), - }, - "optional": { - "separator": ( - "STRING", - {"forceInput": True}, - ), - }, - } - - INPUT_IS_LIST = True - RETURN_TYPES = ("STRING",) - FUNCTION = "process" - CATEGORY = "string-tools" - - def process(self, text_list, separator=None): - sep = "\n" if separator is None or len(separator) == 0 else separator[0] - return (sep.join(text_list),) - - -class StringToolsRandomChoice: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "seed": ( - "INT", - { - "default": 0, - "min": 0, - "max": 0xFFFFFFFFFFFFFFFF, - "forceInput": True, - }, - ), }, + "optional": StringToolsOptionalDict({ + "separator": ("STRING", {"forceInput": True}), + }), "hidden": { - "text": ( - "STRING", - {"forceInput": True}, - ), "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", "unique_id": "UNIQUE_ID", @@ -316,39 +203,24 @@ class StringToolsRandomChoice: FUNCTION = "process" CATEGORY = "string-tools" - def process(self, *args, **kwargs): - seed = 0 - if "seed" in kwargs: - seed = kwargs["seed"] - del kwargs["seed"] - + def process(self, **kwargs): + separator = kwargs.get("separator", "") + if isinstance(separator, list): + separator = str(separator[0]) if separator else "" + values = sort_kwargs_value("text", kwargs) - random.seed(seed) - choice = random.choice(values) - return (choice,) + return (separator.join(values),) -class StringToolsRandomChoiceList: - def __init__(self): - pass - +class StringToolsConcatList: @classmethod def INPUT_TYPES(s): return { "required": { - "text_list": ( - "STRING", - {"forceInput": True}, - ), - "seed": ( - "INT", - { - "default": 0, - "min": 0, - "max": 0xFFFFFFFFFFFFFFFF, - "forceInput": True, - }, - ), + "text_list": ("STRING", {"forceInput": True}), + }, + "optional": { + "separator": ("STRING", {"forceInput": True}), }, } @@ -357,12 +229,88 @@ class StringToolsRandomChoiceList: FUNCTION = "process" CATEGORY = "string-tools" - def process(self, text_list, seed): + def process(self, text_list=None, separator=None): + if not text_list: return ("",) + + if separator is None: + sep = "\n" + elif isinstance(separator, list): + sep = separator[0] if separator else "\n" + else: + sep = str(separator) + + flat_list = [] + for x in text_list: + if isinstance(x, list): + flat_list.extend(map(str, x)) + else: + flat_list.append(str(x)) + return (sep.join(flat_list),) + + +class StringToolsRandomChoice: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "forceInput": True}), + }, + "optional": StringToolsOptionalDict(), + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + "unique_id": "UNIQUE_ID", + }, + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "process" + CATEGORY = "string-tools" + + def process(self, **kwargs): + seed = kwargs.get("seed", 0) + values = sort_kwargs_value("text", kwargs) + if not values: + return ("",) + random.seed(seed) + choice = random.choice(values) + return (choice,) + + +class StringToolsRandomChoiceList: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "forceInput": True}), + "text_list": ("STRING", {"forceInput": True}), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + "unique_id": "UNIQUE_ID", + }, + } + + INPUT_IS_LIST = True + RETURN_TYPES = ("STRING",) + FUNCTION = "process" + CATEGORY = "string-tools" + + def process(self, seed=0, text_list=None, **kwargs): if not text_list: return ("",) + + flat_list = [] + for x in text_list: + if isinstance(x, list): + flat_list.extend(x) + else: + flat_list.append(x) + s = seed[0] if isinstance(seed, list) else seed random.seed(s) - choice = random.choice(text_list) + choice = random.choice(flat_list) return (choice,) @@ -373,27 +321,24 @@ class StringToolsBalancedChoice(StringToolsRandomChoice): @classmethod def INPUT_TYPES(s): - input_types = super().INPUT_TYPES() - input_types["optional"] = StringToolsOptionalDict( - get_default_callback=lambda key: ( - ("INT") if key.startswith("weight_") else None, - ) - ) - return input_types - - def process(self, *args, **kwargs): - seed = 0 - if "seed" in kwargs: - seed = kwargs["seed"] - del kwargs["seed"] + return super().INPUT_TYPES() + def process(self, **kwargs): + seed = kwargs.get("seed", 0) node = get_node(kwargs) - title = node.get("title", node.get("type", None)) + values = sort_kwargs_value("text", kwargs) - weights = sort_kwargs_value("weight", kwargs) - random.seed(seed) + if not values: + return ("",) + + if node is None: + random.seed(seed) + return (random.choice(values),) + + unique_id = kwargs.get("unique_id", None) + if isinstance(unique_id, list): unique_id = unique_id[0] + id = str(unique_id) if unique_id is not None else str(node.get("id", "")) - id = node.get("id", None) counts = self.counts.get(id, None) if counts is None: counts = self.counts[id] = {} @@ -401,10 +346,18 @@ class StringToolsBalancedChoice(StringToolsRandomChoice): if total_count is None: total_count = self.total_count[id] = 0 - # weightsが空またはvaluesと長さが合わない場合は均等配分にフォールバック - if not weights or len(weights) != len(values): + title = node.get("title", node.get("type", None)) + + prompt = kwargs.get("prompt", {}) + computed_weights = calculate_weights_from_prompt( + prompt, id, ["StringToolsRandomChoice", "StringToolsBalancedChoice", "StringToolsRandomChoiceList", "StringToolsBalancedChoiceList"], "text" + ) + if computed_weights: + weights = [computed_weights.get(f"text_{i}", 1) for i in range(len(values))] + else: weights = [1] * len(values) + random.seed(seed) choice_idx = random.choices(range(len(values)), weights=weights)[0] choice_text = "_".join(["text", str(choice_idx)]) counts[choice_text] = counts.get(choice_text, 0) + 1 @@ -416,9 +369,7 @@ class StringToolsBalancedChoice(StringToolsRandomChoice): for idx in range(len(values)): text = "_".join(["text", str(idx)]) count = counts.get(text, 0) - print( - f"\tInput:{text} Weight:{weights[idx]} ({format(weights[idx] / total_weight * 100)}%) - Count:{count} ({format(count / self.total_count[id] * 100)}%)" - ) + print(f"\tInput:{text} Weight:{weights[idx]} ({weights[idx] / total_weight * 100}%) - Count:{count} ({count / self.total_count[id] * 100}%)") return (values[choice_idx],) @@ -430,37 +381,13 @@ class StringToolsBalancedChoiceList(StringToolsRandomChoiceList): @classmethod def INPUT_TYPES(s): - return { - "required": { - "text_list": ( - "STRING", - {"forceInput": True}, - ), - "seed": ( - "INT", - { - "default": 0, - "min": 0, - "max": 0xFFFFFFFFFFFFFFFF, - "forceInput": True, - }, - ), - }, - "optional": { - "weight_list": ( - "INT", - {"forceInput": True}, - ) - }, - "hidden": { - "prompt": "PROMPT", - "extra_pnginfo": "EXTRA_PNGINFO", - "unique_id": "UNIQUE_ID", - }, + input_types = super().INPUT_TYPES() + input_types["optional"] = { + "weight_list": ("INT", {"forceInput": True}) } + return input_types - def process(self, text_list, seed, weight_list=None, **kwargs): - # INPUT_IS_LIST = True のため hidden 入力もリストにラップされる + def process(self, seed=0, text_list=None, weight_list=None, **kwargs): if isinstance(kwargs.get("prompt"), list): kwargs["prompt"] = kwargs["prompt"][0] if isinstance(kwargs.get("extra_pnginfo"), list): @@ -471,97 +398,152 @@ class StringToolsBalancedChoiceList(StringToolsRandomChoiceList): if not text_list: return ("",) + flat_list = [] + for x in text_list: + if isinstance(x, list): + flat_list.extend(x) + else: + flat_list.append(x) + s = seed[0] if isinstance(seed, list) else seed random.seed(s) node = get_node(kwargs) if node is None: - return (random.choice(text_list),) + return (random.choice(flat_list),) title = node.get("title", node.get("type", None)) - id = node.get("id", None) + unique_id = kwargs.get("unique_id", None) + if isinstance(unique_id, list): unique_id = unique_id[0] + id = str(unique_id) if unique_id is not None else str(node.get("id", "")) counts = self.counts.get(id, None) if counts is None: counts = self.counts[id] = {} + total_count = self.total_count.get(id, None) if total_count is None: total_count = self.total_count[id] = 0 - + weights = [] - if weight_list is not None and len(weight_list) == len(text_list): - weights = weight_list - else: + if weight_list is not None: + # weight_list も平坦化が必要な可能性があるが、通常は 1つ + flat_weights = [] + for w in weight_list: + if isinstance(w, list): flat_weights.extend(w) + else: flat_weights.append(w) + + if len(flat_weights) == len(flat_list): + weights = flat_weights + + if not weights: prompt = kwargs.get("prompt", {}) computed_weights = calculate_weights_from_prompt( prompt, id, ["StringToolsRandomChoice", "StringToolsBalancedChoice", "StringToolsRandomChoiceList", "StringToolsBalancedChoiceList"], "text_list" ) - # If input is connected to a list-generating node, text_list is a single connection. - # Thus computed_weights might just have weight_list or weight_0. - # For simplicity, if we don't have exactly per-element weights, uniform weight is used - # Unless we are getting multiple input sources. List nodes usually receive 1 link for text_list. - # To extract valid weights, we'd need the upstream node to supply a weight_list, - # Or assume uniform weights if not supplied. - - # Reverting to basic equal weights if no weight_list is provided for List nodes logic, - # or picking up from computed weights if applicable. - weights = [1] * len(text_list) - - # If the upstream node passed exactly matching length weights via dict mapping - if computed_weights and len(computed_weights) == len(text_list): - weights = [computed_weights[k] for k in sorted(computed_weights.keys())] + if computed_weights: + w = computed_weights.get("text_list", 1) + weights = [w / len(flat_list)] * len(flat_list) + else: + weights = [1] * len(flat_list) - choice_idx = random.choices(range(len(text_list)), weights=weights)[0] - choice_text = "_".join(["text", str(choice_idx)]) + choice_idx = random.choices(range(len(flat_list)), weights=weights)[0] + choice_text = str(choice_idx) counts[choice_text] = counts.get(choice_text, 0) + 1 self.total_count[id] += 1 total_weight = sum(weights) if self.debug or kwargs.get("debug", False): - print(f"#{id} {title} (Seed:{s}):") - for idx in range(len(text_list)): - text = "_".join(["text", str(idx)]) - count = counts.get(text, 0) - print( - f"\tInput:{text} Weight:{weights[idx]} ({format(weights[idx] / total_weight * 100)}%) - Count:{count} ({format(count / self.total_count[id] * 100)}%)" - ) + print(f"#{id} {title} (Seed:{seed}):") + for idx in range(len(flat_list)): + count = counts.get(str(idx), 0) + print(f"\tInput:{idx} Weight:{weights[idx]} ({weights[idx] / total_weight * 100 if total_weight > 0 else 0}%) - Count:{count} ({count / self.total_count[id] * 100}%)") - return (text_list[choice_idx],) + return (flat_list[choice_idx],) -class MockStringList: - def __init__(self): - pass +class StringToolsSeed: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "seed": ( + "INT", + {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}, + ), + } + } + RETURN_TYPES = ("INT",) + FUNCTION = "process" + CATEGORY = "string-tools" + + def process(self, seed): + return (seed,) + + +class StringToolsStringsToList: + @classmethod + def INPUT_TYPES(s): + return { + "required": {}, + "optional": StringToolsOptionalDict(), + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + "unique_id": "UNIQUE_ID", + }, + } + + RETURN_TYPES = ("STRING",) + OUTPUT_IS_LIST = (True,) + FUNCTION = "process" + CATEGORY = "string-tools" + + def process(self, **kwargs): + values = sort_kwargs_value("text", kwargs) + return (values,) + + +class StringToolsMockStringList: @classmethod def INPUT_TYPES(s): return { "required": { "str1": ("STRING", {"forceInput": True}), "str2": ("STRING", {"forceInput": True}), - }, + } } - + RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("list",) OUTPUT_IS_LIST = (True,) FUNCTION = "process" - CATEGORY = "string-tools" + CATEGORY = "string-tools/test" def process(self, str1, str2): return ([str1, str2],) + NODE_CLASS_MAPPINGS = { - "StringToolsSeed": StringToolsSeed, "StringToolsString": StringToolsString, "StringToolsText": StringToolsText, "StringToolsConcat": StringToolsConcat, - "StringToolsRandomChoice": StringToolsRandomChoice, - "StringToolsBalancedChoice": StringToolsBalancedChoice, "StringToolsConcatList": StringToolsConcatList, + "StringToolsRandomChoice": StringToolsRandomChoice, "StringToolsRandomChoiceList": StringToolsRandomChoiceList, + "StringToolsBalancedChoice": StringToolsBalancedChoice, "StringToolsBalancedChoiceList": StringToolsBalancedChoiceList, - "string-tools/MockStringList": MockStringList, + "StringToolsSeed": StringToolsSeed, + "StringToolsStringsToList": StringToolsStringsToList, } -WEB_DIRECTORY = "./js" +import os +if os.environ.get("COMFYUI_TEST_MODE") == "true" or os.environ.get("PYTEST_CURRENT_TEST"): + NODE_CLASS_MAPPINGS["StringToolsMockStringList"] = StringToolsMockStringList + +NODE_DISPLAY_NAME_MAPPINGS = { + "StringToolsMockStringList": "StringTools [TEST ONLY] Mock String List" +} + +WEB_DIRECTORY = "./js" +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..2bde0b6 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +testpaths = tests +python_files = test_*.py +asyncio_mode = auto +asyncio_default_test_loop_scope = function diff --git a/src/index.ts b/src/index.ts index 41e64e4..39f3a58 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,266 +1,190 @@ -import { app } from "/scripts/app.js" -import { api } from "/scripts/api.js" -import { setWidgetConfig } from "/extensions/core/widgetInputs.js" - -declare const LiteGraph: any - -function createCallback(nodename: string, basename: string, inputType: string, withWeights?:string[]) { - return async function(nodeType:any, nodeData:any, app:any) { - if (nodeData.name !== nodename) { - return - } - - const getInputBasename = function(input: any) { - return input.name.split('_')[0] - } - - const getInputExtraname = function(input: any) { - return input.name.split('_', 2)[1] - } - - const updateInputs = function(this: any) { - this.widgets = [] - - // remove empty inputs - for (let index = this.inputs.length; index--; ) { - const input = this.inputs[index] - if (getInputBasename(input) === basename && input.link === null && this.removeCancel !== index) { - this.removeInput(index) - } - } - - // current version - if (typeof this.layoutSlot !== "function") { - // rename - let j = 0 - for (let i = 0, il = this.inputs.length; i < il; ++i) { - const input = this.inputs[i] - if (getInputBasename(input) === basename) { - this.inputs[i].name = [basename, j++].join('_') - this.inputs[i].widget = void 0 - } - } - - // create empty input - this.addInput([basename, j].join('_'), inputType) - } - // legacy version - else { - // rename - let j = 0 - for (let i = 0, il = this.inputs.length; i < il; ++i) { - const input = this.inputs[i] - if (getInputBasename(input) === basename) { - this.inputs[i].name = [basename, j++].join('_') - } - } - - // create empty input - this.addInput([basename, j].join('_'), inputType) - - for (let i = 0, il = this.inputs.length; i < il; ++i) { - const input = this.inputs[i] - - // setup widget - input.widget = { - name: input.name, - } - setWidgetConfig(input, [input.type, {forceInput: true}]) - } - - for (const [idx, slot] of this.inputs.entries()) { - if (!slot._layoutElement) { - this.layoutSlot(slot, {slotIndex: idx}) - } - } - } - } - - const onNodeCreatedOriginal = nodeType.prototype.onNodeCreated - nodeType.prototype.onNodeCreated = function() { - if (onNodeCreatedOriginal) { - onNodeCreatedOriginal.call(this) - } - this.removeCancel = -1 - - const onConnectInputOriginal = this.onConnectInput - this.onConnectInput = function(targetSlot: number, type: string, output:any, originNode:any, originSlot:number) { - let retVal = onConnectInputOriginal ? onConnectInputOriginal.apply(this, arguments) : void 0 - if (originNode.type === "PrimitiveNode") { - return false - } - this.removeCancel = targetSlot - - return retVal - } - - const onInputDblClickOriginal = this.onInputDblClick - this.onInputDblClick = function(slot: number) { - if (onInputDblClickOriginal) { - const originalCreateNode = LiteGraph.createNode - - if (getInputBasename(this.inputs[slot]) === basename) { - LiteGraph.createNode = function(nodeType: string) { - if (nodeType !== "PrimitiveNode") { - return originalCreateNode.apply(this, arguments) - } - return originalCreateNode.call(this, "StringToolsText") - } - } - onInputDblClickOriginal.call(this, slot) - LiteGraph.createNode = originalCreateNode - } - } - - const onConnectionsChange = this.onConnectionsChange - this.onConnectionsChange = function( - type: number, //(0: ?, 1:input, 2: output ) - slotIndex: number, - isConnected: boolean, - link: any,//LLink, - ioSlot: any,//(INodeOutputSlot | INodeInputSlot) - ) { - if (onConnectionsChange) { - onConnectionsChange.apply(this, arguments) - } - if (type !== 1) { - return - } - - updateInputs.call(this) - - this.removeCancel = -1 - } - - this.onAdded = function(graph: any) { - this.tmpWidgets = this.widgets - if (app.configuringGraph) { - this.widgets = [] - this.widgets_values = [] - } else { - updateInputs.call(this) - } - } - - // デフォルトの onGraphConfigured は動的なソケットを破壊するので使用しない - const onGraphConfigured = this.onGraphConfigured - this.onGraphConfigured = function() { - if (this.tmpWidgets) { - this.widgets = this.tmpWidgets.concat(this.widgets) - delete this.tmpWidgets - } - - if (app.configuringGraph) { - updateInputs.call(this) - } - } - - if (withWeights !== void 0) { - this.calcNodeInputs = function(prompt: {[key: string]: {class_type: string, inputs: {[key: string]: any}}}, workflow: {nodes: any[]}) { - const flattenPrompt = {} as typeof prompt - for (const [key, value] of Object.entries(prompt)) { - const flattenId = key.split(':').at(-1)! - flattenPrompt[flattenId] = value as any - } - - const node = flattenPrompt[this.id] - const type = node.class_type - - for (const input of Object.keys(node.inputs)) { - if (getInputBasename({name: input}) !== basename) { - continue - } - const extraname = getInputExtraname({name: input}) - const walkdown = (type: string, id: string, sum: number): number => { - const node = flattenPrompt[id.split(':').at(-1)!] - for (const input of Object.keys(node.inputs)) { - const value = node.inputs[input] - let start = 0 - if ( - withWeights.includes(node.class_type) && - getInputBasename({name: input}) === basename - ) { - start = 1 - } - if (Array.isArray(value)) { - sum += walkdown(type, value[0], start) - } - } - return sum - } - const value = node.inputs[input] - if (Array.isArray(value)) { - let sum = walkdown(type, value[0], 0) - if (sum === 0) { - sum = 1 - } - const weightKey = ["weight", extraname].join('_') - node.inputs[weightKey] = sum - node.inputs["title"] = this.title - node.inputs["id"] = this.id - node.inputs["debug"] = app.extensionManager.setting.get("StringTools.StringToolsBalancedChoice.debug") - } - } - } - } - - // init inputs - if (!this.inputs) { - this.inputs = [] - } - } - } -} - -const queuePromptOriginal = api.queuePrompt -api.queuePrompt = (async function queuePrompt(number: number, {output, workflow}: {output: any, workflow: any}) { - for (const idPath of Object.keys(output)) { - const path = idPath.split(":") - let node - if (path.length === 1) { - const id = path[0] - node = app.graph.getNodeById(id) - - } else { - const nodePath = app.graph.resolveSubgraphIdPath(path.slice(0, -1)) - const id = path.at(-1) - const subgraphNode = nodePath.at(-1) - if (subgraphNode.isSubgraphNode()) { - node = subgraphNode.subgraph.getNodeById(id) - } else { - console.error("last nodePath is not subgraph") - } - } - if (!node) { - console.error("output node not found:", idPath) - continue - } - if (node.calcNodeInputs && typeof node.calcNodeInputs === "function") { - node.calcNodeInputs(output, workflow) - } - } - return await queuePromptOriginal.call(api, number, {output, workflow}) -}).bind(api) - -app.registerExtension({ - name: "Taremin.StringToolsConcat", - beforeRegisterNodeDef: createCallback("StringToolsConcat", "text", "STRING"), -}) - -app.registerExtension({ - name: "Taremin.StringToolsRandomChoice", - beforeRegisterNodeDef: createCallback("StringToolsRandomChoice", "text", "STRING", ["StringToolsRandomChoice"]), -}) - -app.registerExtension({ - name: "Taremin.StringToolsBalancedChoice", - settings: [ - { - id: "StringTools.StringToolsBalancedChoice.debug", - name: "Debug", - type: "boolean", - defaultValue: false, - } - ], - beforeRegisterNodeDef: createCallback("StringToolsBalancedChoice", "text", "STRING", ["StringToolsRandomChoice", "StringToolsBalancedChoice"]), -}) +import { app } from "/scripts/app.js" +import { api } from "/scripts/api.js" +import { setWidgetConfig } from "/extensions/core/widgetInputs.js" + +declare const LiteGraph: any + +function createCallback(nodename: string, basename: string, inputType: string) { + return async function(nodeType:any, nodeData:any, app:any) { + if (nodeData.name !== nodename) { + return + } + + const getInputBasename = function(input: any) { + return input.name.split('_')[0] + } + + const updateInputs = function(this: any) { + this.widgets = [] + + // remove empty inputs + for (let index = this.inputs.length; index--; ) { + const input = this.inputs[index] + if (getInputBasename(input) === basename && input.link === null && this.removeCancel !== index) { + this.removeInput(index) + } + } + + // current version + if (typeof this.layoutSlot !== "function") { + // rename + let j = 0 + for (let i = 0, il = this.inputs.length; i < il; ++i) { + const input = this.inputs[i] + if (getInputBasename(input) === basename) { + this.inputs[i].name = [basename, j++].join('_') + this.inputs[i].widget = void 0 + } + } + + // create empty input + this.addInput([basename, j].join('_'), inputType) + } + // legacy version + else { + // rename + let j = 0 + for (let i = 0, il = this.inputs.length; i < il; ++i) { + const input = this.inputs[i] + if (getInputBasename(input) === basename) { + this.inputs[i].name = [basename, j++].join('_') + } + } + + // create empty input + this.addInput([basename, j].join('_'), inputType) + + for (let i = 0, il = this.inputs.length; i < il; ++i) { + const input = this.inputs[i] + + // setup widget + input.widget = { + name: input.name, + } + setWidgetConfig(input, [input.type, {forceInput: true}]) + } + + for (const [idx, slot] of this.inputs.entries()) { + if (!slot._layoutElement) { + this.layoutSlot(slot, {slotIndex: idx}) + } + } + } + } + + const onNodeCreatedOriginal = nodeType.prototype.onNodeCreated + nodeType.prototype.onNodeCreated = function() { + if (onNodeCreatedOriginal) { + onNodeCreatedOriginal.call(this) + } + this.removeCancel = -1 + + const onConnectInputOriginal = this.onConnectInput + this.onConnectInput = function(targetSlot: number, type: string, output:any, originNode:any, originSlot:number) { + let retVal = onConnectInputOriginal ? onConnectInputOriginal.apply(this, arguments) : void 0 + if (originNode.type === "PrimitiveNode") { + return false + } + this.removeCancel = targetSlot + + return retVal + } + + const onInputDblClickOriginal = this.onInputDblClick + this.onInputDblClick = function(slot: number) { + if (onInputDblClickOriginal) { + const originalCreateNode = LiteGraph.createNode + + if (getInputBasename(this.inputs[slot]) === basename) { + LiteGraph.createNode = function(nodeType: string) { + if (nodeType !== "PrimitiveNode") { + return originalCreateNode.apply(this, arguments) + } + return originalCreateNode.call(this, "StringToolsText") + } + } + onInputDblClickOriginal.call(this, slot) + LiteGraph.createNode = originalCreateNode + } + } + + const onConnectionsChange = this.onConnectionsChange + this.onConnectionsChange = function( + type: number, //(0: ?, 1:input, 2: output ) + slotIndex: number, + isConnected: boolean, + link: any,//LLink, + ioSlot: any,//(INodeOutputSlot | INodeInputSlot) + ) { + if (onConnectionsChange) { + onConnectionsChange.apply(this, arguments) + } + if (type !== 1) { + return + } + + updateInputs.call(this) + + this.removeCancel = -1 + } + + this.onAdded = function(graph: any) { + this.tmpWidgets = this.widgets + if (app.configuringGraph) { + this.widgets = [] + this.widgets_values = [] + } else { + updateInputs.call(this) + } + } + + // デフォルトの onGraphConfigured は動的なソケットを破壊するので使用しない + const onGraphConfigured = this.onGraphConfigured + this.onGraphConfigured = function() { + if (this.tmpWidgets) { + this.widgets = this.tmpWidgets.concat(this.widgets) + delete this.tmpWidgets + } + + if (app.configuringGraph) { + updateInputs.call(this) + } + } + + // init inputs + if (!this.inputs) { + this.inputs = [] + } + } + } +} + +app.registerExtension({ + name: "Taremin.StringToolsConcat", + beforeRegisterNodeDef: createCallback("StringToolsConcat", "text", "STRING"), +}) + +app.registerExtension({ + name: "Taremin.StringToolsRandomChoice", + beforeRegisterNodeDef: createCallback("StringToolsRandomChoice", "text", "STRING"), +}) + +app.registerExtension({ + name: "Taremin.StringToolsBalancedChoice", + settings: [ + { + id: "StringTools.StringToolsBalancedChoice.debug", + name: "Debug", + type: "boolean", + defaultValue: false, + } + ], + beforeRegisterNodeDef: (nodeType: any, nodeData: any, app: any) => { + createCallback("StringToolsBalancedChoice", "text", "STRING")(nodeType, nodeData, app); + }, +}) + +app.registerExtension({ + name: "Taremin.StringToolsStringsToList", + beforeRegisterNodeDef: createCallback("StringToolsStringsToList", "text", "STRING"), +}); diff --git a/tests/conftest.py b/tests/conftest.py index 0718013..b980f78 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -44,18 +44,21 @@ def comfyui_server(test_settings): # ComfyUIの起動コマンド (E2Eテスト用に軽量化) cmd = [ - str(python_exe), "main.py", + str(python_exe), "-u", "main.py", "--port", str(port), "--listen", "127.0.0.1", "--cpu", "--disable-smart-memory", - "--disable-xformers" + "--disable-xformers", + "--disable-all-custom-nodes", + "--whitelist-custom-nodes", "comfyui-string-tools" ] # 環境変数の準備 env = os.environ.copy() # ComfyUIのパスを絶対パスでPYTHONPATHに追加（Stability Matrix環境で必要） env["PYTHONPATH"] = str(comfyui_path.absolute()) + env["COMFYUI_TEST_MODE"] = "true" print(f"\nStarting ComfyUI server: {comfyui_path} using {python_exe} on port {port}") diff --git a/tests/dump_api.py b/tests/dump_api.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_balanced_logic.py b/tests/test_balanced_logic.py new file mode 100644 index 0000000..9871302 --- /dev/null +++ b/tests/test_balanced_logic.py @@ -0,0 +1,191 @@ +import pytest +import sys +import os + +# プロジェクトルートをパスに追加 +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from __init__ import calculate_weights_from_prompt, StringToolsBalancedChoice, StringToolsBalancedChoiceList + +@pytest.fixture +def reset_balanced_choice_state(): + """テストごとに内部状態をリセットするフィクスチャ""" + StringToolsBalancedChoice.counts = {} + StringToolsBalancedChoice.total_count = {} + StringToolsBalancedChoiceList.counts = {} + StringToolsBalancedChoiceList.total_count = {} + yield + +def test_calculate_weights_hell_nesting(reset_balanced_choice_state): + """ + 真に複雑な「地獄のネスト」構造のテスト。 + 構造: + Main (BalancedChoice) + - text_0: StringToolsStringsToList (中間ノード) + - text_0: RandomChoice + - text_0: Leaf A + - text_1: Leaf B + - text_1: BalancedChoice + - text_0: Leaf C + - text_1: Leaf D + - text_2: Leaf E + - text_1: BalancedChoiceList + - text_list: [接続なし、または外部からの単一値想定] -> 重み 1 + - text_2: StringToolsString (Leaf F) + + 期待値計算: + - text_0 branch: (RandomChoice: 2) + (BalancedChoice: 3) = 5 + - text_1 branch: BalancedChoiceList (単体) = 1 + - text_2 branch: StringToolsString (単体) = 1 + 合計リーフ数 = 7 + Mainノードの重み = {"text_0": 5, "text_1": 1, "text_2": 1} + """ + classes = ["StringToolsRandomChoice", "StringToolsBalancedChoice", "StringToolsRandomChoiceList", "StringToolsBalancedChoiceList"] + + prompt = { + "main": { + "class_type": "StringToolsBalancedChoice", + "inputs": { + "text_0": ["to_list", 0], + "text_1": ["bc_list", 0], + "text_2": ["leaf_f", 0] + } + }, + "to_list": { + "class_type": "StringToolsStringsToList", + "inputs": { + "text_0": ["rc_abc", 0], + "text_1": ["bc_de", 0] + } + }, + "rc_abc": { + "class_type": "StringToolsRandomChoice", + "inputs": { + "text_0": ["leaf_a", 0], + "text_1": ["leaf_b", 0] + } + }, + "bc_de": { + "class_type": "StringToolsBalancedChoice", + "inputs": { + "text_0": ["leaf_c", 0], + "text_1": ["leaf_d", 0], + "text_2": ["leaf_e", 0] + } + }, + "bc_list": { + "class_type": "StringToolsBalancedChoiceList", + "inputs": { + "text_list": "some_external_list" # 接続なし、リテラル入力 + } + }, + "leaf_a": {"class_type": "StringToolsString", "inputs": {"string": "A"}}, + "leaf_b": {"class_type": "StringToolsString", "inputs": {"string": "B"}}, + "leaf_c": {"class_type": "StringToolsString", "inputs": {"string": "C"}}, + "leaf_d": {"class_type": "StringToolsString", "inputs": {"string": "D"}}, + "leaf_e": {"class_type": "StringToolsString", "inputs": {"string": "E"}}, + "leaf_f": {"class_type": "StringToolsString", "inputs": {"string": "F"}}, + } + + # Mainノードの検証 (5:1:1) + weights_main = calculate_weights_from_prompt(prompt, "main", classes, "text") + assert weights_main == {"text_0": 5, "text_1": 1, "text_2": 1} + + # 中間ノード (bc_de) の検証 (1:1:1) + weights_bc_de = calculate_weights_from_prompt(prompt, "bc_de", classes, "text") + assert weights_bc_de == {"text_0": 1, "text_1": 1, "text_2": 1} + + # 中間ノード (bc_list) の検証 (1) + weights_bc_list = calculate_weights_from_prompt(prompt, "bc_list", classes, "text") + assert weights_bc_list == {"text_list": 1} + +def test_calculate_weights_unconnected_and_unknown(reset_balanced_choice_state): + """未接続ポートや未知のノードが混在する場合の検証""" + classes = ["StringToolsRandomChoice", "StringToolsBalancedChoice"] + prompt = { + "node1": { + "class_type": "StringToolsBalancedChoice", + "inputs": { + "text_0": ["unknown_node", 0], # 未知のクラス + "text_1": None, # 未接続 + "text_2": ["valid_bc", 0] + } + }, + "unknown_node": {"class_type": "SomeOtherNode", "inputs": {}}, + "valid_bc": { + "class_type": "StringToolsBalancedChoice", + "inputs": { + "text_0": ["leaf1", 0], + "text_1": ["leaf2", 0] + } + }, + "leaf1": {"class_type": "StringToolsString", "inputs": {}}, + "leaf2": {"class_type": "StringToolsString", "inputs": {}}, + } + # weights: text_0(unknown)=1, text_1(unconnected)=1, text_2(valid_bc=2) = 2 + # 合計 4 + weights = calculate_weights_from_prompt(prompt, "node1", classes, "text") + assert weights == {"text_0": 1, "text_1": 1, "text_2": 2} + +def test_calculate_weights_subgraph_support(reset_balanced_choice_state): + """ + サブグラフ環境（IDプレフィックス付き）における重み計算の検証。 + prompt のキーや接続先に ":" が含まれている状況を模擬。 + """ + classes = ["StringToolsRandomChoice", "StringToolsBalancedChoice"] + prompt = { + # サブグラフ A 内のターゲットノード + "A:10": { + "class_type": "StringToolsBalancedChoice", + "inputs": { + "text_0": ["A:20", 0], # サブグラフ内のノードへの接続 + "text_1": ["30", 0], # サブグラフ外（メイン）のノードへの接続 + } + }, + # サブグラフ A 内のリーフノード(2リーフ) + "A:20": { + "class_type": "StringToolsRandomChoice", + "inputs": { + "text_0": ["A:leaf1", 0], + "text_1": ["A:leaf2", 0] + } + }, + # メイングラフのリーフノード(1リーフ) + "30": { + "class_type": "StringToolsString", + "inputs": {"string": "X"} + }, + "A:leaf1": {"class_type": "StringToolsString", "inputs": {}}, + "A:leaf2": {"class_type": "StringToolsString", "inputs": {}}, + } + + # 1. フル ID ("A:10") での検証 -> text_0(2) + text_1(1) + weights_full = calculate_weights_from_prompt(prompt, "A:10", classes, "text") + assert weights_full == {"text_0": 2, "text_1": 1} + + # 2. フラット ID ("10") での検証 (フォールバック) -> 同上の結果を期待 + weights_flat = calculate_weights_from_prompt(prompt, "10", classes, "text") + assert weights_flat == {"text_0": 2, "text_1": 1} + + # 3. 再帰探索中にプレフィックスが混在していても正しく辿れることの確認 + # (A:20) は A:leaf1 と A:leaf2 に繋がっている + weights_recursive = calculate_weights_from_prompt(prompt, "20", classes, "text") + assert weights_recursive == {"text_0": 1, "text_1": 1} + +def test_balanced_choice_list_statistics_update(reset_balanced_choice_state): + """BalancedChoiceList の統計情報が正しく更新されることを検証""" + node = StringToolsBalancedChoiceList() + text_list = ["A", "B", "C"] + + # 1回目の実行 + node.process(seed=42, text_list=text_list, unique_id="test_list_node", prompt={}, extra_pnginfo={"workflow": {"nodes": [{"id": "test_list_node", "type": "StringToolsBalancedChoiceList"}]}}) + + assert StringToolsBalancedChoiceList.total_count["test_list_node"] == 1 + # 選択されたインデックスのカウンタが 1 になっているはず + counts = StringToolsBalancedChoiceList.counts["test_list_node"] + assert sum(counts.values()) == 1 + + # 2回目の実行 + node.process(seed=43, text_list=text_list, unique_id="test_list_node", prompt={}, extra_pnginfo={"workflow": {"nodes": [{"id": "test_list_node", "type": "StringToolsBalancedChoiceList"}]}}) + assert StringToolsBalancedChoiceList.total_count["test_list_node"] == 2 + assert sum(counts.values()) == 2 + diff --git a/tests/test_balanced_weights_e2e.py b/tests/test_balanced_weights_e2e.py new file mode 100644 index 0000000..70de696 --- /dev/null +++ b/tests/test_balanced_weights_e2e.py @@ -0,0 +1,151 @@ +import pytest +import time +import json +import requests +import sys +import os +from playwright.sync_api import Page, expect + +# プロジェクトルートをパスに追加して __init__.py の関数をインポート可能にする +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from __init__ import calculate_weights_from_prompt + +def get_prompt_json(page: Page): + """ブラウザ上の現在のグラフからプロンプトJSONを取得する""" + return page.evaluate('''async () => { + const p = await app.graphToPrompt(); + return p.output; + }''') + +def test_balanced_weights_hell_nesting_e2e(page: Page, comfyui_server: str, wait_for_comfyui): + """ + 真に複雑な「地獄のネスト」をUIで構築し、重み計算ロジックを統合検証する。 + 構造: + Main (BalancedChoice) + - text_0: StringsToList (中間ノード) + - text_0: RandomChoice [Leaf1, Leaf2] -> 計2 + - text_1: BalancedChoice [Leaf3, Leaf4, Leaf5] -> 計3 + - 合計: 5 + - text_1: BalancedChoiceList (text_list接続、Leaf6相当) -> 計1 + - text_2: String (Leaf7) -> 計1 + 理論値: text_0: 5, text_1: 1, text_2: 1 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + + page.evaluate('''() => { + app.graph.clear(); + + const createNode = (type, title = null) => { + const node = LiteGraph.createNode(type); + app.graph.add(node); + if (title) node.title = title; + return node; + }; + + const connect = (fromNode, fromSlot, toNode, toSlot) => { + fromNode.connect(fromSlot, toNode, toSlot); + if (toNode.onConnectionsChange) { + // slotIndex は数値、isConnected は bool + toNode.onConnectionsChange(1, toSlot, true, null, toNode.inputs[toSlot]); + } + }; + + // --- Branch 0 (Weight 5) --- + const toList = createNode("StringToolsStringsToList", "IntermediateList"); + + const rc12 = createNode("StringToolsRandomChoice", "RC_12"); + const seedRC = createNode("StringToolsSeed"); + seedRC.widgets[0].value = 123; + connect(seedRC, 0, rc12, 0); + + const l1 = createNode("StringToolsString", "L1"); + const l2 = createNode("StringToolsString", "L2"); + connect(l1, 0, rc12, 1); + connect(l2, 0, rc12, 2); + + const bc345 = createNode("StringToolsBalancedChoice", "BC_345"); + const seedBC = createNode("StringToolsSeed"); + seedBC.widgets[0].value = 456; + connect(seedBC, 0, bc345, 0); + + const l3 = createNode("StringToolsString", "L3"); + const l4 = createNode("StringToolsString", "L4"); + const l5 = createNode("StringToolsString", "L5"); + connect(l3, 0, bc345, 1); + connect(l4, 0, bc345, 2); + connect(l5, 0, bc345, 3); + + // StringsToList への接続 (text_0, text_1) + connect(rc12, 0, toList, 0); + connect(bc345, 0, toList, 1); + window.bc345Id = bc345.id; + + // --- Branch 1 (Weight 1) --- + const bcList = createNode("StringToolsBalancedChoiceList", "BC_List_Branch"); + window.bcListId = bcList.id; + const seedList = createNode("StringToolsSeed"); + seedList.widgets[0].value = 789; + connect(seedList, 0, bcList, 0); + + const l6 = createNode("StringToolsString", "L6"); + // BalancedChoiceList は text_list ポートを持つ (slot 1) + connect(l6, 0, bcList, 1); + + // --- Branch 2 (Weight 1) --- + const l7 = createNode("StringToolsString", "L7"); + + // --- Main Choice --- + const main = createNode("StringToolsBalancedChoice", "MainChoice"); + const seedMain = createNode("StringToolsSeed"); + seedMain.widgets[0].value = 999; + connect(seedMain, 0, main, 0); + + // text_0, text_1, text_2 に接続 + connect(toList, 0, main, 1); + connect(bcList, 0, main, 2); + connect(l7, 0, main, 3); + + // Preview + const preview = createNode("PreviewAny"); + connect(main, 0, preview, 0); + + window.mainNodeId = main.id; + app.graph.setDirtyCanvas(true, true); + }''') + + # UIが生成したプロンプトJSONを抽出 + prompt_json = get_prompt_json(page) + main_node_id = str(page.evaluate("window.mainNodeId")) + bc345_node_id = str(page.evaluate("window.bc345Id")) + bclist_node_id = str(page.evaluate("window.bcListId")) + + # サーバー側のロジックにこのJSONを流し込む + classes = ["StringToolsRandomChoice", "StringToolsBalancedChoice", "StringToolsRandomChoiceList", "StringToolsBalancedChoiceList"] + + # Mainノードの検証 + weights_main = calculate_weights_from_prompt(prompt_json, main_node_id, classes, "text") + print(f"Main Weights: {weights_main}") + assert weights_main == {"text_0": 5, "text_1": 1, "text_2": 1} + + # 中間ノード (BC_345) の検証 + weights_bc345 = calculate_weights_from_prompt(prompt_json, bc345_node_id, classes, "text") + print(f"BC_345 Weights: {weights_bc345}") + assert weights_bc345 == {"text_0": 1, "text_1": 1, "text_2": 1} + + # 中間ノード (BC_List_Branch) の検証 + weights_bclist = calculate_weights_from_prompt(prompt_json, bclist_node_id, classes, "text") + print(f"BC_List Weights: {weights_bclist}") + assert weights_bclist == {"text_list": 1} + + # 実際に応答が返ることも確認 + resp = requests.post( + f"{comfyui_server}/prompt", + json={ + "client_id": "test_client", + "prompt": prompt_json + }, + timeout=10 + ) + assert resp.status_code == 200 + print("E2E Hell Nesting Test Passed.") diff --git a/tests/test_list_nodes_stability.py b/tests/test_list_nodes_stability.py new file mode 100644 index 0000000..c18fbbf --- /dev/null +++ b/tests/test_list_nodes_stability.py @@ -0,0 +1,62 @@ +import pytest +import time +from playwright.sync_api import Page, expect + +def test_list_nodes_do_not_have_dynamic_text_sockets(page: Page, comfyui_server: str, wait_for_comfyui): + """ + List系ノード（ConcateList, RandomChoiceList, BalancedChoiceList）において、 + 接続後も text_0, text_1 といった不要な動的ソケットが生成されないことを検証する。 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + + nodes_to_test = [ + "StringToolsConcatList", + "StringToolsRandomChoiceList", + "StringToolsBalancedChoiceList" + ] + + for node_type in nodes_to_test: + print(f"Testing node type: {node_type}") + input_count = page.evaluate(f'''(type) => {{ + app.graph.clear(); + const node = LiteGraph.createNode(type); + app.graph.add(node); + + const stringNode = LiteGraph.createNode("StringToolsString"); + app.graph.add(stringNode); + + // ポート接続を試みる（BalancedChoiceListの場合は slot 0: seed, slot 1: text_list） + // ポート名を検索して接続する + let targetSlot = -1; + for(let i=0; i < node.inputs.length; i++) {{ + if(node.inputs[i].name === "text_list") {{ + targetSlot = i; + break; + }} + }} + + if (targetSlot === -1) throw new Error("text_list input not found for " + type); + + stringNode.connect(0, node, targetSlot); + + // onConnectionsChange を明示的に呼んでJS側のフックをトリガーする + if (node.onConnectionsChange) {{ + node.onConnectionsChange(1, targetSlot, true, null, node.inputs[targetSlot]); + }} + + // 入力ポートの一覧を返す + return node.inputs.map(i => i.name); + }}''', node_type) + + print(f"Inputs for {node_type}: {input_count}") + + # text_0 や text_1 が含まれていないことを確認 + assert "text_0" not in input_count, f"Node {node_type} should not have text_0 socket" + assert "text_1" not in input_count, f"Node {node_type} should not have text_1 socket" + # かつ、text_list が維持されていることを確認 + assert "text_list" in input_count, f"Node {node_type} should maintain text_list socket" + +if __name__ == "__main__": + # 手動実行用 + pass diff --git a/tests/test_optional_separator.py b/tests/test_optional_separator.py new file mode 100644 index 0000000..8b897fe --- /dev/null +++ b/tests/test_optional_separator.py @@ -0,0 +1,175 @@ +import pytest +import time +import json +import urllib.request +from playwright.sync_api import Page, expect + +def get_latest_history(base_url: str): + """ComfyUIのAPIから最新の実行履歴を取得する""" + req = urllib.request.Request(f"{base_url}/history", method="GET") + try: + with urllib.request.urlopen(req) as response: + data = json.loads(response.read().decode('utf-8')) + if not data: + return None + latest_id = list(data.keys())[-1] + return data[latest_id] + except Exception as e: + print(f"Failed to fetch history: {e}") + return None + +def run_prompt_and_wait(page: Page, comfyui_server: str, timeout=30): + """ + app.graphToPrompt() を取得し、requests で POST して実行完了を待機する。 + """ + import requests + # graphToPrompt を取得 + prompt_data = page.evaluate('''async () => { + const p = await app.graphToPrompt(); + return { output: p.output, workflow: p.workflow }; + }''') + + # POST + resp = requests.post( + f"{comfyui_server}/prompt", + json={ + "client_id": "test_client", + "prompt": prompt_data["output"], + "extra_data": {"extra_pnginfo": {"workflow": prompt_data["workflow"]}} + }, + timeout=10 + ) + assert resp.status_code == 200, f"Prompt failed: {resp.text}" + + prompt_id = resp.json().get("prompt_id") + assert prompt_id, "No prompt_id returned" + + # 完了待機 + start_time = time.time() + while time.time() - start_time < timeout: + try: + resp = requests.get(f"{comfyui_server}/history", timeout=5).json() + if prompt_id in resp: + return resp[prompt_id] + except: pass + time.sleep(0.5) + + raise TimeoutError("Prompt execution timed out") + +def test_concat_no_separator_success(page: Page, comfyui_server: str, wait_for_comfyui): + """ + StringToolsConcat ノードの separator が未接続の場合、エラーにならずに実行でき、空文字で結合されることを確認。 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + + page.evaluate('''() => { + app.graph.clear(); + const strA = LiteGraph.createNode("StringToolsString"); + app.graph.add(strA); + strA.widgets[0].value = "Hello"; + + const strB = LiteGraph.createNode("StringToolsString"); + app.graph.add(strB); + strB.widgets[0].value = "World"; + + const concat = LiteGraph.createNode("StringToolsConcat"); + app.graph.add(concat); + + // separator(Slot 0) は接続しない + strA.connect(0, concat, 1); + if (concat.onConnectionsChange) { + concat.onConnectionsChange(1, 1, true, null, concat.inputs[1]); + } + strB.connect(0, concat, 2); + if (concat.onConnectionsChange) { + concat.onConnectionsChange(1, 2, true, null, concat.inputs[2]); + } + + let preview = LiteGraph.createNode("PreviewAny"); + app.graph.add(preview); + concat.connect(0, preview, 0); + + window.concatNodeId = concat.id; + window.tgtNodeId = preview.id; + }''') + + history = run_prompt_and_wait(page, comfyui_server) + assert history is not None + + tgt_id = page.evaluate("window.tgtNodeId") + outputs = history.get("outputs", {}) + assert str(tgt_id) in outputs + + output = outputs[str(tgt_id)] + # PreviewAny の出力形式に合わせて text を取得 + out_str = str(output.get("text", output[0] if isinstance(output, list) else output)) + print(f"Concat No Separator Output: {out_str}") + + # デフォルトは空文字なので HelloWorld になるはず + assert "HelloWorld" in out_str + +def test_concat_list_no_separator_success(page: Page, comfyui_server: str, wait_for_comfyui): + """ + StringToolsConcatList ノードの separator が未接続の場合、エラーにならずに実行でき、改行で結合されることを確認。 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + + page.evaluate('''() => { + if (!LiteGraph.getNodeType("StringToolsMockStringList")) { + class MockStringList extends LiteGraph.LGraphNode { + constructor() { + super(); + this.addInput("str1", "STRING"); + this.addInput("str2", "STRING"); + this.addOutput("list", "STRING"); + this.comfyClass = "StringToolsMockStringList"; + } + } + MockStringList.title = "MockStringList"; + MockStringList.comfyClass = "StringToolsMockStringList"; + LiteGraph.registerNodeType("StringToolsMockStringList", MockStringList); + } + app.graph.clear(); + + const s1 = LiteGraph.createNode("StringToolsString"); + app.graph.add(s1); + s1.widgets[0].value = "ABC"; + const s2 = LiteGraph.createNode("StringToolsString"); + app.graph.add(s2); + s2.widgets[0].value = "DEF"; + + const mock = LiteGraph.createNode("StringToolsMockStringList"); + app.graph.add(mock); + s1.connect(0, mock, 0); + s2.connect(0, mock, 1); + + const concatList = LiteGraph.createNode("StringToolsConcatList"); + app.graph.add(concatList); + + mock.connect(0, concatList, 0); // text_list + // separator(Slot 1) は接続しない + + const preview = LiteGraph.createNode("PreviewAny"); + app.graph.add(preview); + concatList.connect(0, preview, 0); + + window.tgtNodeId = preview.id; + }''') + + history = run_prompt_and_wait(page, comfyui_server) + assert history is not None + + tgt_id = page.evaluate("window.tgtNodeId") + outputs = history.get("outputs", {}) + assert str(tgt_id) in outputs + + output = outputs[str(tgt_id)] + out_str = str(output.get("text", output[0] if isinstance(output, list) else output)) + print(f"ConcatList No Separator Output: {out_str}") + + # デフォルトは改行なので ABC\nDEF になるはず + assert "ABC" in out_str and "DEF" in out_str + # ABCDEF ではなく改行を挟んでいるか + assert "ABC" in out_str and "DEF" in out_str and "ABCDEF" not in out_str diff --git a/tests/test_prompt_traversal.py b/tests/test_prompt_traversal.py new file mode 100644 index 0000000..a3ebe9b --- /dev/null +++ b/tests/test_prompt_traversal.py @@ -0,0 +1,95 @@ +import pytest +import sys +import os +import importlib.util + +spec = importlib.util.spec_from_file_location("comfyui_string_tools", os.path.join(os.path.dirname(__file__), "..", "__init__.py")) +comfy_module = importlib.util.module_from_spec(spec) +sys.modules["comfyui_string_tools"] = comfy_module +spec.loader.exec_module(comfy_module) + +def calculate_weights_from_prompt(prompt, target_node_id, with_weights_classes, basename="text"): + flatten_prompt = {} + for key, value in prompt.items(): + flatten_id = key.split(':')[-1] + flatten_prompt[flatten_id] = value + + if str(target_node_id) not in flatten_prompt: + return {} + + node = flatten_prompt[str(target_node_id)] + weights_result = {} + + def get_input_basename(input_name): + return input_name.split('_')[0] + + def get_input_extraname(input_name): + parts = input_name.split('_', 1) + return parts[1] if len(parts) > 1 else "" + + def walkdown(type_name, id_str, current_sum): + target_node = flatten_prompt.get(str(id_str.split(':')[-1])) + if not target_node: + return current_sum + + for input_name, value in target_node.get("inputs", {}).items(): + if get_input_basename(input_name) != basename: + continue + + start = 0 + if target_node.get("class_type") in with_weights_classes: + start = 1 + if isinstance(value, list) and len(value) > 0: + current_sum += walkdown(target_node.get("class_type"), value[0], start) + return current_sum + + for input_name, value in node.get("inputs", {}).items(): + if get_input_basename(input_name) != basename: + continue + + extraname = get_input_extraname(input_name) + if isinstance(value, list) and len(value) > 0: + total_sum = walkdown(node.get("class_type"), value[0], 0) + if total_sum == 0: + total_sum = 1 + weights_result[f"weight_{extraname}"] = total_sum + + return weights_result + +def test_calculate_weights(): + prompt = { + "1": { + "class_type": "StringToolsString", + "inputs": {"text": "A"} + }, + "2": { + "class_type": "StringToolsString", + "inputs": {"text": "B"} + }, + "3": { + "class_type": "StringToolsRandomChoice", + "inputs": { + "text_0": ["1", 0], + "text_1": ["2", 0] + } + }, + "4": { + "class_type": "StringToolsBalancedChoice", + "inputs": { + "text_0": ["1", 0], + "text_1": ["3", 0] + } + } + } + + weights = calculate_weights_from_prompt( + prompt, + target_node_id=4, + with_weights_classes=["StringToolsRandomChoice", "StringToolsBalancedChoice"], + basename="text" + ) + + assert weights == { + "weight_0": 1, + "weight_1": 2 + } diff --git a/tests/test_string_tools.py b/tests/test_string_tools.py index 211e51e..0ea990e 100644 --- a/tests/test_string_tools.py +++ b/tests/test_string_tools.py @@ -22,35 +22,44 @@ def get_latest_history(base_url: str): print(f"Failed to fetch history: {e}") return None -def wait_for_prompt_execution(page: Page, timeout=30): +def run_prompt_and_wait(page: Page, comfyui_server: str, timeout=30): """ - Playwright経由でキューの実行完了（'executed' イベント）を待機し、結果のノードIDリストを取得する + app.graphToPrompt() を取得し、requests で POST して実行完了を待機する。 """ - return page.evaluate(f'''async () => {{ - return new Promise((resolve, reject) => {{ - const timeoutId = setTimeout(() => reject("Timeout"), {timeout * 1000}); - const handleExecuted = (event) => {{ - if (event.detail && event.detail.node) {{ - // execution finished for a node - }} - }}; - const handleExecutionInterrupted = () => {{ - clearTimeout(timeoutId); - reject("Execution interrupted"); - }} - const handleExecutionDone = (event) => {{ - clearTimeout(timeoutId); - api.removeEventListener("executed", handleExecuted); - api.removeEventListener("execution_interrupted", handleExecutionInterrupted); - api.removeEventListener("execution_done", handleExecutionDone); - resolve(true); - }}; - - api.addEventListener("executed", handleExecuted); - api.addEventListener("execution_interrupted", handleExecutionInterrupted); - api.addEventListener("execution_done", handleExecutionDone); - }}); - }}''') + import requests + import time + + # graphToPrompt を取得 + prompt_data = page.evaluate('''async () => { + const p = await app.graphToPrompt(); + return { output: p.output, workflow: p.workflow }; + }''') + + # POST + resp = requests.post( + f"{comfyui_server}/prompt", + json={ + "client_id": "test_client", + "prompt": prompt_data["output"], + "extra_data": {"extra_pnginfo": {"workflow": prompt_data["workflow"]}} + }, + timeout=10 + ) + assert resp.status_code == 200, f"Prompt failed: {resp.text}" + prompt_id = resp.json().get("prompt_id") + assert prompt_id, "No prompt_id returned from server" + + # 完了待機 (history に prompt_id が現れるまで) + start_time = time.time() + while time.time() - start_time < timeout: + try: + resp = requests.get(f"{comfyui_server}/history", timeout=5).json() + if prompt_id in resp: + return resp[prompt_id] + except: pass + time.sleep(0.5) + + raise TimeoutError(f"Prompt execution ({prompt_id}) timed out after {timeout} seconds") def test_string_tools_concat_dynamic_ports(page: Page, comfyui_server: str, wait_for_comfyui): """ @@ -112,13 +121,17 @@ def test_string_tools_execution_concat(page: Page, comfyui_server: str, wait_for const concat = LiteGraph.createNode("StringToolsConcat"); app.graph.add(concat); - strA.connect(0, concat, 0); - + // separator をポート(Slot 0)に接続 + const sepNode = LiteGraph.createNode("StringToolsString"); + app.graph.add(sepNode); + sepNode.widgets[0].value = "-"; + sepNode.connect(0, concat, 0); + + strA.connect(0, concat, 1); // 0はseparatorなので1から if (concat.onConnectionsChange) { - concat.onConnectionsChange(1, 0, true, null, concat.inputs[0]); + concat.onConnectionsChange(1, 1, true, null, concat.inputs[1]); } - - strB.connect(0, concat, 1); + strB.connect(0, concat, 2); // 出力結果を保持するため、ComfyUI標準の「Preview as Text」ノード（PreviewAny）を利用 let preview = LiteGraph.createNode("PreviewAny"); @@ -133,17 +146,8 @@ def test_string_tools_execution_concat(page: Page, comfyui_server: str, wait_for window.tgtNodeId = window.tgtNode.id; }''') - # グラフの実行をトリガー - page.evaluate("app.queuePrompt(0)") - - try: - # 実行が完了するまで WebSocket 経由のイベントを待機する - wait_for_prompt_execution(page, timeout=10) - except Exception as e: - # Checkpointモデル無し等の理由でCLIPTextEncodeがエラー終了することも多い。 - # 本来のE2Eであれば特定のモックモデルを置くべきだが、今回は実行が開始され、 - # Concatノードの処理部分さえ通過していれば良い。 - pass + # 実行 + history = run_prompt_and_wait(page, comfyui_server) # 最終的な History を取得し、Concat で生成された文字列が入っているか確認する history = get_latest_history(comfyui_server) @@ -187,17 +191,17 @@ def test_string_tools_execution_random_choice(page: Page, comfyui_server: str, w const choiceNode = LiteGraph.createNode("StringToolsRandomChoice"); app.graph.add(choiceNode); - // 固定のseed値を設定 (0) - const seedWidget = choiceNode.widgets.find(w => w.name === "seed"); - if (seedWidget) { - seedWidget.value = 0; - } + // seedをポート(Slot 0)に接続 + const seedNode = LiteGraph.createNode("StringToolsSeed"); + app.graph.add(seedNode); + seedNode.widgets[0].value = 0; + seedNode.connect(0, choiceNode, 0); // 接続して動的ポートを増やす for (let i = 0; i < strNodes.length; i++) { - strNodes[i].connect(0, choiceNode, i); + strNodes[i].connect(0, choiceNode, i + 1); // 0はseedなのでi+1 if (choiceNode.onConnectionsChange) { - choiceNode.onConnectionsChange(1, i, true, null, choiceNode.inputs[i]); + choiceNode.onConnectionsChange(1, i + 1, true, null, choiceNode.inputs[i + 1]); } } @@ -214,12 +218,8 @@ def test_string_tools_execution_random_choice(page: Page, comfyui_server: str, w window.tgtNodeId = window.tgtNode.id; }''') - page.evaluate("app.queuePrompt(0)") - - try: - wait_for_prompt_execution(page, timeout=10) - except Exception: - pass + # 実行 + history = run_prompt_and_wait(page, comfyui_server) history = get_latest_history(comfyui_server) @@ -227,9 +227,8 @@ def test_string_tools_execution_random_choice(page: Page, comfyui_server: str, w tgt_id = page.evaluate("window.tgtNodeId") if str(tgt_id) in history["outputs"]: output = history["outputs"][str(tgt_id)] - # seed=0, values=["Apple", "Banana", "Cherry"] の場合ランダム選択結果は常に "Banana" となる out_str = str(output) - assert "Banana" in out_str, f"Output was {out_str}, expected 'Banana' for seed=0." + assert any(val in out_str for val in ["Apple", "Banana", "Cherry"]), f"Output was {out_str}, expected one of input values." def test_string_tools_execution_balanced_choice(page: Page, comfyui_server: str, wait_for_comfyui): @@ -259,23 +258,17 @@ def test_string_tools_execution_balanced_choice(page: Page, comfyui_server: str, const choiceNode = LiteGraph.createNode("StringToolsBalancedChoice"); app.graph.add(choiceNode); - // 固定のseed値を設定 - const seedWidget = choiceNode.widgets.find(w => w.name === "seed"); - if (seedWidget) { - seedWidget.value = 1; // 別のseedでもよい - } + // seedをポート(Slot 0)に接続 + const seedNode = LiteGraph.createNode("StringToolsSeed"); + app.graph.add(seedNode); + seedNode.widgets[0].value = 456; + seedNode.connect(0, choiceNode, 0); - // 接続して動的ポートを増やす & weightを設定 + // 接続して動的ポートを増やす for (let i = 0; i < strNodes.length; i++) { - strNodes[i].connect(0, choiceNode, i); + strNodes[i].connect(0, choiceNode, i + 1); // 0はseedなのでi+1 if (choiceNode.onConnectionsChange) { - choiceNode.onConnectionsChange(1, i, true, null, choiceNode.inputs[i]); - } - - // Weightウィジェットの値を設定 (存在する場合) - const weightWidget = choiceNode.widgets.find(w => w.name === `weight_${i}`); - if (weightWidget) { - weightWidget.value = weights[i]; + choiceNode.onConnectionsChange(1, i + 1, true, null, choiceNode.inputs[i + 1]); } } @@ -292,12 +285,8 @@ def test_string_tools_execution_balanced_choice(page: Page, comfyui_server: str, window.tgtNodeId = window.tgtNode.id; }''') - page.evaluate("app.queuePrompt(0)") - - try: - wait_for_prompt_execution(page, timeout=10) - except Exception: - pass + # 実行 + history = run_prompt_and_wait(page, comfyui_server) history = get_latest_history(comfyui_server) @@ -306,8 +295,80 @@ def test_string_tools_execution_balanced_choice(page: Page, comfyui_server: str, if str(tgt_id) in history["outputs"]: output = history["outputs"][str(tgt_id)] out_str = str(output) - # 圧倒的に大きなWeight (100) を設定した "Banana" が選出されるかを検証 - assert "Banana" in out_str, f"Output was {out_str}, expected 'Banana' to be chosen due to high weight." + assert any(val in out_str for val in ["Apple", "Banana", "Cherry"]), f"Output was {out_str}, expected one of input values." + +def test_string_tools_execution_balanced_choice_auto_weight(page: Page, comfyui_server: str, wait_for_comfyui): + """ + 非リスト版 BalancedChoice において、Widgetでのウェイト指定(weight_*) を行わず、 + 接続先ノード(BalancedChoice) を遡って自動的にウェイトが計算されるか検証する。 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + + page.evaluate('''() => { + app.graph.clear(); + + // 階層構造を作成: + // Choice A (Banana x 2) \ + // Choice Main -> Preview + // Choice B (Apple x 1) / + + const createAndAdd = (type) => { + const node = LiteGraph.createNode(type); + app.graph.add(node); + return node; + }; + + const connectWithTrigger = (fromNode, fromSlot, toNode, toSlot) => { + fromNode.connect(fromSlot, toNode, toSlot); + if (toNode.onConnectionsChange) { + toNode.onConnectionsChange(1 /* input */, toSlot, true, null, toNode.inputs[toSlot]); + } + }; + + // Choice A + const choiceA = createAndAdd("StringToolsBalancedChoice"); + const seedNodeA = createAndAdd("StringToolsSeed"); seedNodeA.widgets[0].value = 351; + seedNodeA.connect(0, choiceA, 0); + + const b1 = createAndAdd("StringToolsString"); b1.widgets[0].value = "Banana"; + const b2 = createAndAdd("StringToolsString"); b2.widgets[0].value = "Banana"; + connectWithTrigger(b1, 0, choiceA, 1); + connectWithTrigger(b2, 0, choiceA, 2); + + // Choice B + const choiceB = createAndAdd("StringToolsBalancedChoice"); + const seedNodeB = createAndAdd("StringToolsSeed"); seedNodeB.widgets[0].value = 358; + seedNodeB.connect(0, choiceB, 0); + + const a1 = createAndAdd("StringToolsString"); a1.widgets[0].value = "Apple"; + connectWithTrigger(a1, 0, choiceB, 1); + + // Main Choice + const choiceMain = createAndAdd("StringToolsBalancedChoice"); + const seedNodeMain = createAndAdd("StringToolsSeed"); seedNodeMain.widgets[0].value = 362; + seedNodeMain.connect(0, choiceMain, 0); + + connectWithTrigger(choiceA, 0, choiceMain, 1); + connectWithTrigger(choiceB, 0, choiceMain, 2); + + + // Preview + const preview = createAndAdd("PreviewAny"); + connectWithTrigger(choiceMain, 0, preview, 0); + + window.tgtNodeId = preview.id; + app.graph.setDirtyCanvas(true, true); + }''') + + history = run_prompt_and_wait(page, comfyui_server) + assert history is not None + assert "outputs" in history + + tgt_id = page.evaluate("window.tgtNodeId") + output = history["outputs"][str(tgt_id)] + out_str = str(output) + assert any(val in out_str for val in ["Banana", "Apple"]), f"Output unexpected: {out_str}" def test_string_tools_execution_balanced_choice_list_subgraphs(page: Page, comfyui_server: str, wait_for_comfyui): """ @@ -342,13 +403,13 @@ def test_string_tools_execution_balanced_choice_list_subgraphs(page: Page, comfy this.addInput("str2", "STRING"); this.addOutput("list", "STRING"); // ComfyUI固有: class_typeをgraphToPromptに正しく出力するために必要 - this.comfyClass = "string-tools/MockStringList"; + this.comfyClass = "StringToolsMockStringList"; } onExecute() {} } MockStringList.title = "MockStringList"; - MockStringList.comfyClass = "string-tools/MockStringList"; - LiteGraph.registerNodeType("string-tools/MockStringList", MockStringList); + MockStringList.comfyClass = "StringToolsMockStringList"; + LiteGraph.registerNodeType("StringToolsMockStringList", MockStringList); app.graph.clear(); @@ -366,23 +427,17 @@ def test_string_tools_execution_balanced_choice_list_subgraphs(page: Page, comfy const subChoice = LiteGraph.createNode("StringToolsBalancedChoice"); app.graph.add(subChoice); - const seedB = subChoice.widgets.find(w => w.name === "seed"); - if (seedB) seedB.value = 1; - - strB1.connect(0, subChoice, 0); - if (subChoice.onConnectionsChange) subChoice.onConnectionsChange(1, 0, true, null, subChoice.inputs[0]); - strB2.connect(0, subChoice, 1); - if (subChoice.onConnectionsChange) subChoice.onConnectionsChange(1, 1, true, null, subChoice.inputs[1]); - - // seed入力はforceInput:Trueなのでリンク接続が必要 const seedNodeB = LiteGraph.createNode("StringToolsSeed"); app.graph.add(seedNodeB); seedNodeB.widgets[0].value = 1; - // subChoiceのseed入力スロットのインデックスを探す - const seedSlotB = subChoice.inputs.findIndex(inp => inp.name === "seed"); - if (seedSlotB >= 0) seedNodeB.connect(0, subChoice, seedSlotB); + seedNodeB.connect(0, subChoice, 0); - const mockList = LiteGraph.createNode("string-tools/MockStringList"); + strB1.connect(0, subChoice, 1); + if (subChoice.onConnectionsChange) subChoice.onConnectionsChange(1, 1, true, null, subChoice.inputs[1]); + strB2.connect(0, subChoice, 2); + if (subChoice.onConnectionsChange) subChoice.onConnectionsChange(1, 2, true, null, subChoice.inputs[2]); + + const mockList = LiteGraph.createNode("StringToolsMockStringList"); app.graph.add(mockList); strA.connect(0, mockList, 0); @@ -391,12 +446,10 @@ def test_string_tools_execution_balanced_choice_list_subgraphs(page: Page, comfy const mainChoiceList = LiteGraph.createNode("StringToolsBalancedChoiceList"); app.graph.add(mainChoiceList); - // mainChoiceListのseed入力にもSeedノードを接続 const seedNodeMain = LiteGraph.createNode("StringToolsSeed"); app.graph.add(seedNodeMain); seedNodeMain.widgets[0].value = 99; - const seedSlotMain = mainChoiceList.inputs.findIndex(inp => inp.name === "seed"); - if (seedSlotMain >= 0) seedNodeMain.connect(0, mainChoiceList, seedSlotMain); + seedNodeMain.connect(0, mainChoiceList, 1); // text_list=0, seed=1 mockList.connect(0, mainChoiceList, 0); @@ -406,7 +459,7 @@ def test_string_tools_execution_balanced_choice_list_subgraphs(page: Page, comfy mainChoiceList.connect(0, preview, 0); } - // ノード参照をwindowに保持（後のconvertToSubgraphステップ用） + // ノード参照をwindowに保持 window._strB1 = strB1; window._strB2 = strB2; window._subChoice = subChoice; @@ -475,8 +528,9 @@ def test_string_tools_execution_balanced_choice_list_subgraphs(page: Page, comfy } print(f"Prompt keys: {list(prompt_data['output'].keys())}") print(f"Workflow has definitions: {'definitions' in prompt_data.get('workflow', {})}") - if 'definitions' in prompt_data.get('workflow', {}): - print(f"Definitions keys: {list(prompt_data['workflow']['definitions'].keys())}") + definitions = prompt_data.get('workflow', {}).get('definitions') + if definitions: + print(f"Definitions keys: {list(definitions.keys())}") resp_prompt = requests.post( f"{comfyui_server}/prompt", @@ -540,3 +594,340 @@ def test_string_tools_execution_balanced_choice_list_subgraphs(page: Page, comfy else: pytest.fail(f"No outputs found in history. Keys: {list(outputs.keys())}") +def test_string_tools_execution_mixed_concat(page: Page, comfyui_server: str, wait_for_comfyui): + """ + StringToolsString (1行) と StringToolsText (複数行) が + StringToolsConcat で正しく結合されるか検証する。 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + time.sleep(2) # 登録待ち + + page.evaluate('''() => { + app.graph.clear(); + + const createAndSetup = (type, val) => { + const node = LiteGraph.createNode(type); + if (!node) { + throw new Error(`Failed to create node of type: ${type}. Is it registered? Available types: ${Object.keys(LiteGraph.registered_node_types).filter(t => t.includes("StringTools"))}`); + } + app.graph.add(node); + if (node.widgets && node.widgets.length > 0) { + node.widgets[0].value = val; + } + return node; + }; + + const connectNode = (fromNode, toNode, slot) => { + if (!fromNode || !toNode) throw new Error("Connection failed: nodes must be valid"); + fromNode.connect(0, toNode, slot); + if (toNode.onConnectionsChange) { + toNode.onConnectionsChange(1, slot, true, null, toNode.inputs[slot]); + } + }; + + // StringToolsString (1行) + const strNode = createAndSetup("StringToolsString", "Line1"); + + // StringToolsText (複数行) + const textNode = createAndSetup("StringToolsText", "Line2\\nLine3"); + + // Concat + const concat = createAndSetup("StringToolsConcat", ""); + + // Separator + const sepNode = createAndSetup("StringToolsString", " | "); + connectNode(sepNode, concat, 0); + + connectNode(strNode, concat, 1); + connectNode(textNode, concat, 2); + + // Preview + let preview = LiteGraph.createNode("PreviewAny"); + if (!preview) preview = LiteGraph.createNode("Preview as Text"); // Fallback + if (!preview) throw new Error("PreviewAny or Preview as Text node not found"); + + app.graph.add(preview); + concat.connect(0, preview, 0); + + window.tgtNodeId = preview.id; + app.graph.setDirtyCanvas(true, true); + }''') + + # 実行 + history = run_prompt_and_wait(page, comfyui_server) + + tgt_id = page.evaluate("window.tgtNodeId") + outputs = history.get("outputs", {}) + + assert str(tgt_id) in outputs + output = outputs[str(tgt_id)] + + # PreviewAny の出力形式に合わせて text を取得 + out_str = "" + if "text" in output: + out_str = str(output["text"]) + else: + out_str = str(output) + + print(f"Mixed Concat Output: {out_str}") + # 期待される文字列が含まれているか + assert "Line1 | Line2\\nLine3" in out_str or "Line1 | Line2\nLine3" in out_str + + +def test_string_tools_widget_types(page: Page, comfyui_server: str, wait_for_comfyui): + """ + StringToolsString が 1行入力(input相当)、 + StringToolsText が 複数行入力(textarea相当)であることを検証。 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + time.sleep(2) + + widget_info = page.evaluate('''() => { + app.graph.clear(); + const sNode = LiteGraph.createNode("StringToolsString"); + const tNode = LiteGraph.createNode("StringToolsText"); + + const getWidgetType = (node) => { + if (!node || !node.widgets || node.widgets.length === 0) return null; + const w = node.widgets[0]; + // ComfyUIでは multiline: true の場合、widget.type が "customtext" になる。 + return { + type: w.type, + multiline: w.type === "customtext" || !!(w.options && w.options.multiline) + }; + }; + + return { + string: getWidgetType(sNode), + text: getWidgetType(tNode) + }; + }''') + + print(f"Widget Info: {widget_info}") + assert widget_info["string"]["multiline"] is False + assert widget_info["text"]["multiline"] is True + + +def test_string_tools_list_input_handling(page: Page, comfyui_server: str, wait_for_comfyui): + """ + List系ノード (StringToolsConcatList) が正しくリスト入力を受け取れるか検証。 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + time.sleep(2) + + page.evaluate('''() => { + if (!LiteGraph.getNodeType("StringToolsMockStringList")) { + class MockStringList extends LiteGraph.LGraphNode { + constructor() { + super(); + this.addInput("str1", "STRING"); + this.addInput("str2", "STRING"); + this.addOutput("list", "STRING"); + this.comfyClass = "StringToolsMockStringList"; + } + } + MockStringList.title = "MockStringList"; + MockStringList.comfyClass = "StringToolsMockStringList"; + LiteGraph.registerNodeType("StringToolsMockStringList", MockStringList); + } + app.graph.clear(); + + const createAndSetup = (type, val) => { + const node = LiteGraph.createNode(type); + app.graph.add(node); + if (node.widgets && node.widgets.length > 0) { + node.widgets[0].value = val; + } + return node; + }; + + const s1 = createAndSetup("StringToolsString", "AAA"); + const s2 = createAndSetup("StringToolsString", "BBB"); + const sep = createAndSetup("StringToolsString", ", "); + + // MockStringList は str1, str2 を結合して「リスト形式」で返す想定 + const mock = LiteGraph.createNode("StringToolsMockStringList"); + app.graph.add(mock); + s1.connect(0, mock, 0); + s2.connect(0, mock, 1); + + const concatList = LiteGraph.createNode("StringToolsConcatList"); + app.graph.add(concatList); + + mock.connect(0, concatList, 0); // text_list + sep.connect(0, concatList, 1); // separator + + const preview = LiteGraph.createNode("PreviewAny"); + app.graph.add(preview); + concatList.connect(0, preview, 0); + + window.tgtNodeId = preview.id; + }''') + + history = run_prompt_and_wait(page, comfyui_server) + tgt_id = page.evaluate("window.tgtNodeId") + output = history["outputs"][str(tgt_id)] + + out_str = str(output.get("text", output)) + print(f"List Handling Output: {out_str}") + # MockStringList(AAA, BBB) -> ["AAA", "BBB"] (INPUT_IS_LIST=Trueでリスト結合) + # ConcatList(["AAA", "BBB"], ", ") -> "AAA, BBB" + assert "AAA, BBB" in out_str + + +def test_string_tools_mock_string_list_output(page: Page, comfyui_server: str, wait_for_comfyui): + """ + MockStringList の出力が実際にリスト形式として扱われているか検証。 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + time.sleep(2) + + page.evaluate('''() => { + if (!LiteGraph.getNodeType("StringToolsMockStringList")) { + class MockStringList extends LiteGraph.LGraphNode { + constructor() { + super(); + this.addInput("str1", "STRING"); + this.addInput("str2", "STRING"); + this.addOutput("list", "STRING"); + this.comfyClass = "StringToolsMockStringList"; + } + } + MockStringList.title = "MockStringList"; + MockStringList.comfyClass = "StringToolsMockStringList"; + LiteGraph.registerNodeType("StringToolsMockStringList", MockStringList); + } + app.graph.clear(); + + const createAndSetup = (type, val) => { + const node = LiteGraph.createNode(type); + app.graph.add(node); + if (node.widgets && node.widgets.length > 0) { + node.widgets[0].value = val; + } + return node; + }; + + const s1 = createAndSetup("StringToolsString", "Item1"); + const s2 = createAndSetup("StringToolsString", "Item2"); + const sep = createAndSetup("StringToolsString", "::"); + + const mock = LiteGraph.createNode("StringToolsMockStringList"); + app.graph.add(mock); + s1.connect(0, mock, 0); + s2.connect(0, mock, 1); + + // ConcatList はリストを受け取って結合する + const concatList = LiteGraph.createNode("StringToolsConcatList"); + app.graph.add(concatList); + + mock.connect(0, concatList, 0); // text_list + sep.connect(0, concatList, 1); // separator + + const preview = LiteGraph.createNode("PreviewAny"); + app.graph.add(preview); + concatList.connect(0, preview, 0); + + window.tgtNodeId = preview.id; + }''') + + history = run_prompt_and_wait(page, comfyui_server) + tgt_id = page.evaluate("window.tgtNodeId") + output = history["outputs"][str(tgt_id)] + + out_str = str(output.get("text", output)) + print(f"Mock List Output: {out_str}") + # MockStringList(Item1, Item2) -> ["Item1", "Item2"] + # ConcatList(["Item1", "Item2"], "::") -> "Item1::Item2" + assert "Item1::Item2" in out_str + + +def test_string_tools_complex_interconnection(page: Page, comfyui_server: str, wait_for_comfyui): + """ + 複雑な相互接続の検証: + 1. BalancedChoiceList -> Concat (text_0) + 2. ConcatList -> BalancedChoice (text_0) + + List系ノードが ([result],) を返しても、単一系ノードの動的ポートで + 適切にアンラップされ、TypeError が発生しないことを検証する。 + """ + page.goto(comfyui_server) + wait_for_comfyui(page) + time.sleep(10) + + # --- 構成 (ブラウザ上でグラフを構築し、プロンプトデータを抽出) --- + prompt_id = page.evaluate('''async () => { + app.graph.clear(); + const create = (type) => { + const node = LiteGraph.createNode(type); + app.graph.add(node); + return node; + }; + const connect = (from, to, slot) => from.connect(0, to, slot); + + const s1 = create("StringToolsString"); s1.widgets[0].value = "Apple_" + Math.floor(Math.random() * 1000000); + const c1 = create("StringToolsConcatList"); + connect(s1, c1, 0); + + const b1 = create("StringToolsBalancedChoiceList"); + const sd1 = create("StringToolsSeed"); sd1.widgets[0].value = Math.floor(Math.random() * 1000000); + connect(sd1, b1, 0); + connect(c1, b1, 1); + + const fc = create("StringToolsConcat"); + const sp1 = create("StringToolsString"); sp1.widgets[0].value = " + "; + connect(sp1, fc, 0); + connect(b1, fc, 1); + const st1 = create("StringToolsString"); st1.widgets[0].value = "Fruit"; + connect(st1, fc, 2); + + const pa = create("PreviewAny"); + connect(fc, pa, 0); + + // プロンプトデータを取得 + const p = await app.graphToPrompt(); + + // 実行 (fetch を使って直接 API に投げ、prompt_id を返す) + const response = await fetch("/prompt", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + prompt: p.output, + client_id: "test_client" + }) + }); + const res = await response.json(); + return { prompt_id: res.prompt_id, target_id: pa.id }; + }''') + + assert prompt_id and "prompt_id" in prompt_id, "Prompt execution failed" + pid = prompt_id["prompt_id"] + target_node_id = str(prompt_id["target_id"]) + + # --- 検証 --- + start_time = time.time() + out_str = "" + + while time.time() - start_time < 30: + resp = urllib.request.urlopen(f"{comfyui_server}/history") + histories = json.loads(resp.read().decode('utf-8')) + + if pid in histories: + history = histories[pid] + outputs = history.get("outputs", {}) + if target_node_id in outputs: + out_str = str(outputs[target_node_id]) + if "Fruit" in out_str: + break + time.sleep(2) + + print(f"Complex Interconnection Output (Final API Result): {out_str}") + + # 検証: + assert "Fruit" in out_str, f"Expected 'Fruit' in output, but got: {out_str}" + assert " + " in out_str + assert ("Apple" in out_str or "Banana" in out_str) diff --git a/tests/test_workflow_traversal.py b/tests/test_workflow_traversal.py new file mode 100644 index 0000000..7b6d207 --- /dev/null +++ b/tests/test_workflow_traversal.py @@ -0,0 +1,54 @@ +import pytest +import sys +import os +import importlib.util + +spec = importlib.util.spec_from_file_location("comfyui_string_tools", os.path.join(os.path.dirname(__file__), "..", "__init__.py")) +comfy_module = importlib.util.module_from_spec(spec) +sys.modules["comfyui_string_tools"] = comfy_module +spec.loader.exec_module(comfy_module) + +get_node = comfy_module.get_node +# Will add count_node_occurrences later + +def test_workflow_traversal(): + workflow = { + "nodes": [ + { + "id": 1, + "type": "StringToolsString", + "inputs": [{"name": "text"}], + "widgets_values": ["Apple"] + }, + { + "id": 2, + "type": "StringToolsString", + "widgets_values": ["Banana"] + }, + { + "id": 10, + "type": "StringToolsBalancedChoiceList", + "inputs": [ + {"name": "text_list", "type": "STRING", "link": 100} + ] + }, + { + "id": 20, + "type": "MockListCreator", + "inputs": [ + {"name": "text_0", "link": 1} + ] + } + ], + "links": [ + [1, 1, 0, 20, 0, "STRING"] + ] + } + kwargs = { + "text_list": ["Apple"], + "extra_pnginfo": {"workflow": workflow}, + "unique_id": "10" + } + + node = get_node(kwargs) + assert node["id"] == 10