refactor: clarify load-balanced dispatch reservations

README.md

@@ -147,7 +147,7 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
 Control your distributed cluster programmatically without opening the browser.
 
 * **Endpoint:** `POST /distributed/queue`
-* **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
+* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
 * **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
 
 > **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
@@ -175,7 +175,7 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
 | **Distributed Value** | Outputs per-worker override values with fallback to default |
 | **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
 | **Image Batch Divider** | Splits image batches for multi-GPU output |
-| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
+| **Audio Batch Divider** | Splits audio batches for multi-GPU output |
 | **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
 | **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
 

api/job_routes.py

@@ -295,24 +295,17 @@ async def job_complete_endpoint(request):
             errors.append("worker_id: expected non-empty string")
         if not isinstance(batch_idx, int) or batch_idx < 0:
             errors.append("batch_idx: expected non-negative integer")
-        image_provided = image_payload is not None
-        audio_provided = audio_payload is not None
-        if not image_provided and not audio_provided:
-            errors.append("expected at least one of image or audio")
-        if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
+        if not isinstance(image_payload, str) or not image_payload.strip():
             errors.append("image: expected non-empty base64 PNG string")
-        if audio_provided and not isinstance(audio_payload, dict):
+        if audio_payload is not None and not isinstance(audio_payload, dict):
             errors.append("audio: expected object when provided")
         if not isinstance(is_last, bool):
             errors.append("is_last: expected boolean")
         if errors:
             return await handle_api_error(request, errors, 400)
 
-        try:
-            tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
-            decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
-        except ValueError as exc:
-            return await handle_api_error(request, exc, 400)
+        tensor = _decode_canonical_png_tensor(image_payload)
+        decoded_audio = _decode_audio_payload(audio_payload) if audio_payload is not None else None
         multi_job_id = job_id.strip()
         worker_id = worker_id.strip()
 

api/orchestration/dispatch.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import time
 import uuid
 
 import aiohttp
@@ -207,9 +208,13 @@ async def _probe_worker_queue(worker, semaphore, probe_timeout):
         payload = await probe_worker(worker_url, timeout=probe_timeout)
         if payload is None:
             return None
+        try:
+            reserved_slots = int(worker.get("reserved_slots", 0) or 0)
+        except (TypeError, ValueError):
+            reserved_slots = 0
         return {
             "worker": worker,
-            "queue_remaining": _extract_queue_remaining(payload),
+            "queue_remaining": _extract_queue_remaining(payload) + max(reserved_slots, 0),
         }
 
 
@@ -222,17 +227,7 @@ def _select_idle_round_robin(statuses):
     return statuses[index]
 
 
-async def select_least_busy_worker(
-    workers,
-    trace_execution_id=None,
-    probe_concurrency=8,
-    probe_timeout=3.0,
-):
-    """Select one worker by queue depth, round-robin among idle workers."""
-    if not workers:
-        return None
-
-    probe_limit = parse_positive_int(probe_concurrency, 8)
+async def _probe_worker_queues(workers, probe_limit, probe_timeout):
     probe_semaphore = asyncio.Semaphore(probe_limit)
     statuses = await asyncio.gather(
         *[
@@ -240,29 +235,66 @@ async def select_least_busy_worker(
             for worker in workers
         ]
     )
-    statuses = [status for status in statuses if status is not None]
-    if not statuses:
-        if trace_execution_id:
-            trace_info(trace_execution_id, "Least-busy selection failed: no worker queue probes succeeded.")
-        else:
-            log("[Distributed] Least-busy selection failed: no worker queue probes succeeded.")
-        return None
+    return [status for status in statuses if status is not None]
 
+
+def _select_worker_status(statuses, require_idle=False):
     idle_statuses = [status for status in statuses if status["queue_remaining"] == 0]
     if idle_statuses:
-        selected = _select_idle_round_robin(idle_statuses)
-    else:
-        selected = min(statuses, key=lambda status: status["queue_remaining"])
+        return _select_idle_round_robin(idle_statuses)
+    if require_idle:
+        return None
+    return min(statuses, key=lambda status: status["queue_remaining"])
 
-    worker = selected["worker"]
-    queue_remaining = selected["queue_remaining"]
-    if trace_execution_id:
-        trace_debug(
-            trace_execution_id,
-            f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}",
-        )
-    else:
-        debug_log(
-            f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}"
-        )
-    return worker
+
+async def select_least_busy_worker(
+    workers,
+    trace_execution_id=None,
+    probe_concurrency=8,
+    probe_timeout=3.0,
+    require_idle=False,
+    idle_poll_interval=1.0,
+    idle_wait_timeout=None,
+):
+    """Select a worker by queue depth, optionally waiting until one is idle."""
+    if not workers:
+        return None
+
+    probe_limit = parse_positive_int(probe_concurrency, 8)
+    poll_interval = max(float(idle_poll_interval or 0), 0.0)
+    deadline = None
+    if idle_wait_timeout is not None:
+        deadline = time.monotonic() + max(float(idle_wait_timeout), 0.0)
+
+    while True:
+        statuses = await _probe_worker_queues(workers, probe_limit, probe_timeout)
+        if not statuses:
+            if trace_execution_id:
+                trace_info(trace_execution_id, "Least-busy selection failed: no worker queue probes succeeded.")
+            else:
+                log("[Distributed] Least-busy selection failed: no worker queue probes succeeded.")
+            return None
+
+        selected = _select_worker_status(statuses, require_idle=require_idle)
+        if selected is not None:
+            worker = selected["worker"]
+            queue_remaining = selected["queue_remaining"]
+            if trace_execution_id:
+                trace_debug(
+                    trace_execution_id,
+                    f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}",
+                )
+            else:
+                debug_log(
+                    f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}"
+                )
+            return worker
+
+        if deadline is not None and time.monotonic() >= deadline:
+            if trace_execution_id:
+                trace_info(trace_execution_id, "No idle load-balance worker became available before timeout.")
+            else:
+                debug_log("No idle load-balance worker became available before timeout.")
+            return None
+
+        await asyncio.sleep(poll_interval)

api/orchestration/prompt_transform.py

@@ -125,209 +125,6 @@ def _find_upstream_nodes(prompt_obj, start_ids):
     return connected
 
 
-_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = {
-    "PrimitiveBoolean",
-    "PrimitiveFloat",
-    "PrimitiveInt",
-    "PrimitiveNode",
-    "PrimitiveString",
-}
-
-_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = {
-    "LoadImage",
-}
-
-_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"}
-_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"}
-
-# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the
-# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility
-# only after its connected inputs recursively prove safe.
-_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"}
-_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = {
-    "CreateList": ("inputs.",),
-}
-
-# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily.
-_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None
-
-
-def _get_delegate_master_node_class_mappings():
-    """Return ComfyUI node-class mappings when available."""
-    if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None:
-        return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS
-    try:
-        import nodes as comfy_nodes  # type: ignore
-    except Exception:  # pragma: no cover - depends on ComfyUI runtime imports
-        return {}
-    return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
-
-
-def _get_delegate_master_node_class(class_type):
-    mappings = _get_delegate_master_node_class_mappings()
-    return mappings.get(class_type) if isinstance(mappings, dict) else None
-
-
-def _normalize_delegate_master_return_type(return_type):
-    if return_type is None:
-        return ""
-    return str(return_type).strip().upper()
-
-
-def _delegate_master_type_is_safe_scalar(type_name):
-    return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES
-
-
-def _delegate_master_type_is_safe_config(type_name):
-    return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES
-
-
-def _delegate_master_output_is_safe_scalar(class_type, output_index):
-    """Return True when a registered node output is lightweight config data."""
-    if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES:
-        return True
-    node_class = _get_delegate_master_node_class(class_type)
-    return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else ()
-    try:
-        output_type = return_types[int(output_index)]
-    except (IndexError, TypeError, ValueError):
-        return False
-    return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type))
-
-
-def _get_delegate_master_input_types(class_type):
-    node_class = _get_delegate_master_node_class(class_type)
-    input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None
-    if callable(input_types):
-        try:
-            input_types = input_types()
-        except TypeError:
-            return {}
-    return input_types if isinstance(input_types, dict) else {}
-
-
-def _normalize_delegate_master_input_type(input_spec):
-    if isinstance(input_spec, (list, tuple)) and input_spec:
-        return _normalize_delegate_master_return_type(input_spec[0])
-    return _normalize_delegate_master_return_type(input_spec)
-
-
-def _delegate_master_input_is_safe_scalar(class_type, input_name):
-    """Return True when a registered downstream input expects config data."""
-    for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()):
-        if input_name.startswith(prefix):
-            return True
-    input_types = _get_delegate_master_input_types(class_type)
-    for section_name in ("required", "optional"):
-        section = input_types.get(section_name, {})
-        if isinstance(section, dict) and input_name in section:
-            input_type = _normalize_delegate_master_input_type(section[input_name])
-            return _delegate_master_type_is_safe_config(input_type)
-    return False
-
-
-def _is_delegate_master_always_retained_upstream_node(node):
-    if not isinstance(node, dict):
-        return False
-    class_type = node.get("class_type")
-    return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES
-
-
-def _is_delegate_master_retained_upstream_node(node, output_index=0):
-    """Return True for lightweight upstream nodes safe to keep on the master."""
-    if not isinstance(node, dict):
-        return False
-    class_type = node.get("class_type")
-    if not isinstance(class_type, str):
-        return False
-    return (
-        class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES
-        or class_type.startswith("Primitive")
-        or _delegate_master_output_is_safe_scalar(class_type, output_index)
-    )
-
-
-def _collect_delegate_master_retained_upstream_branch(
-    prompt_obj,
-    node_id,
-    output_index,
-    memo,
-    visiting,
-):
-    """Return safe retained branch nodes, or None when the branch is not safe."""
-    node_id = str(node_id)
-    cache_key = (node_id, output_index)
-    if cache_key in memo:
-        cached = memo[cache_key]
-        return None if cached is None else set(cached)
-    if cache_key in visiting:
-        memo[cache_key] = None
-        return None
-
-    node = prompt_obj.get(node_id)
-    if not _is_delegate_master_retained_upstream_node(node, output_index):
-        memo[cache_key] = None
-        return None
-
-    visiting.add(cache_key)
-    retained = {node_id}
-    inputs = node.get("inputs", {}) if isinstance(node, dict) else {}
-    class_type = node.get("class_type") if isinstance(node, dict) else None
-    for input_name, value in inputs.items():
-        if not (isinstance(value, list) and len(value) == 2):
-            continue
-        if not _delegate_master_input_is_safe_scalar(class_type, input_name):
-            visiting.remove(cache_key)
-            memo[cache_key] = None
-            return None
-        source_id = str(value[0])
-        branch = _collect_delegate_master_retained_upstream_branch(
-            prompt_obj,
-            source_id,
-            value[1],
-            memo,
-            visiting,
-        )
-        if branch is None:
-            visiting.remove(cache_key)
-            memo[cache_key] = None
-            return None
-        retained.update(branch)
-
-    visiting.remove(cache_key)
-    memo[cache_key] = frozenset(retained)
-    return retained
-
-
-def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
-    """Return lightweight upstream nodes needed by kept delegate-master nodes."""
-    connected = set()
-    memo = {}
-    for node_id in start_ids:
-        node = prompt_obj.get(str(node_id)) or {}
-        inputs = node.get("inputs", {})
-        class_type = node.get("class_type") if isinstance(node, dict) else None
-        for input_name, value in inputs.items():
-            if not (isinstance(value, list) and len(value) == 2):
-                continue
-            source_node = prompt_obj.get(str(value[0]))
-            if _is_delegate_master_always_retained_upstream_node(source_node):
-                connected.add(str(value[0]))
-                continue
-            if not _delegate_master_input_is_safe_scalar(class_type, input_name):
-                continue
-            branch = _collect_delegate_master_retained_upstream_branch(
-                prompt_obj,
-                value[0],
-                value[1],
-                memo,
-                set(),
-            )
-            if branch is not None:
-                connected.update(branch)
-    return connected
-
-
 def prune_prompt_for_worker(prompt_obj):
     """Prune worker prompt to distributed nodes and their upstream dependencies."""
     collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
@@ -351,33 +148,16 @@ def prune_prompt_for_worker(prompt_obj):
         downstream = _find_downstream_nodes(prompt_obj, [dist_id])
         has_removed_downstream = any(node_id != dist_id for node_id in downstream)
         if has_removed_downstream:
-            original_node = prompt_obj.get(str(dist_id), {})
-            class_type = original_node.get("class_type")
-            inputs = original_node.get("inputs", {})
-            image_connected = class_type != "DistributedCollector" or (
-                isinstance(inputs.get("images"), list)
-                and len(inputs["images"]) == 2
-            )
-            audio_connected = (
-                class_type == "DistributedCollector"
-                and isinstance(inputs.get("audio"), list)
-                and len(inputs["audio"]) == 2
-            )
-
-            if image_connected:
-                preview_id = next_id()
-                pruned_prompt[preview_id] = {
-                    "inputs": {"images": [dist_id, 0]},
-                    "class_type": "PreviewImage",
-                    "_meta": {"title": "Preview Image (auto-added)"},
-                }
-            elif audio_connected:
-                preview_id = next_id()
-                pruned_prompt[preview_id] = {
-                    "inputs": {"audio": [dist_id, 1]},
-                    "class_type": "PreviewAudio",
-                    "_meta": {"title": "Preview Audio (auto-added)"},
-                }
+            preview_id = next_id()
+            pruned_prompt[preview_id] = {
+                "inputs": {
+                    "images": [dist_id, 0],
+                },
+                "class_type": "PreviewImage",
+                "_meta": {
+                    "title": "Preview Image (auto-added)",
+                },
+            }
 
     return pruned_prompt
 
@@ -387,9 +167,6 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
     downstream = _find_downstream_nodes(prompt_obj, collector_ids)
     nodes_to_keep = set(collector_ids)
     nodes_to_keep.update(downstream)
-    nodes_to_keep.update(
-        _find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
-    )
 
     pruned_prompt = {}
     for node_id in nodes_to_keep:
@@ -417,10 +194,6 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
         collector_entry = pruned_prompt.get(collector_id)
         if not collector_entry:
             continue
-        original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
-        original_images = original_inputs.get("images")
-        if not (isinstance(original_images, list) and len(original_images) == 2):
-            continue
         placeholder_id = next_id()
         pruned_prompt[placeholder_id] = {
             "class_type": "DistributedEmptyImage",

api/queue_orchestration.py

@@ -46,12 +46,159 @@ def ensure_distributed_state(server_instance=None):
         ps.distributed_pending_jobs = {}
     if not hasattr(ps, "distributed_jobs_lock"):
         ps.distributed_jobs_lock = asyncio.Lock()
+    if not hasattr(ps, "distributed_worker_reservations"):
+        ps.distributed_worker_reservations = {}
 
 
 # Initialize top-level distributed queue state at module import time.
 ensure_distributed_state()
 
 
+def _worker_id(worker):
+    return str(worker.get("id"))
+
+
+def _non_negative_int(value, default=0):
+    try:
+        parsed = int(value if value is not None else default)
+    except (TypeError, ValueError):
+        parsed = default
+    return max(parsed, 0)
+
+
+async def _release_worker_reservation(worker_id):
+    if not worker_id:
+        return
+    ensure_distributed_state()
+    async with prompt_server.distributed_jobs_lock:
+        reservations = prompt_server.distributed_worker_reservations
+        current = _non_negative_int(reservations.get(worker_id, 0))
+        if current <= 1:
+            reservations.pop(worker_id, None)
+        else:
+            reservations[worker_id] = current - 1
+
+
+async def _candidate_workers_with_reservations(candidate_workers):
+    ensure_distributed_state()
+    async with prompt_server.distributed_jobs_lock:
+        reservations = dict(prompt_server.distributed_worker_reservations)
+    return [
+        {
+            **worker,
+            "reserved_slots": reservations.get(_worker_id(worker), 0),
+        }
+        for worker in candidate_workers
+    ]
+
+
+async def _reserve_selected_worker(selected_worker):
+    selected_worker_id = _worker_id(selected_worker)
+    selected_reserved_slots = _non_negative_int(selected_worker.get("reserved_slots", 0))
+    async with prompt_server.distributed_jobs_lock:
+        reservations = prompt_server.distributed_worker_reservations
+        current_reserved_slots = _non_negative_int(reservations.get(selected_worker_id, 0))
+        if current_reserved_slots != selected_reserved_slots:
+            return None
+        reservations[selected_worker_id] = current_reserved_slots + 1
+        return selected_worker_id
+
+
+def _normalized_idle_poll_interval(value):
+    try:
+        poll_interval = float(value or 0)
+    except (TypeError, ValueError):
+        poll_interval = 1.0
+    return max(poll_interval, 0.05)
+
+
+async def _select_and_reserve_load_balance_worker(
+    candidate_workers,
+    execution_trace_id,
+    worker_probe_concurrency,
+    idle_poll_interval,
+):
+    poll_interval = _normalized_idle_poll_interval(idle_poll_interval)
+    while candidate_workers:
+        reserved_candidates = await _candidate_workers_with_reservations(candidate_workers)
+        selected_worker = await select_least_busy_worker(
+            reserved_candidates,
+            trace_execution_id=execution_trace_id,
+            probe_concurrency=worker_probe_concurrency,
+            require_idle=True,
+            idle_poll_interval=0,
+            idle_wait_timeout=0,
+        )
+        if selected_worker is None:
+            await asyncio.sleep(poll_interval)
+            continue
+
+        reserved_worker_id = await _reserve_selected_worker(selected_worker)
+        if reserved_worker_id is not None:
+            return selected_worker
+
+        await asyncio.sleep(0)
+    return None
+
+
+def _load_balance_candidate_workers(active_workers, delegate_master, master_url):
+    candidate_workers = list(active_workers)
+    if not delegate_master:
+        candidate_workers.append(
+            {
+                "id": "master",
+                "name": "Master",
+                "host": master_url,
+                "type": "local",
+            }
+        )
+    return candidate_workers
+
+
+async def _resolve_load_balance_selection(
+    active_workers,
+    delegate_master,
+    master_url,
+    config,
+    execution_trace_id,
+    worker_probe_concurrency,
+):
+    candidate_workers = _load_balance_candidate_workers(active_workers, delegate_master, master_url)
+    selected_worker = None
+    if candidate_workers:
+        settings = (config.get("settings", {}) or {})
+        idle_poll_interval = parse_positive_float(
+            settings.get("load_balance_idle_poll_interval_seconds"),
+            1.0,
+        )
+        selected_worker = await _select_and_reserve_load_balance_worker(
+            candidate_workers,
+            execution_trace_id,
+            worker_probe_concurrency,
+            idle_poll_interval,
+        )
+
+    if selected_worker is not None and _worker_id(selected_worker) == "master":
+        trace_debug(
+            execution_trace_id,
+            "Load-balance selected master for execution (workers skipped).",
+        )
+        return [], False, "master"
+
+    if selected_worker is not None:
+        trace_debug(
+            execution_trace_id,
+            f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).",
+        )
+        return [selected_worker], True, _worker_id(selected_worker)
+
+    trace_debug(
+        execution_trace_id,
+        "Load-balance requested but no idle/probeable execution candidates were available.",
+    )
+    return [], False, None
+
+
 async def _ensure_distributed_queue(job_id):
     """Ensure a queue exists for the given distributed job ID."""
     ensure_distributed_state()
@@ -259,160 +406,123 @@ async def orchestrate_distributed_execution(
         probe_concurrency=worker_probe_concurrency,
     )
 
+    reserved_worker_id = None
     if load_balance_requested:
-        candidate_workers = list(active_workers)
-        if not delegate_master:
-            # Include master in load balancing only when master participation is enabled.
-            candidate_workers.append(
-                {
-                    "id": "master",
-                    "name": "Master",
-                    "host": master_url,
-                    "type": "local",
-                }
+        active_workers, delegate_master, reserved_worker_id = await _resolve_load_balance_selection(
+            active_workers,
+            delegate_master,
+            master_url,
+            config,
+            execution_trace_id,
+            worker_probe_concurrency,
+        )
+
+    try:
+        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_index, discovery_prefix)
+
+        if not job_id_map:
+            trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
+            queue_result = await queue_prompt_payload(
+                prompt_obj,
+                workflow_meta,
+                client_id,
+                include_queue_metadata=True,
+            )
+            return (
+                queue_result["prompt_id"],
+                queue_result["number"],
+                0,
+                queue_result.get("node_errors", {}),
             )
 
-        selected_worker = None
-        if candidate_workers:
-            selected_worker = await select_least_busy_worker(
-                candidate_workers,
-                trace_execution_id=execution_trace_id,
-                probe_concurrency=worker_probe_concurrency,
-            )
-        if selected_worker is None and candidate_workers:
+        for job_id in job_id_map.values():
+            await _ensure_distributed_queue(job_id)
+
+        master_prompt = prompt_index.copy_prompt()
+        master_prompt = apply_participant_overrides(
+            master_prompt,
+            "master",
+            enabled_ids,
+            job_id_map,
+            master_url,
+            delegate_master,
+            prompt_index,
+        )
+
+        if delegate_master:
+            collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
+            upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
+            if upscale_nodes:
+                debug_log(
+                    "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
+                )
+            elif not collector_ids:
+                debug_log(
+                    "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
+                )
+            else:
+                master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
+
+        if active_workers:
             trace_debug(
                 execution_trace_id,
-                "Load-balance selection probe failed; using first available candidate.",
+                "Active distributed workers: "
+                + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
             )
-            selected_worker = candidate_workers[0]
-
-        if selected_worker is not None and str(selected_worker.get("id")) == "master":
-            # Master selected as least busy; run master workload only.
-            active_workers = []
-            delegate_master = False
-            trace_debug(
-                execution_trace_id,
-                "Load-balance selected master for execution (workers skipped).",
+        worker_payloads = []
+        if active_workers:
+            worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
+            media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
+            worker_payloads = await asyncio.gather(
+                *[
+                    _prepare_worker_payload(
+                        worker,
+                        prompt_index,
+                        enabled_ids,
+                        job_id_map,
+                        master_url,
+                        config,
+                        delegate_master,
+                        execution_trace_id,
+                        worker_prep_semaphore,
+                        media_sync_semaphore,
+                        media_sync_timeout_seconds,
+                    )
+                    for worker in active_workers
+                ]
             )
-        elif selected_worker is not None:
-            active_workers = [selected_worker]
-            # Worker selected as least busy; keep master orchestrator-only for this run.
-            delegate_master = True
-            trace_debug(
-                execution_trace_id,
-                f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).",
+
+        if worker_payloads:
+            await asyncio.gather(
+                *[
+                    dispatch_worker_prompt(
+                        worker,
+                        wprompt,
+                        workflow_meta,
+                        client_id,
+                        use_websocket=use_websocket,
+                        trace_execution_id=execution_trace_id,
+                    )
+                    for worker, wprompt in worker_payloads
+                ]
             )
-        else:
-            trace_debug(
-                execution_trace_id,
-                "Load-balance requested but no execution candidates were available.",
-            )
-            active_workers = []
-            delegate_master = False
 
-    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_index, discovery_prefix)
-
-    if not job_id_map:
-        trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
         queue_result = await queue_prompt_payload(
-            prompt_obj,
+            master_prompt,
             workflow_meta,
             client_id,
             include_queue_metadata=True,
         )
-        return (
-            queue_result["prompt_id"],
-            queue_result["number"],
-            0,
-            queue_result.get("node_errors", {}),
-        )
-
-    for job_id in job_id_map.values():
-        await _ensure_distributed_queue(job_id)
-
-    master_prompt = prompt_index.copy_prompt()
-    master_prompt = apply_participant_overrides(
-        master_prompt,
-        "master",
-        enabled_ids,
-        job_id_map,
-        master_url,
-        delegate_master,
-        prompt_index,
-    )
-
-    if delegate_master:
-        collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
-        upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
-        if upscale_nodes:
-            debug_log(
-                "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
-            )
-        elif not collector_ids:
-            debug_log(
-                "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
-            )
-        else:
-            master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
-
-    if active_workers:
+        prompt_id = queue_result["prompt_id"]
+        prompt_number = queue_result["number"]
+        node_errors = queue_result.get("node_errors", {})
         trace_debug(
             execution_trace_id,
-            "Active distributed workers: "
-            + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
+            f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
         )
-    worker_payloads = []
-    if active_workers:
-        worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
-        media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
-        worker_payloads = await asyncio.gather(
-            *[
-                _prepare_worker_payload(
-                    worker,
-                    prompt_index,
-                    enabled_ids,
-                    job_id_map,
-                    master_url,
-                    config,
-                    delegate_master,
-                    execution_trace_id,
-                    worker_prep_semaphore,
-                    media_sync_semaphore,
-                    media_sync_timeout_seconds,
-                )
-                for worker in active_workers
-            ]
-        )
-
-    if worker_payloads:
-        await asyncio.gather(
-            *[
-                dispatch_worker_prompt(
-                    worker,
-                    wprompt,
-                    workflow_meta,
-                    client_id,
-                    use_websocket=use_websocket,
-                    trace_execution_id=execution_trace_id,
-                )
-                for worker, wprompt in worker_payloads
-            ]
-        )
-
-    queue_result = await queue_prompt_payload(
-        master_prompt,
-        workflow_meta,
-        client_id,
-        include_queue_metadata=True,
-    )
-    prompt_id = queue_result["prompt_id"]
-    prompt_number = queue_result["number"]
-    node_errors = queue_result.get("node_errors", {})
-    trace_debug(
-        execution_trace_id,
-        f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
-    )
-    return prompt_id, prompt_number, len(worker_payloads), node_errors
+        return prompt_id, prompt_number, len(worker_payloads), node_errors
+    finally:
+        await _release_worker_reservation(reserved_worker_id)

docs/comfyui-distributed-api.md

@@ -19,8 +19,8 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
 - `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
   - Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
   - Resolves enabled/selected workers.
-  - Probes and dispatches workers through `/distributed/worker_ws` by default.
-  - If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
+  - Pings workers (`GET /prompt`) to include only reachable ones.
+  - Dispatches the workflow to workers (`POST /prompt`).
   - Queues the master workflow in ComfyUI’s prompt queue.
   - If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
 
@@ -61,8 +61,7 @@ Queue a workflow for distributed execution.
 #### Fields
 
 - `prompt` (required unless `workflow.prompt` is present, object)
-  - A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
-  - This is not the normal visual workflow export from the ComfyUI editor.
+  - The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
 - `workflow` (optional, object)
   - Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
   - If you don’t care about UI metadata, you can omit it.
@@ -124,14 +123,10 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
 
 ## Worker requirements (important)
 
-For a worker to participate, it must be reachable from the master. By default:
-
-- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
-
-If `settings.websocket_orchestration=false`:
+For a worker to participate, it must be reachable from the master:
 
 - Health check: `GET <worker-base>/prompt` must return HTTP 200.
-- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
+- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
 
 Also, for collector-based flows:
 
@@ -206,7 +201,7 @@ Proxy endpoint on master that fetches logs from a configured remote/cloud worker
 
 ## Examples
 
-### 1) Minimal `curl` request envelope
+### 1) Minimal `curl`
 
 ```bash
 curl -X POST "http://127.0.0.1:8188/distributed/queue" \
@@ -214,23 +209,18 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
   -d @payload.json
 ```
 
-`payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
+Where `payload.json` contains at least:
 
 ```json
 {
   "prompt": {
-    "<node_id>": {
-      "class_type": "<node class>",
-      "inputs": {"<required_input>": "<value or connection>"}
-    }
+    "1": {"class_type": "KSampler", "inputs": {} }
   },
   "enabled_worker_ids": [],
   "client_id": "external-client"
 }
 ```
 
-Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
-
 ### 2) Python (`requests`)
 
 ```python

docs/worker-setup-guides.md

@@ -79,7 +79,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
 📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
 
 **On Runpod:**
-> If using your own template, launch ComfyUI with `--listen --enable-cors-header` and clone `ComfyUI-Distributed` into `custom_nodes`. ⚠️ **Required!**
+> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!**
 
 1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
 2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
@@ -92,7 +92,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
 	- SAGE_ATTENTION: optional optimisation (set to true/false)
 5. Deploy your pod.
 6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
-7. Download models into `/workspace/ComfyUI/models/` (these will remain on your network drive even after you terminate the pod). Example commands below:
+7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
 ```
 # Download from CivitAI
 comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN

nodes/__init__.py

@@ -26,6 +26,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "DistributedModelName": "Distributed Model Name",
     "DistributedValue": "Distributed Value",
     "ImageBatchDivider": "Image Batch Divider",
-    "AudioBatchDivider": "Audio Segment Divider",
+    "AudioBatchDivider": "Audio Batch Divider",
     "DistributedEmptyImage": "Distributed Empty Image",
 }

nodes/collector.py

@@ -22,13 +22,13 @@ prompt_server = _server.PromptServer.instance
 
 
 class DistributedCollectorNode:
-    INPUT_IS_LIST = True
     EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
 
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
+                "images": ("IMAGE",),
                 "load_balance": (
                     "BOOLEAN",
                     {
@@ -37,10 +37,7 @@ class DistributedCollectorNode:
                     },
                 ),
             },
-            "optional": {
-                "images": ("IMAGE",),
-                "audio": ("AUDIO",),
-            },
+            "optional": { "audio": ("AUDIO",) },
             "hidden": {
                 "multi_job_id": ("STRING", {"default": ""}),
                 "is_worker": ("BOOLEAN", {"default": False}),
@@ -58,74 +55,7 @@ class DistributedCollectorNode:
     FUNCTION = "run"
     CATEGORY = "image"
     
-    @staticmethod
-    def _unwrap_list_input(value):
-        """Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST."""
-        if isinstance(value, (list, tuple)) and len(value) == 1:
-            return value[0]
-        return value
-
-    def _normalize_images_input(self, images):
-        """Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
-        if isinstance(images, (list, tuple)):
-            if not images:
-                raise ValueError("Collector received an empty image list")
-            if not all(isinstance(image, torch.Tensor) for image in images):
-                raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
-            if len(images) == 1:
-                return ensure_contiguous(images[0])
-            return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
-        return ensure_contiguous(images)
-
-    def _normalize_audio_input(self, audio):
-        """Collapse ComfyUI list AUDIO inputs into a single AUDIO payload when present."""
-        if not isinstance(audio, (list, tuple)):
-            return audio
-
-        audio_items = [item for item in audio if item is not None]
-        if not audio_items:
-            return None
-        if len(audio_items) == 1:
-            return audio_items[0]
-
-        waveforms = []
-        sample_rate = 44100
-        for item in audio_items:
-            if not isinstance(item, dict):
-                raise TypeError("Collector expected AUDIO list items to be dictionaries")
-            waveform = item.get("waveform")
-            if waveform is None or waveform.numel() == 0:
-                continue
-            waveforms.append(waveform)
-            if sample_rate == 44100:
-                sample_rate = item.get("sample_rate", 44100)
-
-        if not waveforms:
-            return None
-        return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
-
-    def run(self, images=None, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
-        if images is not None:
-            images = self._normalize_images_input(images)
-        audio = self._normalize_audio_input(audio)
-        load_balance = self._unwrap_list_input(load_balance)
-        multi_job_id = self._unwrap_list_input(multi_job_id)
-        is_worker = self._unwrap_list_input(is_worker)
-        master_url = self._unwrap_list_input(master_url)
-        enabled_worker_ids = self._unwrap_list_input(enabled_worker_ids)
-        worker_batch_size = self._unwrap_list_input(worker_batch_size)
-        worker_id = self._unwrap_list_input(worker_id)
-        pass_through = self._unwrap_list_input(pass_through)
-        delegate_only = self._unwrap_list_input(delegate_only)
-
-        remote_only_master = (
-            bool(multi_job_id)
-            and not is_worker
-            and (delegate_only or is_master_delegate_only())
-        )
-        if images is None and audio is None and not remote_only_master:
-            raise ValueError("DistributedCollector requires at least one image or audio input")
-
+    def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
         # Create empty audio if not provided
         empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
 
@@ -152,56 +82,41 @@ class DistributedCollectorNode:
         return result
 
     async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
-        """Send an image batch, optionally with audio, or an audio-only completion."""
+        """Send image batch to master via canonical JSON envelopes."""
+        batch_size = image_batch.shape[0]
+        if batch_size == 0:
+            return
+
         encoded_audio = encode_audio_payload(audio)
+
         session = await get_client_session()
         url = f"{master_url}/distributed/job_complete"
+        for batch_idx in range(batch_size):
+            img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
+            byte_io = io.BytesIO()
+            img.save(byte_io, format='PNG', compress_level=0)
+            encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
+            payload = {
+                "job_id": str(multi_job_id),
+                "worker_id": str(worker_id),
+                "batch_idx": int(batch_idx),
+                "image": f"data:image/png;base64,{encoded_image}",
+                "is_last": bool(batch_idx == batch_size - 1),
+            }
+            if payload["is_last"] and encoded_audio is not None:
+                payload["audio"] = encoded_audio
 
-        payloads = []
-        batch_size = 0 if image_batch is None else image_batch.shape[0]
-        if batch_size == 0:
-            if encoded_audio is None:
-                raise ValueError("Worker completion requires image or audio data")
-            payloads.append(
-                {
-                    "job_id": str(multi_job_id),
-                    "worker_id": str(worker_id),
-                    "batch_idx": 0,
-                    "audio": encoded_audio,
-                    "is_last": True,
-                }
-            )
-        else:
-            for batch_idx in range(batch_size):
-                img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
-                byte_io = io.BytesIO()
-                img.save(byte_io, format='PNG', compress_level=0)
-                encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
-                payload = {
-                    "job_id": str(multi_job_id),
-                    "worker_id": str(worker_id),
-                    "batch_idx": int(batch_idx),
-                    "image": f"data:image/png;base64,{encoded_image}",
-                    "is_last": bool(batch_idx == batch_size - 1),
-                }
-                if payload["is_last"] and encoded_audio is not None:
-                    payload["audio"] = encoded_audio
-                payloads.append(payload)
-
-        for payload in payloads:
-            timeout_seconds = 60 if "image" in payload else 600
             try:
                 async with session.post(
                     url,
                     json=payload,
-                    timeout=aiohttp.ClientTimeout(total=timeout_seconds),
+                    timeout=aiohttp.ClientTimeout(total=60),
                 ) as response:
                     response.raise_for_status()
             except Exception as e:
-                media_type = "image/audio" if "image" in payload else "audio-only"
-                log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
+                log(f"Worker - Failed to send canonical image envelope to master: {e}")
                 debug_log(f"Worker - Full error details: URL={url}")
-                raise
+                raise  # Re-raise to handle at caller level
 
     def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
         """Combine audio from master and workers into a single audio output.
@@ -283,8 +198,8 @@ class DistributedCollectorNode:
         images_on_cpu,
         delegate_mode: bool,
         fallback_images,
-    ):
-        """Assemble final tensor, or return None when the job contains only audio."""
+    ) -> torch.Tensor:
+        """Assemble final tensor: master first, then workers in enabled order."""
         ordered_tensors = []
         if not delegate_mode and images_on_cpu is not None:
             for i in range(master_batch_size):
@@ -315,15 +230,15 @@ class DistributedCollectorNode:
 
         if cpu_tensors:
             return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
-        if fallback_images is not None:
+        elif fallback_images is not None:
             return ensure_contiguous(fallback_images)
-        return None
+        else:
+            raise ValueError("No image data collected from master or workers")
 
     async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
         if is_worker:
             # Worker mode: send images and audio to master in a single batch
-            image_count = 0 if images is None else images.shape[0]
-            debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
+            debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master")
             await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
             return (images, audio if audio is not None else self.EMPTY_AUDIO)
         else:
@@ -358,15 +273,14 @@ class DistributedCollectorNode:
                 master_audio = None
                 debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
             else:
-                if images is None:
-                    images_on_cpu = None
-                    master_batch_size = 0
-                else:
-                    images_on_cpu = ensure_contiguous(images.cpu())
-                    master_batch_size = images.shape[0]
+                images_on_cpu = images.cpu()
+                master_batch_size = images.shape[0]
                 master_audio = audio  # Keep master's audio for later
                 debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
 
+                # Ensure master images are contiguous
+                images_on_cpu = ensure_contiguous(images_on_cpu)
+
 
             # Initialize storage for collected images and audio
             worker_images = {}  # Dict to store images by worker_id and index
@@ -538,19 +452,18 @@ class DistributedCollectorNode:
                 if multi_job_id in prompt_server.distributed_pending_jobs:
                     del prompt_server.distributed_pending_jobs[multi_job_id]
 
-            combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
             try:
                 combined = self._reorder_and_combine_tensors(
                     worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
                 )
-                if combined is None:
-                    debug_log(f"Master - Job {multi_job_id} complete with audio only")
-                else:
-                    debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
-                              f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
+                debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
+                          f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
+
+                # Combine audio from master and workers
+                combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
 
                 return (combined, combined_audio)
             except Exception as e:
                 log(f"Master - Error combining images: {e}")
-                # Preserve collected audio even when image assembly fails.
-                return (images, combined_audio)
+                # Return just the master images as fallback
+                return (images, audio if audio is not None else self.EMPTY_AUDIO)

nodes/distributed_upscale.py

@@ -56,22 +56,22 @@ class UltimateSDUpscaleDistributed(
     """
     Distributed version of Ultimate SD Upscale (No Upscale).
     
-    Supports two currently selected processing modes:
+    Supports three processing modes:
     1. Single GPU: No workers available, process everything locally
-    2. Distributed tile queue: Workers pull tiles from a shared queue
+    2. Static Mode: Small batches, distributes tiles across workers (flattened)
+    3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
     
     Features:
-    - Tile-based batch handling for video/image upscaling
+    - Multi-mode batch handling for efficient video/image upscaling
     - Tiled VAE support for memory efficiency
-    - Shared work queue so faster workers can process more tiles
+    - Dynamic load balancing for large batches
     - Backward compatible with single-image workflows
     
     Environment Variables:
     - COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
     - COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
     
-    The hidden dynamic_threshold input is retained for workflow compatibility but
-    does not affect the current mode-selection policy.
+    Threshold: dynamic_threshold input controls mode switch (default 8)
     """
 
     def __init__(self):

nodes/utilities.py

@@ -269,7 +269,7 @@ class ImageBatchDivider:
 
 
 class AudioBatchDivider:
-    """Divides an audio waveform into sequential segments along the time dimension."""
+    """Divides an audio waveform into multiple parts along the time/samples dimension."""
 
     @classmethod
     def INPUT_TYPES(s):
@@ -282,7 +282,7 @@ class AudioBatchDivider:
                     "max": 10,
                     "step": 1,
                     "display": "number",
-                    "tooltip": "Number of sequential time segments to create"
+                    "tooltip": "Number of parts to divide the audio into"
                 }),
             }
         }

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-Distributed"
 description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
-version = "1.4.7"
+version = "1.4.4"
 license = {file = "LICENSE"}
 dependencies  = []
 

tests/api/test_distributed_queue.py

@@ -302,77 +302,6 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(queued["audio"]["sample_rate"], 44100)
         self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
 
-    async def test_job_complete_accepts_audio_without_image(self):
-        queue = asyncio.Queue()
-        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
-        job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
-        request = _FakeRequest(
-            {
-                "job_id": "audio-only-job",
-                "worker_id": "worker-1",
-                "batch_idx": 0,
-                "audio": self._encoded_audio_payload(),
-                "is_last": True,
-            }
-        )
-
-        with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
-            response = await job_routes.job_complete_endpoint(request)
-
-        self.assertEqual(response.status, 200)
-        decode_image.assert_not_called()
-        queued = await queue.get()
-        self.assertIsNone(queued["tensor"])
-        self.assertEqual(queued["audio"]["sample_rate"], 44100)
-
-    async def test_job_complete_rejects_payload_without_image_or_audio(self):
-        request = _FakeRequest(
-            {
-                "job_id": "job-1",
-                "worker_id": "worker-1",
-                "batch_idx": 0,
-                "is_last": True,
-            }
-        )
-
-        response = await job_routes.job_complete_endpoint(request)
-
-        self.assertEqual(response.status, 400)
-        self.assertIn("image or audio", response.payload.get("message", "").lower())
-
-    async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
-        request = _FakeRequest(
-            {
-                "job_id": "job-1",
-                "worker_id": "worker-1",
-                "batch_idx": 0,
-                "image": 123,
-                "audio": self._encoded_audio_payload(),
-                "is_last": True,
-            }
-        )
-
-        response = await job_routes.job_complete_endpoint(request)
-
-        self.assertEqual(response.status, 400)
-        self.assertIn("image", response.payload.get("message", "").lower())
-
-    async def test_job_complete_returns_400_for_malformed_audio(self):
-        request = _FakeRequest(
-            {
-                "job_id": "job-1",
-                "worker_id": "worker-1",
-                "batch_idx": 0,
-                "audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
-                "is_last": True,
-            }
-        )
-
-        response = await job_routes.job_complete_endpoint(request)
-
-        self.assertEqual(response.status, 400)
-        self.assertIn("audio.shape", response.payload.get("message", "").lower())
-
     def test_decode_audio_payload_rejects_bad_shape(self):
         bad = {
             "sample_rate": 44100,

tests/api/test_media_sync.py

@@ -164,11 +164,26 @@ class ConvertPathsForPlatformTests(unittest.TestCase):
 # ---------------------------------------------------------------------------
 
 class FindMediaReferencesTests(unittest.TestCase):
+    def test_finds_image_input(self):
+        prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
+        refs = ms._find_media_references(prompt)
+        self.assertIn("photo.png", refs)
+
+    def test_finds_video_input(self):
+        prompt = {"1": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}}
+        refs = ms._find_media_references(prompt)
+        self.assertIn("clip.mp4", refs)
+
     def test_finds_file_input_for_load_video(self):
         prompt = {"1": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}}
         refs = ms._find_media_references(prompt)
         self.assertIn("1 - Copy.mp4", refs)
 
+    def test_finds_audio_input(self):
+        prompt = {"1": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}}}
+        refs = ms._find_media_references(prompt)
+        self.assertIn("track.wav", refs)
+
     def test_strips_annotation_suffix(self):
         prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
         refs = ms._find_media_references(prompt)

tests/test_collector_list_inputs.py

@@ -1,364 +0,0 @@
-import importlib.util
-import sys
-import types
-import asyncio
-from pathlib import Path
-
-import torch
-
-
-def _load_collector_module():
-    module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
-    package_name = "dist_collector_list_testpkg"
-
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
-
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
-
-    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
-    nodes_pkg.__path__ = []
-    sys.modules[f"{package_name}.nodes"] = nodes_pkg
-
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
-
-    class _Routes:
-        def post(self, _path):
-            return lambda fn: fn
-
-        def get(self, _path):
-            return lambda fn: fn
-
-    prompt_server = types.SimpleNamespace(
-        routes=_Routes(),
-        distributed_jobs_lock=None,
-        distributed_pending_jobs={},
-    )
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
-    sys.modules["server"] = server_module
-
-    comfy_module = types.ModuleType("comfy")
-    model_management = types.ModuleType("comfy.model_management")
-
-    class InterruptProcessingException(Exception):
-        pass
-
-    model_management.InterruptProcessingException = InterruptProcessingException
-    model_management.throw_exception_if_processing_interrupted = lambda: None
-    comfy_module.model_management = model_management
-
-    comfy_utils = types.ModuleType("comfy.utils")
-
-    class ProgressBar:
-        def __init__(self, _total):
-            self.total = _total
-            self.updates = []
-
-        def update(self, value):
-            self.updates.append(value)
-
-    comfy_utils.ProgressBar = ProgressBar
-    comfy_module.utils = comfy_utils
-    sys.modules["comfy"] = comfy_module
-    sys.modules["comfy.model_management"] = model_management
-    sys.modules["comfy.utils"] = comfy_utils
-
-    aiohttp_module = types.ModuleType("aiohttp")
-    aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
-    sys.modules["aiohttp"] = aiohttp_module
-
-    logging_module = types.ModuleType(f"{package_name}.utils.logging")
-    logging_module.debug_log = lambda *_args, **_kwargs: None
-    logging_module.log = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.utils.logging"] = logging_module
-
-    config_module = types.ModuleType(f"{package_name}.utils.config")
-    config_module.get_worker_timeout_seconds = lambda: 0.1
-    config_module.load_config = lambda: {"workers": []}
-    config_module.is_master_delegate_only = lambda: False
-    sys.modules[f"{package_name}.utils.config"] = config_module
-
-    constants_module = types.ModuleType(f"{package_name}.utils.constants")
-    constants_module.HEARTBEAT_INTERVAL = 1.0
-    sys.modules[f"{package_name}.utils.constants"] = constants_module
-
-    image_module = types.ModuleType(f"{package_name}.utils.image")
-    def _ensure_contiguous(tensor):
-        return tensor.contiguous() if hasattr(tensor, "contiguous") else tensor
-
-    image_module.ensure_contiguous = _ensure_contiguous
-    image_module.tensor_to_pil = lambda *_args, **_kwargs: None
-    image_module.pil_to_tensor = lambda value: value
-    sys.modules[f"{package_name}.utils.image"] = image_module
-
-    network_module = types.ModuleType(f"{package_name}.utils.network")
-    network_module.build_worker_url = lambda worker: "http://worker"
-    network_module.get_client_session = lambda: None
-    network_module.probe_worker = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.utils.network"] = network_module
-
-    audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
-    audio_payload_module.encode_audio_payload = lambda audio: audio
-    sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
-
-    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
-    async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
-    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
-
-    spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", module_path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec is not None and spec.loader is not None
-    spec.loader.exec_module(module)
-    return module
-
-
-def test_collector_opts_into_comfyui_list_inputs():
-    collector = _load_collector_module().DistributedCollectorNode
-
-    assert collector.INPUT_IS_LIST is True
-
-
-def test_collector_exposes_images_as_optional_input():
-    input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
-
-    assert "images" not in input_types["required"]
-    assert input_types["optional"]["images"] == ("IMAGE",)
-
-
-def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
-    collector = _load_collector_module().DistributedCollectorNode()
-    audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
-
-    images, returned_audio = collector.run(images=None, audio=[audio])
-
-    assert images is None
-    assert returned_audio is audio
-
-
-def test_collector_rejects_missing_images_and_audio():
-    collector = _load_collector_module().DistributedCollectorNode()
-
-    try:
-        collector.run(images=None, audio=None)
-    except ValueError as exc:
-        assert "image or audio" in str(exc).lower()
-    else:
-        raise AssertionError("Expected collector to reject a run with no media input")
-
-
-def test_delegate_only_master_allows_no_local_media_input():
-    collector = _load_collector_module().DistributedCollectorNode()
-
-    images, audio = collector.run(
-        images=None,
-        audio=None,
-        multi_job_id=["delegate-audio-job"],
-        delegate_only=[True],
-        enabled_worker_ids=["[]"],
-    )
-
-    assert images is None
-    assert tuple(audio["waveform"].shape) == (1, 2, 1)
-
-
-def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
-    collector = _load_collector_module().DistributedCollectorNode()
-    first = torch.zeros(1, 2, 2, 3)
-    second = torch.ones(1, 2, 2, 3)
-
-    images, audio = collector.run(
-        images=[first, second],
-        load_balance=[False],
-        audio=[None],
-        multi_job_id=[""],
-        is_worker=[False],
-        master_url=[""],
-        enabled_worker_ids=["[]"],
-        worker_batch_size=[1],
-        worker_id=[""],
-        pass_through=[False],
-        delegate_only=[False],
-    )
-
-    assert tuple(images.shape) == (2, 2, 2, 3)
-    assert torch.equal(images[0:1], first)
-    assert torch.equal(images[1:2], second)
-    assert tuple(audio["waveform"].shape) == (1, 2, 1)
-
-
-def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
-    module = _load_collector_module()
-    collector = module.DistributedCollectorNode()
-    first = torch.zeros(1, 2, 2, 3)
-    second = torch.ones(1, 2, 2, 3)
-    posted_payloads = []
-
-    class _FakeImage:
-        def save(self, fp, format=None, compress_level=None):
-            fp.write(b"png-bytes")
-
-    class _FakeResponse:
-        async def __aenter__(self):
-            return self
-
-        async def __aexit__(self, exc_type, exc, tb):
-            return False
-
-        def raise_for_status(self):
-            return None
-
-    class _FakeSession:
-        def post(self, url, json, timeout):
-            posted_payloads.append(json)
-            return _FakeResponse()
-
-    async def _fake_get_client_session():
-        return _FakeSession()
-
-    module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
-    module.get_client_session = _fake_get_client_session
-    module.encode_audio_payload = lambda _audio: None
-
-    images, audio = collector.run(
-        images=[first, second],
-        load_balance=[False],
-        audio=[None],
-        multi_job_id=["job-list-1"],
-        is_worker=[True],
-        master_url=["http://master"],
-        enabled_worker_ids=["[]"],
-        worker_batch_size=[1],
-        worker_id=["worker-a"],
-        pass_through=[False],
-        delegate_only=[False],
-    )
-
-    assert tuple(images.shape) == (2, 2, 2, 3)
-    assert tuple(audio["waveform"].shape) == (1, 2, 1)
-    assert len(posted_payloads) == 2
-    assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
-    assert [payload["is_last"] for payload in posted_payloads] == [False, True]
-    assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
-    assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
-
-
-def test_audio_only_worker_sends_one_completion_without_image():
-    module = _load_collector_module()
-    collector = module.DistributedCollectorNode()
-    audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
-    posted = []
-
-    class _FakeResponse:
-        async def __aenter__(self):
-            return self
-
-        async def __aexit__(self, exc_type, exc, tb):
-            return False
-
-        def raise_for_status(self):
-            return None
-
-    class _FakeSession:
-        def post(self, url, json, timeout):
-            posted.append((url, json, timeout.total))
-            return _FakeResponse()
-
-    async def _fake_get_client_session():
-        return _FakeSession()
-
-    module.get_client_session = _fake_get_client_session
-    module.encode_audio_payload = lambda value: {"encoded": value is audio}
-
-    images, returned_audio = collector.run(
-        images=None,
-        audio=[audio],
-        multi_job_id=["audio-job"],
-        is_worker=[True],
-        master_url=["http://master"],
-        worker_id=["worker-a"],
-    )
-
-    assert images is None
-    assert returned_audio is audio
-    assert len(posted) == 1
-    assert posted[0][0] == "http://master/distributed/job_complete"
-    assert posted[0][2] == 600
-    assert posted[0][1] == {
-        "job_id": "audio-job",
-        "worker_id": "worker-a",
-        "batch_idx": 0,
-        "audio": {"encoded": True},
-        "is_last": True,
-    }
-
-
-def test_audio_only_master_combines_local_and_worker_audio():
-    module = _load_collector_module()
-    collector = module.DistributedCollectorNode()
-    master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
-    worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
-    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
-    queue = asyncio.Queue()
-    queue.put_nowait(
-        {
-            "worker_id": "worker-a",
-            "image_index": 0,
-            "tensor": None,
-            "audio": worker_audio,
-            "is_last": True,
-        }
-    )
-    module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
-
-    images, combined_audio = asyncio.run(
-        collector.execute(
-            images=None,
-            audio=master_audio,
-            multi_job_id="audio-job",
-            enabled_worker_ids='["worker-a"]',
-        )
-    )
-
-    assert images is None
-    assert combined_audio["sample_rate"] == 48000
-    assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
-    assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
-    assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
-
-
-def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
-    module = _load_collector_module()
-    collector = module.DistributedCollectorNode()
-    worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
-    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
-    queue = asyncio.Queue()
-    queue.put_nowait(
-        {
-            "worker_id": "worker-a",
-            "image_index": 0,
-            "tensor": None,
-            "audio": worker_audio,
-            "is_last": True,
-        }
-    )
-    module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
-
-    images, combined_audio = asyncio.run(
-        collector.execute(
-            images=None,
-            audio=None,
-            multi_job_id="delegate-audio-job",
-            enabled_worker_ids='["worker-a"]',
-            delegate_only=True,
-        )
-    )
-
-    assert images is None
-    assert combined_audio["sample_rate"] == 48000
-    assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])

tests/test_dispatch_selection.py

@@ -250,5 +250,63 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
         self.assertIsNone(selected)
 
 
+    async def test_select_least_busy_worker_waits_for_idle_when_required(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+        ]
+        queue_sequences = {
+            "w1": [1, 1],
+            "w2": [1, 0],
+        }
+
+        async def fake_probe(worker_url, timeout=3.0):
+            worker_id = worker_url.rsplit("/", 1)[-1]
+            sequence = queue_sequences[worker_id]
+            value = sequence.pop(0) if sequence else 0
+            return {"exec_info": {"queue_remaining": value}}
+
+        with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
+            dispatch,
+            "probe_worker",
+            side_effect=fake_probe,
+        ):
+            dispatch._least_busy_rr_index = 0
+            selected = await dispatch.select_least_busy_worker(
+                workers,
+                probe_concurrency=2,
+                require_idle=True,
+                idle_poll_interval=0,
+            )
+
+        self.assertEqual(selected["id"], "w2")
+
+
+    async def test_select_least_busy_worker_counts_reserved_slots_as_busy(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1", "reserved_slots": 1},
+            {"id": "w2", "name": "Worker 2"},
+        ]
+
+        async def fake_probe(worker_url, timeout=3.0):
+            return {"exec_info": {"queue_remaining": 0}}
+
+        with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
+            dispatch,
+            "probe_worker",
+            side_effect=fake_probe,
+        ):
+            dispatch._least_busy_rr_index = 0
+            selected = await dispatch.select_least_busy_worker(
+                workers,
+                probe_concurrency=2,
+                require_idle=True,
+                idle_poll_interval=0,
+                idle_wait_timeout=0,
+            )
+
+        self.assertEqual(selected["id"], "w2")
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_load_balance_queueing.py

@@ -0,0 +1,221 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+
+def _load_queue_orchestration_module():
+    module_path = Path(__file__).resolve().parents[1] / "api" / "queue_orchestration.py"
+    package_name = "dist_queue_orchestration_testpkg"
+
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = []
+    sys.modules[f"{package_name}.api"] = api_pkg
+
+    orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration")
+    orchestration_pkg.__path__ = []
+    sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    prompt_server_instance = types.SimpleNamespace(
+        distributed_pending_jobs={},
+        distributed_worker_reservations={},
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
+    sys.modules["server"] = server_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.queue_prompt_payload = AsyncMock(
+        return_value={"prompt_id": "prompt-master", "number": 1, "node_errors": {}}
+    )
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {
+        "workers": [{"id": "w1", "name": "Worker 1", "enabled": True, "host": "worker-1", "port": 8188}],
+        "settings": {"load_balance_idle_poll_interval_seconds": 0},
+    }
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 4
+    constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 5.0
+    constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8
+    constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.build_master_url = lambda *_args, **_kwargs: "http://master:8188"
+    network_module.build_master_callback_url = lambda *_args, **_kwargs: "http://master:8188"
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    trace_module = types.ModuleType(f"{package_name}.utils.trace_logger")
+    trace_module.trace_debug = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.trace_logger"] = trace_module
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+    schemas_module.parse_positive_int = lambda value, default: int(value if value is not None else default)
+    schemas_module.parse_positive_float = lambda value, default: float(value if value is not None else default)
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    dispatch_module = types.ModuleType(f"{package_name}.api.orchestration.dispatch")
+    dispatch_module.dispatch_worker_prompt = AsyncMock()
+    dispatch_module.select_active_workers = AsyncMock(return_value=([{"id": "w1", "name": "Worker 1"}], False))
+    dispatch_module.select_least_busy_worker = AsyncMock()
+    sys.modules[f"{package_name}.api.orchestration.dispatch"] = dispatch_module
+
+    media_sync_module = types.ModuleType(f"{package_name}.api.orchestration.media_sync")
+    media_sync_module.convert_paths_for_platform = lambda prompt, _sep: prompt
+    media_sync_module.fetch_worker_path_separator = AsyncMock(return_value=None)
+    media_sync_module.sync_worker_media = AsyncMock()
+    sys.modules[f"{package_name}.api.orchestration.media_sync"] = media_sync_module
+
+    prompt_transform_module = types.ModuleType(f"{package_name}.api.orchestration.prompt_transform")
+
+    class _PromptIndex:
+        def __init__(self, prompt):
+            self.prompt = prompt
+            self.inputs_by_node = {node_id: node.get("inputs", {}) for node_id, node in prompt.items()}
+
+        def nodes_for_class(self, class_type):
+            return [
+                node_id
+                for node_id, node in self.prompt.items()
+                if node.get("class_type") == class_type
+            ]
+
+        def copy_prompt(self):
+            return {node_id: dict(node) for node_id, node in self.prompt.items()}
+
+    prompt_transform_module.PromptIndex = _PromptIndex
+    prompt_transform_module.apply_participant_overrides = lambda prompt, *_args, **_kwargs: prompt
+    prompt_transform_module.find_nodes_by_class = lambda prompt, class_type: [
+        node_id for node_id, node in prompt.items() if node.get("class_type") == class_type
+    ]
+    prompt_transform_module.generate_job_id_map = lambda _prompt_index, _prefix: {"10": "job-10"}
+    prompt_transform_module.prepare_delegate_master_prompt = lambda prompt, _collector_ids: prompt
+    prompt_transform_module.prune_prompt_for_worker = lambda prompt: prompt
+    sys.modules[f"{package_name}.api.orchestration.prompt_transform"] = prompt_transform_module
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.api.queue_orchestration",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+queue_orchestration = _load_queue_orchestration_module()
+
+
+class LoadBalanceQueueingTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        queue_orchestration.ensure_distributed_state()
+        queue_orchestration.prompt_server.distributed_worker_reservations.clear()
+        queue_orchestration.prompt_server.distributed_pending_jobs.clear()
+
+    async def test_load_balance_reservation_selection_refreshes_reservations_between_polls(self):
+        queue_orchestration.prompt_server.distributed_worker_reservations["w1"] = 1
+        calls = []
+
+        async def fake_select(candidates, **_kwargs):
+            reserved_slots = candidates[0].get("reserved_slots")
+            calls.append(reserved_slots)
+            if reserved_slots:
+                await queue_orchestration._release_worker_reservation("w1")
+                return None
+            return candidates[0]
+
+        with patch.object(queue_orchestration, "select_least_busy_worker", side_effect=fake_select):
+            selected = await queue_orchestration._select_and_reserve_load_balance_worker(
+                [{"id": "w1", "name": "Worker 1"}],
+                "exec-test",
+                worker_probe_concurrency=1,
+                idle_poll_interval=0,
+            )
+
+        self.assertEqual(selected["id"], "w1")
+        self.assertEqual(calls, [1, 0])
+        self.assertEqual(queue_orchestration.prompt_server.distributed_worker_reservations, {"w1": 1})
+
+    async def test_load_balance_does_not_fallback_to_first_candidate_when_idle_selection_fails(self):
+        prompt = {
+            "10": {
+                "class_type": "DistributedCollector",
+                "inputs": {"load_balance": True},
+            }
+        }
+
+        with patch.object(
+            queue_orchestration,
+            "_select_and_reserve_load_balance_worker",
+            new=AsyncMock(return_value=None),
+        ), patch.object(
+            queue_orchestration,
+            "dispatch_worker_prompt",
+            new=AsyncMock(),
+        ) as dispatch_mock:
+            prompt_id, prompt_number, worker_count, node_errors = await queue_orchestration.orchestrate_distributed_execution(
+                prompt,
+                workflow_meta={},
+                client_id="client-1",
+                trace_execution_id="exec-test",
+            )
+
+        self.assertEqual((prompt_id, prompt_number, worker_count, node_errors), ("prompt-master", 1, 0, {}))
+        dispatch_mock.assert_not_awaited()
+
+    async def test_load_balance_releases_worker_reservation_when_dispatch_fails(self):
+        prompt = {
+            "10": {
+                "class_type": "DistributedCollector",
+                "inputs": {"load_balance": True},
+            }
+        }
+
+        async def select_first_candidate(candidates, **_kwargs):
+            return candidates[0]
+
+        with patch.object(
+            queue_orchestration,
+            "select_least_busy_worker",
+            side_effect=select_first_candidate,
+        ), patch.object(
+            queue_orchestration,
+            "dispatch_worker_prompt",
+            new=AsyncMock(side_effect=RuntimeError("dispatch failed")),
+        ):
+            with self.assertRaisesRegex(RuntimeError, "dispatch failed"):
+                await queue_orchestration.orchestrate_distributed_execution(
+                    prompt,
+                    workflow_meta={},
+                    client_id="client-1",
+                    trace_execution_id="exec-test",
+                )
+
+        self.assertEqual(queue_orchestration.prompt_server.distributed_worker_reservations, {})
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_payload_parsers.py

@@ -86,25 +86,27 @@ class ParseTilesFromFormTests(unittest.TestCase):
 
     # --- happy paths ---
 
-    def test_single_tile_returns_image_and_metadata(self):
+    def test_single_tile_returns_one_entry(self):
         tiles = pp._parse_tiles_from_form(_make_form(1))
         self.assertEqual(len(tiles), 1)
+
+    def test_multiple_tiles_all_returned(self):
+        tiles = pp._parse_tiles_from_form(_make_form(3))
+        self.assertEqual(len(tiles), 3)
+
+    def test_tile_image_is_pil_image(self):
+        tiles = pp._parse_tiles_from_form(_make_form(1))
+        self.assertIsInstance(tiles[0]["image"], PILImage.Image)
+
+    def test_tile_metadata_fields_are_parsed(self):
+        tiles = pp._parse_tiles_from_form(_make_form(1))
         tile = tiles[0]
-        self.assertIsInstance(tile["image"], PILImage.Image)
         self.assertEqual(tile["tile_idx"], 0)
         self.assertEqual(tile["x"], 0)
         self.assertEqual(tile["y"], 0)
         self.assertEqual(tile["extracted_width"], 64)
         self.assertEqual(tile["extracted_height"], 64)
 
-    def test_multiple_tiles_preserve_count_order_and_coordinates(self):
-        tiles = pp._parse_tiles_from_form(_make_form(3))
-        self.assertEqual(len(tiles), 3)
-        for i, tile in enumerate(tiles):
-            self.assertEqual(tile["tile_idx"], i)
-        self.assertEqual(tiles[1]["x"], 64)
-        self.assertEqual(tiles[2]["x"], 128)
-
     def test_padding_is_parsed_from_form(self):
         tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
         self.assertEqual(tiles[0]["padding"], 16)
@@ -136,6 +138,16 @@ class ParseTilesFromFormTests(unittest.TestCase):
         self.assertNotIn("batch_idx", tiles[0])
         self.assertNotIn("global_idx", tiles[0])
 
+    def test_tile_indices_match_metadata_order(self):
+        tiles = pp._parse_tiles_from_form(_make_form(3))
+        for i, tile in enumerate(tiles):
+            self.assertEqual(tile["tile_idx"], i)
+
+    def test_x_coordinates_reflect_metadata(self):
+        tiles = pp._parse_tiles_from_form(_make_form(3))
+        self.assertEqual(tiles[1]["x"], 64)
+        self.assertEqual(tiles[2]["x"], 128)
+
     # --- error cases ---
 
     def test_missing_tiles_metadata_raises_value_error(self):

tests/test_prompt_transform.py

@@ -71,15 +71,6 @@ def _collector_only_prompt():
     }
 
 
-def _audio_only_prompt():
-    """1(Audio source) → 2(DistributedCollector) → 3(Audio sink)."""
-    return {
-        "1": {"class_type": "LoadAudio", "inputs": {}},
-        "2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
-        "3": {"class_type": "SaveAudio", "inputs": {"audio": ["2", 1]}},
-    }
-
-
 def _delegate_prompt():
     """1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
     return {
@@ -226,35 +217,6 @@ class PrunePromptForWorkerTests(unittest.TestCase):
         self.assertEqual(len(preview_nodes), 1)
         self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
 
-    def test_injects_preview_audio_for_audio_only_collector(self):
-        result = pt.prune_prompt_for_worker(_audio_only_prompt())
-        preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewAudio"]
-        self.assertEqual(len(preview_nodes), 1)
-        self.assertEqual(preview_nodes[0]["inputs"]["audio"], ["2", 1])
-        self.assertFalse(any(n.get("class_type") == "PreviewImage" for n in result.values()))
-
-    def test_prefers_image_preview_when_collector_has_images_and_audio(self):
-        prompt = _linear_prompt()
-        prompt["6"] = {"class_type": "LoadAudio", "inputs": {}}
-        prompt["4"]["inputs"]["audio"] = ["6", 0]
-        result = pt.prune_prompt_for_worker(prompt)
-        self.assertEqual(
-            len([n for n in result.values() if n.get("class_type") == "PreviewImage"]),
-            1,
-        )
-        self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
-
-    def test_preserves_image_preview_for_distributed_upscale(self):
-        prompt = {
-            "1": {"class_type": "LoadImage", "inputs": {}},
-            "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
-            "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
-        }
-        result = pt.prune_prompt_for_worker(prompt)
-        preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
-        self.assertEqual(len(preview_nodes), 1)
-        self.assertEqual(preview_nodes[0]["inputs"]["images"], ["2", 0])
-
     def test_no_preview_image_when_no_downstream(self):
         result = pt.prune_prompt_for_worker(_collector_only_prompt())
         preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
@@ -280,7 +242,7 @@ class PrunePromptForWorkerTests(unittest.TestCase):
     def test_upscale_node_is_treated_as_distributed(self):
         prompt = {
             "1": {"class_type": "KSampler", "inputs": {}},
-            "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
+            "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}},
             "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
         }
         result = pt.prune_prompt_for_worker(prompt)
@@ -302,8 +264,8 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
         self.assertNotIn("1", result)
         self.assertNotIn("2", result)
 
-    def test_replaces_dangling_upstream_ref_with_one_empty_image_placeholder(self):
-        """Collector must point to exactly one valid placeholder after pruning."""
+    def test_removes_dangling_upstream_refs(self):
+        """Collector must not retain dangling refs to pruned upstream nodes."""
         prompt = _delegate_prompt()
         result = pt.prepare_delegate_master_prompt(prompt, ["3"])
         collector_inputs = result["3"].get("inputs", {})
@@ -314,520 +276,26 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
         self.assertNotEqual(source_id, "2")
         self.assertIn(source_id, result)
         self.assertEqual(result[source_id].get("class_type"), "DistributedEmptyImage")
+
+    def test_injects_empty_image_placeholder(self):
+        prompt = _delegate_prompt()
+        result = pt.prepare_delegate_master_prompt(prompt, ["3"])
         empty_nodes = [(nid, n) for nid, n in result.items() if n.get("class_type") == "DistributedEmptyImage"]
         self.assertEqual(len(empty_nodes), 1)
         placeholder_id = empty_nodes[0][0]
         self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
 
-    def test_audio_only_collector_does_not_get_image_placeholder(self):
-        prompt = _audio_only_prompt()
-        result = pt.prepare_delegate_master_prompt(prompt, ["2"])
-        empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
-        self.assertEqual(empty_nodes, [])
-        self.assertNotIn("images", result["2"].get("inputs", {}))
-        self.assertNotIn("audio", result["2"].get("inputs", {}))
-
     def test_one_placeholder_per_collector(self):
         """Two collectors → two placeholders."""
         prompt = {
-            "1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0]}},
-            "2": {"class_type": "DistributedCollector", "inputs": {"images": ["11", 0]}},
+            "1": {"class_type": "DistributedCollector", "inputs": {}},
+            "2": {"class_type": "DistributedCollector", "inputs": {}},
             "3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
-            "10": {"class_type": "LoadImage", "inputs": {}},
-            "11": {"class_type": "LoadImage", "inputs": {}},
         }
         result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
         empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
         self.assertEqual(len(empty_nodes), 2)
 
-    def test_preserves_primitive_string_for_downstream_required_input(self):
-        """Delegate-only master keeps primitive inputs needed by SaveImage."""
-        save_image = type(
-            "SaveImage",
-            (),
-            {
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "images": ("IMAGE",),
-                            "filename_prefix": ("STRING",),
-                        }
-                    }
-                )
-            },
-        )
-        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
-        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {"SaveImage": save_image}
-        try:
-            prompt = {
-                "8": {"class_type": "KSampler", "inputs": {}},
-                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
-                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test/input_bug"}},
-                "9": {
-                    "class_type": "SaveImage",
-                    "inputs": {
-                        "images": ["11", 0],
-                        "filename_prefix": ["15", 0],
-                    },
-                },
-            }
-
-            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
-        finally:
-            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
-
-        self.assertIn("15", result)
-        self.assertEqual(result["15"], prompt["15"])
-        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["15", 0])
-
-    def test_preserves_load_image_for_switch_alternate_required_input(self):
-        """Delegate-only master keeps LoadImage inputs needed by switches."""
-        prompt = {
-            "848": {"class_type": "LoadImage", "inputs": {"image": "input_bug2_00003_.png"}},
-            "862": {"class_type": "VAEDecode", "inputs": {}},
-            "854": {"class_type": "DistributedCollector", "inputs": {"images": ["862", 0]}},
-            "850": {
-                "class_type": "ComfySwitchNode",
-                "inputs": {
-                    "on_false": ["848", 0],
-                    "on_true": ["854", 0],
-                },
-            },
-            "851": {"class_type": "PreviewImage", "inputs": {"images": ["850", 0]}},
-        }
-
-        result = pt.prepare_delegate_master_prompt(prompt, ["854"])
-
-        self.assertIn("848", result)
-        self.assertEqual(result["850"]["inputs"]["on_false"], ["848", 0])
-        self.assertEqual(result["850"]["inputs"]["on_true"], ["854", 0])
-        self.assertNotIn("862", result)
-
-    def test_preserves_registered_string_utility_subgraph_for_downstream_required_input(self):
-        """Delegate-only master keeps scalar utility chains used by SaveImage."""
-        string_concat = type(
-            "StringConcatenate",
-            (),
-            {
-                "RETURN_TYPES": ("STRING",),
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "string_a": ("STRING",),
-                            "string_b": ("STRING",),
-                        }
-                    }
-                ),
-            },
-        )
-        save_image = type(
-            "SaveImage",
-            (),
-            {
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "images": ("IMAGE",),
-                            "filename_prefix": ("STRING",),
-                        }
-                    }
-                )
-            },
-        )
-        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
-        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
-            "SaveImage": save_image,
-            "StringConcatenate": string_concat,
-        }
-        try:
-            prompt = {
-                "8": {"class_type": "KSampler", "inputs": {}},
-                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
-                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
-                "16": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug3"}},
-                "17": {
-                    "class_type": "StringConcatenate",
-                    "inputs": {
-                        "string_a": ["15", 0],
-                        "string_b": ["16", 0],
-                    },
-                },
-                "9": {
-                    "class_type": "SaveImage",
-                    "inputs": {
-                        "images": ["11", 0],
-                        "filename_prefix": ["17", 0],
-                    },
-                },
-            }
-
-            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
-        finally:
-            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
-
-        self.assertIn("15", result)
-        self.assertIn("16", result)
-        self.assertIn("17", result)
-        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["17", 0])
-        self.assertEqual(result["17"]["inputs"]["string_a"], ["15", 0])
-        self.assertEqual(result["17"]["inputs"]["string_b"], ["16", 0])
-
-    def test_preserves_registered_multi_string_join_subgraph_for_downstream_required_input(self):
-        """Delegate-only master keeps multi-input scalar utility chains."""
-        join_string_multi = type(
-            "JoinStringMulti",
-            (),
-            {
-                "RETURN_TYPES": ("STRING",),
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {"string_1": ("STRING",)},
-                        "optional": {
-                            "string_2": ("STRING",),
-                            "string_3": ("STRING",),
-                        },
-                    }
-                ),
-            },
-        )
-        save_image = type(
-            "SaveImage",
-            (),
-            {
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "images": ("IMAGE",),
-                            "filename_prefix": ("STRING",),
-                        }
-                    }
-                )
-            },
-        )
-        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
-        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
-            "JoinStringMulti": join_string_multi,
-            "SaveImage": save_image,
-        }
-        try:
-            prompt = {
-                "8": {"class_type": "KSampler", "inputs": {}},
-                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
-                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
-                "16": {"class_type": "PrimitiveString", "inputs": {"value": "input"}},
-                "17": {"class_type": "PrimitiveString", "inputs": {"value": "bug4"}},
-                "18": {
-                    "class_type": "JoinStringMulti",
-                    "inputs": {
-                        "string_1": ["15", 0],
-                        "string_2": ["16", 0],
-                        "string_3": ["17", 0],
-                    },
-                },
-                "9": {
-                    "class_type": "SaveImage",
-                    "inputs": {
-                        "images": ["11", 0],
-                        "filename_prefix": ["18", 0],
-                    },
-                },
-            }
-
-            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
-        finally:
-            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
-
-        self.assertIn("15", result)
-        self.assertIn("16", result)
-        self.assertIn("17", result)
-        self.assertIn("18", result)
-        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["18", 0])
-
-    def test_does_not_preserve_scalar_utility_with_heavy_upstream_dependency(self):
-        """Scalar utility nodes are retained only when their full input branch is safe."""
-        string_concat = type(
-            "StringConcatenate",
-            (),
-            {
-                "RETURN_TYPES": ("STRING",),
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "string_a": ("STRING",),
-                            "string_b": ("STRING",),
-                        }
-                    }
-                ),
-            },
-        )
-        save_image = type(
-            "SaveImage",
-            (),
-            {
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "images": ("IMAGE",),
-                            "filename_prefix": ("STRING",),
-                        }
-                    }
-                )
-            },
-        )
-        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
-        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
-            "SaveImage": save_image,
-            "StringConcatenate": string_concat,
-        }
-        try:
-            prompt = {
-                "8": {"class_type": "KSampler", "inputs": {}},
-                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
-                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
-                "17": {
-                    "class_type": "StringConcatenate",
-                    "inputs": {
-                        "string_a": ["15", 0],
-                        "string_b": ["8", 0],
-                    },
-                },
-                "9": {
-                    "class_type": "SaveImage",
-                    "inputs": {
-                        "images": ["11", 0],
-                        "filename_prefix": ["17", 0],
-                    },
-                },
-            }
-
-            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
-        finally:
-            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
-
-        self.assertNotIn("8", result)
-        self.assertNotIn("17", result)
-        self.assertNotIn("filename_prefix", result["9"]["inputs"])
-
-    def test_does_not_preserve_scalar_output_for_non_scalar_downstream_input(self):
-        """Scalar outputs are retained only for scalar/config downstream inputs."""
-        string_provider = type("StringProvider", (), {"RETURN_TYPES": ("STRING",)})
-        image_consumer = type(
-            "ImageConsumer",
-            (),
-            {
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "images": ("IMAGE",),
-                            "mask": ("IMAGE",),
-                        }
-                    }
-                )
-            },
-        )
-        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
-        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
-            "ImageConsumer": image_consumer,
-            "StringProvider": string_provider,
-        }
-        try:
-            prompt = {
-                "8": {"class_type": "KSampler", "inputs": {}},
-                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
-                "17": {"class_type": "StringProvider", "inputs": {}},
-                "9": {
-                    "class_type": "ImageConsumer",
-                    "inputs": {
-                        "images": ["11", 0],
-                        "mask": ["17", 0],
-                    },
-                },
-            }
-
-            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
-        finally:
-            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
-
-        self.assertNotIn("17", result)
-        self.assertNotIn("mask", result["9"]["inputs"])
-
-    def test_preserves_scalar_list_join_subgraph_for_downstream_required_input(self):
-        """Delegate-only master keeps list-shaped scalar config chains."""
-        create_list = type(
-            "CreateList",
-            (),
-            {
-                "RETURN_TYPES": ("LIST",),
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {"inputs.input0": ("STRING",)},
-                        "optional": {
-                            "inputs.input1": ("STRING",),
-                            "inputs.input2": ("STRING",),
-                        },
-                    }
-                ),
-            },
-        )
-        string_data_list_join = type(
-            "StringDataListJoin",
-            (),
-            {
-                "RETURN_TYPES": ("STRING",),
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "strings": ("LIST",),
-                            "delimiter": ("STRING",),
-                        }
-                    }
-                ),
-            },
-        )
-        save_image = type(
-            "SaveImage",
-            (),
-            {
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "images": ("IMAGE",),
-                            "filename_prefix": ("STRING",),
-                        }
-                    }
-                )
-            },
-        )
-        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
-        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
-            "CreateList": create_list,
-            "SaveImage": save_image,
-            "StringDataListJoin": string_data_list_join,
-        }
-        try:
-            prompt = {
-                "8": {"class_type": "KSampler", "inputs": {}},
-                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
-                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
-                "17": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
-                "29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
-                "28": {
-                    "class_type": "CreateList",
-                    "inputs": {
-                        "inputs.input0": ["15", 0],
-                        "inputs.input1": ["17", 0],
-                        "inputs.input2": ["29", 0],
-                    },
-                },
-                "32": {
-                    "class_type": "StringDataListJoin",
-                    "inputs": {
-                        "strings": ["28", 0],
-                        "delimiter": "/",
-                    },
-                },
-                "9": {
-                    "class_type": "SaveImage",
-                    "inputs": {
-                        "images": ["11", 0],
-                        "filename_prefix": ["32", 0],
-                    },
-                },
-            }
-
-            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
-        finally:
-            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
-
-        self.assertIn("15", result)
-        self.assertIn("17", result)
-        self.assertIn("28", result)
-        self.assertIn("29", result)
-        self.assertIn("32", result)
-        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
-        self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
-        self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["15", 0])
-        self.assertNotIn("8", result)
-
-    def test_preserves_builtin_create_list_string_data_list_join_subgraph(self):
-        """Delegate-only master handles ComfyUI 0.23 CreateList data-list output."""
-        string_data_list_join = type(
-            "StringDataListJoin",
-            (),
-            {
-                "RETURN_TYPES": ("STRING",),
-                "INPUT_IS_LIST": True,
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "strings": ("STRING", {"forceInput": True}),
-                            "sep": ("STRING", {"default": " "}),
-                        }
-                    }
-                ),
-            },
-        )
-        save_image = type(
-            "SaveImage",
-            (),
-            {
-                "INPUT_TYPES": classmethod(
-                    lambda cls: {
-                        "required": {
-                            "images": ("IMAGE",),
-                            "filename_prefix": ("STRING",),
-                        }
-                    }
-                )
-            },
-        )
-        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
-        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
-            "Basic data handling: StringDataListJoin": string_data_list_join,
-            "SaveImage": save_image,
-        }
-        try:
-            prompt = {
-                "8": {"class_type": "KSampler", "inputs": {}},
-                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
-                "15": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
-                "17": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
-                "29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
-                "28": {
-                    "class_type": "CreateList",
-                    "inputs": {
-                        "inputs.input0": ["17", 0],
-                        "inputs.input1": ["15", 0],
-                        "inputs.input2": ["29", 0],
-                    },
-                },
-                "32": {
-                    "class_type": "Basic data handling: StringDataListJoin",
-                    "inputs": {
-                        "strings": ["28", 0],
-                        "sep": "/",
-                    },
-                },
-                "9": {
-                    "class_type": "SaveImage",
-                    "inputs": {
-                        "images": ["11", 0],
-                        "filename_prefix": ["32", 0],
-                    },
-                },
-            }
-
-            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
-        finally:
-            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
-
-        for node_id in ("15", "17", "28", "29", "32"):
-            self.assertIn(node_id, result)
-        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
-        self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
-        self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["17", 0])
-        self.assertNotIn("8", result)
-
     def test_result_is_independent_copy(self):
         prompt = _delegate_prompt()
         result = pt.prepare_delegate_master_prompt(prompt, ["3"])

web/sidebar/settingsSection.js

@@ -72,7 +72,7 @@ export function renderSettingsSection(extension) {
         createCheckboxSetting(
             "setting-debug",
             "Debug Mode",
-            "Enable verbose logging in the browser console and ComfyUI server output.",
+            "Enable verbose logging in the browser console.",
             extension.config?.settings?.debug || false,
             (event) => extension._updateSetting("debug", event.target.checked)
         )
@@ -100,7 +100,7 @@ export function renderSettingsSection(extension) {
         createNumberSetting(
             "setting-worker-timeout",
             "Worker Timeout",
-            "Maximum result-wait and heartbeat inactivity period before recovery begins. Busy workers may receive additional grace. Default: 60 seconds.",
+            "Seconds without a heartbeat before a worker is considered timed out. Default 60.",
             extension.config?.settings?.worker_timeout_seconds ?? 60,
             10,
             1,

web/tests/executionUtils.test.js

@@ -0,0 +1,15 @@
+import { describe, expect, it } from "vitest";
+
+import { buildWorkerWebSocketUrl } from "../urlUtils.js";
+
+
+describe("execution decision helpers", () => {
+    it("buildWorkerWebSocketUrl converts http/https to ws/wss", () => {
+        expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
+            "ws://worker.local:8188/distributed/worker_ws"
+        );
+        expect(buildWorkerWebSocketUrl("https://worker.example.com")).toBe(
+            "wss://worker.example.com/distributed/worker_ws"
+        );
+    });
+});

web/tests/urlUtils.test.js

@@ -106,6 +106,11 @@ describe("buildWorkerWebSocketUrl", () => {
             "wss://worker.example.com/distributed/worker_ws"
         );
     });
+
+    it("always appends /distributed/worker_ws", () => {
+        const url = buildWorkerWebSocketUrl("http://worker.local:8188");
+        expect(url.endsWith("/distributed/worker_ws")).toBe(true);
+    });
 });
 
 

web/tests/workerLifecycle.test.js

@@ -3,7 +3,7 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest";
 import { getWorkerUrl } from "../workerLifecycle.js";
 
 
-describe("workerLifecycle URL wiring", () => {
+describe("workerLifecycle URL construction", () => {
     let originalWindow;
 
     beforeEach(() => {
@@ -22,9 +22,31 @@ describe("workerLifecycle URL wiring", () => {
         globalThis.window = originalWindow;
     });
 
-    // URL variants are covered in urlUtils.test.js; keep the wrapper wiring here.
-    it("forwards the worker and endpoint using window.location", () => {
+    it("builds local worker URL with explicit local port", () => {
         const worker = { id: "w1", port: 8189, type: "local" };
         expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://127.0.0.1:8189/prompt");
     });
+
+    it("builds remote worker URL with host:port", () => {
+        const worker = { id: "w2", host: "worker.example.com", port: 9000, type: "remote" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://worker.example.com:9000/prompt");
+    });
+
+    it("builds cloud worker URL as https", () => {
+        const worker = { id: "w3", host: "cloud.example.com", port: 443, type: "cloud" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://cloud.example.com/prompt");
+    });
+
+    it("rewrites runpod proxy hostname for local worker ports", () => {
+        globalThis.window = {
+            location: {
+                hostname: "podabc.proxy.runpod.net",
+                protocol: "https:",
+                port: "",
+                origin: "https://podabc.proxy.runpod.net",
+            },
+        };
+        const worker = { id: "w4", port: 8189, type: "local" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://podabc-8189.proxy.runpod.net/prompt");
+    });
 });

workflows/distributed-upscale.json

@@ -814,7 +814,7 @@
       "outputs": [],
       "properties": {},
       "widgets_values": [
-        "Workers pull tiles from a shared queue, so faster workers can process more tiles automatically.\n"
+        "If all your GPUs are the same/similar, set static_distribution to true\n"
       ],
       "color": "#432",
       "bgcolor": "#653"