# ComfyUI remote node execution nodes. # Copyright (C) 2024 LatentRat # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with this program. If not, see . # TODO: option to actually run IS_CHANGED checks remotely # TODO: torch weights_only option # TODO: clear remote history item afterwards option # TODO: IS_CHANGED return _ignore_ when disconnected option or always? # TODO: forward interruptions option import asyncio import base64 import collections import concurrent.futures import copy import enum import functools import gzip import hashlib import math import queue import re import threading import time import os import weakref from io import BytesIO from uuid import uuid4 from zipfile import ZipFile import execution import nodes from comfy_execution.graph import DynamicPrompt import requests import torch import websockets.sync.client from .libs.safe_load import torch_safe_load_dict from .libs.fancy_safetensors import torch as fst_torch import json import logging as _logging from comfy.comfy_types import IO as CoIO logger = _logging.getLogger(__name__) if (IS_DEV := os.environ.get("DEV") == "1"): IS_DEV_PROMPTSAVE = True IS_DEV_PROMPTSAVE = False def DEBUG_PROMPT_SAVE(label, prompt): from pathlib import Path dir = Path(f"./dev/jsons/") if not dir.exists(): dir.mkdir(parents = False, exist_ok = True) from comfy.cli_args import args import datetime now = datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%d_%H-%M-%S.%f") Path(f"./dev/jsons/port={args.port}____{label}__{now}__{time.time_ns()}.json").write_text(json.dumps(prompt)) jdumps = lambda data: json.dumps(data).replace("\n", "") _NUM_OUTPUTS = int(os.environ.get("RAT_REMOTE_RUN_NUM_OUTPUTS") or 10) logger.info("%s", f"RemoteRun nodes using _NUM_OUTPUTS={_NUM_OUTPUTS}") TOGGLE_CHOICES = ("only_locally", "only_remotely") INPUTS_OFF_OPTIONS = ["lazy", "disconnected"] BLOCK_OPTIONS = ["error", "block_execution_silent", "block_execution_verbose", "passthrough", "return_none"] SERIALIZATION_OPTIONS = ["safe_torch_pt", "unsafe_torch_pt", "fancy_safetensors"] class BinaryResponseMessageID(enum.IntEnum): BINARY_RESPONSE_MESSAGE_ID = 12341119 GZIPPED_BINARY_RESPONSE_MESSAGE_ID = 12341120 def serialization_load_input(name = "serialization"): return { name: (SERIALIZATION_OPTIONS, { "default": "fancy_safetensors" }) } def response_input(name = "response"): return { name: (["base64_result", "binary"], { "default": "binary" }) } def ws_settings_inputs(): return { "ws_compression": ("BOOLEAN", { "default": False }), # strings to make it easier to add more in case of wanting a partial option without breaking backwards compatibility "forward_progress_messages": (["on", "off"], { "default": "on" }), } def shared_input_output_settings(): return { "max_size_mb": ("INT", { "default": 32, "min": 0, "max": 8 * 1024, "step": 1 }), "gzip_compression": ("BOOLEAN", { "default": False }), "gzip_level": ("INT", { "default": 9, "min": 1, "max": 9, "step": 1 }), } def lazy_input_settings(): return { "lazily_transfer": ("BOOLEAN", { "default": False }), "skip_lazy_transfer_under_mb": ("FLOAT", { "default": 0.25, "min": 0.0, "max": 10 * 1024.0, "step": 0.1, "tooltip": "Even when lazily transferring inputs is enabled still always transfer inputs smaller than this size in MB directly.", }), } class RemoteRunSerializerOutNode(): TYPE_NAME = "RAT_RemoteRunSerializerOut" DISPLAY_NAME = "RemRun Serializer Output (Internal)" CATEGORY = "Remote Run/__Internal__/" FUNCTION = "run" OUTPUT_NODE = True @classmethod def IS_CHANGED(cls, is_changed: bool = False, **_kwargs): return time.time_ns() if is_changed else None @classmethod def INPUT_TYPES(cls): inputs = { "required": { **serialization_load_input(), }, "optional": { **dict((f"input_{num}", (CoIO.ANY, { })) for num in range(_NUM_OUTPUTS)), **shared_input_output_settings(), **response_input(), "is_changed": ("BOOLEAN", { "default": False }), }, } return inputs RETURN_TYPES = ("STRING",) def run(self, serialization: str, max_size_mb: int | None = None, gzip_compression: bool = False, gzip_level: int = 9, response = None, **kwargs, ): inputs = { k: v for k, v in kwargs.items() if k.startswith("input_") } data = serialize_obj(serialization, inputs) if response == "binary": import server instance = server.PromptServer.instance client_id = instance.client_id if not client_id: raise ValueError("No client connected to send binary response to") data_type = BinaryResponseMessageID.BINARY_RESPONSE_MESSAGE_ID.value if gzip_compression: start = time.monotonic_ns() comp = gzip.compress(data, compresslevel = gzip_level) end = time.monotonic_ns() took_ns = (end - start) mb_s = (len(data) * 1e9) / (took_ns * 1024 * 1024) took_sec = took_ns / 1e9 factor = len(comp) / len(data) if factor < 0.99: logger.info("%s", f"RemoteRunSerializerOutNode.run serialization={serialization!r} {response=!r} " f"gzip compressed {len(data)} to {len(comp)} bytes, factor {factor:.3f}, " f"took {took_sec:.3f} sec, {mb_s:.2f} MB/s") data = comp data_type = BinaryResponseMessageID.GZIPPED_BINARY_RESPONSE_MESSAGE_ID.value else: logger.info("%s", f"RemoteRunSerializerOutNode.run serialization={serialization!r} {response=!r} " f"gzip compression not effective, kept original {len(data)} bytes " f"over compressed {len(comp)} bytes, factor {factor:.3f}, " f"took {took_sec:.3f} sec, {mb_s:.2f} MB/s") logger.info("%s", f"RemoteRunSerializerOutNode.run serialization={serialization!r} {response=!r} " f"data_size={len(data)}, {max_size_mb=} {client_id=!r}") if max_size_mb and len(data) > max_size_mb * 1024 * 1024: raise ValueError(f"Serialized data too large: {len(data)} bytes, max_size_mb={max_size_mb}") instance.send_sync(data_type, data, client_id) return ("",) results = base64.b64encode(data).decode() logger.info("%s", f"RemoteRunSerializerOutNode.run serialization={serialization!r} {response=!r} data_size={len(data)}, " f"b64encode_size={len(results)}, {max_size_mb=}") if max_size_mb and len(results) > max_size_mb * 1024 * 1024: raise ValueError(f"Serialized b64 data too large: {len(results)} bytes, max_size_mb={max_size_mb}") results = (results,) return { "results": results, "ui": { "results": results } } def setup_remote_run_api_route(expected_set_up: bool): import server from aiohttp import web ins = server.PromptServer.instance url = "/remote_run/data/" cur_routes = [ins.routes[i] for i in range(len(ins.routes))] routes = [i for i in cur_routes if i.path == url and (i.method or "").upper() == "POST"] if routes: if len(routes) > 1: raise ValueError("Multiple routes for remote run data endpoint???", routes) route_def = routes[0] handler = route_def.handler data_requests = handler._data_requests else: if expected_set_up: raise ValueError("No existing route for remote run data endpoint, disabled or server startup changed") data_requests = weakref.WeakValueDictionary() @ins.routes.post(url) async def receive_data(request): key = request.query.get("key") status = request.query.get("status") logger.info("%s", f"RemoteRunDeserializerOutNode.receive_data /remote_run/data/ POST handler key={key!r} status={status!r}") if not key: return web.Response(status = 400, text = "Missing key") notif = data_requests.get(key) if notif is None: return web.Response(status = 400, text = "No data requested for this key") if status != "data": notif.update("remote_error", f"remote side error status: {status!r}") return web.Response(status = 200, text = f"Error status: {status!r}") notif.update("started", None) try: data = await request.read() notif.update("data", data) except Exception as ex: notif.update("error", ex) return web.Response(status = 200, text = "OK") # ins.app.add_routes(api_routes) receive_data._data_requests = data_requests if len(data_requests) > 5: logger.warning("%s", f"RemoteRunDeserializerOutNode.setup_server has {len(data_requests)} data requests pending, shouldn't happen") return ins, data_requests class RemoteRunDeserializerOutNode(): TYPE_NAME = "RAT_RemoteRunDeserializerOut" DISPLAY_NAME = "RemRun Deserializer Output (Internal)" CATEGORY = "Remote Run/__Internal__/" FUNCTION = "run" @classmethod def INPUT_TYPES(cls): inputs = { } for num in range(_NUM_OUTPUTS): inputs[f"input_{num}_info"] = (CoIO.ANY, { "lazy": True }) inputs[f"input_{num}_data"] = (CoIO.ANY, { "lazy": True }) return { "required": { }, "optional": { **inputs, "total_timeout": ("FLOAT", { }), "request_timeout": ("FLOAT", { }), "is_changed": ("BOOLEAN", { "default": False }), }, } RETURN_TYPES = tuple([CoIO.ANY] * _NUM_OUTPUTS) def run(self, total_timeout: float = None, request_timeout: float = None, **kwargs): get_timeout = _make_adjusted_timeout(total_timeout, allow_empty = True) info_keys = [i for i in kwargs.keys() if i.startswith("input_") and i.endswith("_info")] nums = sorted(int(i[6:-5]) for i in info_keys) outputs = { } for in_num in nums: info_str = kwargs[f"input_{in_num}_info"] data = kwargs[f"input_{in_num}_data"] if info_str == "INPUT": outputs[in_num] = data continue info = json.loads(info_str) data_type = info["type"] if data_type == "serialized": outputs[in_num] = deserialize_response(info["serialization"], data, info.get("compression")) continue if data_type == "request": key = info["key"] serialization = info["serialization"] if serialization not in SERIALIZATION_OPTIONS: raise ValueError("Invalid serialization option", serialization, SERIALIZATION_OPTIONS, info) server_ins, weak_data_requests = setup_remote_run_api_route(expected_set_up = True) client_id = server_ins.client_id if not client_id: raise ValueError("No client connected to request data from") class Notif(): def __init__(self): self.resq = queue.Queue() def update(self, status, data): self.resq.put((status, data)) def cleanup(self): _res = weak_data_requests.pop(key, None) # print("data_pop", key, _res) notif = Notif() try: weak_data_requests[key] = notif logger.info("%s", f"RemoteRunDeserializerOutNode.run sending NEED_DATA data key={key!r}") server_ins.send_sync("NEED_DATA", dict(key = key), client_id) status, _data = notif.resq.get(timeout = get_timeout(request_timeout)) if status == "remote_error": raise ValueError("remote side error requesting data", _data, key, info) if status != "started": raise ValueError("unexpected status waiting for data", status, key, info, _data) result, data = notif.resq.get(timeout = get_timeout(request_timeout)) logger.info("%s", f"RemoteRunDeserializerOutNode.run received data key={key!r} {data and len(data)} bytes") finally: notif.cleanup() if result != "data": raise ValueError("error receiving data", result, key, info, data) deserialized = deserialize_response(serialization, data, info.get("compression")) outputs[in_num] = deserialized continue raise ValueError("Invalid input data type", in_num, info) return tuple(outputs.get(i, None) for i in range(_NUM_OUTPUTS)) class RemoteRunTogglerNode(): TYPE_NAME = "RAT_RemoteRunToggler" DISPLAY_NAME = "RemRun Toggle" """ This is a Input Toggle Switch that can be used to enable/disable certain parts of a graph in the exact opposite way between the local and remote side using lazy evaluation or prompt preprocessing. If the switch run_side is set to enabled it will run all inputs, otherwise set them all to off/lazy or disconnected before prompt execution, depending on [inputs_when_off] setting. When a RemoteRun runner node encounters one of these switches in a prompt it just toggles the run_side value so the remote side will do the exact opposite of the local one, without even knowing it's the remote side. """ CATEGORY = "Remote Run" FUNCTION = "run" @classmethod def INPUT_TYPES(cls): inputs = { "required": { "enabled": (TOGGLE_CHOICES, { }), "inputs_when_off": (INPUTS_OFF_OPTIONS, { "default": "lazy" }), "when_off": (BLOCK_OPTIONS, { "default": "passthrough" }), }, "optional": { **dict((f"input_{num}", (CoIO.ANY, { "lazy": True })) for num in range(_NUM_OUTPUTS)), "_ignore_": ("INTERNAL", { "default": "" }), }, # "hidden": { "own_id": "UNIQUE_ID" }, # "hidden": { "prompt": "PROMPT" }, "hidden": { "own_id": "UNIQUE_ID", # "prompt": "PROMPT", } } return inputs RETURN_TYPES = tuple([CoIO.ANY] * _NUM_OUTPUTS) # @classmethod # def IS_CHANGED(self, **_kwargs): # return None def check_lazy_status(self, enabled: str, own_id = None, **kwargs): if enabled not in TOGGLE_CHOICES: raise ValueError("Invalid enabled value", own_id, enabled, TOGGLE_CHOICES) if enabled == TOGGLE_CHOICES[0]: inputs_with_connections = [i for i in kwargs.keys() if i.startswith("input_")] # print("RemoteRunToggler.check_lazy_status", enabled, own_id, inputs_with_connections, "all:", { k: (v is not None) for k, v in kwargs.items() }) return inputs_with_connections # print("RemoteRunToggler.check_lazy_status", enabled, own_id, [], "all:", { k: (v is not None) for k, v in kwargs.items() }) return [] def run(self, enabled: str, when_off: str, _ignore_ = None, own_id: str = None, **kwargs): disconnected_input_ids = _ignore_ # print(own_id, enabled, _ignore_) print(f"{self.__class__.__name__}.run {own_id=} {enabled=} {when_off=} {disconnected_input_ids=}") # print("RemoteRunToggler.prompt", prompt) if enabled == TOGGLE_CHOICES[1]: error_msg = f"{self.DISPLAY_NAME!r} node output is off and set to 'error'" return block_option_return(when_off or "passthrough", kwargs, _NUM_OUTPUTS, error_msg, "RemoteRunToggler") return tuple(kwargs.get(f"input_{num}", None) for num in range(_NUM_OUTPUTS)) def _shared_input_types(): return { "remote_url": ("STRING", { "default": "http://127.0.0.1:8189/", "tooltip": "ComfyUI instance URL", }), "total_timeout": ("FLOAT", { "tooltip": "Timeout for everything total (sending prompt and running) in seconds. 0 to disable.", "default": 90.0, "min": 0, "max": 24 * 3600.0, "step": 0.1 }), "request_timeout": ("FLOAT", { "tooltip": "Prompt & data post timeout in seconds", "default": 45.0, "min": 0.1, "max": 24 * 3600.0, "step": 0.1 }), **serialization_load_input(), # **output_nodes_input(), } def _get_max_size(max_size_mb: int | None): if not max_size_mb or max_size_mb <= 0: return None, None max_bytes = max_size_mb * 1024 * 1024 ws = math.ceil(max_bytes * 1.1) return max_bytes, ws class RemoteRunInputNode(): TYPE_NAME = "RAT_RemoteRunInput" DISPLAY_NAME = "RemRun Input Graph(s)" NUM_OUTPUTS = _NUM_OUTPUTS CATEGORY = "Remote Run" FUNCTION = "run" @classmethod def IS_CHANGED(cls, is_changed: bool | str = "auto", _ignore_ = None, inputs_when_local = None, **_kwargs): if isinstance(is_changed, bool): return time.time_ns() if is_changed else None if is_changed == "run_every_time": return time.time_ns() if is_changed == "default_behaviour": return None if is_changed == "auto": if inputs_when_local == "disconnected": return _ignore_ return None return None @classmethod def INPUT_TYPES(cls): inputs = { "required": { **_shared_input_types(), }, "optional": { **dict((f"input_{num}", (CoIO.ANY, { "lazy": True })) for num in range(_NUM_OUTPUTS)), **shared_input_output_settings(), **lazy_input_settings(), **response_input(), **ws_settings_inputs(), "run_outputs_connected_to_inputs": (["ignore", "run"], { "default": "ignore" }), "inputs_when_local": (INPUTS_OFF_OPTIONS, { "default": "lazy" }), "is_changed": (["run_every_time", "auto", "default_behaviour"], { "default": "auto", "tooltip": """ run_every_time: Always treat as changed and re-run remote prompt every time. auto: When inputs_when_local is set to "disconnected" this should fix it not re-rerunning if the inputs changed. default_behaviour: Only rerun when inputs changed, this will not work correctly when inputs_when_local is "disconnected" as it can't detect changes then. """ }), "run_sync": ("BOOLEAN", { "default": False }), "_ignore_": ("INTERNAL", { "default": "" }), }, "hidden": { "dynprompt": "DYNPROMPT", "own_id": "UNIQUE_ID", } } return inputs RETURN_TYPES = tuple([CoIO.ANY] * _NUM_OUTPUTS) RETURN_NAMES = tuple([f"out_{num}" for num in range(_NUM_OUTPUTS)]) def check_lazy_status(self, **_kwargs): return [] def run(self, remote_url: str, total_timeout: float, request_timeout: float, serialization: str, max_size_mb: int | None = None, gzip_compression: bool = False, gzip_level: int = 9, response = None, run_outputs_connected_to_inputs = "ignore", is_changed: bool = False, ws_compression: bool = False, forward_progress_messages: str = "off", lazily_transfer: bool = False, skip_lazy_transfer_under_mb: float = 0.25, run_sync: bool = False, dynprompt: DynamicPrompt = None, own_id = None, _ignore_ = None, **_kwargs, ): prompt = { i: dynprompt.get_node(i) for i in dynprompt.all_node_ids() } return partial_json_expansion( prompt, own_id, remote_url, serialization, total_timeout, request_timeout, max_size_mb, gzip_compression, gzip_level, ws_compression, forward_progress_messages, is_changed, response, lazily_transfer, skip_lazy_transfer_under_mb, run_sync, run_outputs_connected_to_inputs, ) class RemoteRunInputOutputNode(RemoteRunInputNode): TYPE_NAME = "RAT_RemoteRunInputOutput" DISPLAY_NAME = "RemRun Input Graph(s) Output" OUTPUT_NODE = True class RemoteRunJsonNode(): TYPE_NAME = "RAT_RemoteRunJson" DISPLAY_NAME = "RemRun JSON" NUM_OUTPUTS = _NUM_OUTPUTS CATEGORY = "Remote Run" FUNCTION = "run" @classmethod def IS_CHANGED(cls, is_changed: bool = False, **_kwargs): return time.time_ns() if is_changed else None @classmethod def INPUT_TYPES(cls): inputs = { "required": { **_shared_input_types(), "JSON": ("STRING", { }), }, "optional": { **dict((f"input_{num}", (CoIO.ANY, { })) for num in range(_NUM_OUTPUTS)), **shared_input_output_settings(), **lazy_input_settings(), **response_input(), **ws_settings_inputs(), "is_changed": ("BOOLEAN", { "default": False }), }, } return inputs RETURN_TYPES = tuple([CoIO.ANY] * _NUM_OUTPUTS) def preprocess(self, remote_url: str, total_timeout: float, request_timeout: float, JSON: str, serialization: str, max_size_mb: int | None = None, gzip_compression: bool = False, gzip_level: int = 9, response = None, ws_compression: bool = False, forward_progress_messages: str = "off", lazily_transfer: bool = False, skip_lazy_transfer_under_mb: float = 0.25, **kwargs, ): max_size_mb = max_size_mb or None run_obj = json.loads(JSON) extra_output_ids = run_obj.get("extra_output_ids") or None run_prompt = run_obj["prompt"] data_total = 0 lazy_data = { } skip_lazy_transfer_under_mb = skip_lazy_transfer_under_mb * 1024 ** 2 if skip_lazy_transfer_under_mb else None for node_id, node in run_prompt.items(): # serialize needed inputs to this node and set up # Deserializer data output nodes with the serialized data in the remote prompt if node["class_type"] == RemoteRunDeserializerOutNode.TYPE_NAME: config = node.pop("deserializer_config") outputs = config["outputs"] other_inputs = config["other_inputs"] node_lazily_transfer = other_inputs.get("lazily_transfer", None) if node_lazily_transfer is not None and node_lazily_transfer != "from_input_node": node_lazily_transfer = { "on": True, "off": False }[node_lazily_transfer.lower()] node_skip_under = other_inputs["skip_lazy_transfer_under_mb"] else: node_lazily_transfer = lazily_transfer node_skip_under = skip_lazy_transfer_under_mb inputs = { } for output_name, output_data in outputs.items(): if not isinstance(output_data, list) or not len(output_data) == 2: raise ValueError("Invalid output data", output_data, config) if not output_name.startswith("output_") and output_name[7:].isdigit(): raise ValueError("Invalid output data", output_data, config) output_num = int(output_name[7:]) output_data_type, output_data_value = output_data if output_data_type == "RunJsonInput": if not output_data_value.startswith("input_") and output_data_value[6:].isdigit(): raise ValueError("Invalid output data", output_data, config) local_input_name = output_data_value input = kwargs[local_input_name] # TODO: skip serialization if input is basic python types input_data = serialize_obj(serialization, input) compression = None if gzip_compression: start = time.monotonic_ns() comp = gzip.compress(input_data, compresslevel = gzip_level) end = time.monotonic_ns() took_ns = (end - start) mb_s = (len(input_data) * 1e9) / (took_ns * 1024 * 1024) took_sec = took_ns / 1e9 factor = len(comp) / len(input_data) if factor < 0.99: logger.info("%s", f"RemoteRunJsonNode.run serialization={serialization!r} {response=!r} " f"gzip compressed {len(input_data)} to {len(comp)} bytes, factor {factor:.3f}, " f"took {took_sec:.3f} sec, {mb_s:.2f} MB/s") input_data = comp compression = "gzip" else: logger.info("%s", f"RemoteRunJsonNode.run serialization={serialization!r} {response=!r} " f"gzip compression not effective, kept original {len(input_data)} bytes " f"over compressed {len(comp)} bytes, factor {factor:.3f}, " f"took {took_sec:.3f} sec, {mb_s:.2f} MB/s") if node_lazily_transfer and node_skip_under and len(input_data) < node_skip_under: node_lazily_transfer = False if node_lazily_transfer: data_total += len(input_data) if max_size_mb and data_total > max_size_mb * 1024 * 1024: raise ValueError(f"Serialized data too large: {data_total} bytes, max_size_mb={max_size_mb}") key = hashlib.sha256(hashlib.sha256(input_data).digest()).hexdigest() lazy_data[key] = input_data inputs[f"input_{output_num}_info"] = json.dumps({ "type": "request", "key": key, "serialization": serialization, "compression": compression, }) inputs[f"input_{output_num}_data"] = True else: input_data_b64 = base64.b64encode(input_data).decode() data_total += len(input_data_b64) if max_size_mb and data_total > max_size_mb * 1024 * 1024: raise ValueError(f"Serialized data too large: {data_total} bytes, max_size_mb={max_size_mb}") inputs[f"input_{output_num}_info"] = json.dumps({ "type": "serialized", "serialization": serialization, "compression": compression, }) inputs[f"input_{output_num}_data"] = input_data_b64 elif output_data_type == "CONSTANT": inputs[f"input_{output_num}_info"] = "INPUT" inputs[f"input_{output_num}_data"] = output_data_value else: raise ValueError("Invalid output data", output_data, config) node["inputs"] = { "serialization": serialization, **inputs, } serializers = { k: v for k, v in run_prompt.items() if v["class_type"] == RemoteRunSerializerOutNode.TYPE_NAME } if len(serializers) != 1: raise ValueError(f"expected exactly one {RemoteRunSerializerOutNode.TYPE_NAME!r} serializer node, got {len(serializers)}", serializers, run_prompt) serialized_id = list(serializers.keys())[0] post_max, ws_max = _get_max_size(max_size_mb) binary_response = response == "binary" fn = lambda: remote_execute_prompt( remote_url, run_prompt, serialized_id, extra_output_ids = extra_output_ids, ws_compression = ws_compression, forward_progress_messages = forward_progress_messages, binary_response = binary_response, ws_max_size = ws_max, post_max_size = post_max, total_timeout = total_timeout, short_timeout = request_timeout, lazy_data = lazy_data, ) return fn async def run(self, *args, **kwargs): fn = self.preprocess(*args, **kwargs) thread_exec = concurrent.futures.ThreadPoolExecutor() raw_result = await asyncio.get_event_loop().run_in_executor(thread_exec, fn) serialization = kwargs["serialization"] obj = deserialize_response(serialization, raw_result) result = tuple(obj.get(f"input_{num}", None) for num in range(_NUM_OUTPUTS)) return result class RemoteRunJsonSyncNode(RemoteRunJsonNode): TYPE_NAME = "RAT_RemoteRunJsonSync" DISPLAY_NAME = "RemRun JSON Sync" def run(self, *args, **kwargs): fn = self.preprocess(*args, **kwargs) raw_result = fn() serialization = kwargs["serialization"] obj = deserialize_response(serialization, raw_result) result = tuple(obj.get(f"input_{num}", None) for num in range(_NUM_OUTPUTS)) return result def block_option_return(option: str, kwargs, num_outputs, error_msg, block_msg): option = option.lower() if option == "passthrough": return tuple(kwargs.get(f"input_{num}", None) for num in range(num_outputs)) elif option == "block_execution_silent": return tuple(execution.ExecutionBlocker(None) for _ in range(num_outputs)) elif option == "block_execution_verbose": return tuple(execution.ExecutionBlocker(f"{block_msg} out={num}") for num in range(num_outputs)) elif option == "error": raise ValueError(error_msg) elif option == "return_none": return tuple(None for _ in range(num_outputs)) else: raise ValueError(f"Invalid option value: {option!r}") class RemoteRunStartNode(): TYPE_NAME = "RAT_RemoteRunStart" DISPLAY_NAME = "RemRun Start From Here ->" CATEGORY = "Remote Run" FUNCTION = "run" @classmethod def INPUT_TYPES(cls): inputs = { "optional": { **dict((f"input_{num}", (CoIO.ANY, { })) for num in range(_NUM_OUTPUTS)), "remote_run_dependent_outputs": (["ignore", "run"], { "default": "ignore" }), "inputs_when_local": (INPUTS_OFF_OPTIONS, { "default": "disconnected" }), "outputs_when_local": (BLOCK_OPTIONS, { "default": "error" }), "lazily_transfer": (["from_input_node", "on", "off"], { "default": "from_input_node" }), **{ k: v for k, v in lazy_input_settings().items() if k == "skip_lazy_transfer_under_mb" }, }, } return inputs @classmethod def VALIDATE_INPUTS(cls, outputs_when_local: str = None, remote_run_dependent_outputs: str = None, *_args, **_kwargs): outputs_when_local = outputs_when_local or "error" if outputs_when_local.lower() not in BLOCK_OPTIONS: return f"Invalid outputs_when_local value: {outputs_when_local!r}" remote_run_dependent_outputs = remote_run_dependent_outputs or "ignore" if remote_run_dependent_outputs.lower() not in ("ignore", "run"): return f"Invalid remote_run_dependent_outputs value: {remote_run_dependent_outputs!r}" return True RETURN_TYPES = tuple([CoIO.ANY] * _NUM_OUTPUTS) RETURN_NAMES = tuple([f"out_{num}" for num in range(_NUM_OUTPUTS)]) def run(self, outputs_when_local: str = None, **kwargs): error_msg = f"{self.DISPLAY_NAME!r} node output used locally, but outputs_when_local is set to 'error'" return block_option_return(outputs_when_local or "error", kwargs, _NUM_OUTPUTS, error_msg = error_msg, block_msg = "RemoteRunStartNode") def _max_id(prompt): def toint(i): try: return int(i) except ValueError: return 0 if not len(prompt): return 0 return max(toint(i) for i in prompt.keys()) def make_counter(start: int, map = None): def get(): nonlocal start ret = start start += 1 return map(ret) if map is not None else ret return get def get_remote_run_input_graph_nodes(prompt: dict, root_node_id: str) -> tuple[set[str], set[str]]: start_nodes = set() input_nodes = set() def _nodes(node_id: str, seen): if node_id in seen: return seen.add(node_id) node = prompt[node_id] if node.get("class_type") == RemoteRunStartNode.TYPE_NAME: start_nodes.add(node_id) return inputs = node["inputs"] for inp in inputs.values(): if isinstance(inp, list): input_nodes.add(inp[0]) _nodes(inp[0], seen) _nodes(root_node_id, set()) return start_nodes, input_nodes def get_input_graph_nodes(prompt: dict, root_node_id: str) -> set[str]: input_nodes = set() def _nodes(node_id: str, seen): if node_id in seen: return seen.add(node_id) node = prompt[node_id] inputs = node["inputs"] for inp in inputs.values(): if isinstance(inp, list): input_nodes.add(inp[0]) _nodes(inp[0], seen) _nodes(root_node_id, set()) return input_nodes def get_full_dependents_of_outputs(prompt: dict, of_main_node_ids: set[str], stop_at_node_ids = None) -> dict[str, set[str]]: # For each main_node_id get all nodes that are directly or indirectly linked to its outputs # but stop going down the graph when hitting another stop_at_node_ids (or main_node_id if stop_at_node_ids not given). of_main_node_ids = set(of_main_node_ids) stop_at_node_ids = set(stop_at_node_ids if stop_at_node_ids is not None else of_main_node_ids) inputs_map = { } for nid, node in prompt.items(): inputs = set() for v in node["inputs"].values(): if isinstance(v, list) and len(v) == 2: inputs.add(v[0]) inputs_map[nid] = inputs @functools.cache def get_all_parents(nid): if nid in stop_at_node_ids: return set() parents = set() for inp in inputs_map.get(nid) or []: parents.add(inp) parents.update(get_all_parents(inp)) return parents res = { i: set() for i in of_main_node_ids } for nid in prompt.keys(): all_parents_until_parent_node_id_or_end = get_all_parents(nid) important_parents = all_parents_until_parent_node_id_or_end & of_main_node_ids for p in important_parents: res[p].add(nid) return res def get_extra_output_ids_dependent_on(prompt: dict, node_ids: set[str], stop_at_node_ids = None): # get all nodes that are directly or indirectly linked to the outputs of the given node_ids # (generally all RemoteRunStartNodes), stop going down the chain when hitting another remote_run # start or input/json node. node_ids = set(node_ids) stop_at_node_ids = set(stop_at_node_ids if stop_at_node_ids is not None else node_ids) deps = get_full_dependents_of_outputs(prompt, node_ids, stop_at_node_ids) extra_output_ids = set() for parent_id, all_deps in deps.items(): if not all_deps or parent_id not in node_ids: continue extra_output_ids.update(all_deps) return sorted(extra_output_ids) def partial_json_expansion( prompt: dict, root_node_id: str, remote_url: str, serialization: str, total_timeout: float, request_timeout: float, max_size_mb: int | None, gzip_compression: bool, gzip_level: int, ws_compression: bool, forward_progress_messages: str, is_changed: bool, response: str | None, lazily_transfer: bool, skip_lazy_transfer_under_mb: float, run_sync: bool, run_outputs_connected_to_inputs: str, ): """ When there are RemoteStart nodes then go from: A -> B -> C -> RemoteStart -> D -> E -> F -> RemoteRunInput -> G -> H -> I to A -> B -> C -> RemoteRunJson(run_remotely="D -> E -> F") -> G -> H -> I Build an expanded graph to replace the RemoteRunInput node with a RemoteRunJson node copying the whole input graph going to the RemoteRunInput node up until RemoteStart nodes are hit, and replace all those nodes with a RemoteRunJson node that will run that part as prompt remotely instead. The RemoteRunJson gets linked to all the inputs of the RemoteStart nodes if any so it can serialize them and send them to the remote instance as part of the remote prompt. Optionally if remote_run_dependent_outputs==run then also include all other nodes that are connected to the outputs of any input graph nodes but don't lead to the RemoteRunInput node and run outputs in those too, see below for example, this can be set either for everything if the RemoteRunInput node has it set or per RemoteRunStart node for anything going off a node in between that and the RemoteRunInput node. If there are no RemoteStart nodes then take all the nodes going into RemoteRunInput. Optionally also all nodes that go off those inputs nodes too if run_outputs_connected_to_inputs is set to "run", so if there are any outputs as part of the input graph but that doesn't lead to RemoteRunInput those get run too. So from: A -> B -> C -> D -> RemoteRunInput -> G -> H -> I A -> SaveImage[F] to: RemoteRunJson(run_remotely="A -> B -> C -> D") -> G -> H -> I or with run_outputs_connected_to_inputs = "run" to: RemoteRunJson(run_remotely="A -> SaveImage[F]; A -> B -> C -> D", run_extra_outputs=["F"]) -> G -> H -> I """ remote_prompt = copy.deepcopy(prompt) next_node_id = make_counter(_max_id(remote_prompt) + 1, str) rcon, dscon = update_toggles_inplace(remote_prompt, True, True, "remote_run_input_toggle", False, True) add_outputs_of_inputs = (run_outputs_connected_to_inputs or "").lower() == "run" if rcon or dscon: logger.info("%s", f"partial_json_expansion {root_node_id=} toggled {rcon + dscon} nodes, " f"reconnected {rcon}, disconnected {dscon}") extra_output_ids = None local_expanded_inputs = { } remote_run_dependent_outputs_of_nodes = [] start_node_ids, other_node_ids = get_remote_run_input_graph_nodes(remote_prompt, root_node_id) # The D -> E -> F part(s), anything before the RemoteRunInput node back up to any RemoteRunStart nodes # and anything starting on its own. small_prompt = { i: remote_prompt[i] for i in (start_node_ids | other_node_ids | { root_node_id }) } if start_node_ids: # For the remote prompt: just replace each RemoteStart node with a Deserializer node that outputs # the serialized and sent over inputs. # On the local side as there can be multiple RemoteStart nodes all the inputs to the RemoteStart nodes have to also # be linked up to the newly created/expanded RemoteRunJson node and kept track of which serialized inputs goes to which # Deserializer node output on the remote side that replace the RemoteStart nodes in the remote prompts. next_json_input_num = make_counter(0) # On remote prompt change all RemoteRunStart nodes to Deserializer nodes for start_id in start_node_ids: # TODO: ignore inputs for outputs that aren't used, currently serialized/sent for nothing start_node = small_prompt[start_id] start_inputs = start_node["inputs"] node_remote_run_dependent_outputs = start_inputs.get("remote_run_dependent_outputs") if (node_remote_run_dependent_outputs or "").lower() == "run": remote_run_dependent_outputs_of_nodes.append(start_id) des = { } other_inputs = { } for input_name, input_value in start_inputs.items(): if not input_name.startswith("input_"): other_inputs[input_name] = input_value continue output_name = "output_" + input_name[6:] if isinstance(input_value, list): local_input_name = f"input_{next_json_input_num()}" local_expanded_inputs[local_input_name] = input_value des[output_name] = ("RunJsonInput", local_input_name) else: des[output_name] = ("CONSTANT", input_value) # The deserializer_config part will be removed by the RemoteRunJson node and rewritten # to use the serialized inputs as given. small_prompt[start_id] = { "class_type": RemoteRunDeserializerOutNode.TYPE_NAME, "deserializer_config": { "outputs": des, "other_inputs": other_inputs, }, } if add_outputs_of_inputs: input_node_ids = start_node_ids | other_node_ids stop_at = start_node_ids | { root_node_id } extra_output_ids = get_extra_output_ids_dependent_on(remote_prompt, input_node_ids, stop_at_node_ids = stop_at) logger.info("partial_json_expansion %s: add_outputs_of_inputs -> extra_output_ids=%s", root_node_id, (len(extra_output_ids), extra_output_ids)) elif remote_run_dependent_outputs_of_nodes: # Add all output nodes that go off away from the input graph, so between a start or StartNode to RemoteRunInput node, # but don't themselves lead to the RemoteRunInput node. extra_output_ids = get_extra_output_ids_dependent_on(remote_prompt, set(remote_run_dependent_outputs_of_nodes), stop_at_node_ids = { root_node_id }) logger.info("partial_json_expansion %s: remote_run_dependent_outputs_of_nodes=%s -> extra_output_ids=%s", root_node_id, remote_run_dependent_outputs_of_nodes, (len(extra_output_ids), extra_output_ids)) # This is really all dependent nodes not just Output ones but no need to limit to only output ones since any normal non output # nodes given will just be ignored by comfy prompt handling (currently). extra_output_ids = list(extra_output_ids) if extra_output_ids else None for i in extra_output_ids or []: if i not in small_prompt: small_prompt[i] = remote_prompt[i] remote_prompt = small_prompt root_node = remote_prompt.pop(root_node_id) root_inputs = { k: v for k, v in root_node["inputs"].items() if isinstance(v, list) } remote_prompt[root_node_id] = { "class_type": RemoteRunSerializerOutNode.TYPE_NAME, "inputs": { "serialization": serialization, "max_size_mb": max_size_mb, "is_changed": is_changed, **root_inputs, }, } DEBUG_PROMPT_SAVE("exported_remote_prompt", remote_prompt) if IS_DEV and IS_DEV_PROMPTSAVE else None remote_obj = dict( prompt = remote_prompt, extra_output_ids = extra_output_ids, ) json_node_id_str = next_node_id() new_graph = { json_node_id_str: { "class_type": RemoteRunJsonSyncNode.TYPE_NAME if run_sync else RemoteRunJsonNode.TYPE_NAME, "inputs": dict( JSON = json.dumps(remote_obj), remote_url = remote_url, total_timeout = total_timeout, request_timeout = request_timeout, serialization = serialization, max_size_mb = max_size_mb, gzip_compression = gzip_compression, gzip_level = gzip_level, ws_compression = ws_compression, forward_progress_messages = forward_progress_messages, # is_changed = is_changed, is_changed = True, response = response, lazily_transfer = lazily_transfer, skip_lazy_transfer_under_mb = skip_lazy_transfer_under_mb, **local_expanded_inputs, ), } } # print("new_graph", json.dumps(new_graph, indent = 4)) replaced_output_ids = tuple([json_node_id_str, i] for i in range(_NUM_OUTPUTS)) return { "result": replaced_output_ids, "expand": new_graph, } def serialize_obj(serialization: str, obj): if serialization in ("torch_pt", "unsafe_torch_pt", "safe_torch_pt"): buf = BytesIO() torch.save(obj, f = buf) buf = buf.getvalue() elif serialization == "fancy_safetensors": buf = fst_torch.save_bytes(obj, transform_tensor_fn = lambda t: t.detach().clone().contiguous()) else: raise ValueError("Unknown serialization method", serialization) return buf def deserialize_response(serialization: str, result: dict | str | bytes, compression: str | None = None): if isinstance(result, dict): result_str = result["results"][0] data = base64.b64decode(result_str) elif isinstance(result, str): data = base64.b64decode(result) elif isinstance(result, bytes): data = result else: raise ValueError("Invalid result type", type(result), result) if compression == "gzip": data = gzip.decompress(data) elif compression is not None: raise ValueError("Unknown compression", compression) return deserialize_data(serialization, data) def deserialize_data(serialization: str, data: bytes): if serialization == "unsafe_torch_pt": return torch.load(BytesIO(data)) elif serialization == "safe_torch_pt": return torch_safe_load_dict(ZipFile(BytesIO(data), mode = "r")) elif serialization == "fancy_safetensors": return fst_torch.load(data) else: raise ValueError("Unknown serialization method", serialization) def update_toggles_inplace(prompt: dict, toggle: bool, reconnect: bool, label: str, raise_on_existing: bool, input_node_is_local: bool | None = None): org_prompt = copy.deepcopy(prompt) to_switch = { RemoteRunTogglerNode.TYPE_NAME: "inputs_when_off", RemoteRunInputNode.TYPE_NAME: "inputs_when_local", RemoteRunInputOutputNode.TYPE_NAME: "inputs_when_local", } reconnected_cnt = disconnected_cnt = 0 for node_id, node in prompt.items(): class_type = node["class_type"] input_name = to_switch.get(class_type) if not input_name: continue inputs = node.get("inputs") if not inputs: continue if class_type == RemoteRunTogglerNode.TYPE_NAME: enabled = inputs.get("enabled") if enabled not in TOGGLE_CHOICES: raise ValueError("Invalid enabled value", node_id, enabled, TOGGLE_CHOICES) if toggle: enabled = inputs["enabled"] = TOGGLE_CHOICES[1] if enabled == TOGGLE_CHOICES[0] else TOGGLE_CHOICES[0] enabled = enabled == "only_locally" else: enabled = input_node_is_local if ( reconnect and enabled is True and inputs.get(input_name) == "disconnected" and (disconnected_info_str := inputs.get("_ignore_")) ): # restore original disconnected inputs for remote side try: disconnected_info = json.loads(disconnected_info_str) disconnected_inputs = disconnected_info["inputs"] except Exception as e: logger.error("%s RemoteRunToggler.run invalid disconnected_inputs %r: %s", label, disconnected_info_str, e) raise ValueError("Invalid disconnected_inputs", node_id, disconnected_info_str, e) if disconnected_inputs: for inp_name, inp_val in disconnected_inputs.items(): if inp_name in inputs: raise ValueError("RemoteRunToggler.run disconnected_inputs input already set", node_id, inp_name, inp_val) inputs[inp_name] = inp_val inputs.pop("_ignore_") logger.debug("%s reconnected %s RemoteRunToggler inputs on node %r: %s", label, len(disconnected_inputs), node_id, disconnected_inputs) reconnected_cnt += 1 elif enabled is False and inputs.get(input_name) == "disconnected": given_inputs = { k: v for k, v in inputs.items() if k.startswith("input_") } for k in given_inputs.keys(): inputs.pop(k) if not given_inputs: continue cur = inputs.get("_ignore_") if cur: if raise_on_existing: raise ValueError("remote_run update_toggles_inplace overwriting existing _ignore_", node_id, cur, given_inputs) else: logger.warning("%s overwriting existing _ignore_ on node_id=%r, cur=%r, new=%r", label, node_id, cur, given_inputs) disconnected_info = { "inputs": given_inputs } if class_type in (RemoteRunInputNode.TYPE_NAME, RemoteRunInputOutputNode.TYPE_NAME): if inputs.get(input_name) == "disconnected": input_ids = get_input_graph_nodes(org_prompt, node_id) input_prompt = { i: org_prompt[i] for i in input_ids } disconnected_info["input_prompt"] = input_prompt inputs["_ignore_"] = json.dumps(disconnected_info) logger.debug("%s disconnected inputs for node_id=%r, inputs=%s", label, node_id, given_inputs) disconnected_cnt += 1 return reconnected_cnt, disconnected_cnt def retry_dec(total_tries: int = 1, log_exc = True, sleep = None, stop_on = None): if isinstance(stop_on, type): stop_on = (stop_on,) stop_on = stop_on or [] def make_dec(func): def wrapper(*args, **kwargs): for i in range(total_tries): try: return func(*args, **kwargs) except Exception as e: if log_exc: logger.exception("function error: fn:%r, retrying %d", func, i) if i == total_tries - 1 or (stop_on and type(e) in stop_on): raise e if sleep: time.sleep(sleep) return wrapper return make_dec def _make_adjusted_timeout(total_timeout: float | None, allow_empty: bool = True): if not total_timeout: if not allow_empty: raise ValueError("total_timeout must be > 0") def get(wanted_timeout: float, raise_on_timeout = False): if raise_on_timeout and wanted_timeout <= 0: raise TimeoutError("wanted timeout reached", wanted_timeout) return wanted_timeout return get if total_timeout < 0: raise ValueError("total_timeout must be >= 0", total_timeout) # ensure current timeout wouldn't exceed total_timeout start_ts = time.monotonic() def get(wanted_timeout: float, raise_on_timeout = False): left_secs = total_timeout - (time.monotonic() - start_ts) if raise_on_timeout and left_secs <= 0: raise TimeoutError("total timeout reached", total_timeout) return min(left_secs, wanted_timeout) return get def remote_execute_prompt( base_url: str, prompt: dict, own_id: str, extra_output_ids: list[str] | None, ws_compression: bool, forward_progress_messages: str, binary_response: bool, ws_max_size: int | None, post_max_size: int | None, total_timeout: float, short_timeout: float, lazy_data: dict | None = None, total_tries = 3, ): # TODI: make async ws_max_size = ws_max_size or 64 * 1024 * 1024 total_timeout = total_timeout # can be 0 to disable short_timeout = short_timeout or 5 # expected to always be set, so default to low ones in case of unset base_url = base_url.rstrip("/") base_ws_url = base_url.replace("http://", "ws://") client_id = str(uuid4()) post_nonce = str(uuid4()) adjusted_timeout_fn = _make_adjusted_timeout(total_timeout or None, allow_empty = True) @retry_dec(total_tries, stop_on = TimeoutError) def post(timeout): post_url = base_url + "/api/prompt" apicall = dict( client_id = client_id, prompt = prompt, nonce = post_nonce, # for safe retry, assumes nonce duplicate node/extension installed partial_execution_targets = list(extra_output_ids or []) + [own_id], ) DEBUG_PROMPT_SAVE("post_prompt", prompt) if IS_DEV and IS_DEV_PROMPTSAVE else None body = json.dumps(apicall) if isinstance(body, str): body = body.encode() if post_max_size and len(body) > post_max_size: raise ValueError(f"Post body too large: {len(body)} bytes, post_max_size={post_max_size}") logger.info("Posting prompt to url=%r, size=%d, timeout=%.1f", post_url, len(body), timeout) headers = { "Content-Type": "application/json" } resp = requests.post(post_url, data = body, headers = headers, timeout = adjusted_timeout_fn(timeout, True)) if resp.status_code != 200: raise ValueError("Failed to post prompt", post_url, resp.status_code, resp.content, apicall) obj = resp.json() return obj["prompt_id"] @retry_dec(total_tries, stop_on = (TimeoutError, ValueError)) def post_lazy_data(key, status, body: bytes | None, timeout): params = dict( key = key, status = status, ) post_url = base_url + "/api/remote_run/data/" headers = { } logger.info("Posting requested lazy data to url=%r, params=%r, size=%d, timeout=%.1f", post_url, params, len(body or ""), timeout) resp = requests.post(post_url, params = params, data = body, headers = headers, timeout = adjusted_timeout_fn(timeout, True)) if resp.status_code != 200: raise ValueError("Failed to post prompt", post_url, resp.status_code, resp.content) return resp.content ws_url = base_ws_url + f"/ws?clientId={client_id}" websocket = None # create websocket connection first, then send prompt POST and wait for result/error/timeout on ws try: @retry_dec(total_tries, stop_on = TimeoutError) def connect(): return websockets.sync.client.connect( ws_url, open_timeout = adjusted_timeout_fn(short_timeout), close_timeout = 5, max_size = ws_max_size, compression = "deflate" if ws_compression else None, ) logger.debug("Connecting to websocket url=%r", ws_url) websocket = connect() prompt_id = post(short_timeout) forward_ws_messages = forward_progress_messages == "on" forwarded_types = { "executing", "progress_state" } while True: # TODO: option to reconnect websocket? message = websocket.recv(timeout = adjusted_timeout_fn(total_timeout, True)) if isinstance(message, bytes): if not binary_response: logger.warning("Ignoring binary websocket message, expected text: size=%d", len(message)) else: if len(message) < 4: logger.warning("Ignoring too short binary websocket message, expected longer: size=%d", len(message)) continue message_type = int.from_bytes(message[:4], "big") if message_type == BinaryResponseMessageID.BINARY_RESPONSE_MESSAGE_ID.value: logger.info("Got binary response websocket message: size=%d", len(message)) rest = message[4:] return rest elif message_type == BinaryResponseMessageID.GZIPPED_BINARY_RESPONSE_MESSAGE_ID.value: logger.info("Got gzipped binary response websocket message: size=%d", len(message)) comp = message[4:] gzrest = gzip.decompress(comp) return gzrest else: logger.warning("Ignoring unknown binary websocket message type: type=%d, size=%d", message_type, len(message)) continue continue message_len = len(message or "") message_obj = json.loads(message) # print("ws message", message_obj) if IS_DEV else None msg_type = message_obj.get("type") if msg_type == "NEED_DATA": logger.info("Got NEED_DATA message for: %s", message_obj) key = None try: req_obj = message_obj["data"] key = req_obj["key"] if lazy_data and key in lazy_data: data = lazy_data.pop(key) post_lazy_data(key, "data", data, short_timeout) continue else: logger.error("No lazy data found for key %r, sending empty response", key) post_lazy_data(key, "not_found", None, short_timeout) except Exception: logger.exception("Invalid NEED_DATA message data") if key: post_lazy_data(key, "error", None, short_timeout) raise ValueError("Invalid NEED_DATA message", message_obj) own_prompt_id = message_obj.get("data", { }).get("prompt_id") == prompt_id if own_prompt_id: if ( forward_ws_messages and msg_type in forwarded_types and (data := message_obj.get("data")) ): node = data.get("node") ignore = msg_type == "executing" and node == own_id if not ignore: try: import server instance = server.PromptServer.instance client_id = instance.client_id logger.debug("Forwarding remote prompt server message type=%r, node=%r, client_id=%r, size=%d", msg_type, node, client_id, message_len) instance.send_sync(msg_type, data, client_id) except Exception: logger.exception("Failed to forward remote prompt server message type=%r, node=%r, client_id=%r", msg_type, node, client_id) if msg_type == "executed": logger.debug("got executed message for own prompt_id=%r, size=%d", prompt_id, message_len) data = message_obj["data"] if data["node"] == own_id: result = data["output"] return result elif own_prompt_id and msg_type == "execution_error": raise ValueError("remote comfy prompt execution error", message_obj, base_url, prompt) except TimeoutError: # TODO: interrupt execution of that prompt in case of timeout?? always or option # also not sure if specific prompt interrupt is easily possible raise finally: if websocket is not None: threading.Thread(target = websocket.close, daemon = True).start() class RemoteRunSetNumOutputsNode(): DISPLAY_NAME = "RemRun Set Num Outputs" CATEGORY = "Remote Run" FUNCTION = "run" OUTPUT_NODE = True @classmethod def INPUT_TYPES(cls): return { "required": { "num": ("INT", { "default": 10, "min": 1, "max": 100, "step": 1 }), }, } RETURN_TYPES = () def run(self, num: int): global _NUM_OUTPUTS logger.info("Changing RemoteRun number of inputs/outputs from %d to %d", _NUM_OUTPUTS, num) _NUM_OUTPUTS = num return { } def filter_definitions(definitions: dict, kwargs): type_regexes = [v for k, v in kwargs.items() if k.startswith("class_type_include_regex_") and v] display_name_regexes = [v for k, v in kwargs.items() if k.startswith("display_name_include_regex_") and v] ignore_type_regexes = [v for k, v in kwargs.items() if k.startswith("class_type_ignore_regex_") and v] ignore_display_name_regexes = [v for k, v in kwargs.items() if k.startswith("display_name_ignore_regex_") and v] category_name_regexes = [v for k, v in kwargs.items() if k.startswith("category_name_regex_") and v] category_name_ignore_regexes = [v for k, v in kwargs.items() if k.startswith("category_name_ignore_regex_") and v] type_regexes = [re.compile(r, flags = re.IGNORECASE) for r in type_regexes] display_name_regexes = [re.compile(r, flags = re.IGNORECASE) for r in display_name_regexes] ignore_type_regexes = [re.compile(r, flags = re.IGNORECASE) for r in ignore_type_regexes] ignore_display_name_regexes = [re.compile(r, flags = re.IGNORECASE) for r in ignore_display_name_regexes] category_name_regexes = [re.compile(r, flags = re.IGNORECASE) for r in category_name_regexes] category_name_ignore_regexes = [re.compile(r, flags = re.IGNORECASE) for r in category_name_ignore_regexes] matching = { } for class_type, node_class in definitions.items(): if type_regexes and not any(r.search(class_type) for r in type_regexes): continue if ignore_type_regexes and any(r.search(class_type) for r in ignore_type_regexes): continue display_name = node_class.get("display_name", None) if display_name and display_name_regexes and not any(r.search(display_name) for r in display_name_regexes): continue if display_name and ignore_display_name_regexes and any(r.search(display_name) for r in ignore_display_name_regexes): continue category = node_class.get("category", None) if category and category_name_regexes and not any(r.search(category) for r in category_name_regexes): continue if category and category_name_ignore_regexes and any(r.search(category) for r in category_name_ignore_regexes): continue matching[class_type] = node_class return matching class RemoteRunAddRemoteNodeDefinitionsNode(): DISPLAY_NAME = "RemRun Add Remote Node Definitions" CATEGORY = "Remote Run" FUNCTION = "run" OUTPUT_NODE = True @classmethod def INPUT_TYPES(cls): return { "required": { "remote_url": ("STRING", { "default": "http://127.0.0.1:8189/", "tooltip": "ComfyUI instance URL", }), "request_timeout": ("FLOAT", { "tooltip": "Request timeout in seconds", "default": 15.0, "min": 0.1, "max": 24 * 3600.0, "step": 0.1 }), "dry_run": ("BOOLEAN", { "default": False }), "display_name_prefix": ("STRING", { "default": "[RemoteNd] " }), "display_name_suffix": ("STRING", { "default": "" }), "category_name_prefix": ("STRING", { "default": "[RemoteNd]/" }), "category_name_suffix": ("STRING", { "default": "" }), "overwrite_existing": ("BOOLEAN", { "default": False }), "is_changed": ("BOOLEAN", { "default": False }), }, "optional": { **{ f"class_type_include_regex_{i}": ("STRING", { }) for i in range(5) }, **{ f"display_name_include_regex_{i}": ("STRING", { }) for i in range(5) }, **{ f"class_type_ignore_regex_{i}": ("STRING", { }) for i in range(5) }, **{ f"display_name_ignore_regex_{i}": ("STRING", { }) for i in range(5) }, **{ f"category_name_regex_{i}": ("STRING", { }) for i in range(5) }, **{ f"category_name_ignore_regex_{i}": ("STRING", { }) for i in range(5) }, }, } RETURN_TYPES = () @classmethod def IS_CHANGED(cls, is_changed: bool = False, **_kwargs): return time.time_ns() if is_changed else None def run(self, remote_url: str, request_timeout: float, is_changed: float | None = None, **kwargs): if not remote_url.endswith("/"): remote_url += "/" if not remote_url.startswith("http"): raise ValueError("Invalid remote_url, must start with http:// or https://", remote_url) url = remote_url + "object_info" resp = requests.get(url, timeout = request_timeout) if not resp.status_code == 200: raise ValueError("Failed to get remote object info", url, resp.status_code, resp.content) obj = resp.json() logger.info("Got %s node definitions from remote %r", len(obj or { }), remote_url) self.add_definitions(obj, kwargs) return { } def add_definitions(self, definitions: dict, kwargs): matching = filter_definitions(definitions, kwargs) logger.info("Found %d matching node definitions of %s", len(matching), len(definitions)) if not matching: return stats = collections.defaultdict(int) for class_type, node_obj in matching.items(): try: self.add_definition(class_type, node_obj, stats, **kwargs) except Exception as e: logger.error(f"Error adding node definition for class {class_type!r}: {e}") stats["errored"] += 1 added, overwrote, skipped, errored = stats["added"], stats["overwrote"], stats["skipped"], stats["errored"] dry_run_label = "DRYRUN: " if kwargs.get("dry_run") else "" logger.info("%s", f"{dry_run_label}Processed {len(matching)} node definitions " f"(added: {added}, overwrote: {overwrote}, skipped: {skipped}, errors: {errored}) of {len(definitions)} total definitions") def add_definition(self, class_type: str, defn: dict, stats: dict, dry_run: bool = False, overwrite_existing: bool = False, display_name_prefix: str = "", display_name_suffix: str = "", category_name_prefix: str = "", category_name_suffix: str = "", **_kwargs, ): dry_run_label = "DRYRUN: " if dry_run else "" stats = stats if stats is not None else { } existed = class_type in nodes.NODE_CLASS_MAPPINGS if existed and not overwrite_existing: logger.debug("%s", f"{dry_run_label}Skipped existing node class: {class_type}") stats["skipped"] += 1 return display_name = defn.get("display_name", class_type) if display_name_prefix: display_name = f"{display_name_prefix}{display_name}" if display_name_suffix: display_name = f"{display_name}{display_name_suffix}" category = defn.get("category") or "" if category_name_prefix: category = f"{category_name_prefix}{category}" if category_name_suffix: category = f"{category}{category_name_suffix}" def tup_get(key): val = defn.get(key) if val is None: return None if isinstance(val, list): return tuple(val) raise ValueError(f"Invalid {key} value, must be list or None", val) input_types = defn.get("input") @classmethod def INPUT_TYPES(cls): return input_types new_class = type(class_type, (object,), { }) setattr(new_class, "INPUT_TYPES", INPUT_TYPES) for name, val in ( ("DISPLAY_NAME", display_name), ("CATEGORY", category), ("RETURN_TYPES", tup_get("output")), ("RETURN_NAMES", tup_get("output_name")), ("DESCRIPTION", defn.get("description")), ("OUTPUT_NODE", defn.get("output_node")), ): if val is not None: setattr(new_class, name, val) if not dry_run: nodes.NODE_CLASS_MAPPINGS[class_type] = new_class if display_name is not None: nodes.NODE_DISPLAY_NAME_MAPPINGS[class_type] = display_name stat_name = "overwrote" if existed else "added" logger.info("%s", f"{dry_run_label}{stat_name.capitalize()} node class: {class_type} ({display_name})") stats[stat_name] += 1 def add_server_prompt_nonce_handler(): logger.info("Adding promptserver remote_run nonce no duplicate prompt handler") import server nonce_ordered = collections.deque() nonce_set = set() NONCE_MAX = 1000 def prompt_handler(original_request: dict): nonce = original_request.get("nonce") if nonce is None: return original_request if nonce in nonce_set: logger.info("Duplicate nonce found, ignoring prompt: %r", nonce) return { } nonce_set.add(nonce) while len(nonce_set) > NONCE_MAX: oldest = nonce_ordered.popleft() nonce_set.remove(oldest) return original_request if os.environ.get("HOTRELOAD"): to_del = [i for i in server.PromptServer.instance.on_prompt_handlers if getattr(i, "rr_nonce_nodup_clearable", False)] for i in to_del: print("hotreload clearing old remote run nonce prompt handler", i) server.PromptServer.instance.on_prompt_handlers.remove(i) prompt_handler.rr_nonce_nodup_clearable = True server.PromptServer.instance.add_on_prompt_handler(prompt_handler) def add_server_prompt_input_disconnect_handler(): logger.info("Adding promptserver remote_run input disconnect handler") import server def prompt_handler(original_request: dict): if not isinstance(original_request.get("prompt"), dict): return original_request original_request = copy.deepcopy(original_request) prompt = original_request["prompt"] DEBUG_PROMPT_SAVE("rr_prompt_handler__pre_rewrite_prompt__", prompt) if IS_DEV and IS_DEV_PROMPTSAVE else None rcon, dscon = update_toggles_inplace(prompt, False, False, "input_disconnect_handler", raise_on_existing = True, input_node_is_local = False) # False so it gets disconnected if rcon or dscon: logger.info("remote_run input disconnect handler, reconnected %d, disconnected %d nodes", rcon, dscon) DEBUG_PROMPT_SAVE("rr_prompt_handler__rerewritten_prompt", prompt) if IS_DEV and IS_DEV_PROMPTSAVE else None return original_request if os.environ.get("HOTRELOAD"): to_del = [i for i in server.PromptServer.instance.on_prompt_handlers if getattr(i, "rr_input_discon_clearable", False)] for i in to_del: print("hotreload clearing old remote run input disconnect prompt handler", i) server.PromptServer.instance.on_prompt_handlers.remove(i) prompt_handler.rr_input_discon_clearable = True server.PromptServer.instance.on_prompt_handlers.insert(0, prompt_handler) else: server.PromptServer.instance.add_on_prompt_handler(prompt_handler) if os.environ.get("ADD_REMOTE_RUN_NONCE_HANDLER") == "1": logger.info("ADD_REMOTE_RUN_NONCE_HANDLER=1, adding remote_run no duplicate nonce handler!") try: add_server_prompt_nonce_handler() except Exception as e: logger.exception("Failed to add remote_run no duplicate nonce handler") else: logger.info("skipping remote run_no duplicate nonce handler, ADD_REMOTE_RUN_NONCE_HANDLER=1 to enable.") if os.environ.get("SKIP_REMOTE_RUN_INPUT_DISCONNECT_HANDLER") != "1": logger.info("SKIP_REMOTE_RUN_INPUT_DISCONNECT_HANDLER=1 not given, adding remote_run input disconnect handler!") try: add_server_prompt_input_disconnect_handler() except Exception as e: logger.exception("Failed to add remote_run input disconnect handler") else: logger.info("skipping remote_run input disconnect handler due to env var.") if os.environ.get("SKIP_REMOTE_RUN_ADD_DATA_ROUTE") != "1": logger.info("SKIP_REMOTE_RUN_ADD_DATA_ROUTE=1 not given, adding remote_run data route handler!") try: setup_remote_run_api_route(expected_set_up = False) except Exception as e: logger.exception("Failed to add remote_run data route handler") else: logger.info("skipping remote_run data route handler due to env var.") def node_mappings(classes): make_type_name = lambda cn: f"RAT_{cn.removesuffix('Node')}" class_map, name_map = { }, { } for i in classes: type_name = getattr(i, "TYPE_NAME", make_type_name(i.__name__)) class_map[type_name] = i name_map[type_name] = getattr(i, "DISPLAY_NAME", type_name) return class_map, name_map REMOTE_RUN_RUN_NODES = [ RemoteRunInputNode, RemoteRunInputOutputNode, RemoteRunJsonNode, RemoteRunJsonSyncNode, ] NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = node_mappings( ( RemoteRunStartNode, *REMOTE_RUN_RUN_NODES, RemoteRunTogglerNode, RemoteRunSerializerOutNode, RemoteRunDeserializerOutNode, RemoteRunSetNumOutputsNode, RemoteRunAddRemoteNodeDefinitionsNode, ) )