From 697792dcdd984d7368cbfe53f7b985f9951cce14 Mon Sep 17 00:00:00 2001 From: Sirshendu Ganguly Date: Fri, 19 Jun 2026 09:42:01 +0530 Subject: [PATCH 1/3] Switching into WS instead of REST to log versions --- modules/mediaUpload.py | 93 ++------- modules/utils/runwareUtils.py | 351 +++++++++++++++++++++++++++++++++- modules/videoInference.py | 1 + requirements.txt | 3 +- 4 files changed, 371 insertions(+), 77 deletions(-) diff --git a/modules/mediaUpload.py b/modules/mediaUpload.py index ffa6fef..f44f610 100644 --- a/modules/mediaUpload.py +++ b/modules/mediaUpload.py @@ -1,16 +1,15 @@ import base64 import os import time -import requests import torch import soundfile as sf import numpy as np from io import BytesIO from .utils.runwareUtils import ( RUNWARE_API_BASE_URL, - generalRequestWrapper, genRandUUID, - getRunwareApiHeaders, + inferenecRequest, + pollVideoResult, sanitize_for_logging, safe_json_dumps, sendMediaUUID, @@ -239,37 +238,14 @@ class RunwareMediaUpload: def _makeUploadRequest(self, uploadConfig): """Make upload request to Runware API""" - def recaller(): - print(f"[Debug] Making POST request to Runware API...") - return requests.post( - RUNWARE_API_BASE_URL, - headers=getRunwareApiHeaders(), - json=uploadConfig, - timeout=self.REQUEST_TIMEOUT, - allow_redirects=False, - stream=True, - ) - - response = generalRequestWrapper(recaller) - return self._parseJsonResponse(response) + print(f"[Debug] Making upload request to Runware API...") + return inferenecRequest(uploadConfig) def _parseJsonResponse(self, response): - """Parse JSON response from API""" - if response.status_code != 200: - errorText = response.text[:500] - print(f"[Debug] Request failed with status {response.status_code}") - print(f"[Debug] Response text: {errorText}") - raise Exception(f"Request failed with status {response.status_code}: {response.text[:200]}") - - if not response.text or not response.text.strip(): - raise Exception("Request failed: Empty response from server") - - try: - return response.json() - except ValueError as e: - print(f"[Debug] Failed to parse JSON response: {str(e)}") - print(f"[Debug] Response text: {response.text[:500]}") - raise Exception(f"Request failed: Invalid JSON response from server: {str(e)}") + """Validate JSON response from API""" + if not isinstance(response, dict): + raise Exception("Request failed: Invalid response from server") + return response def _validateUploadResponse(self, uploadResult): """Validate upload response""" @@ -281,26 +257,11 @@ class RunwareMediaUpload: def _pollForResult(self, taskUuid): """Poll for media upload result""" - pollConfig = [{ - "taskType": "getResponse", - "taskUUID": taskUuid, - }] - for attempt in range(self.MAX_POLL_ATTEMPTS): try: print(f"[Debug] Poll attempt {attempt + 1}/{self.MAX_POLL_ATTEMPTS}") - - def recaller(): - return requests.post( - RUNWARE_API_BASE_URL, - headers=getRunwareApiHeaders(), - json=pollConfig, - timeout=self.REQUEST_TIMEOUT, - allow_redirects=False, - stream=True, - ) - - pollResult = generalRequestWrapper(recaller) + + pollResult = pollVideoResult(taskUuid) pollData = self._parsePollResponse(pollResult) if pollData: @@ -324,32 +285,18 @@ class RunwareMediaUpload: def _parsePollResponse(self, pollResult): """Parse polling response""" - if pollResult.status_code != 200: - print(f"[Debug] Poll failed with status {pollResult.status_code}") - return None - - if not pollResult.text or not pollResult.text.strip(): + if pollResult is None: print(f"[Debug] Poll returned empty response") return None - - try: - pollResultJson = pollResult.json() - except ValueError as e: - print(f"[Debug] Failed to parse JSON in poll: {str(e)}") + + print(f"[Debug] Poll response: {safe_json_dumps(pollResult, indent=2) if isinstance(pollResult, (dict, list)) else pollResult}") + if "errors" in pollResult: + print(f"[Debug] Poll error: {safe_json_dumps(pollResult, indent=2) if isinstance(pollResult, (dict, list)) else pollResult}") return None - - print(f"[Debug] Poll response: {safe_json_dumps(pollResultJson, indent=2) if isinstance(pollResultJson, (dict, list)) else pollResultJson}") - - if "errors" in pollResultJson: - print(f"[Debug] Poll error: {safe_json_dumps(pollResultJson, indent=2) if isinstance(pollResultJson, (dict, list)) else pollResultJson}") - return None - - if "data" in pollResultJson and len(pollResultJson["data"]) > 0: - print(f"[Debug] Poll data: {sanitize_for_logging(pollResultJson['data'][0])}") - return pollResultJson["data"][0] - + + if "data" in pollResult and len(pollResult["data"]) > 0: + print(f"[Debug] Poll data: {sanitize_for_logging(pollResult['data'][0])}") + return pollResult["data"][0] + return None - - -# Export the class directly runwareMediaUpload = RunwareMediaUpload diff --git a/modules/utils/runwareUtils.py b/modules/utils/runwareUtils.py index d16c907..0046295 100644 --- a/modules/utils/runwareUtils.py +++ b/modules/utils/runwareUtils.py @@ -20,6 +20,8 @@ import io import threading import re +import websocket + from ..version import __version__ load_dotenv() @@ -177,7 +179,17 @@ def getCustomEndpoint(): custom_endpoint = os.getenv("RUNWARE_CUSTOM_ENDPOINT") if custom_endpoint and isinstance(custom_endpoint, str): return custom_endpoint - return "https://api.runware.ai/v1" + return "wss://ws-api.runware.ai/v1" + +def usesWebSocketTransport(endpoint=None): + endpoint = endpoint or getCustomEndpoint() + return endpoint.startswith("ws://") or endpoint.startswith("wss://") + +def getRunwareWsHeaders(): + return [ + f"X-SDK-Name: comfyui", + f"X-SDK-Version: {__version__}", + ] SESSION_TIMEOUT = getTimeout() RUNWARE_API_KEY = getAPIKey() @@ -214,6 +226,7 @@ def setAPIKey(apiKey: str): if envSetRes: RUNWARE_API_KEY = apiKey os.environ["RUNWARE_API_KEY"] = apiKey + _reset_ws_client() return True def setTimeout(timeout: int): @@ -268,7 +281,289 @@ def getOrdinal(num): ordinals = {1: "first", 2: "second", 3: "third", 4: "fourth"} return ordinals.get(num, f"{num}th") + +class RunwareWebSocketClient: + """Persistent WebSocket client for Runware API task requests.""" + + def __init__(self): + self._conn_lock = threading.RLock() + self._pending_lock = threading.Lock() + self._ws = None + self._connection_session_uuid = None + self._reader_thread = None + self._stop_reader = threading.Event() + self._pending = {} + + def close(self): + with self._conn_lock: + self._stop_reader.set() + if self._ws is not None: + try: + self._ws.close() + except Exception: + pass + self._ws = None + if self._reader_thread and self._reader_thread.is_alive(): + self._reader_thread.join(timeout=2) + self._reader_thread = None + with self._pending_lock: + for pending in self._pending.values(): + pending["response"] = {"errors": [{"message": "WebSocket connection closed"}]} + pending["event"].set() + self._pending.clear() + + def _connect(self): + url = getCustomEndpoint() + self._stop_reader.set() + if self._ws is not None: + try: + self._ws.close() + except Exception: + pass + self._ws = None + if self._reader_thread and self._reader_thread.is_alive(): + self._reader_thread.join(timeout=2) + + self._stop_reader.clear() + self._ws = websocket.create_connection( + url, + header=getRunwareWsHeaders(), + ) + self._ws.settimeout(None) + self._reader_thread = threading.Thread( + target=self._reader_loop, + name="RunwareWebSocketReader", + daemon=True, + ) + self._reader_thread.start() + self._authenticate() + + def _ensure_connected(self): + reader_alive = self._reader_thread is not None and self._reader_thread.is_alive() + if self._ws is None or not reader_alive: + self._connect() + return + try: + self._ws.ping() + except Exception: + self._connect() + + def _authenticate(self, api_key=None): + key = api_key or getAPIKey() + if not key: + raise Exception( + "Runware API key is not configured. Add your API key in the Runware API Manager node." + ) + auth_task = { + "taskType": "authentication", + "apiKey": key, + } + if self._connection_session_uuid: + auth_task["connectionSessionUUID"] = self._connection_session_uuid + response = self._send_and_wait([auth_task], timeout=10, is_auth=True) + if "errors" in response: + error_message = response["errors"][0].get("message", "Authentication failed") + raise Exception(error_message) + auth_data = response.get("data", [{}])[0] + session_uuid = auth_data.get("connectionSessionUUID") + if session_uuid: + self._connection_session_uuid = session_uuid + + def _reader_loop(self): + while not self._stop_reader.is_set() and self._ws is not None: + try: + raw = self._ws.recv() + except websocket.WebSocketTimeoutException: + continue + except websocket.WebSocketConnectionClosedException: + break + except Exception as e: + if not self._stop_reader.is_set(): + print(f"[Runware] WebSocket reader error: {e}") + break + + if not raw: + continue + + try: + message = json.loads(raw) + except json.JSONDecodeError: + print(f"[Runware] WebSocket received invalid JSON") + continue + + if self._is_ping_response(message): + continue + + self._dispatch_message(message) + + with self._pending_lock: + for pending in self._pending.values(): + if pending["response"] is None: + pending["response"] = {"errors": [{"message": "WebSocket connection closed"}]} + pending["event"].set() + + @staticmethod + def _is_ping_response(message): + data = message.get("data") or [] + return bool(data) and data[0].get("taskType") == "ping" and data[0].get("pong") + + def _dispatch_message(self, message): + with self._pending_lock: + for pending in list(self._pending.values()): + if pending["response"] is not None: + continue + + if pending["is_auth"]: + if message.get("data") and any( + item.get("taskType") == "authentication" for item in message["data"] + ): + pending["response"] = message + pending["event"].set() + continue + if message.get("errors") and any( + item.get("taskType") == "authentication" for item in message["errors"] + ): + pending["response"] = message + pending["event"].set() + continue + + for item in message.get("data", []): + task_uuid = item.get("taskUUID") + if task_uuid and task_uuid in pending["task_uuids"]: + pending["collected_data"].append(item) + pending["task_uuids"].discard(task_uuid) + + for item in message.get("errors", []): + task_uuid = item.get("taskUUID") + if task_uuid and task_uuid in pending["task_uuids"]: + pending["collected_errors"].append(item) + pending["task_uuids"].discard(task_uuid) + elif pending["task_uuids"] and len(pending["task_uuids"]) == 1: + pending["collected_errors"].append(item) + pending["task_uuids"].clear() + + if pending["task_uuids"]: + continue + + response = {} + if pending["collected_data"]: + response["data"] = pending["collected_data"] + if pending["collected_errors"]: + response["errors"] = pending["collected_errors"] + if not response: + response = message + pending["response"] = response + pending["event"].set() + + def _send_and_wait(self, gen_config, timeout, is_auth=False): + wait_event = threading.Event() + task_uuids = {task.get("taskUUID") for task in gen_config if task.get("taskUUID")} + pending = { + "event": wait_event, + "response": None, + "task_uuids": set(task_uuids), + "is_auth": is_auth, + "collected_data": [], + "collected_errors": [], + } + pending_id = id(pending) + + with self._pending_lock: + self._pending[pending_id] = pending + + try: + payload = json.dumps(gen_config) + with self._conn_lock: + self._ensure_connected() + self._ws.send(payload) + + if not wait_event.wait(timeout): + raise TimeoutError(f"WebSocket request timed out after {timeout} seconds") + + if pending["response"] is None: + raise Exception("WebSocket request failed: empty response") + return pending["response"] + finally: + with self._pending_lock: + self._pending.pop(pending_id, None) + + def request(self, gen_config, timeout): + def recaller(): + return self._send_and_wait(gen_config, timeout=timeout) + + return wsRequestWrapper(recaller) + + +_ws_client = None +_ws_client_lock = threading.Lock() + + +def _get_ws_client(): + global _ws_client + with _ws_client_lock: + if _ws_client is None: + _ws_client = RunwareWebSocketClient() + return _ws_client + + +def _reset_ws_client(): + global _ws_client + with _ws_client_lock: + if _ws_client is not None: + _ws_client.close() + _ws_client = None + + +def wsRequestWrapper(recaller): + for attempt in range(MAX_RETRIES + 1): + try: + return recaller() + except ( + websocket.WebSocketException, + ConnectionError, + OSError, + ): + if attempt == MAX_RETRIES: + raise + cooldown = RETRY_COOLDOWNS[attempt] + print( + f"[Runware] WebSocket request failed! Retrying in {cooldown} seconds..." + f" (Attempt {attempt + 1}/{MAX_RETRIES})" + ) + time.sleep(cooldown) + _reset_ws_client() + continue + return False + + +def checkAPIKeyWithWebSocket(apiKey): + url = getCustomEndpoint() + ws = None + try: + ws = websocket.create_connection( + url, + timeout=10, + header=getRunwareWsHeaders(), + ) + ws.send(json.dumps([{"taskType": "authentication", "apiKey": apiKey}])) + response = json.loads(ws.recv()) + if "errors" in response: + return str(response["errors"][0]["message"]) + return True + except Exception: + return False + finally: + if ws is not None: + try: + ws.close() + except Exception: + pass + + def checkAPIKey(apiKey): + if usesWebSocketTransport(): + return checkAPIKeyWithWebSocket(apiKey) + headers = getRunwareApiHeaders(include_auth=False) genConfig = [ { @@ -373,6 +668,43 @@ def inferenecRequest(genConfig): global RUNWARE_API_KEY, RUNWARE_API_BASE_URL, SESSION_TIMEOUT RUNWARE_API_KEY = os.getenv("RUNWARE_API_KEY") SESSION_TIMEOUT = int(os.getenv("RUNWARE_TIMEOUT")) + RUNWARE_API_BASE_URL = getCustomEndpoint() + + if usesWebSocketTransport(RUNWARE_API_BASE_URL): + try: + genResult = _get_ws_client().request(genConfig, timeout=SESSION_TIMEOUT) + except TimeoutError: + raise Exception( + f"Error: Request Timed Out After {SESSION_TIMEOUT} Seconds - Please Try Again!" + ) + except Exception as e: + if "invalid api key" in str(e).lower(): + PromptServer.instance.send_sync( + "runwareError", + { + "success": False, + "errorMessage": str(e), + "errorCode": 401, + }, + ) + raise InterruptProcessingException() + raise Exception(f"Error: {e}") + + if "errors" in genResult: + print(f"[DEBUG] API Error Response: {safe_json_dumps(genResult, indent=2) if isinstance(genResult, dict) else genResult}") + error_obj = genResult["errors"][0] + error_message = error_obj.get("message", "Unknown error") + + if "allowedValues" in error_obj: + allowed_values = error_obj["allowedValues"] + error_message += f"\n\nAllowed values:\n" + "\n".join(f" - {val}" for val in allowed_values) + + if "documentation" in error_obj: + error_message += f"\n\nDocumentation: {error_obj['documentation']}" + + raise Exception(error_message) + return genResult + headers = getRunwareApiHeaders() try: @@ -722,14 +1054,27 @@ def convertVideoB64List(videoDataObject, width=None, height=None): def pollVideoResult(taskUUID): """Poll async task result with taskType getResponse (video, audio, text inference, etc.).""" global RUNWARE_API_KEY, RUNWARE_API_BASE_URL, SESSION_TIMEOUT - + RUNWARE_API_BASE_URL = getCustomEndpoint() + pollConfig = [ { "taskType": "getResponse", "taskUUID": taskUUID, } ] - + + if usesWebSocketTransport(RUNWARE_API_BASE_URL): + try: + pollResult = _get_ws_client().request(pollConfig, timeout=30) + except Exception as e: + print(f"[Debugging] Poll request failed: {str(e)}") + return None + + if "errors" in pollResult: + print(f"[Debugging] Poll error: {safe_json_dumps(pollResult, indent=2) if isinstance(pollResult, (dict, list)) else pollResult}") + return pollResult + return pollResult + headers = getRunwareApiHeaders() try: diff --git a/modules/videoInference.py b/modules/videoInference.py index 747719e..a85dbbd 100644 --- a/modules/videoInference.py +++ b/modules/videoInference.py @@ -304,6 +304,7 @@ class txt2vid: "model": model, "outputFormat": outputFormat, "includeCost": True, + "deliveryMethod": "async", } ] diff --git a/requirements.txt b/requirements.txt index 78d667f..1d35d46 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ python-dotenv brotli zstd librosa -imageio \ No newline at end of file +imageio +websocket-client \ No newline at end of file From a06184b67364ac61f37fed7a1540499a15c43733 Mon Sep 17 00:00:00 2001 From: Sirshendu Ganguly Date: Fri, 19 Jun 2026 14:23:50 +0530 Subject: [PATCH 2/3] Removed dead code --- modules/mediaUpload.py | 6 ------ modules/utils/runwareUtils.py | 13 ++++++++++--- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/modules/mediaUpload.py b/modules/mediaUpload.py index f44f610..9f7ee46 100644 --- a/modules/mediaUpload.py +++ b/modules/mediaUpload.py @@ -241,12 +241,6 @@ class RunwareMediaUpload: print(f"[Debug] Making upload request to Runware API...") return inferenecRequest(uploadConfig) - def _parseJsonResponse(self, response): - """Validate JSON response from API""" - if not isinstance(response, dict): - raise Exception("Request failed: Invalid response from server") - return response - def _validateUploadResponse(self, uploadResult): """Validate upload response""" print(f"[Debug] Upload response: {safe_json_dumps(uploadResult, indent=2) if isinstance(uploadResult, (dict, list)) else uploadResult}") diff --git a/modules/utils/runwareUtils.py b/modules/utils/runwareUtils.py index 0046295..c318e7d 100644 --- a/modules/utils/runwareUtils.py +++ b/modules/utils/runwareUtils.py @@ -327,6 +327,7 @@ class RunwareWebSocketClient: self._stop_reader.clear() self._ws = websocket.create_connection( url, + timeout=10, header=getRunwareWsHeaders(), ) self._ws.settimeout(None) @@ -438,7 +439,12 @@ class RunwareWebSocketClient: if task_uuid and task_uuid in pending["task_uuids"]: pending["collected_errors"].append(item) pending["task_uuids"].discard(task_uuid) - elif pending["task_uuids"] and len(pending["task_uuids"]) == 1: + elif ( + not task_uuid + and len(self._pending) == 1 + and pending["task_uuids"] + and len(pending["task_uuids"]) == 1 + ): pending["collected_errors"].append(item) pending["task_uuids"].clear() @@ -547,8 +553,9 @@ def checkAPIKeyWithWebSocket(apiKey): ) ws.send(json.dumps([{"taskType": "authentication", "apiKey": apiKey}])) response = json.loads(ws.recv()) - if "errors" in response: - return str(response["errors"][0]["message"]) + errors = response.get("errors") or [] + if errors: + return str(errors[0].get("message", "Authentication failed")) return True except Exception: return False From 57e667fabb2c6554d75b8ba163d479ee0f59314e Mon Sep 17 00:00:00 2001 From: Sirshendu Ganguly Date: Fri, 19 Jun 2026 14:44:24 +0530 Subject: [PATCH 3/3] Quick fixes --- modules/utils/runwareUtils.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/modules/utils/runwareUtils.py b/modules/utils/runwareUtils.py index c318e7d..18d0914 100644 --- a/modules/utils/runwareUtils.py +++ b/modules/utils/runwareUtils.py @@ -191,6 +191,9 @@ def getRunwareWsHeaders(): f"X-SDK-Version: {__version__}", ] +WS_CONNECT_TIMEOUT = 10 +WS_READ_TIMEOUT = 1 + SESSION_TIMEOUT = getTimeout() RUNWARE_API_KEY = getAPIKey() OUTPUT_FORMAT = getOutputFormat() @@ -199,6 +202,11 @@ ENABLE_IMAGES_CACHING = getEnableImagesCaching() MIN_IMAGE_CACHE_SIZE = getMinImageCacheSize() RUNWARE_API_BASE_URL = getCustomEndpoint() +def refreshRunwareEndpoint(): + global RUNWARE_API_BASE_URL + RUNWARE_API_BASE_URL = getCustomEndpoint() + return RUNWARE_API_BASE_URL + def setEnvKey(keyName, keyValue): comfyNodeRoot = Path(__file__).parent.parent.parent envFilePath = comfyNodeRoot / ".env" @@ -327,10 +335,10 @@ class RunwareWebSocketClient: self._stop_reader.clear() self._ws = websocket.create_connection( url, - timeout=10, + timeout=WS_CONNECT_TIMEOUT, header=getRunwareWsHeaders(), ) - self._ws.settimeout(None) + self._ws.settimeout(WS_READ_TIMEOUT) self._reader_thread = threading.Thread( target=self._reader_loop, name="RunwareWebSocketReader", @@ -548,7 +556,7 @@ def checkAPIKeyWithWebSocket(apiKey): try: ws = websocket.create_connection( url, - timeout=10, + timeout=WS_CONNECT_TIMEOUT, header=getRunwareWsHeaders(), ) ws.send(json.dumps([{"taskType": "authentication", "apiKey": apiKey}])) @@ -568,7 +576,8 @@ def checkAPIKeyWithWebSocket(apiKey): def checkAPIKey(apiKey): - if usesWebSocketTransport(): + endpoint = refreshRunwareEndpoint() + if usesWebSocketTransport(endpoint): return checkAPIKeyWithWebSocket(apiKey) headers = getRunwareApiHeaders(include_auth=False) @@ -672,12 +681,12 @@ def sendVideoOutputs(draftId, videoId, nodeID): ) def inferenecRequest(genConfig): - global RUNWARE_API_KEY, RUNWARE_API_BASE_URL, SESSION_TIMEOUT + global RUNWARE_API_KEY, SESSION_TIMEOUT RUNWARE_API_KEY = os.getenv("RUNWARE_API_KEY") SESSION_TIMEOUT = int(os.getenv("RUNWARE_TIMEOUT")) - RUNWARE_API_BASE_URL = getCustomEndpoint() + endpoint = refreshRunwareEndpoint() - if usesWebSocketTransport(RUNWARE_API_BASE_URL): + if usesWebSocketTransport(endpoint): try: genResult = _get_ws_client().request(genConfig, timeout=SESSION_TIMEOUT) except TimeoutError: @@ -1060,8 +1069,7 @@ def convertVideoB64List(videoDataObject, width=None, height=None): def pollVideoResult(taskUUID): """Poll async task result with taskType getResponse (video, audio, text inference, etc.).""" - global RUNWARE_API_KEY, RUNWARE_API_BASE_URL, SESSION_TIMEOUT - RUNWARE_API_BASE_URL = getCustomEndpoint() + endpoint = refreshRunwareEndpoint() pollConfig = [ { @@ -1070,7 +1078,7 @@ def pollVideoResult(taskUUID): } ] - if usesWebSocketTransport(RUNWARE_API_BASE_URL): + if usesWebSocketTransport(endpoint): try: pollResult = _get_ws_client().request(pollConfig, timeout=30) except Exception as e: