From a2fc98293d7ebbd366af96a0e7630b9ecd285df5 Mon Sep 17 00:00:00 2001 From: duanyll Date: Thu, 13 Nov 2025 04:29:15 +0000 Subject: [PATCH] add remote function call --- comfyui_functional/comfy_client.py | 345 +++++++++++++++++++++++++++++ comfyui_functional/core.py | 124 ++++++----- comfyui_functional/nodes.py | 29 +-- comfyui_functional/remote.py | 105 +++++++++ comfyui_functional/transform.py | 11 +- comfyui_functional/utils.py | 252 ++++++++++++++++----- 6 files changed, 735 insertions(+), 131 deletions(-) create mode 100644 comfyui_functional/comfy_client.py create mode 100644 comfyui_functional/remote.py diff --git a/comfyui_functional/comfy_client.py b/comfyui_functional/comfy_client.py new file mode 100644 index 0000000..239c3a9 --- /dev/null +++ b/comfyui_functional/comfy_client.py @@ -0,0 +1,345 @@ +import asyncio +import json +import uuid +import aiohttp +from PIL import Image +import io +import hashlib +from typing import Dict, Any, List, Union + +class WorkflowExecutionError(Exception): + """自定义异常，用于表示 ComfyUI 工作流执行期间的错误。""" + def __init__(self, message, node_id, node_type, exception_message): + super().__init__(message) + self.node_id = node_id + self.node_type = node_type + self.exception_message = exception_message + + def __str__(self): + return f"{super().__str__()} - Node ID: {self.node_id}, Type: {self.node_type}, Details: {self.exception_message}" + + +class ComfyClient: + """ + 一个用于与 ComfyUI 服务器交互的客户端。 + """ + def __init__(self, server_base: str, timeout: int = 600): + """ + 初始化 ComfyClient。 + + :param server_base: ComfyUI 服务器的基础地址 (例如 "http://127.0.0.1:8188")。 + :param timeout: WebSocket 接收消息和执行工作流的超时时间（秒）。 + """ + self.server_base = server_base + self.timeout = timeout + self.base_url = server_base.rstrip('/') + ws_base = self.base_url.replace("http://", "ws://").replace("https://", "wss://") + self.ws_url = f"{ws_base}/ws" + + async def _upload_image(self, image: Image.Image, session: aiohttp.ClientSession) -> str: + """ + 将 PIL.Image 对象上传到 ComfyUI 服务器。 + + :param image: 要上传的 PIL.Image 对象。 + :param session: aiohttp.ClientSession 对象。 + :return: 服务器上保存的文件名。 + """ + image_bytes = io.BytesIO() + image.save(image_bytes, format='PNG') + image_bytes.seek(0) + + hash = hashlib.sha256(image_bytes.getvalue()).hexdigest() + filename = f"{hash}.png" + image_bytes.seek(0) + + form_data = aiohttp.FormData() + form_data.add_field('image', image_bytes, filename=filename, content_type='image/png') + form_data.add_field('type', 'input') + form_data.add_field('overwrite', 'true') + + async with session.post(f"{self.base_url}/upload/image", data=form_data) as response: + response.raise_for_status() + result = await response.json() + return result['name'] + + async def _get_image(self, filename: str, subfolder: str, folder_type: str, session: aiohttp.ClientSession) -> Image.Image: + """ + 从 ComfyUI 服务器下载图片。 + + :param filename: 文件名。 + :param subfolder: 子文件夹。 + :param folder_type: 文件夹类型 ('input', 'output', 'temp')。 + :param session: aiohttp.ClientSession 对象。 + :return: 下载的 PIL.Image 对象。 + """ + params = {"filename": filename, "subfolder": subfolder, "type": folder_type} + async with session.get(f"{self.base_url}/view", params=params) as response: + response.raise_for_status() + image_data = await response.read() + return Image.open(io.BytesIO(image_data)) + + def _find_nodes(self, workflow: dict) -> tuple[dict, dict]: + """ + 在工作流中查找用 % 标记的输入和输出节点。 + + :param workflow: 工作流字典。 + :return: (输入节点映射, 输出节点映射) 元组。 + """ + input_map = {} # key: placeholder_name, value: node_id + output_map = {} # key: node_id, value: placeholder_name + + for node_id, node_data in workflow.items(): + meta_title = node_data.get("_meta", {}).get("title") + if meta_title and meta_title.startswith('%') and meta_title.endswith('%'): + name = meta_title.strip('%') + # 假设所有带 % 标记的节点都可能是输出 + output_map[node_id] = name + # 如果节点类型是基本类型或 LoadImage，则认为是输入 + if node_data['class_type'] in [ + "PrimitiveBoolean", "PrimitiveFloat", "PrimitiveInt", + "PrimitiveString", "PrimitiveStringMultiline", "LoadImage" + ]: + input_map[name] = node_id + + return input_map, output_map + + async def _prepare_workflow( + self, + workflow: dict, + session: aiohttp.ClientSession, + **kwargs + ) -> tuple[dict, dict]: + """ + 准备工作流：验证并注入参数，处理图片上传。 + + :param workflow: 原始工作流模板。 + :param session: aiohttp.ClientSession 对象。 + :param kwargs: 用户提供的参数。 + :return: (准备好的工作流, 输出节点映射) 元组。 + """ + prepared_workflow = json.loads(json.dumps(workflow)) + input_map, output_map = self._find_nodes(prepared_workflow) + + for name, value in kwargs.items(): + if name not in input_map: + raise ValueError(f"输入参数 '{name}' 在工作流中没有找到对应的 %{name}% 节点。") + + node_id = input_map[name] + node = prepared_workflow[node_id] + + if node['class_type'] == 'LoadImage': + if isinstance(value, str): # 路径 + with open(value, 'rb') as f: + pil_image = Image.open(io.BytesIO(f.read())) + filename = await self._upload_image(pil_image, session) + elif isinstance(value, Image.Image): + filename = await self._upload_image(value, session) + else: + raise TypeError(f"LoadImage 输入 '{name}' 的类型不受支持: {type(value).__name__}。需要 str (路径) 或 PIL.Image 对象。") + node['inputs']['image'] = filename + else: + if 'value' not in node['inputs']: + raise ValueError(f"节点 '{node_id}' ('{name}') 被标记为输入，但在其 inputs 中没有 'value' 字段。") + + default_value = node['inputs']['value'] + if not isinstance(value, type(default_value)) and default_value is not None: + raise TypeError(f"输入 '{name}' 的类型不匹配。期望类型: {type(default_value).__name__}, 得到类型: {type(value).__name__}。") + + node['inputs']['value'] = value + + return prepared_workflow, output_map + + async def run_workflow(self, workflow: dict, **kwargs) -> Dict[str, List[Any]]: + """ + 异步执行一个 ComfyUI 工作流。 + + :param workflow: 工作流模板 (Python 字典)。 + :param kwargs: 要注入工作流的参数。键对应模板中被 % 包围的节点标题。 + :return: 一个字典，键是输出节点的标题，值是输出内容的列表。 + """ + client_id = str(uuid.uuid4()) + + async with aiohttp.ClientSession() as session: + prepared_workflow, output_map = await self._prepare_workflow(workflow, session, **kwargs) + + async with session.ws_connect(f"{self.ws_url}?clientId={client_id}", timeout=self.timeout, max_msg_size=128*1024*1024) as ws: + prompt_payload = {"prompt": prepared_workflow, "client_id": client_id} + async with session.post(f"{self.base_url}/prompt", json=prompt_payload) as response: + if response.status != 200: + error_text = await response.text() + raise ConnectionError(f"提交工作流失败: {response.status} {error_text}") + queue_data = await response.json() + prompt_id = queue_data['prompt_id'] + + outputs: Dict[str, list] = {name: [] for name in output_map.values()} + image_download_tasks = [] + current_executing_node = None + is_done = False + + while not is_done: + try: + msg_data = await ws.receive(timeout=self.timeout) + except asyncio.TimeoutError: + raise TimeoutError(f"工作流执行超时（{self.timeout}秒）。") + + if msg_data.type == aiohttp.WSMsgType.TEXT: + message = msg_data.json() + msg_type = message.get('type') + data = message.get('data', {}) + + if data.get('prompt_id') != prompt_id: + continue + + print(f"[*] 收到消息: {msg_type} - {data.get('display_node', '')}") + + if msg_type == 'execution_error': + raise WorkflowExecutionError( + "工作流执行出错", + node_id=data.get('node_id'), + node_type=data.get('node_type'), + exception_message=data.get('exception_message') + ) + elif msg_type == 'execution_success': + is_done = True + elif msg_type == 'executing': + current_executing_node = data.get('display_node') or data.get('node') + elif msg_type == 'executed': + node_id = data.get('display_node') or data.get('node') + if node_id in output_map: + output_name = output_map[node_id] + output_data = data.get('output', {}) + if 'text' in output_data: + outputs[output_name].extend(output_data['text']) + if 'images' in output_data: + for img_info in output_data['images']: + task = self._get_image( + img_info['filename'], + img_info.get('subfolder', ''), + img_info['type'], + session + ) + image_download_tasks.append((output_name, task)) + + elif msg_data.type == aiohttp.WSMsgType.BINARY: + if current_executing_node and current_executing_node in output_map: + output_name = output_map[current_executing_node] + image_data = msg_data.data[8:] # 去掉8字节头部 + image = Image.open(io.BytesIO(image_data)) + outputs[output_name].append(image) + + elif msg_data.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR): + is_done = True + + if image_download_tasks: + results = await asyncio.gather(*(task for _, task in image_download_tasks)) + for i, (output_name, _) in enumerate(image_download_tasks): + outputs[output_name].append(results[i]) + + return {k: v for k, v in outputs.items() if v} + + def run_workflow_sync(self, workflow: dict, **kwargs) -> Dict[str, List[Any]]: + """ + 同步执行一个 ComfyUI 工作流。 + 这是 run_workflow 的同步版本。 + + :param workflow: 工作流模板 (Python 字典)。 + :param kwargs: 要注入工作流的参数。 + :return: 一个字典，键是输出节点的标题，值是输出内容的列表。 + """ + return asyncio.run(self.run_workflow(workflow, **kwargs)) + + +def _str_to_primitive(s: str) -> Union[bool, int, float, str]: + """尝试将字符串转换为布尔值、整数或浮点数。""" + s_lower = s.lower() + if s_lower == 'true': + return True + if s_lower == 'false': + return False + try: + return int(s) + except ValueError: + try: + return float(s) + except ValueError: + return s + +def main_cli(): + """命令行接口的主函数。""" + + import argparse + import sys + from pathlib import Path + + parser = argparse.ArgumentParser( + description="ComfyUI Python 客户端命令行接口。", + epilog="工作流的输入参数请使用 '--key=value' 的格式，例如: --prompt=\"a cat\" --steps=25" + ) + parser.add_argument("workflow_file", help="要执行的工作流 JSON 文件路径。") + parser.add_argument( + "--output_dir", + default="workflow_outputs", + help="用于保存输出图片的目录。(默认: workflow_outputs)" + ) + parser.add_argument( + "--server", + default="http://127.0.0.1:8188", + help="ComfyUI 服务器地址。(默认: http://127.0.0.1:8188)" + ) + + args, unknown_args = parser.parse_known_args() + + # --- 解析动态工作流参数 --- + workflow_kwargs = {} + for arg in unknown_args: + if not arg.startswith('--'): + continue + try: + key, value = arg.lstrip('-').split("=", 1) + workflow_kwargs[key] = _str_to_primitive(value) + except ValueError: + parser.error(f"无法解析参数 '{arg}'。请使用 '--key=value' 格式。") + + # --- 日志函数，输出到 stderr --- + log = lambda *a, **kw: print(*a, **kw, file=sys.stderr) + + try: + # --- 准备和执行工作流 --- + log(f"[*] 正在从 {args.workflow_file} 加载工作流...") + with open(args.workflow_file, 'r', encoding='utf-8') as f: + workflow = json.load(f) + + client = ComfyClient(args.server) + + log(f"[*] 正在连接到 ComfyUI 服务器: {args.server}") + log(f"[*] 使用以下输入参数执行工作流: {workflow_kwargs}") + + results = client.run_workflow_sync(workflow, **workflow_kwargs) + + log("[*] 工作流执行完毕。") + + # --- 处理输出 --- + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + log(f"[*] 正在处理输出结果...") + + for output_name, output_items in results.items(): + for i, item in enumerate(output_items, 1): + if isinstance(item, Image.Image): + filename = f"{output_name}_{i}.png" + filepath = output_dir / filename + item.save(filepath) + log(f" - 图片已保存到: {filepath}") + else: + # 将文本和其他原始类型输出到 stdout + print(item) + + except FileNotFoundError: + log(f"[错误] 工作流文件未找到: '{args.workflow_file}'") + sys.exit(1) + except Exception as e: + log(f"[错误] 执行过程中发生错误: {e}") + sys.exit(1) + +if __name__ == "__main__": + main_cli() \ No newline at end of file diff --git a/comfyui_functional/core.py b/comfyui_functional/core.py index 3a3924f..03f1747 100644 --- a/comfyui_functional/core.py +++ b/comfyui_functional/core.py @@ -7,6 +7,7 @@ from .utils import AnyType, ContainsDynamicDict, Closure, create_graph_from_clos RECURISON_LIMIT = 50 COROUTINE_LIMIT = 100 + class FunctionParam: @classmethod def INPUT_TYPES(cls): @@ -19,16 +20,19 @@ class FunctionParam: RETURN_TYPES = (AnyType("*"),) FUNCTION = "run" CATEGORY = "duanyll/functional" - + def run(self, index): - return (ExecutionBlocker( - "All outputs of nodes that directly or indirectly depend on the Function " - "Parameter node must ultimately connect to a Function End node or an Inspect " - "node. An output node is incorrectly connected, as it depends on the Function " - "Parameter node without passing through a Function End or Inspect node. To " - "view temporary results within the function, please use an Inspect node before " - "the output node. Check your node connections and try again." - ),) + return ( + ExecutionBlocker( + "All outputs of nodes that directly or indirectly depend on the Function " + "Parameter node must ultimately connect to a Function End node or an Inspect " + "node. An output node is incorrectly connected, as it depends on the Function " + "Parameter node without passing through a Function End or Inspect node. To " + "view temporary results within the function, please use an Inspect node before " + "the output node. Check your node connections and try again." + ), + ) + class FunctionEnd: @classmethod @@ -36,27 +40,29 @@ class FunctionEnd: return { "required": { "return_value": (AnyType("*"),), + "capture": ("BOOLEAN", {"default": True}), }, } - - RETURN_TYPES = ("CLOSURE", ) + + RETURN_TYPES = ("CLOSURE",) FUNCTION = "run" CATEGORY = "duanyll/functional" - - def run(self, return_value): + + def run(self, return_value, capture): raise NotImplementedError( "The Function End node is a marker and should never be executed directly. " "If you see this error, please file a bug report." ) + class CreateClosure: @classmethod def INPUT_TYPES(cls): return { "required": { - "body": ("STRING", {"multiline": True} ), - "output": ("STRING", ), - "side_effects": ("FLOAT", ) + "body": ("STRING", {"multiline": True}), + "output": ("STRING",), + "side_effects": ("FLOAT",), }, "optional": ContainsDynamicDict( { @@ -65,7 +71,7 @@ class CreateClosure: ), } - RETURN_TYPES = ("CLOSURE", ) + RETURN_TYPES = ("CLOSURE",) FUNCTION = "run" CATEGORY = "duanyll/functional/internal" @@ -76,7 +82,14 @@ class CreateClosure: captures.append(kwargs[f"capture_{i}"]) body = json.loads(body) output = json.loads(output) - return (Closure(body=body, output=output, captures=captures), ) + return ( + Closure( + body=body, + output=output, + captures=captures, + ), + ) + class CallClosure: @classmethod @@ -90,12 +103,10 @@ class CallClosure: "param_0": (AnyType("*"), {"_dynamic": "number"}), } ), - "hidden": { - "unique_id": "UNIQUE_ID" - }, + "hidden": {"unique_id": "UNIQUE_ID"}, } - RETURN_TYPES = (AnyType("*"), ) + RETURN_TYPES = (AnyType("*"),) FUNCTION = "run" CATEGORY = "duanyll/functional" @@ -110,15 +121,14 @@ class CallClosure: params = [] for i in range(len(kwargs)): params.append(kwargs[f"param_{i}"]) - graph, output = create_graph_from_closure(closure, params, caller_unique_id=unique_id) + graph, output = create_graph_from_closure( + closure, params, caller_unique_id=unique_id + ) if len(graph) == 0: - return (output, ) - return { - "result": (output, ), - "expand": graph - } - - + return (output,) + return {"result": (output,), "expand": graph} + + class IntermidiateCoroutine: @classmethod def INPUT_TYPES(cls): @@ -127,28 +137,26 @@ class IntermidiateCoroutine: "return_value": (AnyType("*"),), "coroutine": (AnyType("*"),), }, - "hidden": { - "unique_id": "UNIQUE_ID" - }, + "hidden": {"unique_id": "UNIQUE_ID"}, } - - RETURN_TYPES = (AnyType("*"), ) + + RETURN_TYPES = (AnyType("*"),) FUNCTION = "run" CATEGORY = "duanyll/functional/internal" - + def run(self, return_value, coroutine, unique_id): (base_id, index) = unique_id.rsplit("_", 1) graph = {} - + while len(graph) == 0: try: (closure, params) = coroutine.send(return_value) except StopIteration as e: - return (e.value, ) - + return (e.value,) + index = int(index) + 1 next_node_id = f"{base_id}_{index}" - + if index > COROUTINE_LIMIT: raise RecursionError( "Coroutine execution limit exceeded. Possible infinite loop " @@ -157,13 +165,15 @@ class IntermidiateCoroutine: "in core.py." ) - graph, output = create_graph_from_closure(closure, params, caller_unique_id=unique_id) + graph, output = create_graph_from_closure( + closure, params, caller_unique_id=unique_id + ) if len(graph) == 0: return_value = output - + return { - "result": ([next_node_id, 0], ), + "result": ([next_node_id, 0],), "expand": { **graph, next_node_id: { @@ -172,33 +182,33 @@ class IntermidiateCoroutine: "coroutine": coroutine, }, "class_type": "__IntermidiateCoroutine__", - } - } + }, + }, } - - + + class CoroutineNodeBase: @classmethod def INPUT_TYPES(cls): raise NotImplementedError( "This is an abstract base class and should not be instantiated directly." ) - - RETURN_TYPES = (AnyType("*"), ) + + RETURN_TYPES = (AnyType("*"),) FUNCTION = "run" - + def run_coroutine(self, **kwargs): raise NotImplementedError( "Subclasses must implement the run_coroutine method as a generator." ) - yield # This is just to make this function a generator - + yield # This is just to make this function a generator + def run(self, unique_id, **kwargs): if isinstance(unique_id, list): unique_id = unique_id[0] coroutine = self.run_coroutine(**kwargs) return { - "result": ([f"{unique_id}_0", 0], ), + "result": ([f"{unique_id}_0", 0],), "expand": { f"{unique_id}_0": { "inputs": { @@ -207,20 +217,20 @@ class CoroutineNodeBase: }, "class_type": "__IntermidiateCoroutine__", } - } + }, } - + class RecoverList: @classmethod def INPUT_TYPES(cls): return { "required": { - "values": (AnyType("*"), ), + "values": (AnyType("*"),), }, } - RETURN_TYPES = (AnyType("*"), ) + RETURN_TYPES = (AnyType("*"),) FUNCTION = "run" CATEGORY = "duanyll/functional/internal" diff --git a/comfyui_functional/nodes.py b/comfyui_functional/nodes.py index e8c42ed..e08f303 100644 --- a/comfyui_functional/nodes.py +++ b/comfyui_functional/nodes.py @@ -1,26 +1,13 @@ -from .core import (NODE_CLASS_MAPPINGS as CORE_NODE_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS as CORE_NODE_DISPLAY_NAME_MAPPINGS) -from .high_level import (NODE_CLASS_MAPPINGS as HIGH_LEVEL_NODE_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS as HIGH_LEVEL_NODE_DISPLAY_NAME_MAPPINGS) -from .side_effects import (NODE_CLASS_MAPPINGS as SIDE_EFFECTS_NODE_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS as SIDE_EFFECTS_NODE_DISPLAY_NAME_MAPPINGS) -from .logic import (NODE_CLASS_MAPPINGS as LOGIC_NODE_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS as LOGIC_NODE_DISPLAY_NAME_MAPPINGS) +from . import core, high_level, logic, side_effects, remote -NODE_CLASS_MAPPINGS = { - **CORE_NODE_CLASS_MAPPINGS, - **HIGH_LEVEL_NODE_CLASS_MAPPINGS, - **SIDE_EFFECTS_NODE_CLASS_MAPPINGS, - **LOGIC_NODE_CLASS_MAPPINGS, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - **CORE_NODE_DISPLAY_NAME_MAPPINGS, - **HIGH_LEVEL_NODE_DISPLAY_NAME_MAPPINGS, - **SIDE_EFFECTS_NODE_DISPLAY_NAME_MAPPINGS, - **LOGIC_NODE_DISPLAY_NAME_MAPPINGS, -} +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} +for module in [core, high_level, logic, side_effects, remote]: + NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS) + NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS) + +# Mangle ids to avoid conflicts PREFIX = "duanyll::" NODE_CLASS_MAPPINGS = { (k if k.startswith("__") else PREFIX + k): v diff --git a/comfyui_functional/remote.py b/comfyui_functional/remote.py new file mode 100644 index 0000000..d5132bf --- /dev/null +++ b/comfyui_functional/remote.py @@ -0,0 +1,105 @@ +from .utils import ( + serialize, + deserialize, + AnyType, + ContainsDynamicDict, + Closure, + create_remote_workflow_from_closure, +) +from .comfy_client import ComfyClient + + +class Serialize: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "data": (AnyType("*"),), + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "run" + CATEGORY = "duanyll/functional/internal" + OUTPUT_NODE = True + + @classmethod + def IS_CHANGED(cls, data): + return float("NaN") + + def run(self, data): + try: + text = serialize(data) + print(f"Serialized {type(data)} to string of length {len(text)}") + return { + "ui": {"text": [text]}, + "result": (text,), + } + except Exception as e: + raise ValueError( + f"Cannot serialize function output. It may contain unserializable types. Error: {e}" + ) + + +class Deserialize: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "data": ("STRING", {"multiline": True}), + } + } + + RETURN_TYPES = (AnyType("*"),) + FUNCTION = "run" + CATEGORY = "duanyll/functional/internal" + + def run(self, data): + return (deserialize(data),) + + +class CallClosureRemote: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "closure": ("CLOSURE",), + "base_url": ("STRING",), + "timeout": ("FLOAT", {"default": 600.0, "min": 1.0}), + }, + "optional": ContainsDynamicDict( + { + "param_0": (AnyType("*"), {"_dynamic": "number"}), + } + ) + } + + RETURN_TYPES = (AnyType("*"),) + FUNCTION = "run" + CATEGORY = "duanyll/functional" + + async def run(self, closure, base_url, timeout, **kwargs): + params = [] + for i in range(len(kwargs)): + params.append(kwargs[f"param_{i}"]) + + client = ComfyClient(server_base=base_url, timeout=timeout) + workflow = create_remote_workflow_from_closure(closure, params) + client_response = await client.run_workflow(workflow) + output = client_response.get("__output__", None) + if output is None: + raise ValueError("No output received from remote workflow execution.") + return (deserialize(output[0]),) + + +NODE_CLASS_MAPPINGS = { + "__Serialize__": Serialize, + "__Deserialize__": Deserialize, + "CallClosureRemote": CallClosureRemote +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "__Serialize__": "[Internal] Serialize", + "__Deserialize__": "[Internal] Deserialize", + "CallClosureRemote": "Call Remote Function", +} \ No newline at end of file diff --git a/comfyui_functional/transform.py b/comfyui_functional/transform.py index 56418a9..b767e4d 100644 --- a/comfyui_functional/transform.py +++ b/comfyui_functional/transform.py @@ -108,10 +108,15 @@ def _transform_single_function(end_node_id: str, workflow: Dict, graph: Workflow """转换单个 __FunctionEnd__ 节点及其关联的子图。""" subgraph_nodes = _find_function_subgraph(end_node_id, graph) original_end_node = workflow[end_node_id] + enable_capture = original_end_node["inputs"].get("capture", True) - param_nodes = _find_param_nodes(subgraph_nodes, workflow, graph) - param_nodes.add(end_node_id) - capture_nodes = subgraph_nodes - param_nodes + if enable_capture: + param_nodes = _find_param_nodes(subgraph_nodes, workflow, graph) + param_nodes.add(end_node_id) + capture_nodes = subgraph_nodes - param_nodes + else: + param_nodes = subgraph_nodes + capture_nodes = set() # 恢复原始节点类型以便深拷贝 workflow[end_node_id]["class_type"] = original_end_node["class_type"] diff --git a/comfyui_functional/utils.py b/comfyui_functional/utils.py index ba5e796..bcadf62 100644 --- a/comfyui_functional/utils.py +++ b/comfyui_functional/utils.py @@ -1,5 +1,11 @@ from typing import TypedDict import copy +import json +import io +import base64 +import numpy as np +import torch +from PIL import Image class AnyType(str): def __ne__(self, __value: object) -> bool: @@ -43,63 +49,209 @@ class ContainsDynamicDict(dict): class Closure(TypedDict): body: dict - captures: dict + captures: list output: list -def create_graph_from_closure(self, params, caller_unique_id=None): - body = copy.deepcopy(self.body) +def create_graph_from_closure(closure: Closure, params, caller_unique_id=None): + body = copy.deepcopy(closure["body"]) graph = {} - - def add_recover_node(spec): - prefix = f"{caller_unique_id}_" if caller_unique_id is not None else "" - recover_id = f"{prefix}{spec[0]}_{spec[1]}" - if recover_id in graph: + + def transform_input(spec): + if not isinstance(spec, list): + return spec + # Return normal spec if not a capture or param + if spec[0] != "__capture" and spec[0] != "__param": + return spec + if spec[0] == "__param": + if spec[1] >= len(params): + raise IndexError( + f"Parameter index {spec[1]} out of range for function call. " + f"Only {len(params)} parameters were provided." + ) + value = params[spec[1]] + else: + value = closure["captures"][spec[1]] + if isinstance(value, list): + prefix = f"{caller_unique_id}_" if caller_unique_id is not None else "" + recover_id = f"{prefix}{spec[0]}_{spec[1]}" + if recover_id not in graph: + graph[recover_id] = { + "inputs": { + # Wrap the list in a tuple so it's treated as a literal, not a link + "values": (value,), + }, + "class_type": "__RecoverList__", + } return [recover_id, 0] - value = self.captures[spec[1]] if spec[0] == "__capture" else params[spec[1]] - graph[recover_id] = { - "inputs": { - # Wrap the list in a tuple so it's treated as a literal, not a link - "values": (value,), - }, - "class_type": "__RecoverList__", - } - return [recover_id, 0] + else: + return value for node_id, node_data in body.items(): inputs = node_data["inputs"] for key in inputs.keys(): - spec = inputs[key] - if isinstance(spec, list): - if spec[0] == "__capture": - value = self.captures[spec[1]] - if isinstance(value, list): - inputs[key] = add_recover_node(spec) - else: - inputs[key] = value - elif spec[0] == "__param": - if spec[1] >= len(params): - raise ValueError(f"Parameter index {spec[1]} out of range for provided {len(params)} params.") - value = params[spec[1]] - if isinstance(value, list): - inputs[key] = add_recover_node(spec) - else: - inputs[key] = value - else: - spec[0] = f"{caller_unique_id}_{spec[0]}" + inputs[key] = transform_input(inputs[key]) node_data["override_display_id"] = node_id graph[f"{caller_unique_id}_{node_id}"] = node_data - spec = self.output - if spec[0] == "__capture": - value = self.captures[spec[1]] - if isinstance(value, list): - return graph, add_recover_node(spec) - return graph, value - elif spec[0] == "__param": - if spec[1] >= len(params): - raise ValueError(f"Parameter index {spec[1]} out of range for provided {len(params)} params.") - value = params[spec[1]] - if isinstance(value, list): - return graph, add_recover_node(spec) - return graph, value - else: - return graph, [f"{caller_unique_id}_{spec[0]}", spec[1]] + + output = transform_input(closure["output"]) + return graph, output + + +def create_remote_workflow_from_closure(closure, params): + body = copy.deepcopy(closure["body"]) + graph = {} + + def transform_input(spec): + if not isinstance(spec, list): + return spec + # Return normal spec if not a capture or param + if spec[0] != "__capture" and spec[0] != "__param": + return spec + + # Check if we have already created a deserialize node for this spec + deserialize_node_id = f"deserialize_{spec[0]}_{spec[1]}" + if deserialize_node_id in graph: + return [deserialize_node_id, 0] + + # Try to create a deserialize node + if spec[0] == "__param": + if spec[1] >= len(params): + raise IndexError( + f"Parameter index {spec[1]} out of range for function call. " + f"Only {len(params)} parameters were provided." + ) + value = params[spec[1]] + else: + value = closure["captures"][spec[1]] + + try: + serialized_value = serialize(value) + graph[deserialize_node_id] = { + "inputs": { + "data": serialized_value, + }, + "class_type": "__Deserialize__", + } + return [deserialize_node_id, 0] + except TypeError as e: + if spec[0] == "__capture": + raise ValueError( + f"Cannot serialize captured value at index {spec[1]}. " + "Captured values must be serializable to be used in remote execution. " + "Consider disabling capture for this function to initialize the value remotely. " + f"Error: {e}" + ) + else: + raise ValueError( + f"Cannot serialize parameter at index {spec[1]}. " + "Parameters must be serializable to be used in remote execution. " + f"Error: {e}" + ) + + for node_id, node_data in body.items(): + inputs = node_data["inputs"] + for key in inputs.keys(): + inputs[key] = transform_input(inputs[key]) + graph[node_id] = node_data + + output = transform_input(closure["output"]) + graph["__output__"] = { + "inputs": { + "data": output, + }, + "class_type": "__Serialize__", + "_meta": { + "title": "%__output__%" + } + } + + return graph + +class WhitelistEncoder(json.JSONEncoder): + def default(self, obj): + # 区分 set 和 tuple + if isinstance(obj, set): + return {'__type__': 'set', 'data': list(obj)} + if isinstance(obj, tuple): + return {'__type__': 'tuple', 'data': list(obj)} + + if isinstance(obj, bytes): + return {'__type__': 'bytes', 'data': base64.b64encode(obj).decode('utf-8')} + + if isinstance(obj, np.ndarray): + buffer = io.BytesIO() + np.save(buffer, obj, allow_pickle=False) + return { + '__type__': 'numpy.ndarray', + 'data': base64.b64encode(buffer.getvalue()).decode('utf-8') + } + + if isinstance(obj, torch.Tensor): + buffer = io.BytesIO() + torch.save(obj.cpu(), buffer) + return { + '__type__': 'torch.Tensor', + 'data': base64.b64encode(buffer.getvalue()).decode('utf-8') + } + + if isinstance(obj, Image.Image): + buffer = io.BytesIO() + obj.save(buffer, format='PNG') + return { + '__type__': 'PIL.Image', + 'data': base64.b64encode(buffer.getvalue()).decode('utf-8') + } + + raise TypeError( + f"Object of type '{type(obj).__name__}' is not serializable. " + "Only basic types, set, tuple, bytes, numpy.ndarray, " + "torch.Tensor, and PIL.Image are allowed." + ) + +# 2. 自定义解码器 (Deserialization) +def whitelist_decoder(dct): + """ + 用于 json.load/loads 的 object_hook，用于恢复自定义类型。 + """ + if '__type__' in dct: + obj_type = dct['__type__'] + data = dct['data'] + + if obj_type == 'set': + return set(data) + if obj_type == 'tuple': + return tuple(data) + if obj_type == 'bytes': + return base64.b64decode(data) + + if obj_type == 'numpy.ndarray': + buffer = io.BytesIO(base64.b64decode(data)) + # allow_pickle=True 是为了恢复可能包含的复杂dtype + return np.load(buffer, allow_pickle=True) + + if obj_type == 'torch.Tensor': + buffer = io.BytesIO(base64.b64decode(data)) + return torch.load(buffer) + + if obj_type == 'PIL.Image': + buffer = io.BytesIO(base64.b64decode(data)) + img = Image.open(buffer) + # .copy() 是个好习惯，确保数据已从缓冲区完全加载 + return img.copy() + + return dct + + +def serialize(obj): + """ + 将对象序列化为 JSON 字符串。 + 仅支持基本类型、set、tuple、bytes、numpy.ndarray、torch.Tensor 和 PIL.Image。 + """ + return json.dumps(obj, cls=WhitelistEncoder) + + +def deserialize(json_str): + """ + 从 JSON 字符串反序列化对象。 + """ + return json.loads(json_str, object_hook=whitelist_decoder) \ No newline at end of file