Switching into WS instead of REST to log versions

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

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:

modules/videoInference.py

@@ -304,6 +304,7 @@ class txt2vid:
                 "model": model,
                 "outputFormat": outputFormat,
                 "includeCost": True,
+                "deliveryMethod": "async",
             }
         ]
         

requirements.txt

@@ -2,4 +2,5 @@ python-dotenv
 brotli
 zstd
 librosa
-imageio
+imageio
+websocket-client