Claude/add lora loading node w xg9u

.github/workflows/build-custom.yml

@@ -0,0 +1,48 @@
+name: Build Custom Docker Image
+
+on:
+  # Trigger on push to main branch when Dockerfile.custom changes
+  push:
+    branches: [main]
+    paths:
+      - 'Dockerfile.custom'
+  # Allow manual trigger
+  workflow_dispatch:
+
+jobs:
+  build-and-push:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+          sudo docker image prune --all --force
+          df -h
+
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          file: Dockerfile.custom
+          platforms: linux/amd64
+          push: true
+          tags: |
+            ${{ secrets.DOCKERHUB_USERNAME }}/runpod-comfy-worker:latest
+            ${{ secrets.DOCKERHUB_USERNAME }}/runpod-comfy-worker:${{ github.sha }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max

Dockerfile.custom

@@ -0,0 +1,12 @@
+# Custom ComfyUI worker with video/animation nodes
+# No models included - uses network volume
+FROM runpod/worker-comfyui:5.5.0-base
+
+# Install custom nodes for video/frame processing workflows
+RUN comfy-node-install \
+    comfyui-kjnodes \
+    comfyui-videohelpersuite \
+    comfyui-frame-interpolation \
+    comfyui_essentials \
+    comfyui-gimm-vfi \
+    loadloramodelonlywithurl

docs/configuration.md

@@ -21,6 +21,7 @@ This document outlines the environment variables available for configuring the `
 
 | Environment Variable           | Description                                                                                                            | Default |
 | ------------------------------ | ---------------------------------------------------------------------------------------------------------------------- | ------- |
+| `WEBSOCKET_MESSAGE_TIMEOUT`    | Maximum seconds to wait for a single WebSocket message before checking ComfyUI health. Heavy workloads (e.g. video frame interpolation) can produce long silences on the socket. Set to `0` to disable (recv blocks indefinitely). | `600`   |
 | `WEBSOCKET_RECONNECT_ATTEMPTS` | Number of websocket reconnection attempts when connection drops during job execution.                                  | `5`     |
 | `WEBSOCKET_RECONNECT_DELAY_S`  | Delay in seconds between websocket reconnection attempts.                                                              | `3`     |
 | `WEBSOCKET_TRACE`              | Enable low-level websocket frame tracing for protocol debugging. Set to `true` only when diagnosing connection issues. | `false` |

handler.py

@@ -38,6 +38,11 @@
 # If the respective env-vars are not supplied we fall back to sensible defaults ("5" and "3").
 WEBSOCKET_RECONNECT_ATTEMPTS = int(os.environ.get("WEBSOCKET_RECONNECT_ATTEMPTS", 5))
 WEBSOCKET_RECONNECT_DELAY_S = int(os.environ.get("WEBSOCKET_RECONNECT_DELAY_S", 3))
+# Maximum seconds to wait for a single WebSocket message before checking
+# whether ComfyUI is still alive. Heavy workloads (e.g. video frame
+# interpolation) can produce long silences on the socket. Default: 600 s
+# (10 minutes). Set to 0 to disable (recv blocks indefinitely).
+WEBSOCKET_MESSAGE_TIMEOUT = int(os.environ.get("WEBSOCKET_MESSAGE_TIMEOUT", 600))
 
 # Extra verbose websocket trace logs (set WEBSOCKET_TRACE=true to enable)
 if os.environ.get("WEBSOCKET_TRACE", "false").lower() == "true":
@@ -68,6 +73,31 @@ def _comfy_server_status():
         return {"reachable": False, "error": str(exc)}
 
 
+def _is_job_in_queue(prompt_id):
+    """Check if a prompt is still running or pending in ComfyUI's queue.
+
+    Polls the ``/queue`` endpoint and looks for *prompt_id* in both
+    ``queue_running`` and ``queue_pending``.  Returns ``True`` if found,
+    ``False`` if not found, and ``None`` if the check itself failed (e.g.
+    network error).
+    """
+    try:
+        resp = requests.get(f"http://{COMFY_HOST}/queue", timeout=5)
+        resp.raise_for_status()
+        queue_data = resp.json()
+        for item in queue_data.get("queue_running", []):
+            # Each item is a list; the prompt_id is at index 1
+            if len(item) >= 2 and item[1] == prompt_id:
+                return True
+        for item in queue_data.get("queue_pending", []):
+            if len(item) >= 2 and item[1] == prompt_id:
+                return True
+        return False
+    except Exception as exc:
+        logger.warning("Error polling /queue: %s", exc)
+        return None
+
+
 def _attempt_websocket_reconnect(ws_url, max_attempts, delay_s, initial_error):
     """
     Attempts to reconnect to the WebSocket server after a disconnect.
@@ -564,7 +594,9 @@ def handler(job):
         print(f"worker-comfyui - Connecting to websocket: {ws_url}")
         ws = websocket.WebSocket()
         ws.connect(ws_url, timeout=10)
-        print(f"worker-comfyui - Websocket connected")
+        if WEBSOCKET_MESSAGE_TIMEOUT > 0:
+            ws.settimeout(WEBSOCKET_MESSAGE_TIMEOUT)
+        print(f"worker-comfyui - Websocket connected (message timeout: {WEBSOCKET_MESSAGE_TIMEOUT}s)")
 
         # Queue the workflow
         try:
@@ -627,8 +659,37 @@ def handler(job):
                 else:
                     continue
             except websocket.WebSocketTimeoutException:
-                print(f"worker-comfyui - Websocket receive timed out. Still waiting...")
-                continue
+                # No message received within WEBSOCKET_MESSAGE_TIMEOUT seconds.
+                # 1) Check whether ComfyUI HTTP is alive at all.
+                srv_status = _comfy_server_status()
+                if not srv_status["reachable"]:
+                    raise ValueError(
+                        f"ComfyUI became unreachable while waiting for WebSocket messages "
+                        f"(no message for {WEBSOCKET_MESSAGE_TIMEOUT}s). "
+                        f"Server status: {srv_status}"
+                    )
+
+                # 2) Check whether our job is still in the queue (running/pending).
+                in_queue = _is_job_in_queue(prompt_id)
+                if in_queue:
+                    print(
+                        f"worker-comfyui - No WebSocket message for {WEBSOCKET_MESSAGE_TIMEOUT}s, "
+                        f"but job {prompt_id} is still in ComfyUI queue. Continuing to wait..."
+                    )
+                    continue
+                elif in_queue is None:
+                    # Queue check failed (network blip) — give benefit of the doubt.
+                    print(
+                        f"worker-comfyui - No WebSocket message for {WEBSOCKET_MESSAGE_TIMEOUT}s. "
+                        f"Queue check failed; ComfyUI HTTP is reachable so continuing to wait..."
+                    )
+                    continue
+                else:
+                    # Job is NOT in queue and NOT completed via WS — truly stuck.
+                    raise ValueError(
+                        f"WebSocket message timeout: no message for {WEBSOCKET_MESSAGE_TIMEOUT}s "
+                        f"and job {prompt_id} is no longer in ComfyUI queue."
+                    )
             except websocket.WebSocketConnectionClosedException as closed_err:
                 try:
                     # Attempt to reconnect
@@ -639,6 +700,8 @@ def handler(job):
                         closed_err,
                     )
 
+                    if WEBSOCKET_MESSAGE_TIMEOUT > 0:
+                        ws.settimeout(WEBSOCKET_MESSAGE_TIMEOUT)
                     print(
                         "worker-comfyui - Resuming message listening after successful reconnect."
                     )