Refactor: reduce duplication and simplify arg passing across USDU modules

- Pass UpscaleCoreArgs through static mode instead of 11 individual args
- Extract _post_with_retry helper from duplicate retry loops in worker_comms
- Remove 4 passthrough wrapper pairs in job_store
- Extract _process_image_tiles from duplicate master/worker tile loops
- Merge 4 override functions into 2 in prompt_transform (seed+value, collector+upscale)
- Split _async_collect_results into focused static/dynamic methods

Net reduction of ~256 lines across 10 files.

.gitignore

@@ -0,0 +1,14 @@
+bin/
+logs/
+gpu_config.json
+__pycache__/
+**/__pycache__/
+*.py[cod]
+node_modules/
+npm-debug.log*
+AGENTS.md
+.claude/
+
+# Desloppify artifacts
+.desloppify/
+scorecard.png

.nvmrc

@@ -0,0 +1 @@
+20

README.md

@@ -5,7 +5,7 @@
 <a href="/docs/worker-setup-guides.md"><img src="https://img.shields.io/badge/Setup_Guides-grey?style=flat&logo=gitbook&logoColor=white" alt="Setup Guides"></a>
 <a href="/workflows"><img src="https://img.shields.io/badge/Workflows-grey?style=flat&logo=json&logoColor=white" alt="Workflows"></a>
 <a href="https://buymeacoffee.com/robertvoy"><img src="https://img.shields.io/badge/Donation-grey?style=flat&logo=buymeacoffee&logoColor=white" alt="Donation"></a>
-<a href="https://x.com/robertvoy3D"><img src="https://img.shields.io/twitter/follow/robertvoy3D" alt="Twitter"></a>
+<a href="https://x.com/rbw_ai"><img src="https://img.shields.io/twitter/follow/rbw_ai" alt="Twitter"></a>
 <br><br>
 </div>
 
@@ -15,7 +15,7 @@
 
 ---
 
-## Key Features
+## Key Features
 
 #### Parallel Workflow Processing
 - Run your workflow on multiple GPUs simultaneously with varied seeds, collect results on the master
@@ -25,14 +25,24 @@
 #### Distributed Upscaling
 - Accelerate Ultimate SD Upscale by distributing tiles across GPUs
 - Intelligent distribution
-- Handles single images and batches
+- Handles single images and videos
 
-#### Ease of Use
-- Auto-setup local workers; easily add remote/cloud ones
-- Convert any workflow to distributed with 2 nodes
-- JSON configuration with UI controls
-
----
+#### Ease of Use
+- Auto-setup local workers; easily add remote/cloud ones
+- Convert any workflow to distributed with 2 nodes
+- JSON configuration with UI controls
+
+---
+
+## Current Architecture
+
+- Workflow-level load balancing is controlled by **Distributed Collector** via the `load_balance` toggle.
+- There is **no Distributed Queue node** anymore.
+- With `load_balance=true`, orchestration selects one least-busy execution participant:
+  - If master participation is enabled, master is included as a candidate.
+  - If master is in orchestrator-only mode, only workers are considered.
+
+---
 
 ## Worker Types
 
@@ -82,18 +92,19 @@ Join Runpod with [this link](https://get.runpod.io/0bw29uf3ug0p) and unlock a sp
 
 ## Workflow Examples
 
-### Basic Parallel Generation
-Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
+### Basic Parallel Generation
+Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
 
 ![Clipboard Image (6)](https://github.com/user-attachments/assets/9598c94c-d9b4-4ccf-ab16-a21398220aeb)
 
 > [Download workflow](/workflows/distributed-txt2img.json)
 
 1. Open your ComfyUI workflow
-2. Add **Distributed Seed** → connect to sampler's seed
-3. Add **Distributed Collector** → after VAE Decode
-4. Enable workers in the UI
-5. Run the workflow!
+2. Add **Distributed Seed** → connect to sampler's seed
+3. Add **Distributed Collector** → after VAE Decode
+4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
+5. Enable workers in the UI
+6. Run the workflow!
 
 ### Parallel WAN Generation
 Generate multiple videos in the time it takes to generate one. Each worker uses a different seed.
@@ -111,12 +122,12 @@ Generate multiple videos in the time it takes to generate one. Each worker uses
 7. Enable workers in the UI
 8. Run the workflow!
 
-### Distributed Image Upscaling
-Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
+### Distributed Image Upscaling
+Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
 
 ![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a)
 
-> [Download workflow](/workflows/distributed-txt2img.json)
+> [Download workflow](/workflows/distributed-upscale.json)
 
 1. Load your image
 2. Upscale with ESRGAN or similar
@@ -142,6 +153,45 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
 
 ---
 
+## Developer API
+
+Control your distributed cluster programmatically without opening the browser.
+
+* **Endpoint:** `POST /distributed/queue`
+* **Functionality:** Accepts a distributed queue payload (`prompt`, `client_id`, `enabled_worker_ids`, optional `workflow`/`delegate_master`/`trace_execution_id`), dispatches to healthy workers, and returns `{prompt_id, worker_count}`.
+* **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).
+
+---
+
+## Distributed Value
+
+Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
+
+- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
+- The node shows only currently enabled workers.
+- If worker enablement changes, worker fields update automatically.
+- When disconnected, it resets to default string mode and clears per-worker overrides.
+- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
+
+---
+
+## Nodes
+
+| Node | Description |
+|------|-------------|
+| **Distributed Seed** | Generates unique seeds for each worker |
+| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
+| **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 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 |
+
+---
+
 ## FAQ
 
 <details>
@@ -161,7 +211,7 @@ Currently, it is not compatible with the ComfyUI desktop app.
 
 <details>
 <summary>Can I combine my RTX 5090 with a GTX 980 to get faster results?</summary>
-Yes, you can combine different GPUs, but performance is optimized when using similar GPUs. A significant performance imbalance between GPUs may cause bottlenecks. For upscaling, setting `static_distribution` to `false` allows the faster GPU to handle more processing, which can mitigate some bottlenecks. Note that this setting only applies to upscaling tasks.
+Yes, you can combine different GPUs, but performance is optimized when using similar GPUs. A significant performance imbalance between GPUs may cause bottlenecks.
 </details>
 
 <details>
@@ -169,6 +219,11 @@ Yes, you can combine different GPUs, but performance is optimized when using sim
 Yes, it is compatible with cloud providers. Refer to the setup guides for detailed instructions.
 </details>
 
+<details>
+<summary>Can I use my main machine just to coordinate workers without rendering?</summary>
+Yes. Open the Distributed panel and uncheck the master toggle to run in orchestrator-only mode. The master will distribute work to workers but won't render locally. If all workers become unavailable, the master automatically re-enables to ensure your workflow still runs.
+</details>
+
 <details>
 <summary>Can I make this work with my Docker setup?</summary>
 Yes, it is compatible with Docker setups, but you will need to configure your Docker environment yourself. Unfortunately, assistance with Docker configuration is not provided.
@@ -194,8 +249,3 @@ Buy me a coffee at: https://buymeacoffee.com/robertvoy
 
 
 
-
-
-
-
-

__init__.py

@@ -1,65 +1,11 @@
-import os
-import sys
+"""ComfyUI-Distributed package entrypoint."""
+from __future__ import annotations
 
-# Add the directory to Python path
-sys.path.append(os.path.dirname(__file__))
-
-# Patch ComfyUI execution validation for DistributedVideoCollector
-try:
-    import execution
-    from .distributed import ImageBatchDivider
-    
-    # Store original validate function if it exists
-    if hasattr(execution, 'validate_outputs'):
-        original_validate_outputs = execution.validate_outputs
-        
-        def patched_validate_outputs(executor, node_id, result, node_class):
-            if node_class == ImageBatchDivider:
-                return  # Skip validation for our dynamic output node
-            return original_validate_outputs(executor, node_id, result, node_class)
-        
-        execution.validate_outputs = patched_validate_outputs
-    else:
-        # Fallback: patch the main execute method if validate_outputs doesn't exist
-        if hasattr(execution.PromptExecutor, 'execute'):
-            original_execute = execution.PromptExecutor.execute
-            
-            def patched_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
-                # This is a more complex patch - for now just call original
-                # The ByPassTypeTuple should handle most validation issues
-                return original_execute(self, prompt, prompt_id, extra_data, execute_outputs)
-            
-            execution.PromptExecutor.execute = patched_execute
-            
-except ImportError:
-    pass  # ComfyUI execution module not available during import
-
-# Import everything needed from the main module
-from .distributed import (
-    NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS, 
-    NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
-)
-
-# Import utilities
-from .utils.config import ensure_config_exists, CONFIG_FILE
-from .utils.logging import debug_log
-
-# Import distributed upscale nodes
-from .distributed_upscale import (
-    NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
-    NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
-)
+from .bootstrap.entrypoint import build_node_mappings, initialize_runtime
 
 WEB_DIRECTORY = "./web"
 
-ensure_config_exists()
+NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = build_node_mappings()
+initialize_runtime()
 
-# Merge node mappings
-NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
-NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
-
-__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
-
-debug_log("Loaded Distributed nodes.")
-debug_log(f"Config file: {CONFIG_FILE}")
-debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

api/__init__.py

@@ -0,0 +1,21 @@
+"""Distributed API package with explicit route bootstrap."""
+from __future__ import annotations
+
+from importlib import import_module
+
+_ROUTE_MODULES = (
+    "config_routes",
+    "tunnel_routes",
+    "worker_routes",
+    "job_routes",
+    "usdu_routes",
+)
+
+
+def bootstrap_routes() -> None:
+    """Import route modules once so aiohttp decorators register endpoints."""
+    for module_name in _ROUTE_MODULES:
+        import_module(f"{__name__}.{module_name}")
+
+
+__all__ = ["bootstrap_routes"]

api/config_routes.py

@@ -0,0 +1,295 @@
+from __future__ import annotations
+
+from aiohttp import web
+import server
+
+from ..utils.config import config_transaction, load_config
+from ..utils.logging import debug_log
+from ..utils.network import handle_api_error, normalize_host
+from .endpoint_policy import run_authorized_endpoint
+
+
+def _positive_int(value: int) -> bool:
+    return value > 0
+
+
+CONFIG_SCHEMA = {
+    "workers": (list, None),
+    "master": (dict, None),
+    "settings": (dict, None),
+    "tunnel": (dict, None),
+    "managed_processes": (dict, None),
+    "worker_timeout_seconds": (int, _positive_int),
+    "debug": (bool, None),
+    "auto_launch_workers": (bool, None),
+    "stop_workers_on_master_exit": (bool, None),
+    "master_delegate_only": (bool, None),
+    "websocket_orchestration": (bool, None),
+    "has_auto_populated_workers": (bool, None),
+}
+
+_SETTINGS_FIELDS = {
+    "worker_timeout_seconds",
+    "debug",
+    "auto_launch_workers",
+    "stop_workers_on_master_exit",
+    "master_delegate_only",
+    "websocket_orchestration",
+    "has_auto_populated_workers",
+}
+
+_WORKER_FIELDS = [
+    ("enabled", None, False),
+    ("name", None, False),
+    ("port", None, False),
+    ("host", normalize_host, True),
+    ("cuda_device", None, True),
+    ("extra_args", None, True),
+    ("type", None, False),
+]
+
+_MASTER_FIELDS = [
+    ("name", None, False),
+    ("host", normalize_host, True),
+    ("port", None, False),
+    ("cuda_device", None, False),
+    ("extra_args", None, False),
+]
+
+
+def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
+    """Apply a partial update to a target dict based on field rules."""
+    for key, normalizer, remove_on_none in field_rules:
+        if key not in data:
+            continue
+        value = data[key]
+        if value is None and remove_on_none:
+            target.pop(key, None)
+        else:
+            target[key] = normalizer(value) if (normalizer and value is not None) else value
+
+
+def _validate_workers_payload(workers_value: object) -> list[str]:
+    """Validate top-level workers payload shape before persisting config."""
+    if not isinstance(workers_value, list):
+        return ["workers: expected list"]
+
+    errors: list[str] = []
+    seen_ids: set[str] = set()
+    for index, worker in enumerate(workers_value):
+        if not isinstance(worker, dict):
+            errors.append(f"workers[{index}]: expected object")
+            continue
+
+        worker_id = str(worker.get("id", "")).strip()
+        if not worker_id:
+            errors.append(f"workers[{index}].id: expected non-empty string")
+        elif worker_id in seen_ids:
+            errors.append(f"workers[{index}].id: duplicate id '{worker_id}'")
+        else:
+            seen_ids.add(worker_id)
+
+        if "port" in worker and worker.get("port") is not None:
+            try:
+                port = int(worker.get("port"))
+            except (TypeError, ValueError):
+                errors.append(f"workers[{index}].port: expected integer")
+            else:
+                if port <= 0:
+                    errors.append(f"workers[{index}].port: expected positive integer")
+
+    return errors
+
+
+@server.PromptServer.instance.routes.get("/distributed/config")
+async def get_config_endpoint(request: web.Request) -> web.StreamResponse:
+    async def _operation() -> web.StreamResponse:
+        return web.json_response(load_config())
+
+    return await run_authorized_endpoint(request, _operation)
+
+
+@server.PromptServer.instance.routes.post("/distributed/config")
+async def update_config_endpoint(request: web.Request) -> web.StreamResponse:
+    """Bulk config update with schema validation."""
+    async def _operation() -> web.StreamResponse:
+        try:
+            data = await request.json()
+        except Exception as exc:
+            return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
+
+        if not isinstance(data, dict):
+            return await handle_api_error(request, "Config payload must be an object", 400)
+
+        validated_settings = {}
+        validated_root = {}
+        errors = []
+
+        for key, value in data.items():
+            if key not in CONFIG_SCHEMA:
+                errors.append(f"Unknown field: {key}")
+                continue
+
+            expected_type, validator = CONFIG_SCHEMA[key]
+            if not isinstance(value, expected_type):
+                errors.append(f"{key}: expected {expected_type.__name__}")
+                continue
+
+            if validator and not validator(value):
+                errors.append(f"{key}: value {value!r} failed validation")
+                continue
+
+            if key == "workers":
+                errors.extend(_validate_workers_payload(value))
+                if errors:
+                    continue
+
+            if key in _SETTINGS_FIELDS:
+                validated_settings[key] = value
+            else:
+                validated_root[key] = value
+
+        if errors:
+            return await handle_api_error(request, errors, 400)
+
+        async with config_transaction() as config:
+            settings = config.setdefault("settings", {})
+            settings.update(validated_settings)
+            for key, value in validated_root.items():
+                config[key] = value
+            return web.json_response({"status": "success", "config": config})
+
+    return await run_authorized_endpoint(request, _operation)
+
+
+@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
+async def queue_status_endpoint(request: web.Request) -> web.StreamResponse:
+    """Check if a job queue is initialized."""
+    async def _operation() -> web.StreamResponse:
+        job_id = request.match_info['job_id']
+        
+        # Import to ensure initialization
+        from ..upscale.job_store import ensure_tile_jobs_initialized
+        prompt_server = ensure_tile_jobs_initialized()
+        
+        async with prompt_server.distributed_tile_jobs_lock:
+            exists = job_id in prompt_server.distributed_pending_tile_jobs
+        
+        debug_log(f"Queue status check for job {job_id}: {'exists' if exists else 'not found'}")
+        return web.json_response({"exists": exists, "job_id": job_id})
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
+
+@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
+async def update_worker_endpoint(request: web.Request) -> web.StreamResponse:
+    async def _operation() -> web.StreamResponse:
+        data = await request.json()
+        worker_id = data.get("worker_id")
+        
+        if worker_id is None:
+            return await handle_api_error(request, "Missing worker_id", 400)
+
+        async with config_transaction() as config:
+            worker_found = False
+            workers = config.setdefault("workers", [])
+
+            for worker in workers:
+                if worker["id"] == worker_id:
+                    _apply_field_patch(worker, data, _WORKER_FIELDS)
+                    worker_found = True
+                    break
+
+            if not worker_found:
+                # If worker not found and all required fields are provided, create new worker
+                if all(key in data for key in ["name", "port", "cuda_device"]):
+                    new_worker = {
+                        "id": worker_id,
+                        "name": data["name"],
+                        "host": normalize_host(data.get("host", "localhost")),
+                        "port": data["port"],
+                        "cuda_device": data["cuda_device"],
+                        "enabled": data.get("enabled", False),
+                        "extra_args": data.get("extra_args", ""),
+                        "type": data.get("type", "local")
+                    }
+                    workers.append(new_worker)
+                else:
+                    return await handle_api_error(
+                        request,
+                        f"Worker {worker_id} not found and missing required fields for creation",
+                        404,
+                    )
+
+            return web.json_response({"status": "success"})
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
+
+@server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
+async def delete_worker_endpoint(request: web.Request) -> web.StreamResponse:
+    async def _operation() -> web.StreamResponse:
+        data = await request.json()
+        worker_id = data.get("worker_id")
+        
+        if worker_id is None:
+            return await handle_api_error(request, "Missing worker_id", 400)
+            
+        async with config_transaction() as config:
+            workers = config.get("workers", [])
+
+            # Find and remove the worker
+            worker_index = -1
+            for i, worker in enumerate(workers):
+                if worker["id"] == worker_id:
+                    worker_index = i
+                    break
+
+            if worker_index == -1:
+                return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+
+            # Remove the worker
+            removed_worker = workers.pop(worker_index)
+
+            return web.json_response({
+                "status": "success",
+                "message": f"Worker {removed_worker.get('name', worker_id)} deleted"
+            })
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
+
+@server.PromptServer.instance.routes.post("/distributed/config/update_setting")
+async def update_setting_endpoint(request: web.Request) -> web.StreamResponse:
+    """Updates a specific key in the settings object."""
+    async def _operation() -> web.StreamResponse:
+        data = await request.json()
+        key = data.get("key")
+        value = data.get("value")
+
+        if not key or value is None:
+            return await handle_api_error(request, "Missing 'key' or 'value' in request", 400)
+        if key not in _SETTINGS_FIELDS:
+            return await handle_api_error(request, f"Unknown setting: {key}", 400)
+
+        async with config_transaction() as config:
+            if 'settings' not in config:
+                config['settings'] = {}
+
+            config['settings'][key] = value
+
+            return web.json_response({"status": "success", "message": f"Setting '{key}' updated."})
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
+
+@server.PromptServer.instance.routes.post("/distributed/config/update_master")
+async def update_master_endpoint(request: web.Request) -> web.StreamResponse:
+    """Updates master configuration."""
+    async def _operation() -> web.StreamResponse:
+        data = await request.json()
+        
+        async with config_transaction() as config:
+            if 'master' not in config:
+                config['master'] = {}
+            _apply_field_patch(config['master'], data, _MASTER_FIELDS)
+
+            return web.json_response({"status": "success", "message": "Master configuration updated."})
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)

api/endpoint_policy.py

@@ -0,0 +1,25 @@
+"""Shared endpoint policy helpers for distributed API routes."""
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+
+from aiohttp import web
+
+from ..utils.network import handle_api_error
+from .request_guards import authorization_error_or_none
+
+
+async def run_authorized_endpoint(
+    request: web.Request,
+    operation: Callable[[], Awaitable[web.StreamResponse]],
+    *,
+    unexpected_status: int = 500,
+) -> web.StreamResponse:
+    """Run endpoint operation with shared auth guard and fallback error mapping."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        return await operation()
+    except Exception as exc:
+        return await handle_api_error(request, exc, unexpected_status)

api/job_routes.py

@@ -0,0 +1,390 @@
+from __future__ import annotations
+
+import json
+import asyncio
+import io
+import os
+import base64
+import binascii
+import time
+from typing import Any
+
+from aiohttp import web
+import server
+import torch
+from PIL import Image
+
+from ..utils.logging import debug_log
+from ..utils.image import pil_to_tensor, ensure_contiguous
+from ..utils.network import handle_api_error
+from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
+from ..utils.runtime_state import ensure_distributed_runtime_state
+from .request_guards import authorization_error_or_none
+from .schemas import require_bool_literal
+from .queue_orchestration import orchestrate_distributed_execution
+from .queue_request import parse_queue_request_payload
+
+# Canonical worker result envelope accepted by POST /distributed/job_complete:
+# { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool }
+
+
+def _runtime_state():
+    return ensure_distributed_runtime_state()
+
+
+def _decode_image_sync(image_path: str) -> dict[str, Any]:
+    """Decode image/video file and compute hash in a threadpool worker."""
+    import base64
+    import hashlib
+    import folder_paths
+
+    full_path = folder_paths.get_annotated_filepath(image_path)
+    if not os.path.exists(full_path):
+        raise FileNotFoundError(image_path)
+
+    hash_md5 = hashlib.md5()
+    with open(full_path, 'rb') as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_md5.update(chunk)
+    file_hash = hash_md5.hexdigest()
+
+    video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.webm'}
+    file_ext = os.path.splitext(full_path)[1].lower()
+
+    if file_ext in video_extensions:
+        with open(full_path, 'rb') as f:
+            file_data = f.read()
+        mime_types = {
+            '.mp4': 'video/mp4',
+            '.avi': 'video/x-msvideo',
+            '.mov': 'video/quicktime',
+            '.mkv': 'video/x-matroska',
+            '.webm': 'video/webm'
+        }
+        mime_type = mime_types.get(file_ext, 'video/mp4')
+        image_data = f"data:{mime_type};base64,{base64.b64encode(file_data).decode('utf-8')}"
+    else:
+        with Image.open(full_path) as img:
+            if img.mode not in ('RGB', 'RGBA'):
+                img = img.convert('RGB')
+            buffer = io.BytesIO()
+            img.save(buffer, format='PNG', compress_level=1)
+            image_data = f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
+
+    return {
+        "status": "success",
+        "image_data": image_data,
+        "hash": file_hash,
+    }
+
+
+def _check_file_sync(filename: str, expected_hash: str) -> dict[str, Any]:
+    """Check file presence and hash in a threadpool worker."""
+    import hashlib
+    import folder_paths
+
+    full_path = folder_paths.get_annotated_filepath(filename)
+    if not os.path.exists(full_path):
+        return {
+            "status": "success",
+            "exists": False,
+        }
+
+    hash_md5 = hashlib.md5()
+    with open(full_path, 'rb') as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_md5.update(chunk)
+    file_hash = hash_md5.hexdigest()
+
+    return {
+        "status": "success",
+        "exists": True,
+        "hash_matches": file_hash == expected_hash,
+    }
+
+
+def _decode_canonical_png_tensor(image_payload: str) -> torch.Tensor:
+    """Decode canonical base64 PNG payload into a contiguous IMAGE tensor."""
+    if not isinstance(image_payload, str) or not image_payload.strip():
+        raise ValueError("Field 'image' must be a non-empty base64 PNG string.")
+
+    encoded = image_payload.strip()
+    if encoded.startswith("data:"):
+        header, sep, data_part = encoded.partition(",")
+        if not sep:
+            raise ValueError("Field 'image' data URL is malformed.")
+        if not header.lower().startswith("data:image/png;base64"):
+            raise ValueError("Field 'image' must be a PNG data URL when using data:* format.")
+        encoded = data_part
+
+    try:
+        png_bytes = base64.b64decode(encoded, validate=True)
+    except (binascii.Error, ValueError) as exc:
+        raise ValueError("Field 'image' is not valid base64 PNG data.") from exc
+
+    if not png_bytes:
+        raise ValueError("Field 'image' decoded to empty PNG data.")
+
+    try:
+        with Image.open(io.BytesIO(png_bytes)) as img:
+            img = img.convert("RGB")
+            tensor = pil_to_tensor(img)
+        return ensure_contiguous(tensor)
+    except Exception as exc:
+        raise ValueError(f"Failed to decode PNG image payload: {exc}") from exc
+
+
+def _decode_audio_payload(audio_payload: dict[str, Any]) -> dict[str, Any]:
+    """Decode canonical audio payload into an AUDIO dict."""
+    from ..utils.audio_payload import decode_audio_payload
+
+    return decode_audio_payload(audio_payload)
+
+
+@server.PromptServer.instance.routes.post("/distributed/prepare_job")
+async def prepare_job_endpoint(request: web.Request) -> web.StreamResponse:
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        data = await request.json()
+        multi_job_id = data.get('multi_job_id')
+        if not multi_job_id:
+            return await handle_api_error(request, "Missing multi_job_id", 400)
+
+        runtime_state = _runtime_state()
+        async with runtime_state.distributed_jobs_lock:
+            if multi_job_id not in runtime_state.distributed_pending_jobs:
+                runtime_state.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
+        
+        debug_log(f"Prepared queue for job {multi_job_id}")
+        return web.json_response({"status": "success"})
+    except Exception as e:
+        return await handle_api_error(request, e)
+
+@server.PromptServer.instance.routes.post("/distributed/clear_memory")
+async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
+    debug_log("Received request to clear VRAM.")
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        warnings: list[str] = []
+        # Use ComfyUI's prompt server queue system like the /free endpoint does
+        if hasattr(server.PromptServer.instance, 'prompt_queue'):
+            server.PromptServer.instance.prompt_queue.set_flag("unload_models", True)
+            server.PromptServer.instance.prompt_queue.set_flag("free_memory", True)
+            debug_log("Set queue flags for memory clearing.")
+        
+        # Wait a bit for the queue to process
+        await asyncio.sleep(MEMORY_CLEAR_DELAY)
+        
+        # Also do direct cleanup as backup, but with error handling
+        import gc
+        import comfy.model_management as mm
+        
+        try:
+            mm.unload_all_models()
+        except AttributeError as e:
+            warning = f"Model unload warning: {e}"
+            warnings.append(warning)
+            debug_log(warning)
+        
+        try:
+            mm.soft_empty_cache()
+        except Exception as e:
+            warning = f"Cache clear warning: {e}"
+            warnings.append(warning)
+            debug_log(warning)
+        
+        for _ in range(3):
+            gc.collect()
+        
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        
+        if warnings:
+            debug_log("VRAM cleared with warnings.")
+            return web.json_response(
+                {
+                    "status": "partial",
+                    "message": "GPU memory cleared with warnings.",
+                    "warnings": warnings,
+                }
+            )
+
+        debug_log("VRAM cleared successfully.")
+        return web.json_response({"status": "success", "message": "GPU memory cleared."})
+    except Exception as e:
+        # Even if there's an error, try to do basic cleanup
+        import gc
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        debug_log(f"VRAM clear failed: {e}")
+        return await handle_api_error(request, f"GPU memory clear failed: {e}", 500)
+
+
+@server.PromptServer.instance.routes.post("/distributed/queue")
+async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse:
+    """Queue a distributed workflow, mirroring the UI orchestration pipeline."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        raw_payload = await request.json()
+    except Exception as exc:
+        return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
+
+    try:
+        payload = parse_queue_request_payload(raw_payload)
+    except ValueError as exc:
+        return await handle_api_error(request, exc, 400)
+
+    try:
+        prompt_id, worker_count = await orchestrate_distributed_execution(
+            payload.prompt,
+            payload.workflow_meta,
+            payload.client_id,
+            enabled_worker_ids=payload.enabled_worker_ids,
+            delegate_master=payload.delegate_master,
+            trace_execution_id=payload.trace_execution_id,
+        )
+        return web.json_response({
+            "prompt_id": prompt_id,
+            "worker_count": worker_count,
+        })
+    except Exception as exc:
+        return await handle_api_error(request, exc, 500)
+
+@server.PromptServer.instance.routes.post("/distributed/load_image")
+async def load_image_endpoint(request: web.Request) -> web.StreamResponse:
+    """Load an image or video file and return it as base64 data with hash."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        data = await request.json()
+        image_path = data.get("image_path")
+        
+        if not image_path:
+            return await handle_api_error(request, "Missing image_path", 400)
+        loop = asyncio.get_running_loop()
+        payload = await loop.run_in_executor(None, _decode_image_sync, image_path)
+        return web.json_response(payload)
+    except FileNotFoundError:
+        return await handle_api_error(request, f"File not found: {image_path}", 404)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+@server.PromptServer.instance.routes.post("/distributed/check_file")
+async def check_file_endpoint(request: web.Request) -> web.StreamResponse:
+    """Check if a file exists and matches the given hash."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        data = await request.json()
+        filename = data.get("filename")
+        expected_hash = data.get("hash")
+        
+        if not filename or not expected_hash:
+            return await handle_api_error(request, "Missing filename or hash", 400)
+        loop = asyncio.get_running_loop()
+        payload = await loop.run_in_executor(None, _check_file_sync, filename, expected_hash)
+        return web.json_response(payload)
+        
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.post("/distributed/job_complete")
+async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        data = await request.json()
+    except Exception as exc:
+        return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
+
+    if not isinstance(data, dict):
+        return await handle_api_error(request, "Expected a JSON object body", 400)
+
+    try:
+        job_id = data.get("job_id")
+        worker_id = data.get("worker_id")
+        batch_idx = data.get("batch_idx")
+        image_payload = data.get("image")
+        audio_payload = data.get("audio")
+        is_last_raw = data.get("is_last")
+
+        errors = []
+        if not isinstance(job_id, str) or not job_id.strip():
+            errors.append("job_id: expected non-empty string")
+        if not isinstance(worker_id, str) or not worker_id.strip():
+            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")
+        if not isinstance(image_payload, str) or not image_payload.strip():
+            errors.append("image: expected non-empty base64 PNG string")
+        if audio_payload is not None and not isinstance(audio_payload, dict):
+            errors.append("audio: expected object when provided")
+
+        try:
+            is_last = require_bool_literal(is_last_raw, field_name="is_last")
+        except ValueError as exc:
+            errors.append(str(exc))
+            is_last = False
+        if errors:
+            return await handle_api_error(request, errors, 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()
+
+        runtime_state = _runtime_state()
+        allowed_workers = runtime_state.distributed_job_allowed_workers.get(multi_job_id)
+        if allowed_workers is not None and worker_id not in allowed_workers:
+            return await handle_api_error(
+                request,
+                f"Unauthorized worker_id for job {multi_job_id}",
+                403,
+            )
+
+        pending = None
+        queue_size = 0
+        deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
+        while pending is None:
+            async with runtime_state.distributed_jobs_lock:
+                pending = runtime_state.distributed_pending_jobs.get(multi_job_id)
+                if pending is not None:
+                    queue_item = {
+                        "tensor": tensor,
+                        "worker_id": worker_id,
+                        "image_index": int(batch_idx),
+                        "is_last": is_last,
+                    }
+                    if decoded_audio is not None:
+                        queue_item["audio"] = decoded_audio
+                    await pending.put(
+                        queue_item
+                    )
+                    queue_size = pending.qsize()
+                    break
+
+            if time.monotonic() > deadline:
+                return await handle_api_error(request, "job not initialized", 404)
+            await asyncio.sleep(0.05)
+
+        debug_log(
+            f"job_complete received canonical envelope - job_id: {multi_job_id}, "
+            f"worker: {worker_id}, batch_idx: {batch_idx}, is_last: {is_last}, "
+            f"queue_size: {queue_size}"
+        )
+
+        return web.json_response({"status": "success"})
+    except Exception as e:
+        return await handle_api_error(request, e)

api/orchestration/__init__.py

@@ -0,0 +1,7 @@
+"""Orchestration helpers used by distributed queue execution."""
+
+__all__ = [
+    "dispatch",
+    "media_sync",
+    "prompt_transform",
+]

api/orchestration/dispatch.py

@@ -0,0 +1,316 @@
+import asyncio
+import json
+import uuid
+from typing import Any
+
+import aiohttp
+
+from ...utils.auth import distributed_auth_headers
+from ...utils.config import load_config
+from ...utils.logging import debug_log, log
+from ...utils.network import build_worker_url, get_client_session, probe_worker
+from ...utils.trace_logger import trace_debug, trace_info
+from ..schemas import coerce_positive_int
+
+_least_busy_rr_index = 0
+
+
+async def worker_is_active(worker: dict[str, Any]) -> bool:
+    """Ping worker's /prompt endpoint to confirm it's reachable."""
+    url = build_worker_url(worker)
+    return await probe_worker(url, timeout=3.0) is not None
+
+
+async def worker_ws_is_active(worker: dict[str, Any]) -> bool:
+    """Ping worker's websocket endpoint to confirm it's reachable."""
+    session = await get_client_session()
+    url = build_worker_url(worker, "/distributed/worker_ws")
+    headers = distributed_auth_headers(load_config())
+    try:
+        ws = await session.ws_connect(url, heartbeat=20, timeout=3, headers=headers)
+        await ws.close()
+        return True
+    except asyncio.TimeoutError:
+        debug_log(f"[Distributed] Worker WS probe timed out: {url}")
+        return False
+    except aiohttp.ClientConnectorError:
+        debug_log(f"[Distributed] Worker WS unreachable: {url}")
+        return False
+    except Exception as e:
+        debug_log(f"[Distributed] Worker WS probe unexpected error: {e}")
+        return False
+
+
+async def _probe_worker_active(worker, use_websocket, semaphore):
+    async with semaphore:
+        is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker))
+        return worker, is_active
+
+
+async def _dispatch_via_websocket(
+    worker_url,
+    payload,
+    client_id,
+    headers: dict[str, str],
+    timeout=60.0,
+):
+    """Open a fresh worker websocket, dispatch one prompt, wait for ack, then close."""
+    request_id = uuid.uuid4().hex
+    ws_payload = {
+        "type": "dispatch_prompt",
+        "request_id": request_id,
+        "prompt": payload.get("prompt"),
+        "workflow": payload.get("workflow"),
+        "client_id": client_id,
+    }
+    ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
+    ws_url = f"{ws_url}/distributed/worker_ws"
+    session = await get_client_session()
+
+    async with session.ws_connect(
+        ws_url,
+        heartbeat=20,
+        timeout=timeout,
+        headers=headers,
+    ) as ws:
+        await ws.send_json(ws_payload)
+        async for msg in ws:
+            if msg.type == aiohttp.WSMsgType.TEXT:
+                data = json.loads(msg.data or "{}")
+                if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
+                    if data.get("ok"):
+                        return
+                    error_text = data.get("error") or "Worker rejected websocket dispatch."
+                    validation_error = data.get("validation_error")
+                    node_errors = data.get("node_errors")
+                    if validation_error:
+                        error_text = f"{error_text} | validation_error={validation_error}"
+                    if node_errors:
+                        error_text = f"{error_text} | node_errors={node_errors}"
+                    raise RuntimeError(error_text)
+            elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
+                raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}")
+
+    raise RuntimeError("Worker websocket closed before dispatch_ack was received.")
+
+
+async def dispatch_worker_prompt(
+    worker: dict[str, Any],
+    prompt_obj: dict[str, Any],
+    workflow_meta: dict[str, Any] | None,
+    client_id: str | None = None,
+    use_websocket: bool = False,
+    trace_execution_id: str | None = None,
+) -> None:
+    """Send the prepared prompt to a worker ComfyUI instance."""
+    worker_url = build_worker_url(worker)
+    url = build_worker_url(worker, "/prompt")
+    payload = {"prompt": prompt_obj}
+    extra_data = {}
+    if workflow_meta:
+        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
+    if extra_data:
+        payload["extra_data"] = extra_data
+
+    if use_websocket:
+        ws_headers = distributed_auth_headers(load_config())
+        try:
+            await _dispatch_via_websocket(
+                worker_url,
+                {
+                    "prompt": prompt_obj,
+                    "workflow": workflow_meta,
+                },
+                client_id,
+                ws_headers,
+            )
+            return
+        except Exception as exc:
+            worker_id = worker.get("id")
+            if trace_execution_id:
+                trace_info(trace_execution_id, f"Websocket dispatch failed for worker {worker_id}: {exc}")
+            else:
+                log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
+            raise
+
+    session = await get_client_session()
+    async with session.post(
+        url,
+        json=payload,
+        timeout=aiohttp.ClientTimeout(total=60),
+    ) as resp:
+        resp.raise_for_status()
+
+
+async def select_active_workers(
+    workers: list[dict[str, Any]],
+    use_websocket: bool,
+    delegate_master: bool,
+    trace_execution_id: str | None = None,
+    probe_concurrency: int = 8,
+) -> tuple[list[dict[str, Any]], bool]:
+    """Probe workers and return (active_workers, updated_delegate_master)."""
+    probe_limit = coerce_positive_int(probe_concurrency, 8)
+    probe_semaphore = asyncio.Semaphore(probe_limit)
+
+    if trace_execution_id and workers:
+        trace_debug(
+            trace_execution_id,
+            f"Probing {len(workers)} workers with probe_concurrency={probe_limit}",
+        )
+
+    probe_results = await asyncio.gather(
+        *[
+            _probe_worker_active(worker, use_websocket, probe_semaphore)
+            for worker in workers
+        ]
+    )
+
+    active_workers = []
+    for worker, is_active in probe_results:
+        if is_active:
+            active_workers.append(worker)
+        else:
+            if trace_execution_id:
+                trace_info(trace_execution_id, f"Worker {worker['name']} ({worker['id']}) is offline, skipping.")
+            else:
+                log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
+
+    if trace_execution_id and workers:
+        trace_debug(
+            trace_execution_id,
+            f"Worker probe complete: active={len(active_workers)}/{len(workers)}",
+        )
+
+    if not active_workers and delegate_master:
+        if trace_execution_id:
+            trace_debug(trace_execution_id, "All workers offline while delegate-only requested; enabling master participation.")
+        else:
+            debug_log("All workers offline while delegate-only requested; enabling master participation.")
+        delegate_master = False
+
+    return active_workers, delegate_master
+
+
+def _extract_queue_remaining(payload):
+    if not isinstance(payload, dict):
+        return 0
+    try:
+        queue_remaining = int(payload.get("exec_info", {}).get("queue_remaining", 0))
+    except (TypeError, ValueError):
+        queue_remaining = 0
+    return max(queue_remaining, 0)
+
+
+async def _probe_worker_queue(worker, semaphore, probe_timeout):
+    async with semaphore:
+        worker_url = build_worker_url(worker)
+        payload = await probe_worker(worker_url, timeout=probe_timeout)
+        if payload is None:
+            return None
+        return {
+            "worker": worker,
+            "queue_remaining": _extract_queue_remaining(payload),
+        }
+
+
+def _select_idle_round_robin(statuses):
+    global _least_busy_rr_index
+    if not statuses:
+        return None
+    index = _least_busy_rr_index % len(statuses)
+    _least_busy_rr_index += 1
+    return statuses[index]
+
+
+async def select_least_busy_worker(
+    workers: list[dict[str, Any]],
+    trace_execution_id: str | None = None,
+    probe_concurrency: int = 8,
+    probe_timeout: float = 3.0,
+) -> dict[str, Any] | None:
+    """Select one worker by queue depth, round-robin among idle workers."""
+    if not workers:
+        return None
+
+    probe_limit = coerce_positive_int(probe_concurrency, 8)
+    probe_semaphore = asyncio.Semaphore(probe_limit)
+    statuses = await asyncio.gather(
+        *[
+            _probe_worker_queue(worker, probe_semaphore, probe_timeout)
+            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
+
+    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"])
+
+    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 rank_workers_by_load(
+    workers: list[dict[str, Any]],
+    trace_execution_id: str | None = None,
+    probe_concurrency: int = 8,
+    probe_timeout: float = 3.0,
+) -> list[dict[str, Any]]:
+    """Return all workers sorted by queue depth (ascending), preserving unreachable workers at the end."""
+    if not workers:
+        return []
+
+    probe_limit = coerce_positive_int(probe_concurrency, 8)
+    probe_semaphore = asyncio.Semaphore(probe_limit)
+    statuses = await asyncio.gather(
+        *[
+            _probe_worker_queue(worker, probe_semaphore, probe_timeout)
+            for worker in workers
+        ]
+    )
+
+    ranked_statuses = []
+    unreachable_workers = []
+    for worker, status in zip(workers, statuses):
+        if status is None:
+            unreachable_workers.append(worker)
+            continue
+        ranked_statuses.append(status)
+
+    ranked_statuses.sort(key=lambda status: status["queue_remaining"])
+    ranked_workers = [status["worker"] for status in ranked_statuses] + unreachable_workers
+
+    if trace_execution_id:
+        trace_debug(
+            trace_execution_id,
+            "Ranked workers by load: "
+            + ", ".join(
+                f"{status['worker'].get('id')}={status['queue_remaining']}"
+                for status in ranked_statuses
+            )
+            + (
+                f"; unreachable={','.join(worker.get('id', '') for worker in unreachable_workers)}"
+                if unreachable_workers
+                else ""
+            ),
+        )
+    return ranked_workers

api/orchestration/media_sync.py

@@ -0,0 +1,269 @@
+import asyncio
+import hashlib
+import mimetypes
+import os
+import re
+from typing import Any
+
+import aiohttp
+
+from ...utils.auth import distributed_auth_headers
+from ...utils.config import load_config
+from ...utils.logging import debug_log, log
+from ...utils.network import build_worker_url, get_client_session
+from ...utils.trace_logger import trace_debug, trace_info
+
+
+LIKELY_FILENAME_RE = re.compile(
+    r"\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm|"
+    r"wav|mp3|flac|m4a|aac|ogg|opus|aiff|aif|wma|latent|txt|vae|lora|embedding)"
+    r"(\s*\[\w+\])?$",
+    re.IGNORECASE,
+)
+MEDIA_FILE_RE = re.compile(
+    r"\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm|wav|mp3|flac|m4a|aac|ogg|opus|aiff|aif|wma)(\s*\[\w+\])?$",
+    re.IGNORECASE,
+)
+
+
+def _normalize_media_reference(value):
+    """Normalize one media string value to a path-like reference or None."""
+    if not isinstance(value, str):
+        return None
+    cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/")
+    if MEDIA_FILE_RE.search(cleaned):
+        return cleaned
+    return None
+
+
+def convert_paths_for_platform(obj: Any, target_separator: str) -> Any:
+    """Recursively normalize likely file paths for the worker platform separator."""
+    if target_separator not in ("/", "\\"):
+        return obj
+
+    def _convert(value):
+        if isinstance(value, str):
+            if ("/" in value or "\\" in value) and LIKELY_FILENAME_RE.search(value):
+                trimmed = value.strip()
+                has_drive = bool(re.match(r"^[A-Za-z]:(\\\\|/)", trimmed))
+                is_absolute = trimmed.startswith("/") or trimmed.startswith("\\\\")
+                has_protocol = bool(re.match(r"^\w+://", trimmed))
+
+                # URLs are not local paths and should never be separator-normalized.
+                if has_protocol:
+                    return trimmed
+
+                # Keep relative media-style paths in forward-slash form (Comfy-style annotated paths).
+                if not has_drive and not is_absolute and not has_protocol and MEDIA_FILE_RE.search(trimmed):
+                    return re.sub(r"[\\]+", "/", trimmed)
+
+                if target_separator == "\\":
+                    return re.sub(r"[\\/]+", r"\\", trimmed)
+                return re.sub(r"[\\/]+", "/", trimmed)
+            return value
+        if isinstance(value, list):
+            return [_convert(item) for item in value]
+        if isinstance(value, dict):
+            return {key: _convert(item) for key, item in value.items()}
+        return value
+
+    return _convert(obj)
+
+
+def _find_media_references(prompt_obj):
+    """Find media file references in image/video/audio/file inputs used by worker prompts."""
+    media_refs = set()
+    for node in prompt_obj.values():
+        if not isinstance(node, dict):
+            continue
+        inputs = node.get("inputs", {})
+        for key in ("image", "video", "audio", "file"):
+            cleaned = _normalize_media_reference(inputs.get(key))
+            if cleaned:
+                media_refs.add(cleaned)
+    return sorted(media_refs)
+
+
+def _rewrite_prompt_media_inputs(prompt_obj, worker_media_paths):
+    """Rewrite media string inputs to worker-local uploaded paths."""
+    if not isinstance(worker_media_paths, dict) or not worker_media_paths:
+        return
+
+    for node in prompt_obj.values():
+        if not isinstance(node, dict):
+            continue
+        inputs = node.get("inputs", {})
+        if not isinstance(inputs, dict):
+            continue
+        for key in ("image", "video", "audio", "file"):
+            value = inputs.get(key)
+            cleaned = _normalize_media_reference(value)
+            if not cleaned:
+                continue
+            worker_path = worker_media_paths.get(cleaned)
+            if worker_path:
+                inputs[key] = worker_path
+
+
+def _load_media_file_sync(filename):
+    """Load local media bytes and hash for worker upload sync."""
+    import folder_paths
+
+    full_path = folder_paths.get_annotated_filepath(filename)
+    if not os.path.exists(full_path):
+        raise FileNotFoundError(filename)
+
+    with open(full_path, "rb") as f:
+        file_bytes = f.read()
+
+    file_hash = hashlib.md5(file_bytes).hexdigest()
+    mime_type = mimetypes.guess_type(full_path)[0]
+    if not mime_type:
+        ext = os.path.splitext(full_path)[1].lower()
+        if ext in {".mp4", ".avi", ".mov", ".mkv", ".webm"}:
+            mime_type = "video/mp4"
+        else:
+            mime_type = "image/png"
+    return file_bytes, file_hash, mime_type
+
+
+async def fetch_worker_path_separator(
+    worker: dict[str, Any],
+    trace_execution_id: str | None = None,
+) -> str | None:
+    """Best-effort fetch of a worker's path separator from /distributed/system_info."""
+    url = build_worker_url(worker, "/distributed/system_info")
+    session = await get_client_session()
+    headers = distributed_auth_headers(load_config())
+    try:
+        async with session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=5)) as resp:
+            if resp.status != 200:
+                return None
+            payload = await resp.json()
+            separator = ((payload or {}).get("platform") or {}).get("path_separator")
+            return separator if separator in ("/", "\\") else None
+    except Exception as exc:
+        if trace_execution_id:
+            trace_debug(trace_execution_id, f"Failed to fetch worker system info ({worker.get('id')}): {exc}")
+        else:
+            debug_log(f"[Distributed] Failed to fetch worker system info ({worker.get('id')}): {exc}")
+        return None
+
+
+async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None):
+    """Upload one media file to worker iff missing or hash-mismatched."""
+    session = await get_client_session()
+    headers = distributed_auth_headers(load_config())
+    normalized = filename.replace("\\", "/")
+
+    check_url = build_worker_url(worker, "/distributed/check_file")
+    try:
+        async with session.post(
+            check_url,
+            json={"filename": normalized, "hash": file_hash},
+            headers=headers,
+            timeout=aiohttp.ClientTimeout(total=6),
+        ) as resp:
+            if resp.status == 200:
+                payload = await resp.json()
+                if payload.get("exists") and payload.get("hash_matches"):
+                    return False, normalized
+    except Exception as exc:
+        if trace_execution_id:
+            trace_debug(trace_execution_id, f"Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}")
+        else:
+            debug_log(f"[Distributed] Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}")
+
+    parts = normalized.split("/")
+    clean_name = parts[-1]
+    subfolder = "/".join(parts[:-1])
+
+    form = aiohttp.FormData()
+    form.add_field("image", file_bytes, filename=clean_name, content_type=mime_type)
+    form.add_field("type", "input")
+    form.add_field("subfolder", subfolder)
+    form.add_field("overwrite", "true")
+
+    upload_url = build_worker_url(worker, "/upload/image")
+    async with session.post(
+        upload_url,
+        data=form,
+        timeout=aiohttp.ClientTimeout(total=30),
+    ) as resp:
+        resp.raise_for_status()
+        try:
+            payload = await resp.json()
+        except Exception:
+            payload = {}
+
+    name = str((payload or {}).get("name") or clean_name).strip()
+    subfolder = str((payload or {}).get("subfolder") or "").strip().replace("\\", "/").strip("/")
+    worker_path = f"{subfolder}/{name}" if subfolder else name
+    return True, worker_path
+
+
+async def sync_worker_media(
+    worker: dict[str, Any],
+    prompt_obj: dict[str, Any],
+    trace_execution_id: str | None = None,
+) -> None:
+    """Sync referenced media files from master to a remote worker before dispatch."""
+    media_refs = _find_media_references(prompt_obj)
+    if not media_refs:
+        return
+
+    loop = asyncio.get_running_loop()
+    uploaded = 0
+    skipped = 0
+    missing = 0
+    worker_media_paths = {}
+    for filename in media_refs:
+        try:
+            file_bytes, file_hash, mime_type = await loop.run_in_executor(
+                None, _load_media_file_sync, filename
+            )
+        except FileNotFoundError:
+            missing += 1
+            if trace_execution_id:
+                trace_info(trace_execution_id, f"Media file '{filename}' not found on master; worker may fail to load it.")
+            else:
+                log(f"[Distributed] Media file '{filename}' not found on master; worker may fail to load it.")
+            continue
+        except Exception as exc:
+            if trace_execution_id:
+                trace_info(trace_execution_id, f"Failed to load media '{filename}' for worker sync: {exc}")
+            else:
+                log(f"[Distributed] Failed to load media '{filename}' for worker sync: {exc}")
+            continue
+
+        try:
+            changed, worker_path = await _upload_media_to_worker(
+                worker,
+                filename,
+                file_bytes,
+                file_hash,
+                mime_type,
+                trace_execution_id=trace_execution_id,
+            )
+            if worker_path:
+                worker_media_paths[filename] = worker_path
+            if changed:
+                uploaded += 1
+            else:
+                skipped += 1
+        except Exception as exc:
+            if trace_execution_id:
+                trace_info(trace_execution_id, f"Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}")
+            else:
+                log(f"[Distributed] Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}")
+
+    _rewrite_prompt_media_inputs(prompt_obj, worker_media_paths)
+
+    summary = (
+        f"Media sync for worker {worker.get('id')}: "
+        f"uploaded={uploaded}, skipped={skipped}, missing={missing}, referenced={len(media_refs)}"
+    )
+    if trace_execution_id:
+        trace_debug(trace_execution_id, summary)
+    else:
+        debug_log(f"[Distributed] {summary}")

api/orchestration/prompt_transform.py

@@ -0,0 +1,676 @@
+import json
+from collections import deque
+from typing import Any
+
+from ...utils.logging import debug_log
+
+
+class PromptIndex:
+    """Cache prompt metadata for faster worker/master prompt preparation."""
+
+    def __init__(self, prompt_obj):
+        self._prompt_json = json.dumps(prompt_obj)
+        self.nodes_by_class = {}
+        self.class_by_node = {}
+        self.inputs_by_node = {}
+        for node_id, node in _iter_prompt_nodes(prompt_obj):
+            class_type = node.get("class_type")
+            node_id_str = str(node_id)
+            if class_type:
+                self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
+            self.class_by_node[node_id_str] = class_type
+            self.inputs_by_node[node_id_str] = node.get("inputs", {})
+        self._upstream_cache = {}
+
+    def copy_prompt(self):
+        return json.loads(self._prompt_json)
+
+    def nodes_for_class(self, class_name):
+        return self.nodes_by_class.get(class_name, [])
+
+    def has_upstream(self, start_node_id: str, target_class: str) -> bool:
+        cache_key = (str(start_node_id), target_class)
+        if cache_key in self._upstream_cache:
+            return self._upstream_cache[cache_key]
+
+        visited = set()
+        stack = [str(start_node_id)]
+        while stack:
+            node_id = stack.pop()
+            if node_id in visited:
+                continue
+            visited.add(node_id)
+            inputs = self.inputs_by_node.get(node_id, {})
+            for value in inputs.values():
+                if isinstance(value, list) and len(value) == 2:
+                    upstream_id = str(value[0])
+                    if self.class_by_node.get(upstream_id) == target_class:
+                        self._upstream_cache[cache_key] = True
+                        return True
+                    if upstream_id in self.inputs_by_node:
+                        stack.append(upstream_id)
+
+        self._upstream_cache[cache_key] = False
+        return False
+
+
+def _iter_prompt_nodes(prompt_obj):
+    for node_id, node in prompt_obj.items():
+        if isinstance(node, dict):
+            yield str(node_id), node
+
+
+def find_nodes_by_class(prompt_obj, class_name):
+    nodes = []
+    for node_id, node in _iter_prompt_nodes(prompt_obj):
+        if node.get("class_type") == class_name:
+            nodes.append(node_id)
+    return nodes
+
+
+def _find_downstream_nodes(prompt_obj, start_ids):
+    """Return all nodes reachable downstream from the provided IDs."""
+    adjacency = {}
+    for node_id, node in _iter_prompt_nodes(prompt_obj):
+        inputs = node.get("inputs", {})
+        for value in inputs.values():
+            if isinstance(value, list) and len(value) == 2:
+                source_id = str(value[0])
+                adjacency.setdefault(source_id, set()).add(str(node_id))
+
+    connected = set(start_ids)
+    queue = deque(start_ids)
+    while queue:
+        current = queue.popleft()
+        for dependent in adjacency.get(current, ()):  # pragma: no branch - simple iteration
+            if dependent not in connected:
+                connected.add(dependent)
+                queue.append(dependent)
+    return connected
+
+
+def _find_downstream_of_output_slot(prompt_obj, node_id, slot_index):
+    """Return all nodes reachable from a specific output slot on one node."""
+    source_id = str(node_id)
+    try:
+        slot = int(slot_index)
+    except (TypeError, ValueError):
+        return set()
+
+    direct_consumers = set()
+    for candidate_id, candidate_node in _iter_prompt_nodes(prompt_obj):
+        inputs = candidate_node.get("inputs", {})
+        for value in inputs.values():
+            if not (isinstance(value, list) and len(value) == 2):
+                continue
+            if str(value[0]) != source_id:
+                continue
+            try:
+                input_slot = int(value[1])
+            except (TypeError, ValueError):
+                debug_log(
+                    f"Prompt transform: skipping malformed slot reference for node {candidate_id}: {value}"
+                )
+                continue
+            if input_slot == slot:
+                direct_consumers.add(str(candidate_id))
+                break
+
+    if not direct_consumers:
+        return set()
+    return _find_downstream_nodes(prompt_obj, list(direct_consumers))
+
+
+def _create_numeric_id_generator(prompt_obj):
+    """Return a closure that yields new numeric string IDs."""
+    max_id = 0
+    for node_id in prompt_obj.keys():
+        try:
+            numeric = int(node_id)
+        except (TypeError, ValueError):
+            debug_log(f"Prompt transform: ignoring non-numeric node id while generating ids: {node_id}")
+            continue
+        max_id = max(max_id, numeric)
+
+    counter = max_id
+
+    def _next_id():
+        nonlocal counter
+        counter += 1
+        return str(counter)
+
+    return _next_id
+
+
+def _find_upstream_nodes(prompt_obj, start_ids):
+    """Return all nodes reachable upstream from start_ids, including start nodes."""
+    connected = set(str(node_id) for node_id in start_ids)
+    queue = deque(connected)
+    while queue:
+        node_id = queue.popleft()
+        node = prompt_obj.get(node_id) or {}
+        inputs = node.get("inputs", {})
+        for value in inputs.values():
+            if isinstance(value, list) and len(value) == 2:
+                source_id = str(value[0])
+                if source_id in prompt_obj and source_id not in connected:
+                    connected.add(source_id)
+                    queue.append(source_id)
+    return connected
+
+
+def _resolve_participants(enabled_worker_ids, delegate_master):
+    worker_ids = [str(worker_id) for worker_id in (enabled_worker_ids or [])]
+    if delegate_master:
+        return worker_ids
+    return ["master"] + worker_ids
+
+
+def _remove_dangling_input_refs(prompt_obj):
+    """Drop input links that point to nodes no longer present in prompt_obj."""
+    existing_ids = set(prompt_obj.keys())
+    for _node_id, node in _iter_prompt_nodes(prompt_obj):
+        inputs = node.get("inputs", {})
+        for input_name, input_value in list(inputs.items()):
+            if isinstance(input_value, list) and len(input_value) == 2:
+                source_id = str(input_value[0])
+                if source_id not in existing_ids:
+                    inputs.pop(input_name, None)
+
+
+def _has_terminal_output_nodes(prompt_obj):
+    """Return True when the prompt already has at least one terminal output node."""
+    for _node_id, node in _iter_prompt_nodes(prompt_obj):
+        if node.get("class_type") in {"PreviewImage", "SaveImage"}:
+            return True
+    return False
+
+
+def _add_preview_node(prompt_obj, next_id_fn, source_node_id, slot_index, title_suffix):
+    """Attach a PreviewImage output node to the given source output slot."""
+    preview_id = next_id_fn()
+    prompt_obj[preview_id] = {
+        "inputs": {
+            "images": [str(source_node_id), int(slot_index)],
+        },
+        "class_type": "PreviewImage",
+        "_meta": {
+            "title": f"Preview Image ({title_suffix})",
+        },
+    }
+    return preview_id
+
+
+def _ensure_worker_output_node(prompt_obj):
+    """Guarantee worker prompts retain at least one output node after pruning.
+
+    Branch pruning can remove all explicit output nodes for worker prompts, which
+    causes ComfyUI validation to reject the prompt (`prompt_no_outputs`).
+    """
+    if _has_terminal_output_nodes(prompt_obj):
+        return prompt_obj
+
+    next_id = _create_numeric_id_generator(prompt_obj)
+
+    # Prefer a DistributedBranchCollector output tied to the assigned branch for this worker.
+    for node_id, node in _iter_prompt_nodes(prompt_obj):
+        if node.get("class_type") != "DistributedBranchCollector":
+            continue
+        inputs = node.get("inputs", {})
+        try:
+            assigned_branch = int(inputs.get("assigned_branch", -1))
+        except (TypeError, ValueError):
+            assigned_branch = -1
+        if assigned_branch >= 0:
+            _add_preview_node(prompt_obj, next_id, node_id, assigned_branch, "auto-added worker output")
+            return prompt_obj
+
+    # Otherwise, anchor to the assigned DistributedBranch slot if available.
+    for node_id, node in _iter_prompt_nodes(prompt_obj):
+        if node.get("class_type") != "DistributedBranch":
+            continue
+        inputs = node.get("inputs", {})
+        try:
+            assigned_branch = int(inputs.get("assigned_branch", -1))
+        except (TypeError, ValueError):
+            assigned_branch = -1
+        if assigned_branch >= 0:
+            _add_preview_node(prompt_obj, next_id, node_id, assigned_branch, "auto-added worker output")
+            return prompt_obj
+
+    # Idle participants (assigned_branch=-1) still need a terminal node to pass
+    # validation; keep this cheap by emitting a tiny synthetic image preview.
+    empty_id = next_id()
+    prompt_obj[empty_id] = {
+        "class_type": "DistributedEmptyImage",
+        "inputs": {
+            "height": 64,
+            "width": 64,
+            "channels": 3,
+        },
+        "_meta": {
+            "title": "Distributed Empty Image (auto-added worker output)",
+        },
+    }
+    _add_preview_node(prompt_obj, next_id, empty_id, 0, "auto-added worker output")
+    return prompt_obj
+
+
+def _prune_worker_downstream_of_branch_collectors(prompt_obj):
+    """Drop worker-side nodes downstream of branch collectors."""
+    collector_ids = find_nodes_by_class(prompt_obj, "DistributedBranchCollector")
+    if not collector_ids:
+        return prompt_obj
+
+    downstream = _find_downstream_nodes(prompt_obj, collector_ids)
+    for node_id in downstream:
+        if node_id in collector_ids:
+            continue
+        prompt_obj.pop(node_id, None)
+
+    _remove_dangling_input_refs(prompt_obj)
+    return prompt_obj
+
+
+def prune_prompt_for_branch_worker(
+    prompt_obj: dict[str, Any],
+    branch_node_id: str,
+    assigned_branch: int | list[int] | tuple[int, ...] | set[int],
+    num_branches: int,
+) -> dict[str, Any]:
+    """Prune non-assigned branch paths while keeping shared downstream nodes."""
+    branch_id = str(branch_node_id)
+
+    assigned_slots = set()
+    if isinstance(assigned_branch, (list, tuple, set)):
+        for value in assigned_branch:
+            try:
+                idx = int(value)
+            except (TypeError, ValueError):
+                debug_log(f"Prompt transform: invalid assigned branch slot ignored: {value}")
+                continue
+            if idx >= 0:
+                assigned_slots.add(idx)
+    else:
+        try:
+            idx = int(assigned_branch)
+        except (TypeError, ValueError):
+            idx = -1
+        if idx >= 0:
+            assigned_slots.add(idx)
+
+    try:
+        total_slots = int(num_branches)
+    except (TypeError, ValueError):
+        total_slots = 2
+    total_slots = max(2, min(total_slots, 10))
+
+    downstream_by_slot = {}
+    for slot_idx in range(total_slots):
+        downstream = _find_downstream_of_output_slot(prompt_obj, branch_id, slot_idx)
+        downstream.discard(branch_id)
+        downstream_by_slot[slot_idx] = downstream
+
+    keep = {branch_id}
+    for slot_idx in assigned_slots:
+        keep.update(downstream_by_slot.get(slot_idx, set()))
+
+    remove = set()
+    for slot_idx in range(total_slots):
+        if slot_idx in assigned_slots:
+            continue
+        remove.update(downstream_by_slot.get(slot_idx, set()))
+
+    remove -= keep
+    for node_id in remove:
+        prompt_obj.pop(node_id, None)
+
+    _remove_dangling_input_refs(prompt_obj)
+    return prompt_obj
+
+
+def prune_prompt_for_worker(prompt_obj: dict[str, Any]) -> dict[str, Any]:
+    """Prune worker prompt to distributed nodes and their upstream dependencies."""
+    collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
+    branch_collector_ids = find_nodes_by_class(prompt_obj, "DistributedBranchCollector")
+    upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
+    branch_ids = find_nodes_by_class(prompt_obj, "DistributedBranch")
+    distributed_ids = collector_ids + branch_collector_ids + upscale_ids + branch_ids
+    if not distributed_ids:
+        return prompt_obj
+
+    connected = _find_upstream_nodes(prompt_obj, distributed_ids)
+    if branch_ids:
+        connected.update(_find_downstream_nodes(prompt_obj, branch_ids))
+    if branch_collector_ids:
+        downstream_of_collectors = _find_downstream_nodes(prompt_obj, branch_collector_ids)
+        connected -= (downstream_of_collectors - set(branch_collector_ids))
+
+    pruned_prompt = {}
+    for node_id in connected:
+        node = prompt_obj.get(node_id)
+        if node is not None:
+            pruned_prompt[node_id] = json.loads(json.dumps(node))
+
+    # Generate IDs from the original prompt so we never reuse IDs from pruned downstream nodes.
+    next_id = _create_numeric_id_generator(prompt_obj)
+    for dist_id in distributed_ids:
+        if dist_id not in pruned_prompt:
+            continue
+        downstream = _find_downstream_nodes(prompt_obj, [dist_id])
+        has_removed_downstream = any(node_id != dist_id and node_id not in connected for node_id in downstream)
+        if has_removed_downstream:
+            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
+
+
+def prepare_delegate_master_prompt(
+    prompt_obj: dict[str, Any],
+    collector_ids: list[str],
+) -> dict[str, Any]:
+    """Prune master prompt so it only executes post-collector nodes in delegate mode."""
+    downstream = _find_downstream_nodes(prompt_obj, collector_ids)
+    nodes_to_keep = set(collector_ids)
+    nodes_to_keep.update(downstream)
+
+    pruned_prompt = {}
+    for node_id in nodes_to_keep:
+        node = prompt_obj.get(node_id)
+        if node is not None:
+            pruned_prompt[node_id] = json.loads(json.dumps(node))
+
+    pruned_ids = set(pruned_prompt.keys())
+    for node_id, node in pruned_prompt.items():
+        inputs = node.get("inputs")
+        if not inputs:
+            continue
+        for input_name, input_value in list(inputs.items()):
+            if isinstance(input_value, list) and len(input_value) == 2:
+                source_id = str(input_value[0])
+                if source_id not in pruned_ids:
+                    inputs.pop(input_name, None)
+                    debug_log(
+                        f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
+                    )
+
+    # Generate IDs from the original prompt to avoid ID collisions with pruned nodes.
+    next_id = _create_numeric_id_generator(prompt_obj)
+    for collector_id in collector_ids:
+        collector_entry = pruned_prompt.get(collector_id)
+        if not collector_entry:
+            continue
+        collector_class = collector_entry.get("class_type")
+        if collector_class == "DistributedBranchCollector":
+            # Branch collector does not require an image input in delegate-only mode.
+            continue
+        if collector_class == "UltimateSDUpscaleDistributed":
+            # Swap to lightweight delegate collector — no model inputs needed.
+            collector_entry["class_type"] = "USDUDelegateCollector"
+            debug_log(
+                f"Swapped USDU node {collector_id} to USDUDelegateCollector for delegate-only master prompt."
+            )
+            continue
+        placeholder_id = next_id()
+        pruned_prompt[placeholder_id] = {
+            "class_type": "DistributedEmptyImage",
+            "inputs": {
+                "height": 64,
+                "width": 64,
+                "channels": 3,
+            },
+            "_meta": {
+                "title": "Distributed Empty Image (auto-added)",
+            },
+        }
+        collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
+        debug_log(
+            f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
+        )
+
+    return pruned_prompt
+
+
+def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str]:
+    """Create stable per-node job IDs for distributed nodes."""
+    job_map = {}
+    distributed_nodes = (
+        prompt_index.nodes_for_class("DistributedCollector")
+        + prompt_index.nodes_for_class("DistributedBranchCollector")
+        + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
+        + prompt_index.nodes_for_class("DistributedBranch")
+    )
+    for node_id in distributed_nodes:
+        job_map[node_id] = f"{prefix}_{node_id}"
+    return job_map
+
+
+def _override_simple_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json, class_names):
+    """Configure simple distributed nodes (DistributedSeed, DistributedValue)."""
+    for class_name in class_names:
+        for node_id in prompt_index.nodes_for_class(class_name):
+            node = prompt_copy.get(node_id)
+            if not isinstance(node, dict):
+                continue
+            inputs = node.setdefault("inputs", {})
+            inputs["is_worker"] = not is_master
+            inputs["enabled_worker_ids"] = enabled_json
+            inputs["worker_id"] = "" if is_master else str(participant_id)
+
+
+def _override_job_nodes(
+    prompt_copy, prompt_index, is_master, participant_id,
+    job_id_map, master_url, enabled_json, delegate_master,
+    class_names, skip_if_upstream=None,
+):
+    """Configure job-aware distributed nodes (collector, upscale)."""
+    for class_name in class_names:
+        for node_id in prompt_index.nodes_for_class(class_name):
+            node = prompt_copy.get(node_id)
+            if not isinstance(node, dict):
+                continue
+
+            if skip_if_upstream and prompt_index.has_upstream(node_id, skip_if_upstream):
+                node.setdefault("inputs", {})["pass_through"] = True
+                continue
+
+            inputs = node.setdefault("inputs", {})
+            inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
+            inputs["is_worker"] = not is_master
+            inputs["enabled_worker_ids"] = enabled_json
+            if is_master:
+                inputs["delegate_only"] = bool(delegate_master)
+                inputs.pop("master_url", None)
+                inputs.pop("worker_id", None)
+            else:
+                inputs["master_url"] = master_url
+                inputs["worker_id"] = participant_id
+                inputs["delegate_only"] = False
+
+
+def _override_branch_nodes(
+    prompt_copy,
+    prompt_index,
+    is_master,
+    participant_id,
+    enabled_worker_ids,
+    job_id_map,
+    delegate_master,
+):
+    """Assign branch slots per participant and prune non-assigned branch paths."""
+    participants = _resolve_participants(enabled_worker_ids, delegate_master)
+    participant_id = str(participant_id)
+    participant_count = len(participants)
+    participant_pos = participants.index(participant_id) if participant_id in participants else None
+
+    for node_id in prompt_index.nodes_for_class("DistributedBranch"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+
+        inputs = node.setdefault("inputs", {})
+        try:
+            num_branches = int(inputs.get("num_branches", 2))
+        except (TypeError, ValueError):
+            num_branches = 2
+        num_branches = max(2, min(num_branches, 10))
+
+        assigned_slots = []
+        if participant_pos is not None and participant_count > 0:
+            assigned_slots = [slot for slot in range(num_branches) if slot % participant_count == participant_pos]
+
+        if len(assigned_slots) == 1:
+            assigned_branch = assigned_slots[0]
+        else:
+            assigned_branch = -1
+
+        inputs["is_worker"] = not is_master
+        inputs["worker_id"] = "" if is_master else participant_id
+        inputs["assigned_branch"] = assigned_branch
+        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
+
+        if is_master and delegate_master:
+            continue
+
+        prompt_copy = prune_prompt_for_branch_worker(
+            prompt_copy,
+            node_id,
+            assigned_slots,
+            num_branches,
+        )
+
+    return prompt_copy
+
+
+def _find_upstream_branch_node(prompt_copy, start_node_id):
+    """Locate the nearest upstream DistributedBranch node for the provided node."""
+    visited = set()
+    stack = [str(start_node_id)]
+    while stack:
+        node_id = stack.pop()
+        if node_id in visited:
+            continue
+        visited.add(node_id)
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+
+        inputs = node.get("inputs", {})
+        for input_value in inputs.values():
+            if not (isinstance(input_value, list) and len(input_value) == 2):
+                continue
+            source_id = str(input_value[0])
+            source_node = prompt_copy.get(source_id)
+            if not isinstance(source_node, dict):
+                continue
+            if source_node.get("class_type") == "DistributedBranch":
+                return source_id
+            stack.append(source_id)
+    return None
+
+
+def _override_branch_collector_nodes(
+    prompt_copy,
+    prompt_index,
+    is_master,
+    participant_id,
+    job_id_map,
+    master_url,
+    enabled_json,
+    delegate_master,
+):
+    """Configure DistributedBranchCollector nodes for branch convergence."""
+    for node_id in prompt_index.nodes_for_class("DistributedBranchCollector"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+
+        upstream_branch_id = _find_upstream_branch_node(prompt_copy, node_id)
+        assigned_branch = -1
+        if upstream_branch_id is not None:
+            branch_node = prompt_copy.get(upstream_branch_id, {})
+            branch_inputs = branch_node.get("inputs", {}) if isinstance(branch_node, dict) else {}
+            try:
+                assigned_branch = int(branch_inputs.get("assigned_branch", -1))
+            except (TypeError, ValueError):
+                assigned_branch = -1
+
+        inputs = node.setdefault("inputs", {})
+        # Group all branch collectors under the same DistributedBranch into one queue.
+        job_source_id = upstream_branch_id or node_id
+        inputs["multi_job_id"] = job_id_map.get(job_source_id, job_source_id)
+        inputs["is_worker"] = not is_master
+        inputs["enabled_worker_ids"] = enabled_json
+        inputs["assigned_branch"] = assigned_branch
+        if is_master:
+            inputs["delegate_only"] = bool(delegate_master)
+            inputs.pop("master_url", None)
+            inputs.pop("worker_id", None)
+        else:
+            inputs["master_url"] = master_url
+            inputs["worker_id"] = str(participant_id)
+            inputs["delegate_only"] = False
+
+
+def apply_participant_overrides(
+    prompt_copy: dict[str, Any],
+    participant_id: str,
+    enabled_worker_ids: list[str],
+    job_id_map: dict[str, str],
+    master_url: str,
+    delegate_master: bool,
+    prompt_index: PromptIndex,
+) -> dict[str, Any]:
+    """Return a prompt copy with hidden inputs configured for master/worker."""
+    is_master = participant_id == "master"
+    enabled_json = json.dumps(enabled_worker_ids)
+
+    _override_simple_nodes(
+        prompt_copy, prompt_index, is_master, participant_id, enabled_json,
+        ["DistributedSeed", "DistributedValue"],
+    )
+    _override_job_nodes(
+        prompt_copy, prompt_index, is_master, participant_id,
+        job_id_map, master_url, enabled_json, delegate_master,
+        ["DistributedCollector"],
+        skip_if_upstream="UltimateSDUpscaleDistributed",
+    )
+    _override_job_nodes(
+        prompt_copy, prompt_index, is_master, participant_id,
+        job_id_map, master_url, enabled_json, delegate_master,
+        ["UltimateSDUpscaleDistributed"],
+    )
+    prompt_copy = _override_branch_nodes(
+        prompt_copy,
+        prompt_index,
+        is_master,
+        participant_id,
+        enabled_worker_ids,
+        job_id_map,
+        delegate_master,
+    )
+    _override_branch_collector_nodes(
+        prompt_copy,
+        prompt_index,
+        is_master,
+        participant_id,
+        job_id_map,
+        master_url,
+        enabled_json,
+        delegate_master,
+    )
+
+    if not is_master:
+        prompt_copy = _prune_worker_downstream_of_branch_collectors(prompt_copy)
+        prompt_copy = _ensure_worker_output_node(prompt_copy)
+
+    return prompt_copy

api/queue_orchestration.py

@@ -0,0 +1,498 @@
+import asyncio
+import time
+import uuid
+from typing import Any
+
+from ..utils.async_helpers import queue_prompt_payload
+from ..utils.config import load_config
+from ..utils.constants import (
+    ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
+    ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
+    ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
+    ORCHESTRATION_WORKER_PREP_CONCURRENCY,
+)
+from ..utils.logging import debug_log, log
+from ..utils.network import build_master_url
+from ..utils.runtime_state import ensure_distributed_runtime_state, get_prompt_server_instance
+from ..utils.trace_logger import trace_debug
+from .schemas import coerce_positive_float, coerce_positive_int
+from .orchestration.dispatch import (
+    dispatch_worker_prompt,
+    rank_workers_by_load,
+    select_active_workers,
+    select_least_busy_worker,
+)
+from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
+from .orchestration.prompt_transform import (
+    PromptIndex,
+    apply_participant_overrides,
+    find_nodes_by_class,
+    generate_job_id_map,
+    prepare_delegate_master_prompt,
+    prune_prompt_for_worker,
+)
+
+def _generate_execution_trace_id():
+    return f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
+
+
+def ensure_distributed_state(server_instance=None):
+    """Ensure prompt_server has the state used by distributed queue orchestration."""
+    ensure_distributed_runtime_state(server_instance)
+
+
+async def _ensure_distributed_queue(job_id):
+    """Ensure a queue exists for the given distributed job ID."""
+    runtime_state = ensure_distributed_runtime_state()
+    async with runtime_state.distributed_jobs_lock:
+        if job_id not in runtime_state.distributed_pending_jobs:
+            runtime_state.distributed_pending_jobs[job_id] = asyncio.Queue()
+
+
+def _resolve_enabled_workers(config, requested_ids=None):
+    """Return a list of worker configs that should participate."""
+    workers = []
+    for worker in config.get("workers", []):
+        worker_id = str(worker.get("id") or "").strip()
+        if not worker_id:
+            continue
+
+        if requested_ids is not None:
+            if worker_id not in requested_ids:
+                continue
+        elif not worker.get("enabled", False):
+            continue
+
+        raw_port = worker.get("port", worker.get("listen_port", 8188))
+        try:
+            port = int(raw_port or 8188)
+        except (TypeError, ValueError):
+            log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.")
+            port = 8188
+
+        workers.append(
+            {
+                "id": worker_id,
+                "name": worker.get("name", worker_id),
+                "host": worker.get("host"),
+                "port": port,
+                "type": worker.get("type", "local"),
+            }
+        )
+    return workers
+
+
+def _resolve_orchestration_limits(config):
+    """Resolve bounded concurrency/timeouts for worker preparation pipeline."""
+    settings = (config or {}).get("settings", {}) or {}
+    worker_probe_concurrency = coerce_positive_int(
+        settings.get("worker_probe_concurrency"),
+        ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
+    )
+    worker_prep_concurrency = coerce_positive_int(
+        settings.get("worker_prep_concurrency"),
+        ORCHESTRATION_WORKER_PREP_CONCURRENCY,
+    )
+    media_sync_concurrency = coerce_positive_int(
+        settings.get("media_sync_concurrency"),
+        ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
+    )
+    media_sync_timeout_seconds = coerce_positive_float(
+        settings.get("media_sync_timeout_seconds"),
+        ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
+    )
+    return (
+        worker_probe_concurrency,
+        worker_prep_concurrency,
+        media_sync_concurrency,
+        media_sync_timeout_seconds,
+    )
+
+
+def _is_load_balance_enabled(value):
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, (int, float)):
+        return bool(value)
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return False
+
+
+def _prompt_requests_load_balance(prompt_index):
+    for node_id in prompt_index.nodes_for_class("DistributedCollector"):
+        inputs = prompt_index.inputs_by_node.get(node_id, {})
+        if _is_load_balance_enabled(inputs.get("load_balance", False)):
+            return True
+    return False
+
+
+async def _prepare_worker_payload(
+    worker,
+    prompt_index,
+    enabled_ids,
+    job_id_map,
+    master_url,
+    delegate_master,
+    trace_execution_id,
+    worker_prep_semaphore,
+    media_sync_semaphore,
+    media_sync_timeout_seconds,
+):
+    """Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
+    async with worker_prep_semaphore:
+        worker_prompt = prompt_index.copy_prompt()
+
+        worker_type = str(worker.get("type") or "local").strip().lower()
+        is_remote_like = bool(worker.get("host")) and worker_type != "local"
+        if is_remote_like:
+            path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
+            if path_separator:
+                worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
+
+        worker_prompt = prune_prompt_for_worker(worker_prompt)
+        worker_prompt = apply_participant_overrides(
+            worker_prompt,
+            worker["id"],
+            enabled_ids,
+            job_id_map,
+            master_url,
+            delegate_master,
+            prompt_index,
+        )
+
+        if is_remote_like:
+            async with media_sync_semaphore:
+                try:
+                    await asyncio.wait_for(
+                        sync_worker_media(worker, worker_prompt, trace_execution_id=trace_execution_id),
+                        timeout=media_sync_timeout_seconds,
+                    )
+                except asyncio.TimeoutError:
+                    trace_debug(
+                        trace_execution_id,
+                        (
+                            f"Media sync timed out after {media_sync_timeout_seconds:.1f}s "
+                            f"for worker {worker.get('name')} ({worker.get('id')}); continuing dispatch."
+                        ),
+                    )
+
+        return worker, worker_prompt
+
+
+async def _select_execution_workers(
+    workers,
+    use_websocket,
+    delegate_master,
+    load_balance_requested,
+    has_branch_nodes,
+    master_url,
+    execution_trace_id,
+    worker_probe_concurrency,
+):
+    active_workers, delegate_master = await select_active_workers(
+        workers,
+        use_websocket,
+        delegate_master,
+        trace_execution_id=execution_trace_id,
+        probe_concurrency=worker_probe_concurrency,
+    )
+
+    if load_balance_requested:
+        candidate_workers = list(active_workers)
+        if not delegate_master:
+            candidate_workers.append(
+                {
+                    "id": "master",
+                    "name": "Master",
+                    "host": master_url,
+                    "type": "local",
+                }
+            )
+
+        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:
+            trace_debug(
+                execution_trace_id,
+                "Load-balance selection probe failed; using first available candidate.",
+            )
+            selected_worker = candidate_workers[0]
+
+        if selected_worker is not None and str(selected_worker.get("id")) == "master":
+            active_workers = []
+            delegate_master = False
+            trace_debug(
+                execution_trace_id,
+                "Load-balance selected master for execution (workers skipped).",
+            )
+        elif selected_worker is not None:
+            active_workers = [selected_worker]
+            delegate_master = True
+            trace_debug(
+                execution_trace_id,
+                f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).",
+            )
+        else:
+            trace_debug(
+                execution_trace_id,
+                "Load-balance requested but no execution candidates were available.",
+            )
+            active_workers = []
+            delegate_master = False
+
+    if has_branch_nodes and len(active_workers) > 1:
+        active_workers = await rank_workers_by_load(
+            active_workers,
+            trace_execution_id=execution_trace_id,
+            probe_concurrency=worker_probe_concurrency,
+        )
+
+    return active_workers, delegate_master
+
+
+def _distributed_queue_nodes(prompt_index):
+    return (
+        prompt_index.nodes_for_class("DistributedCollector")
+        + prompt_index.nodes_for_class("DistributedBranchCollector")
+        + prompt_index.nodes_for_class("DistributedBranch")
+        + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
+    )
+
+
+def _register_job_allowed_workers(job_id_map, enabled_ids):
+    runtime_state = ensure_distributed_runtime_state()
+    allowed = {str(worker_id) for worker_id in (enabled_ids or []) if str(worker_id).strip()}
+    for job_id in job_id_map.values():
+        if job_id:
+            runtime_state.distributed_job_allowed_workers[str(job_id)] = set(allowed)
+
+
+async def _ensure_job_queues(prompt_index, job_id_map):
+    for node_id in _distributed_queue_nodes(prompt_index):
+        job_id = job_id_map.get(node_id)
+        if job_id:
+            await _ensure_distributed_queue(job_id)
+
+
+def _build_master_prompt(
+    prompt_index,
+    enabled_ids,
+    job_id_map,
+    master_url,
+    delegate_master,
+):
+    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 not delegate_master:
+        return master_prompt
+
+    collector_ids = (
+        find_nodes_by_class(master_prompt, "DistributedCollector")
+        + find_nodes_by_class(master_prompt, "DistributedBranchCollector")
+    )
+    upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
+    # Include USDU nodes as collector-like for delegate pruning
+    collector_ids.extend(upscale_nodes)
+    if not collector_ids:
+        debug_log(
+            "Delegate-only master mode requested but no collector/branch-collector nodes found in master prompt. Running full prompt on master."
+        )
+        return master_prompt
+    return prepare_delegate_master_prompt(master_prompt, collector_ids)
+
+
+async def _prepare_worker_payloads_for_dispatch(
+    active_workers,
+    prompt_index,
+    enabled_ids,
+    job_id_map,
+    master_url,
+    delegate_master,
+    execution_trace_id,
+    worker_prep_concurrency,
+    media_sync_concurrency,
+    media_sync_timeout_seconds,
+):
+    if not active_workers:
+        return []
+
+    worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
+    media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
+    return await asyncio.gather(
+        *[
+            _prepare_worker_payload(
+                worker,
+                prompt_index,
+                enabled_ids,
+                job_id_map,
+                master_url,
+                delegate_master,
+                execution_trace_id,
+                worker_prep_semaphore,
+                media_sync_semaphore,
+                media_sync_timeout_seconds,
+            )
+            for worker in active_workers
+        ]
+    )
+
+
+async def _dispatch_worker_payloads(
+    worker_payloads,
+    workflow_meta,
+    client_id,
+    use_websocket,
+    execution_trace_id,
+):
+    if not worker_payloads:
+        return
+
+    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
+        ]
+    )
+
+
+async def orchestrate_distributed_execution(
+    prompt_obj: dict[str, Any],
+    workflow_meta: dict[str, Any] | None,
+    client_id: str | None,
+    enabled_worker_ids: list[str] | set[str] | None = None,
+    delegate_master: bool | None = None,
+    trace_execution_id: str | None = None,
+) -> tuple[str, int]:
+    """Core orchestration logic for the /distributed/queue endpoint.
+
+    Returns:
+        tuple[str, int]: (prompt_id, worker_count)
+    """
+    ensure_distributed_state()
+    execution_trace_id = trace_execution_id or _generate_execution_trace_id()
+
+    config = load_config()
+    use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
+    master_url = build_master_url(config=config, prompt_server_instance=get_prompt_server_instance())
+    (
+        worker_probe_concurrency,
+        worker_prep_concurrency,
+        media_sync_concurrency,
+        media_sync_timeout_seconds,
+    ) = _resolve_orchestration_limits(config)
+    requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
+    workers = _resolve_enabled_workers(config, requested_ids)
+    prompt_index = PromptIndex(prompt_obj)
+    load_balance_requested = _prompt_requests_load_balance(prompt_index)
+    has_branch_nodes = bool(prompt_index.nodes_for_class("DistributedBranch"))
+    trace_debug(
+        execution_trace_id,
+        (
+            f"Orchestration start: requested_workers={len(workers)}, "
+            f"requested_ids={requested_ids if requested_ids is not None else 'enabled_only'}, "
+            f"websocket={use_websocket}, "
+            f"probe_concurrency={worker_probe_concurrency}, "
+            f"prep_concurrency={worker_prep_concurrency}, "
+            f"media_sync_concurrency={media_sync_concurrency}, "
+            f"media_sync_timeout={media_sync_timeout_seconds:.1f}s, "
+            f"load_balance={load_balance_requested}, has_branch_nodes={has_branch_nodes}"
+        ),
+    )
+
+    # Respect master delegate-only configuration
+    if delegate_master is None:
+        delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
+
+    if not workers and delegate_master:
+        trace_debug(
+            execution_trace_id,
+            "Delegate-only requested but no workers are enabled. Falling back to master execution.",
+        )
+        delegate_master = False
+
+    active_workers, delegate_master = await _select_execution_workers(
+        workers=workers,
+        use_websocket=use_websocket,
+        delegate_master=delegate_master,
+        load_balance_requested=load_balance_requested,
+        has_branch_nodes=has_branch_nodes,
+        master_url=master_url,
+        execution_trace_id=execution_trace_id,
+        worker_probe_concurrency=worker_probe_concurrency,
+    )
+
+    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.")
+        prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
+        return prompt_id, 0
+
+    _register_job_allowed_workers(job_id_map, enabled_ids)
+    await _ensure_job_queues(prompt_index, job_id_map)
+
+    master_prompt = _build_master_prompt(
+        prompt_index=prompt_index,
+        enabled_ids=enabled_ids,
+        job_id_map=job_id_map,
+        master_url=master_url,
+        delegate_master=delegate_master,
+    )
+
+    if active_workers:
+        trace_debug(
+            execution_trace_id,
+            "Active distributed workers: "
+            + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
+        )
+    worker_payloads = await _prepare_worker_payloads_for_dispatch(
+        active_workers=active_workers,
+        prompt_index=prompt_index,
+        enabled_ids=enabled_ids,
+        job_id_map=job_id_map,
+        master_url=master_url,
+        delegate_master=delegate_master,
+        execution_trace_id=execution_trace_id,
+        worker_prep_concurrency=worker_prep_concurrency,
+        media_sync_concurrency=media_sync_concurrency,
+        media_sync_timeout_seconds=media_sync_timeout_seconds,
+    )
+    await _dispatch_worker_payloads(
+        worker_payloads=worker_payloads,
+        workflow_meta=workflow_meta,
+        client_id=client_id,
+        use_websocket=use_websocket,
+        execution_trace_id=execution_trace_id,
+    )
+
+    prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
+    trace_debug(
+        execution_trace_id,
+        f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
+    )
+    return prompt_id, len(worker_payloads)

api/queue_request.py

@@ -0,0 +1,84 @@
+from dataclasses import dataclass
+from typing import Any
+
+try:
+    from ..utils.worker_ids import require_enabled_worker_ids
+except ImportError:  # pragma: no cover - supports direct module loading in isolated tests.
+    from utils.worker_ids import require_enabled_worker_ids
+
+
+@dataclass(frozen=True)
+class QueueRequestPayload:
+    prompt: dict[str, Any]
+    workflow_meta: dict[str, Any] | None
+    client_id: str
+    delegate_master: bool | None
+    enabled_worker_ids: list[str]
+    trace_execution_id: str | None
+
+
+def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
+    """Parse and validate /distributed/queue payload into a normalized shape."""
+    if not isinstance(data, dict):
+        raise ValueError("Expected a JSON object body")
+
+    prompt = data.get("prompt")
+    # Auto-prepare is always on server-side; keep the field for wire compatibility.
+    if prompt is None:
+        workflow_payload = data.get("workflow")
+        if isinstance(workflow_payload, dict):
+            candidate_prompt = workflow_payload.get("prompt")
+            if isinstance(candidate_prompt, dict):
+                prompt = candidate_prompt
+
+    if not isinstance(prompt, dict):
+        raise ValueError("Field 'prompt' must be an object")
+
+    workflow_meta = data.get("workflow")
+    if workflow_meta is not None and not isinstance(workflow_meta, dict):
+        raise ValueError("Field 'workflow' must be an object when provided")
+
+    enabled_ids_raw = data.get("enabled_worker_ids")
+    workers_field = data.get("workers")
+    if enabled_ids_raw is None and workers_field is not None:
+        if not isinstance(workers_field, list):
+            raise ValueError("Field 'workers' must be a list when provided")
+        enabled_ids_raw = []
+        for entry in workers_field:
+            worker_id = entry.get("id") if isinstance(entry, dict) else entry
+            if worker_id is not None:
+                enabled_ids_raw.append(str(worker_id))
+
+    if enabled_ids_raw is None:
+        raise ValueError("enabled_worker_ids required")
+    else:
+        if not isinstance(enabled_ids_raw, list):
+            raise ValueError("enabled_worker_ids must be a list of worker IDs")
+        try:
+            enabled_ids = require_enabled_worker_ids(enabled_ids_raw)
+        except ValueError as exc:
+            raise ValueError(f"enabled_worker_ids invalid: {exc}") from exc
+
+    delegate_master = data.get("delegate_master")
+    if delegate_master is not None and not isinstance(delegate_master, bool):
+        raise ValueError("delegate_master must be a boolean when provided")
+
+    client_id = data.get("client_id")
+    if not isinstance(client_id, str) or not client_id.strip():
+        raise ValueError("client_id required")
+    client_id = client_id.strip()
+
+    trace_execution_id = data.get("trace_execution_id")
+    if trace_execution_id is not None:
+        if not isinstance(trace_execution_id, str):
+            raise ValueError("trace_execution_id must be a string when provided")
+        trace_execution_id = trace_execution_id.strip() or None
+
+    return QueueRequestPayload(
+        prompt=prompt,
+        workflow_meta=workflow_meta,
+        client_id=client_id,
+        delegate_master=delegate_master,
+        enabled_worker_ids=enabled_ids,
+        trace_execution_id=trace_execution_id,
+    )

api/request_guards.py

@@ -0,0 +1,15 @@
+"""Shared request guards for distributed API endpoints."""
+from __future__ import annotations
+
+from aiohttp import web
+
+from ..utils.config import load_config
+from ..utils.network import handle_api_error
+from .schemas import is_authorized_request
+
+
+async def authorization_error_or_none(request: web.Request) -> web.StreamResponse | None:
+    """Return a 403 response when distributed API auth fails, else None."""
+    if is_authorized_request(request, load_config()):
+        return None
+    return await handle_api_error(request, "Unauthorized", 403)

api/schemas.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from typing import Any
+
+from ..utils.auth import (
+    AUTH_HEADER_NAME,
+    distributed_auth_headers,
+    is_authorized_request,
+)
+from ..utils.parsing import coerce_bool, require_bool_literal
+
+__all__ = [
+    "AUTH_HEADER_NAME",
+    "coerce_bool",
+    "coerce_positive_float",
+    "coerce_positive_int",
+    "distributed_auth_headers",
+    "is_authorized_request",
+    "require_bool_literal",
+    "require_fields",
+    "require_worker_id",
+    "require_positive_int",
+    "validate_worker_id",
+]
+
+
+def require_fields(data: dict, *fields) -> list[str]:
+    """Return field names that are missing or empty in a JSON object."""
+    if not isinstance(data, dict):
+        return list(fields)
+
+    missing = []
+    for field in fields:
+        if field not in data:
+            missing.append(field)
+            continue
+        value = data.get(field)
+        if value is None:
+            missing.append(field)
+            continue
+        if isinstance(value, str) and not value.strip():
+            missing.append(field)
+
+    return missing
+
+
+def validate_worker_id(worker_id: str, config: dict[str, Any]) -> bool:
+    """Return True when worker_id exists in config['workers']."""
+    try:
+        require_worker_id(worker_id, config)
+        return True
+    except ValueError:
+        return False
+
+
+def require_worker_id(worker_id: str, config: dict[str, Any], field_name: str = "worker_id") -> str:
+    """Strict worker-id validator for request boundaries."""
+    worker_id_str = str(worker_id).strip()
+    if not worker_id_str:
+        raise ValueError(f"Field '{field_name}' must be a non-empty worker id.")
+
+    workers = (config or {}).get("workers", [])
+    for worker in workers:
+        if not isinstance(worker, dict):
+            continue
+        if str(worker.get("id")).strip() == worker_id_str:
+            return worker_id_str
+    raise ValueError(f"Worker {worker_id_str} not found.")
+
+
+def require_positive_int(value, field_name: str = "value") -> int:
+    """Strict positive-int parser for request boundary validation."""
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Field '{field_name}' must be a positive integer.") from exc
+    if parsed <= 0:
+        raise ValueError(f"Field '{field_name}' must be a positive integer.")
+    return parsed
+
+
+def coerce_positive_int(value, default: int) -> int:
+    """Permissive positive-int coercion used for non-boundary defaults."""
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError):
+        return max(1, int(default))
+    return max(1, parsed)
+
+
+def coerce_positive_float(value, default: float) -> float:
+    """Permissive positive-float coercion used for non-boundary defaults."""
+    fallback = float(default)
+    if fallback <= 0.0:
+        fallback = 0.1
+    try:
+        parsed = float(value)
+    except (TypeError, ValueError):
+        return fallback
+    return parsed if parsed > 0.0 else fallback

api/tunnel_routes.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from aiohttp import web
+import server
+
+from .endpoint_policy import run_authorized_endpoint
+from ..utils.cloudflare import cloudflare_tunnel_manager
+from ..utils.config import load_config
+
+
+@server.PromptServer.instance.routes.get("/distributed/tunnel/status")
+async def tunnel_status_endpoint(request: web.Request) -> web.StreamResponse:
+    """Return Cloudflare tunnel status and last known details."""
+
+    async def _operation() -> web.StreamResponse:
+        status = cloudflare_tunnel_manager.get_status()
+        config = load_config()
+        master_host = (config.get("master") or {}).get("host")
+        return web.json_response(
+            {
+                "status": "success",
+                "tunnel": status,
+                "master_host": master_host,
+            }
+        )
+
+    return await run_authorized_endpoint(request, _operation)
+
+
+@server.PromptServer.instance.routes.post("/distributed/tunnel/start")
+async def tunnel_start_endpoint(request: web.Request) -> web.StreamResponse:
+    """Start a Cloudflare tunnel pointing at the current ComfyUI server."""
+
+    async def _operation() -> web.StreamResponse:
+        result = await cloudflare_tunnel_manager.start_tunnel()
+        config = load_config()
+        return web.json_response(
+            {
+                "status": "success",
+                "tunnel": result,
+                "master_host": (config.get("master") or {}).get("host"),
+            }
+        )
+
+    return await run_authorized_endpoint(request, _operation)
+
+
+@server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
+async def tunnel_stop_endpoint(request: web.Request) -> web.StreamResponse:
+    """Stop the managed Cloudflare tunnel if running."""
+
+    async def _operation() -> web.StreamResponse:
+        result = await cloudflare_tunnel_manager.stop_tunnel()
+        config = load_config()
+        return web.json_response(
+            {
+                "status": "success",
+                "tunnel": result,
+                "master_host": (config.get("master") or {}).get("host"),
+            }
+        )
+
+    return await run_authorized_endpoint(request, _operation)

api/usdu_routes.py

@@ -0,0 +1,344 @@
+from __future__ import annotations
+
+import asyncio
+import io
+import time
+
+from aiohttp import web
+from PIL import Image
+import server
+
+from .request_guards import authorization_error_or_none
+from .schemas import require_bool_literal
+from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
+from ..upscale.job_store import ensure_tile_jobs_initialized, init_dynamic_job
+from ..upscale.payload_parsers import parse_tiles_from_form
+from ..utils.logging import debug_log
+from ..utils.network import handle_api_error
+from ..utils.usdu_management import MAX_PAYLOAD_SIZE
+
+
+def _parse_int_field(value, field_name: str, *, minimum: int | None = None) -> int:
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"{field_name} must be an integer") from exc
+    if minimum is not None and parsed < minimum:
+        raise ValueError(f"{field_name} must be >= {minimum}")
+    return parsed
+
+
+def _parse_bool_field(
+    value,
+    field_name: str,
+    *,
+    default: bool = False,
+) -> bool:
+    if value is None:
+        return default
+    return require_bool_literal(value, field_name=field_name)
+
+
+def _is_worker_allowed(job_data: BaseJobState, worker_id: str) -> bool:
+    known_workers = {str(wid) for wid in getattr(job_data, "worker_status", {}).keys()}
+    known_workers.update(str(wid) for wid in getattr(job_data, "assigned_to_workers", {}).keys())
+    if not known_workers:
+        return True
+    return str(worker_id) in known_workers
+
+
+@server.PromptServer.instance.routes.post("/distributed/init_dynamic_job")
+async def init_dynamic_job_endpoint(request: web.Request) -> web.StreamResponse:
+    """Allow workers to initialize a dynamic job queue on the master.
+
+    Called by the first worker to reach the USDU node in delegate-only mode.
+    Idempotent — subsequent calls are no-ops if the queue already exists.
+    """
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        data = await request.json()
+        multi_job_id = data.get("multi_job_id")
+        if not multi_job_id:
+            return await handle_api_error(request, "Missing multi_job_id", 400)
+
+        batch_size = _parse_int_field(data.get("batch_size"), "batch_size", minimum=1)
+        enabled_workers = data.get("enabled_workers") or []
+        if isinstance(enabled_workers, str):
+            import json as _json
+            enabled_workers = _json.loads(enabled_workers)
+
+        await init_dynamic_job(multi_job_id, batch_size, enabled_workers)
+        return web.json_response({"status": "success"})
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.post("/distributed/heartbeat")
+async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse:
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        data = await request.json()
+        worker_id = data.get('worker_id')
+        multi_job_id = data.get('multi_job_id')
+
+        if not worker_id or not multi_job_id:
+            return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
+
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                if isinstance(job_data, BaseJobState):
+                    if not _is_worker_allowed(job_data, worker_id):
+                        return await handle_api_error(request, "Unauthorized worker_id", 403)
+                    job_data.worker_status[worker_id] = time.time()
+                    debug_log(f"Heartbeat from worker {worker_id}")
+                    return web.json_response({"status": "success"})
+                return await handle_api_error(request, "Worker status tracking not available", 400)
+            return await handle_api_error(request, "Job not found", 404)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
+async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
+    """Endpoint for workers to submit processed tiles in static mode."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        content_length = request.headers.get('content-length')
+        if content_length:
+            try:
+                payload_size = int(content_length)
+            except ValueError:
+                return await handle_api_error(request, "Invalid content-length header", 400)
+            if payload_size > MAX_PAYLOAD_SIZE:
+                return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
+
+        data = await request.post()
+        multi_job_id = data.get('multi_job_id')
+        worker_id = data.get('worker_id')
+        try:
+            is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False)
+        except ValueError as e:
+            return await handle_api_error(request, str(e), 400)
+
+        if multi_job_id is None or worker_id is None:
+            return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
+
+        prompt_server = ensure_tile_jobs_initialized()
+
+        try:
+            batch_size = _parse_int_field(data.get('batch_size', 0), "batch_size", minimum=0)
+        except ValueError as e:
+            return await handle_api_error(request, str(e), 400)
+
+        # Handle completion signal
+        if batch_size == 0 and is_last:
+            async with prompt_server.distributed_tile_jobs_lock:
+                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                    if not isinstance(job_data, TileJobState):
+                        return await handle_api_error(request, "Job not configured for tile submissions", 400)
+                    if not _is_worker_allowed(job_data, worker_id):
+                        return await handle_api_error(request, "Unauthorized worker_id", 403)
+                    await job_data.queue.put({
+                        'worker_id': worker_id,
+                        'is_last': True,
+                        'tiles': [],
+                    })
+                    debug_log(f"Received completion signal from worker {worker_id}")
+                    return web.json_response({"status": "success"})
+
+        try:
+            tiles = parse_tiles_from_form(data)
+        except ValueError as e:
+            return await handle_api_error(request, str(e), 400)
+
+        # Submit tiles to queue
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                if not isinstance(job_data, TileJobState):
+                    return await handle_api_error(request, "Job not configured for tile submissions", 400)
+                if not _is_worker_allowed(job_data, worker_id):
+                    return await handle_api_error(request, "Unauthorized worker_id", 403)
+
+                q = job_data.queue
+                if batch_size > 0 or len(tiles) > 0:
+                    await q.put({
+                        'worker_id': worker_id,
+                        'tiles': tiles,
+                        'is_last': is_last,
+                    })
+                    debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
+                else:
+                    await q.put({
+                        'worker_id': worker_id,
+                        'is_last': True,
+                        'tiles': [],
+                    })
+
+                return web.json_response({"status": "success"})
+            return await handle_api_error(request, "Job not found", 404)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.post("/distributed/submit_image")
+async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
+    """Endpoint for workers to submit processed images in dynamic mode."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        content_length = request.headers.get('content-length')
+        if content_length:
+            try:
+                payload_size = int(content_length)
+            except ValueError:
+                return await handle_api_error(request, "Invalid content-length header", 400)
+            if payload_size > MAX_PAYLOAD_SIZE:
+                return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
+
+        data = await request.post()
+        multi_job_id = data.get('multi_job_id')
+        worker_id = data.get('worker_id')
+        try:
+            is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False)
+        except ValueError as e:
+            return await handle_api_error(request, str(e), 400)
+
+        if multi_job_id is None or worker_id is None:
+            return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
+
+        prompt_server = ensure_tile_jobs_initialized()
+
+        # Handle image submission
+        if 'full_image' in data and 'image_idx' in data:
+            try:
+                image_idx = _parse_int_field(data.get('image_idx'), "image_idx", minimum=0)
+            except ValueError as e:
+                return await handle_api_error(request, str(e), 400)
+            img_data = data['full_image'].file.read()
+            img = Image.open(io.BytesIO(img_data)).convert("RGB")
+
+            debug_log(f"Received full image {image_idx} from worker {worker_id}")
+
+            async with prompt_server.distributed_tile_jobs_lock:
+                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                    if not isinstance(job_data, ImageJobState):
+                        return await handle_api_error(request, "Job not configured for image submissions", 400)
+                    if not _is_worker_allowed(job_data, worker_id):
+                        return await handle_api_error(request, "Unauthorized worker_id", 403)
+                    await job_data.queue.put({
+                        'worker_id': worker_id,
+                        'image_idx': image_idx,
+                        'image': img,
+                        'is_last': is_last,
+                    })
+                    return web.json_response({"status": "success"})
+
+        # Handle completion signal (no image data)
+        elif is_last:
+            async with prompt_server.distributed_tile_jobs_lock:
+                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                    if not isinstance(job_data, ImageJobState):
+                        return await handle_api_error(request, "Job not configured for image submissions", 400)
+                    if not _is_worker_allowed(job_data, worker_id):
+                        return await handle_api_error(request, "Unauthorized worker_id", 403)
+                    await job_data.queue.put({
+                        'worker_id': worker_id,
+                        'is_last': True,
+                    })
+                    debug_log(f"Received completion signal from worker {worker_id}")
+                    return web.json_response({"status": "success"})
+        else:
+            return await handle_api_error(request, "Missing image data or invalid request", 400)
+
+        return await handle_api_error(request, "Job not found", 404)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.post("/distributed/request_image")
+async def request_image_endpoint(request: web.Request) -> web.StreamResponse:
+    """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        data = await request.json()
+        worker_id = data.get('worker_id')
+        multi_job_id = data.get('multi_job_id')
+
+        if not worker_id or not multi_job_id:
+            return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
+
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                if not isinstance(job_data, BaseJobState):
+                    return await handle_api_error(request, "Invalid job data structure", 500)
+
+                mode = job_data.mode
+                if isinstance(job_data, ImageJobState):
+                    pending_queue = job_data.pending_images
+                elif isinstance(job_data, TileJobState):
+                    pending_queue = job_data.pending_tasks
+                else:
+                    return await handle_api_error(request, "Invalid job configuration", 400)
+                if not _is_worker_allowed(job_data, worker_id):
+                    return await handle_api_error(request, "Unauthorized worker_id", 403)
+
+                try:
+                    task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
+                    job_data.assigned_to_workers.setdefault(worker_id, []).append(task_idx)
+                    job_data.worker_status[worker_id] = time.time()
+                    remaining = pending_queue.qsize()
+
+                    if mode == 'dynamic':
+                        debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
+                        return web.json_response(
+                            {
+                                "kind": "image",
+                                "task_idx": task_idx,
+                                "estimated_remaining": remaining,
+                            }
+                        )
+                    debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
+                    return web.json_response({
+                        "kind": "tile",
+                        "task_idx": task_idx,
+                        "estimated_remaining": remaining,
+                        "batched_static": job_data.batched_static,
+                    })
+                except asyncio.TimeoutError:
+                    return web.json_response({"kind": "none", "task_idx": None, "estimated_remaining": 0})
+            return await handle_api_error(request, "Job not found", 404)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.get("/distributed/job_status")
+async def job_status_endpoint(request: web.Request) -> web.StreamResponse:
+    """Endpoint to check if a job is ready."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    multi_job_id = request.query.get('multi_job_id')
+    if not multi_job_id:
+        return web.json_response({"ready": False})
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        ready = bool(isinstance(job_data, BaseJobState) and job_data.queue is not None)
+        return web.json_response({"ready": ready})

api/worker_routes.py

@@ -0,0 +1,756 @@
+from __future__ import annotations
+
+import json
+import asyncio
+import os
+import time
+import platform
+import subprocess  # nosec B404 - commands are fixed and never shell-expanded
+import socket
+from typing import Any
+
+import torch
+import aiohttp
+from aiohttp import web
+import server
+
+from ..utils.config import load_config
+from ..utils.logging import debug_log, log
+from ..utils.network import (
+    build_worker_url,
+    get_client_session,
+    handle_api_error,
+    normalize_host,
+    probe_worker,
+)
+from ..utils.constants import CHUNK_SIZE
+from ..workers import get_worker_manager
+from .request_guards import authorization_error_or_none
+from .schemas import (
+    distributed_auth_headers,
+    require_fields,
+    require_worker_id,
+)
+from ..workers.detection import (
+    get_machine_id,
+    is_docker_environment,
+    is_runpod_environment,
+)
+from ..utils.async_helpers import PromptValidationError, queue_prompt_payload
+
+
+@server.PromptServer.instance.routes.get("/distributed/worker_ws")
+async def worker_ws_endpoint(request: web.Request) -> web.StreamResponse:
+    """WebSocket endpoint for worker prompt dispatch."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+
+    ws = web.WebSocketResponse(heartbeat=30)
+    await ws.prepare(request)
+
+    async for msg in ws:
+        if msg.type == aiohttp.WSMsgType.TEXT:
+            try:
+                data = json.loads(msg.data or "{}")
+            except json.JSONDecodeError:
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": None,
+                    "ok": False,
+                    "error": "Invalid JSON payload.",
+                })
+                continue
+
+            if data.get("type") != "dispatch_prompt":
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": False,
+                    "error": "Unsupported websocket message type.",
+                })
+                continue
+
+            prompt = data.get("prompt")
+            if not isinstance(prompt, dict):
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": False,
+                    "error": "Field 'prompt' must be an object.",
+                })
+                continue
+
+            try:
+                prompt_id = await queue_prompt_payload(
+                    prompt,
+                    workflow_meta=data.get("workflow"),
+                    client_id=data.get("client_id"),
+                )
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": True,
+                    "prompt_id": prompt_id,
+                })
+            except PromptValidationError as exc:
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": False,
+                    "error": str(exc),
+                    "validation_error": exc.validation_error,
+                    "node_errors": exc.node_errors,
+                })
+            except Exception as exc:
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": False,
+                    "error": str(exc),
+                })
+        elif msg.type == aiohttp.WSMsgType.ERROR:
+            log(f"[Distributed] Worker websocket error: {ws.exception()}")
+
+    return ws
+
+
+@server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
+async def clear_launching_state(request: web.Request) -> web.StreamResponse:
+    """Clear the launching flag when worker is confirmed running."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        wm = get_worker_manager()
+        data = await request.json()
+        missing = require_fields(data, "worker_id")
+        if missing:
+            return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
+
+        worker_id = str(data.get("worker_id")).strip()
+        config = load_config()
+        try:
+            worker_id = require_worker_id(worker_id, config)
+        except ValueError as exc:
+            return await handle_api_error(request, exc, 404)
+        
+        # Clear launching flag in managed processes
+        if worker_id in wm.processes:
+            if 'launching' in wm.processes[worker_id]:
+                del wm.processes[worker_id]['launching']
+                wm.save_processes()
+                debug_log(f"Cleared launching state for worker {worker_id}")
+        
+        return web.json_response({"status": "success"})
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+def get_network_ips() -> list[str]:
+    """Get all network IPs, trying multiple methods."""
+    command_timeout = 5.0
+    ips = []
+    hostname = socket.gethostname()
+
+    # Method 1: Try socket.getaddrinfo
+    try:
+        addr_info = socket.getaddrinfo(hostname, None)
+        for info in addr_info:
+            ip = info[4][0]
+            if ip and ip not in ips and not ip.startswith('::'):  # Skip IPv6 for now
+                ips.append(ip)
+    except (socket.gaierror, OSError) as exc:
+        debug_log(f"get_network_ips: getaddrinfo failed for hostname {hostname}: {exc}")
+
+    # Method 2: Try to connect to external server and get local IP
+    try:
+        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        s.connect(("8.8.8.8", 80))  # Google DNS
+        local_ip = s.getsockname()[0]
+        s.close()
+        if local_ip not in ips:
+            ips.append(local_ip)
+    except (OSError, socket.error) as exc:
+        debug_log(f"get_network_ips: UDP local IP probe failed: {exc}")
+
+    # Method 3: Platform-specific commands
+    try:
+        if platform.system() == "Windows":
+            # Windows ipconfig
+            result = subprocess.run(  # nosec B603 - static command, no user input
+                ["ipconfig"],
+                capture_output=True,
+                text=True,
+                timeout=command_timeout,
+            )
+            lines = result.stdout.split('\n')
+            for i, line in enumerate(lines):
+                if 'IPv4' in line and i + 1 < len(lines):
+                    ip = lines[i].split(':')[-1].strip()
+                    if ip and ip not in ips:
+                        ips.append(ip)
+        else:
+            # Unix/Linux/Mac ifconfig or ip addr
+            try:
+                result = subprocess.run(  # nosec B603 - static command, no user input
+                    ["ip", "addr"],
+                    capture_output=True,
+                    text=True,
+                    timeout=command_timeout,
+                )
+            except (FileNotFoundError, OSError):
+                try:
+                    result = subprocess.run(  # nosec B603 - static command, no user input
+                        ["ifconfig"],
+                        capture_output=True,
+                        text=True,
+                        timeout=command_timeout,
+                    )
+                except (FileNotFoundError, OSError):
+                    result = None
+
+            import re
+            ip_pattern = re.compile(r'inet\s+(\d+\.\d+\.\d+\.\d+)')
+            if result is not None:
+                for match in ip_pattern.finditer(result.stdout):
+                    ip = match.group(1)
+                    if ip and ip not in ips:
+                        ips.append(ip)
+    except (OSError, subprocess.SubprocessError) as exc:
+        debug_log(f"get_network_ips: platform command probe failed: {exc}")
+
+    return ips
+
+
+def get_recommended_ip(ips: list[str]) -> str | None:
+    """Choose the best IP for master-worker communication."""
+    # Priority order:
+    # 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x)
+    # 2. Other non-localhost IPs
+    # 3. Localhost as last resort
+
+    private_ips = []
+    public_ips = []
+
+    for ip in ips:
+        if ip.startswith('127.') or ip == 'localhost':
+            continue
+        elif (ip.startswith('192.168.')
+                or ip.startswith('10.')
+                or (ip.startswith('172.') and 16 <= int(ip.split('.')[1]) <= 31)):
+            private_ips.append(ip)
+        else:
+            public_ips.append(ip)
+
+    # Prefer private IPs
+    if private_ips:
+        # Prefer 192.168 range as it's most common
+        for ip in private_ips:
+            if ip.startswith('192.168.'):
+                return ip
+        return private_ips[0]
+    elif public_ips:
+        return public_ips[0]
+    elif ips:
+        return ips[0]
+    else:
+        return None
+
+
+def _get_cuda_info() -> tuple[int | None, int, int]:
+    """Detect CUDA device index and total physical GPU count.
+
+    Returns (cuda_device, cuda_device_count, physical_device_count).
+    All three are 0/None if CUDA is unavailable.
+    """
+    if not torch.cuda.is_available():
+        return None, 0, 0
+    try:
+        cuda_device_count = torch.cuda.device_count()
+        cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES', '')
+        if cuda_visible and cuda_visible.strip():
+            visible_devices = [int(d.strip()) for d in cuda_visible.split(',') if d.strip().isdigit()]
+            if visible_devices:
+                cuda_device = visible_devices[0]
+                try:
+                    result = subprocess.run(  # nosec B603 - static command, no user input
+                        ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
+                        capture_output=True,
+                        text=True,
+                        timeout=5,
+                    )
+                    physical_device_count = (
+                        len(result.stdout.strip().split('\n'))
+                        if result.returncode == 0
+                        else max(visible_devices) + 1
+                    )
+                except (FileNotFoundError, OSError, subprocess.SubprocessError):
+                    physical_device_count = max(visible_devices) + 1
+                return cuda_device, cuda_device_count, physical_device_count
+            else:
+                return 0, cuda_device_count, cuda_device_count
+        else:
+            cuda_device = torch.cuda.current_device()
+            return cuda_device, cuda_device_count, cuda_device_count
+    except Exception as e:
+        debug_log(f"CUDA detection error: {e}")
+        return None, 0, 0
+
+
+def _collect_network_info_sync() -> dict[str, Any]:
+    """Collect network/cuda info in a worker thread to avoid blocking route handlers."""
+    cuda_device, cuda_device_count, physical_device_count = _get_cuda_info()
+    hostname = socket.gethostname()
+    all_ips = get_network_ips()
+    recommended_ip = get_recommended_ip(all_ips)
+    return {
+        "hostname": hostname,
+        "all_ips": all_ips,
+        "recommended_ip": recommended_ip,
+        "cuda_device": cuda_device,
+        "cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count,
+    }
+
+
+def _read_worker_log_sync(log_file: str, lines_to_read: int) -> dict[str, Any]:
+    """Read worker log content from disk in a threadpool worker."""
+    file_size = os.path.getsize(log_file)
+
+    with open(log_file, 'r', encoding='utf-8', errors='replace') as f:
+        if lines_to_read > 0 and file_size > 1024 * 1024:
+            # Read last N lines efficiently from end of file.
+            lines = []
+            f.seek(0, 2)
+            file_length = f.tell()
+            chunk_size = min(CHUNK_SIZE, file_length)
+
+            while len(lines) < lines_to_read and f.tell() > 0:
+                current_pos = max(0, f.tell() - chunk_size)
+                f.seek(current_pos)
+                chunk = f.read(chunk_size)
+                chunk_lines = chunk.splitlines()
+                if current_pos > 0:
+                    chunk_lines = chunk_lines[1:]
+                lines = chunk_lines + lines
+                f.seek(current_pos)
+
+            content = '\n'.join(lines[-lines_to_read:])
+            truncated = len(lines) > lines_to_read
+        else:
+            content = f.read()
+            truncated = False
+
+    return {
+        "content": content,
+        "file_size": file_size,
+        "truncated": truncated,
+        "lines_shown": lines_to_read if truncated else content.count('\n') + 1,
+    }
+
+
+def _parse_positive_int_query(
+    value: Any,
+    default: int,
+    minimum: int = 1,
+    maximum: int | None = 10000,
+) -> int:
+    """Parse bounded positive integer query params with sane fallback."""
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError):
+        return default
+    parsed = max(minimum, parsed)
+    if maximum is not None:
+        parsed = min(maximum, parsed)
+    return parsed
+
+
+def _find_worker_by_id(config: dict[str, Any], worker_id: str) -> dict[str, Any] | None:
+    worker_id_str = str(worker_id).strip()
+    for worker in config.get("workers", []):
+        if str(worker.get("id")).strip() == worker_id_str:
+            return worker
+    return None
+
+
+@server.PromptServer.instance.routes.get("/distributed/local_log")
+async def get_local_log_endpoint(request: web.Request) -> web.StreamResponse:
+    """Return this instance's in-memory ComfyUI log buffer."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        from app.logger import get_logs
+    except Exception as e:
+        return await handle_api_error(request, f"Failed to import app.logger: {e}", 500)
+
+    try:
+        lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000)
+        logs = get_logs()
+        if logs is None:
+            return web.json_response(
+                {
+                    "status": "success",
+                    "content": "",
+                    "entries": 0,
+                    "source": "memory",
+                    "truncated": False,
+                    "lines_shown": 0,
+                }
+            )
+
+        entries = list(logs)
+        selected_entries = entries[-lines_to_read:]
+        content = "".join(
+            entry.get("m", "") if isinstance(entry, dict) else str(entry)
+            for entry in selected_entries
+        )
+        lines_shown = content.count("\n") + (1 if content else 0)
+
+        return web.json_response(
+            {
+                "status": "success",
+                "content": content,
+                "entries": len(selected_entries),
+                "source": "memory",
+                "truncated": len(entries) > len(selected_entries),
+                "lines_shown": lines_shown,
+            }
+        )
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.get("/distributed/network_info")
+async def get_network_info_endpoint(request: web.Request) -> web.StreamResponse:
+    """Get network interfaces and recommend best IP for master."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        loop = asyncio.get_running_loop()
+        info = await loop.run_in_executor(None, _collect_network_info_sync)
+        
+        return web.json_response({
+            "status": "success",
+            **info,
+            "message": "Auto-detected network configuration"
+        })
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+@server.PromptServer.instance.routes.get("/distributed/system_info")
+async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse:
+    """Get system information including machine ID for local worker detection."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        import socket
+        
+        return web.json_response({
+            "status": "success",
+            "hostname": socket.gethostname(),
+            "machine_id": get_machine_id(),
+            "platform": {
+                "system": platform.system(),
+                "machine": platform.machine(),
+                "node": platform.node(),
+                "path_separator": os.sep,  # Add path separator
+                "os_name": os.name  # Add OS name (posix, nt, etc.)
+            },
+            "is_docker": is_docker_environment(),
+            "is_runpod": is_runpod_environment(),
+            "runpod_pod_id": os.environ.get('RUNPOD_POD_ID')
+        })
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+@server.PromptServer.instance.routes.post("/distributed/launch_worker")
+async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
+    """Launch a worker process from the UI."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        wm = get_worker_manager()
+        data = await request.json()
+        missing = require_fields(data, "worker_id")
+        if missing:
+            return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
+
+        worker_id = str(data.get("worker_id")).strip()
+        
+        # Find worker config
+        config = load_config()
+        try:
+            worker_id = require_worker_id(worker_id, config)
+        except ValueError as exc:
+            return await handle_api_error(request, exc, 404)
+        worker = next((w for w in config.get("workers", []) if str(w.get("id")) == worker_id), None)
+        if not worker:
+            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+        
+        # Ensure consistent string ID
+        worker_id_str = worker_id
+        
+        # Check if already running (managed by this instance)
+        if worker_id_str in wm.processes:
+            proc_info = wm.processes[worker_id_str]
+            process = proc_info.get('process')
+            
+            # Check if still running
+            is_running = False
+            if process:
+                is_running = process.poll() is None
+            else:
+                # Restored process without subprocess object
+                is_running = wm.is_process_running(proc_info['pid'])
+            
+            if is_running:
+                return await handle_api_error(request, "Worker already running (managed by UI)", 409)
+            else:
+                # Process is dead, remove it
+                del wm.processes[worker_id_str]
+                wm.save_processes()
+        
+        # Launch the worker
+        try:
+            loop = asyncio.get_running_loop()
+            pid = await loop.run_in_executor(None, wm.launch_worker, worker)
+            log_file = wm.processes[worker_id_str].get('log_file')
+            return web.json_response({
+                "status": "success",
+                "pid": pid,
+                "message": f"Worker {worker['name']} launched",
+                "log_file": log_file
+            })
+        except Exception as e:
+            return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500)
+            
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
+
+
+@server.PromptServer.instance.routes.post("/distributed/stop_worker")
+async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse:
+    """Stop a worker process that was launched from the UI."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        wm = get_worker_manager()
+        data = await request.json()
+        missing = require_fields(data, "worker_id")
+        if missing:
+            return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
+
+        worker_id = str(data.get("worker_id")).strip()
+        config = load_config()
+        try:
+            worker_id = require_worker_id(worker_id, config)
+        except ValueError as exc:
+            return await handle_api_error(request, exc, 404)
+
+        success, message = wm.stop_worker(worker_id)
+        
+        if success:
+            return web.json_response({"status": "success", "message": message})
+        else:
+            return await handle_api_error(
+                request,
+                message,
+                404 if "not managed" in message else 409,
+            )
+            
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
+
+
+@server.PromptServer.instance.routes.get("/distributed/managed_workers")
+async def get_managed_workers_endpoint(request: web.Request) -> web.StreamResponse:
+    """Get list of workers managed by this UI instance."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        managed = get_worker_manager().get_managed_workers()
+        return web.json_response({
+            "status": "success",
+            "managed_workers": managed
+        })
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.get("/distributed/local-worker-status")
+async def get_local_worker_status_endpoint(request: web.Request) -> web.StreamResponse:
+    """Check status of all local workers (localhost/no host specified)."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        config = load_config()
+        worker_statuses = {}
+        
+        for worker in config.get("workers", []):
+            # Only check local workers
+            host = normalize_host(worker.get("host")) or ""
+            if not host or host in ["localhost", "127.0.0.1"]:
+                worker_id = worker["id"]
+                port = worker["port"]
+                
+                # Check if worker is enabled
+                if not worker.get("enabled", False):
+                    worker_statuses[worker_id] = {
+                        "online": False,
+                        "enabled": False,
+                        "processing": False,
+                        "queue_count": 0
+                    }
+                    continue
+                
+                # Try to connect to worker
+                try:
+                    worker_url = build_worker_url(worker)
+                    data = await probe_worker(worker_url, timeout=2.0)
+                    if data is None:
+                        worker_statuses[worker_id] = {
+                            "online": False,
+                            "enabled": True,
+                            "processing": False,
+                            "queue_count": 0,
+                            "error": "Unavailable",
+                        }
+                        continue
+                    queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0)
+                    worker_statuses[worker_id] = {
+                        "online": True,
+                        "enabled": True,
+                        "processing": queue_remaining > 0,
+                        "queue_count": queue_remaining
+                    }
+                except asyncio.TimeoutError:
+                    worker_statuses[worker_id] = {
+                        "online": False,
+                        "enabled": True,
+                        "processing": False,
+                        "queue_count": 0,
+                        "error": "Timeout"
+                    }
+                except Exception as e:
+                    worker_statuses[worker_id] = {
+                        "online": False,
+                        "enabled": True,
+                        "processing": False,
+                        "queue_count": 0,
+                        "error": str(e)
+                    }
+        
+        return web.json_response({
+            "status": "success",
+            "worker_statuses": worker_statuses
+        })
+    except Exception as e:
+        debug_log(f"Error checking local worker status: {e}")
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}")
+async def get_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
+    """Get log content for a specific worker."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        wm = get_worker_manager()
+        worker_id = request.match_info['worker_id']
+        
+        # Ensure worker_id is string
+        worker_id = str(worker_id)
+        
+        # Check if we manage this worker
+        if worker_id not in wm.processes:
+            return await handle_api_error(request, f"Worker {worker_id} not managed by UI", 404)
+        
+        proc_info = wm.processes[worker_id]
+        log_file = proc_info.get('log_file')
+        
+        if not log_file or not os.path.exists(log_file):
+            return await handle_api_error(request, "Log file not found", 404)
+        
+        # Read last N lines (or full file if small)
+        lines_to_read = _parse_positive_int_query(request.query.get('lines'), default=1000)
+        
+        try:
+            loop = asyncio.get_running_loop()
+            payload = await loop.run_in_executor(None, _read_worker_log_sync, log_file, lines_to_read)
+            
+            return web.json_response({
+                "status": "success",
+                "content": payload["content"],
+                "log_file": log_file,
+                "file_size": payload["file_size"],
+                "truncated": payload["truncated"],
+                "lines_shown": payload["lines_shown"],
+            })
+            
+        except Exception as e:
+            return await handle_api_error(request, f"Error reading log file: {str(e)}", 500)
+            
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.get("/distributed/remote_worker_log/{worker_id}")
+async def get_remote_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
+    """Proxy a remote worker log request to the worker's local in-memory log endpoint."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        worker_id = str(request.match_info["worker_id"]).strip()
+        config = load_config()
+        worker = _find_worker_by_id(config, worker_id)
+        if not worker:
+            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+
+        # Remote log proxy is only meaningful for remote/cloud workers.
+        host = normalize_host(worker.get("host")) or ""
+        if not host:
+            return await handle_api_error(
+                request,
+                f"Worker {worker_id} is local; use /distributed/worker_log/{worker_id} instead.",
+                400,
+            )
+
+        lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000)
+        worker_url = build_worker_url(worker, "/distributed/local_log")
+        session = await get_client_session()
+        async with session.get(
+            worker_url,
+            params={"lines": str(lines_to_read)},
+            headers=distributed_auth_headers(config),
+            timeout=aiohttp.ClientTimeout(total=5),
+        ) as resp:
+            if resp.status >= 400:
+                body = await resp.text()
+                return await handle_api_error(
+                    request,
+                    f"Remote worker {worker_id} returned HTTP {resp.status}: {body[:400]}",
+                    resp.status,
+                )
+
+            try:
+                data = await resp.json()
+            except Exception as e:
+                return await handle_api_error(
+                    request,
+                    f"Remote worker {worker_id} returned invalid JSON: {e}",
+                    502,
+                )
+
+        return web.json_response(data)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)

bootstrap/__init__.py

@@ -0,0 +1,5 @@
+"""Bootstrap entrypoints for ComfyUI-Distributed."""
+
+from .entrypoint import build_node_mappings, initialize_runtime
+
+__all__ = ["build_node_mappings", "initialize_runtime"]

bootstrap/entrypoint.py

@@ -0,0 +1,66 @@
+"""Single startup owner for node surface and runtime initialization."""
+from __future__ import annotations
+
+import atexit
+import os
+from functools import lru_cache
+from typing import Any
+
+
+@lru_cache(maxsize=1)
+def build_node_mappings() -> tuple[dict[str, Any], dict[str, str]]:
+    """Build and cache the full exported node mapping surface."""
+    from ..nodes import (
+        NODE_CLASS_MAPPINGS as distributed_class_mappings,
+        NODE_DISPLAY_NAME_MAPPINGS as distributed_display_name_mappings,
+    )
+    from ..nodes.distributed_upscale import (
+        NODE_CLASS_MAPPINGS as upscale_class_mappings,
+        NODE_DISPLAY_NAME_MAPPINGS as upscale_display_name_mappings,
+    )
+
+    return (
+        {**distributed_class_mappings, **upscale_class_mappings},
+        {**distributed_display_name_mappings, **upscale_display_name_mappings},
+    )
+
+
+@lru_cache(maxsize=1)
+def _runtime_initialized() -> dict[str, bool]:
+    return {"done": False}
+
+
+def initialize_runtime(prompt_server: Any | None = None) -> None:
+    """Initialize distributed runtime side effects once."""
+    state = _runtime_initialized()
+    if state["done"]:
+        return
+
+    import server
+
+    from ..api import bootstrap_routes
+    from ..api.queue_orchestration import ensure_distributed_state
+    from ..upscale.job_store import ensure_tile_jobs_initialized
+    from ..utils.config import CONFIG_FILE, ensure_config_exists
+    from ..utils.logging import debug_log
+    from ..workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
+
+    ensure_config_exists()
+    bootstrap_routes()
+
+    server_instance = prompt_server if prompt_server is not None else server.PromptServer.instance
+    ensure_distributed_state(server_instance)
+    ensure_tile_jobs_initialized()
+
+    atexit.register(lambda: None)  # placeholder; real cleanup in sync_cleanup
+    if not os.environ.get("COMFYUI_IS_WORKER"):
+        atexit.register(sync_cleanup)
+        delayed_auto_launch()
+        register_async_signals()
+
+    state["done"] = True
+
+    node_mappings, _ = build_node_mappings()
+    debug_log("Loaded Distributed nodes.")
+    debug_log(f"Config file: {CONFIG_FILE}")
+    debug_log(f"Available nodes: {list(node_mappings.keys())}")

conftest.py

@@ -0,0 +1,7 @@
+"""Project-level pytest collection rules for plugin-style package layout."""
+
+collect_ignore = [
+    "__init__.py",
+    "distributed.py",
+    "distributed_upscale.py",
+]

docs/comfyui-distributed-api.md

@@ -0,0 +1,273 @@
+# ComfyUI-Distributed API (Experimental)
+
+This document describes the **public HTTP API** added to ComfyUI-Distributed to allow queueing *distributed* workflows from external tools (scripts, services, CI jobs, render farms, etc.) without using the ComfyUI web UI.
+
+## Demo
+
+- Video walkthrough: https://youtu.be/yiQlPd0MzLk
+
+## Examples Repository
+
+- Examples repo: https://github.com/umanets/ComfyUI-Distributed-API-examples.git
+
+---
+
+## Overview
+
+### What this adds
+
+- `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.
+  - 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.
+
+### What it does *not* add
+
+- Authentication/authorization.
+- A separate “job status” API for distributed results (you still use ComfyUI’s normal prompt history / websocket flow, and the existing `/distributed/queue_status/{job_id}` behavior for collector queues).
+
+---
+
+## Endpoint: `POST /distributed/queue`
+
+Queue a workflow for distributed execution.
+
+### URL
+
+- `http://<master-host>:<master-port>/distributed/queue`
+
+### Headers
+
+- `Content-Type: application/json`
+
+### Request Body
+
+```json
+{
+  "prompt": { "<node_id>": { "class_type": "...", "inputs": { } } },
+  "workflow": { },
+  "client_id": "external-client",
+  "delegate_master": false,
+  "enabled_worker_ids": ["1", "2"],
+  "workers": ["1", "2"],
+  "trace_execution_id": "exec_1700000000_ab12cd"
+}
+```
+
+#### Fields
+
+- `prompt` (required unless `workflow.prompt` is present, object)
+  - 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.
+- `client_id` (required, string)
+  - Passed through as `extra_data.client_id` (useful if you consume ComfyUI websocket events).
+- `delegate_master` (optional, boolean)
+  - If `true`, attempts “workers-only” execution for workflows based on `DistributedCollector`.
+  - Current limitation: delegate-only mode **does not support** `UltimateSDUpscaleDistributed` and will fall back to running the full prompt on master.
+- `enabled_worker_ids` (required, array of strings)
+  - The explicit worker IDs to consider for this run.
+- `workers` (optional, array of strings or objects with `id`)
+  - Transitional alias for `enabled_worker_ids` used by older clients.
+- `trace_execution_id` (optional, string)
+  - Passed through to orchestration logs.
+  - Server log lines include the marker as `[exec:<trace_execution_id>]`.
+
+##### How to get `enabled_worker_ids`
+
+Worker IDs come from the plugin config (`GET /distributed/config`) under `workers[].id`.
+
+Example (bash + `jq`):
+
+```bash
+curl -s "http://127.0.0.1:8188/distributed/config" \
+  | jq -r '.workers[] | "id=\(.id)\tname=\(.name)\tenabled=\(.enabled)\thost=\(.host)\tport=\(.port)\ttype=\(.type)"'
+```
+
+Example (PowerShell):
+
+```powershell
+$cfg = Invoke-RestMethod "http://127.0.0.1:8188/distributed/config"
+$cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -AutoSize
+```
+
+### Response Body
+
+```json
+{
+  "prompt_id": "<uuid>",
+  "worker_count": 2
+}
+```
+
+- `prompt_id` — the master prompt id queued into ComfyUI.
+- `worker_count` — number of workers that received a dispatched prompt (only those that passed the health check).
+
+### Status Codes
+
+- `200` — queued.
+- `400` — invalid JSON or invalid body.
+- `500` — orchestration/dispatch failure (see server logs for details).
+
+---
+
+## Worker requirements (important)
+
+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 workflow.
+
+Also, for collector-based flows:
+
+- Workers will send results back to the master via `POST /distributed/job_complete` (that route must be reachable from workers).
+
+---
+
+## Endpoint: `POST /distributed/job_complete`
+
+Submit one completed worker image back to the master collector queue.
+
+### URL
+
+- `http://<master-host>:<master-port>/distributed/job_complete`
+
+### Request Body
+
+```json
+{
+  "job_id": "exec_1234567890_17",
+  "worker_id": "worker-1",
+  "batch_idx": 0,
+  "image": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAA...",
+  "is_last": false
+}
+```
+
+### Canonical envelope (required fields)
+
+- `job_id` (string, required)
+- `worker_id` (string, required)
+- `batch_idx` (integer >= 0, required)
+- `image` (string, required)
+  - PNG payload as either:
+    - data URL: `data:image/png;base64,...`
+    - raw base64 PNG bytes
+- `is_last` (boolean, required)
+
+Legacy multipart/tensor payload formats are no longer accepted on this endpoint.
+
+### CORS note
+
+If you call the API from a browser (not from a backend), ensure the master ComfyUI is started with `--enable-cors-header`.
+
+---
+
+## Log Endpoints
+
+### `GET /distributed/worker_log/{worker_id}`
+
+Read log files for workers launched locally by the master UI process manager.
+
+- Intended for managed local workers.
+- Query param: `lines` (optional, default `1000`).
+
+### `GET /distributed/local_log`
+
+Read this ComfyUI instance's in-memory runtime log buffer.
+
+- Available on any ComfyUI-Distributed instance (master or worker).
+- Query param: `lines` (optional, default `300`, max `3000`).
+
+### `GET /distributed/remote_worker_log/{worker_id}`
+
+Proxy endpoint on master that fetches logs from a configured remote/cloud worker's
+`/distributed/local_log`.
+
+- Intended for remote/cloud workers in master config.
+- Query param: `lines` (optional, default `300`, max `3000`).
+
+---
+
+## Examples
+
+### 1) Minimal `curl`
+
+```bash
+curl -X POST "http://127.0.0.1:8188/distributed/queue" \
+  -H "Content-Type: application/json" \
+  -d @payload.json
+```
+
+Where `payload.json` contains at least:
+
+```json
+{
+  "prompt": {
+    "1": {"class_type": "KSampler", "inputs": {} }
+  },
+  "enabled_worker_ids": [],
+  "client_id": "external-client"
+}
+```
+
+### 2) Python (`requests`)
+
+```python
+import requests
+
+url = "http://127.0.0.1:8188/distributed/queue"
+payload = {
+    "prompt": {...},
+    "workflow": {...},
+    "client_id": "external-client",
+    "delegate_master": False,
+    "enabled_worker_ids": ["1", "2"],
+}
+
+r = requests.post(url, json=payload, timeout=60)
+r.raise_for_status()
+print(r.json())
+```
+
+### 3) JavaScript (`fetch`)
+
+```js
+const url = "http://127.0.0.1:8188/distributed/queue";
+
+const payload = {
+  prompt: {/* ... */},
+  workflow: {/* ... */},
+  client_id: "external-client",
+  delegate_master: false,
+  enabled_worker_ids: ["1", "2"],
+};
+
+const resp = await fetch(url, {
+  method: "POST",
+  headers: { "Content-Type": "application/json" },
+  body: JSON.stringify(payload),
+});
+
+if (!resp.ok) throw new Error(await resp.text());
+console.log(await resp.json());
+```
+
+---
+
+## Operational notes / gotchas
+
+- If the workflow contains **no distributed nodes**, the endpoint falls back to normal master queueing and returns `worker_count: 0`.
+- Worker selection is “best-effort”: offline workers are skipped.
+- For public URLs/tunnels: prefer configuring `master.host` with an explicit scheme (`https://...`) to avoid ambiguity.
+
+---
+
+## Changelog (this feature)
+
+- Added `POST /distributed/queue` endpoint.
+- Added orchestration module used by the endpoint.

docs/video-upscaler-runpod-preset.md

@@ -15,5 +15,6 @@
 7. In ComfyUI, open the GPU panel on the left.
 > If you set SAGE_ATTENTION to true, add "--use-sage-attention" to Extra Args on the workers.
 8. Launch the workers.
-9. Upload video, add prompt and run workflow.
-10. Right-click the Video Combine node and click Save Preview to save the video.
+9. [Load the workflow.](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/workflows/distributed-upscale-video.json)
+10. Upload video, add prompt and run workflow.
+11. Right-click the Video Combine node and click Save Preview to save the video.

docs/worker-setup-guides.md

@@ -6,6 +6,14 @@
 
 <img width="600" src="https://github.com/user-attachments/assets/609c42aa-8a1c-4a3f-939e-f3552fa1d54f" />
 
+### Master participation modes
+
+The master can either contribute GPU work or stay in **orchestrator-only** mode:
+
+- **Participating**: Master renders alongside workers, useful when you want every available GPU.
+- **Orchestrator-only**: Master sends jobs to selected workers but skips local rendering. Enable this by opening the Distributed panel and unchecking the master toggle. The master card will display *“Master disabled: running as orchestrator only.”*
+- **Fallback**: If orchestrator-only is enabled but no workers remain selected, the master automatically re-enables execution to guarantee the workflow still runs. The UI shows a green *“Master fallback execution active”* badge so you know work is executing locally again.
+
 ### Types of Workers
 
 - **Local workers**: Additional GPUs on the same machine as the master

nodes/__init__.py

@@ -0,0 +1,37 @@
+from .utilities import (
+    DistributedSeed,
+    DistributedModelName,
+    DistributedValue,
+    ImageBatchDivider,
+    AudioBatchDivider,
+    DistributedEmptyImage,
+    AnyType,
+    ByPassTypeTuple,
+    any_type,
+)
+from .collector import DistributedCollectorNode
+from .branch import DistributedBranch
+from .branch_collector import DistributedBranchCollector
+
+NODE_CLASS_MAPPINGS = {
+    "DistributedCollector": DistributedCollectorNode,
+    "DistributedBranch": DistributedBranch,
+    "DistributedBranchCollector": DistributedBranchCollector,
+    "DistributedSeed": DistributedSeed,
+    "DistributedModelName": DistributedModelName,
+    "DistributedValue": DistributedValue,
+    "ImageBatchDivider": ImageBatchDivider,
+    "AudioBatchDivider": AudioBatchDivider,
+    "DistributedEmptyImage": DistributedEmptyImage,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DistributedCollector": "Distributed Collector",
+    "DistributedBranch": "Distributed Branch",
+    "DistributedBranchCollector": "Distributed Branch Collector",
+    "DistributedSeed": "Distributed Seed",
+    "DistributedModelName": "Distributed Model Name",
+    "DistributedValue": "Distributed Value",
+    "ImageBatchDivider": "Image Batch Divider",
+    "AudioBatchDivider": "Audio Batch Divider",
+    "DistributedEmptyImage": "Distributed Empty Image",
+}

nodes/branch.py

@@ -0,0 +1,33 @@
+from typing import Any
+
+from .utilities import any_type
+
+
+class DistributedBranch:
+    @classmethod
+    def INPUT_TYPES(cls: type["DistributedBranch"]) -> dict[str, Any]:
+        return {
+            "required": {
+                "input": (any_type,),
+                "num_branches": ("INT", {"default": 2, "min": 2, "max": 10, "step": 1}),
+            },
+            "hidden": {
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "worker_id": ("STRING", {"default": ""}),
+                "assigned_branch": ("INT", {"default": -1, "min": -1, "max": 9}),
+                "multi_job_id": ("STRING", {"default": ""}),
+            },
+        }
+
+    RETURN_TYPES = tuple([any_type] * 10)
+    RETURN_NAMES = tuple([f"branch_{idx + 1}" for idx in range(10)])
+    FUNCTION = "branch"
+    CATEGORY = "utils"
+
+    def branch(
+        self,
+        input: Any,
+        num_branches: int = 2,
+        **_kwargs: Any,
+    ) -> tuple[Any, ...]:
+        return tuple([input] * 10)

nodes/branch_collector.py

@@ -0,0 +1,298 @@
+import asyncio
+from dataclasses import dataclass, field
+from typing import Any
+
+import aiohttp
+import torch
+
+from ..utils.auth import distributed_auth_headers
+from ..utils.async_helpers import run_async_in_server_loop
+from ..utils.config import is_master_delegate_only, load_config
+from ..utils.image import encode_tensor_png_data_url, ensure_contiguous
+from ..utils.logging import log
+from ..utils.network import get_client_session
+from ..utils.worker_ids import coerce_enabled_worker_ids
+from .context_kwargs import parse_distributed_hidden_context
+from .hidden_inputs import build_distributed_hidden_inputs
+from .queue_wait import collect_worker_queue_results
+from .runtime_helpers import (
+    get_prompt_server_instance as _get_prompt_server_instance,
+    throw_if_processing_interrupted as _throw_if_processing_interrupted,
+)
+from .utilities import any_type
+
+
+MAX_BRANCH_OUTPUTS = 10
+
+
+@dataclass(frozen=True)
+class BranchRunContext:
+    multi_job_id: str = ""
+    is_worker: bool = False
+    master_url: str = ""
+    enabled_worker_ids: list[str] = field(default_factory=list)
+    worker_id: str = ""
+    assigned_branch: int = -1
+    delegate_only: bool = False
+
+class DistributedBranchCollector:
+    @classmethod
+    def INPUT_TYPES(cls: type["DistributedBranchCollector"]) -> dict[str, Any]:
+        optional_inputs = {
+            f"branch_{idx + 1}": (any_type,)
+            for idx in range(MAX_BRANCH_OUTPUTS)
+        }
+        return {
+            "required": {
+                "num_branches": ("INT", {"default": 2, "min": 2, "max": MAX_BRANCH_OUTPUTS, "step": 1}),
+            },
+            "optional": optional_inputs,
+            "hidden": build_distributed_hidden_inputs(
+                include_assigned_branch=True,
+                assigned_branch_max=MAX_BRANCH_OUTPUTS - 1,
+            ),
+        }
+
+    RETURN_TYPES = tuple([any_type] * MAX_BRANCH_OUTPUTS)
+    RETURN_NAMES = tuple([f"branch_{idx + 1}" for idx in range(MAX_BRANCH_OUTPUTS)])
+    FUNCTION = "run"
+    CATEGORY = "utils"
+
+    def _branch_values(self, **kwargs):
+        values = []
+        for idx in range(MAX_BRANCH_OUTPUTS):
+            values.append(kwargs.get(f"branch_{idx + 1}", None))
+        return values
+
+    def _build_run_context(self, **kwargs: Any) -> BranchRunContext:
+        try:
+            assigned_branch = int(kwargs.get("assigned_branch", -1))
+        except (TypeError, ValueError):
+            assigned_branch = -1
+
+        common_context = parse_distributed_hidden_context(kwargs)
+        return BranchRunContext(
+            assigned_branch=assigned_branch,
+            **common_context,
+        )
+
+    def _resolve_local_branch_value(self, branch_values, assigned_branch):
+        try:
+            assigned_idx = int(assigned_branch)
+        except (TypeError, ValueError):
+            assigned_idx = -1
+
+        if 0 <= assigned_idx < MAX_BRANCH_OUTPUTS:
+            assigned_value = branch_values[assigned_idx]
+            if assigned_value is not None:
+                return assigned_value, assigned_idx
+
+            # In participant-pruned prompts, each participant should typically have
+            # exactly one local branch input. If socket names drifted (e.g. branch_2
+            # connected for a participant assigned to slot 0), remap that sole value
+            # to the assigned slot so the participant still contributes correctly.
+            non_none_indices = [idx for idx, value in enumerate(branch_values) if value is not None]
+            if len(non_none_indices) == 1:
+                return branch_values[non_none_indices[0]], assigned_idx
+
+        for idx, value in enumerate(branch_values):
+            if value is not None:
+                return value, idx
+        return None, assigned_idx
+
+    def _build_outputs(self, branch_values, num_branches):
+        outputs = [None] * MAX_BRANCH_OUTPUTS
+        try:
+            branch_count = int(num_branches)
+        except (TypeError, ValueError):
+            branch_count = 2
+        branch_count = max(2, min(branch_count, MAX_BRANCH_OUTPUTS))
+
+        for idx in range(branch_count):
+            value = branch_values[idx]
+            if value is not None:
+                outputs[idx] = value
+        return tuple(outputs)
+
+    def run(
+        self,
+        num_branches: int = 2,
+        **kwargs: Any,
+    ) -> tuple[Any, ...]:
+        branch_values = self._branch_values(**kwargs)
+        context = self._build_run_context(**kwargs)
+
+        if not context.multi_job_id:
+            return self._build_outputs(branch_values, num_branches)
+
+        return run_async_in_server_loop(
+            self.execute(
+                branch_values,
+                num_branches=num_branches,
+                context=context,
+            )
+        )
+
+    def _participant_branch_slots(self, enabled_workers, delegate_mode, num_branches):
+        participants = list(enabled_workers) if delegate_mode else ["master"] + list(enabled_workers)
+        participant_count = len(participants)
+        if participant_count <= 0:
+            return {}
+
+        slots_by_participant = {}
+        for branch_slot in range(int(num_branches)):
+            participant_id = str(participants[branch_slot % participant_count])
+            slots_by_participant.setdefault(participant_id, []).append(branch_slot)
+        return slots_by_participant
+
+    def _expected_worker_ids_for_branches(self, enabled_workers, delegate_mode, num_branches):
+        slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, num_branches)
+        expected = []
+        for participant_id, slots in slots_by_participant.items():
+            if participant_id == "master" or not slots:
+                continue
+            expected.append(str(participant_id))
+        return expected
+
+    def _fallback_for_missing_branch(self, source_value):
+        if isinstance(source_value, torch.Tensor):
+            tensor = source_value
+            if tensor.ndim == 3:
+                tensor = tensor.unsqueeze(0)
+            if tensor.ndim == 4:
+                fallback = torch.zeros_like(tensor)
+                if fallback.is_cuda:
+                    fallback = fallback.cpu()
+                return ensure_contiguous(fallback)
+        return torch.zeros((1, 64, 64, 3), dtype=torch.float32)
+
+    async def _send_branch_result_to_master(self, value, branch_idx, multi_job_id, master_url, worker_id):
+        if not isinstance(value, torch.Tensor):
+            raise ValueError("DistributedBranchCollector currently supports torch.Tensor results for worker transfer.")
+
+        image_batch = value
+        if image_batch.ndim == 3:
+            image_batch = image_batch.unsqueeze(0)
+        if image_batch.ndim != 4 or image_batch.shape[0] <= 0:
+            raise ValueError(
+                "DistributedBranchCollector tensor result must be IMAGE-like with shape [B,H,W,C] or [H,W,C]."
+            )
+
+        payload = {
+            "job_id": str(multi_job_id),
+            "worker_id": str(worker_id),
+            "batch_idx": int(branch_idx),
+            "image": encode_tensor_png_data_url(image_batch, 0),
+            "is_last": True,
+        }
+
+        session = await get_client_session()
+        url = f"{master_url}/distributed/job_complete"
+        async with session.post(
+            url,
+            json=payload,
+            headers=distributed_auth_headers(load_config()),
+            timeout=aiohttp.ClientTimeout(total=60),
+        ) as response:
+            response.raise_for_status()
+
+    async def execute(self, *args: Any, **kwargs: Any) -> tuple[Any, ...]:
+        """Compatibility wrapper with a uniform execute() signature across collectors."""
+        return await self._execute_branch(*args, **kwargs)
+
+    async def _execute_branch(
+        self,
+        branch_values: list[Any],
+        num_branches: int = 2,
+        context: BranchRunContext | None = None,
+    ) -> tuple[Any, ...]:
+        run_context = context or BranchRunContext()
+        try:
+            branch_count = int(num_branches)
+        except (TypeError, ValueError):
+            branch_count = 2
+        branch_count = max(2, min(branch_count, MAX_BRANCH_OUTPUTS))
+
+        outputs = [None] * MAX_BRANCH_OUTPUTS
+        for idx in range(branch_count):
+            if branch_values[idx] is not None:
+                outputs[idx] = branch_values[idx]
+
+        local_value, local_branch_idx = self._resolve_local_branch_value(branch_values, run_context.assigned_branch)
+        if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:
+            outputs[local_branch_idx] = local_value
+
+        if run_context.is_worker:
+            if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:
+                try:
+                    await self._send_branch_result_to_master(
+                        local_value,
+                        local_branch_idx,
+                        run_context.multi_job_id,
+                        run_context.master_url,
+                        run_context.worker_id,
+                    )
+                    outputs[local_branch_idx] = local_value
+                except Exception as exc:
+                    log(f"Worker - DistributedBranchCollector failed to send branch result: {exc}")
+            return tuple(outputs)
+
+        delegate_mode = bool(run_context.delegate_only or is_master_delegate_only())
+        enabled_workers = coerce_enabled_worker_ids(run_context.enabled_worker_ids)
+        slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, branch_count)
+        expected_worker_ids = self._expected_worker_ids_for_branches(enabled_workers, delegate_mode, branch_count)
+        expected_workers = set(expected_worker_ids)
+
+        if not expected_workers:
+            return tuple(outputs)
+
+        prompt_server = _get_prompt_server_instance()
+        async with prompt_server.distributed_jobs_lock:
+            if run_context.multi_job_id not in prompt_server.distributed_pending_jobs:
+                prompt_server.distributed_pending_jobs[run_context.multi_job_id] = asyncio.Queue()
+
+        workers_done: set[str] = set()
+
+        def _handle_queue_result(result: dict[str, Any]) -> None:
+            image_index = result.get("image_index")
+            tensor = result.get("tensor")
+            try:
+                slot_idx = int(image_index)
+            except (TypeError, ValueError):
+                slot_idx = -1
+
+            if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and tensor is not None:
+                if isinstance(tensor, torch.Tensor):
+                    if tensor.is_cuda:
+                        tensor = tensor.cpu()
+                    tensor = ensure_contiguous(tensor)
+                outputs[slot_idx] = tensor
+
+        try:
+            workers_done = await collect_worker_queue_results(
+                prompt_server=prompt_server,
+                multi_job_id=run_context.multi_job_id,
+                expected_workers=expected_workers,
+                on_result=_handle_queue_result,
+                timeout_log_prefix=(
+                    "Master - DistributedBranchCollector heartbeat timeout. "
+                    "Still waiting for workers: "
+                ),
+                throw_if_interrupted=_throw_if_processing_interrupted,
+            )
+
+        finally:
+            async with prompt_server.distributed_jobs_lock:
+                prompt_server.distributed_pending_jobs.pop(run_context.multi_job_id, None)
+                if hasattr(prompt_server, "distributed_job_allowed_workers"):
+                    prompt_server.distributed_job_allowed_workers.pop(run_context.multi_job_id, None)
+
+        missing_workers = sorted(expected_workers - workers_done)
+        if missing_workers:
+            fallback = self._fallback_for_missing_branch(local_value)
+            for missing_worker_id in missing_workers:
+                for slot_idx in slots_by_participant.get(str(missing_worker_id), []):
+                    if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and outputs[slot_idx] is None:
+                        outputs[slot_idx] = fallback
+
+        return tuple(outputs)

nodes/collector.py

@@ -0,0 +1,533 @@
+import asyncio
+import base64
+import io
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+import aiohttp
+import server as _server
+import comfy.model_management
+import torch
+from comfy.utils import ProgressBar
+
+from ..utils.auth import distributed_auth_headers
+from ..utils.logging import debug_log, log
+from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only
+from ..utils.constants import HEARTBEAT_INTERVAL
+from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
+from ..utils.network import build_worker_url, get_client_session, probe_worker
+from ..utils.worker_ids import coerce_enabled_worker_ids
+from ..utils.audio_payload import encode_audio_payload
+from ..utils.async_helpers import run_async_in_server_loop
+from .context_kwargs import parse_distributed_hidden_context
+from .hidden_inputs import build_distributed_hidden_inputs
+
+prompt_server = _server.PromptServer.instance
+
+
+@dataclass(frozen=True)
+class CollectorRunContext:
+    multi_job_id: str = ""
+    is_worker: bool = False
+    master_url: str = ""
+    enabled_worker_ids: list[str] = field(default_factory=list)
+    worker_batch_size: int = 1
+    worker_id: str = ""
+    pass_through: bool = False
+    delegate_only: bool = False
+
+
+class DistributedCollectorNode:
+    EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
+
+    @classmethod
+    def INPUT_TYPES(cls: type["DistributedCollectorNode"]) -> dict[str, Any]:
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "load_balance": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Run this workflow on one least-busy participant (master included when participating).",
+                    },
+                ),
+            },
+            "optional": { "audio": ("AUDIO",) },
+            "hidden": build_distributed_hidden_inputs(
+                include_worker_batch_size=True,
+                include_pass_through=True,
+            ),
+        }
+
+    RETURN_TYPES = ("IMAGE", "AUDIO")
+    RETURN_NAMES = ("images", "audio")
+    FUNCTION = "run"
+    CATEGORY = "image"
+
+    def _build_run_context(self, **kwargs: Any) -> CollectorRunContext:
+        try:
+            worker_batch_size = int(kwargs.get("worker_batch_size", 1))
+        except (TypeError, ValueError):
+            worker_batch_size = 1
+
+        common_context = parse_distributed_hidden_context(kwargs)
+        return CollectorRunContext(
+            worker_batch_size=max(worker_batch_size, 1),
+            pass_through=bool(kwargs.get("pass_through", False)),
+            **common_context,
+        )
+
+    def run(
+        self,
+        images: torch.Tensor,
+        load_balance: bool = False,
+        audio: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        context = self._build_run_context(**kwargs)
+        # Create empty audio if not provided
+        empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
+
+        if not context.multi_job_id or context.pass_through:
+            if context.pass_through:
+                debug_log("Collector: pass-through mode enabled, returning images unchanged")
+            return (images, audio if audio is not None else empty_audio)
+
+        # Use async helper to run in server loop
+        result = run_async_in_server_loop(
+            self.execute(
+                images=images,
+                audio=audio,
+                load_balance=load_balance,
+                context=context,
+            )
+        )
+        return result
+
+    async def send_batch_to_master(
+        self,
+        image_batch: torch.Tensor,
+        audio: dict[str, Any] | None,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+    ) -> None:
+        """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
+
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers=distributed_auth_headers(load_config()),
+                    timeout=aiohttp.ClientTimeout(total=60),
+                ) as response:
+                    response.raise_for_status()
+            except Exception as e:
+                log(f"Worker - Failed to send canonical image envelope to master: {e}")
+                debug_log(f"Worker - Full error details: URL={url}")
+                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.
+
+        Ordering: master first, then workers in `worker_order` (if provided),
+        then any unexpected worker ids in sorted order.
+        """
+        audio_pieces = []
+        sample_rate = 44100
+
+        # Add master audio first if present
+        if master_audio is not None:
+            waveform = master_audio.get("waveform")
+            if waveform is not None and waveform.numel() > 0:
+                audio_pieces.append(waveform)
+                sample_rate = master_audio.get("sample_rate", 44100)
+
+        # Add worker audio in configured enabled-worker order first.
+        ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
+        seen = set()
+        for worker_id_str in ordered_worker_ids:
+            seen.add(worker_id_str)
+            w_audio = worker_audio.get(worker_id_str)
+            if w_audio is not None:
+                waveform = w_audio.get("waveform")
+                if waveform is not None and waveform.numel() > 0:
+                    audio_pieces.append(waveform)
+                    # Use first available sample rate
+                    if sample_rate == 44100:
+                        sample_rate = w_audio.get("sample_rate", 44100)
+
+        # Append any audio from unexpected worker ids deterministically.
+        for worker_id_str in sorted(worker_audio.keys()):
+            if worker_id_str in seen:
+                continue
+            w_audio = worker_audio[worker_id_str]
+            if w_audio is not None:
+                waveform = w_audio.get("waveform")
+                if waveform is not None and waveform.numel() > 0:
+                    audio_pieces.append(waveform)
+                    if sample_rate == 44100:
+                        sample_rate = w_audio.get("sample_rate", 44100)
+
+        if not audio_pieces:
+            return empty_audio
+
+        try:
+            # Concatenate along the samples dimension (dim=-1)
+            # Ensure all pieces have same batch and channel dimensions
+            combined_waveform = torch.cat(audio_pieces, dim=-1)
+            debug_log(f"Master - Combined audio: {len(audio_pieces)} pieces, final shape={combined_waveform.shape}")
+            return {"waveform": combined_waveform, "sample_rate": sample_rate}
+        except Exception as e:
+            log(f"[Distributed] Master - Audio combination failed, returning silence: {e}")
+            return empty_audio
+
+    def _store_worker_result(self, worker_images: dict, item: dict) -> int:
+        """Store one canonical queue item in worker_images in-place.
+
+        Canonical format:
+        - item has 'worker_id', 'image_index', and 'tensor'
+        Returns 1 when stored, otherwise 0.
+        """
+        worker_id = item['worker_id']
+        tensor = item.get('tensor')
+        image_index = item.get('image_index')
+        if tensor is None or image_index is None:
+            return 0
+
+        worker_images.setdefault(worker_id, {})
+        worker_images[worker_id][image_index] = tensor
+        return 1
+
+    def _reorder_and_combine_tensors(
+        self,
+        worker_images: dict,
+        worker_order: list,
+        master_batch_size: int,
+        images_on_cpu,
+        delegate_mode: bool,
+        fallback_images,
+    ) -> 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):
+                ordered_tensors.append(images_on_cpu[i:i+1])
+
+        ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
+        seen = set()
+        for worker_id_str in ordered_worker_ids:
+            seen.add(worker_id_str)
+            if worker_id_str not in worker_images:
+                continue
+            for idx in sorted(worker_images[worker_id_str].keys()):
+                ordered_tensors.append(worker_images[worker_id_str][idx])
+
+        # Append any unexpected worker ids deterministically.
+        for worker_id_str in sorted(worker_images.keys()):
+            if worker_id_str in seen:
+                continue
+            for idx in sorted(worker_images[worker_id_str].keys()):
+                ordered_tensors.append(worker_images[worker_id_str][idx])
+
+        cpu_tensors = []
+        for t in ordered_tensors:
+            if t.is_cuda:
+                t = t.cpu()
+            t = ensure_contiguous(t)
+            cpu_tensors.append(t)
+
+        if cpu_tensors:
+            return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
+        elif fallback_images is not None:
+            return ensure_contiguous(fallback_images)
+        else:
+            raise ValueError("No image data collected from master or workers")
+
+    async def _ensure_pending_queue(self, multi_job_id):
+        async with prompt_server.distributed_jobs_lock:
+            if multi_job_id not in prompt_server.distributed_pending_jobs:
+                prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
+                debug_log(f"Master - Initialized queue early for job {multi_job_id}")
+            else:
+                existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize()
+                debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})")
+
+    async def _cleanup_pending_queue(self, multi_job_id):
+        async with prompt_server.distributed_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_jobs:
+                del prompt_server.distributed_pending_jobs[multi_job_id]
+            if hasattr(prompt_server, "distributed_job_allowed_workers"):
+                prompt_server.distributed_job_allowed_workers.pop(multi_job_id, None)
+
+    def _build_master_inputs(self, images, audio, delegate_mode, multi_job_id, num_workers):
+        if delegate_mode:
+            debug_log(
+                f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers"
+            )
+            return 0, None, None
+
+        images_on_cpu = ensure_contiguous(images.cpu())
+        master_batch_size = images.shape[0]
+        debug_log(
+            f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers..."
+        )
+        return master_batch_size, images_on_cpu, audio
+
+    async def _probe_missing_workers_busy(self, missing_workers):
+        any_busy = False
+        try:
+            cfg = load_config()
+            cfg_workers = cfg.get('workers', [])
+            for wid in list(missing_workers):
+                wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
+                if not wrec:
+                    debug_log(f"Collector probe: worker {wid} not found in config")
+                    continue
+                worker_url = build_worker_url(wrec)
+                try:
+                    payload = await probe_worker(worker_url, timeout=2.0)
+                    queue_remaining = None
+                    if payload is not None:
+                        queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
+                    debug_log(
+                        "Collector probe: worker "
+                        f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
+                    )
+                    if payload is not None and queue_remaining and queue_remaining > 0:
+                        any_busy = True
+                        log(
+                            f"Master - Probe grace: worker {wid} appears busy "
+                            f"(queue_remaining={queue_remaining}). Continuing to wait."
+                        )
+                        break
+                except Exception as e:
+                    debug_log(f"Collector probe failed for worker {wid}: {e}")
+        except Exception as e:
+            debug_log(f"Collector probe setup error: {e}")
+        return any_busy
+
+    async def _drain_remaining_queue_items(
+        self,
+        multi_job_id: str,
+        worker_images: dict[str, dict[int, torch.Tensor]],
+        mark_worker_done: Callable[[str], None],
+    ) -> int:
+        collected_count = 0
+        async with prompt_server.distributed_jobs_lock:
+            if multi_job_id not in prompt_server.distributed_pending_jobs:
+                log(f"Master - Queue {multi_job_id} no longer exists!")
+                return collected_count
+
+            final_q = prompt_server.distributed_pending_jobs[multi_job_id]
+            remaining_items = []
+            while not final_q.empty():
+                try:
+                    remaining_items.append(final_q.get_nowait())
+                except asyncio.QueueEmpty:
+                    break
+
+        for item in remaining_items:
+            worker_id = item['worker_id']
+            is_last = item.get('is_last', False)
+            collected_count += self._store_worker_result(worker_images, item)
+            if is_last:
+                mark_worker_done(worker_id)
+        return collected_count
+
+    async def _collect_worker_results(
+        self,
+        multi_job_id: str,
+        enabled_workers: list[str],
+        expected_workers: set[str],
+        worker_images: dict[str, dict[int, torch.Tensor]],
+        worker_audio: dict[str, dict[str, Any]],
+    ) -> int:
+        num_workers = len(expected_workers)
+        workers_done = set()
+        collected_count = 0
+        base_timeout = float(get_worker_timeout_seconds())
+        slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
+        last_activity = time.time()
+        progress = ProgressBar(num_workers)
+
+        def mark_worker_done(done_worker_id: str) -> None:
+            done_worker_id = str(done_worker_id)
+            if done_worker_id not in expected_workers:
+                debug_log(
+                    f"Master - Ignoring completion from unexpected worker {done_worker_id} for job {multi_job_id}"
+                )
+                return
+            if done_worker_id in workers_done:
+                debug_log(
+                    f"Master - Ignoring duplicate completion from worker {done_worker_id} for job {multi_job_id}"
+                )
+                return
+            workers_done.add(done_worker_id)
+            progress.update(1)
+
+        while len(workers_done) < num_workers:
+            comfy.model_management.throw_exception_if_processing_interrupted()
+            try:
+                async with prompt_server.distributed_jobs_lock:
+                    q = prompt_server.distributed_pending_jobs[multi_job_id]
+
+                result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
+                worker_id = result['worker_id']
+                is_last = result.get('is_last', False)
+                collected_count += self._store_worker_result(worker_images, result)
+                debug_log(
+                    f"Master - Got canonical result from worker {worker_id}, "
+                    f"image {result.get('image_index', 0)}, is_last={is_last}"
+                )
+
+                result_audio = result.get('audio')
+                if result_audio is not None:
+                    worker_audio[worker_id] = result_audio
+                    debug_log(f"Master - Got audio from worker {worker_id}")
+
+                last_activity = time.time()
+                base_timeout = float(get_worker_timeout_seconds())
+                if is_last:
+                    mark_worker_done(worker_id)
+            except asyncio.TimeoutError:
+                if (time.time() - last_activity) < base_timeout:
+                    comfy.model_management.throw_exception_if_processing_interrupted()
+                    continue
+
+                comfy.model_management.throw_exception_if_processing_interrupted()
+                missing_workers = set(str(w) for w in enabled_workers) - workers_done
+                elapsed = time.time() - last_activity
+                for missing_worker_id in sorted(missing_workers):
+                    log(
+                        "Master - Heartbeat timeout: "
+                        f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
+                    )
+                log(
+                    f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} "
+                    f"(elapsed={elapsed:.1f}s)"
+                )
+
+                if await self._probe_missing_workers_busy(missing_workers):
+                    last_activity = time.time()
+                    base_timeout = float(get_worker_timeout_seconds())
+                    continue
+
+                collected_count += await self._drain_remaining_queue_items(
+                    multi_job_id,
+                    worker_images,
+                    mark_worker_done,
+                )
+                break
+
+        return collected_count
+
+    async def _execute_master(
+        self,
+        images,
+        audio,
+        multi_job_id,
+        enabled_worker_ids,
+        delegate_only,
+    ):
+        delegate_mode = delegate_only or is_master_delegate_only()
+        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+        expected_workers = set(enabled_workers)
+        num_workers = len(expected_workers)
+        if num_workers == 0:
+            return (images, audio if audio is not None else self.EMPTY_AUDIO)
+
+        await self._ensure_pending_queue(multi_job_id)
+        master_batch_size, images_on_cpu, master_audio = self._build_master_inputs(
+            images, audio, delegate_mode, multi_job_id, num_workers
+        )
+
+        worker_images = {}
+        worker_audio = {}
+        try:
+            await self._collect_worker_results(
+                multi_job_id=multi_job_id,
+                enabled_workers=enabled_workers,
+                expected_workers=expected_workers,
+                worker_images=worker_images,
+                worker_audio=worker_audio,
+            )
+        except comfy.model_management.InterruptProcessingException:
+            await self._cleanup_pending_queue(multi_job_id)
+            raise
+
+        await self._cleanup_pending_queue(multi_job_id)
+        try:
+            combined = self._reorder_and_combine_tensors(
+                worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
+            )
+            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})"
+            )
+            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}")
+            return (images, audio if audio is not None else self.EMPTY_AUDIO)
+
+    async def execute(self, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, dict[str, Any]]:
+        """Compatibility wrapper with a uniform execute() signature across collectors."""
+        return await self._execute_collector(*args, **kwargs)
+
+    async def _execute_collector(
+        self,
+        images: torch.Tensor,
+        audio: dict[str, Any] | None,
+        load_balance: bool = False,
+        context: CollectorRunContext | None = None,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        run_context = context or CollectorRunContext()
+        _ = load_balance
+        _ = run_context.worker_batch_size
+        if run_context.is_worker:
+            debug_log(
+                "Worker - Job "
+                f"{run_context.multi_job_id} complete. Sending {images.shape[0]} image(s) to master"
+            )
+            await self.send_batch_to_master(
+                images,
+                audio,
+                run_context.multi_job_id,
+                run_context.master_url,
+                run_context.worker_id,
+            )
+            return (images, audio if audio is not None else self.EMPTY_AUDIO)
+
+        return await self._execute_master(
+            images=images,
+            audio=audio,
+            multi_job_id=run_context.multi_job_id,
+            enabled_worker_ids=run_context.enabled_worker_ids,
+            delegate_only=run_context.delegate_only,
+        )

nodes/context_kwargs.py

@@ -0,0 +1,26 @@
+from typing import Any
+
+from ..utils.parsing import coerce_bool
+from ..utils.worker_ids import coerce_enabled_worker_ids
+
+
+def _parse_bool(value: Any, default: bool = False) -> bool:
+    return coerce_bool(value, default=default)
+
+
+def _parse_enabled_workers(value: Any) -> list[str]:
+    if value is None:
+        return []
+    return coerce_enabled_worker_ids(value)
+
+
+def parse_distributed_hidden_context(kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Parse common distributed hidden inputs from node kwargs."""
+    return {
+        "multi_job_id": str(kwargs.get("multi_job_id", "") or ""),
+        "is_worker": _parse_bool(kwargs.get("is_worker", False)),
+        "master_url": str(kwargs.get("master_url", "") or ""),
+        "enabled_worker_ids": _parse_enabled_workers(kwargs.get("enabled_worker_ids", "[]")),
+        "worker_id": str(kwargs.get("worker_id", "") or ""),
+        "delegate_only": _parse_bool(kwargs.get("delegate_only", False)),
+    }

nodes/distributed_upscale.py

@@ -0,0 +1,517 @@
+import math
+from functools import wraps
+from typing import Any, Callable
+
+import comfy.samplers
+
+from ..utils.logging import debug_log, log
+from ..utils.async_helpers import run_async_in_server_loop
+from ..upscale.job_store import ensure_tile_jobs_initialized
+from .hidden_inputs import build_distributed_hidden_inputs
+
+from ..upscale.tile_ops import TileOpsMixin
+from ..upscale.result_collector import ResultCollectorMixin
+from ..upscale.worker_comms import WorkerCommsMixin
+from ..upscale.job_state import JobStateMixin
+from ..upscale.processing_args import UpscaleCoreArgs
+from ..upscale.modes.single_gpu import SingleGpuModeMixin
+from ..upscale.modes.static import StaticModeMixin
+from ..upscale.modes.dynamic import DynamicModeMixin
+from ..utils.worker_ids import coerce_enabled_worker_ids
+
+def sync_wrapper(async_func: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorator to wrap async methods for synchronous execution."""
+    @wraps(async_func)
+    def sync_func(self, *args: Any, **kwargs: Any) -> Any:
+        # Use run_async_in_server_loop for ComfyUI compatibility
+        return run_async_in_server_loop(
+            async_func(self, *args, **kwargs),
+            timeout=600.0  # 10 minute timeout for long operations
+        )
+    return sync_func
+
+
+def _parse_enabled_worker_ids(enabled_worker_ids: str | list[str] | None) -> list[str]:
+    """Backward-compatible alias for enabled-worker normalization."""
+    return coerce_enabled_worker_ids(enabled_worker_ids)
+
+class UltimateSDUpscaleDistributed(
+    DynamicModeMixin,
+    StaticModeMixin,
+    SingleGpuModeMixin,
+    ResultCollectorMixin,
+    WorkerCommsMixin,
+    JobStateMixin,
+    TileOpsMixin,
+):
+
+    """
+    Distributed version of Ultimate SD Upscale (No Upscale).
+    
+    Supports three processing modes:
+    1. Single GPU: No workers available, process everything locally
+    2. Static Mode: Small batches, distributes tiles across workers (flattened)
+    3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
+    
+    Features:
+    - Multi-mode batch handling for efficient video/image upscaling
+    - Tiled VAE support for memory efficiency
+    - 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)
+    
+    Threshold: dynamic_threshold input controls mode switch (default 8)
+    """
+
+    def __init__(self):
+        """Initialize the node and ensure persistent storage exists."""
+        # Pre-initialize the persistent storage on node creation
+        ensure_tile_jobs_initialized()
+        debug_log("UltimateSDUpscaleDistributed - Node initialized")
+
+    @classmethod
+    def INPUT_TYPES(cls: type["UltimateSDUpscaleDistributed"]) -> dict[str, Any]:
+        hidden_inputs = build_distributed_hidden_inputs()
+        hidden_inputs.update(
+            {
+                "tile_indices": ("STRING", {"default": ""}),  # Unused - kept for compatibility
+                "dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
+            }
+        )
+        return {
+            "required": {
+                "upscaled_image": ("IMAGE",),
+                "model": ("MODEL",),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "vae": ("VAE",),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
+                "denoise": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "tile_width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
+                "tile_height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
+                "padding": ("INT", {"default": 32, "min": 0, "max": 256, "step": 8}),
+                "mask_blur": ("INT", {"default": 8, "min": 0, "max": 256}),
+                "force_uniform_tiles": ("BOOLEAN", {"default": True}),
+                "tiled_decode": ("BOOLEAN", {"default": False}),
+            },
+            "hidden": hidden_inputs,
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "run"
+    CATEGORY = "image/upscaling"
+
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        """Force re-execution."""
+        return float("nan")  # Always re-execute
+
+    def run(
+        self,
+        upscaled_image: Any,
+        model: Any,
+        positive: Any,
+        negative: Any,
+        vae: Any,
+        seed: int,
+        steps: int,
+        cfg: float,
+        sampler_name: str,
+        scheduler: str,
+        denoise: float,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        tiled_decode: bool,
+        multi_job_id: str = "",
+        is_worker: bool = False,
+        master_url: str = "",
+        enabled_worker_ids: str = "[]",
+        worker_id: str = "",
+        tile_indices: str = "",
+        dynamic_threshold: int = 8,
+    ) -> tuple[Any, ...]:
+        """Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler."""
+        core_args = UpscaleCoreArgs(
+            model=model,
+            positive=positive,
+            negative=negative,
+            vae=vae,
+            seed=seed,
+            steps=steps,
+            cfg=cfg,
+            sampler_name=sampler_name,
+            scheduler=scheduler,
+            denoise=denoise,
+            tiled_decode=tiled_decode,
+        )
+        # Strict WAN/FLOW batching: error if batch is not 4n+1 (except allow 1)
+        try:
+            batch_size = int(getattr(upscaled_image, 'shape', [1])[0])
+        except Exception:
+            batch_size = 1
+        # Enforce 4n+1 batches globally for any model when batch > 1 (master only)
+        if not is_worker and batch_size != 1 and (batch_size % 4 != 1):
+            raise ValueError(
+                f"Batch size {batch_size} is not of the form 4n+1. "
+                "This node requires batch sizes of 1 or 4n+1 (1, 5, 9, 13, ...). "
+                "Please adjust the batch size."
+            )
+        if not multi_job_id:
+            # No distributed processing, run single GPU version
+            return self.process_single_gpu(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+            )
+        
+        if is_worker:
+            # Worker mode: process tiles synchronously
+            return self.process_worker(upscaled_image, model, positive, negative, vae,
+                                      seed, steps, cfg, sampler_name, scheduler, denoise,
+                                      tile_width, tile_height, padding, mask_blur,
+                                      force_uniform_tiles, tiled_decode, multi_job_id, master_url,
+                                      worker_id, enabled_worker_ids, dynamic_threshold)
+        else:
+            # Master mode: distribute and collect synchronously
+            return self.process_master(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+                multi_job_id=multi_job_id,
+                enabled_worker_ids=enabled_worker_ids,
+                dynamic_threshold=dynamic_threshold,
+            )
+
+    def process_worker(
+        self,
+        upscaled_image: Any,
+        model: Any,
+        positive: Any,
+        negative: Any,
+        vae: Any,
+        seed: int,
+        steps: int,
+        cfg: float,
+        sampler_name: str,
+        scheduler: str,
+        denoise: float,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        tiled_decode: bool,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+        enabled_worker_ids: str,
+        dynamic_threshold: int,
+    ) -> tuple[Any, ...]:
+        """Unified worker processing - handles both static and dynamic modes."""
+        core_args = UpscaleCoreArgs(
+            model=model,
+            positive=positive,
+            negative=negative,
+            vae=vae,
+            seed=seed,
+            steps=steps,
+            cfg=cfg,
+            sampler_name=sampler_name,
+            scheduler=scheduler,
+            denoise=denoise,
+            tiled_decode=tiled_decode,
+        )
+        # Get batch size to determine mode
+        batch_size = upscaled_image.shape[0]
+        
+        # Ensure mode consistency across master/workers via shared threshold
+        # Determine mode (must match master's logic)
+        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+        num_workers = len(enabled_workers)
+        # Compute number of tiles for this image to decide if tile distribution makes sense
+        _, height, width, _ = upscaled_image.shape
+        all_tiles = self.calculate_tiles(width, height, self.round_to_multiple(tile_width), self.round_to_multiple(tile_height), force_uniform_tiles)
+        num_tiles_per_image = len(all_tiles)
+
+        mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
+        # For USDU-style processing, we want tile distribution whenever workers are available
+        # and there is more than one tile to process for single-image runs.
+        if num_workers > 0 and batch_size <= 1 and num_tiles_per_image > 1:
+            mode = "static"
+            
+        debug_log(f"USDU Dist Worker - Batch size {batch_size}")
+        
+        if mode == "dynamic":
+            return self.process_worker_dynamic(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+                multi_job_id=multi_job_id,
+                master_url=master_url,
+                worker_id=worker_id,
+                enabled_worker_ids=enabled_worker_ids,
+                dynamic_threshold=dynamic_threshold,
+            )
+        
+        # Static mode - enhanced with health monitoring and retry logic
+        return self._process_worker_static_sync(upscaled_image, core_args,
+                                               tile_width, tile_height, padding, mask_blur,
+                                               force_uniform_tiles, multi_job_id, master_url,
+                                               worker_id, enabled_workers)
+
+    def process_master(
+        self,
+        upscaled_image: Any,
+        core_args: UpscaleCoreArgs,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        multi_job_id: str,
+        enabled_worker_ids: str,
+        dynamic_threshold: int,
+    ) -> tuple[Any, ...]:
+        """Unified master processing with enhanced monitoring and failure handling."""
+        # Round tile dimensions
+        tile_width = self.round_to_multiple(tile_width)
+        tile_height = self.round_to_multiple(tile_height)
+        
+        # Get image dimensions and batch size
+        batch_size, height, width, _ = upscaled_image.shape
+        
+        # Calculate all tiles and grid
+        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+        num_tiles_per_image = len(all_tiles)
+        rows = math.ceil(height / tile_height)
+        cols = math.ceil(width / tile_width)
+        log(
+            f"USDU Dist: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({num_tiles_per_image} tiles/image) | Batch {batch_size}"
+        )
+        
+        # Parse enabled workers
+        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+        num_workers = len(enabled_workers)
+        
+        # Determine processing mode
+        mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
+        # Prefer tile-based static distribution when workers are available and there are multiple tiles,
+        # for single-image jobs to spread tiles across GPUs like the legacy tile queue.
+        if num_workers > 0 and batch_size <= 1 and num_tiles_per_image > 1:
+            mode = "static"
+        
+        log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
+
+        if mode == "single_gpu":
+            # No workers, process all tiles locally
+            return self.process_single_gpu(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+            )
+        
+        elif mode == "dynamic":
+            # Dynamic mode for large batches
+            return self.process_master_dynamic(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+                multi_job_id=multi_job_id,
+                enabled_workers=enabled_workers,
+            )
+        
+        # Static mode - enhanced with unified job management
+        return self._process_master_static_sync(upscaled_image, core_args,
+                                               tile_width, tile_height, padding, mask_blur,
+                                               force_uniform_tiles, multi_job_id, enabled_workers,
+                                               all_tiles, num_tiles_per_image)
+
+    def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
+        """Determine mode from worker availability and configured threshold."""
+        if num_workers == 0:
+            return "single_gpu"
+        threshold = max(1, int(dynamic_threshold))
+        if int(batch_size) >= threshold:
+            return "dynamic"
+        return "static"
+
+class USDUDelegateCollector:
+    """Lightweight stand-in used automatically in delegate-only master prompts.
+
+    When the master is in delegate-only mode, the orchestration code swaps the
+    full ``UltimateSDUpscaleDistributed`` class for this one so that no upstream
+    model/image nodes execute on the master.  Workers initialise the dynamic job
+    queue via ``/distributed/init_dynamic_job`` and this node simply waits for
+    all images to arrive, then assembles the output tensor.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {},
+            "hidden": {
+                "multi_job_id": ("STRING", {"default": ""}),
+                "enabled_worker_ids": ("STRING", {"default": "[]"}),
+                "delegate_only": ("BOOLEAN", {"default": True}),
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "worker_id": ("STRING", {"default": ""}),
+                "master_url": ("STRING", {"default": ""}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "run"
+    CATEGORY = "image/upscaling"
+
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        return float("nan")
+
+    def run(self, multi_job_id="", enabled_worker_ids="[]", **kwargs):
+        import time
+        import torch
+        from ..utils.image import pil_to_tensor
+        from ..upscale.job_models import ImageJobState
+
+        if not multi_job_id:
+            log("USDUDelegateCollector: no multi_job_id, returning empty image")
+            return (torch.zeros(1, 64, 64, 3),)
+
+        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+        num_workers = len(enabled_workers)
+        log(f"USDU delegate-only: waiting for workers to init job {multi_job_id}")
+
+        prompt_server = ensure_tile_jobs_initialized()
+
+        # Wait for workers to create the job via /distributed/init_dynamic_job
+        max_wait = 120.0
+        poll_interval = 1.0
+        start = time.time()
+        job_data = None
+        while time.time() - start < max_wait:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, ImageJobState) and job_data.batch_size > 0:
+                break
+            job_data = None
+            time.sleep(poll_interval)
+
+        if job_data is None:
+            log(f"USDUDelegateCollector: job {multi_job_id} was not initialized by workers within {max_wait}s")
+            return (torch.zeros(1, 64, 64, 3),)
+
+        batch_size = job_data.batch_size
+        log(f"USDU delegate-only: job ready, collecting {batch_size} images from {num_workers} workers")
+
+        # Collect all images using the existing result collector
+        collected = run_async_in_server_loop(
+            self._collect_all_images(multi_job_id, batch_size, num_workers, prompt_server),
+            timeout=600.0,
+        )
+
+        # Assemble tensor from collected PIL images
+        result_images = []
+        for idx in range(batch_size):
+            img = collected.get(idx)
+            if img is not None:
+                result_images.append(pil_to_tensor(img))
+            else:
+                log(f"USDUDelegateCollector: missing image {idx}, using blank")
+                result_images.append(torch.zeros(1, 64, 64, 3))
+
+        result_tensor = torch.cat(result_images, dim=0)
+        log(f"USDU delegate-only: collected all {batch_size} images")
+        return (result_tensor,)
+
+    @staticmethod
+    async def _collect_all_images(multi_job_id, batch_size, num_workers, prompt_server):
+        """Wait for all images to arrive from workers."""
+        import asyncio
+        from ..upscale.job_models import ImageJobState
+        from ..upscale.job_timeout import check_and_requeue_timed_out_workers
+
+        timeout_seconds = 300.0
+        poll_interval = 2.0
+        start = asyncio.get_event_loop().time()
+
+        while True:
+            # Drain results from the queue
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, ImageJobState):
+                # Drain queue items into completed_images
+                while True:
+                    try:
+                        result = job_data.queue.get_nowait()
+                        worker_id = result.get("worker_id")
+                        if "image_idx" in result and "image" in result:
+                            idx = result["image_idx"]
+                            if idx not in job_data.completed_images:
+                                job_data.completed_images[idx] = result["image"]
+                                debug_log(f"Delegate collected image {idx} from worker {worker_id}")
+                    except asyncio.QueueEmpty:
+                        break
+
+                completed_count = len(job_data.completed_images)
+                if completed_count >= batch_size:
+                    debug_log(f"Delegate: all {batch_size} images collected")
+                    completed = dict(job_data.completed_images)
+                    # Cleanup
+                    async with prompt_server.distributed_tile_jobs_lock:
+                        prompt_server.distributed_pending_tile_jobs.pop(multi_job_id, None)
+                    return completed
+
+                # Check for timeouts periodically
+                await check_and_requeue_timed_out_workers(multi_job_id, batch_size)
+
+            elapsed = asyncio.get_event_loop().time() - start
+            if elapsed > timeout_seconds:
+                log(f"USDUDelegateCollector: timed out after {elapsed:.0f}s with {len(getattr(job_data, 'completed_images', {}))} of {batch_size} images")
+                if isinstance(job_data, ImageJobState):
+                    completed = dict(job_data.completed_images)
+                    async with prompt_server.distributed_tile_jobs_lock:
+                        prompt_server.distributed_pending_tile_jobs.pop(multi_job_id, None)
+                    return completed
+                return {}
+
+            await asyncio.sleep(poll_interval)
+
+
+# Node registration
+NODE_CLASS_MAPPINGS = {
+    "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
+    "USDUDelegateCollector": USDUDelegateCollector,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
+    # No display name for USDUDelegateCollector — it's an internal node
+}

nodes/hidden_inputs.py

@@ -0,0 +1,38 @@
+from typing import Any
+
+
+def build_worker_identity_hidden_inputs() -> dict[str, tuple[str, dict[str, Any]]]:
+    """Hidden-input contract for worker identity context."""
+    return {
+        "is_worker": ("BOOLEAN", {"default": False}),
+        "enabled_worker_ids": ("STRING", {"default": "[]"}),
+        "worker_id": ("STRING", {"default": ""}),
+    }
+
+
+def build_distributed_hidden_inputs(
+    *,
+    include_assigned_branch: bool = False,
+    assigned_branch_max: int = 9,
+    include_worker_batch_size: bool = False,
+    include_pass_through: bool = False,
+) -> dict[str, tuple[str, dict[str, Any]]]:
+    hidden_inputs: dict[str, tuple[str, dict[str, Any]]] = {
+        "multi_job_id": ("STRING", {"default": ""}),
+        "master_url": ("STRING", {"default": ""}),
+        "delegate_only": ("BOOLEAN", {"default": False}),
+    }
+    hidden_inputs.update(build_worker_identity_hidden_inputs())
+    if include_assigned_branch:
+        hidden_inputs["assigned_branch"] = (
+            "INT",
+            {"default": -1, "min": -1, "max": assigned_branch_max},
+        )
+    if include_worker_batch_size:
+        hidden_inputs["worker_batch_size"] = (
+            "INT",
+            {"default": 1, "min": 1, "max": 1024},
+        )
+    if include_pass_through:
+        hidden_inputs["pass_through"] = ("BOOLEAN", {"default": False})
+    return hidden_inputs

nodes/queue_wait.py

@@ -0,0 +1,46 @@
+import asyncio
+import time
+from collections.abc import Callable
+from typing import Any
+
+from ..utils.config import get_worker_timeout_seconds
+from ..utils.constants import HEARTBEAT_INTERVAL
+from ..utils.logging import log
+
+
+async def collect_worker_queue_results(
+    *,
+    prompt_server: Any,
+    multi_job_id: str,
+    expected_workers: set[str],
+    on_result: Callable[[dict[str, Any]], None],
+    timeout_log_prefix: str,
+    throw_if_interrupted: Callable[[], None],
+) -> set[str]:
+    workers_done: set[str] = set()
+    base_timeout = float(get_worker_timeout_seconds())
+    slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
+    last_activity = time.time()
+
+    while len(workers_done) < len(expected_workers):
+        throw_if_interrupted()
+        try:
+            async with prompt_server.distributed_jobs_lock:
+                queue = prompt_server.distributed_pending_jobs[multi_job_id]
+            result = await asyncio.wait_for(queue.get(), timeout=slice_timeout)
+        except asyncio.TimeoutError:
+            if (time.time() - last_activity) < base_timeout:
+                continue
+            missing_workers = sorted(expected_workers - workers_done)
+            log(f"{timeout_log_prefix}{missing_workers}")
+            break
+
+        on_result(result)
+        worker_id_value = str(result.get("worker_id", ""))
+        is_last = bool(result.get("is_last", False))
+        last_activity = time.time()
+        base_timeout = float(get_worker_timeout_seconds())
+        if is_last and worker_id_value in expected_workers:
+            workers_done.add(worker_id_value)
+
+    return workers_done

nodes/runtime_helpers.py

@@ -0,0 +1,16 @@
+from typing import Any
+
+
+def get_prompt_server_instance() -> Any:
+    import server as _server
+
+    return _server.PromptServer.instance
+
+
+def throw_if_processing_interrupted() -> None:
+    try:
+        import comfy.model_management as model_management
+
+        model_management.throw_exception_if_processing_interrupted()
+    except Exception:
+        return

nodes/utilities.py

@@ -0,0 +1,379 @@
+import torch
+import json
+from typing import Any
+
+from ..utils.logging import debug_log, log
+from ..utils.worker_ids import coerce_enabled_worker_ids, parse_worker_index, worker_value_key
+from .hidden_inputs import build_worker_identity_hidden_inputs
+
+
+def _chunk_bounds(total_items: int, n_splits: int) -> list[tuple[int, int]]:
+    """Return contiguous [start, end) bounds for n_splits chunks."""
+    split_count = max(1, int(n_splits))
+    total = max(0, int(total_items))
+    base, remainder = divmod(total, split_count)
+
+    bounds: list[tuple[int, int]] = []
+    start = 0
+    for idx in range(split_count):
+        size = base + (1 if idx < remainder else 0)
+        end = start + size
+        bounds.append((start, end))
+        start = end
+    return bounds
+
+
+class DistributedSeed:
+    """
+    Distributes seed values across multiple GPUs.
+    On master: passes through the original seed.
+    On workers: adds offset based on worker ID.
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls: type["DistributedSeed"]) -> dict[str, Any]:
+        return {
+            "required": {
+                "seed": ("INT", {
+                    "default": 1125899906842, 
+                    "min": 0,
+                    "max": 1125899906842624,
+                    "forceInput": False  # Widget by default, can be converted to input
+                }),
+            },
+            "hidden": {
+                **build_worker_identity_hidden_inputs(),
+            },
+        }
+    
+    RETURN_TYPES = ("INT",)
+    RETURN_NAMES = ("seed",)
+    FUNCTION = "distribute"
+    CATEGORY = "utils"
+    
+    def distribute(
+        self,
+        seed: int,
+        is_worker: bool = False,
+        worker_id: str = "",
+        enabled_worker_ids: str = "[]",
+    ) -> tuple[int]:
+        if not is_worker:
+            # Master node: pass through original values
+            debug_log(f"Distributor - Master: seed={seed}")
+            return (seed,)
+        else:
+            enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+            worker_index = parse_worker_index(worker_id, enabled_workers)
+            if worker_index is not None:
+                offset = worker_index + 1
+                new_seed = seed + offset
+                debug_log(f"Distributor - Worker {worker_index}: seed={seed} → {new_seed}")
+                return (new_seed,)
+
+            debug_log(f"Distributor - Error parsing worker_id '{worker_id}': no worker index resolved")
+            # Fallback: return original seed
+            return (seed,)
+
+
+# Define ByPassTypeTuple for flexible return types
+class AnyType(str):
+    def __ne__(self, __value: object) -> bool:
+        return False
+
+any_type = AnyType("*")
+
+
+class DistributedValue:
+    """
+    Outputs a different value per worker.
+    On master: returns default_value.
+    On workers: looks up the worker-specific value from a JSON map,
+    falling back to default_value if not set.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls: type["DistributedValue"]) -> dict[str, Any]:
+        return {
+            "required": {
+                "default_value": ("STRING", {"default": ""}),
+                "worker_values": ("STRING", {"default": "{}"}),
+            },
+            "hidden": {
+                **build_worker_identity_hidden_inputs(),
+            },
+        }
+
+    RETURN_TYPES = (any_type,)
+    RETURN_NAMES = ("value",)
+    FUNCTION = "distribute"
+    CATEGORY = "utils"
+
+    @staticmethod
+    def _coerce(value: Any, value_type: str) -> Any:
+        """Convert a string value to the requested type."""
+        if value_type == "INT":
+            return int(float(value))
+        if value_type == "FLOAT":
+            return float(value)
+        return value  # STRING and COMBO stay as strings
+
+    @staticmethod
+    def _coerce_safe(value: Any, value_type: str) -> Any:
+        """Best-effort coercion with graceful fallback to original value."""
+        try:
+            return DistributedValue._coerce(value, value_type)
+        except (TypeError, ValueError):
+            return value
+
+    @staticmethod
+    def _infer_value_type(value: Any) -> str:
+        """Infer coercion type from the provided default value."""
+        if isinstance(value, bool):
+            return "STRING"
+        if isinstance(value, int):
+            return "INT"
+        if isinstance(value, float):
+            return "FLOAT"
+        return "STRING"
+
+    def distribute(
+        self,
+        default_value: Any,
+        worker_values: str | dict[str, Any] = "{}",
+        is_worker: bool = False,
+        worker_id: str = "",
+        enabled_worker_ids: str = "[]",
+    ) -> tuple[Any]:
+        values = {}
+        value_type = "STRING"
+
+        try:
+            raw_values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
+            values = dict(raw_values) if isinstance(raw_values, dict) else {}
+            if not isinstance(values, dict):
+                values = {}
+        except json.JSONDecodeError as e:
+            debug_log(f"DistributedValue - Error parsing worker_values: {e}")
+            values = {}
+
+        inferred_type = self._infer_value_type(default_value)
+        value_type = values.get("_type", inferred_type)
+        if value_type not in {"STRING", "COMBO", "INT", "FLOAT"}:
+            value_type = inferred_type
+        coerced_default = self._coerce_safe(default_value, value_type)
+
+        if not is_worker:
+            debug_log(f"DistributedValue - Master: returning default '{coerced_default}'")
+            return (coerced_default,)
+
+        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+        direct_key = str(worker_id).strip()
+        lookup_key = direct_key if direct_key in values else worker_value_key(worker_id, enabled_workers)
+        raw = values.get(lookup_key)
+        if raw is not None and raw != "":
+            try:
+                coerced = self._coerce(raw, value_type)
+                debug_log(f"DistributedValue - Worker key {lookup_key}: returning '{coerced}'")
+                return (coerced,)
+            except (TypeError, ValueError) as e:
+                debug_log(f"DistributedValue - Error coercing worker value for key {lookup_key}: {e}")
+        debug_log(f"DistributedValue - Worker fallback: returning default '{coerced_default}'")
+        return (coerced_default,)
+
+class DistributedModelName:
+    @classmethod
+    def INPUT_TYPES(cls: type["DistributedModelName"]) -> dict[str, Any]:
+        return {
+            "required": {
+                "text": ("STRING", {"default": ""}),
+            },
+            "hidden": {
+                "unique_id": "UNIQUE_ID",
+                "extra_pnginfo": "EXTRA_PNGINFO",
+            },
+        }
+
+    RETURN_TYPES = (any_type,)
+    RETURN_NAMES = ("output",)
+    FUNCTION = "log_input"
+    OUTPUT_NODE = True
+    CATEGORY = "utils"
+
+    def _stringify(self, value: Any) -> str:
+        if isinstance(value, str):
+            return value
+        if isinstance(value, (int, float, bool)):
+            return str(value)
+        try:
+            return json.dumps(value, indent=4)
+        except Exception:
+            return str(value)
+
+    def _update_workflow(self, extra_pnginfo: Any, unique_id: Any, values: list[str]) -> None:
+        if not extra_pnginfo:
+            return
+        info = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo
+        if not isinstance(info, dict) or "workflow" not in info:
+            return
+        node_id = None
+        if isinstance(unique_id, list) and unique_id:
+            node_id = str(unique_id[0])
+        elif unique_id is not None:
+            node_id = str(unique_id)
+        if not node_id:
+            return
+        workflow = info["workflow"]
+        node = next((x for x in workflow["nodes"] if str(x.get("id")) == node_id), None)
+        if node:
+            node["widgets_values"] = [values]
+
+    def log_input(
+        self,
+        text: Any,
+        unique_id: Any = None,
+        extra_pnginfo: Any = None,
+    ) -> dict[str, Any]:
+        values = []
+        if isinstance(text, list):
+            for val in text:
+                values.append(self._stringify(val))
+        else:
+            values.append(self._stringify(text))
+
+        # Keep widget display in workflow metadata if available.
+        self._update_workflow(extra_pnginfo, unique_id, values)
+
+        if isinstance(values, list) and len(values) == 1:
+            return {"ui": {"text": values}, "result": (values[0],)}
+        return {"ui": {"text": values}, "result": (values,)}
+
+class ByPassTypeTuple(tuple):
+    def __getitem__(self, index: int) -> Any:
+        if index > 0:
+            index = 0
+        item = super().__getitem__(index)
+        if isinstance(item, str):
+            return any_type
+        return item
+
+class ImageBatchDivider:
+    @classmethod
+    def INPUT_TYPES(cls: type["ImageBatchDivider"]) -> dict[str, Any]:
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "divide_by": ("INT", {
+                    "default": 2, 
+                    "min": 1, 
+                    "max": 10, 
+                    "step": 1,
+                    "display": "number",
+                    "tooltip": "Number of parts to divide the batch into"
+                }),
+            }
+        }
+    
+    RETURN_TYPES = ByPassTypeTuple(("IMAGE", ))  # Flexible for variable outputs
+    RETURN_NAMES = ByPassTypeTuple(tuple([f"batch_{i+1}" for i in range(10)]))
+    FUNCTION = "divide_batch"
+    OUTPUT_NODE = True
+    CATEGORY = "image"
+    
+    def divide_batch(self, images: torch.Tensor, divide_by: int) -> tuple[torch.Tensor, ...]:
+        total_splits = max(1, min(int(divide_by), 10))
+        total_frames = images.shape[0]
+        empty_tensor = images[:0]
+        bounds = _chunk_bounds(total_frames, total_splits)
+        outputs = [images[start:end] if end > start else empty_tensor for start, end in bounds]
+
+        while len(outputs) < 10:
+            outputs.append(empty_tensor)
+
+        return tuple(outputs[:10])
+
+
+class AudioBatchDivider:
+    """Divides an audio waveform into multiple parts along the time/samples dimension."""
+
+    @classmethod
+    def INPUT_TYPES(cls: type["AudioBatchDivider"]) -> dict[str, Any]:
+        return {
+            "required": {
+                "audio": ("AUDIO",),
+                "divide_by": ("INT", {
+                    "default": 2,
+                    "min": 1,
+                    "max": 10,
+                    "step": 1,
+                    "display": "number",
+                    "tooltip": "Number of parts to divide the audio into"
+                }),
+            }
+        }
+
+    RETURN_TYPES = ByPassTypeTuple(("AUDIO",))  # Flexible for variable outputs
+    RETURN_NAMES = ByPassTypeTuple(tuple([f"audio_{i+1}" for i in range(10)]))
+    FUNCTION = "divide_audio"
+    OUTPUT_NODE = True
+    CATEGORY = "audio"
+
+    def divide_audio(self, audio: dict[str, Any], divide_by: int) -> tuple[dict[str, Any], ...]:
+        import torch
+
+        waveform = audio.get("waveform")
+        sample_rate = audio.get("sample_rate", 44100)
+
+        if waveform is None or waveform.numel() == 0:
+            # Return empty audio for all outputs
+            empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": sample_rate}
+            return tuple([empty_audio] * 10)
+
+        total_splits = max(1, min(int(divide_by), 10))
+        total_samples = int(waveform.shape[-1])
+        bounds = _chunk_bounds(total_samples, total_splits)
+
+        outputs = []
+        empty_waveform = waveform[..., :0]
+        for start, end in bounds:
+            split_waveform = waveform[..., start:end] if end > start else empty_waveform
+            outputs.append({
+                "waveform": split_waveform,
+                "sample_rate": sample_rate
+            })
+
+        # Pad with empty audio up to max (10) to match RETURN_TYPES length
+        empty_audio = {
+            "waveform": empty_waveform,
+            "sample_rate": sample_rate
+        }
+
+        while len(outputs) < 10:
+            outputs.append(empty_audio)
+
+        return tuple(outputs)
+
+
+class DistributedEmptyImage:
+    """Produces an empty IMAGE batch used when the master delegates all work."""
+
+    @classmethod
+    def INPUT_TYPES(cls: type["DistributedEmptyImage"]) -> dict[str, Any]:
+        return {
+            "required": {
+                "height": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}),
+                "width": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}),
+                "channels": ("INT", {"default": 3, "min": 1, "max": 4, "step": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "create"
+    CATEGORY = "image"
+
+    def create(self, height: int, width: int, channels: int) -> tuple[torch.Tensor]:
+        import torch
+
+        shape = (0, height, width, channels)
+        tensor = torch.zeros(shape, dtype=torch.float32)
+        return (tensor,)

package.json

@@ -0,0 +1,12 @@
+{
+  "name": "comfyui-distributed-web-tests",
+  "private": true,
+  "type": "module",
+  "scripts": {
+    "test:web": "bash ./scripts/test-web.sh",
+    "test:web:watch": "bash ./scripts/test-web.sh --watch"
+  },
+  "devDependencies": {
+    "vitest": "^2.1.9"
+  }
+}

pyproject.toml

@@ -1,9 +1,12 @@
 [project]
 name = "ComfyUI-Distributed"
 description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
-version = "1.1.0"
+version = "1.4.1"
 license = {file = "LICENSE"}
-dependencies  = []
+dependencies = [
+    "aiohttp>=3.9,<4",
+    "Pillow>=10,<12",
+]
 
 [project.urls]
 Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
@@ -13,3 +16,7 @@ Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
 PublisherId = "robertvoy"
 DisplayName = "ComfyUI-Distributed"
 Icon = "https://raw.githubusercontent.com/robertvoy/ComfyUI-Distributed/refs/heads/main/web/distributed-logo-icon.png"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["."]

scripts/test-web.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"
+
+export NVM_DIR="${NVM_DIR:-$HOME/.nvm}"
+if [[ -s "${NVM_DIR}/nvm.sh" ]]; then
+    # shellcheck source=/dev/null
+    . "${NVM_DIR}/nvm.sh"
+fi
+
+if ! command -v node >/dev/null 2>&1 || ! command -v npm >/dev/null 2>&1; then
+    echo "[test-web] node/npm are not available." >&2
+    echo "[test-web] Install nvm and Node, or ensure node/npm are on PATH." >&2
+    exit 1
+fi
+
+if [[ -f "${REPO_ROOT}/.nvmrc" ]] && command -v nvm >/dev/null 2>&1; then
+    nvm use >/dev/null
+fi
+
+cd "${REPO_ROOT}"
+
+if [[ "${1:-}" == "--watch" ]]; then
+    exec npx vitest web/tests
+fi
+
+exec npx vitest run web/tests

tests/api/harness.py

@@ -0,0 +1,163 @@
+"""Shared stubs/utilities for API route module unit tests."""
+from __future__ import annotations
+
+import sys
+import types
+from typing import Any, Callable
+
+
+class RoutesStub:
+    """Simple decorator-style route registrar used in unit tests."""
+
+    def get(self, _path: str):
+        return lambda fn: fn
+
+    def post(self, _path: str):
+        return lambda fn: fn
+
+
+class ClientTimeoutStub:
+    def __init__(self, total: float | None = None):
+        self.total = total
+
+
+class WSMsgTypeStub:
+    TEXT = "TEXT"
+    ERROR = "ERROR"
+    CLOSED = "CLOSED"
+
+
+class WebSocketResponseStub:
+    def __init__(self, *args: Any, **kwargs: Any):
+        self.args = args
+        self.kwargs = kwargs
+
+    async def prepare(self, _request: Any) -> None:
+        return None
+
+    async def send_json(self, _payload: Any) -> None:
+        return None
+
+    def __aiter__(self):
+        async def _empty():
+            if False:
+                yield None
+
+        return _empty()
+
+
+class FormDataStub:
+    def add_field(self, *_args: Any, **_kwargs: Any) -> None:
+        return None
+
+
+def reset_package_namespace(package_name: str) -> None:
+    """Clear test package modules from sys.modules."""
+    prefix = f"{package_name}."
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(prefix):
+            del sys.modules[mod_name]
+
+
+def ensure_namespace_package(module_name: str) -> types.ModuleType:
+    """Create a namespace-like package module in sys.modules."""
+    module = types.ModuleType(module_name)
+    module.__path__ = []
+    sys.modules[module_name] = module
+    return module
+
+
+def bootstrap_test_package(
+    package_name: str,
+    *,
+    with_api: bool = True,
+    with_utils: bool = True,
+    with_workers: bool = False,
+    with_upscale: bool = False,
+    with_orchestration: bool = False,
+) -> None:
+    """Create a clean package scaffold for module-level import tests."""
+    reset_package_namespace(package_name)
+    ensure_namespace_package(package_name)
+
+    if with_api:
+        ensure_namespace_package(f"{package_name}.api")
+    if with_utils:
+        ensure_namespace_package(f"{package_name}.utils")
+    if with_workers:
+        ensure_namespace_package(f"{package_name}.workers")
+    if with_upscale:
+        ensure_namespace_package(f"{package_name}.upscale")
+    if with_orchestration:
+        ensure_namespace_package(f"{package_name}.api.orchestration")
+
+
+def install_server_stub(prompt_server_instance: Any | None = None) -> Any:
+    """Install a minimal `server` module exposing PromptServer.instance.routes."""
+    if prompt_server_instance is None:
+        prompt_server_instance = types.SimpleNamespace()
+    if not hasattr(prompt_server_instance, "routes"):
+        prompt_server_instance.routes = RoutesStub()
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
+    sys.modules["server"] = server_module
+    return prompt_server_instance
+
+
+def install_aiohttp_stub(
+    json_response_factory: Callable[..., Any],
+) -> bool:
+    """Install/augment `aiohttp` with common members used by API tests."""
+    created = False
+    if "aiohttp" not in sys.modules:
+        created = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        sys.modules["aiohttp"] = aiohttp_module
+
+    aiohttp_module = sys.modules["aiohttp"]
+    if not hasattr(aiohttp_module, "ClientTimeout"):
+        aiohttp_module.ClientTimeout = ClientTimeoutStub
+    if not hasattr(aiohttp_module, "WSMsgType"):
+        aiohttp_module.WSMsgType = WSMsgTypeStub
+    if not hasattr(aiohttp_module, "FormData"):
+        aiohttp_module.FormData = FormDataStub
+
+    web_obj = getattr(aiohttp_module, "web", None)
+    if web_obj is None:
+        web_obj = types.SimpleNamespace()
+        aiohttp_module.web = web_obj
+    web_obj.json_response = json_response_factory
+    if not hasattr(web_obj, "WebSocketResponse"):
+        web_obj.WebSocketResponse = WebSocketResponseStub
+
+    return created
+
+
+def cleanup_optional_module(module_name: str, created: bool) -> None:
+    """Remove stubbed module only when this test created it."""
+    if created:
+        sys.modules.pop(module_name, None)
+
+
+def install_request_guards_stub(package_name: str) -> None:
+    """Install permissive request guard for isolated route tests."""
+    request_guards_module = types.ModuleType(f"{package_name}.api.request_guards")
+
+    async def _authorization_error_or_none(_request):
+        return None
+
+    request_guards_module.authorization_error_or_none = _authorization_error_or_none
+    sys.modules[f"{package_name}.api.request_guards"] = request_guards_module
+
+
+def install_endpoint_policy_passthrough_stub(package_name: str) -> None:
+    """Install endpoint policy stub that executes operation directly."""
+    endpoint_policy_module = types.ModuleType(f"{package_name}.api.endpoint_policy")
+
+    async def _run_authorized_endpoint(_request, operation, unexpected_status=500):
+        _ = unexpected_status
+        return await operation()
+
+    endpoint_policy_module.run_authorized_endpoint = _run_authorized_endpoint
+    sys.modules[f"{package_name}.api.endpoint_policy"] = endpoint_policy_module

tests/api/test_auth_guard_integration.py

@@ -0,0 +1,423 @@
+import asyncio
+import copy
+import importlib.util
+import sys
+import types
+import unittest
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_server_stub,
+)
+
+
+AUTH_TOKEN = "integration-secret"
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, *, headers=None, json_payload=None, post_payload=None, query=None):
+        self.headers = headers or {}
+        self._json_payload = json_payload or {}
+        self._post_payload = post_payload or {}
+        self.query = query or {}
+
+    async def json(self):
+        return self._json_payload
+
+    async def post(self):
+        return self._post_payload
+
+
+def _load_module(package_name: str, rel_path: str, module_name: str):
+    module_path = Path(__file__).resolve().parents[2] / rel_path
+    spec = importlib.util.spec_from_file_location(f"{package_name}.{module_name}", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _install_common_utils(package_name: str, config_data: dict):
+    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.load_config = lambda: copy.deepcopy(config_data)
+    config_module.save_config = lambda _cfg: True
+
+    @asynccontextmanager
+    async def _config_transaction():
+        yield config_module.load_config()
+
+    config_module.config_transaction = _config_transaction
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    network_module.normalize_host = lambda value: value
+    network_module.build_worker_url = lambda worker, path="": f"http://{worker.get('host', '127.0.0.1')}{path}"
+    network_module.get_client_session = lambda: None
+    network_module.probe_worker = lambda *_args, **_kwargs: {"ok": True}
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+
+def _load_real_guard_stack(package_name: str):
+    _load_module(package_name, "utils/auth.py", "utils.auth")
+    _load_module(package_name, "utils/parsing.py", "utils.parsing")
+    _load_module(package_name, "api/schemas.py", "api.schemas")
+    _load_module(package_name, "api/request_guards.py", "api.request_guards")
+    _load_module(package_name, "api/endpoint_policy.py", "api.endpoint_policy")
+
+
+def _load_request_guards_only_module():
+    package_name = "dist_auth_guard_testpkg"
+    config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+    request_guards = sys.modules[f"{package_name}.api.request_guards"]
+    schemas = sys.modules[f"{package_name}.api.schemas"]
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    return request_guards, schemas
+
+
+def _load_config_routes_real_guard_module():
+    package_name = "dist_config_auth_integration_testpkg"
+    config_data = {
+        "workers": [],
+        "master": {"host": "127.0.0.1"},
+        "settings": {"distributed_api_token": AUTH_TOKEN},
+        "tunnel": {},
+    }
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+    module = _load_module(package_name, "api/config_routes.py", "api.config_routes")
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    return module
+
+
+def _load_tunnel_routes_real_guard_module():
+    package_name = "dist_tunnel_auth_integration_testpkg"
+    config_data = {
+        "workers": [],
+        "master": {"host": "127.0.0.1"},
+        "settings": {"distributed_api_token": AUTH_TOKEN},
+        "tunnel": {},
+    }
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+
+    cloudflare_module = types.ModuleType(f"{package_name}.utils.cloudflare")
+    cloudflare_module.cloudflare_tunnel_manager = types.SimpleNamespace(
+        get_status=lambda: {"active": False},
+        start_tunnel=lambda: {"active": True},
+        stop_tunnel=lambda: {"active": False},
+    )
+    sys.modules[f"{package_name}.utils.cloudflare"] = cloudflare_module
+
+    module = _load_module(package_name, "api/tunnel_routes.py", "api.tunnel_routes")
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    return module
+
+
+def _load_usdu_routes_real_guard_module():
+    package_name = "dist_usdu_auth_integration_testpkg"
+    config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+
+    usdu_management_module = types.ModuleType(f"{package_name}.utils.usdu_management")
+    usdu_management_module.MAX_PAYLOAD_SIZE = 1024 * 1024
+    sys.modules[f"{package_name}.utils.usdu_management"] = usdu_management_module
+
+    prompt_server_holder = {
+        "value": types.SimpleNamespace(
+            distributed_tile_jobs_lock=asyncio.Lock(),
+            distributed_pending_tile_jobs={},
+        )
+    }
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
+    job_store_module.init_dynamic_job = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
+
+    class _BaseJobState:
+        pass
+
+    class _ImageJobState(_BaseJobState):
+        def __init__(self):
+            self.mode = "dynamic"
+            self.queue = asyncio.Queue()
+            self.pending_images = asyncio.Queue()
+            self.completed_images = {}
+            self.worker_status = {}
+            self.assigned_to_workers = {}
+
+    class _TileJobState(_BaseJobState):
+        def __init__(self):
+            self.mode = "static"
+            self.queue = asyncio.Queue()
+            self.pending_tasks = asyncio.Queue()
+            self.completed_tasks = {}
+            self.worker_status = {}
+            self.assigned_to_workers = {}
+            self.batched_static = False
+
+    job_models_module.BaseJobState = _BaseJobState
+    job_models_module.ImageJobState = _ImageJobState
+    job_models_module.TileJobState = _TileJobState
+    sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
+
+    parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers")
+    parsers_module.parse_tiles_from_form = lambda _data: []
+    sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
+
+    module = _load_module(package_name, "api/usdu_routes.py", "api.usdu_routes")
+    module._prompt_server_holder = prompt_server_holder
+    module._ImageJobState = _ImageJobState
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    return module
+
+
+def _install_torch_stub_if_missing():
+    if "torch" in sys.modules:
+        return False
+
+    torch_module = types.ModuleType("torch")
+    torch_module.cuda = types.SimpleNamespace(
+        is_available=lambda: False,
+        empty_cache=lambda: None,
+        ipc_collect=lambda: None,
+    )
+    sys.modules["torch"] = torch_module
+    return True
+
+
+def _load_worker_routes_real_guard_module():
+    package_name = "dist_worker_auth_integration_testpkg"
+    config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}, "workers": []}
+    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_workers=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    created_torch_stub = _install_torch_stub_if_missing()
+
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.CHUNK_SIZE = 65536
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    workers_module = types.ModuleType(f"{package_name}.workers")
+    workers_module.get_worker_manager = lambda: types.SimpleNamespace(
+        processes={},
+        save_processes=lambda: None,
+        is_process_running=lambda _pid: False,
+    )
+    sys.modules[f"{package_name}.workers"] = workers_module
+
+    detection_module = types.ModuleType(f"{package_name}.workers.detection")
+    detection_module.get_machine_id = lambda: "machine-id"
+    detection_module.is_docker_environment = lambda: False
+    detection_module.is_runpod_environment = lambda: False
+    sys.modules[f"{package_name}.workers.detection"] = detection_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+
+    class _PromptValidationError(Exception):
+        def __init__(self, message="validation error"):
+            super().__init__(message)
+            self.validation_error = {}
+            self.node_errors = {}
+
+    async def _queue_prompt_payload(*_args, **_kwargs):
+        return "prompt-id"
+
+    async_helpers_module.PromptValidationError = _PromptValidationError
+    async_helpers_module.queue_prompt_payload = _queue_prompt_payload
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    module = _load_module(package_name, "api/worker_routes.py", "api.worker_routes")
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    if created_torch_stub:
+        sys.modules.pop("torch", None)
+    return module
+
+
+def _load_job_routes_real_guard_module():
+    package_name = "dist_job_auth_integration_testpkg"
+    config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    created_torch_stub = _install_torch_stub_if_missing()
+
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+    image_module.pil_to_tensor = lambda _img: None
+    image_module.ensure_contiguous = lambda tensor: tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.JOB_INIT_GRACE_PERIOD = 0.1
+    constants_module.MEMORY_CLEAR_DELAY = 0.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
+    runtime_state = types.SimpleNamespace(
+        distributed_jobs_lock=asyncio.Lock(),
+        distributed_pending_jobs={},
+    )
+    runtime_state_module.ensure_distributed_runtime_state = lambda: runtime_state
+    sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
+
+    orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
+    orchestration_module.orchestrate_distributed_execution = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.api.queue_orchestration"] = orchestration_module
+
+    queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
+    queue_request_module.parse_queue_request_payload = lambda data: data
+    sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
+
+    module = _load_module(package_name, "api/job_routes.py", "api.job_routes")
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    if created_torch_stub:
+        sys.modules.pop("torch", None)
+    return module
+
+
+request_guards_module, schemas_module = _load_request_guards_only_module()
+config_routes_module = _load_config_routes_real_guard_module()
+tunnel_routes_module = _load_tunnel_routes_real_guard_module()
+usdu_routes_module = _load_usdu_routes_real_guard_module()
+worker_routes_module = _load_worker_routes_real_guard_module()
+job_routes_module = _load_job_routes_real_guard_module()
+
+
+class AuthGuardContractTests(unittest.IsolatedAsyncioTestCase):
+    async def test_is_authorized_request_missing_token_fails(self):
+        request = _FakeRequest(headers={})
+        config = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+        self.assertFalse(schemas_module.is_authorized_request(request, config))
+
+    async def test_is_authorized_request_matching_token_succeeds(self):
+        request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN})
+        config = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+        self.assertTrue(schemas_module.is_authorized_request(request, config))
+
+    async def test_authorization_error_or_none_returns_403_for_invalid_token(self):
+        request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: "wrong-token"})
+        response = await request_guards_module.authorization_error_or_none(request)
+        self.assertIsNotNone(response)
+        self.assertEqual(response.status, 403)
+
+    async def test_authorization_error_or_none_returns_none_for_valid_token(self):
+        request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN})
+        response = await request_guards_module.authorization_error_or_none(request)
+        self.assertIsNone(response)
+
+
+class RouteAuthIntegrationTests(unittest.IsolatedAsyncioTestCase):
+    def _auth_headers(self):
+        return {schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN}
+
+    async def test_config_endpoint_enforces_real_guard(self):
+        unauthorized = await config_routes_module.get_config_endpoint(_FakeRequest(headers={}))
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await config_routes_module.get_config_endpoint(_FakeRequest(headers=self._auth_headers()))
+        self.assertEqual(authorized.status, 200)
+        self.assertIn("settings", authorized.payload)
+
+    async def test_tunnel_endpoint_enforces_real_guard(self):
+        unauthorized = await tunnel_routes_module.tunnel_status_endpoint(_FakeRequest(headers={}))
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await tunnel_routes_module.tunnel_status_endpoint(_FakeRequest(headers=self._auth_headers()))
+        self.assertEqual(authorized.status, 200)
+        self.assertEqual(authorized.payload.get("status"), "success")
+
+    async def test_usdu_endpoint_enforces_real_guard(self):
+        unauthorized = await usdu_routes_module.job_status_endpoint(_FakeRequest(headers={}))
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await usdu_routes_module.job_status_endpoint(
+            _FakeRequest(headers=self._auth_headers(), query={"multi_job_id": "job-1"})
+        )
+        self.assertEqual(authorized.status, 200)
+        self.assertIn("ready", authorized.payload)
+
+    async def test_worker_endpoint_enforces_real_guard(self):
+        worker_routes_module._collect_network_info_sync = lambda: {
+            "interfaces": [],
+            "recommended_host": "127.0.0.1",
+        }
+        unauthorized = await worker_routes_module.get_network_info_endpoint(_FakeRequest(headers={}))
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await worker_routes_module.get_network_info_endpoint(
+            _FakeRequest(headers=self._auth_headers())
+        )
+        self.assertEqual(authorized.status, 200)
+        self.assertEqual(authorized.payload.get("status"), "success")
+
+    async def test_job_endpoint_enforces_real_guard(self):
+        unauthorized = await job_routes_module.prepare_job_endpoint(
+            _FakeRequest(headers={}, json_payload={"multi_job_id": "job-1"})
+        )
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await job_routes_module.prepare_job_endpoint(
+            _FakeRequest(headers=self._auth_headers(), json_payload={"multi_job_id": "job-1"})
+        )
+        self.assertEqual(authorized.status, 200)
+        self.assertEqual(authorized.payload.get("status"), "success")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_config_routes.py

@@ -0,0 +1,143 @@
+import copy
+import importlib.util
+import sys
+import types
+import unittest
+from contextlib import asynccontextmanager
+from pathlib import Path
+from unittest.mock import patch
+
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_endpoint_policy_passthrough_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, payload=None):
+        self._payload = payload
+
+    async def json(self):
+        return self._payload
+
+
+def _load_config_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "config_routes.py"
+    package_name = "dist_api_config_testpkg"
+
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+    schemas_module.is_authorized_request = lambda _request, _config: True
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    install_request_guards_stub(package_name)
+    install_endpoint_policy_passthrough_stub(package_name)
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    install_server_stub()
+
+    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")
+
+    async def _handle_api_error(_request, error, status=500):
+        if isinstance(error, list):
+            message = "; ".join(str(item) for item in error)
+            error_payload = [str(item) for item in error]
+        else:
+            message = str(error)
+            error_payload = str(error)
+        return _FakeResponse(
+            {"status": "error", "error": error_payload, "message": message},
+            status=status,
+        )
+
+    network_module.handle_api_error = _handle_api_error
+    network_module.normalize_host = lambda value: value
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    default_config = {
+        "workers": [],
+        "master": {"host": ""},
+        "settings": {"debug": False},
+        "tunnel": {},
+    }
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: copy.deepcopy(default_config)
+    config_module.save_config = lambda _cfg: True
+
+    @asynccontextmanager
+    async def _config_transaction():
+        yield config_module.load_config()
+
+    config_module.config_transaction = _config_transaction
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.config_routes", 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)
+
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+
+    return module
+
+
+config_routes = _load_config_routes_module()
+
+
+class ConfigRoutesTests(unittest.IsolatedAsyncioTestCase):
+    async def test_get_config_returns_core_sections(self):
+        cfg = {"workers": [], "master": {}, "settings": {}, "tunnel": {}}
+        with patch.object(config_routes, "load_config", return_value=cfg):
+            response = await config_routes.get_config_endpoint(_FakeRequest())
+
+        self.assertEqual(response.status, 200)
+        self.assertIn("workers", response.payload)
+        self.assertIn("master", response.payload)
+        self.assertIn("settings", response.payload)
+
+    async def test_update_config_valid_field_persists(self):
+        cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
+        with patch.object(config_routes, "load_config", return_value=cfg):
+            response = await config_routes.update_config_endpoint(_FakeRequest({"debug": True}))
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload["status"], "success")
+        self.assertTrue(response.payload["config"]["settings"]["debug"])
+
+    async def test_update_config_unknown_field_returns_400(self):
+        cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
+        with patch.object(config_routes, "load_config", return_value=cfg):
+            response = await config_routes.update_config_endpoint(_FakeRequest({"unknown_field": 1}))
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("unknown_field", " ".join(response.payload.get("error", [])).lower())
+
+    async def test_update_config_wrong_type_returns_400(self):
+        cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
+        with patch.object(config_routes, "load_config", return_value=cfg):
+            response = await config_routes.update_config_endpoint(_FakeRequest({"debug": "true"}))
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("debug", " ".join(response.payload.get("error", [])).lower())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_job_routes.py

@@ -0,0 +1,360 @@
+import importlib.util
+import sys
+import types
+import unittest
+import asyncio
+import base64
+from dataclasses import dataclass
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+import numpy as np
+import torch
+
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, payload, headers=None):
+        self._payload = payload
+        self.headers = headers or {}
+
+    async def json(self):
+        return self._payload
+
+
+def _load_job_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
+    package_name = "dist_api_queue_testpkg"
+
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+
+    # aiohttp.web stub
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+
+    # server module stub with route decorators
+    prompt_server_instance = types.SimpleNamespace(
+        distributed_jobs_lock=None,
+        distributed_pending_jobs={},
+    )
+    install_server_stub(prompt_server_instance)
+
+    # torch stub (only needed to satisfy import)
+    created_torch_stub = False
+    if "torch" not in sys.modules:
+        created_torch_stub = True
+        torch_module = types.ModuleType("torch")
+        torch_module.cuda = types.SimpleNamespace(
+            is_available=lambda: False,
+            empty_cache=lambda: None,
+            ipc_collect=lambda: None,
+        )
+        sys.modules["torch"] = torch_module
+
+    # PIL stub (only needed to satisfy import)
+    created_pil_stub = False
+    if "PIL" not in sys.modules:
+        created_pil_stub = True
+        pil_module = types.ModuleType("PIL")
+        image_module = types.ModuleType("PIL.Image")
+        pil_module.Image = image_module
+        sys.modules["PIL"] = pil_module
+        sys.modules["PIL.Image"] = image_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
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+    image_module.pil_to_tensor = lambda *_args, **_kwargs: None
+    image_module.ensure_contiguous = lambda tensor: tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
+
+    def _decode_audio_payload(payload):
+        if payload is None:
+            return None
+        if not isinstance(payload, dict):
+            raise ValueError("Field 'audio' must be an object when provided.")
+
+        encoded = payload.get("data")
+        shape = payload.get("shape")
+        dtype = payload.get("dtype", "float32")
+        sample_rate = payload.get("sample_rate", 44100)
+        if not isinstance(encoded, str) or not encoded.strip():
+            raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
+        if not isinstance(shape, list) or len(shape) != 3:
+            raise ValueError("Field 'audio.shape' must be a 3-item list.")
+        if dtype != "float32":
+            raise ValueError("Field 'audio.dtype' must be 'float32'.")
+        try:
+            shape_tuple = tuple(int(dim) for dim in shape)
+        except Exception as exc:
+            raise ValueError("Field 'audio.shape' must contain integers.") from exc
+
+        raw = base64.b64decode(encoded, validate=True)
+        expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
+        if len(raw) != expected_bytes:
+            raise ValueError("Field 'audio.data' byte size mismatch.")
+
+        waveform = torch.from_numpy(np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple).copy())
+        return {"waveform": waveform, "sample_rate": int(sample_rate)}
+
+    audio_payload_module.decode_audio_payload = _decode_audio_payload
+    sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.MEMORY_CLEAR_DELAY = 0.0
+    constants_module.JOB_INIT_GRACE_PERIOD = 10.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"settings": {}}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
+
+    def _ensure_distributed_runtime_state(server_instance=None):
+        ps = server_instance or prompt_server_instance
+        if not hasattr(ps, "distributed_pending_jobs"):
+            ps.distributed_pending_jobs = {}
+        if not hasattr(ps, "distributed_jobs_lock") or ps.distributed_jobs_lock is None:
+            ps.distributed_jobs_lock = asyncio.Lock()
+        if not hasattr(ps, "distributed_job_allowed_workers"):
+            ps.distributed_job_allowed_workers = {}
+        if not hasattr(ps, "distributed_pending_tile_jobs"):
+            ps.distributed_pending_tile_jobs = {}
+        if not hasattr(ps, "distributed_tile_jobs_lock"):
+            ps.distributed_tile_jobs_lock = asyncio.Lock()
+        return types.SimpleNamespace(
+            distributed_pending_jobs=ps.distributed_pending_jobs,
+            distributed_jobs_lock=ps.distributed_jobs_lock,
+            distributed_job_allowed_workers=ps.distributed_job_allowed_workers,
+            distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs,
+            distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock,
+        )
+
+    runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state
+    runtime_state_module.get_prompt_server_instance = lambda: prompt_server_instance
+    sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+    schemas_module.is_authorized_request = lambda _request, _config: True
+    schemas_module.require_bool_literal = (
+        lambda value, field_name="value": value
+        if isinstance(value, bool)
+        else (_ for _ in ()).throw(ValueError(f"Field '{field_name}' must be a boolean literal."))
+    )
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    install_request_guards_stub(package_name)
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
+    queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1))
+    sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
+
+    @dataclass(frozen=True)
+    class _QueuePayload:
+        prompt: dict
+        workflow_meta: object
+        client_id: str
+        delegate_master: object
+        enabled_worker_ids: list
+        trace_execution_id: object
+
+    def _parse_queue_request_payload(data):
+        if not isinstance(data, dict):
+            raise ValueError("Expected a JSON object body")
+        prompt = data.get("prompt")
+        if not isinstance(prompt, dict):
+            raise ValueError("Field 'prompt' must be an object")
+        enabled = data.get("enabled_worker_ids")
+        if not isinstance(enabled, list):
+            raise ValueError("enabled_worker_ids required")
+        client_id = data.get("client_id")
+        if not isinstance(client_id, str) or not client_id.strip():
+            raise ValueError("client_id required")
+        return _QueuePayload(
+            prompt=prompt,
+            workflow_meta=data.get("workflow"),
+            client_id=client_id,
+            delegate_master=data.get("delegate_master"),
+            enabled_worker_ids=enabled,
+            trace_execution_id=data.get("trace_execution_id"),
+        )
+
+    queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
+    queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
+    sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_routes", 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)
+
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    cleanup_optional_module("torch", created_torch_stub)
+    if created_pil_stub:
+        cleanup_optional_module("PIL.Image", True)
+        cleanup_optional_module("PIL", True)
+
+    return module
+
+
+job_routes = _load_job_routes_module()
+
+
+class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
+    async def test_distributed_queue_happy_path_returns_prompt_id(self):
+        request = _FakeRequest(
+            {
+                "prompt": {"1": {"class_type": "Node"}},
+                "enabled_worker_ids": ["w1"],
+                "client_id": "client-1",
+                "auto_prepare": True,
+            }
+        )
+        with patch.object(
+            job_routes,
+            "orchestrate_distributed_execution",
+            new=AsyncMock(return_value=("prompt_123", 2)),
+        ):
+            response = await job_routes.distributed_queue_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
+
+    async def test_distributed_queue_missing_prompt_returns_400(self):
+        request = _FakeRequest(
+            {
+                "enabled_worker_ids": ["w1"],
+                "client_id": "client-1",
+            }
+        )
+        response = await job_routes.distributed_queue_endpoint(request)
+        self.assertEqual(response.status, 400)
+        self.assertIn("prompt", response.payload.get("message", "").lower())
+
+    async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
+        request = _FakeRequest(
+            {
+                "prompt": {"1": {"class_type": "Node"}},
+                "client_id": "client-1",
+            }
+        )
+        response = await job_routes.distributed_queue_endpoint(request)
+        self.assertEqual(response.status, 400)
+        self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
+
+
+class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
+    def _encoded_audio_payload(self):
+        waveform = np.arange(8, dtype=np.float32).reshape(1, 2, 4)
+        return {
+            "sample_rate": 44100,
+            "shape": [1, 2, 4],
+            "dtype": "float32",
+            "data": base64.b64encode(waveform.tobytes()).decode("ascii"),
+        }
+
+    async def test_job_complete_accepts_audio_payload(self):
+        queue = asyncio.Queue()
+        prompt_server = job_routes.server.PromptServer.instance
+        prompt_server.distributed_jobs_lock = asyncio.Lock()
+        prompt_server.distributed_pending_jobs = {"job-1": queue}
+        prompt_server.distributed_job_allowed_workers = {"job-1": {"worker-1"}}
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "image": "data:image/png;base64,AAAA",
+                "audio": self._encoded_audio_payload(),
+                "is_last": True,
+            }
+        )
+
+        with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
+            response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        queued = await queue.get()
+        self.assertEqual(queued["worker_id"], "worker-1")
+        self.assertTrue(queued["is_last"])
+        self.assertIsNotNone(queued["audio"])
+        self.assertEqual(queued["audio"]["sample_rate"], 44100)
+        self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
+
+    async def test_job_complete_rejects_worker_not_in_job_allowlist(self):
+        queue = asyncio.Queue()
+        prompt_server = job_routes.server.PromptServer.instance
+        prompt_server.distributed_jobs_lock = asyncio.Lock()
+        prompt_server.distributed_pending_jobs = {"job-allow": queue}
+        prompt_server.distributed_job_allowed_workers = {"job-allow": {"worker-expected"}}
+        request = _FakeRequest(
+            {
+                "job_id": "job-allow",
+                "worker_id": "worker-unexpected",
+                "batch_idx": 0,
+                "image": "data:image/png;base64,AAAA",
+                "is_last": True,
+            }
+        )
+
+        with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
+            response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 403)
+        self.assertIn("unauthorized", response.payload.get("message", "").lower())
+
+    def test_decode_audio_payload_rejects_bad_shape(self):
+        bad = {
+            "sample_rate": 44100,
+            "shape": [1, 2],
+            "dtype": "float32",
+            "data": base64.b64encode(b"\x00\x00\x00\x00").decode("ascii"),
+        }
+        with self.assertRaises(ValueError):
+            job_routes._decode_audio_payload(bad)
+
+    def test_decode_audio_payload_rejects_bad_dtype(self):
+        payload = {
+            "sample_rate": 44100,
+            "shape": [1, 2, 4],
+            "dtype": "float16",
+            "data": base64.b64encode((np.zeros((1, 2, 4), dtype=np.float32)).tobytes()).decode("ascii"),
+        }
+        with self.assertRaises(ValueError):
+            job_routes._decode_audio_payload(payload)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_media_sync.py

@@ -0,0 +1,259 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+)
+
+
+class _AiohttpResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+def _load_media_sync_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py"
+    package_name = "dist_ms_testpkg"
+
+    bootstrap_test_package(
+        package_name,
+        with_api=True,
+        with_utils=True,
+        with_orchestration=True,
+    )
+
+    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_worker_url = lambda worker, endpoint="": f"http://localhost{endpoint}"
+
+    async def _fake_session():
+        raise RuntimeError("network calls not used in pure-function tests")
+
+    network_module.get_client_session = _fake_session
+    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
+    trace_module.trace_info = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.trace_logger"] = trace_module
+
+    auth_module = types.ModuleType(f"{package_name}.utils.auth")
+    auth_module.distributed_auth_headers = lambda _config=None: {}
+    sys.modules[f"{package_name}.utils.auth"] = auth_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"settings": {}}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _AiohttpResponse(payload, status=status)
+    )
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.api.orchestration.media_sync",
+        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)
+
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+
+    return module
+
+
+ms = _load_media_sync_module()
+
+
+# ---------------------------------------------------------------------------
+# convert_paths_for_platform
+# ---------------------------------------------------------------------------
+
+class ConvertPathsForPlatformTests(unittest.TestCase):
+    def test_forward_slash_target_normalises_backslashes(self):
+        obj = {"ckpt_name": "C:\\Models\\model.safetensors"}
+        result = ms.convert_paths_for_platform(obj, "/")
+        self.assertEqual(result["ckpt_name"], "C:/Models/model.safetensors")
+
+    def test_backslash_target_normalises_forward_slashes(self):
+        obj = {"ckpt_name": "/models/checkpoints/model.safetensors"}
+        result = ms.convert_paths_for_platform(obj, "\\")
+        self.assertIn("\\", result["ckpt_name"])
+        self.assertNotIn("/", result["ckpt_name"])
+
+    def test_relative_media_paths_always_stay_forward_slash(self):
+        """Relative image/video/audio paths (Comfy annotated style) must not be backslash-ified."""
+        obj = {"image": "subfolder/my_photo.png"}
+        result = ms.convert_paths_for_platform(obj, "\\")
+        self.assertEqual(result["image"], "subfolder/my_photo.png")
+
+    def test_relative_audio_paths_stay_forward_slash(self):
+        obj = {"audio": "subfolder/my_track.wav"}
+        result = ms.convert_paths_for_platform(obj, "\\")
+        self.assertEqual(result["audio"], "subfolder/my_track.wav")
+
+    def test_annotated_relative_media_path_stays_forward_slash(self):
+        obj = {"image": "input/frame.jpg [abc123]"}
+        result = ms.convert_paths_for_platform(obj, "\\")
+        self.assertIn("/", result["image"])
+        self.assertNotIn("\\", result["image"].split("[")[0])
+
+    def test_non_filename_strings_are_untouched(self):
+        obj = {"prompt": "a beautiful cat", "count": 5}
+        result = ms.convert_paths_for_platform(obj, "\\")
+        self.assertEqual(result["prompt"], "a beautiful cat")
+        self.assertEqual(result["count"], 5)
+
+    def test_url_strings_are_untouched(self):
+        obj = {"url": "https://example.com/model.safetensors"}
+        result = ms.convert_paths_for_platform(obj, "\\")
+        self.assertEqual(result["url"], "https://example.com/model.safetensors")
+
+    def test_invalid_separator_returns_obj_unchanged(self):
+        obj = {"ckpt_name": "/models/model.safetensors"}
+        result = ms.convert_paths_for_platform(obj, "|")
+        self.assertEqual(result, obj)
+
+    def test_nested_dict_is_processed_recursively(self):
+        obj = {"node": {"ckpt_name": "C:\\Models\\model.safetensors"}}
+        result = ms.convert_paths_for_platform(obj, "/")
+        self.assertEqual(result["node"]["ckpt_name"], "C:/Models/model.safetensors")
+
+    def test_list_items_are_processed_recursively(self):
+        obj = [{"ckpt_name": "C:\\Models\\model.safetensors"}, "plain string"]
+        result = ms.convert_paths_for_platform(obj, "/")
+        self.assertEqual(result[0]["ckpt_name"], "C:/Models/model.safetensors")
+        self.assertEqual(result[1], "plain string")
+
+    def test_non_string_scalar_values_are_untouched(self):
+        obj = {"seed": 42, "enabled": True, "ratio": 1.5}
+        result = ms.convert_paths_for_platform(obj, "/")
+        self.assertEqual(result["seed"], 42)
+        self.assertTrue(result["enabled"])
+
+    def test_absolute_unix_path_to_windows(self):
+        obj = {"lora": "/home/user/loras/my_lora.safetensors"}
+        result = ms.convert_paths_for_platform(obj, "\\")
+        self.assertNotIn("/", result["lora"])
+
+    def test_already_normalised_path_is_idempotent(self):
+        obj = {"ckpt": "C:/Models/model.safetensors"}
+        result = ms.convert_paths_for_platform(obj, "/")
+        self.assertEqual(result["ckpt"], "C:/Models/model.safetensors")
+
+
+# ---------------------------------------------------------------------------
+# _find_media_references
+# ---------------------------------------------------------------------------
+
+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)
+        self.assertIn("photo.jpg", refs)
+        self.assertFalse(any("[" in r for r in refs))
+
+    def test_normalises_backslashes_in_path(self):
+        prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "sub\\img.png"}}}
+        refs = ms._find_media_references(prompt)
+        self.assertIn("sub/img.png", refs)
+
+    def test_ignores_non_media_text_inputs(self):
+        prompt = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}
+        refs = ms._find_media_references(prompt)
+        self.assertEqual(refs, [])
+
+    def test_ignores_node_link_values(self):
+        """Inputs that are [node_id, slot] lists should be ignored."""
+        prompt = {"1": {"class_type": "Anything", "inputs": {"image": ["2", 0]}}}
+        refs = ms._find_media_references(prompt)
+        self.assertEqual(refs, [])
+
+    def test_deduplicates_same_file_across_nodes(self):
+        prompt = {
+            "1": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
+            "2": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
+        }
+        refs = ms._find_media_references(prompt)
+        self.assertEqual(len(refs), 1)
+
+    def test_returns_sorted_list(self):
+        prompt = {
+            "1": {"class_type": "LoadImage", "inputs": {"image": "z_image.png"}},
+            "2": {"class_type": "LoadImage", "inputs": {"image": "a_image.jpg"}},
+        }
+        refs = ms._find_media_references(prompt)
+        self.assertEqual(refs, sorted(refs))
+
+    def test_ignores_non_dict_nodes(self):
+        prompt = {"1": "not a node dict", "2": {"class_type": "LoadImage", "inputs": {"image": "img.png"}}}
+        refs = ms._find_media_references(prompt)
+        self.assertIn("img.png", refs)
+
+    def test_empty_prompt_returns_empty_list(self):
+        self.assertEqual(ms._find_media_references({}), [])
+
+    def test_multiple_media_types_all_found(self):
+        prompt = {
+            "1": {"class_type": "LoadImage", "inputs": {"image": "frame.png"}},
+            "2": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}},
+            "3": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}},
+        }
+        refs = ms._find_media_references(prompt)
+        self.assertIn("frame.png", refs)
+        self.assertIn("clip.mp4", refs)
+        self.assertIn("track.wav", refs)
+
+
+class RewritePromptMediaInputsTests(unittest.TestCase):
+    def test_rewrites_video_file_input_to_worker_path(self):
+        prompt = {
+            "79": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}},
+        }
+        ms._rewrite_prompt_media_inputs(prompt, {"1 - Copy.mp4": "videos/1 - Copy.mp4"})
+        self.assertEqual(prompt["79"]["inputs"]["file"], "videos/1 - Copy.mp4")
+
+    def test_rewrites_audio_input_and_strips_annotation_when_matching(self):
+        prompt = {
+            "1": {"class_type": "LoadAudio", "inputs": {"audio": "song.wav [input]"}},
+        }
+        ms._rewrite_prompt_media_inputs(prompt, {"song.wav": "song.wav"})
+        self.assertEqual(prompt["1"]["inputs"]["audio"], "song.wav")
+
+
+if __name__ == "__main__":
+    unittest.main()
+    class _AiohttpResponse:
+        def __init__(self, payload, status=200):
+            self.payload = payload
+            self.status = status

tests/api/test_queue_orchestration.py

@@ -0,0 +1,195 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+from tests.api.harness import (
+    bootstrap_test_package,
+    install_server_stub,
+)
+
+
+def _bootstrap_package(package_name):
+    bootstrap_test_package(
+        package_name,
+        with_api=True,
+        with_utils=True,
+        with_orchestration=True,
+    )
+
+    prompt_server_instance = types.SimpleNamespace(address="127.0.0.1", port=8188)
+    install_server_stub(prompt_server_instance)
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+
+    async def _queue_prompt_payload(prompt_obj, workflow_meta, client_id):
+        _ = (prompt_obj, workflow_meta, client_id)
+        return "prompt-id"
+
+    async_helpers_module.queue_prompt_payload = _queue_prompt_payload
+    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: {"settings": {}, "workers": []}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 2
+    constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 120.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 **_kwargs: "http://127.0.0.1:8188"
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    trace_logger_module = types.ModuleType(f"{package_name}.utils.trace_logger")
+    trace_logger_module.trace_debug = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.trace_logger"] = trace_logger_module
+
+    runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
+
+    def _ensure_distributed_runtime_state(server_instance=None):
+        ps = server_instance or prompt_server_instance
+        if not hasattr(ps, "distributed_pending_jobs"):
+            ps.distributed_pending_jobs = {}
+        if not hasattr(ps, "distributed_jobs_lock"):
+            ps.distributed_jobs_lock = asyncio.Lock()
+        if not hasattr(ps, "distributed_job_allowed_workers"):
+            ps.distributed_job_allowed_workers = {}
+        if not hasattr(ps, "distributed_pending_tile_jobs"):
+            ps.distributed_pending_tile_jobs = {}
+        if not hasattr(ps, "distributed_tile_jobs_lock"):
+            ps.distributed_tile_jobs_lock = asyncio.Lock()
+        return types.SimpleNamespace(
+            distributed_pending_jobs=ps.distributed_pending_jobs,
+            distributed_jobs_lock=ps.distributed_jobs_lock,
+            distributed_job_allowed_workers=ps.distributed_job_allowed_workers,
+            distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs,
+            distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock,
+        )
+
+    runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state
+    runtime_state_module.get_prompt_server_instance = lambda: prompt_server_instance
+    sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+    schemas_module.coerce_positive_int = (
+        lambda value, default: int(value) if str(value).isdigit() and int(value) > 0 else default
+    )
+    schemas_module.coerce_positive_float = (
+        lambda value, default: float(value) if isinstance(value, (int, float, str)) and float(value) > 0 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 = lambda *args, **kwargs: None
+    dispatch_module.rank_workers_by_load = lambda workers, **kwargs: workers
+    dispatch_module.select_active_workers = lambda workers, *_args, **_kwargs: (workers, False)
+    dispatch_module.select_least_busy_worker = lambda workers, **_kwargs: workers[0] if workers else None
+    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 = lambda *_args, **_kwargs: "/"
+    media_sync_module.sync_worker_media = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.api.orchestration.media_sync"] = media_sync_module
+
+    prompt_transform_module = types.ModuleType(f"{package_name}.api.orchestration.prompt_transform")
+    prompt_transform_module.PromptIndex = object
+    prompt_transform_module.apply_participant_overrides = lambda prompt, *_args, **_kwargs: prompt
+    prompt_transform_module.find_nodes_by_class = lambda _prompt, _class_name: []
+    prompt_transform_module.generate_job_id_map = lambda _index, _prefix: {}
+    prompt_transform_module.prepare_delegate_master_prompt = lambda prompt, _ids: prompt
+    prompt_transform_module.prune_prompt_for_worker = lambda prompt: prompt
+    sys.modules[f"{package_name}.api.orchestration.prompt_transform"] = prompt_transform_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[2] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_queue_orchestration_module():
+    package_name = "dist_queue_orch_testpkg"
+    _bootstrap_package(package_name)
+    return _load_module(package_name, "api/queue_orchestration.py", "api.queue_orchestration")
+
+
+queue_orchestration = _load_queue_orchestration_module()
+
+
+class QueueOrchestrationHelpersTests(unittest.TestCase):
+    def test_resolve_enabled_workers_filters_disabled_and_requested(self):
+        config = {
+            "workers": [
+                {"id": "1", "enabled": True, "host": "a", "port": 8188},
+                {"id": "2", "enabled": False, "host": "b", "port": 8189},
+                {"id": "3", "enabled": True, "host": "c", "port": "bad-port"},
+            ]
+        }
+
+        enabled_only = queue_orchestration._resolve_enabled_workers(config)
+        requested = queue_orchestration._resolve_enabled_workers(config, requested_ids={"2", "3"})
+
+        self.assertEqual([worker["id"] for worker in enabled_only], ["1", "3"])
+        self.assertEqual([worker["id"] for worker in requested], ["2", "3"])
+        self.assertEqual(requested[1]["port"], 8188)
+
+    def test_resolve_orchestration_limits_uses_defaults_for_invalid_values(self):
+        config = {
+            "settings": {
+                "worker_probe_concurrency": "x",
+                "worker_prep_concurrency": "3",
+                "media_sync_concurrency": "-1",
+                "media_sync_timeout_seconds": "42.5",
+            }
+        }
+
+        limits = queue_orchestration._resolve_orchestration_limits(config)
+
+        self.assertEqual(limits[0], 8)   # default
+        self.assertEqual(limits[1], 3)   # parsed
+        self.assertEqual(limits[2], 2)   # default
+        self.assertEqual(limits[3], 42.5)
+
+    def test_is_load_balance_enabled_accepts_common_forms(self):
+        self.assertTrue(queue_orchestration._is_load_balance_enabled(True))
+        self.assertTrue(queue_orchestration._is_load_balance_enabled(1))
+        self.assertTrue(queue_orchestration._is_load_balance_enabled(" yes "))
+        self.assertFalse(queue_orchestration._is_load_balance_enabled(0))
+        self.assertFalse(queue_orchestration._is_load_balance_enabled("off"))
+
+    def test_prompt_requests_load_balance_detects_collector_flag(self):
+        prompt_index = types.SimpleNamespace(
+            nodes_for_class=lambda class_name: ["10"] if class_name == "DistributedCollector" else [],
+            inputs_by_node={"10": {"load_balance": "true"}},
+        )
+        self.assertTrue(queue_orchestration._prompt_requests_load_balance(prompt_index))
+
+    def test_ensure_distributed_state_initializes_attrs(self):
+        prompt_server = types.SimpleNamespace()
+        queue_orchestration.ensure_distributed_state(prompt_server)
+        self.assertTrue(hasattr(prompt_server, "distributed_pending_jobs"))
+        self.assertTrue(hasattr(prompt_server, "distributed_jobs_lock"))
+        self.assertIsInstance(prompt_server.distributed_jobs_lock, asyncio.Lock)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_tunnel_routes.py

@@ -0,0 +1,95 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+from tests.api.harness import (
+    bootstrap_test_package,
+    install_aiohttp_stub,
+    install_endpoint_policy_passthrough_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+def _load_tunnel_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "tunnel_routes.py"
+    package_name = "dist_tunnel_routes_testpkg"
+
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    install_aiohttp_stub(lambda payload, status=200: _FakeResponse(payload, status=status))
+
+    cloudflare_module = types.ModuleType(f"{package_name}.utils.cloudflare")
+
+    class _TunnelManager:
+        def get_status(self):
+            return {"status": "running"}
+
+        async def start_tunnel(self):
+            return {"status": "running", "public_url": "https://example.trycloudflare.com"}
+
+        async def stop_tunnel(self):
+            return {"status": "stopped"}
+
+    cloudflare_module.cloudflare_tunnel_manager = _TunnelManager()
+    sys.modules[f"{package_name}.utils.cloudflare"] = cloudflare_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"master": {"host": "https://master.example.com"}}
+    sys.modules[f"{package_name}.utils.config"] = config_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")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    install_request_guards_stub(package_name)
+    install_endpoint_policy_passthrough_stub(package_name)
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.tunnel_routes", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+tunnel_routes = _load_tunnel_routes_module()
+
+
+class TunnelRoutesTests(unittest.IsolatedAsyncioTestCase):
+    async def test_status_endpoint_returns_tunnel_and_master_host(self):
+        response = await tunnel_routes.tunnel_status_endpoint(request=None)
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload["status"], "success")
+        self.assertEqual(response.payload["tunnel"]["status"], "running")
+        self.assertEqual(response.payload["master_host"], "https://master.example.com")
+
+    async def test_start_and_stop_endpoints_return_success_payload(self):
+        start_response = await tunnel_routes.tunnel_start_endpoint(request=None)
+        stop_response = await tunnel_routes.tunnel_stop_endpoint(request=None)
+
+        self.assertEqual(start_response.status, 200)
+        self.assertEqual(start_response.payload["tunnel"]["status"], "running")
+        self.assertEqual(stop_response.status, 200)
+        self.assertEqual(stop_response.payload["tunnel"]["status"], "stopped")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_usdu_routes.py

@@ -0,0 +1,316 @@
+import asyncio
+import importlib.util
+import io
+import sys
+import types
+import unittest
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from PIL import Image
+
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, json_payload=None, post_payload=None, headers=None, query=None):
+        self._json_payload = json_payload
+        self._post_payload = post_payload or {}
+        self.headers = headers or {}
+        self.query = query or {}
+
+    async def json(self):
+        return self._json_payload
+
+    async def post(self):
+        return self._post_payload
+
+
+def _load_usdu_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py"
+    package_name = "dist_api_usdu_testpkg"
+
+    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True)
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+
+    def _require_bool_literal(value, field_name: str = "value"):
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, str):
+            normalized = value.strip().lower()
+            if normalized == "true":
+                return True
+            if normalized == "false":
+                return False
+        raise ValueError(f"Field '{field_name}' must be a boolean literal.")
+
+    schemas_module.require_bool_literal = _require_bool_literal
+    schemas_module.is_authorized_request = lambda _request, _config: True
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    install_request_guards_stub(package_name)
+
+    prompt_server_holder = {
+        "value": types.SimpleNamespace(
+            distributed_tile_jobs_lock=asyncio.Lock(),
+            distributed_pending_tile_jobs={},
+        )
+    }
+
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    install_server_stub()
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_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.load_config = lambda: {"settings": {}}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    usdu_management_module = types.ModuleType(f"{package_name}.utils.usdu_management")
+    usdu_management_module.MAX_PAYLOAD_SIZE = 1024
+    sys.modules[f"{package_name}.utils.usdu_management"] = usdu_management_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.MAX_PAYLOAD_SIZE = 1024
+    job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
+
+    async def _init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices=None):
+        prompt_server = prompt_server_holder["value"]
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                return
+            job_data = ImageJobState(multi_job_id=multi_job_id)
+            job_data.worker_status = {str(worker_id): 0 for worker_id in (enabled_workers or [])}
+            job_data.assigned_to_workers = {str(worker_id): [] for worker_id in (enabled_workers or [])}
+            indices = all_indices if all_indices is not None else list(range(int(batch_size or 0)))
+            for idx in indices:
+                await job_data.pending_images.put(int(idx))
+            prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
+
+    job_store_module.init_dynamic_job = _init_dynamic_job
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
+
+    class BaseJobState:
+        pass
+
+    @dataclass
+    class TileJobState(BaseJobState):
+        multi_job_id: str
+        mode: str = field(default="static", init=False)
+        queue: asyncio.Queue = field(default_factory=asyncio.Queue)
+        pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
+        completed_tasks: dict = field(default_factory=dict)
+        worker_status: dict = field(default_factory=dict)
+        assigned_to_workers: dict = field(default_factory=dict)
+        batch_size: int = 0
+        num_tiles_per_image: int = 0
+        batched_static: bool = False
+
+    @dataclass
+    class ImageJobState(BaseJobState):
+        multi_job_id: str
+        mode: str = field(default="dynamic", init=False)
+        queue: asyncio.Queue = field(default_factory=asyncio.Queue)
+        pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
+        completed_images: dict = field(default_factory=dict)
+        worker_status: dict = field(default_factory=dict)
+        assigned_to_workers: dict = field(default_factory=dict)
+        batch_size: int = 0
+        num_tiles_per_image: int = 0
+        batched_static: bool = False
+
+        @property
+        def pending_tasks(self):
+            return self.pending_images
+
+        @property
+        def completed_tasks(self):
+            return self.completed_images
+
+    job_models_module.BaseJobState = BaseJobState
+    job_models_module.TileJobState = TileJobState
+    job_models_module.ImageJobState = ImageJobState
+    sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
+
+    parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers")
+    parsers_module.parse_tiles_from_form = lambda _data: []
+    parsers_module._parse_tiles_from_form = lambda _data: []
+    sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.usdu_routes", 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)
+
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+
+    module.web = types.SimpleNamespace(
+        json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+
+    module._prompt_server_holder = prompt_server_holder
+    module._TileJobState = TileJobState
+    module._ImageJobState = ImageJobState
+    return module
+
+
+usdu_routes = _load_usdu_routes_module()
+
+
+class _UploadField:
+    def __init__(self, data):
+        self.file = io.BytesIO(data)
+
+
+def _tiny_png_bytes():
+    image = Image.new("RGB", (1, 1), (255, 0, 0))
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        usdu_routes._prompt_server_holder["value"] = types.SimpleNamespace(
+            distributed_tile_jobs_lock=asyncio.Lock(),
+            distributed_pending_tile_jobs={},
+        )
+
+    async def test_heartbeat_updates_worker_status(self):
+        prompt_server = usdu_routes._prompt_server_holder["value"]
+        job_data = usdu_routes._TileJobState("job-1")
+        prompt_server.distributed_pending_tile_jobs["job-1"] = job_data
+
+        request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-1"})
+        response = await usdu_routes.heartbeat_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("status"), "success")
+        self.assertIn("worker-a", job_data.worker_status)
+
+    async def test_heartbeat_missing_fields_returns_400(self):
+        request = _FakeRequest(json_payload={"worker_id": "worker-a"})
+        response = await usdu_routes.heartbeat_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("missing", response.payload.get("message", "").lower())
+
+    async def test_request_image_dynamic_assigns_next_image(self):
+        prompt_server = usdu_routes._prompt_server_holder["value"]
+        job_data = usdu_routes._ImageJobState("job-2")
+        await job_data.pending_images.put(7)
+        prompt_server.distributed_pending_tile_jobs["job-2"] = job_data
+
+        request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-2"})
+        response = await usdu_routes.request_image_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("kind"), "image")
+        self.assertEqual(response.payload.get("task_idx"), 7)
+        self.assertEqual(response.payload.get("estimated_remaining"), 0)
+        self.assertEqual(job_data.assigned_to_workers["worker-a"], [7])
+        self.assertIn("worker-a", job_data.worker_status)
+
+    async def test_request_image_static_assigns_tile_and_batched_flag(self):
+        prompt_server = usdu_routes._prompt_server_holder["value"]
+        job_data = usdu_routes._TileJobState("job-3")
+        job_data.batched_static = True
+        await job_data.pending_tasks.put(4)
+        prompt_server.distributed_pending_tile_jobs["job-3"] = job_data
+
+        request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-3"})
+        response = await usdu_routes.request_image_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("kind"), "tile")
+        self.assertEqual(response.payload.get("task_idx"), 4)
+        self.assertTrue(response.payload.get("batched_static"))
+        self.assertEqual(job_data.assigned_to_workers["worker-a"], [4])
+
+    async def test_submit_tiles_completion_signal_enqueues_last_marker(self):
+        prompt_server = usdu_routes._prompt_server_holder["value"]
+        job_data = usdu_routes._TileJobState("job-4")
+        prompt_server.distributed_pending_tile_jobs["job-4"] = job_data
+
+        request = _FakeRequest(
+            post_payload={
+                "multi_job_id": "job-4",
+                "worker_id": "worker-a",
+                "batch_size": "0",
+                "is_last": "true",
+            },
+            headers={"content-length": "128"},
+        )
+
+        response = await usdu_routes.submit_tiles_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        queued = await job_data.queue.get()
+        self.assertEqual(queued["worker_id"], "worker-a")
+        self.assertTrue(queued["is_last"])
+        self.assertEqual(queued["tiles"], [])
+
+    async def test_submit_image_enqueues_processed_image_payload(self):
+        prompt_server = usdu_routes._prompt_server_holder["value"]
+        job_data = usdu_routes._ImageJobState("job-5")
+        prompt_server.distributed_pending_tile_jobs["job-5"] = job_data
+
+        request = _FakeRequest(
+            post_payload={
+                "multi_job_id": "job-5",
+                "worker_id": "worker-a",
+                "image_idx": "2",
+                "full_image": _UploadField(_tiny_png_bytes()),
+                "is_last": "false",
+            },
+            headers={"content-length": "256"},
+        )
+
+        response = await usdu_routes.submit_image_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        queued = await job_data.queue.get()
+        self.assertEqual(queued["worker_id"], "worker-a")
+        self.assertEqual(queued["image_idx"], 2)
+        self.assertIn("image", queued)
+
+    async def test_job_status_endpoint_reports_ready(self):
+        prompt_server = usdu_routes._prompt_server_holder["value"]
+        prompt_server.distributed_pending_tile_jobs["job-6"] = usdu_routes._TileJobState("job-6")
+
+        request = _FakeRequest(query={"multi_job_id": "job-6"})
+        response = await usdu_routes.job_status_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertTrue(response.payload.get("ready"))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_worker_routes.py

@@ -0,0 +1,358 @@
+import importlib.util
+import os
+import sys
+import tempfile
+import types
+import unittest
+from collections import deque
+from pathlib import Path
+from unittest.mock import patch
+
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, payload=None, match_info=None, query=None, headers=None):
+        self._payload = payload
+        self.match_info = match_info or {}
+        self.query = query or {}
+        self.headers = headers or {}
+
+    async def json(self):
+        return self._payload
+
+
+class _FakeHTTPClientResponse:
+    def __init__(self, payload, status=200):
+        self._payload = payload
+        self.status = status
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, _exc_type, _exc, _tb):
+        return False
+
+    async def json(self):
+        return self._payload
+
+    async def text(self):
+        return str(self._payload)
+
+
+class _FakeHTTPClientSession:
+    def __init__(self, payload, status=200):
+        self._payload = payload
+        self._status = status
+        self.calls = []
+
+    def get(self, url, params=None, headers=None, timeout=None):
+        self.calls.append({"url": url, "params": params, "headers": headers, "timeout": timeout})
+        return _FakeHTTPClientResponse(self._payload, status=self._status)
+
+
+class _DummyWorkerManager:
+    def __init__(self):
+        self.processes = {}
+
+    def launch_worker(self, worker):
+        worker_id = str(worker["id"])
+        self.processes[worker_id] = {
+            "pid": 12345,
+            "log_file": f"/tmp/distributed_worker_{worker_id}.log",  # nosec B108 - deterministic fake test path
+            "process": None,
+        }
+        return 12345
+
+    def is_process_running(self, _pid):
+        return False
+
+    def save_processes(self):
+        return None
+
+    def stop_worker(self, _worker_id):
+        return True, "Stopped"
+
+    def get_managed_workers(self):
+        return []
+
+
+class _ImmediateLoop:
+    async def run_in_executor(self, _executor, func, *args):
+        return func(*args)
+
+
+def _load_worker_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py"
+    package_name = "dist_api_worker_testpkg"
+
+    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_workers=True)
+
+    workers_pkg = types.ModuleType(f"{package_name}.workers")
+    workers_pkg.__path__ = []
+    workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
+    sys.modules[f"{package_name}.workers"] = workers_pkg
+
+    detection_module = types.ModuleType(f"{package_name}.workers.detection")
+    detection_module.is_local_worker = lambda *_args, **_kwargs: True
+    detection_module.is_same_physical_host = lambda *_args, **_kwargs: True
+    detection_module.get_machine_id = lambda: "machine-id"
+    detection_module.is_docker_environment = lambda: False
+    detection_module.is_runpod_environment = lambda: False
+    detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan"
+    sys.modules[f"{package_name}.workers.detection"] = detection_module
+
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    install_server_stub()
+
+    created_torch_stub = False
+    if "torch" not in sys.modules:
+        created_torch_stub = True
+        torch_module = types.ModuleType("torch")
+        torch_module.cuda = types.SimpleNamespace(
+            is_available=lambda: False,
+            empty_cache=lambda: None,
+            ipc_collect=lambda: None,
+            current_device=lambda: 0,
+            device_count=lambda: 0,
+        )
+        sys.modules["torch"] = torch_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.load_config = lambda: {"workers": []}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    network_module.normalize_host = lambda value: value
+    network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost:{worker.get('port', 8188)}{endpoint}"
+
+    async def _probe_worker(*_args, **_kwargs):
+        return None
+
+    network_module.probe_worker = _probe_worker
+
+    async def _get_client_session():
+        raise RuntimeError("not used in these tests")
+
+    network_module.get_client_session = _get_client_session
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.CHUNK_SIZE = 8192
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+
+    async def _queue_prompt_payload(*_args, **_kwargs):
+        return "prompt-id"
+
+    async_helpers_module.queue_prompt_payload = _queue_prompt_payload
+
+    class _PromptValidationError(RuntimeError):
+        def __init__(self, message="invalid prompt", validation_error=None, node_errors=None):
+            super().__init__(message)
+            self.validation_error = validation_error if isinstance(validation_error, dict) else {}
+            self.node_errors = node_errors if isinstance(node_errors, dict) else {}
+
+    async_helpers_module.PromptValidationError = _PromptValidationError
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+
+    def _require_fields(data, *fields):
+        missing = []
+        for field in fields:
+            value = data.get(field) if isinstance(data, dict) else None
+            if value is None or (isinstance(value, str) and not value.strip()):
+                missing.append(field)
+        return missing
+
+    def _validate_worker_id(worker_id, config):
+        return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", []))
+
+    def _require_worker_id(worker_id, config, field_name="worker_id"):
+        _ = field_name
+        worker_id_str = str(worker_id).strip()
+        if any(str(worker.get("id")) == worker_id_str for worker in config.get("workers", [])):
+            return worker_id_str
+        raise ValueError(f"Worker {worker_id_str} not found")
+
+    schemas_module.require_fields = _require_fields
+    schemas_module.validate_worker_id = _validate_worker_id
+    schemas_module.require_worker_id = _require_worker_id
+    schemas_module.is_authorized_request = lambda _request, _config: True
+    schemas_module.distributed_auth_headers = lambda _config: {}
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    install_request_guards_stub(package_name)
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", 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)
+
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    cleanup_optional_module("torch", created_torch_stub)
+
+    return module
+
+
+worker_routes = _load_worker_routes_module()
+
+
+class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
+    async def test_launch_worker_valid_id_returns_200(self):
+        manager = _DummyWorkerManager()
+        config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
+        request = _FakeRequest({"worker_id": "worker-a"})
+
+        with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
+            worker_routes, "load_config", return_value=config
+        ), patch.object(
+            worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
+        ):
+            response = await worker_routes.launch_worker_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("status"), "success")
+        self.assertEqual(response.payload.get("pid"), 12345)
+
+    async def test_launch_worker_unknown_id_returns_404(self):
+        manager = _DummyWorkerManager()
+        config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
+        request = _FakeRequest({"worker_id": "missing-worker"})
+
+        with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
+            worker_routes, "load_config", return_value=config
+        ):
+            response = await worker_routes.launch_worker_endpoint(request)
+
+        self.assertEqual(response.status, 404)
+        self.assertIn("not found", response.payload.get("message", "").lower())
+
+    async def test_worker_log_returns_content_json(self):
+        manager = _DummyWorkerManager()
+        with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as handle:
+            handle.write("line-1\nline-2\nline-3\n")
+            log_path = handle.name
+
+        manager.processes["worker-a"] = {
+            "pid": 9999,
+            "log_file": log_path,
+            "process": None,
+        }
+
+        request = _FakeRequest(match_info={"worker_id": "worker-a"}, query={"lines": "2"})
+        try:
+            with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
+                worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
+            ):
+                response = await worker_routes.get_worker_log_endpoint(request)
+        finally:
+            if os.path.exists(log_path):
+                os.remove(log_path)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("status"), "success")
+        self.assertIn("content", response.payload)
+        self.assertIn("line-3", response.payload["content"])
+
+    async def test_local_log_reads_memory_buffer(self):
+        request = _FakeRequest(query={"lines": "2"})
+        fake_logs = deque(
+            [
+                {"m": "line-1\n"},
+                {"m": "line-2\n"},
+                {"m": "line-3\n"},
+            ],
+            maxlen=300,
+        )
+        app_module = types.ModuleType("app")
+        app_module.__path__ = []
+        logger_module = types.ModuleType("app.logger")
+        logger_module.get_logs = lambda: fake_logs
+        app_module.logger = logger_module
+
+        with patch.dict(sys.modules, {"app": app_module, "app.logger": logger_module}):
+            response = await worker_routes.get_local_log_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("status"), "success")
+        self.assertEqual(response.payload.get("source"), "memory")
+        self.assertEqual(response.payload.get("entries"), 2)
+        self.assertIn("line-3", response.payload.get("content", ""))
+
+    async def test_remote_worker_log_proxies_to_worker_local_log_endpoint(self):
+        config = {
+            "workers": [
+                {
+                    "id": "worker-remote",
+                    "name": "Remote Worker",
+                    "host": "worker.example.com",
+                    "port": 8188,
+                    "type": "remote",
+                }
+            ]
+        }
+        request = _FakeRequest(match_info={"worker_id": "worker-remote"}, query={"lines": "120"})
+        proxied_payload = {
+            "status": "success",
+            "content": "remote-log-content\n",
+            "entries": 1,
+            "source": "memory",
+            "truncated": False,
+            "lines_shown": 1,
+        }
+        fake_session = _FakeHTTPClientSession(proxied_payload)
+
+        async def _fake_get_client_session():
+            return fake_session
+
+        with patch.object(worker_routes, "load_config", return_value=config), patch.object(
+            worker_routes, "get_client_session", side_effect=_fake_get_client_session
+        ):
+            response = await worker_routes.get_remote_worker_log_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("content"), "remote-log-content\n")
+        self.assertEqual(len(fake_session.calls), 1)
+        self.assertEqual(fake_session.calls[0]["params"], {"lines": "120"})
+        self.assertTrue(fake_session.calls[0]["url"].endswith("/distributed/local_log"))
+
+    async def test_remote_worker_log_rejects_local_workers(self):
+        config = {"workers": [{"id": "worker-local", "name": "Local Worker", "port": 8188}]}
+        request = _FakeRequest(match_info={"worker_id": "worker-local"})
+
+        with patch.object(worker_routes, "load_config", return_value=config):
+            response = await worker_routes.get_remote_worker_log_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("local", response.payload.get("message", "").lower())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/__init__.py

@@ -0,0 +1 @@
+"""Unit tests."""

tests/unit/nodes/__init__.py

@@ -0,0 +1 @@
+"""Unit tests."""

tests/unit/nodes/test_batch_dividers.py

@@ -0,0 +1,102 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+
+
+def _load_utilities_module():
+    module_path = Path(__file__).resolve().parents[3] / "nodes" / "utilities.py"
+    package_name = "dist_divider_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
+
+    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
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.nodes.utilities",
+        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
+
+
+utils = _load_utilities_module()
+
+
+class ImageBatchDividerTests(unittest.TestCase):
+    def test_divides_images_into_contiguous_chunks(self):
+        divider = utils.ImageBatchDivider()
+        images = torch.arange(10, dtype=torch.float32).reshape(10, 1, 1, 1)
+
+        outputs = divider.divide_batch(images, 3)
+
+        self.assertEqual(outputs[0].shape[0], 4)
+        self.assertEqual(outputs[1].shape[0], 3)
+        self.assertEqual(outputs[2].shape[0], 3)
+        self.assertEqual(outputs[0][:, 0, 0, 0].tolist(), [0.0, 1.0, 2.0, 3.0])
+        self.assertEqual(outputs[1][:, 0, 0, 0].tolist(), [4.0, 5.0, 6.0])
+        self.assertEqual(outputs[2][:, 0, 0, 0].tolist(), [7.0, 8.0, 9.0])
+
+    def test_unused_image_outputs_are_empty(self):
+        divider = utils.ImageBatchDivider()
+        images = torch.arange(4, dtype=torch.float32).reshape(4, 1, 1, 1)
+
+        outputs = divider.divide_batch(images, 2)
+
+        self.assertEqual(len(outputs), 10)
+        for idx in range(2, 10):
+            self.assertEqual(outputs[idx].shape[0], 0)
+
+
+class AudioBatchDividerTests(unittest.TestCase):
+    def test_divides_audio_samples_into_contiguous_chunks(self):
+        divider = utils.AudioBatchDivider()
+        audio = {
+            "waveform": torch.arange(10, dtype=torch.float32).reshape(1, 1, 10),
+            "sample_rate": 24000,
+        }
+
+        outputs = divider.divide_audio(audio, 3)
+
+        self.assertEqual(outputs[0]["waveform"][0, 0].tolist(), [0.0, 1.0, 2.0, 3.0])
+        self.assertEqual(outputs[1]["waveform"][0, 0].tolist(), [4.0, 5.0, 6.0])
+        self.assertEqual(outputs[2]["waveform"][0, 0].tolist(), [7.0, 8.0, 9.0])
+
+    def test_unused_audio_outputs_are_empty(self):
+        divider = utils.AudioBatchDivider()
+        audio = {
+            "waveform": torch.arange(8, dtype=torch.float32).reshape(1, 1, 8),
+            "sample_rate": 24000,
+        }
+
+        outputs = divider.divide_audio(audio, 2)
+
+        self.assertEqual(len(outputs), 10)
+        for idx in range(2, 10):
+            self.assertEqual(outputs[idx]["waveform"].shape[-1], 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/nodes/test_branch.py

@@ -0,0 +1,80 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    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
+
+    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
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_branch_module():
+    package_name = "dist_branch_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
+    return _load_module(package_name, "nodes/branch.py", "nodes.branch")
+
+
+branch_module = _load_branch_module()
+
+
+class DistributedBranchTests(unittest.TestCase):
+    def test_node_declares_ten_output_slots(self):
+        self.assertEqual(len(branch_module.DistributedBranch.RETURN_TYPES), 10)
+        self.assertEqual(len(branch_module.DistributedBranch.RETURN_NAMES), 10)
+
+    def test_no_distribution_outputs_all_branches(self):
+        node = branch_module.DistributedBranch()
+        token = {"k": "v"}
+
+        outputs = node.branch(token, num_branches=2, assigned_branch=-1)
+
+        self.assertEqual(len(outputs), 10)
+        self.assertTrue(all(output is token for output in outputs))
+
+    def test_assigned_branch_still_passes_input(self):
+        node = branch_module.DistributedBranch()
+        token = "payload"
+
+        outputs = node.branch(token, num_branches=3, is_worker=True, worker_id="worker-a", assigned_branch=1)
+
+        self.assertEqual(outputs[0], "payload")
+        self.assertEqual(outputs[1], "payload")
+        self.assertEqual(outputs[9], "payload")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/nodes/test_branch_collector.py

@@ -0,0 +1,282 @@
+import asyncio
+import base64
+import importlib.util
+import io
+import json
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+
+def _bootstrap_package(package_name):
+    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
+
+    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.2
+    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()
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    def encode_tensor_png_data_url(image_batch, batch_index=0):
+        image = tensor_to_pil(image_batch, batch_index)
+        byte_io = io.BytesIO()
+        image.save(byte_io, format="PNG", compress_level=0)
+        encoded = base64.b64encode(byte_io.getvalue()).decode("utf-8")
+        return f"data:image/png;base64,{encoded}"
+
+    image_module.ensure_contiguous = ensure_contiguous
+    image_module.tensor_to_pil = tensor_to_pil
+    image_module.encode_tensor_png_data_url = encode_tensor_png_data_url
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def get_client_session():
+        raise RuntimeError("test should patch get_client_session")
+
+    network_module.get_client_session = get_client_session
+    sys.modules[f"{package_name}.utils.network"] = network_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
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_branch_collector_module():
+    package_name = "dist_join_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
+    _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
+    _load_module(package_name, "nodes/runtime_helpers.py", "nodes.runtime_helpers")
+    _load_module(package_name, "nodes/queue_wait.py", "nodes.queue_wait")
+    _load_module(package_name, "utils/worker_ids.py", "utils.worker_ids")
+    return _load_module(package_name, "nodes/branch_collector.py", "nodes.branch_collector")
+
+
+branch_collector_module = _load_branch_collector_module()
+
+
+def _context(**overrides):
+    defaults = {
+        "multi_job_id": "",
+        "is_worker": False,
+        "master_url": "",
+        "enabled_worker_ids": "[]",
+        "worker_id": "",
+        "assigned_branch": -1,
+        "delegate_only": False,
+    }
+    defaults.update(overrides)
+    return branch_collector_module.BranchRunContext(**defaults)
+
+
+class _FakeResponseCtx:
+    def __init__(self, recorder, payload, url):
+        self._recorder = recorder
+        self._payload = payload
+        self._url = url
+
+    async def __aenter__(self):
+        self._recorder.append({"url": self._url, "payload": self._payload})
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+    def raise_for_status(self):
+        return None
+
+
+class _FakeSession:
+    def __init__(self):
+        self.calls = []
+
+    def post(self, url, json=None, timeout=None):
+        return _FakeResponseCtx(self.calls, json, url)
+
+
+class _FakePromptServer:
+    def __init__(self):
+        self.distributed_pending_jobs = {}
+        self.distributed_jobs_lock = asyncio.Lock()
+
+
+class DistributedBranchCollectorTests(unittest.IsolatedAsyncioTestCase):
+    def setUp(self):
+        self.node = branch_collector_module.DistributedBranchCollector()
+
+    async def test_worker_mode_sends_result_with_assigned_branch_index(self):
+        fake_session = _FakeSession()
+
+        async def _get_session():
+            return fake_session
+
+        branch_collector_module.get_client_session = _get_session
+
+        tensor = torch.ones((1, 4, 4, 3), dtype=torch.float32)
+        outputs = await self.node.execute(
+            [None, tensor, None, None, None, None, None, None, None, None],
+            num_branches=2,
+            context=_context(
+                multi_job_id="join-job",
+                is_worker=True,
+                master_url="http://master.local:8188",
+                worker_id="worker-a",
+                assigned_branch=1,
+            ),
+        )
+
+        self.assertEqual(len(fake_session.calls), 1)
+        payload = fake_session.calls[0]["payload"]
+        self.assertEqual(payload["batch_idx"], 1)
+        self.assertTrue(payload["is_last"])
+        self.assertEqual(payload["worker_id"], "worker-a")
+        self.assertTrue(torch.equal(outputs[1], tensor))
+
+    async def test_master_mode_collects_worker_results_into_branch_slots(self):
+        fake_server = _FakePromptServer()
+        branch_collector_module._get_prompt_server_instance = lambda: fake_server
+
+        queue = asyncio.Queue()
+        worker_1 = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32)
+        worker_2 = torch.full((1, 4, 4, 3), 0.8, dtype=torch.float32)
+        await queue.put({"tensor": worker_1, "worker_id": "worker-a", "image_index": 1, "is_last": True})
+        await queue.put({"tensor": worker_2, "worker_id": "worker-b", "image_index": 2, "is_last": True})
+        fake_server.distributed_pending_jobs["join-job"] = queue
+
+        master_tensor = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
+        outputs = await self.node.execute(
+            [master_tensor, None, None, None, None, None, None, None, None, None],
+            num_branches=3,
+            context=_context(
+                multi_job_id="join-job",
+                enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
+                assigned_branch=0,
+            ),
+        )
+
+        self.assertTrue(torch.equal(outputs[0], master_tensor))
+        self.assertTrue(torch.equal(outputs[1], worker_1))
+        self.assertTrue(torch.equal(outputs[2], worker_2))
+
+    async def test_unassigned_slots_remain_none(self):
+        fake_server = _FakePromptServer()
+        branch_collector_module._get_prompt_server_instance = lambda: fake_server
+
+        queue = asyncio.Queue()
+        worker_1 = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32)
+        await queue.put({"tensor": worker_1, "worker_id": "worker-a", "image_index": 1, "is_last": True})
+        fake_server.distributed_pending_jobs["join-job"] = queue
+
+        master_tensor = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
+        outputs = await self.node.execute(
+            [master_tensor, None, None, None, None, None, None, None, None, None],
+            num_branches=2,
+            context=_context(
+                multi_job_id="join-job",
+                enabled_worker_ids=json.dumps(["worker-a"]),
+                assigned_branch=0,
+            ),
+        )
+
+        self.assertIsNone(outputs[2])
+        self.assertIsNone(outputs[9])
+
+    async def test_missing_worker_results_use_safe_fallback_for_expected_slots(self):
+        fake_server = _FakePromptServer()
+        branch_collector_module._get_prompt_server_instance = lambda: fake_server
+
+        queue = asyncio.Queue()
+        # No worker payloads enqueued; master should timeout and fill expected worker slots.
+        fake_server.distributed_pending_jobs["join-job"] = queue
+
+        master_tensor = torch.full((1, 4, 4, 3), 0.7, dtype=torch.float32)
+        outputs = await self.node.execute(
+            [master_tensor, None, None, None, None, None, None, None, None, None],
+            num_branches=3,
+            context=_context(
+                multi_job_id="join-job",
+                enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
+                assigned_branch=0,
+            ),
+        )
+
+        self.assertTrue(torch.equal(outputs[0], master_tensor))
+        self.assertIsInstance(outputs[1], torch.Tensor)
+        self.assertIsInstance(outputs[2], torch.Tensor)
+        self.assertEqual(float(outputs[1].abs().sum().item()), 0.0)
+        self.assertEqual(float(outputs[2].abs().sum().item()), 0.0)
+
+    async def test_assigned_slot_can_be_filled_from_single_mismatched_local_input(self):
+        fake_server = _FakePromptServer()
+        branch_collector_module._get_prompt_server_instance = lambda: fake_server
+
+        # No worker results expected for this check.
+        fake_server.distributed_pending_jobs["join-job"] = asyncio.Queue()
+
+        local_tensor = torch.full((1, 4, 4, 3), 0.33, dtype=torch.float32)
+        outputs = await self.node.execute(
+            # Local value arrives on slot 1 even though participant is assigned slot 0.
+            [None, local_tensor, None, None, None, None, None, None, None, None],
+            num_branches=3,
+            context=_context(
+                multi_job_id="join-job",
+                enabled_worker_ids=json.dumps([]),
+                assigned_branch=0,
+            ),
+        )
+
+        self.assertTrue(torch.equal(outputs[0], local_tensor))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/nodes/test_collector.py

@@ -0,0 +1,228 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def _bootstrap_package(package_name):
+    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
+
+    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.2
+    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()
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    def pil_to_tensor(image):
+        array = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(array).unsqueeze(0)
+
+    image_module.ensure_contiguous = ensure_contiguous
+    image_module.tensor_to_pil = tensor_to_pil
+    image_module.pil_to_tensor = pil_to_tensor
+    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.local:8188"
+    network_module.get_client_session = lambda: None
+
+    async def probe_worker(_url, timeout=2.0):
+        _ = timeout
+        return None
+
+    network_module.probe_worker = probe_worker
+    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 payload: payload
+    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
+
+    if "aiohttp" not in sys.modules:
+        try:
+            import aiohttp as _aiohttp  # noqa: F401
+        except Exception:
+            aiohttp_module = types.ModuleType("aiohttp")
+
+            class _ClientTimeout:
+                def __init__(self, total=None):
+                    self.total = total
+
+            aiohttp_module.ClientTimeout = _ClientTimeout
+            sys.modules["aiohttp"] = aiohttp_module
+
+    class _FakePromptServer:
+        def __init__(self):
+            self.distributed_pending_jobs = {}
+            self.distributed_jobs_lock = asyncio.Lock()
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=_FakePromptServer())
+    sys.modules["server"] = server_module
+
+    comfy_module = types.ModuleType("comfy")
+    model_mgmt = types.ModuleType("comfy.model_management")
+
+    class _InterruptProcessingException(Exception):
+        pass
+
+    model_mgmt.throw_exception_if_processing_interrupted = lambda: None
+    model_mgmt.InterruptProcessingException = _InterruptProcessingException
+    comfy_module.model_management = model_mgmt
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.model_management"] = model_mgmt
+
+    comfy_utils_module = types.ModuleType("comfy.utils")
+
+    class _ProgressBar:
+        def __init__(self, _total):
+            self.total = _total
+
+        def update(self, _step):
+            return None
+
+    comfy_utils_module.ProgressBar = _ProgressBar
+    sys.modules["comfy.utils"] = comfy_utils_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_collector_module():
+    package_name = "dist_collector_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
+    _load_module(package_name, "utils/worker_ids.py", "utils.worker_ids")
+    return _load_module(package_name, "nodes/collector.py", "nodes.collector")
+
+
+collector_module = _load_collector_module()
+
+
+class DistributedCollectorTests(unittest.TestCase):
+    def setUp(self):
+        self.node = collector_module.DistributedCollectorNode()
+
+    def test_store_worker_result_tracks_by_worker_and_index(self):
+        worker_images = {}
+        stored = self.node._store_worker_result(
+            worker_images,
+            {
+                "worker_id": "worker-a",
+                "image_index": 2,
+                "tensor": torch.full((1, 1, 1, 1), 0.25),
+            },
+        )
+
+        self.assertEqual(stored, 1)
+        self.assertIn("worker-a", worker_images)
+        self.assertIn(2, worker_images["worker-a"])
+
+    def test_combine_audio_honors_master_then_worker_order(self):
+        master_audio = {"waveform": torch.ones((1, 2, 1)), "sample_rate": 48000}
+        worker_audio = {
+            "worker-b": {"waveform": torch.full((1, 2, 1), 3.0), "sample_rate": 48000},
+            "worker-a": {"waveform": torch.full((1, 2, 2), 2.0), "sample_rate": 48000},
+        }
+        empty_audio = {"waveform": torch.zeros((1, 2, 1)), "sample_rate": 44100}
+
+        combined = self.node._combine_audio(
+            master_audio=master_audio,
+            worker_audio=worker_audio,
+            empty_audio=empty_audio,
+            worker_order=["worker-a", "worker-b"],
+        )
+
+        expected = torch.cat(
+            [
+                master_audio["waveform"],
+                worker_audio["worker-a"]["waveform"],
+                worker_audio["worker-b"]["waveform"],
+            ],
+            dim=-1,
+        )
+        self.assertEqual(combined["sample_rate"], 48000)
+        self.assertTrue(torch.equal(combined["waveform"], expected))
+
+    def test_reorder_and_combine_tensors_uses_enabled_worker_priority(self):
+        worker_images = {
+            "worker-b": {
+                1: torch.full((1, 1, 1, 1), 4.0),
+                0: torch.full((1, 1, 1, 1), 3.0),
+            },
+            "worker-a": {
+                0: torch.full((1, 1, 1, 1), 2.0),
+            },
+        }
+        master_images = torch.full((1, 1, 1, 1), 1.0)
+
+        combined = self.node._reorder_and_combine_tensors(
+            worker_images=worker_images,
+            worker_order=["worker-a", "worker-b"],
+            master_batch_size=1,
+            images_on_cpu=master_images,
+            delegate_mode=False,
+            fallback_images=None,
+        )
+
+        self.assertEqual(combined.shape, (4, 1, 1, 1))
+        self.assertEqual(combined[0, 0, 0, 0].item(), 1.0)
+        self.assertEqual(combined[1, 0, 0, 0].item(), 2.0)
+        self.assertEqual(combined[2, 0, 0, 0].item(), 3.0)
+        self.assertEqual(combined[3, 0, 0, 0].item(), 4.0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/nodes/test_distributed_upscale.py

@@ -0,0 +1,248 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+
+
+def _bootstrap_package(package_name):
+    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
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
+    modes_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
+
+    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
+
+    worker_ids_module = types.ModuleType(f"{package_name}.utils.worker_ids")
+    worker_ids_module.coerce_enabled_worker_ids = lambda value: (
+        [str(v) for v in value]
+        if isinstance(value, list)
+        else []
+    )
+    sys.modules[f"{package_name}.utils.worker_ids"] = worker_ids_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.run_async_in_server_loop = lambda coro, timeout=None: (_ for _ in ()).throw(
+        RuntimeError(f"run_async_in_server_loop should not be used in these tests (timeout={timeout})")
+    )
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: None
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    tile_ops_module = types.ModuleType(f"{package_name}.upscale.tile_ops")
+
+    class _TileOpsMixin:
+        def round_to_multiple(self, value):
+            return int(value)
+
+        def calculate_tiles(self, width, height, tile_width, tile_height, force_uniform_tiles):
+            _ = force_uniform_tiles
+            tiles = []
+            for y in range(0, int(height), max(1, int(tile_height))):
+                for x in range(0, int(width), max(1, int(tile_width))):
+                    tiles.append((x, y))
+            return tiles or [(0, 0)]
+
+    tile_ops_module.TileOpsMixin = _TileOpsMixin
+    sys.modules[f"{package_name}.upscale.tile_ops"] = tile_ops_module
+
+    result_collector_module = types.ModuleType(f"{package_name}.upscale.result_collector")
+    result_collector_module.ResultCollectorMixin = type("ResultCollectorMixin", (), {})
+    sys.modules[f"{package_name}.upscale.result_collector"] = result_collector_module
+
+    worker_comms_module = types.ModuleType(f"{package_name}.upscale.worker_comms")
+    worker_comms_module.WorkerCommsMixin = type("WorkerCommsMixin", (), {})
+    sys.modules[f"{package_name}.upscale.worker_comms"] = worker_comms_module
+
+    job_state_module = types.ModuleType(f"{package_name}.upscale.job_state")
+    job_state_module.JobStateMixin = type("JobStateMixin", (), {})
+    sys.modules[f"{package_name}.upscale.job_state"] = job_state_module
+
+    single_gpu_module = types.ModuleType(f"{package_name}.upscale.modes.single_gpu")
+
+    class _SingleGpuModeMixin:
+        def process_single_gpu(self, *_args, **_kwargs):
+            return ("single_gpu_result",)
+
+    single_gpu_module.SingleGpuModeMixin = _SingleGpuModeMixin
+    sys.modules[f"{package_name}.upscale.modes.single_gpu"] = single_gpu_module
+
+    static_mode_module = types.ModuleType(f"{package_name}.upscale.modes.static")
+
+    class _StaticModeMixin:
+        def _process_worker_static_sync(self, *_args, **_kwargs):
+            return ("worker_static_result",)
+
+        def _process_master_static_sync(self, *_args, **_kwargs):
+            return ("master_static_result",)
+
+    static_mode_module.StaticModeMixin = _StaticModeMixin
+    sys.modules[f"{package_name}.upscale.modes.static"] = static_mode_module
+
+    dynamic_mode_module = types.ModuleType(f"{package_name}.upscale.modes.dynamic")
+
+    class _DynamicModeMixin:
+        def process_worker_dynamic(self, *_args, **_kwargs):
+            return ("worker_dynamic_result",)
+
+        def process_master_dynamic(self, *_args, **_kwargs):
+            return ("master_dynamic_result",)
+
+    dynamic_mode_module.DynamicModeMixin = _DynamicModeMixin
+    sys.modules[f"{package_name}.upscale.modes.dynamic"] = dynamic_mode_module
+
+    comfy_module = types.ModuleType("comfy")
+    samplers_module = types.ModuleType("comfy.samplers")
+
+    class _KSampler:
+        SAMPLERS = ("euler",)
+        SCHEDULERS = ("normal",)
+
+    samplers_module.KSampler = _KSampler
+    comfy_module.samplers = samplers_module
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.samplers"] = samplers_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_distributed_upscale_module():
+    package_name = "dist_distributed_upscale_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
+    _load_module(package_name, "upscale/mode_contexts.py", "upscale.mode_contexts")
+    _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args")
+    return _load_module(package_name, "nodes/distributed_upscale.py", "nodes.distributed_upscale")
+
+
+upscale_module = _load_distributed_upscale_module()
+
+
+class DistributedUpscaleTests(unittest.TestCase):
+    def setUp(self):
+        self.node = upscale_module.UltimateSDUpscaleDistributed()
+        self.base_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
+        self.common_kwargs = {
+            "model": object(),
+            "positive": object(),
+            "negative": object(),
+            "vae": object(),
+            "seed": 1,
+            "steps": 10,
+            "cfg": 7.5,
+            "sampler_name": "euler",
+            "scheduler": "normal",
+            "denoise": 0.5,
+            "tile_width": 32,
+            "tile_height": 32,
+            "padding": 16,
+            "mask_blur": 4,
+            "force_uniform_tiles": True,
+            "tiled_decode": False,
+        }
+
+    def test_parse_enabled_worker_ids_supports_json_and_lists(self):
+        self.assertEqual(
+            upscale_module._parse_enabled_worker_ids(["worker-a", 2]),
+            ["worker-a", "2"],
+        )
+        self.assertEqual(
+            upscale_module._parse_enabled_worker_ids('["worker-a","worker-b"]'),
+            ["worker-a", "worker-b"],
+        )
+        self.assertEqual(upscale_module._parse_enabled_worker_ids("invalid-json"), [])
+        self.assertEqual(upscale_module._parse_enabled_worker_ids(None), [])
+
+    def test_determine_processing_mode(self):
+        self.assertEqual(self.node._determine_processing_mode(batch_size=1, num_workers=0, dynamic_threshold=8), "single_gpu")
+        self.assertEqual(self.node._determine_processing_mode(batch_size=9, num_workers=2, dynamic_threshold=8), "static")
+
+    def test_run_raises_for_non_4n_plus_1_master_batches(self):
+        bad_batch = torch.zeros((2, 64, 64, 3), dtype=torch.float32)
+        with self.assertRaises(ValueError):
+            self.node.run(
+                upscaled_image=bad_batch,
+                multi_job_id="job-1",
+                is_worker=False,
+                **self.common_kwargs,
+            )
+
+    def test_run_dispatches_single_gpu_when_not_distributed(self):
+        self.node.process_single_gpu = lambda *_args, **_kwargs: ("single",)
+        result = self.node.run(
+            upscaled_image=self.base_image,
+            multi_job_id="",
+            is_worker=False,
+            **self.common_kwargs,
+        )
+        self.assertEqual(result, ("single",))
+
+    def test_run_dispatches_worker_or_master_paths(self):
+        self.node.process_worker = lambda *_args, **_kwargs: ("worker",)
+        self.node.process_master = lambda *_args, **_kwargs: ("master",)
+
+        worker_result = self.node.run(
+            upscaled_image=self.base_image,
+            multi_job_id="job-2",
+            is_worker=True,
+            master_url="http://master.local:8188",
+            enabled_worker_ids='["worker-a"]',
+            worker_id="worker-a",
+            tile_indices="",
+            dynamic_threshold=8,
+            **self.common_kwargs,
+        )
+        master_result = self.node.run(
+            upscaled_image=self.base_image,
+            multi_job_id="job-2",
+            is_worker=False,
+            enabled_worker_ids='["worker-a"]',
+            worker_id="",
+            tile_indices="",
+            dynamic_threshold=8,
+            **self.common_kwargs,
+        )
+
+        self.assertEqual(worker_result, ("worker",))
+        self.assertEqual(master_result, ("master",))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/nodes/test_distributed_value.py

@@ -0,0 +1,230 @@
+import json
+import unittest
+
+
+class DistributedValueTests(unittest.TestCase):
+    """Unit tests for the DistributedValue node's distribute() method."""
+
+    def _make_node(self):
+        # Import inline to avoid plugin-level imports
+        import importlib.util
+        import sys
+        import types
+        from pathlib import Path
+        from unittest.mock import MagicMock
+
+        module_path = Path(__file__).resolve().parents[3] / "nodes" / "utilities.py"
+        pkg_name = "dv_test_pkg"
+
+        for mod_name in list(sys.modules):
+            if mod_name == pkg_name or mod_name.startswith(f"{pkg_name}."):
+                del sys.modules[mod_name]
+
+        # Mock torch if not available
+        if "torch" not in sys.modules:
+            sys.modules["torch"] = MagicMock()
+
+        root_pkg = types.ModuleType(pkg_name)
+        root_pkg.__path__ = []
+        sys.modules[pkg_name] = root_pkg
+
+        nodes_pkg = types.ModuleType(f"{pkg_name}.nodes")
+        nodes_pkg.__path__ = []
+        sys.modules[f"{pkg_name}.nodes"] = nodes_pkg
+
+        hidden_inputs_mod = types.ModuleType(f"{pkg_name}.nodes.hidden_inputs")
+        hidden_inputs_mod.build_worker_identity_hidden_inputs = lambda: {
+            "is_worker": ("BOOLEAN", {"default": False}),
+            "enabled_worker_ids": ("STRING", {"default": "[]"}),
+            "worker_id": ("STRING", {"default": ""}),
+        }
+        sys.modules[f"{pkg_name}.nodes.hidden_inputs"] = hidden_inputs_mod
+
+        utils_pkg = types.ModuleType(f"{pkg_name}.utils")
+        utils_pkg.__path__ = []
+        sys.modules[f"{pkg_name}.utils"] = utils_pkg
+
+        logging_mod = types.ModuleType(f"{pkg_name}.utils.logging")
+        logging_mod.debug_log = lambda *_a, **_k: None
+        logging_mod.log = lambda *_a, **_k: None
+        sys.modules[f"{pkg_name}.utils.logging"] = logging_mod
+
+        worker_ids_path = Path(__file__).resolve().parents[3] / "utils" / "worker_ids.py"
+        worker_ids_spec = importlib.util.spec_from_file_location(
+            f"{pkg_name}.utils.worker_ids",
+            worker_ids_path,
+        )
+        worker_ids_mod = importlib.util.module_from_spec(worker_ids_spec)
+        assert worker_ids_spec is not None and worker_ids_spec.loader is not None
+        worker_ids_spec.loader.exec_module(worker_ids_mod)
+        sys.modules[f"{pkg_name}.utils.worker_ids"] = worker_ids_mod
+
+        spec = importlib.util.spec_from_file_location(
+            f"{pkg_name}.nodes.utilities", module_path
+        )
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+        return mod.DistributedValue()
+
+    def setUp(self):
+        self.node = self._make_node()
+
+    def test_master_returns_default(self):
+        result = self.node.distribute(
+            default_value="model_a",
+            worker_values="{}",
+            is_worker=False,
+            worker_id="",
+        )
+        self.assertEqual(result, ("model_a",))
+
+    def test_master_coerces_default_int(self):
+        values = json.dumps({"_type": "INT"})
+        result = self.node.distribute(
+            default_value="42",
+            worker_values=values,
+            is_worker=False,
+            worker_id="",
+        )
+        self.assertEqual(result, (42,))
+        self.assertIsInstance(result[0], int)
+
+    def test_master_coerces_default_float(self):
+        values = json.dumps({"_type": "FLOAT"})
+        result = self.node.distribute(
+            default_value="2.5",
+            worker_values=values,
+            is_worker=False,
+            worker_id="",
+        )
+        self.assertEqual(result, (2.5,))
+        self.assertIsInstance(result[0], float)
+
+    def test_worker_returns_specific_value(self):
+        values = json.dumps({"1": "model_x", "2": "model_y"})
+        result = self.node.distribute(
+            default_value="default",
+            worker_values=values,
+            is_worker=True,
+            worker_id="worker_0",
+        )
+        self.assertEqual(result, ("model_x",))
+
+    def test_worker_second_index(self):
+        values = json.dumps({"1": "model_x", "2": "model_y"})
+        result = self.node.distribute(
+            default_value="default",
+            worker_values=values,
+            is_worker=True,
+            worker_id="worker_1",
+        )
+        self.assertEqual(result, ("model_y",))
+
+    def test_worker_falls_back_to_default_when_key_missing(self):
+        values = json.dumps({"_type": "INT", "1": "3"})
+        result = self.node.distribute(
+            default_value="9",
+            worker_values=values,
+            is_worker=True,
+            worker_id="worker_5",
+        )
+        self.assertEqual(result, (9,))
+        self.assertIsInstance(result[0], int)
+
+    def test_worker_falls_back_to_default_on_empty_value(self):
+        values = json.dumps({"1": ""})
+        result = self.node.distribute(
+            default_value="fallback",
+            worker_values=values,
+            is_worker=True,
+            worker_id="worker_0",
+        )
+        self.assertEqual(result, ("fallback",))
+
+    def test_worker_falls_back_on_invalid_json(self):
+        result = self.node.distribute(
+            default_value="safe",
+            worker_values="not-json",
+            is_worker=True,
+            worker_id="worker_0",
+        )
+        self.assertEqual(result, ("safe",))
+
+    def test_worker_falls_back_on_invalid_worker_id(self):
+        values = json.dumps({"1": "model_x"})
+        result = self.node.distribute(
+            default_value="safe",
+            worker_values=values,
+            is_worker=True,
+            worker_id="bad_id",
+        )
+        self.assertEqual(result, ("safe",))
+
+    def test_worker_id_as_direct_integer(self):
+        values = json.dumps({"1": "model_x"})
+        result = self.node.distribute(
+            default_value="default",
+            worker_values=values,
+            is_worker=True,
+            worker_id="0",
+        )
+        self.assertEqual(result, ("model_x",))
+
+    def test_type_int_coerces_value(self):
+        values = json.dumps({"_type": "INT", "1": "42"})
+        result = self.node.distribute(
+            default_value="0",
+            worker_values=values,
+            is_worker=True,
+            worker_id="worker_0",
+        )
+        self.assertEqual(result, (42,))
+        self.assertIsInstance(result[0], int)
+
+    def test_type_float_coerces_value(self):
+        values = json.dumps({"_type": "FLOAT", "1": "3.14"})
+        result = self.node.distribute(
+            default_value="0",
+            worker_values=values,
+            is_worker=True,
+            worker_id="worker_0",
+        )
+        self.assertAlmostEqual(result[0], 3.14)
+        self.assertIsInstance(result[0], float)
+
+    def test_type_combo_stays_string(self):
+        values = json.dumps({"_type": "COMBO", "1": "model_v2"})
+        result = self.node.distribute(
+            default_value="model_v1",
+            worker_values=values,
+            is_worker=True,
+            worker_id="worker_0",
+        )
+        self.assertEqual(result, ("model_v2",))
+        self.assertIsInstance(result[0], str)
+
+    def test_type_string_default_stays_string(self):
+        values = json.dumps({"1": "hello"})
+        result = self.node.distribute(
+            default_value="default",
+            worker_values=values,
+            is_worker=True,
+            worker_id="worker_0",
+        )
+        self.assertEqual(result, ("hello",))
+        self.assertIsInstance(result[0], str)
+
+    def test_int_coerce_handles_float_string(self):
+        """INT coercion of '3.7' should truncate to 3."""
+        values = json.dumps({"_type": "INT", "1": "3.7"})
+        result = self.node.distribute(
+            default_value="0",
+            worker_values=values,
+            is_worker=True,
+            worker_id="worker_0",
+        )
+        self.assertEqual(result, (3,))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/__init__.py

@@ -0,0 +1 @@
+"""Unit tests."""

tests/unit/system/test_async_helpers.py

@@ -0,0 +1,104 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+class _FakePromptQueue:
+    def __init__(self):
+        self.items = []
+
+    def put(self, item):
+        self.items.append(item)
+
+
+class _FakePromptServer:
+    def __init__(self):
+        self.number = 0
+        self.prompt_queue = _FakePromptQueue()
+
+    def trigger_on_prompt(self, payload):
+        return payload
+
+
+def _bootstrap_package(package_name):
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.get_server_loop = asyncio.get_event_loop
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    fake_prompt_server = _FakePromptServer()
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=fake_prompt_server)
+    sys.modules["server"] = server_module
+
+    execution_module = types.ModuleType("execution")
+
+    async def _validate_prompt(_prompt_id, _prompt, _partial_targets):
+        return (True, None, ["1"], {})
+
+    execution_module.validate_prompt = _validate_prompt
+    execution_module.SENSITIVE_EXTRA_DATA_KEYS = ()
+    sys.modules["execution"] = execution_module
+
+    return fake_prompt_server
+
+
+def _load_async_helpers_module():
+    package_name = "dist_async_helpers_testpkg"
+    fake_prompt_server = _bootstrap_package(package_name)
+    module_path = Path(__file__).resolve().parents[3] / "utils/async_helpers.py"
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.utils.async_helpers",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module, fake_prompt_server
+
+
+class AsyncHelpersQueuePromptPayloadTests(unittest.TestCase):
+    def test_queue_prompt_payload_includes_create_time_metadata(self):
+        async_helpers_module, fake_prompt_server = _load_async_helpers_module()
+
+        prompt_id = asyncio.run(
+            async_helpers_module.queue_prompt_payload(
+                {"1": {"class_type": "KSampler", "inputs": {}}},
+                workflow_meta={"id": "workflow-1"},
+                client_id="client-1",
+            )
+        )
+
+        self.assertIsInstance(prompt_id, str)
+        self.assertTrue(fake_prompt_server.prompt_queue.items)
+
+        queued_item = fake_prompt_server.prompt_queue.items[-1]
+        self.assertEqual(len(queued_item), 6)
+        self.assertEqual(queued_item[1], prompt_id)
+
+        extra_data = queued_item[3]
+        self.assertEqual(extra_data["client_id"], "client-1")
+        self.assertEqual(extra_data["extra_pnginfo"]["workflow"]["id"], "workflow-1")
+        self.assertIn("create_time", extra_data)
+        self.assertIsInstance(extra_data["create_time"], int)
+        self.assertGreater(extra_data["create_time"], 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_config.py

@@ -0,0 +1,293 @@
+import importlib.util
+import json
+import os
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+def _load_config_module():
+    module_path = Path(__file__).resolve().parents[3] / "utils" / "config.py"
+    package_name = "dist_cfg_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
+
+    constants_module = types.ModuleType(f"{package_name}.constants")
+    constants_module.GPU_CONFIG_FILE = "gpu_config.json"
+    constants_module.HEARTBEAT_TIMEOUT = 30
+    constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8
+    constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4
+    constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 2
+    constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 120
+    sys.modules[f"{package_name}.constants"] = constants_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.config", 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
+
+
+config = _load_config_module()
+
+
+# ---------------------------------------------------------------------------
+# _merge_with_defaults
+# ---------------------------------------------------------------------------
+
+class MergeWithDefaultsTests(unittest.TestCase):
+    def test_non_dict_input_returns_defaults(self):
+        result = config._merge_with_defaults("not a dict", {"key": "default"})
+        self.assertEqual(result, {"key": "default"})
+
+    def test_fills_missing_keys_with_defaults(self):
+        result = config._merge_with_defaults({}, {"a": 1, "b": 2})
+        self.assertEqual(result, {"a": 1, "b": 2})
+
+    def test_loaded_value_overrides_default(self):
+        result = config._merge_with_defaults({"a": 99}, {"a": 1, "b": 2})
+        self.assertEqual(result["a"], 99)
+        self.assertEqual(result["b"], 2)
+
+    def test_nested_dict_merges_recursively(self):
+        defaults = {"settings": {"debug": False, "count": 5}}
+        loaded = {"settings": {"debug": True}}
+        result = config._merge_with_defaults(loaded, defaults)
+        self.assertTrue(result["settings"]["debug"])
+        self.assertEqual(result["settings"]["count"], 5)
+
+    def test_preserves_unknown_keys_for_forward_compatibility(self):
+        result = config._merge_with_defaults({"extra_key": "extra"}, {"a": 1})
+        self.assertEqual(result["extra_key"], "extra")
+
+    def test_none_loaded_value_overrides_default(self):
+        """Explicitly set None in config should override non-None default."""
+        result = config._merge_with_defaults({"a": None}, {"a": "default"})
+        self.assertIsNone(result["a"])
+
+    def test_non_dict_nested_loaded_value_replaces_dict_default(self):
+        """If loaded has a scalar where default has a dict, use the scalar."""
+        defaults = {"settings": {"debug": False}}
+        loaded = {"settings": "flat_string"}
+        result = config._merge_with_defaults(loaded, defaults)
+        self.assertEqual(result["settings"], "flat_string")
+
+
+# ---------------------------------------------------------------------------
+# load_config
+# ---------------------------------------------------------------------------
+
+class LoadConfigTests(unittest.TestCase):
+    def setUp(self):
+        config.invalidate_config_cache()
+
+    def tearDown(self):
+        config.invalidate_config_cache()
+
+    def test_returns_defaults_when_file_missing(self):
+        with patch.object(config, "CONFIG_FILE", "/nonexistent/path/config.json"):
+            cfg = config.load_config()
+        defaults = config.get_default_config()
+        self.assertEqual(cfg["settings"]["debug"], defaults["settings"]["debug"])
+        self.assertIn("workers", cfg)
+
+    def test_loads_valid_json_file(self):
+        data = {
+            "workers": [{"id": "w1"}],
+            "master": {"host": "test.host"},
+            "settings": {},
+            "tunnel": {},
+        }
+        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
+            json.dump(data, f)
+            tmp_path = f.name
+        try:
+            with patch.object(config, "CONFIG_FILE", tmp_path):
+                cfg = config.load_config()
+            self.assertEqual(cfg["master"]["host"], "test.host")
+            self.assertEqual(len(cfg["workers"]), 1)
+        finally:
+            os.unlink(tmp_path)
+
+    def test_merges_loaded_file_with_defaults(self):
+        """Loaded file with partial settings should be filled in from defaults."""
+        data = {"master": {"host": "h"}, "workers": [], "settings": {"debug": True}, "tunnel": {}}
+        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
+            json.dump(data, f)
+            tmp_path = f.name
+        try:
+            with patch.object(config, "CONFIG_FILE", tmp_path):
+                cfg = config.load_config()
+            # debug was set to True
+            self.assertTrue(cfg["settings"]["debug"])
+            # auto_launch_workers is a default key and should be present
+            self.assertIn("auto_launch_workers", cfg["settings"])
+        finally:
+            os.unlink(tmp_path)
+
+    def test_falls_back_to_defaults_on_invalid_json(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
+            f.write("{invalid json{{")
+            tmp_path = f.name
+        try:
+            with patch.object(config, "CONFIG_FILE", tmp_path):
+                cfg = config.load_config()
+            self.assertIn("settings", cfg)
+            self.assertIn("workers", cfg)
+        finally:
+            os.unlink(tmp_path)
+
+    def test_second_call_returns_cached_object(self):
+        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
+            json.dump(config.get_default_config(), f)
+            tmp_path = f.name
+        try:
+            with patch.object(config, "CONFIG_FILE", tmp_path):
+                cfg1 = config.load_config()
+                cfg2 = config.load_config()
+            self.assertIs(cfg1, cfg2)
+        finally:
+            os.unlink(tmp_path)
+
+    def test_invalidate_cache_forces_reload(self):
+        data = config.get_default_config()
+        data["master"]["host"] = "first"
+        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
+            json.dump(data, f)
+            tmp_path = f.name
+        try:
+            with patch.object(config, "CONFIG_FILE", tmp_path):
+                cfg1 = config.load_config()
+                config.invalidate_config_cache()
+                data["master"]["host"] = "second"
+                with open(tmp_path, "w", encoding="utf-8") as fh:
+                    json.dump(data, fh)
+                cfg2 = config.load_config()
+            self.assertEqual(cfg1["master"]["host"], "first")
+            self.assertEqual(cfg2["master"]["host"], "second")
+        finally:
+            os.unlink(tmp_path)
+
+
+# ---------------------------------------------------------------------------
+# save_config
+# ---------------------------------------------------------------------------
+
+class SaveConfigTests(unittest.TestCase):
+    def setUp(self):
+        config.invalidate_config_cache()
+
+    def tearDown(self):
+        config.invalidate_config_cache()
+
+    def test_saves_and_reloads_correctly(self):
+        data = config.get_default_config()
+        data["master"]["host"] = "saved.host"
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_path = os.path.join(tmpdir, "config.json")
+            with patch.object(config, "CONFIG_FILE", tmp_path):
+                result = config.save_config(data)
+                self.assertTrue(result)
+                loaded = config.load_config()
+        self.assertEqual(loaded["master"]["host"], "saved.host")
+
+    def test_returns_false_when_path_unwritable(self):
+        with patch.object(config, "CONFIG_FILE", "/nonexistent_dir/config.json"):
+            result = config.save_config({})
+        self.assertFalse(result)
+
+    def test_save_invalidates_cache(self):
+        """After saving, the cache should be cleared so next load re-reads."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_path = os.path.join(tmpdir, "config.json")
+            with patch.object(config, "CONFIG_FILE", tmp_path):
+                data = config.get_default_config()
+                config.save_config(data)
+                # Cache is now None; load_config should re-read
+                self.assertIsNone(config._config_state().cache)
+
+    def test_written_file_is_valid_json(self):
+        data = config.get_default_config()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_path = os.path.join(tmpdir, "config.json")
+            with patch.object(config, "CONFIG_FILE", tmp_path):
+                config.save_config(data)
+            with open(tmp_path, encoding="utf-8") as fh:
+                parsed = json.load(fh)
+        self.assertEqual(parsed["master"], data["master"])
+
+
+# ---------------------------------------------------------------------------
+# get_worker_timeout_seconds
+# ---------------------------------------------------------------------------
+
+class GetWorkerTimeoutSecondsTests(unittest.TestCase):
+    def test_returns_configured_value(self):
+        cfg = config.get_default_config()
+        cfg["settings"]["worker_timeout_seconds"] = 120
+        with patch.object(config, "load_config", return_value=cfg):
+            self.assertEqual(config.get_worker_timeout_seconds(), 120)
+
+    def test_clamps_zero_to_one(self):
+        cfg = config.get_default_config()
+        cfg["settings"]["worker_timeout_seconds"] = 0
+        with patch.object(config, "load_config", return_value=cfg):
+            self.assertEqual(config.get_worker_timeout_seconds(), 1)
+
+    def test_clamps_negative_to_one(self):
+        cfg = config.get_default_config()
+        cfg["settings"]["worker_timeout_seconds"] = -10
+        with patch.object(config, "load_config", return_value=cfg):
+            self.assertEqual(config.get_worker_timeout_seconds(), 1)
+
+    def test_falls_back_to_provided_default_when_key_missing(self):
+        cfg = config.get_default_config()
+        # worker_timeout_seconds is not present in default config
+        cfg["settings"].pop("worker_timeout_seconds", None)
+        with patch.object(config, "load_config", return_value=cfg):
+            result = config.get_worker_timeout_seconds(default=45)
+        self.assertEqual(result, 45)
+
+    def test_fallback_also_clamped_to_one(self):
+        cfg = config.get_default_config()
+        cfg["settings"].pop("worker_timeout_seconds", None)
+        with patch.object(config, "load_config", return_value=cfg):
+            result = config.get_worker_timeout_seconds(default=0)
+        self.assertEqual(result, 1)
+
+
+# ---------------------------------------------------------------------------
+# is_master_delegate_only
+# ---------------------------------------------------------------------------
+
+class IsMasterDelegateOnlyTests(unittest.TestCase):
+    def test_returns_false_by_default(self):
+        cfg = config.get_default_config()
+        with patch.object(config, "load_config", return_value=cfg):
+            self.assertFalse(config.is_master_delegate_only())
+
+    def test_returns_true_when_enabled(self):
+        cfg = config.get_default_config()
+        cfg["settings"]["master_delegate_only"] = True
+        with patch.object(config, "load_config", return_value=cfg):
+            self.assertTrue(config.is_master_delegate_only())
+
+    def test_returns_false_on_exception(self):
+        def _raise():
+            raise RuntimeError("config exploded")
+
+        with patch.object(config, "load_config", side_effect=RuntimeError("boom")):
+            self.assertFalse(config.is_master_delegate_only())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_detection.py

@@ -0,0 +1,197 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+def _load_detection_module():
+    module_path = Path(__file__).resolve().parents[3] / "workers" / "detection.py"
+    package_name = "dist_det_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
+
+    workers_pkg = types.ModuleType(f"{package_name}.workers")
+    workers_pkg.__path__ = []
+    sys.modules[f"{package_name}.workers"] = workers_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_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.normalize_host = lambda value: value
+    network_module.build_worker_url = lambda worker, endpoint="": f"http://{worker.get('host', 'localhost')}{endpoint}"
+
+    async def _fake_session():
+        raise RuntimeError("network calls not used in these tests")
+
+    network_module.get_client_session = _fake_session
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+
+        class _ClientTimeout:
+            def __init__(self, total=None):
+                pass
+
+        aiohttp_module.ClientTimeout = _ClientTimeout
+        sys.modules["aiohttp"] = aiohttp_module
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.workers.detection",
+        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)
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+
+    return module
+
+
+detection = _load_detection_module()
+
+
+# ---------------------------------------------------------------------------
+# is_docker_environment
+# ---------------------------------------------------------------------------
+
+class IsDockerEnvironmentTests(unittest.TestCase):
+    def test_true_when_dockerenv_file_exists(self):
+        with patch.object(detection.os.path, "exists", return_value=True), \
+             patch.dict(detection.os.environ, {}, clear=True), \
+             patch.object(detection.platform, "node", return_value="my-laptop"):
+            self.assertTrue(detection.is_docker_environment())
+
+    def test_true_when_docker_container_env_var_is_set(self):
+        with patch.object(detection.os.path, "exists", return_value=False), \
+             patch.dict(detection.os.environ, {"DOCKER_CONTAINER": "1"}, clear=True), \
+             patch.object(detection.platform, "node", return_value="my-laptop"):
+            self.assertTrue(detection.is_docker_environment())
+
+    def test_true_when_platform_node_contains_docker(self):
+        with patch.object(detection.os.path, "exists", return_value=False), \
+             patch.dict(detection.os.environ, {}, clear=True), \
+             patch.object(detection.platform, "node", return_value="my-docker-host"):
+            self.assertTrue(detection.is_docker_environment())
+
+    def test_false_when_none_of_the_signals_are_present(self):
+        with patch.object(detection.os.path, "exists", return_value=False), \
+             patch.dict(detection.os.environ, {}, clear=True), \
+             patch.object(detection.platform, "node", return_value="my-laptop"):
+            self.assertFalse(detection.is_docker_environment())
+
+    def test_docker_node_name_is_case_insensitive(self):
+        with patch.object(detection.os.path, "exists", return_value=False), \
+             patch.dict(detection.os.environ, {}, clear=True), \
+             patch.object(detection.platform, "node", return_value="My-Docker-Box"):
+            self.assertTrue(detection.is_docker_environment())
+
+    def test_docker_env_var_empty_string_is_falsy(self):
+        """An empty DOCKER_CONTAINER env var should NOT trigger docker detection."""
+        with patch.object(detection.os.path, "exists", return_value=False), \
+             patch.dict(detection.os.environ, {"DOCKER_CONTAINER": ""}, clear=True), \
+             patch.object(detection.platform, "node", return_value="my-laptop"):
+            self.assertFalse(detection.is_docker_environment())
+
+
+# ---------------------------------------------------------------------------
+# is_runpod_environment
+# ---------------------------------------------------------------------------
+
+class IsRunpodEnvironmentTests(unittest.TestCase):
+    def test_true_when_runpod_pod_id_is_set(self):
+        with patch.dict(detection.os.environ, {"RUNPOD_POD_ID": "pod-abc"}, clear=True):
+            self.assertTrue(detection.is_runpod_environment())
+
+    def test_true_when_runpod_api_key_is_set(self):
+        with patch.dict(detection.os.environ, {"RUNPOD_API_KEY": "key-xyz"}, clear=True):
+            self.assertTrue(detection.is_runpod_environment())
+
+    def test_true_when_both_vars_are_set(self):
+        with patch.dict(
+            detection.os.environ,
+            {"RUNPOD_POD_ID": "pod-abc", "RUNPOD_API_KEY": "key-xyz"},
+            clear=True,
+        ):
+            self.assertTrue(detection.is_runpod_environment())
+
+    def test_false_when_neither_var_is_set(self):
+        with patch.dict(detection.os.environ, {}, clear=True):
+            self.assertFalse(detection.is_runpod_environment())
+
+    def test_true_when_pod_id_is_empty_string(self):
+        """is not None check means even empty string counts as detected."""
+        with patch.dict(detection.os.environ, {"RUNPOD_POD_ID": ""}, clear=True):
+            self.assertTrue(detection.is_runpod_environment())
+
+
+# ---------------------------------------------------------------------------
+# is_local_worker (synchronous paths only)
+# ---------------------------------------------------------------------------
+
+class IsLocalWorkerTests(unittest.IsolatedAsyncioTestCase):
+    async def test_true_for_localhost_host(self):
+        result = await detection.is_local_worker({"host": "localhost", "port": 8188})
+        self.assertTrue(result)
+
+    async def test_true_for_127_0_0_1(self):
+        result = await detection.is_local_worker({"host": "127.0.0.1", "port": 8188})
+        self.assertTrue(result)
+
+    async def test_true_for_0_0_0_0(self):
+        result = await detection.is_local_worker({"host": "0.0.0.0", "port": 8188})  # nosec B104 - explicit wildcard host test case
+        self.assertTrue(result)
+
+    async def test_true_when_type_is_local(self):
+        result = await detection.is_local_worker({"type": "local", "host": "remote.example.com"})
+        self.assertTrue(result)
+
+    async def test_false_for_remote_host(self):
+        result = await detection.is_local_worker({"host": "remote.example.com", "port": 8188})
+        self.assertFalse(result)
+
+    async def test_true_when_no_host_key(self):
+        """Missing host defaults to 'localhost'."""
+        result = await detection.is_local_worker({"port": 8188})
+        self.assertTrue(result)
+
+
+# ---------------------------------------------------------------------------
+# get_machine_id
+# ---------------------------------------------------------------------------
+
+class GetMachineIdTests(unittest.TestCase):
+    def test_returns_a_string(self):
+        result = detection.get_machine_id()
+        self.assertIsInstance(result, str)
+
+    def test_returns_non_empty_string(self):
+        result = detection.get_machine_id()
+        self.assertTrue(len(result) > 0)
+
+    def test_stable_across_calls(self):
+        r1 = detection.get_machine_id()
+        r2 = detection.get_machine_id()
+        self.assertEqual(r1, r2)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_dispatch.py

@@ -0,0 +1,298 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+def _load_dispatch_module():
+    module_path = Path(__file__).resolve().parents[3] / "api" / "orchestration" / "dispatch.py"
+
+    package_name = "dist_dispatch_testpkg"
+    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
+
+    orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
+    orch_pkg.__path__ = []
+    sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    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_worker_url = lambda *_args, **_kwargs: "http://example.invalid"
+
+    async def _probe_worker(*_args, **_kwargs):
+        return None
+
+    network_module.probe_worker = _probe_worker
+
+    async def _fake_session():
+        raise RuntimeError("get_client_session should be mocked in these tests")
+
+    network_module.get_client_session = _fake_session
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+
+        class _ClientTimeout:
+            def __init__(self, total=None):
+                self.total = total
+
+        class _ClientConnectorError(Exception):
+            pass
+
+        class _WSMsgType:
+            TEXT = "TEXT"
+            ERROR = "ERROR"
+            CLOSED = "CLOSED"
+
+        class _TCPConnector:
+            def __init__(self, *args, **kwargs):
+                self.args = args
+                self.kwargs = kwargs
+
+        class _ClientSession:
+            def __init__(self, *args, **kwargs):
+                self.closed = False
+
+            async def close(self):
+                self.closed = True
+
+        aiohttp_module.ClientTimeout = _ClientTimeout
+        aiohttp_module.ClientConnectorError = _ClientConnectorError
+        aiohttp_module.WSMsgType = _WSMsgType
+        aiohttp_module.TCPConnector = _TCPConnector
+        aiohttp_module.ClientSession = _ClientSession
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: {"payload": payload, "status": status}
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.api.orchestration.dispatch",
+        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)
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+    return module
+
+
+dispatch = _load_dispatch_module()
+
+
+class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
+    async def test_select_active_workers_filters_offline(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+            {"id": "w3", "name": "Worker 3"},
+        ]
+
+        async def fake_probe(worker):
+            return worker["id"] != "w2"
+
+        with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
+            active_workers, delegate_master = await dispatch.select_active_workers(
+                workers,
+                use_websocket=False,
+                delegate_master=False,
+                probe_concurrency=3,
+            )
+
+        self.assertEqual([w["id"] for w in active_workers], ["w1", "w3"])
+        self.assertFalse(delegate_master)
+
+    async def test_select_active_workers_disables_delegate_when_all_offline(self):
+        workers = [{"id": "w1", "name": "Worker 1"}]
+
+        async def fake_probe(_worker):
+            return False
+
+        with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
+            active_workers, delegate_master = await dispatch.select_active_workers(
+                workers,
+                use_websocket=False,
+                delegate_master=True,
+                probe_concurrency=1,
+            )
+
+        self.assertEqual(active_workers, [])
+        self.assertFalse(delegate_master)
+
+    async def test_select_active_workers_uses_websocket_probe_when_enabled(self):
+        workers = [{"id": "w1", "name": "Worker 1"}, {"id": "w2", "name": "Worker 2"}]
+
+        async def fake_http_probe(_worker):
+            return False
+
+        async def fake_ws_probe(_worker):
+            return True
+
+        with patch.object(dispatch, "worker_is_active", side_effect=fake_http_probe) as http_probe, patch.object(
+            dispatch,
+            "worker_ws_is_active",
+            side_effect=fake_ws_probe,
+        ) as ws_probe:
+            active_workers, _ = await dispatch.select_active_workers(
+                workers,
+                use_websocket=True,
+                delegate_master=False,
+                probe_concurrency=4,
+            )
+
+        self.assertEqual([w["id"] for w in active_workers], ["w1", "w2"])
+        self.assertEqual(ws_probe.call_count, 2)
+        self.assertEqual(http_probe.call_count, 0)
+
+    async def test_probe_concurrency_is_bounded(self):
+        workers = [{"id": f"w{i}", "name": f"Worker {i}"} for i in range(6)]
+        state = {"in_flight": 0, "max_in_flight": 0}
+
+        async def fake_probe(_worker):
+            state["in_flight"] += 1
+            state["max_in_flight"] = max(state["max_in_flight"], state["in_flight"])
+            await asyncio.sleep(0.01)
+            state["in_flight"] -= 1
+            return True
+
+        with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
+            active_workers, _ = await dispatch.select_active_workers(
+                workers,
+                use_websocket=False,
+                delegate_master=False,
+                probe_concurrency=2,
+            )
+
+        self.assertEqual(len(active_workers), len(workers))
+        self.assertLessEqual(state["max_in_flight"], 2)
+        self.assertGreaterEqual(state["max_in_flight"], 2)
+
+    async def test_select_least_busy_worker_round_robins_idle_workers(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+            {"id": "w3", "name": "Worker 3"},
+        ]
+        queue_map = {"w1": 0, "w2": 0, "w3": 2}
+
+        async def fake_probe(worker_url, timeout=3.0):
+            worker_id = worker_url.rsplit("/", 1)[-1]
+            return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
+
+        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
+            selected1 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
+            selected2 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
+            selected3 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
+
+        self.assertEqual(selected1["id"], "w1")
+        self.assertEqual(selected2["id"], "w2")
+        self.assertEqual(selected3["id"], "w1")
+
+    async def test_select_least_busy_worker_chooses_smallest_queue_when_all_busy(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+            {"id": "w3", "name": "Worker 3"},
+        ]
+        queue_map = {"w1": 5, "w2": 2, "w3": 4}
+
+        async def fake_probe(worker_url, timeout=3.0):
+            worker_id = worker_url.rsplit("/", 1)[-1]
+            return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
+
+        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,
+        ):
+            selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=2)
+
+        self.assertEqual(selected["id"], "w2")
+
+    async def test_select_least_busy_worker_returns_none_when_all_probes_fail(self):
+        workers = [{"id": "w1", "name": "Worker 1"}]
+
+        async def fake_probe(_worker_url, timeout=3.0):
+            return None
+
+        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,
+        ):
+            selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=1)
+
+        self.assertIsNone(selected)
+
+    async def test_rank_workers_by_load_returns_sorted_by_queue_depth(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+            {"id": "w3", "name": "Worker 3"},
+        ]
+        queue_map = {"w1": 4, "w2": 0, "w3": 2}
+
+        async def fake_probe(worker_url, timeout=3.0):
+            worker_id = worker_url.rsplit("/", 1)[-1]
+            return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
+
+        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,
+        ):
+            ranked = await dispatch.rank_workers_by_load(workers, probe_concurrency=3)
+
+        self.assertEqual([worker["id"] for worker in ranked], ["w2", "w3", "w1"])
+
+    async def test_rank_workers_by_load_handles_unreachable_workers(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+            {"id": "w3", "name": "Worker 3"},
+        ]
+
+        async def fake_probe(worker_url, timeout=3.0):
+            worker_id = worker_url.rsplit("/", 1)[-1]
+            if worker_id == "w2":
+                return None
+            queue_map = {"w1": 1, "w3": 0}
+            return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
+
+        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,
+        ):
+            ranked = await dispatch.rank_workers_by_load(workers, probe_concurrency=3)
+
+        self.assertEqual([worker["id"] for worker in ranked], ["w3", "w1", "w2"])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_distributed.py

@@ -0,0 +1,71 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _load_distributed_module():
+    module_path = Path(__file__).resolve().parents[3] / "distributed.py"
+    package_name = "dist_distributed_entry_testpkg"
+    calls = {
+        "build_node_mappings": 0,
+        "initialize_runtime": 0,
+        "initialize_runtime_args": [],
+    }
+
+    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
+
+    bootstrap_pkg = types.ModuleType(f"{package_name}.bootstrap")
+    bootstrap_pkg.__path__ = []
+    sys.modules[f"{package_name}.bootstrap"] = bootstrap_pkg
+
+    entrypoint_module = types.ModuleType(f"{package_name}.bootstrap.entrypoint")
+
+    def _build_node_mappings():
+        calls["build_node_mappings"] += 1
+        return ({"MockNode": object}, {"MockNode": "Mock Node"})
+
+    def _initialize_runtime(prompt_server=None):
+        calls["initialize_runtime"] += 1
+        calls["initialize_runtime_args"].append(prompt_server)
+
+    entrypoint_module.build_node_mappings = _build_node_mappings
+    entrypoint_module.initialize_runtime = _initialize_runtime
+    sys.modules[f"{package_name}.bootstrap.entrypoint"] = entrypoint_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.distributed", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module, calls
+
+
+class DistributedEntryTests(unittest.TestCase):
+    def test_import_builds_mappings_without_runtime_bootstrap(self):
+        module, calls = _load_distributed_module()
+
+        self.assertIn("MockNode", module.NODE_CLASS_MAPPINGS)
+        self.assertEqual(module.NODE_DISPLAY_NAME_MAPPINGS["MockNode"], "Mock Node")
+        self.assertEqual(calls["build_node_mappings"], 1)
+        self.assertEqual(calls["initialize_runtime"], 0)
+
+    def test_initialize_runtime_delegates_to_bootstrap_entrypoint(self):
+        module, calls = _load_distributed_module()
+        sentinel = object()
+
+        module.initialize_runtime(sentinel)
+
+        self.assertEqual(calls["initialize_runtime"], 1)
+        self.assertEqual(calls["initialize_runtime_args"], [sentinel])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_exceptions.py

@@ -0,0 +1,40 @@
+import importlib.util
+import unittest
+from pathlib import Path
+
+
+def _load_exceptions_module():
+    module_path = Path(__file__).resolve().parents[3] / "utils" / "exceptions.py"
+    spec = importlib.util.spec_from_file_location("dist_test_exceptions", 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
+
+
+exceptions = _load_exceptions_module()
+
+
+class ExceptionsModuleTests(unittest.TestCase):
+    def test_worker_error_tracks_worker_id(self):
+        err = exceptions.WorkerError("worker failed", worker_id="w-1")
+        self.assertIsInstance(err, exceptions.DistributedError)
+        self.assertEqual(str(err), "worker failed")
+        self.assertEqual(err.worker_id, "w-1")
+
+    def test_process_error_tracks_pid_and_worker_id(self):
+        err = exceptions.ProcessError("process failed", pid=1001, worker_id="worker-x")
+        self.assertIsInstance(err, exceptions.DistributedError)
+        self.assertEqual(err.pid, 1001)
+        self.assertEqual(err.worker_id, "worker-x")
+
+    def test_specialized_errors_inherit_expected_base_types(self):
+        self.assertTrue(issubclass(exceptions.WorkerTimeoutError, exceptions.WorkerError))
+        self.assertTrue(issubclass(exceptions.WorkerNotAvailableError, exceptions.WorkerError))
+        self.assertTrue(issubclass(exceptions.JobQueueError, exceptions.DistributedError))
+        self.assertTrue(issubclass(exceptions.TileCollectionError, exceptions.DistributedError))
+        self.assertTrue(issubclass(exceptions.TunnelError, exceptions.DistributedError))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_lifecycle.py

@@ -0,0 +1,228 @@
+import importlib.util
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    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
+
+    workers_pkg = types.ModuleType(f"{package_name}.workers")
+    workers_pkg.__path__ = []
+    sys.modules[f"{package_name}.workers"] = workers_pkg
+
+    process_pkg = types.ModuleType(f"{package_name}.workers.process")
+    process_pkg.__path__ = []
+    sys.modules[f"{package_name}.workers.process"] = process_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_lifecycle_module():
+    package_name = "dist_lifecycle_testpkg"
+    _bootstrap_package(package_name)
+
+    state = {
+        "config": {"settings": {"stop_workers_on_master_exit": True}, "managed_processes": {}},
+        "saved_configs": [],
+        "launch_calls": [],
+        "terminated": [],
+        "alive_pids": set(),
+    }
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: state["config"]
+    config_module.save_config = lambda cfg: state["saved_configs"].append(dict(cfg))
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.PROCESS_TERMINATION_TIMEOUT = 5.0
+    constants_module.PROCESS_WAIT_TIMEOUT = 1.0
+    constants_module.WORKER_CHECK_INTERVAL = 0.01
+    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
+
+    process_module = types.ModuleType(f"{package_name}.utils.process")
+
+    class _FakeProcess:
+        def __init__(self, pid=1234, poll_result=None):
+            self.pid = pid
+            self._poll_result = poll_result
+            self.returncode = poll_result
+
+        def poll(self):
+            return self._poll_result
+
+    def _launch_process_with_timeout(command, timeout_seconds=10.0, **kwargs):
+        state["launch_calls"].append((list(command), float(timeout_seconds), dict(kwargs)))
+        return _FakeProcess(pid=4321, poll_result=None)
+
+    def _terminate_process(process, timeout=5):
+        state["terminated"].append((process.pid, timeout))
+        process._poll_result = 0
+        process.returncode = 0
+
+    process_module.launch_process_with_timeout = _launch_process_with_timeout
+    process_module.terminate_process = _terminate_process
+    process_module.is_process_alive = lambda pid: int(pid) in state["alive_pids"]
+    process_module.get_python_executable = lambda: "/usr/bin/python3"
+    process_module._FakeProcess = _FakeProcess
+    sys.modules[f"{package_name}.utils.process"] = process_module
+
+    lifecycle_module = _load_module(package_name, "workers/process/lifecycle.py", "workers.process.lifecycle")
+    return lifecycle_module, state, process_module
+
+
+lifecycle_module, test_state, process_stub_module = _load_lifecycle_module()
+
+
+class _FakeManager:
+    def __init__(self, comfy_root):
+        self._comfy_root = comfy_root
+        self.processes = {}
+        self.save_calls = 0
+
+    def find_comfy_root(self):
+        return self._comfy_root
+
+    def build_launch_command(self, worker_config, _comfy_root):
+        return ["python", "main.py", "--port", str(worker_config["port"])]
+
+    def save_processes(self):
+        self.save_calls += 1
+
+
+class ProcessLifecycleTests(unittest.TestCase):
+    def setUp(self):
+        test_state["config"] = {"settings": {"stop_workers_on_master_exit": True}, "managed_processes": {}}
+        test_state["saved_configs"].clear()
+        test_state["launch_calls"].clear()
+        test_state["terminated"].clear()
+        test_state["alive_pids"].clear()
+
+    def test_launch_worker_uses_monitor_wrapper_when_enabled(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+
+            pid = lifecycle.launch_worker(
+                {
+                    "id": "1",
+                    "name": "Worker A",
+                    "port": 8189,
+                    "cuda_device": 2,
+                }
+            )
+
+            self.assertEqual(pid, 4321)
+            self.assertEqual(manager.save_calls, 1)
+            self.assertIn("1", manager.processes)
+            self.assertTrue(manager.processes["1"]["is_monitor"])
+            self.assertIn("logs/workers", manager.processes["1"]["log_file"])
+            self.assertEqual(len(test_state["launch_calls"]), 1)
+            launched_command, timeout_seconds, launch_kwargs = test_state["launch_calls"][0]
+            self.assertEqual(timeout_seconds, 1.0)
+            self.assertEqual(launched_command[0], "/usr/bin/python3")
+            self.assertTrue(launched_command[1].endswith("workers/worker_monitor.py"))
+            self.assertEqual(launch_kwargs["env"]["CUDA_VISIBLE_DEVICES"], "2")
+
+    def test_launch_worker_skips_monitor_when_disabled(self):
+        test_state["config"] = {"settings": {"stop_workers_on_master_exit": False}, "managed_processes": {}}
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+
+            lifecycle.launch_worker(
+                {
+                    "id": "2",
+                    "name": "Worker-B",
+                    "port": 8190,
+                    "cuda_device": 0,
+                }
+            )
+
+            launched_command, _, _ = test_state["launch_calls"][0]
+            self.assertEqual(launched_command[0], "python")
+            self.assertFalse(manager.processes["2"]["is_monitor"])
+
+    def test_stop_worker_uses_fallback_terminate_when_tree_kill_fails(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+            running_process = process_stub_module._FakeProcess(pid=88, poll_result=None)
+            manager.processes["7"] = {
+                "pid": 88,
+                "process": running_process,
+                "started_at": 0.0,
+                "config": {},
+            }
+            lifecycle._kill_process_tree = lambda _pid: False
+
+            ok, message = lifecycle.stop_worker("7")
+
+            self.assertTrue(ok)
+            self.assertIn("fallback", message.lower())
+            self.assertEqual(test_state["terminated"], [(88, 5.0)])
+            self.assertNotIn("7", manager.processes)
+
+    def test_stop_worker_without_process_uses_tree_kill(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+            manager.processes["9"] = {
+                "pid": 99,
+                "process": None,
+                "started_at": 0.0,
+                "config": {},
+            }
+            lifecycle._kill_process_tree = lambda _pid: True
+
+            ok, message = lifecycle.stop_worker("9")
+
+            self.assertTrue(ok)
+            self.assertEqual(message, "Worker stopped")
+            self.assertNotIn("9", manager.processes)
+
+    def test_check_worker_process_falls_back_to_pid_probe(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+            lifecycle._is_process_running = lambda pid: int(pid) == 42
+
+            running, from_subprocess = lifecycle._check_worker_process(
+                "any",
+                {"pid": 42, "process": None},
+            )
+            self.assertTrue(running)
+            self.assertFalse(from_subprocess)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_logging.py

@@ -0,0 +1,32 @@
+import io
+import unittest
+from contextlib import redirect_stdout
+from unittest.mock import patch
+
+from utils import logging as logging_utils
+
+
+class LoggingUtilsTests(unittest.TestCase):
+    def test_is_debug_enabled_delegates_to_config(self):
+        with patch.object(logging_utils, "get_debug_enabled", return_value=True) as mocked:
+            self.assertTrue(logging_utils.is_debug_enabled())
+        mocked.assert_called_once_with(default=False)
+
+    def test_is_debug_enabled_handles_config_errors(self):
+        with patch.object(logging_utils, "get_debug_enabled", side_effect=RuntimeError("boom")):
+            self.assertFalse(logging_utils.is_debug_enabled())
+
+    def test_debug_log_emits_only_when_enabled(self):
+        output = io.StringIO()
+        with patch.object(logging_utils, "is_debug_enabled", return_value=True), redirect_stdout(output):
+            logging_utils.debug_log("hello")
+        self.assertIn("[Distributed] hello", output.getvalue())
+
+        output = io.StringIO()
+        with patch.object(logging_utils, "is_debug_enabled", return_value=False), redirect_stdout(output):
+            logging_utils.debug_log("hidden")
+        self.assertEqual(output.getvalue(), "")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_network.py

@@ -0,0 +1,107 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _load_network_module():
+    module_path = Path(__file__).resolve().parents[3] / "utils" / "network.py"
+
+    package_name = "dist_utils_testpkg"
+    package_module = types.ModuleType(package_name)
+    package_module.__path__ = []  # mark as package
+    sys.modules[package_name] = package_module
+
+    logging_module = types.ModuleType(f"{package_name}.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.logging"] = logging_module
+
+    exceptions_module = types.ModuleType(f"{package_name}.exceptions")
+    exceptions_module.DistributedError = type("DistributedError", (Exception,), {})
+    sys.modules[f"{package_name}.exceptions"] = exceptions_module
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(
+        instance=types.SimpleNamespace(address="127.0.0.1", port=8188, loop=None)
+    )
+    sys.modules["server"] = server_module
+
+    if "aiohttp" not in sys.modules:
+        aiohttp_module = types.ModuleType("aiohttp")
+
+        class _TCPConnector:
+            def __init__(self, *args, **kwargs):
+                self.args = args
+                self.kwargs = kwargs
+
+        class _ClientSession:
+            def __init__(self, *args, **kwargs):
+                self.closed = False
+
+            async def close(self):
+                self.closed = True
+
+        aiohttp_module.TCPConnector = _TCPConnector
+        aiohttp_module.ClientSession = _ClientSession
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: {"payload": payload, "status": status}
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.network", 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
+
+
+network = _load_network_module()
+
+
+class NetworkHelpersTests(unittest.TestCase):
+    def test_normalize_host_strips_protocol_and_path(self):
+        self.assertEqual(network.normalize_host("  https://example.com/a/b  "), "example.com")
+
+    def test_normalize_host_keeps_none(self):
+        self.assertIsNone(network.normalize_host(None))
+
+    def test_build_worker_url_defaults_to_server_address(self):
+        worker = {"id": "w1", "port": 8189}
+        self.assertEqual(network.build_worker_url(worker, "/prompt"), "http://127.0.0.1:8189/prompt")
+
+    def test_build_worker_url_cloud_defaults_to_https(self):
+        worker = {"id": "w2", "host": "foo.proxy.runpod.net", "port": 443}
+        self.assertEqual(network.build_worker_url(worker), "https://foo.proxy.runpod.net")
+
+    def test_build_worker_url_keeps_explicit_scheme(self):
+        worker = {"id": "w3", "host": "https://worker.example.com", "port": 1234}
+        self.assertEqual(network.build_worker_url(worker, "/prompt"), "https://worker.example.com/prompt")
+
+    def test_build_master_url_uses_https_for_cloud_host(self):
+        cfg = {"master": {"host": "demo.proxy.runpod.net"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "https://demo.proxy.runpod.net",
+        )
+
+    def test_build_master_url_keeps_explicit_scheme(self):
+        cfg = {"master": {"host": "https://master.example.com/"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "https://master.example.com",
+        )
+
+    def test_build_master_url_falls_back_to_server_address(self):
+        cfg = {"master": {"host": ""}}
+        prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)  # nosec B104 - wildcard bind normalization test
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "http://127.0.0.1:8190",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_prompt_transform.py

@@ -0,0 +1,705 @@
+import importlib.util
+import json
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _load_prompt_transform_module():
+    module_path = Path(__file__).resolve().parents[3] / "api" / "orchestration" / "prompt_transform.py"
+    package_name = "dist_pt_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
+
+    orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
+    orch_pkg.__path__ = []
+    sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    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
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.api.orchestration.prompt_transform",
+        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
+
+
+pt = _load_prompt_transform_module()
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _linear_prompt():
+    """1 → 2 → 3 → 4(DistributedCollector) → 5(SaveImage)"""
+    return {
+        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
+        "2": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["1", 1]}},
+        "3": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["2", 0]}},
+        "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
+        "5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}},
+    }
+
+
+def _collector_only_prompt():
+    """1(Checkpoint) → 2(DistributedCollector) [no downstream from 2]"""
+    return {
+        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
+        "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
+    }
+
+
+def _delegate_prompt():
+    """1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
+    return {
+        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
+        "2": {"class_type": "KSampler", "inputs": {"model": ["1", 0]}},
+        "3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
+        "4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
+    }
+
+
+def _branch_prompt():
+    """1 → 2(DistributedBranch) -> slot0:3->4, slot1:5->6."""
+    return {
+        "1": {"class_type": "KSampler", "inputs": {}},
+        "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+        "3": {"class_type": "Blur", "inputs": {"image": ["2", 0]}},
+        "4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
+        "5": {"class_type": "Sharpen", "inputs": {"image": ["2", 1]}},
+        "6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
+    }
+
+
+def _branch_collector_prompt():
+    """1 → 2(DistributedBranch) -> 3(branch0) and 4(branch1) -> 5(DistributedBranchCollector) -> 6(SaveImage)."""
+    return {
+        "1": {"class_type": "KSampler", "inputs": {}},
+        "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+        "3": {"class_type": "Blur", "inputs": {"image": ["2", 0]}},
+        "4": {"class_type": "Sharpen", "inputs": {"image": ["2", 1]}},
+        "5": {"class_type": "DistributedBranchCollector", "inputs": {"input": ["3", 0], "num_branches": 2}},
+        "6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
+    }
+
+
+def _apply(prompt, participant_id, enabled_worker_ids=None, delegate_master=False):
+    if enabled_worker_ids is None:
+        enabled_worker_ids = ["worker-a", "worker-b"]
+    prompt_copy = json.loads(json.dumps(prompt))
+    idx = pt.PromptIndex(prompt_copy)
+    job_id_map = pt.generate_job_id_map(idx, "run")
+    return pt.apply_participant_overrides(
+        prompt_copy,
+        participant_id=participant_id,
+        enabled_worker_ids=enabled_worker_ids,
+        job_id_map=job_id_map,
+        master_url="http://master.example.com",
+        delegate_master=delegate_master,
+        prompt_index=idx,
+    )
+
+
+# ---------------------------------------------------------------------------
+# PromptIndex
+# ---------------------------------------------------------------------------
+
+class PromptIndexTests(unittest.TestCase):
+    def test_nodes_by_class_groups_correctly(self):
+        prompt = {
+            "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
+            "2": {"class_type": "DistributedCollector", "inputs": {}},
+            "3": {"class_type": "DistributedCollector", "inputs": {}},
+        }
+        idx = pt.PromptIndex(prompt)
+        self.assertCountEqual(idx.nodes_for_class("DistributedCollector"), ["2", "3"])
+        self.assertEqual(idx.nodes_for_class("CheckpointLoaderSimple"), ["1"])
+
+    def test_nodes_for_class_unknown_returns_empty(self):
+        idx = pt.PromptIndex({"1": {"class_type": "KSampler", "inputs": {}}})
+        self.assertEqual(idx.nodes_for_class("Nonexistent"), [])
+
+    def test_nodes_without_class_type_are_indexed_under_none(self):
+        prompt = {"1": {"inputs": {}}}
+        idx = pt.PromptIndex(prompt)
+        # Should not raise; nodes_for_class with None key or missing class_type
+        self.assertEqual(idx.nodes_for_class("KSampler"), [])
+
+    def test_copy_prompt_is_a_deep_copy(self):
+        prompt = {"1": {"class_type": "KSampler", "inputs": {"seed": 42}}}
+        idx = pt.PromptIndex(prompt)
+        copy = idx.copy_prompt()
+        copy["1"]["inputs"]["seed"] = 999
+        self.assertEqual(prompt["1"]["inputs"]["seed"], 42)
+
+    def test_has_upstream_direct_connection(self):
+        """Node 4 reads directly from node 3 (KSampler)."""
+        idx = pt.PromptIndex(_linear_prompt())
+        self.assertTrue(idx.has_upstream("4", "KSampler"))
+
+    def test_has_upstream_transitive_connection(self):
+        """Node 4 → 3 → 2 → 1 (CheckpointLoaderSimple)."""
+        idx = pt.PromptIndex(_linear_prompt())
+        self.assertTrue(idx.has_upstream("4", "CheckpointLoaderSimple"))
+
+    def test_has_upstream_returns_false_when_no_path(self):
+        idx = pt.PromptIndex(_linear_prompt())
+        # CheckpointLoaderSimple has no upstream nodes
+        self.assertFalse(idx.has_upstream("1", "DistributedCollector"))
+
+    def test_has_upstream_result_is_cached(self):
+        idx = pt.PromptIndex(_linear_prompt())
+        r1 = idx.has_upstream("4", "KSampler")
+        r2 = idx.has_upstream("4", "KSampler")
+        self.assertEqual(r1, r2)
+        self.assertIn(("4", "KSampler"), idx._upstream_cache)
+
+    def test_has_upstream_does_not_infinite_loop_on_cycle(self):
+        """Cyclic references in inputs should not cause infinite recursion."""
+        prompt = {
+            "1": {"class_type": "A", "inputs": {"x": ["2", 0]}},
+            "2": {"class_type": "B", "inputs": {"x": ["1", 0]}},
+        }
+        idx = pt.PromptIndex(prompt)
+        # Should terminate without error
+        result = idx.has_upstream("1", "NonExistent")
+        self.assertFalse(result)
+
+
+# ---------------------------------------------------------------------------
+# find_nodes_by_class
+# ---------------------------------------------------------------------------
+
+class FindNodesByClassTests(unittest.TestCase):
+    def test_finds_matching_nodes(self):
+        prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedCollector", "inputs": {}},
+        }
+        result = pt.find_nodes_by_class(prompt, "KSampler")
+        self.assertEqual(result, ["1"])
+
+    def test_returns_empty_when_no_match(self):
+        prompt = {"1": {"class_type": "KSampler", "inputs": {}}}
+        self.assertEqual(pt.find_nodes_by_class(prompt, "DistributedCollector"), [])
+
+    def test_skips_non_dict_nodes(self):
+        prompt = {"1": "not a dict", "2": {"class_type": "KSampler", "inputs": {}}}
+        result = pt.find_nodes_by_class(prompt, "KSampler")
+        self.assertEqual(result, ["2"])
+
+
+# ---------------------------------------------------------------------------
+# prune_prompt_for_worker
+# ---------------------------------------------------------------------------
+
+class PrunePromptForWorkerTests(unittest.TestCase):
+    def test_no_distributed_nodes_returns_prompt_unchanged(self):
+        prompt = {
+            "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
+            "2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
+        }
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertCountEqual(result.keys(), ["1", "2"])
+
+    def test_keeps_collector_and_upstream(self):
+        prompt = _linear_prompt()
+        result = pt.prune_prompt_for_worker(prompt)
+        for node_id in ("1", "2", "3", "4"):
+            self.assertIn(node_id, result)
+
+    def test_removes_downstream_of_collector(self):
+        prompt = _linear_prompt()
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertNotIn("5", result)
+
+    def test_injects_preview_image_when_downstream_exists(self):
+        prompt = _linear_prompt()
+        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"], ["4", 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"]
+        self.assertEqual(len(preview_nodes), 0)
+
+    def test_unrelated_nodes_are_pruned(self):
+        prompt = {
+            "1": {"class_type": "DistributedCollector", "inputs": {}},
+            "2": {"class_type": "UnrelatedNode", "inputs": {}},  # no connection to 1
+        }
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertIn("1", result)
+        self.assertNotIn("2", result)
+
+    def test_result_is_a_copy_not_same_object(self):
+        prompt = _linear_prompt()
+        result = pt.prune_prompt_for_worker(prompt)
+        # Mutating the result should not affect the original
+        original_keys = set(prompt.keys())
+        result["NEW"] = {"class_type": "Test", "inputs": {}}
+        self.assertEqual(set(prompt.keys()), original_keys)
+
+    def test_upscale_node_is_treated_as_distributed(self):
+        prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}},
+            "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
+        }
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertIn("1", result)
+        self.assertIn("2", result)
+        self.assertNotIn("3", result)
+
+    def test_branch_anchor_keeps_downstream_for_later_branch_pruning(self):
+        prompt = _branch_prompt()
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertIn("2", result)
+        self.assertIn("3", result)
+        self.assertIn("4", result)
+        self.assertIn("5", result)
+        self.assertIn("6", result)
+
+    def test_branch_collector_anchor_keeps_collector_but_prunes_downstream(self):
+        prompt = _branch_collector_prompt()
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertIn("5", result)
+        self.assertNotIn("6", result)
+
+
+# ---------------------------------------------------------------------------
+# prepare_delegate_master_prompt
+# ---------------------------------------------------------------------------
+
+class PrepareDelegateMasterPromptTests(unittest.TestCase):
+    def test_keeps_collector_and_downstream(self):
+        prompt = _delegate_prompt()
+        result = pt.prepare_delegate_master_prompt(prompt, ["3"])
+        self.assertIn("3", result)
+        self.assertIn("4", result)
+        self.assertNotIn("1", result)
+        self.assertNotIn("2", result)
+
+    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", {})
+        # Original "images" pointed at node 2, which is pruned.
+        # It should now point at a newly injected placeholder node.
+        self.assertIn("images", collector_inputs)
+        source_id = str(collector_inputs["images"][0])
+        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_one_placeholder_per_collector(self):
+        """Two collectors → two placeholders."""
+        prompt = {
+            "1": {"class_type": "DistributedCollector", "inputs": {}},
+            "2": {"class_type": "DistributedCollector", "inputs": {}},
+            "3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
+        }
+        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_result_is_independent_copy(self):
+        prompt = _delegate_prompt()
+        result = pt.prepare_delegate_master_prompt(prompt, ["3"])
+        result["3"]["inputs"]["NEW"] = "injected"
+        # Original should be untouched
+        self.assertNotIn("NEW", prompt["3"].get("inputs", {}))
+
+
+# ---------------------------------------------------------------------------
+# generate_job_id_map
+# ---------------------------------------------------------------------------
+
+class GenerateJobIdMapTests(unittest.TestCase):
+    def test_maps_collector_nodes(self):
+        prompt = {
+            "1": {"class_type": "DistributedCollector", "inputs": {}},
+            "2": {"class_type": "KSampler", "inputs": {}},
+        }
+        idx = pt.PromptIndex(prompt)
+        job_map = pt.generate_job_id_map(idx, "prefix")
+        self.assertEqual(job_map["1"], "prefix_1")
+        self.assertNotIn("2", job_map)
+
+    def test_maps_upscale_nodes(self):
+        prompt = {
+            "5": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}},
+        }
+        idx = pt.PromptIndex(prompt)
+        job_map = pt.generate_job_id_map(idx, "run")
+        self.assertEqual(job_map["5"], "run_5")
+
+    def test_maps_branch_collector_nodes(self):
+        prompt = {
+            "10": {"class_type": "DistributedBranchCollector", "inputs": {}},
+        }
+        idx = pt.PromptIndex(prompt)
+        job_map = pt.generate_job_id_map(idx, "run")
+        self.assertEqual(job_map["10"], "run_10")
+
+    def test_maps_branch_nodes(self):
+        prompt = {
+            "9": {"class_type": "DistributedBranch", "inputs": {"num_branches": 2}},
+        }
+        idx = pt.PromptIndex(prompt)
+        job_map = pt.generate_job_id_map(idx, "run")
+        self.assertEqual(job_map["9"], "run_9")
+
+    def test_empty_prompt_returns_empty_map(self):
+        idx = pt.PromptIndex({})
+        self.assertEqual(pt.generate_job_id_map(idx, "prefix"), {})
+
+    def test_stable_ids_across_calls(self):
+        prompt = {"1": {"class_type": "DistributedCollector", "inputs": {}}}
+        idx = pt.PromptIndex(prompt)
+        m1 = pt.generate_job_id_map(idx, "run")
+        m2 = pt.generate_job_id_map(idx, "run")
+        self.assertEqual(m1, m2)
+
+
+# ---------------------------------------------------------------------------
+# apply_participant_overrides – DistributedCollector
+# ---------------------------------------------------------------------------
+
+class ApplyOverridesCollectorTests(unittest.TestCase):
+    def _collector_prompt(self):
+        return {"1": {"class_type": "DistributedCollector", "inputs": {}}}
+
+    def test_worker_sets_is_worker_true(self):
+        result = _apply(self._collector_prompt(), "worker-a")
+        self.assertTrue(result["1"]["inputs"]["is_worker"])
+
+    def test_worker_sets_master_url(self):
+        result = _apply(self._collector_prompt(), "worker-a")
+        self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
+
+    def test_worker_sets_worker_id(self):
+        result = _apply(self._collector_prompt(), "worker-a")
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
+
+    def test_worker_sets_delegate_only_false(self):
+        result = _apply(self._collector_prompt(), "worker-a")
+        self.assertFalse(result["1"]["inputs"]["delegate_only"])
+
+    def test_master_sets_is_worker_false(self):
+        result = _apply(self._collector_prompt(), "master")
+        self.assertFalse(result["1"]["inputs"]["is_worker"])
+
+    def test_master_clears_stale_master_url(self):
+        prompt = {"1": {"class_type": "DistributedCollector", "inputs": {"master_url": "stale"}}}
+        result = _apply(prompt, "master")
+        self.assertNotIn("master_url", result["1"]["inputs"])
+
+    def test_master_clears_stale_worker_id(self):
+        prompt = {"1": {"class_type": "DistributedCollector", "inputs": {"worker_id": "stale"}}}
+        result = _apply(prompt, "master")
+        self.assertNotIn("worker_id", result["1"]["inputs"])
+
+    def test_master_with_delegate_master_sets_delegate_only_true(self):
+        result = _apply(self._collector_prompt(), "master", delegate_master=True)
+        self.assertTrue(result["1"]["inputs"]["delegate_only"])
+
+    def test_master_without_delegate_master_sets_delegate_only_false(self):
+        result = _apply(self._collector_prompt(), "master", delegate_master=False)
+        self.assertFalse(result["1"]["inputs"]["delegate_only"])
+
+    def test_enabled_worker_ids_serialized_as_json(self):
+        enabled = ["worker-a", "worker-b"]
+        result = _apply(self._collector_prompt(), "master", enabled_worker_ids=enabled)
+        self.assertEqual(result["1"]["inputs"]["enabled_worker_ids"], json.dumps(enabled))
+
+    def test_multi_job_id_is_set_from_job_map(self):
+        prompt = {"1": {"class_type": "DistributedCollector", "inputs": {}}}
+        idx = pt.PromptIndex(prompt)
+        job_id_map = {"1": "run_abc_1"}
+        result = pt.apply_participant_overrides(
+            prompt,
+            participant_id="worker-a",
+            enabled_worker_ids=["worker-a"],
+            job_id_map=job_id_map,
+            master_url="http://master",
+            delegate_master=False,
+            prompt_index=idx,
+        )
+        self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_abc_1")
+
+
+# ---------------------------------------------------------------------------
+# apply_participant_overrides – DistributedSeed
+# ---------------------------------------------------------------------------
+
+class ApplyOverridesSeedTests(unittest.TestCase):
+    def _seed_prompt(self):
+        return {"1": {"class_type": "DistributedSeed", "inputs": {}}}
+
+    def test_worker_sets_is_worker_true(self):
+        result = _apply(self._seed_prompt(), "worker-a")
+        self.assertTrue(result["1"]["inputs"]["is_worker"])
+
+    def test_worker_id_uses_canonical_worker_id(self):
+        result = _apply(self._seed_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"])
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-b")
+
+    def test_master_sets_is_worker_false(self):
+        result = _apply(self._seed_prompt(), "master")
+        self.assertFalse(result["1"]["inputs"]["is_worker"])
+
+    def test_master_sets_empty_worker_id(self):
+        result = _apply(self._seed_prompt(), "master")
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "")
+
+    def test_enabled_worker_ids_is_set(self):
+        enabled = ["worker-a", "worker-b"]
+        result = _apply(self._seed_prompt(), "worker-a", enabled_worker_ids=enabled)
+        self.assertEqual(result["1"]["inputs"]["enabled_worker_ids"], json.dumps(enabled))
+
+
+# ---------------------------------------------------------------------------
+# apply_participant_overrides – UltimateSDUpscaleDistributed
+# ---------------------------------------------------------------------------
+
+class ApplyOverridesUpscaleTests(unittest.TestCase):
+    def _upscale_prompt(self):
+        return {"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}}}
+
+    def test_worker_sets_is_worker_true(self):
+        result = _apply(self._upscale_prompt(), "worker-a")
+        self.assertTrue(result["1"]["inputs"]["is_worker"])
+
+    def test_worker_sets_master_url_and_worker_id(self):
+        result = _apply(self._upscale_prompt(), "worker-a")
+        self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
+
+    def test_master_clears_master_url_and_worker_id(self):
+        prompt = {"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"master_url": "x", "worker_id": "y"}}}
+        result = _apply(prompt, "master")
+        self.assertNotIn("master_url", result["1"]["inputs"])
+        self.assertNotIn("worker_id", result["1"]["inputs"])
+
+    def test_collector_downstream_of_upscale_gets_pass_through(self):
+        """A DistributedCollector that is downstream of UltimateSDUpscaleDistributed → pass_through=True."""
+        prompt = {
+            "1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}},
+            "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
+        }
+        result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a"])
+        self.assertTrue(result["2"]["inputs"].get("pass_through"))
+
+
+# ---------------------------------------------------------------------------
+# apply_participant_overrides – DistributedValue
+# ---------------------------------------------------------------------------
+
+class ApplyOverridesValueTests(unittest.TestCase):
+    def _value_prompt(self):
+        return {"1": {"class_type": "DistributedValue", "inputs": {}}}
+
+    def test_worker_sets_is_worker_true(self):
+        result = _apply(self._value_prompt(), "worker-a")
+        self.assertTrue(result["1"]["inputs"]["is_worker"])
+
+    def test_worker_id_uses_canonical_worker_id(self):
+        result = _apply(self._value_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"])
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-b")
+
+    def test_master_sets_is_worker_false(self):
+        result = _apply(self._value_prompt(), "master")
+        self.assertFalse(result["1"]["inputs"]["is_worker"])
+
+    def test_master_sets_empty_worker_id(self):
+        result = _apply(self._value_prompt(), "master")
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "")
+
+    def test_enabled_worker_ids_is_set(self):
+        enabled = ["worker-a", "worker-b"]
+        result = _apply(self._value_prompt(), "worker-a", enabled_worker_ids=enabled)
+        self.assertEqual(result["1"]["inputs"]["enabled_worker_ids"], json.dumps(enabled))
+
+
+class ApplyOverridesBranchTests(unittest.TestCase):
+    def test_master_gets_branch_zero_worker_gets_branch_one(self):
+        prompt = _branch_prompt()
+        master_result = _apply(prompt, "master", enabled_worker_ids=["worker-a"])
+        worker_result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a"])
+
+        self.assertEqual(master_result["2"]["inputs"]["assigned_branch"], 0)
+        self.assertEqual(worker_result["2"]["inputs"]["assigned_branch"], 1)
+        self.assertIn("3", master_result)
+        self.assertIn("4", master_result)
+        self.assertNotIn("5", master_result)
+        self.assertNotIn("6", master_result)
+        self.assertIn("5", worker_result)
+        self.assertIn("6", worker_result)
+        self.assertNotIn("3", worker_result)
+        self.assertNotIn("4", worker_result)
+
+    def test_delegate_mode_assigns_worker_a_to_branch_zero(self):
+        prompt = _branch_prompt()
+        worker_result = _apply(
+            prompt,
+            "worker-a",
+            enabled_worker_ids=["worker-a", "worker-b"],
+            delegate_master=True,
+        )
+        self.assertEqual(worker_result["2"]["inputs"]["assigned_branch"], 0)
+
+    def test_pruning_keeps_shared_nodes_between_branches(self):
+        prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+            "3": {"class_type": "NodeA", "inputs": {"image": ["2", 0]}},
+            "4": {"class_type": "NodeB", "inputs": {"image": ["2", 1]}},
+            "5": {"class_type": "SaveImage", "inputs": {"images": ["3", 0], "aux": ["4", 0]}},
+        }
+        worker_result = _apply(prompt, "master", enabled_worker_ids=["worker-a"])
+        self.assertIn("3", worker_result)
+        self.assertNotIn("4", worker_result)
+        self.assertIn("5", worker_result)
+        self.assertNotIn("aux", worker_result["5"]["inputs"])
+
+    def test_more_participants_than_branches_sets_unassigned_to_minus_one(self):
+        prompt = _branch_prompt()
+        worker_b_result = _apply(
+            prompt,
+            "worker-b",
+            enabled_worker_ids=["worker-a", "worker-b", "worker-c"],
+        )
+        self.assertEqual(worker_b_result["2"]["inputs"]["assigned_branch"], -1)
+        class_types = {node.get("class_type") for node in worker_b_result.values()}
+        self.assertNotIn("Blur", class_types)
+        self.assertNotIn("Sharpen", class_types)
+
+    def test_worker_with_pruned_outputs_gets_auto_preview_for_assigned_branch(self):
+        prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 3}},
+            "3": {"class_type": "Blur", "inputs": {"image": ["2", 0]}},
+            "4": {"class_type": "PreviewImage", "inputs": {"images": ["3", 0]}},
+        }
+
+        worker_result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a", "worker-b"])
+        preview_nodes = [node for node in worker_result.values() if node.get("class_type") == "PreviewImage"]
+
+        self.assertTrue(preview_nodes)
+        self.assertIn(["2", 1], [node.get("inputs", {}).get("images") for node in preview_nodes])
+
+    def test_unassigned_worker_gets_idle_fallback_output(self):
+        prompt = _branch_prompt()
+        worker_result = _apply(
+            prompt,
+            "worker-c",
+            enabled_worker_ids=["worker-a", "worker-b", "worker-c"],
+        )
+
+        preview_nodes = [node for node in worker_result.values() if node.get("class_type") == "PreviewImage"]
+        empty_nodes = [node for node in worker_result.values() if node.get("class_type") == "DistributedEmptyImage"]
+
+        self.assertTrue(preview_nodes)
+        self.assertTrue(empty_nodes)
+
+
+class ApplyOverridesBranchCollectorTests(unittest.TestCase):
+    def test_branch_collector_inherits_assigned_branch_from_upstream_branch_node(self):
+        master_prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+            "3": {"class_type": "DistributedBranchCollector", "inputs": {"input": ["2", 0], "num_branches": 2}},
+        }
+        worker_prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+            "3": {"class_type": "DistributedBranchCollector", "inputs": {"input": ["2", 1], "num_branches": 2}},
+        }
+
+        master_result = _apply(master_prompt, "master", enabled_worker_ids=["worker-a"])
+        worker_result = _apply(worker_prompt, "worker-a", enabled_worker_ids=["worker-a"])
+
+        self.assertEqual(master_result["3"]["inputs"]["assigned_branch"], 0)
+        self.assertEqual(worker_result["3"]["inputs"]["assigned_branch"], 1)
+
+    def test_branch_collector_sets_multi_job_id(self):
+        prompt = {
+            "1": {"class_type": "DistributedBranchCollector", "inputs": {"input": ["2", 0]}},
+            "2": {"class_type": "KSampler", "inputs": {}},
+        }
+        result = _apply(prompt, "master", enabled_worker_ids=["worker-a"])
+        self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_1")
+
+    def test_branch_collector_uses_upstream_branch_job_id_for_grouped_convergence(self):
+        master_prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+            "3": {"class_type": "DistributedBranchCollector", "inputs": {"input": ["2", 0], "num_branches": 2}},
+        }
+        worker_prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+            "3": {"class_type": "DistributedBranchCollector", "inputs": {"input": ["2", 1], "num_branches": 2}},
+        }
+
+        master_result = _apply(master_prompt, "master", enabled_worker_ids=["worker-a"])
+        worker_result = _apply(worker_prompt, "worker-a", enabled_worker_ids=["worker-a"])
+
+        self.assertEqual(master_result["3"]["inputs"]["multi_job_id"], "run_2")
+        self.assertEqual(worker_result["3"]["inputs"]["multi_job_id"], "run_2")
+
+    def test_worker_prunes_nodes_downstream_of_branch_collector(self):
+        prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+            "3": {"class_type": "Blur", "inputs": {"image": ["2", 1]}},
+            "4": {
+                "class_type": "DistributedBranchCollector",
+                "inputs": {"branch_2": ["3", 0], "num_branches": 2},
+            },
+            "5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}},
+        }
+
+        worker_result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a"])
+
+        class_types = {node.get("class_type") for node in worker_result.values()}
+        self.assertIn("DistributedBranchCollector", class_types)
+        self.assertNotIn("SaveImage", class_types)
+        self.assertIn("PreviewImage", class_types)
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_queue_request.py

@@ -0,0 +1,144 @@
+import importlib.util
+import unittest
+from pathlib import Path
+
+
+def _load_queue_request_module():
+    module_path = Path(__file__).resolve().parents[3] / "api" / "queue_request.py"
+    spec = importlib.util.spec_from_file_location("queue_request", 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_request = _load_queue_request_module()
+parse_queue_request_payload = queue_request.parse_queue_request_payload
+
+
+class QueueRequestPayloadTests(unittest.TestCase):
+    def _base_payload(self):
+        return {
+            "prompt": {"1": {"class_type": "Anything"}},
+            "enabled_worker_ids": ["worker-1"],
+            "client_id": "client-1",
+        }
+
+    def test_normalizes_enabled_worker_ids(self):
+        payload_data = self._base_payload()
+        payload_data["enabled_worker_ids"] = [" worker-a ", "worker-b", "worker-a"]
+        payload_data["delegate_master"] = True
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertEqual(payload.enabled_worker_ids, ["worker-a", "worker-b"])
+        self.assertTrue(payload.delegate_master)
+
+    def test_supports_legacy_workers_field(self):
+        payload_data = self._base_payload()
+        payload_data.pop("enabled_worker_ids", None)
+        payload_data["workers"] = [{"id": "w1"}, "w2", {"id": "w3"}, {"name": "no-id"}]
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "w3"])
+
+    def test_supports_workflow_prompt_fallback(self):
+        payload_data = self._base_payload()
+        payload_data.pop("prompt", None)
+        payload_data["workflow"] = {
+            "prompt": {
+                "10": {"class_type": "DistributedCollector"},
+            }
+        }
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertIn("10", payload.prompt)
+
+    def test_normalizes_trace_execution_id(self):
+        payload_data = self._base_payload()
+        payload_data["trace_execution_id"] = "  exec_123  "
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertEqual(payload.trace_execution_id, "exec_123")
+
+    def test_blank_trace_execution_id_normalizes_to_none(self):
+        payload_data = self._base_payload()
+        payload_data["trace_execution_id"] = "   "
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertIsNone(payload.trace_execution_id)
+
+    def test_workers_field_must_be_list(self):
+        payload_data = self._base_payload()
+        payload_data.pop("enabled_worker_ids", None)
+        payload_data["workers"] = "worker-a"
+        with self.assertRaisesRegex(ValueError, "Field 'workers' must be a list"):
+            parse_queue_request_payload(payload_data)
+
+    def test_trace_execution_id_must_be_string(self):
+        payload_data = self._base_payload()
+        payload_data["trace_execution_id"] = 123
+        with self.assertRaisesRegex(ValueError, "trace_execution_id must be a string"):
+            parse_queue_request_payload(payload_data)
+
+    def test_auto_prepare_is_ignored_for_backward_compat(self):
+        payload_data = self._base_payload()
+        payload_data.pop("prompt", None)
+        payload_data["auto_prepare"] = False
+        payload_data["workflow"] = {
+            "prompt": {"10": {"class_type": "DistributedCollector"}},
+        }
+        payload = parse_queue_request_payload(payload_data)
+        self.assertIn("10", payload.prompt)
+
+    def test_invalid_delegate_master_type_raises(self):
+        payload_data = self._base_payload()
+        payload_data["delegate_master"] = "yes"
+        with self.assertRaisesRegex(ValueError, "delegate_master must be a boolean"):
+            parse_queue_request_payload(payload_data)
+
+    def test_invalid_enabled_worker_ids_type_raises(self):
+        payload_data = self._base_payload()
+        payload_data["enabled_worker_ids"] = "worker-a"
+        with self.assertRaisesRegex(ValueError, "enabled_worker_ids must be a list"):
+            parse_queue_request_payload(payload_data)
+
+    def test_legacy_index_worker_tokens_are_rejected(self):
+        payload_data = self._base_payload()
+        payload_data["enabled_worker_ids"] = ["worker-a", "0", "worker_1"]
+        with self.assertRaisesRegex(ValueError, "legacy index token"):
+            parse_queue_request_payload(payload_data)
+
+    def test_invalid_top_level_payload_raises(self):
+        with self.assertRaisesRegex(ValueError, "Expected a JSON object body"):
+            parse_queue_request_payload(["not", "an", "object"])
+
+    def test_missing_prompt_raises(self):
+        with self.assertRaisesRegex(ValueError, "Field 'prompt' must be an object"):
+            parse_queue_request_payload(
+                {
+                    "workflow": {},
+                    "enabled_worker_ids": ["worker-1"],
+                    "client_id": "client-1",
+                }
+            )
+
+    def test_missing_enabled_worker_ids_raises(self):
+        payload_data = self._base_payload()
+        payload_data.pop("enabled_worker_ids", None)
+        with self.assertRaisesRegex(ValueError, "enabled_worker_ids required"):
+            parse_queue_request_payload(payload_data)
+
+    def test_missing_client_id_raises(self):
+        payload_data = self._base_payload()
+        payload_data.pop("client_id", None)
+        with self.assertRaisesRegex(ValueError, "client_id required"):
+            parse_queue_request_payload(payload_data)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/system/test_worker_monitor.py

@@ -0,0 +1,154 @@
+import importlib.util
+import os
+import tempfile
+import types
+import unittest
+from pathlib import Path
+
+
+def _load_worker_monitor_module():
+    module_path = Path(__file__).resolve().parents[3] / "workers" / "worker_monitor.py"
+    spec = importlib.util.spec_from_file_location("dist_test_worker_monitor", 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
+
+
+worker_monitor = _load_worker_monitor_module()
+
+
+class _FakeWorkerProcess:
+    def __init__(self, pid=9001, poll_sequence=None):
+        self.pid = pid
+        self._poll_sequence = list(poll_sequence or [None])
+        self._poll_calls = 0
+        self.returncode = None
+        self.terminate_calls = 0
+        self.kill_calls = 0
+        self.wait_calls = 0
+
+    def poll(self):
+        if self._poll_calls < len(self._poll_sequence):
+            value = self._poll_sequence[self._poll_calls]
+        else:
+            value = self._poll_sequence[-1]
+        self._poll_calls += 1
+        if value is not None:
+            self.returncode = value
+        return value
+
+    def terminate(self):
+        self.terminate_calls += 1
+        self.returncode = 0
+
+    def kill(self):
+        self.kill_calls += 1
+        self.returncode = -9
+
+    def wait(self, timeout=None):
+        _ = timeout
+        self.wait_calls += 1
+        if self.returncode is None:
+            self.returncode = 0
+        return self.returncode
+
+
+class WorkerMonitorTests(unittest.TestCase):
+    def setUp(self):
+        self._orig_launch = worker_monitor.launch_process_with_timeout
+        self._orig_is_alive = worker_monitor.is_process_alive
+        self._orig_terminate = getattr(worker_monitor, "terminate_process", None)
+        self._orig_sleep = worker_monitor.time.sleep
+        self._orig_signal = worker_monitor.signal.signal
+        self._env_backup = dict(os.environ)
+
+        worker_monitor.time.sleep = lambda _seconds: None
+        worker_monitor.signal.signal = lambda *_args, **_kwargs: None
+
+    def tearDown(self):
+        worker_monitor.launch_process_with_timeout = self._orig_launch
+        worker_monitor.is_process_alive = self._orig_is_alive
+        if self._orig_terminate is not None:
+            worker_monitor.terminate_process = self._orig_terminate
+        worker_monitor.time.sleep = self._orig_sleep
+        worker_monitor.signal.signal = self._orig_signal
+        os.environ.clear()
+        os.environ.update(self._env_backup)
+
+    def test_main_validates_required_env_and_args(self):
+        os.environ.pop("COMFYUI_MASTER_PID", None)
+        self.assertEqual(worker_monitor.main(["python", "main.py"]), 1)
+
+        os.environ["COMFYUI_MASTER_PID"] = "not-an-int"
+        self.assertEqual(worker_monitor.main(["python", "main.py"]), 1)
+
+        os.environ["COMFYUI_MASTER_PID"] = "101"
+        self.assertEqual(worker_monitor.main([]), 1)
+
+    def test_main_delegates_to_monitor_and_run(self):
+        calls = []
+
+        def _monitor(master_pid, command):
+            calls.append((master_pid, list(command)))
+            return 7
+
+        os.environ["COMFYUI_MASTER_PID"] = "202"
+        original_monitor = worker_monitor.monitor_and_run
+        worker_monitor.monitor_and_run = _monitor
+        try:
+            exit_code = worker_monitor.main(["python", "worker.py"])
+        finally:
+            worker_monitor.monitor_and_run = original_monitor
+
+        self.assertEqual(exit_code, 7)
+        self.assertEqual(calls, [(202, ["python", "worker.py"])])
+
+    def test_monitor_and_run_returns_worker_exit_code(self):
+        fake_process = _FakeWorkerProcess(poll_sequence=[None, 5])
+        worker_monitor.launch_process_with_timeout = lambda *_args, **_kwargs: fake_process
+        worker_monitor.is_process_alive = lambda _pid: True
+        worker_monitor.terminate_process = lambda *_args, **_kwargs: None
+
+        exit_code = worker_monitor.monitor_and_run(100, ["python", "worker.py"])
+
+        self.assertEqual(exit_code, 5)
+
+    def test_monitor_and_run_terminates_worker_when_master_dies(self):
+        fake_process = _FakeWorkerProcess(poll_sequence=[None, None])
+        term_calls = []
+
+        worker_monitor.launch_process_with_timeout = lambda *_args, **_kwargs: fake_process
+        worker_monitor.is_process_alive = lambda _pid: False
+
+        def _terminate(process, timeout):
+            term_calls.append((process.pid, timeout))
+            process.returncode = 0
+
+        worker_monitor.terminate_process = _terminate
+
+        exit_code = worker_monitor.monitor_and_run(333, ["python", "worker.py"])
+
+        self.assertEqual(exit_code, 0)
+        self.assertEqual(term_calls, [(9001, worker_monitor.PROCESS_TERMINATION_TIMEOUT)])
+
+    def test_monitor_and_run_writes_pid_file_when_requested(self):
+        fake_process = _FakeWorkerProcess(poll_sequence=[0])
+        worker_monitor.launch_process_with_timeout = lambda *_args, **_kwargs: fake_process
+        worker_monitor.is_process_alive = lambda _pid: True
+        worker_monitor.terminate_process = lambda *_args, **_kwargs: None
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            pid_file = Path(tmpdir) / "worker.pid"
+            os.environ["WORKER_PID_FILE"] = str(pid_file)
+            exit_code = worker_monitor.monitor_and_run(444, ["python", "worker.py"])
+
+            self.assertEqual(exit_code, 0)
+            contents = pid_file.read_text(encoding="utf-8")
+            monitor_pid, worker_pid = contents.split(",")
+            self.assertTrue(monitor_pid.isdigit())
+            self.assertEqual(worker_pid, "9001")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/__init__.py

@@ -0,0 +1 @@
+"""Unit tests."""

tests/unit/upscale/test_conditioning.py

@@ -0,0 +1,53 @@
+import unittest
+
+import torch
+
+from upscale.conditioning import clone_conditioning, clone_control_chain
+
+
+class _FakeControl:
+    def __init__(self, hint, previous=None):
+        self.cond_hint_original = hint
+        self.previous_controlnet = previous
+
+
+class ConditioningUtilsTests(unittest.TestCase):
+    def test_clone_control_chain_clones_hints_and_links(self):
+        tail = _FakeControl(torch.ones((1, 2)))
+        head = _FakeControl(torch.zeros((1, 2)), previous=tail)
+
+        cloned = clone_control_chain(head, clone_hint=True)
+
+        self.assertIsNot(cloned, head)
+        self.assertIsNot(cloned.previous_controlnet, tail)
+        self.assertTrue(torch.equal(cloned.cond_hint_original, head.cond_hint_original))
+        self.assertTrue(torch.equal(cloned.previous_controlnet.cond_hint_original, tail.cond_hint_original))
+        self.assertIsNot(cloned.cond_hint_original, head.cond_hint_original)
+
+    def test_clone_conditioning_clones_emb_and_known_fields(self):
+        control = _FakeControl(torch.ones((1, 3)))
+        cond = [
+            [
+                torch.tensor([[1.0, 2.0, 3.0]]),
+                {
+                    "control": control,
+                    "mask": torch.ones((1, 1, 1)),
+                    "pooled_output": torch.zeros((1, 4)),
+                    "area": [0, 1, 2, 3],
+                },
+            ]
+        ]
+
+        cloned = clone_conditioning(cond, clone_hints=True)
+
+        self.assertTrue(torch.equal(cloned[0][0], cond[0][0]))
+        self.assertIsNot(cloned[0][0], cond[0][0])
+        self.assertIsNot(cloned[0][1]["control"], cond[0][1]["control"])
+        self.assertIsNot(cloned[0][1]["mask"], cond[0][1]["mask"])
+        self.assertIsNot(cloned[0][1]["pooled_output"], cond[0][1]["pooled_output"])
+        self.assertEqual(cloned[0][1]["area"], cond[0][1]["area"])
+        self.assertIsNot(cloned[0][1]["area"], cond[0][1]["area"])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_dynamic.py

@@ -0,0 +1,247 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+from PIL import Image
+try:
+    import torch
+except ModuleNotFoundError:  # pragma: no cover - optional dependency in CI/runtime
+    torch = None
+
+
+def _bootstrap_package(package_name):
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
+    modes_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
+
+    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
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    def pil_to_tensor(image):
+        if torch is None:
+            raise RuntimeError("torch is required for this test helper")
+        array = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(array).unsqueeze(0)
+
+    image_module.tensor_to_pil = tensor_to_pil
+    image_module.pil_to_tensor = pil_to_tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.run_async_in_server_loop = lambda _coro, timeout=None: (_ for _ in ()).throw(
+        RuntimeError(f"run_async_in_server_loop unexpectedly called with timeout={timeout}")
+    )
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.get_worker_timeout_seconds = lambda: 5.0
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.JOB_POLL_INTERVAL = 0.0
+    constants_module.JOB_POLL_MAX_ATTEMPTS = 3
+    constants_module.TILE_WAIT_TIMEOUT = 5.0
+    constants_module.TILE_SEND_TIMEOUT = 5.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    if torch is None:
+        torch_module = types.ModuleType("torch")
+        torch_module.cat = lambda *_args, **_kwargs: None
+        sys.modules["torch"] = torch_module
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: None
+
+    async def init_dynamic_job(*_args, **_kwargs):
+        return None
+
+    job_store_module.init_dynamic_job = init_dynamic_job
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    comfy_module = types.ModuleType("comfy")
+    model_mgmt = types.ModuleType("comfy.model_management")
+    model_mgmt.throw_exception_if_processing_interrupted = lambda: None
+    comfy_module.model_management = model_mgmt
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.model_management"] = model_mgmt
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_dynamic_mode_module():
+    package_name = "dist_dynamic_mode_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args")
+    _load_module(package_name, "upscale/mode_contexts.py", "upscale.mode_contexts")
+    return _load_module(package_name, "upscale/modes/dynamic.py", "upscale.modes.dynamic")
+
+
+dynamic_module = _load_dynamic_mode_module()
+
+
+def _run_sync(value, timeout=None):
+    _ = timeout
+    if asyncio.iscoroutine(value):
+        return asyncio.run(value)
+    return value
+
+
+dynamic_module.run_async_in_server_loop = _run_sync
+
+
+class _DummyDynamicNode(dynamic_module.DynamicModeMixin):
+    def __init__(self, *, poll_ready=True, assignments=None):
+        self._poll_ready = poll_ready
+        self._assignments = list(assignments or [])
+        self.sent_images = []
+        self.heartbeat_count = 0
+        self.completion_sent = False
+
+    def round_to_multiple(self, value):
+        return int(value)
+
+    def calculate_tiles(self, _width, _height, _tile_width, _tile_height, _force_uniform_tiles):
+        return [(0, 0)]
+
+    def _poll_job_ready(self, *_args, **_kwargs):
+        return self._poll_ready
+
+    async def check_job_status(self, *_args, **_kwargs):
+        return self._poll_ready
+
+    async def request_assignment(self, *_args, **_kwargs):
+        if self._assignments:
+            image_idx, estimated_remaining = self._assignments.pop(0)
+            return SimpleNamespace(
+                kind="image" if image_idx is not None else "none",
+                task_idx=image_idx,
+                estimated_remaining=estimated_remaining,
+                batched_static=False,
+            )
+        return SimpleNamespace(kind="none", task_idx=None, estimated_remaining=0, batched_static=False)
+
+    def slice_conditioning(self, positive, negative, _image_idx):
+        return positive, negative
+
+    def process_and_blend_tile(self, _tile_idx, _pos, _source_tensor, local_image, *_args, **_kwargs):
+        return local_image
+
+    async def send_heartbeat(self, *_args, **_kwargs):
+        self.heartbeat_count += 1
+
+    async def send_full_image(self, local_image, image_idx, _multi_job_id, _master_url, _worker_id, is_last):
+        self.sent_images.append((image_idx, is_last, local_image.size))
+
+    async def send_worker_complete_signal(self, *_args, **_kwargs):
+        self.completion_sent = True
+
+
+@unittest.skipIf(torch is None, "torch is not installed")
+class DynamicModeTests(unittest.TestCase):
+    def setUp(self):
+        self.image_batch = torch.zeros((1, 8, 8, 3), dtype=torch.float32)
+        self.core_args = dynamic_module.UpscaleCoreArgs(
+            model=None,
+            positive=None,
+            negative=None,
+            vae=None,
+            seed=1,
+            steps=10,
+            cfg=7.0,
+            sampler_name="euler",
+            scheduler="normal",
+            denoise=0.5,
+            tiled_decode=False,
+        )
+
+    def test_worker_dynamic_returns_early_if_job_not_ready(self):
+        node = _DummyDynamicNode(poll_ready=False)
+
+        result = node.process_worker_dynamic(
+            upscaled_image=self.image_batch,
+            core_args=self.core_args,
+            tile_width=8,
+            tile_height=8,
+            padding=4,
+            mask_blur=2,
+            force_uniform_tiles=True,
+            multi_job_id="job-1",
+            master_url="http://master.local:8188",
+            worker_id="worker-a",
+            enabled_worker_ids='["worker-a"]',
+            dynamic_threshold=8,
+        )
+
+        self.assertEqual(result[0].shape, self.image_batch.shape)
+        self.assertEqual(node.sent_images, [])
+        self.assertFalse(node.completion_sent)
+
+    def test_worker_dynamic_processes_assigned_image_and_sends_completion(self):
+        node = _DummyDynamicNode(poll_ready=True, assignments=[(0, 0)])
+
+        result = node.process_worker_dynamic(
+            upscaled_image=self.image_batch,
+            core_args=self.core_args,
+            tile_width=8,
+            tile_height=8,
+            padding=4,
+            mask_blur=2,
+            force_uniform_tiles=True,
+            multi_job_id="job-2",
+            master_url="http://master.local:8188",
+            worker_id="worker-a",
+            enabled_worker_ids='["worker-a"]',
+            dynamic_threshold=8,
+        )
+
+        self.assertEqual(result[0].shape, self.image_batch.shape)
+        self.assertEqual(len(node.sent_images), 1)
+        self.assertEqual(node.sent_images[0][0], 0)
+        self.assertTrue(node.sent_images[0][1])
+        self.assertTrue(node.completion_sent)
+        self.assertGreaterEqual(node.heartbeat_count, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_image.py

@@ -0,0 +1,39 @@
+import unittest
+
+import torch
+from PIL import Image
+
+from utils.image import ensure_contiguous, pil_to_tensor, tensor_to_pil
+
+
+class ImageUtilsTests(unittest.TestCase):
+    def test_tensor_to_pil_converts_batch_item(self):
+        tensor = torch.tensor(
+            [[[[0.0, 0.5, 1.0], [1.0, 0.0, 0.0]]]],
+            dtype=torch.float32,
+        )
+        image = tensor_to_pil(tensor, batch_index=0)
+
+        self.assertEqual(image.size, (2, 1))
+        self.assertEqual(image.getpixel((0, 0)), (0, 127, 255))
+        self.assertEqual(image.getpixel((1, 0)), (255, 0, 0))
+
+    def test_pil_to_tensor_adds_channel_for_grayscale(self):
+        grayscale = Image.new("L", (3, 2), color=128)
+        tensor = pil_to_tensor(grayscale)
+
+        self.assertEqual(tuple(tensor.shape), (1, 2, 3, 1))
+        self.assertAlmostEqual(float(tensor[0, 0, 0, 0]), 128 / 255.0, places=6)
+
+    def test_ensure_contiguous_makes_non_contiguous_tensor_contiguous(self):
+        base = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)
+        non_contiguous = base.transpose(1, 2)
+        self.assertFalse(non_contiguous.is_contiguous())
+
+        contiguous = ensure_contiguous(non_contiguous)
+        self.assertTrue(contiguous.is_contiguous())
+        self.assertTrue(torch.equal(contiguous, non_contiguous))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_job_state.py

@@ -0,0 +1,131 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_job_state_modules():
+    package_name = "dist_job_state_testpkg"
+    _bootstrap_package(package_name)
+    job_models = _load_module(package_name, "upscale/job_models.py", "upscale.job_models")
+
+    prompt_server = types.SimpleNamespace(
+        distributed_pending_tile_jobs={},
+        distributed_tile_jobs_lock=asyncio.Lock(),
+    )
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    call_log = {"args": None}
+    job_timeout_module = types.ModuleType(f"{package_name}.upscale.job_timeout")
+
+    async def _requeue(multi_job_id, batch_size):
+        call_log["args"] = (multi_job_id, batch_size)
+        return 3
+
+    job_timeout_module._check_and_requeue_timed_out_workers = _requeue
+    sys.modules[f"{package_name}.upscale.job_timeout"] = job_timeout_module
+
+    job_state = _load_module(package_name, "upscale/job_state.py", "upscale.job_state")
+    return job_models, job_state, prompt_server, call_log
+
+
+job_models_module, job_state_module, prompt_server, timeout_call_log = _load_job_state_modules()
+
+
+class _Node(job_state_module.JobStateMixin):
+    pass
+
+
+class JobStateMixinTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        prompt_server.distributed_pending_tile_jobs = {}
+        prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
+        self.node = _Node()
+
+    async def test_get_all_completed_tasks_supports_tile_and_image_modes(self):
+        tile_state = job_models_module.TileJobState("tile-job")
+        tile_state.completed_tasks[4] = "tile-4"
+        image_state = job_models_module.ImageJobState("image-job")
+        image_state.completed_images[2] = "image-2"
+        prompt_server.distributed_pending_tile_jobs = {
+            "tile-job": tile_state,
+            "image-job": image_state,
+        }
+
+        self.assertEqual(await self.node._get_all_completed_tasks("tile-job"), {4: "tile-4"})
+        self.assertEqual(await self.node._get_all_completed_tasks("image-job"), {2: "image-2"})
+
+    async def test_next_index_and_pending_count_from_queues(self):
+        image_state = job_models_module.ImageJobState("image-job")
+        await image_state.pending_images.put(7)
+        tile_state = job_models_module.TileJobState("tile-job")
+        await tile_state.pending_tasks.put(9)
+        prompt_server.distributed_pending_tile_jobs = {
+            "image-job": image_state,
+            "tile-job": tile_state,
+        }
+
+        self.assertEqual(await self.node._get_next_image_index("image-job"), 7)
+        self.assertEqual(await self.node._get_next_tile_index("tile-job"), 9)
+        self.assertEqual(await self.node._get_pending_count("image-job"), 0)
+        self.assertEqual(await self.node._get_pending_count("tile-job"), 0)
+
+    async def test_drain_worker_results_queue_collects_completed_images(self):
+        image_state = job_models_module.ImageJobState("image-job")
+        await image_state.queue.put({"worker_id": "w1", "image_idx": 1, "image": "img1"})
+        await image_state.queue.put({"worker_id": "w2", "image_idx": 2, "image": "img2"})
+        prompt_server.distributed_pending_tile_jobs = {"image-job": image_state}
+
+        drained = await self.node._drain_worker_results_queue("image-job")
+
+        self.assertEqual(drained, 2)
+        self.assertEqual(image_state.completed_images[1], "img1")
+        self.assertEqual(image_state.completed_images[2], "img2")
+
+    async def test_check_and_requeue_delegates_to_timeout_helper(self):
+        count = await self.node._check_and_requeue_timed_out_workers("job-abc", 5)
+        self.assertEqual(count, 3)
+        self.assertEqual(timeout_call_log["args"], ("job-abc", 5))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_job_store.py

@@ -0,0 +1,164 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    prompt_server = types.SimpleNamespace()
+
+    runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
+
+    def _ensure_distributed_runtime_state(server_instance=None):
+        ps = server_instance or prompt_server
+        if not hasattr(ps, "distributed_pending_jobs"):
+            ps.distributed_pending_jobs = {}
+        if not hasattr(ps, "distributed_jobs_lock"):
+            ps.distributed_jobs_lock = asyncio.Lock()
+        if not hasattr(ps, "distributed_job_allowed_workers"):
+            ps.distributed_job_allowed_workers = {}
+        if not hasattr(ps, "distributed_pending_tile_jobs"):
+            ps.distributed_pending_tile_jobs = {}
+        if not hasattr(ps, "distributed_tile_jobs_lock"):
+            ps.distributed_tile_jobs_lock = asyncio.Lock()
+        return types.SimpleNamespace(
+            distributed_pending_jobs=ps.distributed_pending_jobs,
+            distributed_jobs_lock=ps.distributed_jobs_lock,
+            distributed_job_allowed_workers=ps.distributed_job_allowed_workers,
+            distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs,
+            distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock,
+        )
+
+    runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state
+    runtime_state_module.get_prompt_server_instance = lambda: prompt_server
+    sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_job_store_modules():
+    package_name = "dist_job_store_testpkg"
+    _bootstrap_package(package_name)
+    job_models = _load_module(package_name, "upscale/job_models.py", "upscale.job_models")
+    job_store = _load_module(package_name, "upscale/job_store.py", "upscale.job_store")
+    return job_models, job_store
+
+
+job_models_module, job_store_module = _load_job_store_modules()
+
+
+class JobStoreTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.prompt_server = job_store_module.get_prompt_server_instance()
+        self.prompt_server.distributed_pending_tile_jobs = {}
+        self.prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
+
+    async def test_ensure_tile_jobs_initialized_prunes_invalid_entries(self):
+        self.prompt_server.distributed_pending_tile_jobs["valid"] = job_models_module.TileJobState("job-valid")
+        self.prompt_server.distributed_pending_tile_jobs["invalid"] = {"not": "state"}
+
+        job_store_module.ensure_tile_jobs_initialized()
+
+        self.assertIn("valid", self.prompt_server.distributed_pending_tile_jobs)
+        self.assertNotIn("invalid", self.prompt_server.distributed_pending_tile_jobs)
+
+    async def test_init_dynamic_job_populates_pending_images(self):
+        await job_store_module.init_dynamic_job(
+            multi_job_id="job-dynamic",
+            batch_size=3,
+            enabled_workers=["worker-a"],
+            all_indices=[2, 0, 1],
+        )
+
+        job_data = self.prompt_server.distributed_pending_tile_jobs["job-dynamic"]
+        self.assertIsInstance(job_data, job_models_module.ImageJobState)
+        self.assertEqual(job_data.batch_size, 3)
+        self.assertIn("worker-a", job_data.worker_status)
+
+        pulled = [job_data.pending_images.get_nowait() for _ in range(3)]
+        self.assertEqual(pulled, [2, 0, 1])
+
+    async def test_init_static_job_batched_queues_tile_ids(self):
+        await job_store_module.init_static_job_batched(
+            multi_job_id="job-static",
+            batch_size=2,
+            num_tiles_per_image=4,
+            enabled_workers=["worker-a", "worker-b"],
+        )
+
+        job_data = self.prompt_server.distributed_pending_tile_jobs["job-static"]
+        self.assertIsInstance(job_data, job_models_module.TileJobState)
+        self.assertTrue(job_data.batched_static)
+
+        pulled = [job_data.pending_tasks.get_nowait() for _ in range(4)]
+        self.assertEqual(pulled, [0, 1, 2, 3])
+
+    async def test_drain_results_queue_collects_images_tiles_and_completion(self):
+        job_data = job_models_module.TileJobState("job-drain")
+        job_data.worker_status = {"worker-a": 1.0, "worker-b": 1.0}
+        self.prompt_server.distributed_pending_tile_jobs["job-drain"] = job_data
+
+        await job_data.queue.put(
+            {
+                "worker_id": "worker-a",
+                "is_last": True,
+                "image_idx": 5,
+                "image": "img-5",
+            }
+        )
+        await job_data.queue.put(
+            {
+                "worker_id": "worker-b",
+                "is_last": False,
+                "tiles": [
+                    {"tile_idx": 0, "global_idx": 0, "tensor": "tile-0"},
+                    {"tile_idx": 1, "global_idx": 1, "tensor": "tile-1"},
+                ],
+            }
+        )
+
+        drained = await job_store_module.drain_results_queue("job-drain")
+
+        self.assertEqual(drained, 3)
+        self.assertEqual(job_data.completed_tasks[5], "img-5")
+        self.assertEqual(job_data.completed_tasks[0]["tensor"], "tile-0")
+        self.assertEqual(job_data.completed_tasks[1]["tensor"], "tile-1")
+        self.assertNotIn("worker-a", job_data.worker_status)
+        self.assertIn("worker-b", job_data.worker_status)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_job_timeout.py

@@ -0,0 +1,197 @@
+import asyncio
+import importlib.util
+import sys
+import time
+import types
+import unittest
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+def _load_job_timeout_module():
+    module_path = Path(__file__).resolve().parents[3] / "upscale" / "job_timeout.py"
+    package_name = "dist_job_timeout_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
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    config_holder = {"value": {"settings": {}, "workers": []}}
+    probe_holder = {"fn": None}
+    prompt_server_holder = {
+        "value": types.SimpleNamespace(
+            distributed_tile_jobs_lock=asyncio.Lock(),
+            distributed_pending_tile_jobs={},
+        )
+    }
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: config_holder["value"]
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.HEARTBEAT_TIMEOUT = 60
+    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_worker_url = lambda worker: f"http://{worker.get('host', '127.0.0.1')}:{worker.get('port', 8188)}"
+
+    async def _probe_worker(url, timeout=2.0):
+        fn = probe_holder["fn"]
+        if fn is None:
+            return None
+        return await fn(url, timeout)
+
+    network_module.probe_worker = _probe_worker
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
+
+    class BaseJobState:
+        pass
+
+    @dataclass
+    class ImageJobState(BaseJobState):
+        multi_job_id: str
+        mode: str = field(default="dynamic", init=False)
+        queue: asyncio.Queue = field(default_factory=asyncio.Queue)
+        pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
+        completed_images: dict = field(default_factory=dict)
+        worker_status: dict = field(default_factory=dict)
+        assigned_to_workers: dict = field(default_factory=dict)
+        batch_size: int = 0
+        num_tiles_per_image: int = 0
+        batched_static: bool = False
+
+        @property
+        def pending_tasks(self):
+            return self.pending_images
+
+        @property
+        def completed_tasks(self):
+            return self.completed_images
+
+    @dataclass
+    class TileJobState(BaseJobState):
+        multi_job_id: str
+        mode: str = field(default="static", init=False)
+        queue: asyncio.Queue = field(default_factory=asyncio.Queue)
+        pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
+        completed_tasks: dict = field(default_factory=dict)
+        worker_status: dict = field(default_factory=dict)
+        assigned_to_workers: dict = field(default_factory=dict)
+        batch_size: int = 0
+        num_tiles_per_image: int = 0
+        batched_static: bool = False
+
+    job_models_module.BaseJobState = BaseJobState
+    job_models_module.ImageJobState = ImageJobState
+    job_models_module.TileJobState = TileJobState
+    sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.job_timeout", 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)
+
+    module._config_holder = config_holder
+    module._probe_holder = probe_holder
+    module._prompt_server_holder = prompt_server_holder
+    module._ImageJobState = ImageJobState
+    module._TileJobState = TileJobState
+    return module
+
+
+jt = _load_job_timeout_module()
+
+
+class JobTimeoutRequeueTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        jt._prompt_server_holder["value"] = types.SimpleNamespace(
+            distributed_tile_jobs_lock=asyncio.Lock(),
+            distributed_pending_tile_jobs={},
+        )
+        jt._config_holder["value"] = {
+            "settings": {"worker_timeout_seconds": 5},
+            "workers": [{"id": "worker-1", "host": "worker.local", "port": 8188}],
+        }
+
+    async def test_requeues_only_incomplete_dynamic_tasks_for_timed_out_worker(self):
+        async def _offline_probe(_url, _timeout):
+            return None
+
+        jt._probe_holder["fn"] = _offline_probe
+        prompt_server = jt._prompt_server_holder["value"]
+        job_data = jt._ImageJobState("job-1")
+        job_data.worker_status["worker-1"] = time.time() - 60.0
+        job_data.assigned_to_workers["worker-1"] = [0, 1]
+        job_data.completed_images[1] = "done"
+        prompt_server.distributed_pending_tile_jobs["job-1"] = job_data
+
+        requeued = await jt._check_and_requeue_timed_out_workers("job-1", total_tasks=2)
+
+        self.assertEqual(requeued, 1)
+        self.assertEqual(await job_data.pending_images.get(), 0)
+        self.assertNotIn("worker-1", job_data.worker_status)
+        self.assertEqual(job_data.assigned_to_workers["worker-1"], [])
+
+    async def test_busy_probe_graces_worker_and_skips_requeue(self):
+        async def _busy_probe(_url, _timeout):
+            return {"exec_info": {"queue_remaining": 3}}
+
+        jt._probe_holder["fn"] = _busy_probe
+        prompt_server = jt._prompt_server_holder["value"]
+        job_data = jt._ImageJobState("job-2")
+        old_heartbeat = time.time() - 60.0
+        job_data.worker_status["worker-1"] = old_heartbeat
+        job_data.assigned_to_workers["worker-1"] = [0]
+        prompt_server.distributed_pending_tile_jobs["job-2"] = job_data
+
+        requeued = await jt._check_and_requeue_timed_out_workers("job-2", total_tasks=1)
+
+        self.assertEqual(requeued, 0)
+        self.assertIn("worker-1", job_data.worker_status)
+        self.assertGreaterEqual(job_data.worker_status["worker-1"], old_heartbeat)
+        self.assertTrue(job_data.pending_images.empty())
+
+    async def test_completed_dynamic_task_is_not_requeued(self):
+        async def _offline_probe(_url, _timeout):
+            return None
+
+        jt._probe_holder["fn"] = _offline_probe
+        prompt_server = jt._prompt_server_holder["value"]
+        job_data = jt._ImageJobState("job-3")
+        job_data.worker_status["worker-1"] = time.time() - 60.0
+        job_data.assigned_to_workers["worker-1"] = [7]
+        job_data.completed_images[7] = "complete"
+        prompt_server.distributed_pending_tile_jobs["job-3"] = job_data
+
+        requeued = await jt._check_and_requeue_timed_out_workers("job-3", total_tasks=1)
+
+        self.assertEqual(requeued, 0)
+        self.assertTrue(job_data.pending_images.empty())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_payload_parsers.py

@@ -0,0 +1,208 @@
+import importlib.util
+import io
+import json
+import sys
+import types
+import unittest
+from pathlib import Path
+
+try:
+    from PIL import Image as PILImage
+    PIL_AVAILABLE = True
+except ImportError:
+    PIL_AVAILABLE = False
+
+
+def _load_payload_parsers_module():
+    # payload_parsers.py has no relative imports; only stdlib + PIL
+    module_path = Path(__file__).resolve().parents[3] / "upscale" / "payload_parsers.py"
+    spec = importlib.util.spec_from_file_location("upscale_payload_parsers", 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
+
+
+if PIL_AVAILABLE:
+    pp = _load_payload_parsers_module()
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_png_bytes(width=64, height=64, color=(128, 64, 32)):
+    """Return raw PNG bytes for a solid-colour image."""
+    img = PILImage.new("RGB", (width, height), color=color)
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+class _MockFileField:
+    """Minimal multipart file-field stub."""
+
+    class _MockFile:
+        def __init__(self, data: bytes):
+            self._buf = io.BytesIO(data)
+
+        def read(self) -> bytes:
+            return self._buf.read()
+
+    def __init__(self, data: bytes):
+        self.file = self._MockFile(data)
+
+
+def _make_form(n_tiles, *, padding=None, extra_meta=None, image_color=(128, 64, 32)):
+    """Build a minimal form-data dict with `n_tiles` tile entries."""
+    image_bytes = _make_png_bytes(color=image_color)
+    metadata = []
+    for i in range(n_tiles):
+        entry = {
+            "tile_idx": i,
+            "x": i * 64,
+            "y": 0,
+            "extracted_width": 64,
+            "extracted_height": 64,
+        }
+        if extra_meta and i < len(extra_meta):
+            entry.update(extra_meta[i])
+        metadata.append(entry)
+
+    form = {"tiles_metadata": json.dumps(metadata)}
+    if padding is not None:
+        form["padding"] = str(padding)
+    for i in range(n_tiles):
+        form[f"tile_{i}"] = _MockFileField(image_bytes)
+    return form
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+@unittest.skipUnless(PIL_AVAILABLE, "PIL not installed")
+class ParseTilesFromFormTests(unittest.TestCase):
+
+    # --- happy paths ---
+
+    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.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_padding_is_parsed_from_form(self):
+        tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
+        self.assertEqual(tiles[0]["padding"], 16)
+
+    def test_default_padding_is_zero(self):
+        form = _make_form(1)
+        form.pop("padding", None)
+        tiles = pp._parse_tiles_from_form(form)
+        self.assertEqual(tiles[0]["padding"], 0)
+
+    def test_invalid_padding_string_falls_back_to_zero(self):
+        form = _make_form(1)
+        form["padding"] = "not_a_number"
+        tiles = pp._parse_tiles_from_form(form)
+        self.assertEqual(tiles[0]["padding"], 0)
+
+    def test_optional_batch_idx_included_when_present(self):
+        extra = [{"batch_idx": 2}]
+        tiles = pp._parse_tiles_from_form(_make_form(1, extra_meta=extra))
+        self.assertEqual(tiles[0]["batch_idx"], 2)
+
+    def test_optional_global_idx_included_when_present(self):
+        extra = [{"global_idx": 5}]
+        tiles = pp._parse_tiles_from_form(_make_form(1, extra_meta=extra))
+        self.assertEqual(tiles[0]["global_idx"], 5)
+
+    def test_batch_idx_and_global_idx_absent_when_not_in_metadata(self):
+        tiles = pp._parse_tiles_from_form(_make_form(1))
+        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):
+        with self.assertRaises(ValueError, msg="Missing tiles_metadata"):
+            pp._parse_tiles_from_form({})
+
+    def test_invalid_json_metadata_raises_value_error(self):
+        form = {"tiles_metadata": "{not valid json}"}
+        with self.assertRaises(ValueError):
+            pp._parse_tiles_from_form(form)
+
+    def test_non_list_metadata_raises_value_error(self):
+        form = {"tiles_metadata": json.dumps({"not": "a list"})}
+        with self.assertRaises(ValueError):
+            pp._parse_tiles_from_form(form)
+
+    def test_missing_tile_file_field_raises_value_error(self):
+        form = {
+            "tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
+            # tile_0 intentionally omitted
+        }
+        with self.assertRaises(ValueError):
+            pp._parse_tiles_from_form(form)
+
+    def test_tile_field_without_file_attr_raises_value_error(self):
+        form = {
+            "tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
+            "tile_0": "plain string without .file",
+        }
+        with self.assertRaises(ValueError):
+            pp._parse_tiles_from_form(form)
+
+    def test_non_image_bytes_raises_value_error(self):
+        class _BadFileField:
+            class _BadFile:
+                def read(self):
+                    return b"this is definitely not image data"
+            file = _BadFile()
+
+        form = {
+            "tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
+            "tile_0": _BadFileField(),
+        }
+        with self.assertRaises(ValueError):
+            pp._parse_tiles_from_form(form)
+
+    def test_invalid_metadata_value_type_raises_value_error(self):
+        """Non-integer metadata fields (x, y, etc.) should raise ValueError."""
+        form = {
+            "tiles_metadata": json.dumps([{"tile_idx": 0, "x": "not_int", "y": 0}]),
+            "tile_0": _MockFileField(_make_png_bytes()),
+        }
+        with self.assertRaises(ValueError):
+            pp._parse_tiles_from_form(form)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_result_collector.py

@@ -0,0 +1,248 @@
+import asyncio
+import importlib.util
+import sys
+import time
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    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.05
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.DYNAMIC_MODE_MAX_POLL_TIMEOUT = 10.0
+    constants_module.HEARTBEAT_INTERVAL = 0.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    comfy_module = types.ModuleType("comfy")
+    model_mgmt = types.ModuleType("comfy.model_management")
+
+    class _InterruptProcessingException(Exception):
+        pass
+
+    model_mgmt.processing_interrupted = lambda: False
+    model_mgmt.InterruptProcessingException = _InterruptProcessingException
+    comfy_module.model_management = model_mgmt
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.model_management"] = model_mgmt
+
+    prompt_server = types.SimpleNamespace(
+        distributed_pending_tile_jobs={},
+        distributed_tile_jobs_lock=asyncio.Lock(),
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
+    sys.modules["server"] = server_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_result_collector_modules():
+    package_name = "dist_result_collector_testpkg"
+    _bootstrap_package(package_name)
+    job_models = _load_module(package_name, "upscale/job_models.py", "upscale.job_models")
+    job_store = _load_module(package_name, "upscale/job_store.py", "upscale.job_store")
+
+    job_timeout = types.ModuleType(f"{package_name}.upscale.job_timeout")
+
+    async def _noop_requeue(_multi_job_id, _batch_size):
+        return 0
+
+    job_timeout._check_and_requeue_timed_out_workers = _noop_requeue
+    job_timeout.check_and_requeue_timed_out_workers = _noop_requeue
+    sys.modules[f"{package_name}.upscale.job_timeout"] = job_timeout
+
+    result_collector = _load_module(
+        package_name,
+        "upscale/result_collector.py",
+        "upscale.result_collector",
+    )
+    return job_models, job_store, result_collector
+
+
+job_models_module, job_store_module, result_collector_module = _load_result_collector_modules()
+
+
+class _Node(result_collector_module.ResultCollectorMixin):
+    def __init__(self):
+        self.requeue_calls = []
+
+    async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
+        self.requeue_calls.append((multi_job_id, batch_size))
+        return 1
+
+
+class ResultCollectorMixinTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.prompt_server = job_store_module.server.PromptServer.instance
+        self.prompt_server.distributed_pending_tile_jobs = {}
+        self.prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
+        self.node = _Node()
+
+    async def test_static_collection_collects_tiles_and_cleans_up_job(self):
+        job_data = job_models_module.TileJobState("job-static")
+        self.prompt_server.distributed_pending_tile_jobs["job-static"] = job_data
+        await job_data.queue.put(
+            {
+                "worker_id": "worker-a",
+                "is_last": True,
+                "tiles": [
+                    {"tile_idx": 3},
+                    {
+                        "tile_idx": 2,
+                        "x": 10,
+                        "y": 20,
+                        "extracted_width": 64,
+                        "extracted_height": 64,
+                        "padding": 8,
+                        "batch_idx": 1,
+                        "global_idx": 9,
+                        "tensor": "tile-9",
+                    },
+                ],
+            }
+        )
+
+        collected = await self.node._async_collect_worker_tiles(
+            "job-static",
+            num_workers=1,
+        )
+
+        self.assertEqual(list(collected.keys()), [9])
+        self.assertEqual(collected[9]["tensor"], "tile-9")
+        self.assertEqual(collected[9]["worker_id"], "worker-a")
+        self.assertNotIn("job-static", self.prompt_server.distributed_pending_tile_jobs)
+
+    async def test_dynamic_collection_stops_at_remaining_target_and_cleans_up(self):
+        job_data = job_models_module.ImageJobState("job-dynamic")
+        self.prompt_server.distributed_pending_tile_jobs["job-dynamic"] = job_data
+        await job_data.queue.put(
+            {
+                "worker_id": "worker-a",
+                "is_last": False,
+                "image_idx": 2,
+                "image": "img-2",
+            }
+        )
+
+        completed = await self.node.collect_dynamic_images(
+            "job-dynamic",
+            remaining_to_collect=1,
+            num_workers=2,
+            batch_size=3,
+            master_processed_count=0,
+        )
+
+        self.assertEqual(completed[2], "img-2")
+        self.assertNotIn("job-dynamic", self.prompt_server.distributed_pending_tile_jobs)
+
+    async def test_dynamic_timeout_calls_requeue_helper(self):
+        job_data = job_models_module.ImageJobState("job-timeout")
+        job_data.worker_status = {"worker-a": time.time() - 5}
+        self.prompt_server.distributed_pending_tile_jobs["job-timeout"] = job_data
+
+        old_timeout = result_collector_module.get_worker_timeout_seconds
+        result_collector_module.get_worker_timeout_seconds = lambda: 0.01
+        try:
+            completed = await self.node.collect_dynamic_images(
+                "job-timeout",
+                remaining_to_collect=None,
+                num_workers=1,
+                batch_size=4,
+                master_processed_count=0,
+            )
+        finally:
+            result_collector_module.get_worker_timeout_seconds = old_timeout
+
+        self.assertEqual(completed, {})
+        self.assertEqual(self.node.requeue_calls, [("job-timeout", 4)])
+        self.assertNotIn("job-timeout", self.prompt_server.distributed_pending_tile_jobs)
+
+    async def test_log_worker_timeout_status_handles_non_job_state(self):
+        worker_ids = self.node._log_worker_timeout_status(
+            {"worker_status": {"worker-a": time.time()}},
+            current_time=time.time(),
+            multi_job_id="job-invalid",
+        )
+        self.assertEqual(worker_ids, [])
+
+    async def test_log_worker_timeout_status_returns_worker_ids_for_job_state(self):
+        job_data = job_models_module.ImageJobState("job-state")
+        job_data.worker_status = {"worker-a": time.time() - 1}
+
+        worker_ids = self.node._log_worker_timeout_status(
+            job_data,
+            current_time=time.time(),
+            multi_job_id="job-state",
+        )
+        self.assertEqual(worker_ids, ["worker-a"])
+
+    async def test_static_mode_mismatch_raises(self):
+        self.prompt_server.distributed_pending_tile_jobs["job-mismatch"] = (
+            job_models_module.ImageJobState("job-mismatch")
+        )
+        with self.assertRaises(RuntimeError):
+            await self.node._async_collect_worker_tiles(
+                "job-mismatch",
+                num_workers=1,
+            )
+
+    async def test_dynamic_mode_mismatch_raises(self):
+        self.prompt_server.distributed_pending_tile_jobs["job-mismatch"] = (
+            job_models_module.TileJobState("job-mismatch")
+        )
+        with self.assertRaises(RuntimeError):
+            await self.node.collect_dynamic_images(
+                "job-mismatch",
+                remaining_to_collect=1,
+                num_workers=1,
+                batch_size=1,
+                master_processed_count=0,
+            )
+
+    async def test_mark_image_completed_updates_job_state(self):
+        job_data = job_models_module.ImageJobState("job-mark")
+        self.prompt_server.distributed_pending_tile_jobs["job-mark"] = job_data
+
+        await self.node.mark_image_completed("job-mark", 7, "img-7")
+
+        self.assertEqual(job_data.completed_images[7], "img-7")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_single_gpu.py

@@ -0,0 +1,207 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+try:
+    import torch
+except ModuleNotFoundError:  # pragma: no cover - optional dependency in CI/runtime
+    torch = None
+
+
+def _bootstrap_package(package_name):
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
+    modes_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
+
+    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
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    def pil_to_tensor(image):
+        if torch is None:
+            raise RuntimeError("torch is required for this test helper")
+        array = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(array).unsqueeze(0)
+
+    def blend_processed_batch_item(
+        result_images,
+        processed_batch,
+        batch_index,
+        blend_fn,
+        x1,
+        y1,
+        ew,
+        eh,
+        tile_mask,
+        padding,
+    ):
+        tile_pil = tensor_to_pil(processed_batch, batch_index)
+        if tile_pil.size != (ew, eh):
+            tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
+        result_images[batch_index] = blend_fn(
+            result_images[batch_index],
+            tile_pil,
+            x1,
+            y1,
+            (ew, eh),
+            tile_mask,
+            padding,
+        )
+
+    image_module.tensor_to_pil = tensor_to_pil
+    image_module.pil_to_tensor = pil_to_tensor
+    image_module.blend_processed_batch_item = blend_processed_batch_item
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    if torch is None:
+        torch_module = types.ModuleType("torch")
+        torch_module.cat = lambda *_args, **_kwargs: None
+        sys.modules["torch"] = torch_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_single_gpu_mode_module():
+    package_name = "dist_single_gpu_mode_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args")
+    _load_module(package_name, "upscale/tile_processing.py", "upscale.tile_processing")
+    _load_module(package_name, "upscale/mode_contexts.py", "upscale.mode_contexts")
+    return _load_module(package_name, "upscale/modes/single_gpu.py", "upscale.modes.single_gpu")
+
+
+single_gpu_module = _load_single_gpu_mode_module()
+
+
+class _DummySingleGpuNode(single_gpu_module.SingleGpuModeMixin):
+    def __init__(self):
+        self.round_calls = []
+        self.process_batch_calls = 0
+
+    def round_to_multiple(self, value):
+        self.round_calls.append(int(value))
+        return int(value)
+
+    def calculate_tiles(self, _width, _height, _tile_width, _tile_height, _force_uniform_tiles):
+        return [(0, 0)]
+
+    def create_tile_mask(self, _width, _height, _tx, _ty, _tile_width, _tile_height, _mask_blur):
+        return None
+
+    def extract_batch_tile_with_padding(self, source_batch, _tx, _ty, _tile_width, _tile_height, _padding, _force_uniform_tiles):
+        _, h, w, _ = source_batch.shape
+        return source_batch, 0, 0, w, h
+
+    def process_tiles_batch(
+        self,
+        tile_batch,
+        _model,
+        _positive,
+        _negative,
+        _vae,
+        _seed,
+        _steps,
+        _cfg,
+        _sampler_name,
+        _scheduler,
+        _denoise,
+        _tiled_decode,
+        _region,
+        _canvas_shape,
+    ):
+        self.process_batch_calls += 1
+        return tile_batch
+
+    def blend_tile(self, _base_image, tile_pil, _x1, _y1, _size, _tile_mask, _padding):
+        return tile_pil
+
+
+@unittest.skipIf(torch is None, "torch is not installed")
+class SingleGpuModeTests(unittest.TestCase):
+    def setUp(self):
+        self.node = _DummySingleGpuNode()
+        self.images = torch.rand((2, 8, 8, 3), dtype=torch.float32)
+        self.core_args = single_gpu_module.UpscaleCoreArgs(
+            model=None,
+            positive=None,
+            negative=None,
+            vae=None,
+            seed=1,
+            steps=10,
+            cfg=7.5,
+            sampler_name="euler",
+            scheduler="normal",
+            denoise=0.5,
+            tiled_decode=False,
+        )
+
+    def test_process_single_gpu_preserves_batch_shape(self):
+        result = self.node.process_single_gpu(
+            upscaled_image=self.images,
+            core_args=self.core_args,
+            tile_width=8,
+            tile_height=8,
+            padding=4,
+            mask_blur=2,
+            force_uniform_tiles=True,
+        )[0]
+
+        self.assertEqual(result.shape, self.images.shape)
+        self.assertTrue(torch.allclose(result, self.images, atol=(1.5 / 255.0)))
+
+    def test_process_single_gpu_invokes_rounding_and_batch_processing(self):
+        self.node.process_single_gpu(
+            upscaled_image=self.images,
+            core_args=self.core_args,
+            tile_width=16,
+            tile_height=24,
+            padding=4,
+            mask_blur=2,
+            force_uniform_tiles=True,
+        )
+
+        self.assertEqual(self.node.round_calls[:2], [16, 24])
+        self.assertEqual(self.node.process_batch_calls, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_static.py

@@ -0,0 +1,323 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+
+
+def _load_static_mode_module():
+    module_path = Path(__file__).resolve().parents[3] / "upscale" / "modes" / "static.py"
+    package_name = "dist_static_mode_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
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
+    modes_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    created_comfy_stub = False
+    if "comfy" not in sys.modules:
+        created_comfy_stub = True
+        comfy_module = types.ModuleType("comfy")
+        model_mgmt = types.ModuleType("comfy.model_management")
+
+        class _InterruptProcessingException(Exception):
+            pass
+
+        model_mgmt.processing_interrupted = lambda: False
+        model_mgmt.throw_exception_if_processing_interrupted = lambda: None
+        model_mgmt.InterruptProcessingException = _InterruptProcessingException
+
+        comfy_module.model_management = model_mgmt
+        sys.modules["comfy"] = comfy_module
+        sys.modules["comfy.model_management"] = model_mgmt
+
+    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
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+    from PIL import Image as PILImage
+    import numpy as np
+
+    def _tensor_to_pil(img_tensor, batch_index=0):
+        return PILImage.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8))
+
+    def _pil_to_tensor(image):
+        arr = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(arr).unsqueeze(0)
+
+    def _blend_processed_batch_item(
+        result_images,
+        processed_batch,
+        batch_index,
+        blend_fn,
+        x1,
+        y1,
+        ew,
+        eh,
+        tile_mask,
+        padding,
+    ):
+        tile_pil = _tensor_to_pil(processed_batch, batch_index)
+        if tile_pil.size != (ew, eh):
+            tile_pil = tile_pil.resize((ew, eh), PILImage.LANCZOS)
+        result_images[batch_index] = blend_fn(
+            result_images[batch_index],
+            tile_pil,
+            x1,
+            y1,
+            (ew, eh),
+            tile_mask,
+            padding,
+        )
+
+    image_module.tensor_to_pil = _tensor_to_pil
+    image_module.pil_to_tensor = _pil_to_tensor
+    image_module.blend_processed_batch_item = _blend_processed_batch_item
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+
+    def _run_async_in_server_loop(coro, timeout=None):
+        if timeout is not None:
+            return asyncio.run(asyncio.wait_for(coro, timeout=timeout))
+        return asyncio.run(coro)
+
+    async_helpers_module.run_async_in_server_loop = _run_async_in_server_loop
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.get_worker_timeout_seconds = lambda: 60
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.HEARTBEAT_INTERVAL = 10.0
+    constants_module.JOB_POLL_INTERVAL = 0.0
+    constants_module.JOB_POLL_MAX_ATTEMPTS = 3
+    constants_module.MAX_BATCH = 20
+    constants_module.TILE_SEND_TIMEOUT = 1.0
+    constants_module.TILE_WAIT_TIMEOUT = 1.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+
+    async def _noop(*_args, **_kwargs):
+        return None
+
+    job_store_module.ensure_tile_jobs_initialized = lambda: types.SimpleNamespace(
+        distributed_tile_jobs_lock=asyncio.Lock(),
+        distributed_pending_tile_jobs={},
+    )
+    job_store_module.init_static_job_batched = _noop
+    job_store_module.mark_task_completed = _noop
+    job_store_module.cleanup_job = _noop
+    job_store_module.drain_results_queue = _noop
+    job_store_module.get_completed_count = _noop
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
+
+    class _TileJobState:
+        pass
+
+    job_models_module.TileJobState = _TileJobState
+    sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
+
+    processing_args_path = Path(__file__).resolve().parents[3] / "upscale" / "processing_args.py"
+    processing_args_spec = importlib.util.spec_from_file_location(
+        f"{package_name}.upscale.processing_args",
+        processing_args_path,
+    )
+    processing_args_module = importlib.util.module_from_spec(processing_args_spec)
+    assert processing_args_spec is not None and processing_args_spec.loader is not None
+    sys.modules[processing_args_spec.name] = processing_args_module
+    processing_args_spec.loader.exec_module(processing_args_module)
+
+    tile_processing_path = Path(__file__).resolve().parents[3] / "upscale" / "tile_processing.py"
+    tile_processing_spec = importlib.util.spec_from_file_location(
+        f"{package_name}.upscale.tile_processing",
+        tile_processing_path,
+    )
+    tile_processing_module = importlib.util.module_from_spec(tile_processing_spec)
+    assert tile_processing_spec is not None and tile_processing_spec.loader is not None
+    sys.modules[tile_processing_spec.name] = tile_processing_module
+    tile_processing_spec.loader.exec_module(tile_processing_module)
+
+    mode_contexts_path = Path(__file__).resolve().parents[3] / "upscale" / "mode_contexts.py"
+    mode_contexts_spec = importlib.util.spec_from_file_location(
+        f"{package_name}.upscale.mode_contexts",
+        mode_contexts_path,
+    )
+    mode_contexts_module = importlib.util.module_from_spec(mode_contexts_spec)
+    assert mode_contexts_spec is not None and mode_contexts_spec.loader is not None
+    sys.modules[mode_contexts_spec.name] = mode_contexts_module
+    mode_contexts_spec.loader.exec_module(mode_contexts_module)
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.modes.static", 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)
+
+    if created_comfy_stub:
+        sys.modules.pop("comfy.model_management", None)
+        sys.modules.pop("comfy", None)
+
+    return module
+
+
+static_mode = _load_static_mode_module()
+
+
+class _FakeStaticWorker(static_mode.StaticModeMixin):
+    def __init__(self):
+        self.sent_batches = []
+        self.request_calls = 0
+        self.heartbeat_calls = 0
+        self.job_ready = True
+        self.tile_sequence = [(0, 0, True), (None, 0, True)]
+
+    def round_to_multiple(self, value):
+        return value
+
+    def calculate_tiles(self, _width, _height, _tile_width, _tile_height, _force_uniform_tiles):
+        return [(0, 0)]
+
+    def _poll_job_ready(self, *_args, **_kwargs):
+        return self.job_ready
+
+    async def request_assignment(self, *_args, **_kwargs):
+        self.request_calls += 1
+        tile_idx, estimated_remaining, batched_static = self.tile_sequence.pop(0)
+        return SimpleNamespace(
+            kind="tile" if tile_idx is not None else "none",
+            task_idx=tile_idx,
+            estimated_remaining=estimated_remaining,
+            batched_static=batched_static,
+        )
+
+    async def send_heartbeat(self, *_args, **_kwargs):
+        self.heartbeat_calls += 1
+
+    async def send_tiles_batch(
+        self,
+        processed_tiles,
+        _multi_job_id,
+        _master_url,
+        _padding,
+        _worker_id,
+        is_final_flush=False,
+    ):
+        self.sent_batches.append(
+            {
+                "tiles": list(processed_tiles),
+                "is_final_flush": bool(is_final_flush),
+            }
+        )
+
+    def _extract_and_process_tile(self, upscaled_image, *_args, **_kwargs):
+        batch_size = upscaled_image.shape[0]
+        processed_batch = torch.zeros((batch_size, 2, 2, 3), dtype=torch.float32)
+        return processed_batch, 0, 0, 2, 2
+
+    def create_tile_mask(self, *_args, **_kwargs):
+        from PIL import Image
+        return Image.new("L", (4, 4), 255)
+
+    def blend_tile(self, base_image, *_args, **_kwargs):
+        return base_image
+
+
+def _call_worker_static(fake_worker):
+    image = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
+    return fake_worker._process_worker_static_sync(
+        image,
+        model=None,
+        positive=None,
+        negative=None,
+        vae=None,
+        seed=1,
+        steps=1,
+        cfg=1.0,
+        sampler_name="euler",
+        scheduler="normal",
+        denoise=0.5,
+        tile_width=4,
+        tile_height=4,
+        padding=8,
+        mask_blur=4,
+        force_uniform_tiles=True,
+        tiled_decode=False,
+        multi_job_id="job-1",
+        master_url="http://master:8188",
+        worker_id="worker-1",
+        enabled_workers=["worker-1"],
+    )
+
+
+class StaticModeWorkerFlowTests(unittest.TestCase):
+    def test_worker_static_aborts_when_job_not_ready(self):
+        worker = _FakeStaticWorker()
+        worker.job_ready = False
+
+        result = _call_worker_static(worker)
+
+        self.assertEqual(result[0].shape[0], 1)
+        self.assertEqual(worker.request_calls, 0)
+        self.assertEqual(worker.heartbeat_calls, 0)
+        self.assertEqual(worker.sent_batches, [])
+
+    def test_worker_static_requests_tiles_and_flushes_final_batch(self):
+        worker = _FakeStaticWorker()
+
+        _call_worker_static(worker)
+
+        self.assertEqual(worker.request_calls, 2)  # one tile, then sentinel
+        self.assertEqual(worker.heartbeat_calls, 1)
+        self.assertEqual(len(worker.sent_batches), 1)
+        self.assertTrue(worker.sent_batches[0]["is_final_flush"])
+        tiles = worker.sent_batches[0]["tiles"]
+        self.assertEqual(len(tiles), 1)
+        self.assertEqual(tiles[0]["tile_idx"], 0)
+        self.assertEqual(tiles[0]["global_idx"], 0)
+        self.assertEqual(tiles[0]["batch_idx"], 0)
+
+    def test_flush_empty_final_still_sends_completion_signal(self):
+        worker = _FakeStaticWorker()
+
+        returned = worker._flush_tiles_to_master(
+            [],
+            "job-1",
+            "http://master:8188",
+            8,
+            "worker-1",
+            is_final_flush=True,
+        )
+
+        self.assertEqual(returned, [])
+        self.assertEqual(len(worker.sent_batches), 1)
+        self.assertEqual(worker.sent_batches[0]["tiles"], [])
+        self.assertTrue(worker.sent_batches[0]["is_final_flush"])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_tile_ops.py

@@ -0,0 +1,179 @@
+import copy
+import importlib.util
+import sys
+import types
+import unittest
+from contextlib import nullcontext
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def _bootstrap_package(package_name):
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    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
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    def pil_to_tensor(image):
+        array = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(array).unsqueeze(0)
+
+    image_module.tensor_to_pil = tensor_to_pil
+    image_module.pil_to_tensor = pil_to_tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    usdu_utils_module = types.ModuleType(f"{package_name}.utils.usdu_utils")
+
+    def crop_cond(cond, *_args, **_kwargs):
+        return cond
+
+    def get_crop_region(mask, padding):
+        bbox = mask.getbbox()
+        if bbox is None:
+            return (0, 0, mask.size[0], mask.size[1])
+        x1, y1, x2, y2 = bbox
+        return (
+            max(0, x1 - int(padding)),
+            max(0, y1 - int(padding)),
+            min(mask.size[0], x2 + int(padding)),
+            min(mask.size[1], y2 + int(padding)),
+        )
+
+    def expand_crop(region, width, height, target_w, target_h):
+        x1, y1, x2, y2 = region
+        crop_w = x2 - x1
+        crop_h = y2 - y1
+        if crop_w >= target_w and crop_h >= target_h:
+            return (region, (crop_w, crop_h))
+        new_x2 = min(width, x1 + max(crop_w, target_w))
+        new_y2 = min(height, y1 + max(crop_h, target_h))
+        return ((x1, y1, new_x2, new_y2), (new_x2 - x1, new_y2 - y1))
+
+    usdu_utils_module.crop_cond = crop_cond
+    usdu_utils_module.get_crop_region = get_crop_region
+    usdu_utils_module.expand_crop = expand_crop
+    sys.modules[f"{package_name}.utils.usdu_utils"] = usdu_utils_module
+
+    crop_model_patch_module = types.ModuleType(f"{package_name}.utils.crop_model_patch")
+    crop_model_patch_module.crop_model_cond = lambda model, *_args, **_kwargs: nullcontext(model)
+    sys.modules[f"{package_name}.utils.crop_model_patch"] = crop_model_patch_module
+
+    conditioning_module = types.ModuleType(f"{package_name}.upscale.conditioning")
+    conditioning_module.clone_conditioning = lambda cond, clone_hints=False: copy.deepcopy(cond)
+    sys.modules[f"{package_name}.upscale.conditioning"] = conditioning_module
+
+    comfy_module = types.ModuleType("comfy")
+    comfy_samplers = types.ModuleType("comfy.samplers")
+    comfy_model_management = types.ModuleType("comfy.model_management")
+    comfy_module.samplers = comfy_samplers
+    comfy_module.model_management = comfy_model_management
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.samplers"] = comfy_samplers
+    sys.modules["comfy.model_management"] = comfy_model_management
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_tile_ops_module():
+    package_name = "dist_tile_ops_testpkg"
+    _bootstrap_package(package_name)
+    return _load_module(package_name, "upscale/tile_ops.py", "upscale.tile_ops")
+
+
+tile_ops_module = _load_tile_ops_module()
+
+
+class _DummyTileOps(tile_ops_module.TileOpsMixin):
+    pass
+
+
+class _FakeControl:
+    def __init__(self, cond_hint_original, previous_controlnet=None):
+        self.cond_hint_original = cond_hint_original
+        self.previous_controlnet = previous_controlnet
+
+
+class TileOpsTests(unittest.TestCase):
+    def setUp(self):
+        self.node = _DummyTileOps()
+
+    def test_round_and_calculate_tiles(self):
+        self.assertEqual(self.node.round_to_multiple(13, 8), 16)
+        tiles = self.node.calculate_tiles(10, 9, 4, 5, force_uniform_tiles=True)
+        self.assertEqual(len(tiles), 6)
+        self.assertEqual(tiles[0], (0, 0))
+        self.assertEqual(tiles[-1], (8, 5))
+
+    def test_create_tile_mask_and_blend_tile(self):
+        base = Image.new("RGB", (8, 8), (0, 0, 0))
+        tile = Image.new("RGB", (4, 4), (255, 255, 255))
+        mask = self.node.create_tile_mask(8, 8, 2, 2, 4, 4, mask_blur=0)
+
+        self.assertEqual(mask.getpixel((0, 0)), 0)
+        self.assertEqual(mask.getpixel((3, 3)), 255)
+
+        blended = self.node.blend_tile(
+            base_image=base,
+            tile_image=tile,
+            x=2,
+            y=2,
+            extracted_size=(4, 4),
+            mask=mask,
+            padding=0,
+        )
+        self.assertEqual(blended.size, (8, 8))
+        self.assertEqual(blended.getpixel((3, 3)), (255, 255, 255))
+
+    def test_slice_conditioning_selects_batch_index(self):
+        control = _FakeControl(cond_hint_original=torch.ones((2, 1, 1)))
+        positive = [[torch.arange(6, dtype=torch.float32).reshape(2, 3), {"control": control, "mask": torch.ones((2, 1, 1))}]]
+        negative = [[torch.arange(6, dtype=torch.float32).reshape(2, 3), {"mask": torch.zeros((2, 1, 1))}]]
+
+        pos_sliced, neg_sliced = self.node._slice_conditioning(positive, negative, batch_idx=1)
+
+        self.assertEqual(tuple(pos_sliced[0][0].shape), (1, 3))
+        self.assertEqual(tuple(neg_sliced[0][0].shape), (1, 3))
+        self.assertEqual(tuple(pos_sliced[0][1]["mask"].shape), (1, 1, 1))
+        self.assertEqual(tuple(neg_sliced[0][1]["mask"].shape), (1, 1, 1))
+        self.assertEqual(tuple(pos_sliced[0][1]["control"].cond_hint_original.shape), (1, 1, 1))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_tile_processing.py

@@ -0,0 +1,115 @@
+import unittest
+
+from upscale.processing_args import UpscaleCoreArgs
+from upscale.tile_processing import TileBatchArgs, extract_and_process_tile_batch
+
+
+class _DummyTileProcessor:
+    def __init__(self):
+        self.extract_calls = []
+        self.process_calls = []
+
+    def extract_batch_tile_with_padding(
+        self,
+        upscaled_image,
+        tx,
+        ty,
+        tile_width,
+        tile_height,
+        padding,
+        force_uniform_tiles,
+    ):
+        self.extract_calls.append(
+            {
+                "upscaled_image": upscaled_image,
+                "tx": tx,
+                "ty": ty,
+                "tile_width": tile_width,
+                "tile_height": tile_height,
+                "padding": padding,
+                "force_uniform_tiles": force_uniform_tiles,
+            }
+        )
+        return "tile-batch", 10, 20, 30, 40
+
+    def process_tiles_batch(
+        self,
+        tile_batch,
+        model,
+        positive,
+        negative,
+        vae,
+        seed,
+        steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        denoise,
+        tiled_decode,
+        region,
+        canvas_shape,
+    ):
+        self.process_calls.append(
+            {
+                "tile_batch": tile_batch,
+                "model": model,
+                "positive": positive,
+                "negative": negative,
+                "vae": vae,
+                "seed": seed,
+                "steps": steps,
+                "cfg": cfg,
+                "sampler_name": sampler_name,
+                "scheduler": scheduler,
+                "denoise": denoise,
+                "tiled_decode": tiled_decode,
+                "region": region,
+                "canvas_shape": canvas_shape,
+            }
+        )
+        return "processed-batch"
+
+
+class TileProcessingTests(unittest.TestCase):
+    def test_extract_and_process_tile_batch_forwards_args_and_computes_region(self):
+        node = _DummyTileProcessor()
+        tile_args = TileBatchArgs(
+            core=UpscaleCoreArgs(
+                model="model",
+                positive="positive",
+                negative="negative",
+                vae="vae",
+                seed=42,
+                steps=20,
+                cfg=7.5,
+                sampler_name="euler",
+                scheduler="normal",
+                denoise=0.4,
+                tiled_decode=False,
+            ),
+            tile_width=256,
+            tile_height=192,
+            padding=32,
+            force_uniform_tiles=True,
+            width=1024,
+            height=768,
+        )
+        result = extract_and_process_tile_batch(
+            node=node,
+            upscaled_image="upscaled-image",
+            tx=3,
+            ty=4,
+            args=tile_args,
+        )
+
+        self.assertEqual(result, ("processed-batch", 10, 20, 30, 40))
+        self.assertEqual(len(node.extract_calls), 1)
+        self.assertEqual(len(node.process_calls), 1)
+        self.assertEqual(node.extract_calls[0]["tx"], 3)
+        self.assertEqual(node.extract_calls[0]["ty"], 4)
+        self.assertEqual(node.process_calls[0]["region"], (10, 20, 40, 60))
+        self.assertEqual(node.process_calls[0]["canvas_shape"], (1024, 768))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/upscale/test_usdu_utils.py

@@ -0,0 +1,128 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+from PIL import Image
+
+
+def _ensure_torch_stubs():
+    if "torch" not in sys.modules:
+        torch_module = types.ModuleType("torch")
+
+        class _Tensor:
+            pass
+
+        torch_module.Tensor = _Tensor
+        torch_module.float32 = object()
+        torch_module.zeros = lambda *_args, **_kwargs: None
+        torch_module.from_numpy = lambda array: array
+        torch_module.cat = lambda *_args, **_kwargs: None
+
+        nn_module = types.ModuleType("torch.nn")
+        functional_module = types.ModuleType("torch.nn.functional")
+        functional_module.interpolate = lambda tensor, **_kwargs: tensor
+        nn_module.functional = functional_module
+
+        torch_module.nn = nn_module
+
+        sys.modules["torch"] = torch_module
+        sys.modules["torch.nn"] = nn_module
+        sys.modules["torch.nn.functional"] = functional_module
+
+    if "torchvision" not in sys.modules:
+        torchvision_module = types.ModuleType("torchvision")
+        transforms_module = types.ModuleType("torchvision.transforms")
+
+        class _GaussianBlur:
+            def __init__(self, *_args, **_kwargs):
+                pass
+
+            def __call__(self, value):
+                return value
+
+        transforms_module.GaussianBlur = _GaussianBlur
+        torchvision_module.transforms = transforms_module
+
+        sys.modules["torchvision"] = torchvision_module
+        sys.modules["torchvision.transforms"] = transforms_module
+
+
+def _load_usdu_utils_module():
+    _ensure_torch_stubs()
+    module_path = Path(__file__).resolve().parents[3] / "utils" / "usdu_utils.py"
+    spec = importlib.util.spec_from_file_location("dist_test_usdu_utils", 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
+
+
+usdu_utils = _load_usdu_utils_module()
+
+
+class UsduUtilsGeometryTests(unittest.TestCase):
+    def test_fix_crop_region_decrements_upper_bounds_when_interior(self):
+        self.assertEqual(
+            usdu_utils.fix_crop_region((2, 3, 9, 9), (10, 10)),
+            (2, 3, 8, 8),
+        )
+        self.assertEqual(
+            usdu_utils.fix_crop_region((1, 1, 10, 10), (10, 10)),
+            (1, 1, 10, 10),
+        )
+
+    def test_get_crop_region_with_padding_and_empty_mask(self):
+        mask = Image.new("L", (10, 10), 0)
+        for x in range(2, 5):
+            for y in range(3, 7):
+                mask.putpixel((x, y), 255)
+
+        region = usdu_utils.get_crop_region(mask, pad=1)
+        self.assertEqual(region, (1, 2, 5, 7))
+
+        empty_region = usdu_utils.get_crop_region(Image.new("L", (8, 6), 0), pad=2)
+        self.assertEqual(empty_region, (6, 4, 1, 1))
+
+    def test_expand_crop_and_resize_region(self):
+        expanded_region, expanded_size = usdu_utils.expand_crop(
+            region=(2, 2, 6, 6),
+            width=12,
+            height=12,
+            target_width=8,
+            target_height=10,
+        )
+        self.assertEqual(expanded_region, (0, 0, 8, 10))
+        self.assertEqual(expanded_size, (8, 10))
+
+        resized = usdu_utils.resize_region(
+            region=(2, 2, 6, 10),
+            init_size=(8, 16),
+            resize_size=(16, 32),
+        )
+        self.assertEqual(resized, (4, 4, 12, 20))
+
+    def test_region_intersection(self):
+        self.assertEqual(
+            usdu_utils.region_intersection((0, 0, 5, 5), (3, 2, 8, 7)),
+            (3, 2, 5, 5),
+        )
+        self.assertIsNone(usdu_utils.region_intersection((0, 0, 1, 1), (2, 2, 4, 4)))
+
+    def test_pad_image2_preserves_expected_output_size(self):
+        image = Image.new("RGB", (4, 3), "black")
+        padded = usdu_utils.pad_image2(
+            image,
+            left_pad=2,
+            right_pad=1,
+            top_pad=3,
+            bottom_pad=2,
+            fill=False,
+            blur=False,
+        )
+        self.assertEqual(padded.size, (7, 8))
+
+
+if __name__ == "__main__":
+    unittest.main()

upscale/conditioning.py

@@ -0,0 +1,36 @@
+import copy
+from typing import Any
+
+
+def clone_control_chain(control: Any, clone_hint: bool = True) -> Any:
+    """Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
+    if control is None:
+        return None
+    new_control = copy.copy(control)
+    if clone_hint and hasattr(control, 'cond_hint_original'):
+        hint = getattr(control, 'cond_hint_original', None)
+        new_control.cond_hint_original = hint.clone() if hint is not None else None
+    if hasattr(control, 'previous_controlnet'):
+        new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
+    return new_control
+
+
+def clone_conditioning(
+    cond_list: list[tuple[Any, dict[str, Any]]] | list[list[Any]],
+    clone_hints: bool = True,
+) -> list[list[Any]]:
+    """Clone conditioning without duplicating ControlNet models."""
+    new_cond = []
+    for emb, cond_dict in cond_list:
+        new_emb = emb.clone() if emb is not None else None
+        new_dict = cond_dict.copy()
+        if 'control' in new_dict:
+            new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
+        if 'mask' in new_dict and new_dict['mask'] is not None:
+            new_dict['mask'] = new_dict['mask'].clone()
+        if 'pooled_output' in new_dict and new_dict['pooled_output'] is not None:
+            new_dict['pooled_output'] = new_dict['pooled_output'].clone()
+        if 'area' in new_dict:
+            new_dict['area'] = new_dict['area'][:]
+        new_cond.append([new_emb, new_dict])
+    return new_cond

upscale/job_models.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass, field
+import asyncio
+import time
+
+
+class BaseJobState:
+    """Marker base class for typed USDU job state containers."""
+
+
+@dataclass
+class TileJobState(BaseJobState):
+    """Typed state container for static (tile) USDU jobs."""
+
+    multi_job_id: str
+    mode: str = field(default="static", init=False)
+    queue: asyncio.Queue = field(default_factory=asyncio.Queue)
+    pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
+    completed_tasks: dict = field(default_factory=dict)
+    worker_status: dict = field(default_factory=dict)
+    assigned_to_workers: dict = field(default_factory=dict)
+    batch_size: int = 0
+    num_tiles_per_image: int = 0
+    batched_static: bool = False
+    created_at: float = field(default_factory=time.monotonic)
+
+
+@dataclass
+class ImageJobState(BaseJobState):
+    """Typed state container for dynamic (per-image) USDU jobs."""
+
+    multi_job_id: str
+    mode: str = field(default="dynamic", init=False)
+    queue: asyncio.Queue = field(default_factory=asyncio.Queue)
+    pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
+    completed_images: dict = field(default_factory=dict)
+    worker_status: dict = field(default_factory=dict)
+    assigned_to_workers: dict = field(default_factory=dict)
+    batch_size: int = 0
+    num_tiles_per_image: int = 0
+    batched_static: bool = False
+    created_at: float = field(default_factory=time.monotonic)
+
+    @property
+    def pending_tasks(self):
+        return self.pending_images
+
+    @property
+    def completed_tasks(self):
+        return self.completed_images

upscale/job_state.py

@@ -0,0 +1,167 @@
+import asyncio
+
+from ..utils.logging import debug_log
+from .job_store import ensure_tile_jobs_initialized
+from .job_timeout import check_and_requeue_timed_out_workers
+from .job_models import ImageJobState, TileJobState
+
+
+class JobStateMixin:
+    async def _get_job_data(self, multi_job_id):
+        """Return current job data reference while holding lock briefly."""
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            return prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+
+    async def _get_all_completed_tasks(self, multi_job_id):
+        """Helper to retrieve all completed tasks from the job data."""
+        job_data = await self._get_job_data(multi_job_id)
+        if isinstance(job_data, TileJobState):
+            return dict(job_data.completed_tasks)
+        if isinstance(job_data, ImageJobState):
+            return dict(job_data.completed_images)
+        return {}
+
+    async def all_completed_tasks(self, multi_job_id):
+        """Public completed-task accessor."""
+        return await self._get_all_completed_tasks(multi_job_id)
+
+    async def _get_next_image_index(self, multi_job_id):
+        """Get next image index from pending queue for master."""
+        prompt_server = ensure_tile_jobs_initialized()
+        pending_queue = None
+        async with prompt_server.distributed_tile_jobs_lock:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, ImageJobState):
+                pending_queue = job_data.pending_images
+
+        if pending_queue is None:
+            return None
+
+        try:
+            return await asyncio.wait_for(pending_queue.get(), timeout=1.0)
+        except asyncio.TimeoutError:
+            return None
+
+    async def next_image_index(self, multi_job_id):
+        """Public next-image accessor."""
+        return await self._get_next_image_index(multi_job_id)
+
+    async def _get_next_tile_index(self, multi_job_id):
+        """Get next tile index from pending queue for master in static mode."""
+        prompt_server = ensure_tile_jobs_initialized()
+        pending_queue = None
+        async with prompt_server.distributed_tile_jobs_lock:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, TileJobState):
+                pending_queue = job_data.pending_tasks
+
+        if pending_queue is None:
+            return None
+
+        try:
+            return await asyncio.wait_for(pending_queue.get(), timeout=0.1)
+        except asyncio.TimeoutError:
+            return None
+
+    async def next_tile_index(self, multi_job_id):
+        """Public next-tile accessor."""
+        return await self._get_next_tile_index(multi_job_id)
+
+    async def _get_total_completed_count(self, multi_job_id):
+        """Get total count of all completed images (master + workers)."""
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, ImageJobState):
+                return len(job_data.completed_images)
+            if isinstance(job_data, TileJobState):
+                return len(job_data.completed_tasks)
+            return 0
+
+    async def total_completed_count(self, multi_job_id):
+        """Public completed-count accessor."""
+        return await self._get_total_completed_count(multi_job_id)
+
+    async def _get_all_completed_images(self, multi_job_id):
+        """Get all completed images."""
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, ImageJobState):
+                return job_data.completed_images.copy()
+            return {}
+
+    async def all_completed_images(self, multi_job_id):
+        """Public completed-image accessor."""
+        return await self._get_all_completed_images(multi_job_id)
+
+    async def _get_pending_count(self, multi_job_id):
+        """Get count of pending images in the queue."""
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, ImageJobState):
+                return job_data.pending_images.qsize()
+            if isinstance(job_data, TileJobState):
+                return job_data.pending_tasks.qsize()
+            return 0
+
+    async def pending_count(self, multi_job_id):
+        """Public pending-count accessor."""
+        return await self._get_pending_count(multi_job_id)
+
+    async def _drain_worker_results_queue(self, multi_job_id):
+        """Drain pending worker results from queue and update completed images."""
+        prompt_server = ensure_tile_jobs_initialized()
+        worker_queue = None
+        async with prompt_server.distributed_tile_jobs_lock:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, ImageJobState):
+                worker_queue = job_data.queue
+
+        if worker_queue is None:
+            return 0
+
+        drained_results = []
+        while True:
+            try:
+                drained_results.append(worker_queue.get_nowait())
+            except asyncio.QueueEmpty:
+                break
+
+        if not drained_results:
+            return 0
+
+        collected = 0
+        async with prompt_server.distributed_tile_jobs_lock:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if not isinstance(job_data, ImageJobState):
+                return 0
+
+            for result in drained_results:
+                worker_id = result.get("worker_id")
+                if "image_idx" in result and "image" in result:
+                    image_idx = result["image_idx"]
+                    image_pil = result["image"]
+                    if image_idx not in job_data.completed_images:
+                        job_data.completed_images[image_idx] = image_pil
+                        collected += 1
+                        debug_log(f"Drained image {image_idx} from worker {worker_id}")
+
+        if collected > 0:
+            debug_log(f"Drained {collected} worker images during retry")
+
+        return collected
+
+    async def drain_worker_results_queue(self, multi_job_id):
+        """Public worker-result drain API."""
+        return await self._drain_worker_results_queue(multi_job_id)
+
+    async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
+        """Check for timed out workers and requeue their assigned images."""
+        return await check_and_requeue_timed_out_workers(multi_job_id, batch_size)
+
+    async def check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
+        """Public timeout-requeue API."""
+        return await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size)

upscale/job_store.py

@@ -0,0 +1,199 @@
+import asyncio
+import os
+import time
+from dataclasses import dataclass
+from typing import Any
+
+from ..utils.logging import debug_log
+from ..utils.runtime_state import ensure_distributed_runtime_state, get_prompt_server_instance
+from .job_models import BaseJobState, ImageJobState, TileJobState
+
+# Configure maximum payload size (50MB default, configurable via environment variable)
+MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
+
+
+@dataclass(frozen=True)
+class JobQueueInitConfig:
+    mode: str
+    batch_size: int = 0
+    num_tiles_per_image: int = 0
+    all_indices: list[int] | None = None
+    enabled_workers: list[str] | None = None
+    batched_static: bool = False
+
+
+def ensure_tile_jobs_initialized() -> Any:
+    """Ensure tile job storage is initialized on the server instance."""
+    prompt_server = get_prompt_server_instance()
+    state = ensure_distributed_runtime_state(prompt_server)
+    if not isinstance(state.distributed_pending_tile_jobs, dict):
+        debug_log("Resetting invalid distributed_pending_tile_jobs state to empty dict.")
+        state.distributed_pending_tile_jobs = {}
+        ensure_distributed_runtime_state(prompt_server)
+
+    invalid_job_ids = [
+        job_id
+        for job_id, job_data in state.distributed_pending_tile_jobs.items()
+        if not isinstance(job_data, BaseJobState)
+    ]
+    for job_id in invalid_job_ids:
+        debug_log(f"Removing invalid job state for {job_id}")
+        del state.distributed_pending_tile_jobs[job_id]
+    return prompt_server
+
+
+async def _init_job_queue(
+    multi_job_id: str,
+    config: JobQueueInitConfig,
+) -> None:
+    """Unified initialization for job queues in static and dynamic modes."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+            debug_log(f"Queue already exists for {multi_job_id}")
+            return
+
+        if config.mode == 'dynamic':
+            job_data = ImageJobState(multi_job_id=multi_job_id)
+        elif config.mode == 'static':
+            job_data = TileJobState(multi_job_id=multi_job_id)
+        else:
+            raise ValueError(f"Unknown mode: {config.mode}")
+
+        job_data.worker_status = {w: time.time() for w in config.enabled_workers or []}
+        job_data.assigned_to_workers = {w: [] for w in config.enabled_workers or []}
+
+        if config.mode == 'dynamic':
+            job_data.batch_size = int(config.batch_size or 0)
+            pending_queue = job_data.pending_images
+            indices = config.all_indices or list(range(int(config.batch_size or 0)))
+            for i in indices:
+                await pending_queue.put(i)
+            debug_log(f"Initialized image queue with {config.batch_size} pending items")
+        elif config.mode == 'static':
+            job_data.num_tiles_per_image = int(config.num_tiles_per_image or 0)
+            job_data.batch_size = int(config.batch_size or 0)
+            job_data.batched_static = bool(config.batched_static)
+            # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
+            pending_queue = job_data.pending_tasks
+            if config.batched_static and config.num_tiles_per_image > 0:
+                for i in range(config.num_tiles_per_image):
+                    await pending_queue.put(i)
+            else:
+                total_tiles = int(config.batch_size or 0) * int(config.num_tiles_per_image or 0)
+                for i in range(total_tiles):
+                    await pending_queue.put(i)
+
+        prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
+
+
+async def init_dynamic_job(
+    multi_job_id: str,
+    batch_size: int,
+    enabled_workers: list[str],
+    all_indices: list[int] | None = None,
+) -> None:
+    """Initialize queue for dynamic mode (per-image), with collector fields."""
+    await _init_job_queue(
+        multi_job_id,
+        JobQueueInitConfig(
+            mode='dynamic',
+            batch_size=batch_size,
+            all_indices=all_indices or list(range(batch_size)),
+            enabled_workers=enabled_workers,
+        ),
+    )
+    debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
+
+
+async def init_static_job_batched(
+    multi_job_id: str,
+    batch_size: int,
+    num_tiles_per_image: int,
+    enabled_workers: list[str],
+) -> None:
+    """Initialize queue for static mode (batched-per-tile)."""
+    await _init_job_queue(
+        multi_job_id,
+        JobQueueInitConfig(
+            mode='static',
+            batch_size=batch_size,
+            num_tiles_per_image=num_tiles_per_image,
+            enabled_workers=enabled_workers,
+            batched_static=True,
+        ),
+    )
+
+
+async def drain_results_queue(multi_job_id: str) -> int:
+    """Drain pending results from queue and update completed_tasks. Returns count drained."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if not isinstance(job_data, BaseJobState):
+            return 0
+        q = job_data.queue
+        completed_tasks = job_data.completed_tasks
+
+        collected = 0
+        while True:
+            try:
+                result = q.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+
+            worker_id = result['worker_id']
+            is_last = result.get('is_last', False)
+
+            if 'image_idx' in result and 'image' in result:
+                task_id = result['image_idx']
+                if task_id not in completed_tasks:
+                    completed_tasks[task_id] = result['image']
+                    collected += 1
+            elif 'tiles' in result:
+                for tile_data in result['tiles']:
+                    task_id = tile_data.get('global_idx', tile_data['tile_idx'])
+                    if task_id not in completed_tasks:
+                        completed_tasks[task_id] = tile_data
+                        collected += 1
+            if is_last:
+                if worker_id in job_data.worker_status:
+                    del job_data.worker_status[worker_id]
+
+        return collected
+
+
+async def get_completed_count(multi_job_id: str) -> int:
+    """Get count of completed tasks."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if isinstance(job_data, BaseJobState):
+            return len(job_data.completed_tasks)
+        return 0
+
+
+async def mark_task_completed(multi_job_id: str, task_id: int, result: Any) -> None:
+    """Mark a task as completed."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if isinstance(job_data, BaseJobState):
+            job_data.completed_tasks[task_id] = result
+
+
+async def cleanup_job(multi_job_id: str) -> None:
+    """Cleanup the job data."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+            del prompt_server.distributed_pending_tile_jobs[multi_job_id]
+            debug_log(f"Cleaned up job {multi_job_id}")
+
+
+async def init_job_queue(
+    multi_job_id: str,
+    config: JobQueueInitConfig,
+) -> None:
+    """Public entry point for unified job queue initialization."""
+    await _init_job_queue(multi_job_id, config)

upscale/job_timeout.py

@@ -0,0 +1,159 @@
+import time
+
+from ..utils.config import load_config
+from ..utils.constants import HEARTBEAT_TIMEOUT
+from ..utils.logging import debug_log, log
+from ..utils.network import build_worker_url, probe_worker
+from .job_models import BaseJobState
+from .job_store import ensure_tile_jobs_initialized
+
+
+def _find_worker_record(worker_id):
+    """Return worker config entry by id, or None when missing."""
+    workers = load_config().get("workers", [])
+    return next((w for w in workers if str(w.get("id")) == str(worker_id)), None)
+
+
+async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
+    """Check timed out workers and requeue their tasks. Returns requeued count."""
+    prompt_server = ensure_tile_jobs_initialized()
+    current_time = time.time()
+
+    # Allow override via config setting 'worker_timeout_seconds'
+    cfg = load_config()
+    hb_timeout = int(cfg.get("settings", {}).get("worker_timeout_seconds", HEARTBEAT_TIMEOUT))
+
+    # Snapshot timed-out workers and job details under lock.
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if not isinstance(job_data, BaseJobState):
+            return 0
+
+        completed_tasks_snapshot = set(job_data.completed_tasks.keys())
+        batched_static_snapshot = bool(job_data.batched_static)
+        num_tiles_per_image_snapshot = int(job_data.num_tiles_per_image or 1)
+        batch_size_snapshot = int(job_data.batch_size or 1)
+
+        timed_out_workers = []
+        for worker, last_heartbeat in list(job_data.worker_status.items()):
+            age = current_time - float(last_heartbeat)
+            debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
+            if age > hb_timeout:
+                timed_out_workers.append(
+                    {
+                        "worker_id": worker,
+                        "last_heartbeat": float(last_heartbeat),
+                        "assigned_tasks": list(job_data.assigned_to_workers.get(worker, [])),
+                    }
+                )
+
+    if not timed_out_workers:
+        return 0
+
+    # Probe outside lock to avoid lock contention on network latency.
+    workers_to_requeue = []
+    workers_graced = []
+    for worker_info in timed_out_workers:
+        worker = worker_info["worker_id"]
+        assigned = worker_info["assigned_tasks"]
+        age = current_time - worker_info["last_heartbeat"]
+
+        incomplete_assigned = 0
+        try:
+            if assigned:
+                if batched_static_snapshot:
+                    for task_id in assigned:
+                        for b in range(batch_size_snapshot):
+                            gidx = b * num_tiles_per_image_snapshot + task_id
+                            if gidx not in completed_tasks_snapshot:
+                                incomplete_assigned += 1
+                                break
+                else:
+                    for task_id in assigned:
+                        if task_id not in completed_tasks_snapshot:
+                            incomplete_assigned += 1
+            debug_log(
+                f"Assigned diagnostics: total_assigned={len(assigned)} "
+                f"incomplete_assigned={incomplete_assigned}"
+            )
+        except Exception as e:
+            debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
+
+        busy = False
+        probe_queue = None
+        try:
+            worker_record = _find_worker_record(worker)
+            if worker_record:
+                worker_url = build_worker_url(worker_record)
+                debug_log(f"Probing worker {worker} at {worker_url}/prompt")
+                payload = await probe_worker(worker_url, timeout=2.0)
+                if payload is not None:
+                    probe_queue = int(payload.get("exec_info", {}).get("queue_remaining", 0))
+                    busy = probe_queue is not None and probe_queue > 0
+            else:
+                debug_log(f"Probe skipped; worker {worker} not found in config")
+        except Exception as e:
+            debug_log(f"Probe failed for worker {worker}: {e}")
+        finally:
+            debug_log(
+                f"Probe diagnostics: online={probe_queue is not None} queue_remaining={probe_queue}"
+            )
+
+        if incomplete_assigned == 0:
+            debug_log(f"Worker {worker} heartbeat stale but all tasks complete; skipping")
+            continue
+
+        if busy:
+            workers_graced.append(worker)
+            debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
+            continue
+
+        log(f"Worker {worker} heartbeat timed out after {age:.1f}s")
+        workers_to_requeue.append((worker, assigned))
+
+    # Re-acquire lock and apply requeue/cleanup decisions.
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if not isinstance(job_data, BaseJobState):
+            return 0
+
+        # Refresh heartbeat for workers that we proved are still busy.
+        for worker in workers_graced:
+            if worker in job_data.worker_status:
+                job_data.worker_status[worker] = current_time
+
+        requeued_count = 0
+        completed_tasks = job_data.completed_tasks
+        batched_static = bool(job_data.batched_static)
+        num_tiles_per_image = int(job_data.num_tiles_per_image or 1)
+        batch_size = int(job_data.batch_size or 1)
+        for worker, assigned_snapshot in workers_to_requeue:
+            # Use current assignments if present, falling back to the snapshot.
+            assigned_tasks = list(job_data.assigned_to_workers.get(worker, assigned_snapshot))
+            for task_id in assigned_tasks:
+                # If batched_static, task_id is a tile_idx; consider it complete only if
+                # all corresponding global_idx entries are present in completed_tasks.
+                if batched_static:
+                    all_done = True
+                    for b in range(batch_size):
+                        gidx = b * num_tiles_per_image + task_id
+                        if gidx not in completed_tasks:
+                            all_done = False
+                            break
+                    if not all_done:
+                        await job_data.pending_tasks.put(task_id)
+                        requeued_count += 1
+                else:
+                    if task_id not in completed_tasks:
+                        await job_data.pending_tasks.put(task_id)
+                        requeued_count += 1
+            job_data.worker_status.pop(worker, None)
+            if worker in job_data.assigned_to_workers:
+                job_data.assigned_to_workers[worker] = []
+
+        return requeued_count
+
+
+async def check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
+    """Public wrapper for worker-timeout requeue checks."""
+    return await _check_and_requeue_timed_out_workers(multi_job_id, total_tasks)

upscale/mode_contexts.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+class TileOpsCollaborator:
+    """Typed facade over tile-related operations used by USDU modes."""
+
+    def __init__(self, delegate: Any) -> None:
+        self._delegate = delegate
+
+    def round_to_multiple(self, value: int) -> int:
+        return self._delegate.round_to_multiple(value)
+
+    def calculate_tiles(
+        self,
+        width: int,
+        height: int,
+        tile_width: int,
+        tile_height: int,
+        force_uniform_tiles: bool,
+    ) -> list[tuple[int, int]]:
+        return self._delegate.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+
+    def create_tile_mask(
+        self,
+        width: int,
+        height: int,
+        tx: int,
+        ty: int,
+        tile_width: int,
+        tile_height: int,
+        mask_blur: int,
+    ) -> Any:
+        return self._delegate.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)
+
+    def blend_tile(
+        self,
+        base_image: Any,
+        tile_pil: Any,
+        x1: int,
+        y1: int,
+        size: tuple[int, int],
+        tile_mask: Any,
+        padding: int,
+    ) -> Any:
+        return self._delegate.blend_tile(base_image, tile_pil, x1, y1, size, tile_mask, padding)
+
+    def slice_conditioning(self, positive: Any, negative: Any, image_idx: int) -> tuple[Any, Any]:
+        return self._delegate.slice_conditioning(positive, negative, image_idx)
+
+    def process_and_blend_tile(
+        self,
+        tile_idx: int,
+        tile_pos: tuple[int, int],
+        upscaled_image: Any,
+        result_image: Any,
+        model: Any,
+        positive: Any,
+        negative: Any,
+        vae: Any,
+        seed: int,
+        steps: int,
+        cfg: float,
+        sampler_name: str,
+        scheduler: str,
+        denoise: float,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        image_width: int,
+        image_height: int,
+        force_uniform_tiles: bool,
+        tiled_decode: bool,
+        batch_idx: int,
+    ) -> Any:
+        return self._delegate.process_and_blend_tile(
+            tile_idx,
+            tile_pos,
+            upscaled_image,
+            result_image,
+            model,
+            positive,
+            negative,
+            vae,
+            seed,
+            steps,
+            cfg,
+            sampler_name,
+            scheduler,
+            denoise,
+            tile_width,
+            tile_height,
+            padding,
+            mask_blur,
+            image_width,
+            image_height,
+            force_uniform_tiles,
+            tiled_decode,
+            batch_idx=batch_idx,
+        )
+
+
+class JobStateCollaborator:
+    """Typed facade over job-state operations used by USDU modes."""
+
+    def __init__(self, delegate: Any) -> None:
+        self._delegate = delegate
+
+    async def next_tile_index(self, multi_job_id: str) -> int | None:
+        return await self._delegate.next_tile_index(multi_job_id)
+
+    async def all_completed_tasks(self, multi_job_id: str) -> dict[int, dict[str, Any]]:
+        return await self._delegate.all_completed_tasks(multi_job_id)
+
+    async def check_and_requeue_timed_out_workers(self, multi_job_id: str, expected_total: int) -> int:
+        return await self._delegate.check_and_requeue_timed_out_workers(multi_job_id, expected_total)
+
+    async def next_image_index(self, multi_job_id: str) -> int | None:
+        return await self._delegate.next_image_index(multi_job_id)
+
+    async def drain_worker_results_queue(self, multi_job_id: str) -> int:
+        return await self._delegate.drain_worker_results_queue(multi_job_id)
+
+    async def total_completed_count(self, multi_job_id: str) -> int:
+        return await self._delegate.total_completed_count(multi_job_id)
+
+    async def pending_count(self, multi_job_id: str) -> int:
+        return await self._delegate.pending_count(multi_job_id)
+
+    async def all_completed_images(self, multi_job_id: str) -> dict[int, Any]:
+        return await self._delegate.all_completed_images(multi_job_id)
+
+
+class WorkerCommsCollaborator:
+    """Typed facade over worker/master communication operations used by USDU modes."""
+
+    def __init__(self, delegate: Any) -> None:
+        self._delegate = delegate
+
+    async def request_assignment(self, multi_job_id: str, master_url: str, worker_id: str) -> Any:
+        return await self._delegate.request_assignment(multi_job_id, master_url, worker_id)
+
+    async def send_tiles_batch(
+        self,
+        processed_tiles: list[dict[str, Any]],
+        multi_job_id: str,
+        master_url: str,
+        padding: int,
+        worker_id: str,
+        is_final_flush: bool = False,
+    ) -> None:
+        await self._delegate.send_tiles_batch(
+            processed_tiles,
+            multi_job_id,
+            master_url,
+            padding,
+            worker_id,
+            is_final_flush=is_final_flush,
+        )
+
+    async def send_heartbeat(self, multi_job_id: str, master_url: str, worker_id: str) -> None:
+        await self._delegate.send_heartbeat(multi_job_id, master_url, worker_id)
+
+    async def check_job_status(self, multi_job_id: str, master_url: str) -> bool:
+        return await self._delegate.check_job_status(multi_job_id, master_url)
+
+    async def async_yield(self) -> None:
+        await self._delegate.async_yield()
+
+    async def send_full_image(
+        self,
+        image_pil: Any,
+        image_idx: int,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+        is_last: bool,
+    ) -> None:
+        await self._delegate.send_full_image(
+            image_pil,
+            image_idx,
+            multi_job_id,
+            master_url,
+            worker_id,
+            is_last,
+        )
+
+    async def send_worker_complete_signal(self, multi_job_id: str, master_url: str, worker_id: str) -> None:
+        await self._delegate.send_worker_complete_signal(multi_job_id, master_url, worker_id)
+
+
+class ResultCollectorCollaborator:
+    """Typed facade over result-collection operations used by dynamic mode."""
+
+    def __init__(self, delegate: Any) -> None:
+        self._delegate = delegate
+
+    async def mark_image_completed(self, multi_job_id: str, image_idx: int, image_pil: Any) -> None:
+        await self._delegate.mark_image_completed(multi_job_id, image_idx, image_pil)
+
+    async def collect_dynamic_images(
+        self,
+        multi_job_id: str,
+        remaining_to_collect: int,
+        num_workers: int,
+        batch_size: int,
+        master_processed_count: int,
+    ) -> dict[int, Any]:
+        return await self._delegate.collect_dynamic_images(
+            multi_job_id,
+            remaining_to_collect,
+            num_workers,
+            batch_size,
+            master_processed_count,
+        )
+
+
+@dataclass(frozen=True)
+class SingleGpuModeContext:
+    tile_ops: TileOpsCollaborator
+
+
+@dataclass(frozen=True)
+class StaticModeContext:
+    tile_ops: TileOpsCollaborator
+    job_state: JobStateCollaborator
+    worker_comms: WorkerCommsCollaborator
+
+
+@dataclass(frozen=True)
+class DynamicModeContext:
+    tile_ops: TileOpsCollaborator
+    job_state: JobStateCollaborator
+    worker_comms: WorkerCommsCollaborator
+    result_collector: ResultCollectorCollaborator

upscale/modes/dynamic.py

@@ -0,0 +1,472 @@
+from __future__ import annotations
+
+import asyncio, time, torch
+from typing import Any
+from PIL import Image
+from ...utils.logging import debug_log, log
+from ...utils.image import tensor_to_pil, pil_to_tensor
+from ...utils.async_helpers import run_async_in_server_loop
+from ...utils.config import get_worker_timeout_seconds
+from ...utils.constants import (
+    JOB_POLL_INTERVAL,
+    JOB_POLL_MAX_ATTEMPTS,
+    TILE_WAIT_TIMEOUT,
+    TILE_SEND_TIMEOUT,
+)
+from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job
+from ..mode_contexts import (
+    DynamicModeContext,
+    JobStateCollaborator,
+    ResultCollectorCollaborator,
+    TileOpsCollaborator,
+    WorkerCommsCollaborator,
+)
+from ..processing_args import UpscaleCoreArgs
+
+
+class DynamicModeMixin:
+    """
+    Dynamic (per-image queue) USDU mode behaviors for master and worker roles.
+
+    Expected co-mixins on `self`:
+    - TileOpsMixin (`calculate_tiles`, `slice_conditioning`, `process_and_blend_tile`).
+    - JobStateMixin (image queue/task completion helpers).
+    - WorkerCommsMixin (`request_assignment`, `send_full_image`, `send_heartbeat`).
+    """
+
+    def _build_dynamic_mode_context(self) -> DynamicModeContext:
+        """Build explicit collaborators for dynamic mode execution."""
+        return DynamicModeContext(
+            tile_ops=TileOpsCollaborator(self),
+            job_state=JobStateCollaborator(self),
+            worker_comms=WorkerCommsCollaborator(self),
+            result_collector=ResultCollectorCollaborator(self),
+        )
+
+    def _poll_job_ready_dynamic(
+        self,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str | None = None,
+        max_attempts: int = JOB_POLL_MAX_ATTEMPTS,
+        mode_context: DynamicModeContext | None = None,
+    ) -> bool:
+        """Poll master for job readiness using explicit worker-comms collaborator."""
+        context = mode_context or self._build_dynamic_mode_context()
+        for attempt in range(max_attempts):
+            ready = run_async_in_server_loop(
+                context.worker_comms.check_job_status(multi_job_id, master_url),
+                timeout=5.0,
+            )
+            if ready:
+                if worker_id:
+                    debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
+                else:
+                    debug_log(f"Job {multi_job_id} ready after {attempt} attempts")
+                return True
+            time.sleep(JOB_POLL_INTERVAL)
+        return False
+
+    def _process_image_tiles(
+        self,
+        image_idx: int,
+        local_image: Image.Image,
+        single_tensor: torch.Tensor,
+        core_args: UpscaleCoreArgs,
+        all_tiles: list,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        width: int,
+        height: int,
+        force_uniform_tiles: bool,
+        context: DynamicModeContext,
+        per_tile_callback: Any | None = None,
+    ) -> Image.Image:
+        """Process all tiles for a single image. Returns the processed PIL image."""
+        positive_sliced, negative_sliced = context.tile_ops.slice_conditioning(
+            core_args.positive, core_args.negative, image_idx,
+        )
+        for tile_idx, pos in enumerate(all_tiles):
+            source_tensor = pil_to_tensor(local_image)
+            if single_tensor.is_cuda:
+                source_tensor = source_tensor.cuda()
+            local_image = context.tile_ops.process_and_blend_tile(
+                tile_idx, pos, source_tensor, local_image,
+                core_args.model,
+                positive_sliced,
+                negative_sliced,
+                core_args.vae,
+                core_args.seed,
+                core_args.steps,
+                core_args.cfg,
+                core_args.sampler_name,
+                core_args.scheduler,
+                core_args.denoise,
+                tile_width,
+                tile_height,
+                padding, mask_blur, width, height, force_uniform_tiles,
+                core_args.tiled_decode,
+                batch_idx=image_idx,
+            )
+            if per_tile_callback:
+                per_tile_callback()
+        return local_image
+
+    def _handle_master_dynamic_idle_state(
+        self,
+        multi_job_id,
+        batch_size,
+        consecutive_retries,
+        max_consecutive_retries,
+        mode_context: DynamicModeContext | None = None,
+    ):
+        """Handle queue-empty state while master waits for worker progress."""
+        context = mode_context or self._build_dynamic_mode_context()
+        drained_count = run_async_in_server_loop(
+            context.job_state.drain_worker_results_queue(multi_job_id),
+            timeout=5.0,
+        )
+        run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
+
+        requeued_count = run_async_in_server_loop(
+            context.job_state.check_and_requeue_timed_out_workers(multi_job_id, batch_size),
+            timeout=5.0,
+        )
+        run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
+
+        if requeued_count > 0:
+            log(f"Requeued {requeued_count} images from timed out workers")
+            return False, True, 0
+
+        completed_now = run_async_in_server_loop(
+            context.job_state.total_completed_count(multi_job_id),
+            timeout=1.0,
+        )
+        log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
+        if completed_now >= batch_size:
+            return True, False, consecutive_retries
+
+        run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
+
+        pending_count = run_async_in_server_loop(
+            context.job_state.pending_count(multi_job_id),
+            timeout=1.0,
+        )
+        if pending_count > 0:
+            return False, True, 0
+
+        consecutive_retries += 1
+        if consecutive_retries >= max_consecutive_retries:
+            log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.")
+            return True, False, consecutive_retries
+
+        if drained_count > 0:
+            debug_log(f"Master idle drain picked up {drained_count} images while waiting for workers")
+        debug_log("Waiting for workers")
+        run_async_in_server_loop(asyncio.sleep(2), timeout=3.0)
+        return False, True, consecutive_retries
+
+    def _finalize_master_dynamic_results(
+        self,
+        multi_job_id,
+        batch_size,
+        num_workers,
+        processed_count,
+        result_images,
+        upscaled_image,
+        mode_context: DynamicModeContext | None = None,
+    ):
+        """Collect remaining worker outputs and convert final images back to tensor."""
+        context = mode_context or self._build_dynamic_mode_context()
+        all_completed = run_async_in_server_loop(
+            context.job_state.all_completed_images(multi_job_id),
+            timeout=5.0,
+        )
+        remaining_to_collect = batch_size - len(all_completed)
+        if remaining_to_collect > 0:
+            debug_log(f"Waiting for {remaining_to_collect} more images from workers")
+            collection_timeout = float(get_worker_timeout_seconds())
+            collected_images = run_async_in_server_loop(
+                context.result_collector.collect_dynamic_images(
+                    multi_job_id,
+                    remaining_to_collect,
+                    num_workers,
+                    batch_size,
+                    processed_count,
+                ),
+                timeout=collection_timeout,
+            )
+            all_completed.update(collected_images)
+
+        for idx, processed_img in all_completed.items():
+            if idx < batch_size:
+                result_images[idx] = processed_img
+
+        result_tensor = (
+            torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
+            if batch_size > 1
+            else pil_to_tensor(result_images[0])
+        )
+        if upscaled_image.is_cuda:
+            result_tensor = result_tensor.cuda()
+        return result_tensor
+
+    def process_master_dynamic(
+        self,
+        upscaled_image: torch.Tensor,
+        core_args: UpscaleCoreArgs,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        multi_job_id: str,
+        enabled_workers: list[str],
+        mode_context: DynamicModeContext | None = None,
+    ) -> tuple[torch.Tensor]:
+        """Dynamic mode for large batches - assigns whole images to workers dynamically, including master."""
+        context = mode_context or self._build_dynamic_mode_context()
+        # Get batch size and dimensions
+        batch_size, height, width, _ = upscaled_image.shape
+        num_workers = len(enabled_workers)
+        
+        log(f"USDU Dist: Image queue distribution | Batch {batch_size} | Workers {num_workers} | Canvas {width}x{height} | Tile {tile_width}x{tile_height}")
+
+        # No fixed share - all images are dynamic
+        all_indices = list(range(batch_size))
+        
+        debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.")
+        
+        # Calculate tiles for processing
+        all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+        
+        # Initialize job queue for communication
+        try:
+            run_async_in_server_loop(
+                init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices),
+                timeout=2.0
+            )
+        except Exception as e:
+            debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}")
+            raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}")
+        
+        # Convert batch to PIL list
+        result_images = [tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB').copy() for b in range(batch_size)]
+        
+        # Process images dynamically with master participating
+        prompt_server = ensure_tile_jobs_initialized()
+        processed_count = 0
+        consecutive_retries = 0
+        max_consecutive_retries = 10
+        
+        # Process loop - master pulls from queue and processes synchronously
+        while processed_count < batch_size:
+            # Try to get an image to process
+            image_idx = run_async_in_server_loop(
+                context.job_state.next_image_index(multi_job_id),
+                timeout=5.0  # Short timeout to allow frequent checks
+            )
+
+            if image_idx is not None:
+                # Reset retry counter and process locally
+                consecutive_retries = 0
+                debug_log(f"Master processing image {image_idx} dynamically")
+                processed_count += 1
+
+                # Process locally
+                single_tensor = upscaled_image[image_idx:image_idx+1]
+                local_image = result_images[image_idx]
+
+                local_image = self._process_image_tiles(
+                    image_idx, local_image, single_tensor, core_args, all_tiles,
+                    tile_width, tile_height, padding, mask_blur, width, height,
+                    force_uniform_tiles, context,
+                    per_tile_callback=lambda: run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1),
+                )
+
+                result_images[image_idx] = local_image
+                
+                # Mark as completed
+                run_async_in_server_loop(
+                    context.result_collector.mark_image_completed(multi_job_id, image_idx, local_image),
+                    timeout=5.0
+                )
+                
+                # NEW: Drain after the full image is marked complete (catches workers who finished during master's processing)
+                drained_count = run_async_in_server_loop(
+                    context.job_state.drain_worker_results_queue(multi_job_id),
+                    timeout=5.0
+                )
+                
+                if drained_count > 0:
+                    debug_log(f"Drained {drained_count} worker images after master's image {image_idx}")
+                
+                # NEW: Log overall progress (includes master's image + any drained workers)
+                completed_now = run_async_in_server_loop(
+                    context.job_state.total_completed_count(multi_job_id),
+                    timeout=1.0
+                )
+                log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
+                
+                # Yield to allow workers to get new images after completing one
+                run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
+            else:
+                should_break, should_continue, consecutive_retries = self._handle_master_dynamic_idle_state(
+                    multi_job_id=multi_job_id,
+                    batch_size=batch_size,
+                    consecutive_retries=consecutive_retries,
+                    max_consecutive_retries=max_consecutive_retries,
+                    mode_context=context,
+                )
+                if should_break:
+                    break
+                if should_continue:
+                    continue
+        
+        debug_log(f"Master processed {processed_count} images locally")
+        result_tensor = self._finalize_master_dynamic_results(
+            multi_job_id=multi_job_id,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            processed_count=processed_count,
+            result_images=result_images,
+            upscaled_image=upscaled_image,
+            mode_context=context,
+        )
+        
+        debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
+        log(f"Completed processing all {batch_size} images")
+        return (result_tensor,)
+
+    def process_worker_dynamic(
+        self,
+        upscaled_image: torch.Tensor,
+        core_args: UpscaleCoreArgs,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+        enabled_worker_ids: list[str] | str,
+        dynamic_threshold: int,
+        mode_context: DynamicModeContext | None = None,
+    ) -> tuple[torch.Tensor]:
+        """Worker processing in dynamic mode - processes whole images."""
+        context = mode_context or self._build_dynamic_mode_context()
+        # Round tile dimensions
+        tile_width = context.tile_ops.round_to_multiple(tile_width)
+        tile_height = context.tile_ops.round_to_multiple(tile_height)
+
+        # Get dimensions and tile grid
+        batch_size, height, width, _ = upscaled_image.shape
+        all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+        log(f"USDU Dist Worker[{worker_id[:8]}]: Processing image queue | Batch {batch_size}")
+
+        # Keep track of processed images for is_last detection
+        processed_count = 0
+
+        # Ensure the dynamic job queue exists on the master.  In delegate-only
+        # mode the master's USDU node is replaced with a lightweight collector
+        # that cannot initialise the queue itself, so workers do it instead.
+        # The call is idempotent — if the master already created the queue this
+        # is a no-op.
+        from ...utils.worker_ids import coerce_enabled_worker_ids as _coerce
+        _enabled = _coerce(enabled_worker_ids)
+        run_async_in_server_loop(
+            context.worker_comms.init_dynamic_job_on_master(
+                multi_job_id, master_url, batch_size, _enabled,
+            ),
+            timeout=10.0,
+        )
+
+        # Poll for job readiness to avoid races during master init
+        max_poll_attempts = JOB_POLL_MAX_ATTEMPTS
+        if not self._poll_job_ready_dynamic(
+            multi_job_id,
+            master_url,
+            worker_id=worker_id,
+            max_attempts=max_poll_attempts,
+            mode_context=context,
+        ):
+            log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
+            return (upscaled_image,)
+
+        # Loop to request and process images
+        while True:
+            # Request an image to process
+            assignment = run_async_in_server_loop(
+                context.worker_comms.request_assignment(multi_job_id, master_url, worker_id),
+                timeout=TILE_WAIT_TIMEOUT,
+            )
+            image_idx = assignment.task_idx if assignment.kind == "image" else None
+            estimated_remaining = assignment.estimated_remaining
+
+            if image_idx is None:
+                debug_log(f"USDU Dist Worker - No more images to process")
+                break
+
+            debug_log(f"Worker[{worker_id[:8]}] - Assigned image {image_idx}")
+            processed_count += 1
+
+            # Determine if this should be marked as last for this worker
+            is_last_for_worker = (estimated_remaining == 0)
+
+            # Extract single image tensor
+            single_tensor = upscaled_image[image_idx:image_idx+1]
+
+            # Convert to PIL for processing
+            local_image = tensor_to_pil(single_tensor, 0).copy()
+
+            # Process all tiles for this image
+            local_image = self._process_image_tiles(
+                image_idx, local_image, single_tensor, core_args, all_tiles,
+                tile_width, tile_height, padding, mask_blur, width, height,
+                force_uniform_tiles, context,
+                per_tile_callback=lambda: run_async_in_server_loop(
+                    context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id),
+                    timeout=5.0,
+                ),
+            )
+
+            # Send processed image back to master
+            try:
+                # Use the estimated remaining to determine if this is the last image
+                is_last = is_last_for_worker
+                run_async_in_server_loop(
+                    context.worker_comms.send_full_image(
+                        local_image,
+                        image_idx,
+                        multi_job_id,
+                        master_url,
+                        worker_id,
+                        is_last,
+                    ),
+                    timeout=TILE_SEND_TIMEOUT
+                )
+                # Send heartbeat after processing
+                run_async_in_server_loop(
+                    context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id),
+                    timeout=5.0
+                )
+                if is_last:
+                    break
+            except Exception as e:
+                log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
+                # Continue processing other images
+
+        # Send final is_last signal
+        debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal")
+        try:
+            run_async_in_server_loop(
+                context.worker_comms.send_worker_complete_signal(multi_job_id, master_url, worker_id),
+                timeout=TILE_SEND_TIMEOUT
+            )
+        except Exception as e:
+            log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending completion signal: {e}")
+
+        return (upscaled_image,)

upscale/modes/single_gpu.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import math, torch
+from typing import Any
+from ...utils.logging import log
+from ...utils.image import blend_processed_batch_item, pil_to_tensor, tensor_to_pil
+from ..mode_contexts import SingleGpuModeContext, TileOpsCollaborator
+from ..processing_args import UpscaleCoreArgs
+from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch
+
+
+class SingleGpuModeMixin:
+    def _build_single_gpu_mode_context(self) -> SingleGpuModeContext:
+        """Build explicit collaborators for single-GPU mode execution."""
+        return SingleGpuModeContext(tile_ops=TileOpsCollaborator(self))
+
+    def process_single_gpu(
+        self,
+        upscaled_image: torch.Tensor,
+        core_args: UpscaleCoreArgs,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        mode_context: SingleGpuModeContext | None = None,
+    ) -> tuple[torch.Tensor]:
+        """Process all tiles on a single GPU (no distribution), batching per tile like USDU."""
+        context = mode_context or self._build_single_gpu_mode_context()
+        ops = context.tile_ops
+        # Round tile dimensions
+        tile_width = ops.round_to_multiple(tile_width)
+        tile_height = ops.round_to_multiple(tile_height)
+
+        # Get image dimensions and batch size
+        batch_size, height, width, _ = upscaled_image.shape
+
+        # Calculate all tiles
+        all_tiles = ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+
+        rows = math.ceil(height / tile_height)
+        cols = math.ceil(width / tile_width)
+        log(
+            f"USDU Dist: Single GPU | Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({len(all_tiles)} tiles/image) | Batch {batch_size}"
+        )
+        tile_batch_args = TileBatchArgs(
+            core=core_args,
+            tile_width=tile_width,
+            tile_height=tile_height,
+            padding=padding,
+            force_uniform_tiles=force_uniform_tiles,
+            width=width,
+            height=height,
+        )
+
+        # Prepare result images list
+        result_images = []
+        for b in range(batch_size):
+            image_pil = tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB')
+            result_images.append(image_pil.copy())
+
+        # Precompute tile masks once
+        tile_masks = []
+        for tx, ty in all_tiles:
+            tile_masks.append(ops.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
+
+        # Process tiles batched across images
+        for tile_idx, (tx, ty) in enumerate(all_tiles):
+            # Progressive state parity: extract each tile from the current updated image batch.
+            source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
+            if upscaled_image.is_cuda:
+                source_batch = source_batch.cuda()
+
+            processed_batch, x1, y1, ew, eh = extract_and_process_tile_batch(
+                node=ops,
+                upscaled_image=source_batch,
+                tx=tx,
+                ty=ty,
+                args=tile_batch_args,
+            )
+
+            # Blend results back into each image using cached mask
+            tile_mask = tile_masks[tile_idx]
+            for b in range(batch_size):
+                blend_processed_batch_item(
+                    result_images,
+                    processed_batch,
+                    b,
+                    ops.blend_tile,
+                    x1,
+                    y1,
+                    ew,
+                    eh,
+                    tile_mask,
+                    padding,
+                )
+
+        # Convert back to tensor
+        result_tensors = [pil_to_tensor(img) for img in result_images]
+        result_tensor = torch.cat(result_tensors, dim=0)
+        if upscaled_image.is_cuda:
+            result_tensor = result_tensor.cuda()
+
+        return (result_tensor,)

upscale/modes/static.py

@@ -0,0 +1,685 @@
+import asyncio, time, torch
+from PIL import Image
+import comfy.model_management
+from ...utils.logging import debug_log, log
+from ...utils.image import blend_processed_batch_item, pil_to_tensor, tensor_to_pil
+from ...utils.async_helpers import run_async_in_server_loop
+from ...utils.config import get_worker_timeout_seconds
+from ...utils.constants import (
+    HEARTBEAT_INTERVAL,
+    JOB_POLL_INTERVAL,
+    JOB_POLL_MAX_ATTEMPTS,
+    MAX_BATCH,
+    TILE_SEND_TIMEOUT,
+    TILE_WAIT_TIMEOUT,
+)
+from ..job_store import (
+    ensure_tile_jobs_initialized, init_static_job_batched,
+    cleanup_job, drain_results_queue, get_completed_count, mark_task_completed,
+)
+from ..job_models import TileJobState
+from ..mode_contexts import (
+    JobStateCollaborator,
+    StaticModeContext,
+    TileOpsCollaborator,
+    WorkerCommsCollaborator,
+)
+from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch
+
+
+class StaticModeMixin:
+    """
+    Static (tile-queue) USDU mode behaviors for master and worker roles.
+
+    Expected co-mixins on `self`:
+    - TileOpsMixin (`calculate_tiles`, tile extract/blend helpers).
+    - JobStateMixin (`_get_next_tile_index`, `_get_all_completed_tasks`, requeue checks).
+    - WorkerCommsMixin (`send_tiles_batch`, `request_assignment`, `send_heartbeat`).
+    """
+
+    def _build_static_mode_context(self) -> StaticModeContext:
+        """Build explicit collaborators for static mode execution."""
+        return StaticModeContext(
+            tile_ops=TileOpsCollaborator(self),
+            job_state=JobStateCollaborator(self),
+            worker_comms=WorkerCommsCollaborator(self),
+        )
+
+    def _poll_job_ready(
+        self,
+        multi_job_id,
+        master_url,
+        worker_id=None,
+        max_attempts=JOB_POLL_MAX_ATTEMPTS,
+        mode_context: StaticModeContext | None = None,
+    ):
+        """Poll master for job readiness to avoid worker/master initialization race."""
+        context = mode_context or self._build_static_mode_context()
+        for attempt in range(max_attempts):
+            ready = run_async_in_server_loop(
+                context.worker_comms.check_job_status(multi_job_id, master_url),
+                timeout=5.0
+            )
+            if ready:
+                if worker_id:
+                    debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
+                else:
+                    debug_log(f"Job {multi_job_id} ready after {attempt} attempts")
+                return True
+            time.sleep(JOB_POLL_INTERVAL)
+        return False
+
+    def _extract_and_process_tile(
+        self,
+        upscaled_image,
+        tile_id,
+        all_tiles,
+        tile_width,
+        tile_height,
+        padding,
+        force_uniform_tiles,
+        core_args,
+        width,
+        height,
+    ):
+        """Extract one tile position for the whole batch and process it."""
+        tx, ty = all_tiles[tile_id]
+        tile_batch_args = TileBatchArgs(
+            core=core_args,
+            tile_width=tile_width,
+            tile_height=tile_height,
+            padding=padding,
+            force_uniform_tiles=force_uniform_tiles,
+            width=width,
+            height=height,
+        )
+        processed_batch, x1, y1, ew, eh = extract_and_process_tile_batch(
+            node=self,
+            upscaled_image=upscaled_image,
+            tx=tx,
+            ty=ty,
+            args=tile_batch_args,
+        )
+        return processed_batch, x1, y1, ew, eh
+
+    def _flush_tiles_to_master(
+        self,
+        processed_tiles,
+        multi_job_id,
+        master_url,
+        padding,
+        worker_id,
+        is_final_flush=False,
+        mode_context: StaticModeContext | None = None,
+    ):
+        """Send accumulated tile payloads to master and return a fresh accumulator."""
+        context = mode_context or self._build_static_mode_context()
+        if not processed_tiles:
+            if is_final_flush:
+                run_async_in_server_loop(
+                    context.worker_comms.send_tiles_batch(
+                        [],
+                        multi_job_id,
+                        master_url,
+                        padding,
+                        worker_id,
+                        is_final_flush=True,
+                    ),
+                    timeout=TILE_SEND_TIMEOUT,
+                )
+            return processed_tiles
+        run_async_in_server_loop(
+            context.worker_comms.send_tiles_batch(
+                processed_tiles,
+                multi_job_id,
+                master_url,
+                padding,
+                worker_id,
+                is_final_flush=is_final_flush,
+            ),
+            timeout=TILE_SEND_TIMEOUT
+        )
+        return []
+
+    def _master_process_one_tile(
+        self,
+        tile_id,
+        all_tiles,
+        upscaled_image,
+        result_images,
+        tile_masks,
+        multi_job_id,
+        batch_size,
+        num_tiles_per_image,
+        tile_width,
+        tile_height,
+        padding,
+        force_uniform_tiles,
+        core_args,
+        width,
+        height,
+        mode_context: StaticModeContext | None = None,
+    ):
+        """Process one tile_id across the batch and blend into result_images."""
+        context = mode_context or self._build_static_mode_context()
+        source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
+        if upscaled_image.is_cuda:
+            source_batch = source_batch.cuda()
+        processed_batch, x1, y1, ew, eh = self._extract_and_process_tile(
+            source_batch,
+            tile_id,
+            all_tiles,
+            tile_width,
+            tile_height,
+            padding,
+            force_uniform_tiles,
+            core_args,
+            width,
+            height,
+        )
+        tile_mask = tile_masks[tile_id]
+        out_bs = processed_batch.shape[0] if hasattr(processed_batch, "shape") else batch_size
+        processed_items = min(batch_size, out_bs)
+        for b in range(processed_items):
+            blend_processed_batch_item(
+                result_images,
+                processed_batch,
+                b,
+                context.tile_ops.blend_tile,
+                x1,
+                y1,
+                ew,
+                eh,
+                tile_mask,
+                padding,
+            )
+            global_idx = b * num_tiles_per_image + tile_id
+            run_async_in_server_loop(
+                mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}),
+                timeout=5.0
+            )
+        return processed_items
+
+    def _process_worker_static_sync(self, upscaled_image, core_args,
+                                    tile_width, tile_height, padding, mask_blur,
+                                    force_uniform_tiles, multi_job_id, master_url,
+                                    worker_id, enabled_workers,
+                                    mode_context: StaticModeContext | None = None):
+        """Worker static mode processing with optional dynamic queue pulling."""
+        context = mode_context or self._build_static_mode_context()
+        # Round tile dimensions
+        tile_width = context.tile_ops.round_to_multiple(tile_width)
+        tile_height = context.tile_ops.round_to_multiple(tile_height)
+        
+        # Get dimensions and calculate tiles
+        _, height, width, _ = upscaled_image.shape
+        all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+        num_tiles_per_image = len(all_tiles)
+        batch_size = upscaled_image.shape[0]
+        total_tiles = batch_size * num_tiles_per_image
+        
+        processed_tiles = []
+        working_images = []
+        for b in range(batch_size):
+            image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
+            working_images.append(image_pil.copy())
+        tile_masks = []
+        for tx, ty in all_tiles:
+            tile_masks.append(context.tile_ops.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
+        
+        # Dynamic queue mode (static processing): process batched-per-tile
+        log(f"USDU Dist Worker[{worker_id[:8]}]: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Tiles/image {num_tiles_per_image} | Batch {batch_size}")
+        processed_count = 0
+
+        max_poll_attempts = JOB_POLL_MAX_ATTEMPTS
+        if not self._poll_job_ready(
+            multi_job_id,
+            master_url,
+            worker_id=worker_id,
+            max_attempts=max_poll_attempts,
+            mode_context=context,
+        ):
+            log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
+            return (upscaled_image,)
+
+        # Main processing loop - pull tile ids from queue
+        while True:
+            # Request a tile to process
+            assignment = run_async_in_server_loop(
+                context.worker_comms.request_assignment(multi_job_id, master_url, worker_id),
+                timeout=TILE_WAIT_TIMEOUT,
+            )
+            tile_idx = assignment.task_idx if assignment.kind == "tile" else None
+
+            if tile_idx is None:
+                debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process")
+                break
+
+            # Always batched-per-tile in static mode
+            debug_log(f"Worker[{worker_id[:8]}] - Assigned tile_id {tile_idx}")
+            processed_count += batch_size
+            tile_id = tile_idx
+            source_batch = torch.cat([pil_to_tensor(img) for img in working_images], dim=0)
+            if upscaled_image.is_cuda:
+                source_batch = source_batch.cuda()
+            processed_batch, x1, y1, ew, eh = self._extract_and_process_tile(
+                source_batch,
+                tile_id,
+                all_tiles,
+                tile_width,
+                tile_height,
+                padding,
+                force_uniform_tiles,
+                core_args,
+                width,
+                height,
+            )
+            # Queue results
+            for b in range(batch_size):
+                tile_pil = tensor_to_pil(processed_batch, b)
+                if tile_pil.size != (ew, eh):
+                    tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
+                working_images[b] = context.tile_ops.blend_tile(
+                    working_images[b],
+                    tile_pil,
+                    x1,
+                    y1,
+                    (ew, eh),
+                    tile_masks[tile_id],
+                    padding,
+                )
+                processed_tiles.append({
+                    'tile': processed_batch[b:b+1],
+                    'tile_idx': tile_id,
+                    'x': x1,
+                    'y': y1,
+                    'extracted_width': ew,
+                    'extracted_height': eh,
+                    'padding': padding,
+                    'batch_idx': b,
+                    'global_idx': b * num_tiles_per_image + tile_id
+                })
+
+            # Send heartbeat
+            try:
+                run_async_in_server_loop(
+                    context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id),
+                    timeout=5.0
+                )
+            except Exception as e:
+                debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}")
+
+            # Send tiles in batches within loop
+            if len(processed_tiles) >= MAX_BATCH:
+                processed_tiles = self._flush_tiles_to_master(
+                    processed_tiles,
+                    multi_job_id,
+                    master_url,
+                    padding,
+                    worker_id,
+                    is_final_flush=False,
+                    mode_context=context,
+                )
+
+        # Send any remaining tiles
+        processed_tiles = self._flush_tiles_to_master(
+            processed_tiles,
+            multi_job_id,
+            master_url,
+            padding,
+            worker_id,
+            is_final_flush=True,
+            mode_context=context,
+        )
+        
+        debug_log(f"Worker {worker_id} completed all assigned and requeued tiles")
+        return (upscaled_image,)
+
+    async def _async_collect_and_monitor_static(
+        self,
+        multi_job_id,
+        total_tiles,
+        expected_total,
+        mode_context: StaticModeContext | None = None,
+    ):
+        """Async helper for collection and monitoring in static mode.
+        Returns collected tasks dict. Caller should check if all tasks are complete."""
+        context = mode_context or self._build_static_mode_context()
+        last_progress_log = time.time()
+        progress_interval = 5.0
+        last_heartbeat_check = time.time()
+        last_completed_count = 0
+        
+        while True:
+            # Check for user interruption
+            if comfy.model_management.processing_interrupted():
+                log("Processing interrupted by user")
+                raise comfy.model_management.InterruptProcessingException()
+            
+            # Drain any pending results
+            collected_count = await drain_results_queue(multi_job_id)
+            
+            # Check and requeue timed-out workers periodically
+            current_time = time.time()
+            if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
+                requeued_count = await context.job_state.check_and_requeue_timed_out_workers(multi_job_id, expected_total)
+                if requeued_count > 0:
+                    log(f"Requeued {requeued_count} tasks from timed-out workers")
+                last_heartbeat_check = current_time
+            
+            # Get current completion count
+            completed_count = await get_completed_count(multi_job_id)
+            
+            # Progress logging
+            if current_time - last_progress_log >= progress_interval:
+                log(f"Progress: {completed_count}/{expected_total} tasks completed")
+                last_progress_log = current_time
+            
+            # Check if all tasks are completed
+            if completed_count >= expected_total:
+                debug_log(f"All {expected_total} tasks completed")
+                break
+            
+            # If no active workers remain and there are pending tasks, return for local processing
+            prompt_server = ensure_tile_jobs_initialized()
+            async with prompt_server.distributed_tile_jobs_lock:
+                job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+                if isinstance(job_data, TileJobState):
+                    pending_queue = job_data.pending_tasks
+                    active_workers = list(job_data.worker_status.keys())
+                    if pending_queue and not pending_queue.empty() and len(active_workers) == 0:
+                        log(f"No active workers remaining with {expected_total - completed_count} tasks pending. Returning for local processing.")
+                        break
+            
+            # Wait a bit before next check
+            await asyncio.sleep(0.1)
+        
+        # Get all completed tasks for return
+        return await context.job_state.all_completed_tasks(multi_job_id)
+
+    def _init_master_static_state(
+        self,
+        upscaled_image,
+        multi_job_id,
+        batch_size,
+        num_tiles_per_image,
+        enabled_workers,
+        all_tiles,
+        tile_width,
+        tile_height,
+        mask_blur,
+        mode_context: StaticModeContext | None = None,
+    ):
+        context = mode_context or self._build_static_mode_context()
+        _, height, width, _ = upscaled_image.shape
+        result_images = [tensor_to_pil(upscaled_image[b:b + 1], 0).copy() for b in range(batch_size)]
+
+        log("USDU Dist: Using tile queue distribution")
+        run_async_in_server_loop(
+            init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers),
+            timeout=10.0,
+        )
+        debug_log(f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}")
+
+        tile_masks = [
+            context.tile_ops.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)
+            for tx, ty in all_tiles
+        ]
+        return result_images, tile_masks, width, height
+
+    def _process_master_static_initial_tiles(
+        self,
+        multi_job_id,
+        total_tiles,
+        all_tiles,
+        upscaled_image,
+        result_images,
+        tile_masks,
+        batch_size,
+        num_tiles_per_image,
+        tile_width,
+        tile_height,
+        padding,
+        force_uniform_tiles,
+        core_args,
+        width,
+        height,
+        mode_context: StaticModeContext | None = None,
+    ):
+        context = mode_context or self._build_static_mode_context()
+        processed_count = 0
+        consecutive_no_tile = 0
+        max_consecutive_no_tile = 2
+
+        while processed_count < total_tiles:
+            comfy.model_management.throw_exception_if_processing_interrupted()
+            tile_id = run_async_in_server_loop(context.job_state.next_tile_index(multi_job_id), timeout=5.0)
+            if tile_id is None:
+                consecutive_no_tile += 1
+                if consecutive_no_tile >= max_consecutive_no_tile:
+                    debug_log(f"Master processed {processed_count} tiles, moving to collection phase")
+                    break
+                time.sleep(0.1)
+                continue
+
+            consecutive_no_tile = 0
+            processed_count += self._master_process_one_tile(
+                tile_id,
+                all_tiles,
+                upscaled_image,
+                result_images,
+                tile_masks,
+                multi_job_id,
+                batch_size,
+                num_tiles_per_image,
+                tile_width,
+                tile_height,
+                padding,
+                force_uniform_tiles,
+                core_args,
+                width,
+                height,
+                mode_context=context,
+            )
+            log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})")
+        return processed_count
+
+    def _collect_remaining_static_tiles(
+        self,
+        multi_job_id,
+        total_tiles,
+        master_processed_count,
+        all_tiles,
+        upscaled_image,
+        result_images,
+        tile_masks,
+        batch_size,
+        num_tiles_per_image,
+        tile_width,
+        tile_height,
+        padding,
+        force_uniform_tiles,
+        core_args,
+        width,
+        height,
+        mode_context: StaticModeContext | None = None,
+    ):
+        context = mode_context or self._build_static_mode_context()
+        remaining_tiles = total_tiles - master_processed_count
+        if remaining_tiles <= 0:
+            return run_async_in_server_loop(context.job_state.all_completed_tasks(multi_job_id), timeout=5.0)
+
+        debug_log(f"Master waiting for {remaining_tiles} tiles from workers")
+        collected_tasks = run_async_in_server_loop(
+            self._async_collect_and_monitor_static(
+                multi_job_id,
+                total_tiles,
+                expected_total=total_tiles,
+                mode_context=context,
+            ),
+            timeout=None,
+        )
+
+        completed_count = len(collected_tasks)
+        if completed_count >= total_tiles:
+            return collected_tasks
+
+        log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures")
+        while True:
+            comfy.model_management.throw_exception_if_processing_interrupted()
+            tile_id = run_async_in_server_loop(context.job_state.next_tile_index(multi_job_id), timeout=5.0)
+            if tile_id is None:
+                break
+            self._master_process_one_tile(
+                tile_id,
+                all_tiles,
+                upscaled_image,
+                result_images,
+                tile_masks,
+                multi_job_id,
+                batch_size,
+                num_tiles_per_image,
+                tile_width,
+                tile_height,
+                padding,
+                force_uniform_tiles,
+                core_args,
+                width,
+                height,
+                mode_context=context,
+            )
+        return collected_tasks
+
+    def _blend_static_collected_tiles(
+        self,
+        collected_tasks,
+        result_images,
+        all_tiles,
+        tile_masks,
+        num_tiles_per_image,
+        batch_size,
+        tile_width,
+        tile_height,
+        padding,
+        mode_context: StaticModeContext | None = None,
+    ):
+        context = mode_context or self._build_static_mode_context()
+        def _sort_key(item):
+            global_idx, tile_data = item
+            batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image)
+            tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image)
+            return (tile_idx, batch_idx, global_idx)
+
+        for global_idx, tile_data in sorted(collected_tasks.items(), key=_sort_key):
+            if 'tensor' not in tile_data and 'image' not in tile_data:
+                continue
+            batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image)
+            tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image)
+            if batch_idx >= batch_size:
+                continue
+
+            x = tile_data.get('x', 0)
+            y = tile_data.get('y', 0)
+            tile_pil = tile_data['image'] if 'image' in tile_data else tensor_to_pil(tile_data['tensor'], 0)
+            tile_mask = tile_masks[tile_idx]
+            extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding)
+            extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding)
+            result_images[batch_idx] = context.tile_ops.blend_tile(
+                result_images[batch_idx],
+                tile_pil,
+                x,
+                y,
+                (extracted_width, extracted_height),
+                tile_mask,
+                padding,
+            )
+
+    def _result_images_to_tensor(self, result_images, batch_size, upscaled_image):
+        if batch_size == 1:
+            result_tensor = pil_to_tensor(result_images[0])
+        else:
+            result_tensor = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
+        if upscaled_image.is_cuda:
+            result_tensor = result_tensor.cuda()
+        return result_tensor
+
+    def _process_master_static_sync(self, upscaled_image, core_args,
+                                    tile_width, tile_height, padding, mask_blur,
+                                    force_uniform_tiles, multi_job_id, enabled_workers,
+                                    all_tiles, num_tiles_per_image,
+                                    mode_context: StaticModeContext | None = None):
+        """Static mode master processing with optional dynamic queue pulling."""
+        context = mode_context or self._build_static_mode_context()
+        batch_size = upscaled_image.shape[0]
+        total_tiles = batch_size * num_tiles_per_image
+        result_images, tile_masks, width, height = self._init_master_static_state(
+            upscaled_image=upscaled_image,
+            multi_job_id=multi_job_id,
+            batch_size=batch_size,
+            num_tiles_per_image=num_tiles_per_image,
+            enabled_workers=enabled_workers,
+            all_tiles=all_tiles,
+            tile_width=tile_width,
+            tile_height=tile_height,
+            mask_blur=mask_blur,
+            mode_context=context,
+        )
+
+        try:
+            master_processed_count = self._process_master_static_initial_tiles(
+                multi_job_id=multi_job_id,
+                total_tiles=total_tiles,
+                all_tiles=all_tiles,
+                upscaled_image=upscaled_image,
+                result_images=result_images,
+                tile_masks=tile_masks,
+                batch_size=batch_size,
+                num_tiles_per_image=num_tiles_per_image,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                force_uniform_tiles=force_uniform_tiles,
+                core_args=core_args,
+                width=width,
+                height=height,
+                mode_context=context,
+            )
+
+            collected_tasks = self._collect_remaining_static_tiles(
+                multi_job_id=multi_job_id,
+                total_tiles=total_tiles,
+                master_processed_count=master_processed_count,
+                all_tiles=all_tiles,
+                upscaled_image=upscaled_image,
+                result_images=result_images,
+                tile_masks=tile_masks,
+                batch_size=batch_size,
+                num_tiles_per_image=num_tiles_per_image,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                force_uniform_tiles=force_uniform_tiles,
+                core_args=core_args,
+                width=width,
+                height=height,
+                mode_context=context,
+            )
+
+            self._blend_static_collected_tiles(
+                collected_tasks=collected_tasks,
+                result_images=result_images,
+                all_tiles=all_tiles,
+                tile_masks=tile_masks,
+                num_tiles_per_image=num_tiles_per_image,
+                batch_size=batch_size,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mode_context=context,
+            )
+
+            result_tensor = self._result_images_to_tensor(result_images, batch_size, upscaled_image)
+            log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
+            return (result_tensor,)
+        finally:
+            run_async_in_server_loop(cleanup_job(multi_job_id), timeout=5.0)

upscale/payload_parsers.py

@@ -0,0 +1,70 @@
+import io
+import json
+from collections.abc import Mapping
+from typing import Any
+
+from PIL import Image
+
+
+def parse_tiles_from_form(data: Mapping[str, Any]) -> list[dict[str, Any]]:
+    """Parse tiles submitted via multipart/form-data into a list of tile dicts."""
+    try:
+        padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
+    except Exception:
+        padding = 0
+
+    meta_raw = data.get('tiles_metadata')
+    if meta_raw is None:
+        raise ValueError("Missing tiles_metadata")
+
+    try:
+        metadata = json.loads(meta_raw)
+    except Exception as e:
+        raise ValueError(f"Invalid tiles_metadata JSON: {e}")
+
+    if not isinstance(metadata, list):
+        raise ValueError("tiles_metadata must be a list")
+
+    tiles = []
+    for i, meta in enumerate(metadata):
+        file_field = data.get(f'tile_{i}')
+        if file_field is None or not hasattr(file_field, 'file'):
+            raise ValueError(f"Missing tile data for index {i}")
+
+        raw = file_field.file.read()
+        try:
+            img = Image.open(io.BytesIO(raw)).convert("RGB")
+        except Exception as e:
+            raise ValueError(f"Invalid image data for tile {i}: {e}")
+
+        try:
+            tile_info = {
+                'image': img,
+                'tile_idx': int(meta.get('tile_idx', i)),
+                'x': int(meta.get('x', 0)),
+                'y': int(meta.get('y', 0)),
+                'extracted_width': int(meta.get('extracted_width', img.width)),
+                'extracted_height': int(meta.get('extracted_height', img.height)),
+                'padding': int(padding),
+            }
+        except Exception as e:
+            raise ValueError(f"Invalid metadata values for tile {i}: {e}")
+
+        if 'batch_idx' in meta:
+            try:
+                tile_info['batch_idx'] = int(meta['batch_idx'])
+            except Exception as exc:
+                raise ValueError(f"Invalid batch_idx for tile {i}: {meta.get('batch_idx')} ({exc})")
+        if 'global_idx' in meta:
+            try:
+                tile_info['global_idx'] = int(meta['global_idx'])
+            except Exception as exc:
+                raise ValueError(f"Invalid global_idx for tile {i}: {meta.get('global_idx')} ({exc})")
+
+        tiles.append(tile_info)
+
+    return tiles
+
+
+# Backward compatibility for existing imports.
+_parse_tiles_from_form = parse_tiles_from_form

upscale/processing_args.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class UpscaleCoreArgs:
+    """Shared denoise/sampling arguments for USDU tile processing."""
+
+    model: Any
+    positive: Any
+    negative: Any
+    vae: Any
+    seed: int
+    steps: int
+    cfg: float
+    sampler_name: str
+    scheduler: str
+    denoise: float
+    tiled_decode: bool

upscale/result_collector.py

@@ -0,0 +1,223 @@
+import asyncio, time
+import comfy.model_management
+import server
+from ..utils.constants import DYNAMIC_MODE_MAX_POLL_TIMEOUT, HEARTBEAT_INTERVAL
+from ..utils.logging import debug_log, log
+from ..utils.config import get_worker_timeout_seconds
+from .job_store import ensure_tile_jobs_initialized, mark_task_completed
+from .job_timeout import check_and_requeue_timed_out_workers
+from .job_models import BaseJobState, ImageJobState, TileJobState
+
+
+class ResultCollectorMixin:
+    """
+    Mixin for master-side result collection in USDU distributed jobs.
+
+    Expected co-mixins/attributes:
+    - JobStateMixin methods for queue/task access.
+    - `self._check_and_requeue_timed_out_workers(...)` coroutine.
+    - `self.async_yield(...)` optional helper from WorkerCommsMixin.
+    """
+
+    def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> list[str]:
+        """Log timeout elapsed seconds for each tracked worker and return worker ids."""
+        if not isinstance(job_data, BaseJobState):
+            return []
+
+        worker_status = dict(job_data.worker_status)
+        for worker_id, last_seen in worker_status.items():
+            elapsed = max(0.0, current_time - float(last_seen))
+            log(
+                "UltimateSDUpscale Master - Heartbeat timeout: "
+                f"job={multi_job_id}, worker={worker_id}, elapsed={elapsed:.1f}s"
+            )
+        return list(worker_status.keys())
+
+    async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
+        """Default timeout requeue hook; override in host mixins when needed."""
+        return await check_and_requeue_timed_out_workers(multi_job_id, batch_size)
+
+    def _get_job_data_snapshot(self, prompt_server, multi_job_id):
+        """Get a snapshot of job data for timeout logging (non-async helper)."""
+        current_job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if isinstance(current_job_data, BaseJobState):
+            return current_job_data
+        return None
+
+    async def _async_collect_worker_tiles(self, multi_job_id, num_workers):
+        """Collect tiles from workers in static mode."""
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
+                raise RuntimeError(f"Job queue not initialized for {multi_job_id}")
+            job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+            if not isinstance(job_data, TileJobState):
+                raise RuntimeError(
+                    f"Mode mismatch: expected static, got {getattr(job_data, 'mode', 'unknown')}"
+                )
+            q = job_data.queue
+            expected_count = len(job_data.completed_tasks) + job_data.pending_tasks.qsize()
+
+        debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} tiles from {num_workers} workers")
+
+        collected_results = {}
+        workers_done = set()
+        timeout = float(get_worker_timeout_seconds())
+        wait_started_at = time.time()
+
+        while len(workers_done) < num_workers:
+            if comfy.model_management.processing_interrupted():
+                log("Processing interrupted by user")
+                raise comfy.model_management.InterruptProcessingException()
+
+            job_data_snapshot = None
+            async with prompt_server.distributed_tile_jobs_lock:
+                job_data_snapshot = self._get_job_data_snapshot(prompt_server, multi_job_id)
+
+            try:
+                result = await asyncio.wait_for(q.get(), timeout=timeout)
+                worker_id = result['worker_id']
+                is_last = result.get('is_last', False)
+
+                tiles = result.get('tiles', [])
+                debug_log(
+                    f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker "
+                    f"'{worker_id}' (is_last={is_last})"
+                )
+
+                for tile_data in tiles:
+                    if 'batch_idx' not in tile_data:
+                        log("UltimateSDUpscale Master - Missing batch_idx in tile data, skipping")
+                        continue
+
+                    tile_idx = tile_data['tile_idx']
+                    key = tile_data.get('global_idx', tile_idx)
+                    entry = {
+                        **tile_data,
+                        'tile_idx': tile_idx,
+                        'worker_id': worker_id,
+                        'global_idx': key,
+                    }
+                    collected_results[entry['global_idx']] = entry
+
+                if is_last:
+                    workers_done.add(worker_id)
+                    debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
+
+            except asyncio.TimeoutError:
+                current_time = time.time()
+                waiting_workers = type(self)._log_worker_timeout_status(
+                    self, job_data_snapshot, current_time, multi_job_id,
+                )
+                elapsed = current_time - wait_started_at
+                log(
+                    f"UltimateSDUpscale Master - Heartbeat timeout waiting for tiles; "
+                    f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
+                )
+                break
+
+        debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_results)} tiles from {len(workers_done)} workers")
+
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                del prompt_server.distributed_pending_tile_jobs[multi_job_id]
+
+        return collected_results
+
+    async def collect_dynamic_images(
+        self,
+        multi_job_id,
+        remaining_to_collect,
+        num_workers,
+        batch_size,
+        master_processed_count,
+    ):
+        """Collect remaining processed images from workers in dynamic mode."""
+        _ = master_processed_count
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
+                raise RuntimeError(f"Job queue not initialized for {multi_job_id}")
+            job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+            if not isinstance(job_data, ImageJobState):
+                raise RuntimeError(
+                    f"Mode mismatch: expected dynamic, got {getattr(job_data, 'mode', 'unknown')}"
+                )
+            q = job_data.queue
+            completed_images = job_data.completed_images
+            expected_count = remaining_to_collect or batch_size
+
+        debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} images from {num_workers} workers")
+
+        workers_done = set()
+        timeout = float(get_worker_timeout_seconds())
+        last_heartbeat_check = time.time()
+        wait_started_at = time.time()
+        collected_count = 0
+
+        while len(workers_done) < num_workers:
+            if comfy.model_management.processing_interrupted():
+                log("Processing interrupted by user")
+                raise comfy.model_management.InterruptProcessingException()
+
+            if remaining_to_collect and collected_count >= remaining_to_collect:
+                break
+
+            job_data_snapshot = None
+            async with prompt_server.distributed_tile_jobs_lock:
+                job_data_snapshot = self._get_job_data_snapshot(prompt_server, multi_job_id)
+
+            try:
+                wait_timeout = min(DYNAMIC_MODE_MAX_POLL_TIMEOUT, timeout)
+                result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
+                worker_id = result['worker_id']
+                is_last = result.get('is_last', False)
+
+                if 'image_idx' in result and 'image' in result:
+                    image_idx = result['image_idx']
+                    image_pil = result['image']
+                    completed_images[image_idx] = image_pil
+                    collected_count += 1
+                    debug_log(f"UltimateSDUpscale Master - Received image {image_idx} from worker {worker_id}")
+
+                if is_last:
+                    workers_done.add(worker_id)
+                    debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
+
+            except asyncio.TimeoutError:
+                current_time = time.time()
+                waiting_workers = type(self)._log_worker_timeout_status(
+                    self, job_data_snapshot, current_time, multi_job_id,
+                )
+                if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
+                    requeued = await type(self)._check_and_requeue_timed_out_workers(
+                        self, multi_job_id, batch_size,
+                    )
+                    if requeued > 0:
+                        log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers")
+                    last_heartbeat_check = current_time
+
+                if current_time - wait_started_at > timeout:
+                    elapsed = current_time - wait_started_at
+                    log(
+                        "UltimateSDUpscale Master - Heartbeat timeout while waiting for images; "
+                        f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
+                    )
+                    break
+
+        debug_log(f"UltimateSDUpscale Master - Collection complete. Got {collected_count} images from {len(workers_done)} workers")
+
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                del prompt_server.distributed_pending_tile_jobs[multi_job_id]
+
+        return completed_images
+
+    async def mark_image_completed(self, multi_job_id, image_idx, image_pil):
+        """Mark an image as completed in the job data."""
+        await mark_task_completed(multi_job_id, image_idx, {'image': image_pil})
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, ImageJobState):
+                job_data.completed_images[image_idx] = image_pil

upscale/tile_ops.py

@@ -0,0 +1,503 @@
+import math, torch
+from contextlib import nullcontext
+from PIL import Image, ImageFilter, ImageDraw
+from typing import List, Tuple
+import comfy.samplers, comfy.model_management
+from ..utils.logging import debug_log, log
+from ..utils.image import tensor_to_pil, pil_to_tensor
+from ..utils.usdu_utils import crop_cond, get_crop_region, expand_crop
+from ..utils.crop_model_patch import crop_model_cond
+from .conditioning import clone_conditioning
+
+
+class TileOpsMixin:
+    def round_to_multiple(self, value: int, multiple: int = 8) -> int:
+        """Round value to nearest multiple."""
+        return round(value / multiple) * multiple
+
+    def calculate_tiles(self, image_width: int, image_height: int,
+                       tile_width: int, tile_height: int, force_uniform_tiles: bool = True) -> List[Tuple[int, int]]:
+        """Calculate tile positions to match Ultimate SD Upscale.
+
+        Positions are a simple grid starting at (0,0) with steps of
+        `tile_width` and `tile_height`, using ceil(rows/cols) to cover edges.
+        Uniform vs non-uniform affects only crop/resize, not positions.
+        """
+        rows = math.ceil(image_height / tile_height)
+        cols = math.ceil(image_width / tile_width)
+        tiles: List[Tuple[int, int]] = []
+        for yi in range(rows):
+            for xi in range(cols):
+                tiles.append((xi * tile_width, yi * tile_height))
+        return tiles
+
+    def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int,
+                                 tile_width: int, tile_height: int, padding: int,
+                                 force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
+        """Extract a tile region and resize to match USDU cropping logic.
+
+        Mirrors ComfyUI_UltimateSDUpscale processing:
+        - Build a mask with a white rectangle at the tile rect
+        - Compute crop_region via get_crop_region(mask, padding)
+        - If force_uniform_tiles: expand by crop/aspect ratio, then resize to
+          fixed processing size of round_to_multiple(tile + padding)
+        - Else: target is ceil(crop_size/8)*8 per dimension
+        - Extract the crop and resize to target tile_size
+        Returns the resized tensor and crop origin/size for blending.
+        """
+        _, h, w, _ = image.shape
+
+        # Create mask and compute initial padded crop region
+        mask = Image.new('L', (w, h), 0)
+        draw = ImageDraw.Draw(mask)
+        draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
+        x1, y1, x2, y2 = get_crop_region(mask, padding)
+
+        # Determine crop + processing size
+        if force_uniform_tiles:
+            process_w = self.round_to_multiple(tile_width + padding, 8)
+            process_h = self.round_to_multiple(tile_height + padding, 8)
+            crop_w = x2 - x1
+            crop_h = y2 - y1
+            crop_ratio = crop_w / crop_h if crop_h != 0 else 1.0
+            process_ratio = process_w / process_h if process_h != 0 else 1.0
+            if crop_ratio > process_ratio:
+                target_w = crop_w
+                target_h = round(crop_w / process_ratio) if process_ratio != 0 else crop_h
+            else:
+                target_w = round(crop_h * process_ratio)
+                target_h = crop_h
+            (x1, y1, x2, y2), _ = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
+            target_w = process_w
+            target_h = process_h
+        else:
+            crop_w = x2 - x1
+            crop_h = y2 - y1
+            target_w = max(8, math.ceil(crop_w / 8) * 8)
+            target_h = max(8, math.ceil(crop_h / 8) * 8)
+            (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
+
+        # Actual extracted size before resizing (for blending)
+        extracted_width = x2 - x1
+        extracted_height = y2 - y1
+
+        # Extract tile and resize to processing size
+        tile = image[:, y1:y2, x1:x2, :]
+        tile_pil = tensor_to_pil(tile, 0)
+        if tile_pil.size != (target_w, target_h):
+            tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)
+
+        tile_tensor = pil_to_tensor(tile_pil)
+        if image.is_cuda:
+            tile_tensor = tile_tensor.cuda()
+
+        return tile_tensor, x1, y1, extracted_width, extracted_height
+
+    def extract_batch_tile_with_padding(self, images: torch.Tensor, x: int, y: int,
+                                        tile_width: int, tile_height: int, padding: int,
+                                        force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
+        """Extract a tile region for the entire batch and resize to USDU logic.
+
+        - Computes a single crop region from a mask at (x,y,w,h) with padding
+        - force_uniform_tiles controls target processing size logic
+        - Returns a batched tensor [B,H',W',C] and crop origin/size for blending
+        """
+        batch, h, w, _ = images.shape
+
+        # Create mask and compute initial padded crop region (same for all images)
+        mask = Image.new('L', (w, h), 0)
+        draw = ImageDraw.Draw(mask)
+        draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
+        x1, y1, x2, y2 = get_crop_region(mask, padding)
+
+        # Determine crop + processing size
+        if force_uniform_tiles:
+            process_w = self.round_to_multiple(tile_width + padding, 8)
+            process_h = self.round_to_multiple(tile_height + padding, 8)
+            crop_w = x2 - x1
+            crop_h = y2 - y1
+            crop_ratio = crop_w / crop_h if crop_h != 0 else 1.0
+            process_ratio = process_w / process_h if process_h != 0 else 1.0
+            if crop_ratio > process_ratio:
+                target_w = crop_w
+                target_h = round(crop_w / process_ratio) if process_ratio != 0 else crop_h
+            else:
+                target_w = round(crop_h * process_ratio)
+                target_h = crop_h
+            (x1, y1, x2, y2), _ = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
+            target_w = process_w
+            target_h = process_h
+        else:
+            crop_w = x2 - x1
+            crop_h = y2 - y1
+            target_w = max(8, math.ceil(crop_w / 8) * 8)
+            target_h = max(8, math.ceil(crop_h / 8) * 8)
+            (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
+
+        extracted_width = x2 - x1
+        extracted_height = y2 - y1
+
+        # Slice batch region
+        tiles = images[:, y1:y2, x1:x2, :]
+
+        # Resize each tile to target size
+        resized_tiles = []
+        for i in range(batch):
+            tile_pil = tensor_to_pil(tiles, i)
+            if tile_pil.size != (target_w, target_h):
+                tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)
+            resized_tiles.append(pil_to_tensor(tile_pil))
+        tile_batch = torch.cat(resized_tiles, dim=0)
+
+        if images.is_cuda:
+            tile_batch = tile_batch.cuda()
+
+        return tile_batch, x1, y1, extracted_width, extracted_height
+
+    def process_tile(self, tile_tensor: torch.Tensor, model, positive, negative, vae,
+                     seed: int, steps: int, cfg: float, sampler_name: str, 
+                     scheduler: str, denoise: float, tiled_decode: bool = False,
+                     batch_idx: int = 0, region: Tuple[int, int, int, int] = None,
+                     image_size: Tuple[int, int] = None) -> torch.Tensor:
+        """Process a single tile through SD sampling. 
+        Note: positive and negative should already be pre-sliced for the current batch_idx."""
+        debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}")
+        
+        
+        # Import here to avoid circular dependencies
+        from nodes import common_ksampler, VAEEncode, VAEDecode
+        
+        # Try to import tiled VAE nodes if available
+        try:
+            from nodes import VAEEncodeTiled, VAEDecodeTiled
+            tiled_vae_available = True
+        except ImportError:
+            tiled_vae_available = False
+            if tiled_decode:
+                debug_log("Tiled VAE nodes not available, falling back to standard VAE")
+        
+        # Convert to PIL and back to ensure clean tensor without gradient tracking
+        tile_pil = tensor_to_pil(tile_tensor, 0)
+        clean_tensor = pil_to_tensor(tile_pil)
+        
+        # Ensure tensor is detached and doesn't require gradients
+        clean_tensor = clean_tensor.detach()
+        if hasattr(clean_tensor, 'requires_grad_'):
+            clean_tensor.requires_grad_(False)
+        
+        # Move to correct device
+        if tile_tensor.is_cuda:
+            clean_tensor = clean_tensor.cuda()
+            clean_tensor = clean_tensor.detach()  # Detach again after device transfer
+        
+        # Clone conditioning per tile (shares models, clones hints for cropping)
+        positive_tile = clone_conditioning(positive, clone_hints=True)
+        negative_tile = clone_conditioning(negative, clone_hints=True)
+        
+        # Crop conditioning to tile region if provided (assumes hints at image resolution)
+        if region is not None and image_size is not None:
+            init_size = image_size  # (width, height) of full image
+            canvas_size = image_size
+            tile_size = (tile_tensor.shape[2], tile_tensor.shape[1])  # (width, height)
+            w_pad = 0  # No extra pad needed; region already includes padding
+            h_pad = 0
+            positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
+            negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
+        else:
+            # No region cropping needed, use cloned conditioning as-is
+            positive_cropped = positive_tile
+            negative_cropped = negative_tile
+        
+        # Encode to latent (always non-tiled, matching original node)
+        latent = VAEEncode().encode(vae, clean_tensor)[0]
+        
+        # Sample with model patch cropping parity (ControlNet patch hints)
+        if region is not None and image_size is not None:
+            model_ctx = crop_model_cond(
+                model,
+                region,
+                image_size,
+                image_size,
+                (clean_tensor.shape[2], clean_tensor.shape[1]),
+            )
+        else:
+            model_ctx = nullcontext(model)
+        with model_ctx as model_for_sampling:
+            samples = common_ksampler(
+                model_for_sampling, seed, steps, cfg, sampler_name, scheduler,
+                positive_cropped, negative_cropped, latent, denoise=denoise
+            )[0]
+        
+        # Decode back to image
+        if tiled_decode and tiled_vae_available:
+            image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
+        else:
+            image = VAEDecode().decode(vae, samples)[0]
+        
+        return image
+
+    def process_tiles_batch(self, tile_batch: torch.Tensor, model, positive, negative, vae,
+                            seed: int, steps: int, cfg: float, sampler_name: str,
+                            scheduler: str, denoise: float, tiled_decode: bool,
+                            region: Tuple[int, int, int, int], image_size: Tuple[int, int]) -> torch.Tensor:
+        """Process a batch of tiles together (USDU behavior).
+
+        tile_batch: [B, H, W, C]
+        Returns image batch tensor [B, H, W, C]
+        """
+        # Import locally to avoid circular deps
+        from nodes import common_ksampler, VAEEncode, VAEDecode
+        try:
+            from nodes import VAEEncodeTiled, VAEDecodeTiled
+            tiled_vae_available = True
+        except ImportError:
+            tiled_vae_available = False
+
+        # Detach and move device
+        clean = tile_batch.detach()
+        if hasattr(clean, 'requires_grad_'):
+            clean.requires_grad_(False)
+        if tile_batch.is_cuda:
+            clean = clean.cuda().detach()
+
+        # Clone/crop conditioning once for the region
+        positive_tile = clone_conditioning(positive, clone_hints=True)
+        negative_tile = clone_conditioning(negative, clone_hints=True)
+
+        init_size = image_size
+        canvas_size = image_size
+        tile_size = (clean.shape[2], clean.shape[1])  # (W,H)
+        w_pad = 0
+        h_pad = 0
+        positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
+        negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
+
+        # Encode -> Sample -> Decode
+        latent = VAEEncode().encode(vae, clean)[0]
+        with crop_model_cond(model, region, image_size, image_size, tile_size) as model_for_sampling:
+            samples = common_ksampler(
+                model_for_sampling, seed, steps, cfg, sampler_name, scheduler,
+                positive_cropped, negative_cropped, latent, denoise=denoise
+            )[0]
+        if tiled_decode and tiled_vae_available:
+            image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
+        else:
+            image = VAEDecode().decode(vae, samples)[0]
+
+        return image
+
+    def create_tile_mask(self, image_width: int, image_height: int,
+                        x: int, y: int, tile_width: int, tile_height: int, 
+                        mask_blur: int) -> Image.Image:
+        """Create a mask for blending tiles - matches Ultimate SD Upscale approach.
+        
+        Creates a black image with a white rectangle at the tile position,
+        then applies blur to create soft edges.
+        """
+        # Create a full-size mask matching the image dimensions
+        mask = Image.new('L', (image_width, image_height), 0)  # Black background
+        
+        # Draw white rectangle at tile position
+        draw = ImageDraw.Draw(mask)
+        draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
+        
+        # Apply blur to soften edges
+        if mask_blur > 0:
+            mask = mask.filter(ImageFilter.GaussianBlur(mask_blur))
+        
+        return mask
+
+    def blend_tile(self, base_image: Image.Image, tile_image: Image.Image,
+                  x: int, y: int, extracted_size: Tuple[int, int], 
+                  mask: Image.Image, padding: int) -> Image.Image:
+        """Blend a processed tile back into the base image using Ultimate SD Upscale's exact approach.
+        
+        This follows the exact method from ComfyUI_UltimateSDUpscale/modules/processing.py
+        """
+        extracted_width, extracted_height = extracted_size
+        
+        # Debug logging (uncomment if needed)
+        # debug_log(f"[Blend] Placing tile at ({x}, {y}), size: {extracted_width}x{extracted_height}")
+        
+        # Calculate the crop region that was used for extraction
+        crop_region = (x, y, x + extracted_width, y + extracted_height)
+        
+        # The mask is already full-size, no need to crop
+        
+        # Resize the processed tile back to the extracted size
+        if tile_image.size != (extracted_width, extracted_height):
+            tile_resized = tile_image.resize((extracted_width, extracted_height), Image.LANCZOS)
+        else:
+            tile_resized = tile_image
+        
+        # Follow Ultimate SD Upscale blending approach:
+        # Put the tile into position
+        image_tile_only = Image.new('RGBA', base_image.size)
+        image_tile_only.paste(tile_resized, crop_region[:2])
+        
+        # Add the mask as an alpha channel
+        # Must make a copy due to the possibility of an edge becoming black
+        temp = image_tile_only.copy()
+        temp.putalpha(mask)  # Use the full image mask
+        image_tile_only.paste(temp, image_tile_only)
+        
+        # Add back the tile to the initial image according to the mask in the alpha channel
+        result = base_image.convert('RGBA')
+        result.alpha_composite(image_tile_only)
+        
+        # Convert back to RGB
+        return result.convert('RGB')
+
+    def _slice_conditioning(self, positive, negative, batch_idx):
+        """Helper to slice conditioning for a specific batch index."""
+        # Clone and slice conditioning properly, including ControlNet hints
+        positive_sliced = clone_conditioning(positive)
+        negative_sliced = clone_conditioning(negative)
+        
+        for cond_list in [positive_sliced, negative_sliced]:
+            for i in range(len(cond_list)):
+                emb, cond_dict = cond_list[i]
+                if emb.shape[0] > 1:
+                    cond_list[i][0] = emb[batch_idx:batch_idx+1]
+                if 'control' in cond_dict:
+                    control = cond_dict['control']
+                    while control is not None:
+                        hint = control.cond_hint_original
+                        if hint.shape[0] > 1:
+                            control.cond_hint_original = hint[batch_idx:batch_idx+1]
+                        control = control.previous_controlnet
+                if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
+                    cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1]
+        
+        return positive_sliced, negative_sliced
+
+    def slice_conditioning(self, positive, negative, batch_idx):
+        """Public conditioning-slice API."""
+        return self._slice_conditioning(positive, negative, batch_idx)
+
+    def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image,
+                               model, positive, negative, vae, seed, steps, cfg,
+                               sampler_name, scheduler, denoise, tile_width, tile_height,
+                               padding, mask_blur, image_width, image_height, force_uniform_tiles,
+                               tiled_decode, batch_idx: int = 0):
+        """Process a single tile and blend it into the result image."""
+        x, y = tile_pos
+        
+        # Extract and process tile
+        tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
+            upscaled_image, x, y, tile_width, tile_height, padding, force_uniform_tiles
+        )
+        
+        processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae,
+                                         seed, steps, cfg, sampler_name, 
+                                         scheduler, denoise, tiled_decode, batch_idx=batch_idx,
+                                         region=(x1, y1, x1 + ew, y1 + eh), image_size=(image_width, image_height))
+        
+        # Convert and blend
+        processed_pil = tensor_to_pil(processed_tile, 0)
+        # Create mask for this specific tile (no cache here; only used in single-tile path)
+        tile_mask = self.create_tile_mask(image_width, image_height, x, y, tile_width, tile_height, mask_blur)
+        # Use extraction position and size for blending
+        result_image = self.blend_tile(result_image, processed_pil, 
+                                     x1, y1, (ew, eh), tile_mask, padding)
+        
+        return result_image
+
+    def process_and_blend_tile(
+        self,
+        tile_idx,
+        tile_pos,
+        upscaled_image,
+        result_image,
+        model,
+        positive,
+        negative,
+        vae,
+        seed,
+        steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        denoise,
+        tile_width,
+        tile_height,
+        padding,
+        mask_blur,
+        image_width,
+        image_height,
+        force_uniform_tiles,
+        tiled_decode,
+        batch_idx: int = 0,
+    ):
+        """Public tile-processing API."""
+        return self._process_and_blend_tile(
+            tile_idx,
+            tile_pos,
+            upscaled_image,
+            result_image,
+            model,
+            positive,
+            negative,
+            vae,
+            seed,
+            steps,
+            cfg,
+            sampler_name,
+            scheduler,
+            denoise,
+            tile_width,
+            tile_height,
+            padding,
+            mask_blur,
+            image_width,
+            image_height,
+            force_uniform_tiles,
+            tiled_decode,
+            batch_idx=batch_idx,
+        )
+
+    def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles,
+                                  model, positive, negative, vae, seed, steps, cfg, sampler_name,
+                                  scheduler, denoise, tiled_decode, tile_width, tile_height, padding,
+                                  width, height, force_uniform_tiles, sliced_conditioning_cache):
+        """Process a single tile."""
+        # Calculate which image and tile this corresponds to
+        batch_idx = global_idx // num_tiles_per_image
+        tile_idx = global_idx % num_tiles_per_image
+        
+        # Skip if batch_idx is out of range
+        if batch_idx >= upscaled_image.shape[0]:
+            debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}")
+            return None
+        
+        # Get or create sliced conditioning for this batch index
+        if batch_idx not in sliced_conditioning_cache:
+            positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx)
+            sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced)
+        else:
+            positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx]
+        
+        x, y = all_tiles[tile_idx]
+        
+        # Extract tile from the specific image in the batch
+        tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
+            upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding, force_uniform_tiles
+        )
+        
+        # Process tile through SD with the exact seed (USDU parity)
+        image_seed = seed
+        processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae,
+                                         image_seed, steps, cfg, sampler_name,
+                                         scheduler, denoise, tiled_decode, batch_idx=batch_idx,
+                                         region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height))
+        
+        return {
+            'tile': processed_tile,
+            'global_idx': global_idx,
+            'batch_idx': batch_idx,
+            'tile_idx': tile_idx,
+            'x': x1,
+            'y': y1,
+            'extracted_width': ew,
+            'extracted_height': eh
+        }