diff --git a/__init__.py b/__init__.py index fcd891f..b259dcd 100644 --- a/__init__.py +++ b/__init__.py @@ -40,6 +40,12 @@ 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] + if isinstance(unique_id, list): + unique_id = unique_id[0] + if extra_pnginfo is None or unique_id is None: print( re.sub( @@ -70,7 +76,9 @@ def get_node(kwargs): nodes_by_id = {} for n in graph["nodes"]: nodes_by_id[n["id"]] = n - node = nodes_by_id[current_path[0]] + 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"]]) @@ -82,6 +90,55 @@ def get_node(kwargs): 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"): + 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 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 format(num): return "{:6.2f}".format(num) @@ -195,6 +252,36 @@ class StringToolsConcat: 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): @@ -241,6 +328,44 @@ class StringToolsRandomChoice: return (choice,) +class StringToolsRandomChoiceList: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "text_list": ( + "STRING", + {"forceInput": True}, + ), + "seed": ( + "INT", + { + "default": 0, + "min": 0, + "max": 0xFFFFFFFFFFFFFFFF, + "forceInput": True, + }, + ), + }, + } + + INPUT_IS_LIST = True + RETURN_TYPES = ("STRING",) + FUNCTION = "process" + CATEGORY = "string-tools" + + def process(self, text_list, seed): + if not text_list: + return ("",) + s = seed[0] if isinstance(seed, list) else seed + random.seed(s) + choice = random.choice(text_list) + return (choice,) + + class StringToolsBalancedChoice(StringToolsRandomChoice): total_count = {} counts = {} @@ -276,6 +401,10 @@ 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): + weights = [1] * len(values) + 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 @@ -294,6 +423,134 @@ class StringToolsBalancedChoice(StringToolsRandomChoice): return (values[choice_idx],) +class StringToolsBalancedChoiceList(StringToolsRandomChoiceList): + total_count = {} + counts = {} + debug = False + + @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", + }, + } + + def process(self, text_list, seed, weight_list=None, **kwargs): + # INPUT_IS_LIST = True のため hidden 入力もリストにラップされる + if isinstance(kwargs.get("prompt"), list): + kwargs["prompt"] = kwargs["prompt"][0] + if isinstance(kwargs.get("extra_pnginfo"), list): + kwargs["extra_pnginfo"] = kwargs["extra_pnginfo"][0] + if isinstance(kwargs.get("unique_id"), list): + kwargs["unique_id"] = kwargs["unique_id"][0] + + if not text_list: + return ("",) + + 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),) + + title = node.get("title", node.get("type", None)) + id = node.get("id", None) + + 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: + 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())] + + choice_idx = random.choices(range(len(text_list)), weights=weights)[0] + choice_text = "_".join(["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)}%)" + ) + + return (text_list[choice_idx],) + + +class MockStringList: + def __init__(self): + pass + + @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" + + def process(self, str1, str2): + return ([str1, str2],) + NODE_CLASS_MAPPINGS = { "StringToolsSeed": StringToolsSeed, "StringToolsString": StringToolsString, @@ -301,5 +558,10 @@ NODE_CLASS_MAPPINGS = { "StringToolsConcat": StringToolsConcat, "StringToolsRandomChoice": StringToolsRandomChoice, "StringToolsBalancedChoice": StringToolsBalancedChoice, + "StringToolsConcatList": StringToolsConcatList, + "StringToolsRandomChoiceList": StringToolsRandomChoiceList, + "StringToolsBalancedChoiceList": StringToolsBalancedChoiceList, + "string-tools/MockStringList": MockStringList, } WEB_DIRECTORY = "./js" + diff --git a/tests/conftest.py b/tests/conftest.py index 419ca26..0718013 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -125,6 +125,15 @@ def comfyui_server(test_settings): yield base_url print("\nShutting down ComfyUI server...") + # Flush remaining logs + time.sleep(1) # wait for any trailing errors + try: + while not output_queue.empty(): + line = output_queue.get_nowait() + print(f"ComfyUI: {line.strip()}") + except queue.Empty: + pass + process.terminate() try: process.wait(timeout=10) diff --git a/tests/test_string_tools.py b/tests/test_string_tools.py index 42d47e0..211e51e 100644 --- a/tests/test_string_tools.py +++ b/tests/test_string_tools.py @@ -120,23 +120,12 @@ def test_string_tools_execution_concat(page: Page, comfyui_server: str, wait_for strB.connect(0, concat, 1); - // 出力結果を保持しやすくするため、ComfyUI標準の「Preview as Text」ノード（内部クラス名: ShowText または PreviewTextNode）を利用する。 - let preview = LiteGraph.createNode("ShowText"); - - // フォールバック（環境によって登録名が異なる可能性があるため） - if (!preview) { - preview = LiteGraph.createNode("PreviewTextNode"); - } - + // 出力結果を保持するため、ComfyUI標準の「Preview as Text」ノード（PreviewAny）を利用 + let preview = LiteGraph.createNode("PreviewAny"); if (preview) { app.graph.add(preview); concat.connect(0, preview, 0); window.tgtNode = preview; - } else { - // 万が一見つからなかった場合はCLIPTextEncodeをフォールバックとして配置 - window.tgtNode = LiteGraph.createNode("CLIPTextEncode"); - app.graph.add(window.tgtNode); - concat.connect(0, window.tgtNode, 0); } app.graph.setDirtyCanvas(true, true); @@ -212,16 +201,12 @@ def test_string_tools_execution_random_choice(page: Page, comfyui_server: str, w } } - // 出力先 (Preview as Text 等) - let preview = LiteGraph.createNode("ShowText") || LiteGraph.createNode("PreviewTextNode"); + // 出力先 (Preview as Text) + let preview = LiteGraph.createNode("PreviewAny"); if (preview) { app.graph.add(preview); choiceNode.connect(0, preview, 0); window.tgtNode = preview; - } else { - window.tgtNode = LiteGraph.createNode("CLIPTextEncode"); - app.graph.add(window.tgtNode); - choiceNode.connect(0, window.tgtNode, 0); } app.graph.setDirtyCanvas(true, true); @@ -239,9 +224,9 @@ def test_string_tools_execution_random_choice(page: Page, comfyui_server: str, w history = get_latest_history(comfyui_server) if history and "outputs" in history: - choice_id = page.evaluate("window.choiceNodeId") - if str(choice_id) in history["outputs"]: - output = history["outputs"][str(choice_id)] + 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." @@ -295,15 +280,11 @@ def test_string_tools_execution_balanced_choice(page: Page, comfyui_server: str, } // 出力先 - let preview = LiteGraph.createNode("ShowText") || LiteGraph.createNode("PreviewTextNode"); + let preview = LiteGraph.createNode("PreviewAny"); if (preview) { app.graph.add(preview); choiceNode.connect(0, preview, 0); window.tgtNode = preview; - } else { - window.tgtNode = LiteGraph.createNode("CLIPTextEncode"); - app.graph.add(window.tgtNode); - choiceNode.connect(0, window.tgtNode, 0); } app.graph.setDirtyCanvas(true, true); @@ -321,11 +302,241 @@ def test_string_tools_execution_balanced_choice(page: Page, comfyui_server: str, history = get_latest_history(comfyui_server) if history and "outputs" in history: - choice_id = page.evaluate("window.choiceNodeId") - if str(choice_id) in history["outputs"]: - output = history["outputs"][str(choice_id)] - # 圧倒的に大きなWeight (100) を設定した "Banana" が選出されるかを検証 + tgt_id = page.evaluate("window.tgtNodeId") + 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." +def test_string_tools_execution_balanced_choice_list_subgraphs(page: Page, comfyui_server: str, wait_for_comfyui): + """ + StringToolsBalancedChoiceList ノードとStringToolsBalancedChoiceを使用し、 + ComfyUIのサブグラフ機能（Convert to Subgraph）を通して、 + 深い階層から重み付けされた選択が固定seedで期待通り出力されるか検証する。 + + グラフ構造: + strA("Apple") -------> MockStringList(str1) + strB1("Banana1") ---> BalancedChoice ----> MockStringList(str2) + strB2("Banana2") --/ | + v + BalancedChoiceList(text_list) ---> ShowText + + strB1, strB2, BalancedChoice を「Convert to Subgraph」でサブグラフ化する。 + """ + import requests + + page.goto(comfyui_server) + wait_for_comfyui(page) + + page.on("console", lambda msg: print(f"Browser Console: {msg.text}")) + page.on("pageerror", lambda err: print(f"Browser Error: {err}")) + + # Step 1: グラフを構築（サブグラフ化の前） + page.evaluate('''() => { + // MockStringList をJS側でも登録（Python側は __init__.py で登録済み） + class MockStringList extends LiteGraph.LGraphNode { + constructor() { + super(); + this.addInput("str1", "STRING"); + this.addInput("str2", "STRING"); + this.addOutput("list", "STRING"); + // ComfyUI固有: class_typeをgraphToPromptに正しく出力するために必要 + this.comfyClass = "string-tools/MockStringList"; + } + onExecute() {} + } + MockStringList.title = "MockStringList"; + MockStringList.comfyClass = "string-tools/MockStringList"; + LiteGraph.registerNodeType("string-tools/MockStringList", MockStringList); + + app.graph.clear(); + + const strA = LiteGraph.createNode("StringToolsString"); + strA.widgets[0].value = "Apple"; + app.graph.add(strA); + + const strB1 = LiteGraph.createNode("StringToolsString"); + strB1.widgets[0].value = "Banana1"; + app.graph.add(strB1); + + const strB2 = LiteGraph.createNode("StringToolsString"); + strB2.widgets[0].value = "Banana2"; + app.graph.add(strB2); + + 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); + + const mockList = LiteGraph.createNode("string-tools/MockStringList"); + app.graph.add(mockList); + + strA.connect(0, mockList, 0); + subChoice.connect(0, mockList, 1); + + 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); + + mockList.connect(0, mainChoiceList, 0); + + const preview = LiteGraph.createNode("PreviewAny"); + if (preview) { + app.graph.add(preview); + mainChoiceList.connect(0, preview, 0); + } + + // ノード参照をwindowに保持（後のconvertToSubgraphステップ用） + window._strB1 = strB1; + window._strB2 = strB2; + window._subChoice = subChoice; + window._seedNodeB = seedNodeB; + window.mainChoiceListId = mainChoiceList.id; + + app.graph.setDirtyCanvas(true, true); + + }''') + + # Step 2: Convert to Subgraph + convert_result = page.evaluate('''() => { + try { + if (typeof app.graph.convertToSubgraph !== 'function') { + return { error: "convertToSubgraph not available" }; + } + + const nodesToConvert = new Set([window._strB1, window._strB2, window._subChoice, window._seedNodeB]); + app.graph.convertToSubgraph(nodesToConvert); + + // サブグラフ変換後のグラフ状態を確認 + const nodeTypes = app.graph._nodes.map(n => n.type); + return { + success: true, + nodeCount: app.graph._nodes.length, + nodeTypes: nodeTypes + }; + } catch (e) { + return { error: e.toString(), stack: e.stack ? e.stack.substring(0, 500) : '' }; + } + }''') + print(f"Convert to Subgraph result: {convert_result}") + + # Step 3: graphToPrompt でプロンプトが生成できるか確認 + prompt_result = page.evaluate('''async () => { + try { + const p = await app.graphToPrompt(); + return { + success: true, + outputKeys: Object.keys(p.output || {}), + workflowNodeCount: (p.workflow && p.workflow.nodes) ? p.workflow.nodes.length : -1, + hasDefinitions: !!(p.workflow && p.workflow.definitions), + promptSnippet: JSON.stringify(p.output).substring(0, 500) + }; + } catch (e) { + return { error: e.toString(), stack: e.stack ? e.stack.substring(0, 500) : '' }; + } + }''') + print(f"graphToPrompt result: {prompt_result}") + + # Step 4: Python側から直接POSTして400エラーの詳細を取得 + prompt_data = page.evaluate('''async () => { + const p = await app.graphToPrompt(); + return { output: p.output, workflow: p.workflow }; + }''') + + import json + payload = { + "client_id": "test", + "prompt": prompt_data["output"], + "extra_data": { + "extra_pnginfo": { + "workflow": prompt_data["workflow"] + } + } + } + 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())}") + + resp_prompt = requests.post( + f"{comfyui_server}/prompt", + json=payload, + headers={"Content-Type": "application/json"}, + timeout=10 + ) + print(f"POST /prompt status: {resp_prompt.status_code}") + print(f"POST /prompt body: {resp_prompt.text[:500]}") + + if resp_prompt.status_code == 200: + page.wait_for_timeout(5000) + else: + # エラーでも継続して詳細を確認 + page.wait_for_timeout(2000) + + # Step 5: ヒストリーから結果を取得して検証 + resp = requests.get(f"{comfyui_server}/history", timeout=5) + history_data = resp.json() + print(f"History entry count: {len(history_data)}") + + assert history_data, "No history entries found after queuePrompt" + + latest_key = list(history_data.keys())[-1] + latest = history_data[latest_key] + status_info = latest.get("status", {}) + print(f"Latest history status: {status_info}") + + # history全体のキーを確認 + print(f"History top-level keys: {list(latest.keys())}") + + outputs = latest.get("outputs", {}) + print(f"Output node IDs: {list(outputs.keys())}") + + # outputs が空の場合はstatus/metaを詳細にダンプ + if not outputs: + import json + print(f"Full history entry: {json.dumps(latest, default=str)[:2000]}") + status_msgs = latest.get("status", {}).get("messages", []) + for msg in status_msgs: + print(f" status_msg: {msg}") + + + # PreviewAnyの出力を探す（OUTPUT_NODEの出力はそのIDで格納される） + found_output = None + for k, v in outputs.items(): + print(f" output[{k}]: {str(v)[:200]}") + # text キーに結果がある場合 (PreviewAnyのui出力形式) + if "text" in v: + found_output = v + break + + if found_output is None and outputs: + # text キーがなくても最初の出力を使う + found_output = list(outputs.values())[0] + + if found_output is not None: + out_str = str(found_output) + print(f"Output: {out_str}") + assert "Banana" in out_str or "Apple" in out_str, f"Output was {out_str}" + else: + pytest.fail(f"No outputs found in history. Keys: {list(outputs.keys())}")