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.

.github/workflows/publish_action.yml

@@ -7,14 +7,18 @@ on:
     paths:
       - "pyproject.toml"
 
+permissions:
+  issues: write
+
 jobs:
   publish-node:
     name: Publish Custom Node to registry
     runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'robertvoy' }}
     steps:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
+        uses: Comfy-Org/publish-node-action@v1
         with:
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.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

@@ -1,36 +1,60 @@
 <div align="center">
-<img width="320" src="https://github.com/user-attachments/assets/537d38cc-2e34-4527-bca7-0d79f4120cce" />
+<img width="250" src="https://github.com/user-attachments/assets/533bb98d-0c4a-499f-9bca-5c937e361087" />
 <br><br>
 <a href="https://www.youtube.com/watch?v=p6eE3IlAbOs"><img src="https://img.shields.io/badge/Video_Tutorial-grey?style=flat&logo=youtube&logoColor=white" alt="Video Tutorial"></a>
 <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>
 
-> **A powerful extension for ComfyUI that enables distributed and parallel processing across multiple GPUs and machines. Generate more images and videos and accelerate your upscaling workflows by leveraging all available GPU resources in your network.**
+> **A powerful extension for ComfyUI that enables distributed and parallel processing across multiple GPUs and machines. Generate more images and videos and accelerate your upscaling workflows by leveraging all available GPU resources in your network and cloud.**
 
 ![Clipboard Image (7)](https://github.com/user-attachments/assets/66aaadef-f195-48a1-a368-17dd0dae477d)
 
 ---
 
-## Key Features
+## Key Features
 
 #### Parallel Workflow Processing
-- Run workflows on multiple GPUs simultaneously with varied seeds
+- Run your workflow on multiple GPUs simultaneously with varied seeds, collect results on the master
 - Scale output with more workers
 - Supports images and videos
 
 #### 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
+
+<img width="200" align="right" alt="ComfyUI_temp_khvcc_00034_@0 25x" src="https://github.com/user-attachments/assets/651e4912-7c23-4e32-bd88-250f5175e129" />
+
+ComfyUI Distributed supports three types of workers:
+
+- **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
+- **Remote Workers** - GPUs on other computers in your network
+- **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
+
+> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
 
 ---
 
@@ -38,7 +62,8 @@
 
 - ComfyUI
 > Note: Desktop app not currently supported
-- Multiple GPUs
+- Multiple NVIDIA GPUs
+> No additional GPUs? Use [Cloud Workers](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/worker-setup-guides.md#cloud-workers)
 - That's it
 
 ---
@@ -50,26 +75,36 @@
    git clone https://github.com/robertvoy/ComfyUI-Distributed.git
    ```
    
-2. **Restart ComfyUI** - If you'll be using remote/cloud workers, add `--enable-cors-header` to your launch arguments on the master
+2. **Restart ComfyUI**
+   - If you'll be using remote/cloud workers, add `--enable-cors-header` to your launch arguments on the master
 
-3. Read the [setup guide](https://github.com/robertvoy/ComfyUI-Distributed/blob/feature/cloud-workers/docs/worker-setup-guides.md) for adding workers
+3. Read the [setup guide](/docs/worker-setup-guides.md) for adding workers
+
+---
+
+## Official Sponsor
+
+[<img width="1500" height="339" src="https://github.com/user-attachments/assets/c5f75e1f-3e19-4c57-b05d-151311cd1cf0" />](https://get.runpod.io/0bw29uf3ug0p)
+
+Join Runpod with [this link](https://get.runpod.io/0bw29uf3ug0p) and unlock a special bonus.
 
 ---
 
 ## 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.
@@ -87,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 Upscaling
-Accelerate the process 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
@@ -100,6 +135,61 @@ Accelerate the process by distributing tiles across multiple workers, with speed
 4. Configure tile settings
 5. Enable workers for faster processing
 
+### Distributed Video Upscaling
+Accelerate Ultimate SD Upscaler by distributing video tiles across multiple workers, with speed scaling as you add more GPUs.
+
+![Video Upscaler workflow](https://github.com/user-attachments/assets/3c3d61b1-0b5f-422e-8c58-7c1555fed765)
+
+> [Download workflow](/workflows/distributed-upscale-video.json)
+
+1. Load your video
+2. Optional: upscale with ESRGAN or similar
+3. Connect to **Ultimate SD Upscale Distributed**
+4. Configure tile settings
+5. Use RES4LYF (bong/res2) to get better results
+6. Enable workers for faster processing
+
+> You can run this workflow entirely on Runpod with minimal setup. [Check out the guide here.](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/video-upscaler-runpod-preset.md)
+
+---
+
+## 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
@@ -121,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>
@@ -129,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.
@@ -143,4 +238,14 @@ This software is provided "as is" without any warranties, express or implied, in
 ---
 
 ## Support the Project
-If my custom nodes have added value to your workflow, consider fueling future development with a coffee! Your support helps keep this project thriving. Buy me a coffee at: https://buymeacoffee.com/robertvoy
+
+<img width="200" align="right" src="https://github.com/user-attachments/assets/84291921-c44e-4556-94f2-a3b16500f4f9" />
+
+If my custom nodes have added value to your workflow, consider fueling future development with a coffee!
+
+Your support helps keep this project thriving.
+
+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/model-download-script.md

@@ -1,15 +1,14 @@
 ## Automating ComfyUI Model Downloads
+> This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM).
 
-This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM).
+1. In ComfyUI (on your local machine), export your workflow as an API workflow
+2. Copy the below prompt and upload the API workflow to an LLM **that has access to the internet**
+
+<details>
+<summary><strong>📋 Click to expand the full prompt</strong></summary>
 
-1. Launched the Pod with these Environment Variables:
-  - CIVITAI_API_TOKEN ([get your token here](https://civitai.com/user/account))
-  - HF_API_TOKEN ([get your token here](https://huggingface.co/settings/tokens))
-2. In ComfyUI, export your workflow as an API workflow
-3. Copy the below prompt and upload the API workflow to a LLM **that has access to the internet**
 ```
 Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths:
-
 base_path: /workspace/ComfyUI
 checkpoints: models/checkpoints/
 clip: models/clip/
@@ -25,32 +24,30 @@ text_encoders: models/text_encoders/
 unet: models/unet/
 upscale_models: models/upscale_models/
 vae: models/vae/
-
 ---
-
 Important:
 Make sure you find the correct URLs for the models online.
 Use comfy cli to download the models: `comfy model download --url <URL> [--relative-path <PATH>] [--set-civitai-api-token <TOKEN>] [--set-hf-api-token <TOKEN>]`
 Make sure you add `--set-civitai-api-token $CIVITAI_API_TOKEN` for CivitAI download and `--set-hf-api-token $HF_API_TOKEN` for Hugging Face downloads.
-
 ---
-
 Example:
-
 #!/bin/bash
-
 # Download from CivitAI
 comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
-
 # Download model from Hugging Face
 comfy model download --url https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors --relative-path /workspace/ComfyUI/models/unet --set-hf-api-token $HF_API_TOKEN
-
 # If a model in the workflow was in a subfolder
 comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints/SDXL --set-civitai-api-token $CIVITAI_API_TOKEN
-
 ```
-4. Review the LLMs output to make sure all download links are correct and save it as a .sh file, for example `download_models.sh`
-5. Upload it onto your Runpod instance, into `/workspace`
+
+</details>
+
+3. Review the LLMs output to make sure all download links are correct and save it as a .sh file, for example `download_models.sh`
+4. Launch the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=ak218p52) with these Environment Variables:
+   - `CIVITAI_API_TOKEN`: [get your token here](https://civitai.com/user/account)
+   - `HF_API_TOKEN`: [get your token here](https://huggingface.co/settings/tokens)
+5. Upload the .sh file to your Runpod instance, into `/workspace`
 6. Then run these commands:
    - `chmod 755 /workspace/download_models.sh`
    - `/workspace/download_models.sh`
+7. Confirm each model name (sometimes you might need to rename them to match the name on your local machine)

docs/video-upscaler-runpod-preset.md

@@ -0,0 +1,20 @@
+![Clipboard Image](https://github.com/user-attachments/assets/5dc5224f-3f47-442c-b94a-116afeb28132)
+
+**Accelerated Creative Video Upscaler On Runpod:**
+
+1. Use the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=0bw29uf3ug0p) template.
+2. Filter instances by CUDA 12.8 (add filter in Additional Filters at the top of the page).
+3. Choose 4x 5090s
+4. Press Edit Template to configure the pod's Environment Variables:
+	- CIVITAI_API_TOKEN: Not necessary for this workflow.
+	- HF_API_TOKEN: [get your token here](https://huggingface.co/settings/tokens)
+	- SAGE_ATTENTION: optional optimisation (set to true/false). Recommended for this workflow.
+	- PRESET_VIDEO_UPSCALER: set to true. This will download everything you need.
+5. Deploy your pod.
+6. Once pod setup is complete, connect to ComfyUI running on your pod.
+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. [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

@@ -1,11 +1,19 @@
-## ComfyUI Distributed Worker Setup Guide
-
-ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs, whether they're in the same machine, on your local network, or in the cloud.
+## Worker Setup Guide
 
 **Master**: The main ComfyUI instance that coordinates and distributes work. This is where you load workflows, manage the queue, and view results.
 
 **Worker**: A ComfyUI instance that receives and processes tasks from the master. Workers handle just the GPU computation and send results back to the master. You can have multiple workers connected to a single master, each utilizing their own GPU.
 
+<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
@@ -13,8 +21,14 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs
 - **Cloud workers**: GPUs hosted on cloud services like Runpod
 
 ## Local workers
+
+<img align="right" width="200" src="https://github.com/user-attachments/assets/651e4912-7c23-4e32-bd88-250f5175e129" />
+
 > These are added automatically on first launch, but you can add them manually if you need to.
 
+
+📺 [Watch Tutorial](https://youtu.be/p6eE3IlAbOs?si=K7Km0_flmPHwRQwz&t=43)
+
 1. **Open** the Distributed GPU panel.
 2. **Click** "Add Worker" in the UI.
 3. **Configure** your local worker:
@@ -25,8 +39,14 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs
 4. **Save** and  launch the local worker.
 
 ## Remote workers
+
+<img align="right" width="200"  src="https://github.com/user-attachments/assets/84291921-c44e-4556-94f2-a3b16500f4f9" />
+
+
 > ComfyUI instances running on completely different computers on your network. These allow you to harness GPU power from other machines. Remote workers must be manually started on their respective computers and are connected via IP address.
 
+📺 [Watch Tutorial](https://youtu.be/p6eE3IlAbOs?si=Oxj3EzPyf4jKDvfG&t=140)
+
 **On the Remote Worker Machine:**
 1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!**
    - This ComfyUI instance will serve as a worker for your main master.
@@ -38,47 +58,66 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs
 3. **Open** the ComfyUI port (e.g., 8188) and any additional worker ports (e.g., 8189, 8190) in the firewall.
   
 **On the Main Machine:**
-1. **Open** the Distributed GPU panel (sidebar on the left).
-2. **Click** "Add Worker."
-3. **Choose** "Remote".
-4. **Configure** your remote worker:
+1. **Launch** ComfyUI with `--enable-cors-header` launch argument.
+2. **Open** the Distributed GPU panel (sidebar on the left).
+3. **Click** "Add Worker."
+4. **Choose** "Remote".
+5. **Configure** your remote worker:
    - **Name**: A descriptive name for the worker (e.g., "Server Rack GPU 0")
    - **Host**: The remote worker's IP address.
    - **Port**: The port number used when launching ComfyUI on the remote master/worker (e.g., 8188).
-5. **Save** the remote worker configuration.
+6. **Save** the remote worker configuration.
   
 ## Cloud workers
-> ComfyUI instances running on a cloud service like [Runpod](https://get.runpod.io/0bw29uf3ug0p). 
+
+<img align="right" width="200"  src="https://github.com/user-attachments/assets/a053f3ae-22f0-4e1c-8f2e-f26a1f660adf" />
+
+> ComfyUI instances running on a cloud service like Runpod. 
 
 ### Deploy Cloud Worker on Runpod
 
+📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
+
 **On Runpod:**
+> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!**
+
 1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
-2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later.
-3. Now go to Pods and find a suitable GPU for your workflows. 
-4. Choose the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=ak218p52) template and make sure your network drive is mounted.
-> To use the ComfyUI Distributed Pod template, you will need to filter instances by CUDA 12.8.
-6. Launch your pod.
-7. Access your pod using JupyterLabs.
-8. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod).
-> You can use [this guide](model-download-script.md) to make this process easy for you. It will generate a shell script that will automatically download the models you need for a given workflow.
-8. If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you git clone ComfyUI-Distributed into custom_nodes. ⚠️ **Required!**
-9. Download any additional custom nodes using the ComfyUI Manager.
+2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
+3. Use the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=0bw29uf3ug0p) template.
+4. Make sure your Network Volume is mounted and choose a suitable GPU.
+> ⚠️ To use the ComfyUI Distributed Pod template, you will need to filter instances by CUDA 12.8 (add filter in Additional Filters).
+6. Press Edit Template to configure the pod's Environment Variables:
+	- CIVITAI_API_TOKEN: [get your token here](https://civitai.com/user/account)
+	- HF_API_TOKEN: [get your token here](https://huggingface.co/settings/tokens)
+	- SAGE_ATTENTION: optional optimisation (set to true/false)
+5. Deploy your pod.
+6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
+7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
+```
+# Download from CivitAI
+comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
+# Download model from Hugging Face
+comfy model download --url https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors --relative-path /workspace/ComfyUI/models/unet --set-hf-api-token $HF_API_TOKEN
+```
+> ℹ️ Use [this guide](model-download-script.md) to make this process easy. It will generate a shell script that automatically downloads the models for a given workflow.
+9. Access ComfyUI through the Runpod URL.
+10. Download any additional custom nodes you need using the ComfyUI Manager.
 
 **On the Main Machine:**
-1. Launch a Cloudflare tunnel
+1. **Launch** a Cloudflare tunnel.
    - Download from here: [https://github.com/cloudflare/cloudflared/releases](https://github.com/cloudflare/cloudflared/releases)
 	- Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188`
-> Cloudflare tunnels create secure connections without exposing ports directly to the internet
-2. Copy the Cloudflare address
-3. **Open** the Distributed GPU panel (sidebar on the left).
-4. **Edit** the Master's host address and replace it with the Cloudflare address.
-   - **Click** "Add Worker."
-   - **Choose** "Cloud".
-   - **Configure** your cloud worker:
-     - **Host**: The Runpod address
-     - **Port**: 443
-   - **Save** the remote worker configuration.
+> ℹ️ Cloudflare tunnels create secure connections without exposing ports directly to the internet and are required for Cloud Workers.
+2. **Copy** the Cloudflare address
+3. **Launch** ComfyUI with `--enable-cors-header` launch argument.
+4. **Open** the Distributed GPU panel (sidebar on the left).
+5. **Edit** the Master's settings to change the host address to the Cloudflare address.
+6. **Click** "Add Worker."
+7. **Choose** "Cloud".
+8. **Configure** your cloud worker:
+	- **Host**: The ComfyUI Runpod address. For example: `wcegfo9tbbml9l-8188.proxy.runpod.net`
+	- **Port**: 443
+9. **Save** the remote worker configuration.
 
 ---
 
@@ -89,7 +128,7 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs
    - If your cloud platform doesn't provide a secure connection, use Cloudflare to create a tunnel for the worker. Each GPU needs their own tunnel for their respective port.
 	- For example: `./cloudflared tunnel --url http://localhost:8188`
 1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!**
-2. **Add** workers in the UI panel if the remote machine has more than one GPU.
+2. **Add** workers in the UI panel if the cloud machine has more than one GPU.
    - Make sure that they also have `--listen` set in `Extra Args`.
    - Then launch them.
   
@@ -98,11 +137,12 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs
    - Download from here: [https://github.com/cloudflare/cloudflared/releases](https://github.com/cloudflare/cloudflared/releases)
    - Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188`
 2. **Copy** the Cloudflare address
-3. **Open** the Distributed GPU panel (sidebar on the left).
-4. **Edit** the Master's host address and replace it with the Cloudflare address.
-5. **Click** "Add Worker."
-6. **Choose** "Cloud".
-7. **Configure** your cloud worker:
+3. **Launch** ComfyUI with `--enable-cors-header` launch argument.
+4. **Open** the Distributed GPU panel (sidebar on the left).
+5. **Edit** the Master's host address and replace it with the Cloudflare address.
+6. **Click** "Add Worker."
+7. **Choose** "Cloud".
+8. **Configure** your cloud worker:
    - **Host**: The remote worker's IP address/domain
    - **Port**: 443
-8. **Save** the remote worker configuration.
+9. **Save** the remote worker configuration.

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 and remotely "
-version = "1.0.1"
+description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
+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"
@@ -12,4 +15,8 @@ Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
 [tool.comfy]
 PublisherId = "robertvoy"
 DisplayName = "ComfyUI-Distributed"
-Icon = ""
+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