merged fork

distributed.py

@@ -10,6 +10,7 @@ import aiohttp
 from aiohttp import web
 import io
 import server
+import execution
 import comfy.model_management
 import subprocess
 import platform
@@ -96,6 +97,102 @@ async def queue_status_endpoint(request):
     except Exception as e:
         return await handle_api_error(request, e, 500)
 
+
+prompt_server = server.PromptServer.instance
+
+
+async def _queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
+    """Validate and queue a prompt via ComfyUI's prompt queue."""
+    payload = {"prompt": prompt_obj}
+    payload = prompt_server.trigger_on_prompt(payload)
+    prompt = payload["prompt"]
+
+    prompt_id = str(uuid.uuid4())
+    valid = await execution.validate_prompt(prompt_id, prompt, None)
+    if not valid[0]:
+        raise RuntimeError(f"Invalid prompt: {valid[1]}")
+
+    extra_data = {}
+    if workflow_meta:
+        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
+    if client_id:
+        extra_data["client_id"] = client_id
+
+    sensitive = {}
+    for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
+        if key in extra_data:
+            sensitive[key] = extra_data.pop(key)
+
+    number = getattr(prompt_server, "number", 0)
+    prompt_server.number = number + 1
+    prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
+    prompt_server.prompt_queue.put(prompt_queue_item)
+    return prompt_id
+
+
+@server.PromptServer.instance.routes.get("/distributed/worker_ws")
+async def worker_ws_endpoint(request):
+    """WebSocket endpoint for worker prompt dispatch."""
+    ws = web.WebSocketResponse(heartbeat=30)
+    await ws.prepare(request)
+
+    async for msg in ws:
+        if msg.type == aiohttp.WSMsgType.TEXT:
+            try:
+                data = json.loads(msg.data or "{}")
+            except json.JSONDecodeError:
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": None,
+                    "ok": False,
+                    "error": "Invalid JSON payload.",
+                })
+                continue
+
+            if data.get("type") != "dispatch_prompt":
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": False,
+                    "error": "Unsupported websocket message type.",
+                })
+                continue
+
+            prompt = data.get("prompt")
+            if not isinstance(prompt, dict):
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": False,
+                    "error": "Field 'prompt' must be an object.",
+                })
+                continue
+
+            try:
+                prompt_id = await _queue_prompt_payload(
+                    prompt,
+                    workflow_meta=data.get("workflow"),
+                    client_id=data.get("client_id"),
+                )
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": True,
+                    "prompt_id": prompt_id,
+                })
+            except Exception as exc:
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": False,
+                    "error": str(exc),
+                })
+        elif msg.type == aiohttp.WSMsgType.ERROR:
+            log(f"[Distributed] Worker websocket error: {ws.exception()}")
+
+    return ws
+
+
 @server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
 async def clear_launching_state(request):
     """Clear the launching flag when worker is confirmed running."""
@@ -2486,17 +2583,165 @@ class DistributedEmptyImage:
         return (tensor,)
 
 
-NODE_CLASS_MAPPINGS = { 
+# --- Distributed Queue Node ---
+class DistributedQueueNode:
+    """Routes entire workflows to the least-busy worker for job scheduling/load balancing."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {},
+            "hidden": {
+                "prompt": "PROMPT",
+                "extra_pnginfo": "EXTRA_PNGINFO",
+                "client_id": "CLIENT_ID",
+                "skip_dispatch": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "STRING")
+    RETURN_NAMES = ("worker_id", "prompt_id")
+    FUNCTION = "queue"
+    OUTPUT_NODE = True
+    CATEGORY = "utils"
+
+    def queue(self, prompt=None, extra_pnginfo=None, client_id=None, skip_dispatch=False):
+        if skip_dispatch:
+            debug_log("DistributedQueue: skip_dispatch enabled; not dispatching.")
+            return ("", "")
+        if not prompt:
+            log("DistributedQueue: No prompt supplied; skipping dispatch.")
+            return ("", "")
+
+        return run_async_in_server_loop(
+            self._queue_on_worker(prompt, extra_pnginfo, client_id)
+        )
+
+    async def _queue_on_worker(self, prompt_obj, workflow_meta, client_id):
+        config = load_config()
+        enabled_workers = [
+            worker for worker in config.get("workers", [])
+            if worker.get("enabled", False)
+        ]
+
+        if not enabled_workers:
+            log("DistributedQueue: No enabled workers found.")
+            return ("", "")
+
+        statuses = await self._fetch_worker_statuses(enabled_workers)
+        if not statuses:
+            log("DistributedQueue: No reachable workers found.")
+            return ("", "")
+
+        selected = self._select_worker(statuses)
+        worker = selected["worker"]
+        queue_remaining = selected["queue_remaining"]
+
+        prompt_copy = json.loads(json.dumps(prompt_obj))
+        self._mark_skip_dispatch(prompt_copy)
+
+        payload = {"prompt": prompt_copy}
+        extra_data = {}
+        if workflow_meta:
+            extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
+        if client_id:
+            extra_data["client_id"] = client_id
+        if extra_data:
+            payload["extra_data"] = extra_data
+
+        url = self._build_worker_url(worker, "/prompt")
+        session = await get_client_session()
+        async with session.post(
+            url,
+            json=payload,
+            timeout=aiohttp.ClientTimeout(total=60),
+        ) as resp:
+            resp.raise_for_status()
+            response_payload = await resp.json()
+
+        prompt_id = str(response_payload.get("prompt_id", ""))
+        worker_id = str(worker.get("id", ""))
+        log(
+            f"DistributedQueue: queued prompt {prompt_id} on worker {worker_id} (queue_remaining={queue_remaining})."
+        )
+        return (worker_id, prompt_id)
+
+    async def _fetch_worker_statuses(self, workers):
+        session = await get_client_session()
+        statuses = []
+        for worker in workers:
+            url = self._build_worker_url(worker, "/prompt")
+            try:
+                async with session.get(
+                    url,
+                    timeout=aiohttp.ClientTimeout(total=3),
+                ) as resp:
+                    if resp.status != 200:
+                        continue
+                    data = await resp.json()
+                    queue_remaining = int(data.get("exec_info", {}).get("queue_remaining", 0))
+                    statuses.append({
+                        "worker": worker,
+                        "queue_remaining": queue_remaining,
+                    })
+            except Exception as e:
+                debug_log(f"DistributedQueue: Worker status check failed ({url}): {e}")
+
+        return statuses
+
+    def _select_worker(self, statuses):
+        idle_workers = [status for status in statuses if status["queue_remaining"] == 0]
+        if idle_workers:
+            return self._select_round_robin(idle_workers)
+        return min(statuses, key=lambda status: status["queue_remaining"])
+
+    def _select_round_robin(self, statuses):
+        prompt_server = server.PromptServer.instance
+        if not hasattr(prompt_server, "distributed_queue_rr_index"):
+            prompt_server.distributed_queue_rr_index = 0
+
+        index = prompt_server.distributed_queue_rr_index % len(statuses)
+        prompt_server.distributed_queue_rr_index += 1
+        return statuses[index]
+
+    def _mark_skip_dispatch(self, prompt_obj):
+        for node in prompt_obj.values():
+            if isinstance(node, dict) and node.get("class_type") == "DistributedQueue":
+                inputs = node.setdefault("inputs", {})
+                inputs["skip_dispatch"] = True
+
+    def _build_worker_url(self, worker, endpoint=""):
+        host = (worker.get("host") or "").strip()
+        port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
+
+        if not host:
+            host = getattr(server.PromptServer.instance, "address", "127.0.0.1") or "127.0.0.1"
+
+        if host.startswith(("http://", "https://")):
+            base = host.rstrip("/")
+        else:
+            is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
+            scheme = "https" if is_cloud else "http"
+            default_port = 443 if scheme == "https" else 80
+            port_part = "" if port == default_port else f":{port}"
+            base = f"{scheme}://{host}{port_part}"
+
+        return f"{base}{endpoint}"
+
+
+NODE_CLASS_MAPPINGS = {
+    "DistributedQueue": DistributedQueueNode,
     "DistributedCollector": DistributedCollectorNode,
     "DistributedSeed": DistributedSeed,
     "DistributedModelName": DistributedModelName,
     "ImageBatchDivider": ImageBatchDivider,
     "DistributedEmptyImage": DistributedEmptyImage,
 }
-NODE_DISPLAY_NAME_MAPPINGS = { 
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DistributedQueue": "Distributed Queue",
     "DistributedCollector": "Distributed Collector",
     "DistributedSeed": "Distributed Seed",
-    "DistributedModelName": "DistributedModelName",
+    "DistributedModelName": "Distributed Model Name",
     "ImageBatchDivider": "Image Batch Divider",
     "DistributedEmptyImage": "Distributed Empty Image",
 }

distributed_queue_api.py

@@ -22,6 +22,8 @@ def ensure_distributed_state():
         prompt_server.distributed_pending_jobs = {}
     if not hasattr(prompt_server, "distributed_jobs_lock"):
         prompt_server.distributed_jobs_lock = asyncio.Lock()
+    if not hasattr(prompt_server, "distributed_worker_ws"):
+        prompt_server.distributed_worker_ws = {}
 
 
 async def _get_client_session():
@@ -31,9 +33,53 @@ async def _get_client_session():
     return prompt_server._distributed_session
 
 
-def _deepcopy_prompt(prompt_obj):
-    """Return a deep copy of the workflow/prompt dictionary."""
-    return json.loads(json.dumps(prompt_obj))
+class PromptIndex:
+    """Cache prompt metadata for faster worker/master prompt preparation."""
+
+    def __init__(self, prompt_obj):
+        self._prompt_json = json.dumps(prompt_obj)
+        self.nodes_by_class = {}
+        self.class_by_node = {}
+        self.inputs_by_node = {}
+        for node_id, node in _iter_prompt_nodes(prompt_obj):
+            class_type = node.get("class_type")
+            node_id_str = str(node_id)
+            if class_type:
+                self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
+            self.class_by_node[node_id_str] = class_type
+            self.inputs_by_node[node_id_str] = node.get("inputs", {})
+        self._upstream_cache = {}
+
+    def copy_prompt(self):
+        return json.loads(self._prompt_json)
+
+    def nodes_for_class(self, class_name):
+        return self.nodes_by_class.get(class_name, [])
+
+    def has_upstream(self, start_node_id, target_class):
+        cache_key = (str(start_node_id), target_class)
+        if cache_key in self._upstream_cache:
+            return self._upstream_cache[cache_key]
+
+        visited = set()
+        stack = [str(start_node_id)]
+        while stack:
+            node_id = stack.pop()
+            if node_id in visited:
+                continue
+            visited.add(node_id)
+            inputs = self.inputs_by_node.get(node_id, {})
+            for value in inputs.values():
+                if isinstance(value, list) and len(value) == 2:
+                    upstream_id = str(value[0])
+                    if self.class_by_node.get(upstream_id) == target_class:
+                        self._upstream_cache[cache_key] = True
+                        return True
+                    if upstream_id in self.inputs_by_node:
+                        stack.append(upstream_id)
+
+        self._upstream_cache[cache_key] = False
+        return False
 
 
 def _iter_prompt_nodes(prompt_obj):
@@ -50,32 +96,6 @@ def _find_nodes_by_class(prompt_obj, class_name):
     return nodes
 
 
-def _has_upstream_node(prompt_obj, start_node_id, target_class):
-    """Depth-first search to find if start node depends on target_class."""
-    visited = set()
-    stack = [start_node_id]
-
-    while stack:
-        node_id = stack.pop()
-        if node_id in visited:
-            continue
-        visited.add(node_id)
-        node = prompt_obj.get(node_id)
-        if not node:
-            continue
-        inputs = node.get("inputs", {})
-        for value in inputs.values():
-            if isinstance(value, list) and len(value) == 2:
-                upstream_id = str(value[0])
-                upstream_node = prompt_obj.get(upstream_id)
-                if not upstream_node:
-                    continue
-                if upstream_node.get("class_type") == target_class:
-                    return True
-                stack.append(upstream_id)
-    return False
-
-
 def _find_downstream_nodes(prompt_obj, start_ids):
     """Return all nodes reachable downstream from the provided IDs."""
     adjacency = {}
@@ -176,11 +196,11 @@ async def _ensure_distributed_queue(job_id):
             prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
 
 
-def _generate_job_id_map(prompt_obj, prefix):
+def _generate_job_id_map(prompt_index, prefix):
     """Create stable per-node job IDs for distributed nodes."""
     job_map = {}
-    distributed_nodes = _find_nodes_by_class(prompt_obj, "DistributedCollector") + _find_nodes_by_class(
-        prompt_obj, "UltimateSDUpscaleDistributed"
+    distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
+        "UltimateSDUpscaleDistributed"
     )
     for node_id in distributed_nodes:
         job_map[node_id] = f"{prefix}_{node_id}"
@@ -201,12 +221,19 @@ def _resolve_enabled_workers(config, requested_ids=None):
         elif not worker.get("enabled", False):
             continue
 
+        raw_port = worker.get("port", worker.get("listen_port", 8188))
+        try:
+            port = int(raw_port or 8188)
+        except (TypeError, ValueError):
+            log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.")
+            port = 8188
+
         workers.append(
             {
                 "id": worker_id,
                 "name": worker.get("name", worker_id),
                 "host": worker.get("host"),
-                "port": int(worker.get("port", worker.get("listen_port", 8188)) or 8188),
+                "port": port,
                 "type": worker.get("type", "local"),
             }
         )
@@ -287,7 +314,40 @@ async def _worker_is_active(worker):
         return False
 
 
-async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta):
+async def _worker_ws_is_active(worker):
+    """Ping worker's websocket endpoint to confirm it's reachable."""
+    session = await _get_client_session()
+    url = _build_worker_url(worker, "/distributed/worker_ws")
+    try:
+        ws = await session.ws_connect(url, heartbeat=20, timeout=3)
+        await ws.close()
+        return True
+    except Exception:
+        return False
+
+
+async def _get_worker_ws_connection(worker):
+    """Return a cached websocket connection to a worker (creating if needed)."""
+    ensure_distributed_state()
+    worker_id = worker.get("id")
+    existing = prompt_server.distributed_worker_ws.get(worker_id)
+    if existing and not existing.closed:
+        return existing
+
+    session = await _get_client_session()
+    url = _build_worker_url(worker, "/distributed/worker_ws")
+    ws = await session.ws_connect(url, heartbeat=20, timeout=10)
+    prompt_server.distributed_worker_ws[worker_id] = ws
+    return ws
+
+
+async def _close_worker_ws(worker_id):
+    ws = prompt_server.distributed_worker_ws.pop(worker_id, None)
+    if ws and not ws.closed:
+        await ws.close()
+
+
+async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta, client_id=None, use_websocket=False):
     """Send the prepared prompt to a worker ComfyUI instance."""
     url = _build_worker_url(worker, "/prompt")
     payload = {"prompt": prompt_obj}
@@ -297,6 +357,35 @@ async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta):
     if extra_data:
         payload["extra_data"] = extra_data
 
+    if use_websocket:
+        request_id = uuid.uuid4().hex
+        ws_payload = {
+            "type": "dispatch_prompt",
+            "request_id": request_id,
+            "prompt": prompt_obj,
+            "workflow": workflow_meta,
+            "client_id": client_id,
+        }
+        try:
+            ws = await _get_worker_ws_connection(worker)
+            await ws.send_json(ws_payload)
+            response = await asyncio.wait_for(ws.receive(), timeout=60)
+            if response.type == aiohttp.WSMsgType.TEXT:
+                data = json.loads(response.data or "{}")
+                if (
+                    data.get("type") == "dispatch_ack"
+                    and data.get("request_id") == request_id
+                    and data.get("ok")
+                ):
+                    return
+                error_msg = data.get("error") or "Worker rejected websocket dispatch."
+                raise RuntimeError(error_msg)
+            raise RuntimeError("Unexpected websocket response from worker.")
+        except Exception as exc:
+            worker_id = worker.get("id")
+            await _close_worker_ws(worker_id)
+            log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
+
     session = await _get_client_session()
     async with session.post(
         url,
@@ -336,21 +425,21 @@ async def _queue_master_prompt(prompt_obj, workflow_meta, client_id):
 
 
 def _apply_participant_overrides(
-    base_prompt,
+    prompt_copy,
     participant_id,
     enabled_worker_ids,
     job_id_map,
     master_url,
     delegate_master,
+    prompt_index,
 ):
     """Return a prompt copy with hidden inputs configured for master/worker."""
-    prompt_copy = _deepcopy_prompt(base_prompt)
     is_master = participant_id == "master"
     worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
     enabled_json = json.dumps(enabled_worker_ids)
 
     # Distributed seed nodes
-    for node_id in _find_nodes_by_class(prompt_copy, "DistributedSeed"):
+    for node_id in prompt_index.nodes_for_class("DistributedSeed"):
         node = prompt_copy.get(node_id, {})
         inputs = node.setdefault("inputs", {})
         inputs["is_worker"] = not is_master
@@ -361,11 +450,11 @@ def _apply_participant_overrides(
             inputs["worker_id"] = ""
 
     # Distributed collectors
-    for node_id in _find_nodes_by_class(prompt_copy, "DistributedCollector"):
+    for node_id in prompt_index.nodes_for_class("DistributedCollector"):
         node = prompt_copy.get(node_id, {})
         inputs = node.setdefault("inputs", {})
 
-        if _has_upstream_node(prompt_copy, node_id, "UltimateSDUpscaleDistributed"):
+        if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
             inputs["pass_through"] = True
             continue
 
@@ -385,7 +474,7 @@ def _apply_participant_overrides(
             inputs["delegate_only"] = False
 
     # Distributed upscaler nodes
-    for node_id in _find_nodes_by_class(prompt_copy, "UltimateSDUpscaleDistributed"):
+    for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
         node = prompt_copy.get(node_id, {})
         inputs = node.setdefault("inputs", {})
         unique_id = job_id_map.get(node_id, node_id)
@@ -418,8 +507,10 @@ async def orchestrate_distributed_execution(
     ensure_distributed_state()
 
     config = load_config()
+    use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
     requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
     workers = _resolve_enabled_workers(config, requested_ids)
+    prompt_index = PromptIndex(prompt_obj)
 
     # Respect master delegate-only configuration
     if delegate_master is None:
@@ -432,7 +523,11 @@ async def orchestrate_distributed_execution(
     # Filter to active workers
     active_workers = []
     for worker in workers:
-        if await _worker_is_active(worker):
+        if use_websocket:
+            is_active = await _worker_ws_is_active(worker)
+        else:
+            is_active = await _worker_is_active(worker)
+        if is_active:
             active_workers.append(worker)
         else:
             log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
@@ -444,7 +539,7 @@ async def orchestrate_distributed_execution(
     enabled_ids = [worker["id"] for worker in active_workers]
 
     discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
-    job_id_map = _generate_job_id_map(prompt_obj, discovery_prefix)
+    job_id_map = _generate_job_id_map(prompt_index, discovery_prefix)
 
     if not job_id_map:
         prompt_id = await _queue_master_prompt(prompt_obj, workflow_meta, client_id)
@@ -454,13 +549,15 @@ async def orchestrate_distributed_execution(
         await _ensure_distributed_queue(job_id)
 
     master_url = _resolve_master_url()
+    master_prompt = prompt_index.copy_prompt()
     master_prompt = _apply_participant_overrides(
-        prompt_obj,
+        master_prompt,
         "master",
         enabled_ids,
         job_id_map,
         master_url,
         delegate_master,
+        prompt_index,
     )
 
     if delegate_master:
@@ -477,21 +574,31 @@ async def orchestrate_distributed_execution(
         else:
             master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids)
 
+    if active_workers:
+        debug_log(
+            "Active distributed workers: "
+            + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers)
+        )
     worker_payloads = []
     for worker in active_workers:
+        worker_prompt = prompt_index.copy_prompt()
         worker_prompt = _apply_participant_overrides(
-            prompt_obj,
+            worker_prompt,
             worker["id"],
             enabled_ids,
             job_id_map,
             master_url,
             delegate_master,
+            prompt_index,
         )
         worker_payloads.append((worker, worker_prompt))
 
     if worker_payloads:
         await asyncio.gather(
-            *[_dispatch_worker_prompt(worker, wprompt, workflow_meta) for worker, wprompt in worker_payloads]
+            *[
+                _dispatch_worker_prompt(worker, wprompt, workflow_meta, client_id, use_websocket=use_websocket)
+                for worker, wprompt in worker_payloads
+            ]
         )
 
     prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id)

distributed_upscale.py

@@ -63,6 +63,24 @@ def sync_wrapper(async_func):
 
 # Note: tensor_to_pil and pil_to_tensor are imported from utils.image
 
+
+def _parse_enabled_worker_ids(enabled_worker_ids):
+    """Parse enabled worker IDs from either JSON or list input."""
+    if isinstance(enabled_worker_ids, list):
+        return [str(worker_id) for worker_id in enabled_worker_ids]
+    if not enabled_worker_ids:
+        return []
+    if isinstance(enabled_worker_ids, str):
+        try:
+            parsed = json.loads(enabled_worker_ids)
+        except json.JSONDecodeError:
+            log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
+            return []
+        if isinstance(parsed, list):
+            return [str(wid) for wid in parsed]
+    return []
+
+
 class UltimateSDUpscaleDistributed:
     """
     Distributed version of Ultimate SD Upscale (No Upscale).
@@ -248,8 +266,8 @@ class UltimateSDUpscaleDistributed:
         if num_workers > 0 and num_tiles_per_image > 1:
             mode = "static"
         
-        log(f"USDU Dist: Workers {num_workers}")
-        
+        log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
+
         if mode == "single_gpu":
             # No workers, process all tiles locally
             return self.process_single_gpu(upscaled_image, model, positive, negative, vae,

utils/config.py

@@ -19,7 +19,8 @@ def get_default_config():
             "debug": False,
             "auto_launch_workers": False,
             "stop_workers_on_master_exit": True,
-            "master_delegate_only": False
+            "master_delegate_only": False,
+            "websocket_orchestration": True
         },
         "tunnel": {
             "status": "stopped",

web/executionUtils.js

@@ -60,7 +60,8 @@ export function setupInterceptor(extension) {
         if (extension.isEnabled) {
             const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0;
             const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0;
-            
+            const hasDistQueue = findNodesByClass(prompt.output, "DistributedQueue").length > 0;
+
             if (hasCollector || hasDistUpscale) {
                 const result = await executeParallelDistributed(extension, prompt);
                 // Immediate status check for instant feedback
@@ -69,6 +70,14 @@ export function setupInterceptor(extension) {
                 setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
                 return result;
             }
+
+            // DistributedQueue: route entire workflow to least-busy worker
+            if (hasDistQueue) {
+                const result = await executeQueueDistributed(extension, prompt);
+                extension.checkAllWorkerStatuses();
+                setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
+                return result;
+            }
         }
         return extension.originalQueuePrompt(number, prompt);
     };
@@ -94,40 +103,39 @@ export async function executeParallelDistributed(extension, promptWrapper) {
         }
         
         extension.log(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`, "debug");
-        
+
         // Check if master host might be unreachable by workers (cloudflare tunnel down)
-        const masterUrl = extension.getMasterUrl();
-        const masterHostForCheck = masterUrl.replace(/^https?:\/\//i, "").split("/")[0];
-        const isCloudflareHost = /\.(trycloudflare\.com|cloudflare\.dev)$/i.test(masterHostForCheck);
+        const masterHost = extension.config?.master?.host || '';
+        const isCloudflareHost = /\.(trycloudflare\.com|cloudflare\.dev)$/i.test(masterHost);
 
         if (isCloudflareHost && activeWorkers.length > 0) {
             // Try to verify if the cloudflare tunnel is actually up
             try {
-                const testUrl = masterUrl.endsWith("/prompt") ? masterUrl : `${masterUrl.replace(/\/$/, '')}/prompt`;
+                const testUrl = `${window.location.protocol}//${masterHost}/prompt`;
                 const response = await fetch(testUrl, {
                     method: 'GET',
                     mode: 'cors',
                     cache: 'no-cache',
                     signal: AbortSignal.timeout(3000) // 3 second timeout
                 });
-                
+
                 if (!response.ok) {
                     throw new Error('Master not reachable');
                 }
             } catch (error) {
                 // Cloudflare tunnel appears to be down
-                extension.log(`Master host ${masterUrl} is not reachable - cloudflare tunnel may be down`, "error");
-                
+                extension.log(`Master host ${masterHost} is not reachable - cloudflare tunnel may be down`, "error");
+
                 if (extension.ui?.showCloudflareWarning) {
-                    extension.ui.showCloudflareWarning(extension, masterUrl);
+                    extension.ui.showCloudflareWarning(extension, masterHost);
                 }
-                
+
                 // Stop execution - workers won't be able to send results back
                 extension.log("Blocking execution - workers cannot reach master at cloudflare domain", "error");
                 return null; // This will prevent the workflow from running
             }
         }
-        
+
         // Find all distributed nodes in the workflow
         const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector");
         const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed");
@@ -201,6 +209,126 @@ export async function executeParallelDistributed(extension, promptWrapper) {
     }
 }
 
+/**
+ * Execute a workflow via DistributedQueue - routes to the least-busy worker.
+ * The workflow runs ONLY on the selected worker, not on the master.
+ */
+export async function executeQueueDistributed(extension, promptWrapper) {
+    try {
+        const enabledWorkers = extension.enabledWorkers;
+
+        if (enabledWorkers.length === 0) {
+            extension.log("DistributedQueue: No enabled workers. Falling back to local execution.", "warn");
+            return extension.originalQueuePrompt(0, promptWrapper);
+        }
+
+        // Fetch queue status from all enabled workers
+        const workerStatuses = await fetchWorkerQueueStatuses(extension, enabledWorkers);
+
+        if (workerStatuses.length === 0) {
+            extension.log("DistributedQueue: No reachable workers. Falling back to local execution.", "warn");
+            return extension.originalQueuePrompt(0, promptWrapper);
+        }
+
+        // Select the least-busy worker
+        const selectedWorker = selectLeastBusyWorker(extension, workerStatuses);
+        const worker = selectedWorker.worker;
+        const queueRemaining = selectedWorker.queueRemaining;
+
+        extension.log(`DistributedQueue: Routing to ${worker.name} (queue_remaining=${queueRemaining})`, "info");
+
+        // Prepare the prompt with skip_dispatch=True for DistributedQueue nodes
+        const promptToSend = JSON.parse(JSON.stringify(promptWrapper.output));
+        markSkipDispatch(promptToSend);
+
+        // Dispatch to the selected worker
+        const workerUrl = extension.getWorkerUrl(worker);
+        try {
+            const response = await fetch(`${workerUrl}/prompt`, {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                mode: 'cors',
+                body: JSON.stringify({
+                    prompt: promptToSend,
+                    extra_data: { extra_pnginfo: { workflow: promptWrapper.workflow } },
+                    client_id: api.clientId
+                })
+            });
+
+            if (!response.ok) {
+                throw new Error(`Worker returned ${response.status}`);
+            }
+
+            const result = await response.json();
+            extension.log(`DistributedQueue: Dispatched to ${worker.name}, prompt_id=${result.prompt_id}`, "info");
+
+            return { prompt_id: result.prompt_id, worker_id: worker.id };
+        } catch (error) {
+            extension.log(`DistributedQueue: Failed to dispatch to ${worker.name}: ${error.message}`, "error");
+            // Fall back to local execution
+            return extension.originalQueuePrompt(0, promptWrapper);
+        }
+    } catch (error) {
+        extension.log("DistributedQueue execution failed: " + error.message, "error");
+        throw error;
+    }
+}
+
+async function fetchWorkerQueueStatuses(extension, workers) {
+    const statuses = [];
+
+    const checkPromises = workers.map(async (worker) => {
+        const url = extension.getWorkerUrl(worker, '/prompt');
+        try {
+            const response = await fetch(url, {
+                method: 'GET',
+                mode: 'cors',
+                cache: 'no-store',
+                signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
+            });
+
+            if (response.ok) {
+                const data = await response.json();
+                const queueRemaining = data.exec_info?.queue_remaining || 0;
+                return { worker, queueRemaining, online: true };
+            }
+        } catch (error) {
+            extension.log(`DistributedQueue: Worker ${worker.name} unreachable: ${error.message}`, "debug");
+        }
+        return null;
+    });
+
+    const results = await Promise.all(checkPromises);
+    return results.filter(r => r !== null);
+}
+
+function selectLeastBusyWorker(extension, statuses) {
+    // Find idle workers (queue_remaining == 0)
+    const idleWorkers = statuses.filter(s => s.queueRemaining === 0);
+
+    if (idleWorkers.length > 0) {
+        // Round-robin among idle workers
+        if (!extension._distributedQueueRRIndex) {
+            extension._distributedQueueRRIndex = 0;
+        }
+        const index = extension._distributedQueueRRIndex % idleWorkers.length;
+        extension._distributedQueueRRIndex++;
+        return idleWorkers[index];
+    }
+
+    // No idle workers - pick the one with the shortest queue
+    return statuses.reduce((min, s) => s.queueRemaining < min.queueRemaining ? s : min);
+}
+
+function markSkipDispatch(promptObj) {
+    for (const node of Object.values(promptObj)) {
+        if (node && typeof node === 'object' && node.class_type === 'DistributedQueue') {
+            node.inputs = node.inputs || {};
+            node.inputs.skip_dispatch = true;
+        }
+    }
+}
+
 export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) {
     let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt));
     const isMaster = participantId === 'master';
@@ -497,24 +625,102 @@ async function dispatchToWorker(extension, worker, prompt, workflow, imageRefere
     
     const promptToSend = {
         prompt,
-        extra_data: { extra_pnginfo: { workflow } },
+        workflow,
         client_id: api.clientId
     };
-    
+
     extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug");
-    
+
     try {
-        await fetch(`${workerUrl}/prompt`, { 
-            method: 'POST', 
-            headers: { 'Content-Type': 'application/json' }, 
+        if (extension.config?.settings?.websocket_orchestration) {
+            const wsUrl = buildWorkerWebSocketUrl(workerUrl);
+            await dispatchPromptViaWebSocket(extension, wsUrl, promptToSend);
+            return;
+        }
+        await fetch(`${workerUrl}/prompt`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
             mode: 'cors',
-            body: JSON.stringify(promptToSend) 
+            body: JSON.stringify({
+                prompt,
+                extra_data: { extra_pnginfo: { workflow } },
+                client_id: api.clientId
+            })
         });
     } catch (e) {
         extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error");
     }
 }
 
+function buildWorkerWebSocketUrl(workerUrl) {
+    const url = new URL(workerUrl);
+    url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:';
+    url.pathname = '/distributed/worker_ws';
+    url.search = '';
+    return url.toString();
+}
+
+function generateRequestId() {
+    if (crypto.randomUUID) {
+        return crypto.randomUUID();
+    }
+    return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => {
+        const r = Math.random() * 16 | 0;
+        const v = c === 'x' ? r : (r & 0x3 | 0x8);
+        return v.toString(16);
+    });
+}
+
+async function dispatchPromptViaWebSocket(extension, wsUrl, payload) {
+    const requestId = generateRequestId();
+    const message = {
+        type: 'dispatch_prompt',
+        request_id: requestId,
+        prompt: payload.prompt,
+        workflow: payload.workflow,
+        client_id: payload.client_id
+    };
+
+    await new Promise((resolve, reject) => {
+        const ws = new WebSocket(wsUrl);
+        const timeoutId = setTimeout(() => {
+            ws.close();
+            reject(new Error('Websocket dispatch timed out.'));
+        }, 60000);
+
+        ws.addEventListener('open', () => {
+            ws.send(JSON.stringify(message));
+        });
+
+        ws.addEventListener('message', (event) => {
+            try {
+                const data = JSON.parse(event.data || '{}');
+                if (data.type !== 'dispatch_ack' || data.request_id !== requestId) {
+                    return;
+                }
+                clearTimeout(timeoutId);
+                ws.close();
+                if (data.ok) {
+                    resolve();
+                } else {
+                    reject(new Error(data.error || 'Worker rejected websocket dispatch.'));
+                }
+            } catch (error) {
+                clearTimeout(timeoutId);
+                ws.close();
+                reject(error);
+            }
+        });
+
+        ws.addEventListener('error', () => {
+            clearTimeout(timeoutId);
+            ws.close();
+            reject(new Error('Websocket connection failed.'));
+        });
+    });
+    extension.log(`[Distributed] Websocket dispatch succeeded for ${wsUrl}`, "debug");
+}
+
 export async function loadImagesForWorker(extension, imageReferences) {
     const images = [];
     
@@ -667,9 +873,10 @@ export async function performPreflightCheck(extension, workers) {
             const response = await fetch(url, {
                 method: 'GET',
                 mode: 'cors',
+                cache: 'no-store',
                 signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
             });
-            
+
             if (response.ok) {
                 extension.log(`Worker ${worker.name} is active`, "debug");
                 return { worker, active: true };
@@ -678,6 +885,10 @@ export async function performPreflightCheck(extension, workers) {
                 return { worker, active: false };
             }
         } catch (error) {
+            if (error?.name === 'AbortError') {
+                extension.log(`Worker ${worker.name} pre-flight check timed out; assuming active`, "debug");
+                return { worker, active: true, uncertain: true };
+            }
             extension.log(`Worker ${worker.name} is offline or unreachable: ${error.message}`, "debug");
             return { worker, active: false };
         }

web/main.js

@@ -331,7 +331,7 @@ class DistributedExtension {
         }
 
         if (this.panelElement) {
-            renderSidebarContent(this, this.panelElement);
+            await renderSidebarContent(this, this.panelElement);
         }
     }
 
@@ -636,21 +636,21 @@ class DistributedExtension {
 
     async launchWorker(workerId) {
         const worker = this.config.workers.find(w => w.id === workerId);
-        const launchBtn = document.querySelector(`#controls-${workerId} button`);
 
         // If worker is disabled, enable it first
         if (!worker.enabled) {
             await this.updateWorkerEnabled(workerId, true);
-            
+
             // Update the checkbox UI
             const checkbox = document.getElementById(`gpu-${workerId}`);
             if (checkbox) {
                 checkbox.checked = true;
             }
-            
-            this.updateSummary();
         }
 
+        // Re-query button AFTER updateWorkerEnabled (which may re-render sidebar)
+        const launchBtn = document.querySelector(`#controls-${workerId} button`);
+
         this.ui.updateStatusDot(workerId, "#f0ad4e", "Launching...", true);
         this.state.setWorkerLaunching(workerId, true);
 

worker_monitor.py

@@ -39,6 +39,7 @@ except ImportError:
 def monitor_and_run(master_pid, command):
     """Run command and monitor master process."""
     # Start the actual worker process
+    print(f"[Distributed] Launching worker command: {' '.join(command)}")
     worker_process = subprocess.Popen(command)
     
     print(f"[Distributed] Started worker PID: {worker_process.pid}")