Bump version from 1.4.7 to 1.5.0

* feat: migrate distributed nodes to native ComfyUI V3 API

* fix: hide Distributed Value backing fields in Nodes 2.0

Use native widget visibility without converting backing fields into sockets. Preserve widget types, serializers and options identity; cover edits, reload, worker refresh and linked widgets.

.github/FUNDING.yml

@@ -1,4 +1,3 @@
 # These are supported funding model platforms
 
 github: robertvoy
-buy_me_a_coffee: robertvoy

.github/workflows/publish_action.yml

@@ -19,6 +19,6 @@ jobs:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@v1
+        uses: Comfy-Org/publish-node-action@main
         with:
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.gitignore

@@ -1,4 +1,9 @@
 bin/
 logs/
 gpu_config.json
-__pycache__/
+__pycache__/
+**/__pycache__/
+*.py[cod]
+node_modules/
+npm-debug.log*
+AGENTS.md

.nvmrc

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

README.md

@@ -51,7 +51,6 @@ ComfyUI Distributed supports three types of workers:
 ## Requirements
 
 - ComfyUI
-> Note: Desktop app not currently supported
 - 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
@@ -92,8 +91,9 @@ Generate multiple images in the time it takes to generate one. Each worker uses
 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!
+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.
@@ -116,7 +116,7 @@ Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, w
 
 ![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
@@ -147,13 +147,40 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
 Control your distributed cluster programmatically without opening the browser.
 
 * **Endpoint:** `POST /distributed/queue`
-* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
+* **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
 * **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
 
 > **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
 
 ---
 
+## 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 Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
+| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
+| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
+
+---
+
 ## FAQ
 
 <details>
@@ -168,12 +195,12 @@ No, it does not speed up the generation of a single image or video. Instead, it
 
 <details>
 <summary>Does it work with the ComfyUI desktop app?</summary>
-Currently, it is not compatible with the ComfyUI desktop app.
+Yes, it does now.
 </details>
 
 <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>
@@ -181,6 +208,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.
@@ -204,3 +236,6 @@ Your support helps keep this project thriving.
 
 Buy me a coffee at: https://buymeacoffee.com/robertvoy
 
+
+
+

__init__.py

@@ -1,65 +1,22 @@
-import os
-import sys
+"""ComfyUI-Distributed's native V3 extension entrypoint."""
+from comfy_api.v0_0_2 import ComfyExtension, io
 
-# Add the directory to Python path
-sys.path.append(os.path.dirname(__file__))
+from .nodes.v3 import NODES
+from .runtime.bootstrap import initialize
 
-# 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
+WEB_DIRECTORY = './web'
 
-# 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
+class DistributedExtension(ComfyExtension):
+    async def on_load(self) -> None:
+        initialize()
 
-# Import distributed upscale nodes
-from .distributed_upscale import (
-    NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
-    NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
-)
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return list(NODES)
 
-WEB_DIRECTORY = "./web"
 
-ensure_config_exists()
+async def comfy_entrypoint() -> DistributedExtension:
+    return DistributedExtension()
 
-# 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__ = ['comfy_entrypoint', 'WEB_DIRECTORY']

api/__init__.py

@@ -0,0 +1,5 @@
+from . import config_routes  # noqa: F401
+from . import tunnel_routes  # noqa: F401
+from . import worker_routes  # noqa: F401
+from . import job_routes  # noqa: F401
+from . import usdu_routes  # noqa: F401

api/config_routes.py

@@ -0,0 +1,277 @@
+import json
+from contextlib import asynccontextmanager
+
+from aiohttp import web
+import server
+
+try:
+    from ..utils.config import config_transaction, load_config, save_config
+except ImportError:
+    from ..utils.config import load_config
+
+    try:
+        from ..utils.config import save_config
+    except ImportError:
+        def save_config(_config):
+            return True
+
+    @asynccontextmanager
+    async def config_transaction():
+        config = load_config()
+        original_snapshot = json.dumps(config, sort_keys=True)
+        yield config
+        if json.dumps(config, sort_keys=True) != original_snapshot:
+            save_config(config)
+from ..utils.logging import debug_log, log
+from ..utils.network import handle_api_error, normalize_host
+
+
+def _positive_int(value):
+    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
+
+
+@server.PromptServer.instance.routes.get("/distributed/config")
+async def get_config_endpoint(request):
+    config = load_config()
+    return web.json_response(config)
+
+
+@server.PromptServer.instance.routes.post("/distributed/config")
+async def update_config_endpoint(request):
+    """Bulk config update with schema validation."""
+    try:
+        data = await request.json()
+    except Exception as e:
+        return await handle_api_error(request, f"Invalid JSON payload: {e}", 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 in _SETTINGS_FIELDS:
+            validated_settings[key] = value
+        else:
+            validated_root[key] = value
+
+    if errors:
+        return web.json_response({
+            "status": "error",
+            "error": errors,
+            "message": "; ".join(errors),
+        }, status=400)
+
+    try:
+        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})
+    except Exception as e:
+        return await handle_api_error(request, e)
+
+
+@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
+async def queue_status_endpoint(request):
+    """Check if a job queue is initialized."""
+    try:
+        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})
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
+async def update_worker_endpoint(request):
+    try:
+        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"})
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
+
+@server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
+async def delete_worker_endpoint(request):
+    try:
+        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"
+            })
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
+
+@server.PromptServer.instance.routes.post("/distributed/config/update_setting")
+async def update_setting_endpoint(request):
+    """Updates a specific key in the settings object."""
+    try:
+        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."})
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
+
+@server.PromptServer.instance.routes.post("/distributed/config/update_master")
+async def update_master_endpoint(request):
+    """Updates master configuration."""
+    try:
+        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."})
+    except Exception as e:
+        return await handle_api_error(request, e, 400)

api/job_routes.py

@@ -0,0 +1,350 @@
+import json
+import asyncio
+import io
+import os
+import base64
+import binascii
+import time
+
+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
+try:
+    from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
+except ImportError:
+    from .queue_orchestration import orchestrate_distributed_execution
+
+    def ensure_distributed_state():
+        return None
+from .queue_request import parse_queue_request_payload
+
+prompt_server = server.PromptServer.instance
+
+# 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 _decode_image_sync(image_path):
+    """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, expected_hash):
+    """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):
+    """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):
+    """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):
+    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)
+
+        ensure_distributed_state()
+        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"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):
+    debug_log("Received request to clear VRAM.")
+    try:
+        # 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:
+            debug_log(f"Warning during model unload: {e}")
+        
+        try:
+            mm.soft_empty_cache()
+        except Exception as e:
+            debug_log(f"Warning during cache clear: {e}")
+        
+        for _ in range(3):
+            gc.collect()
+        
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        
+        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"Partial VRAM clear completed with warning: {e}")
+        return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"})
+
+
+@server.PromptServer.instance.routes.post("/distributed/queue")
+async def distributed_queue_endpoint(request):
+    """Queue a distributed workflow, mirroring the UI orchestration pipeline."""
+    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, prompt_number, worker_count, node_errors = 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,
+            "number": prompt_number,
+            "node_errors": node_errors,
+            "worker_count": worker_count,
+            "auto_prepare_supported": True,
+        })
+    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):
+    """Load an image or video file and return it as base64 data with hash."""
+    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):
+    """Check if a file exists and matches the given hash."""
+    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):
+    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 = 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")
+        image_provided = image_payload is not None
+        audio_provided = audio_payload is not None
+        if not image_provided and not audio_provided:
+            errors.append("expected at least one of image or audio")
+        if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
+            errors.append("image: expected non-empty base64 PNG string")
+        if audio_provided and not isinstance(audio_payload, dict):
+            errors.append("audio: expected object when provided")
+        if not isinstance(is_last, bool):
+            errors.append("is_last: expected boolean")
+        if errors:
+            return await handle_api_error(request, errors, 400)
+
+        try:
+            tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
+            decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
+        except ValueError as exc:
+            return await handle_api_error(request, exc, 400)
+        multi_job_id = job_id.strip()
+        worker_id = worker_id.strip()
+
+        pending = None
+        queue_size = 0
+        deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
+        while pending is None:
+            async with prompt_server.distributed_jobs_lock:
+                pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
+                if pending is not None:
+                    await pending.put(
+                        {
+                            "tensor": tensor,
+                            "worker_id": worker_id,
+                            "image_index": int(batch_idx),
+                            "is_last": is_last,
+                            "audio": decoded_audio,
+                        }
+                    )
+                    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,4 @@
+# Orchestration helpers split out from queue_orchestration.py:
+# - prompt_transform.py: graph pruning + hidden input overrides
+# - media_sync.py: remote media/path normalization
+# - dispatch.py: worker probe + prompt dispatch

api/orchestration/dispatch.py

@@ -0,0 +1,268 @@
+import asyncio
+import json
+import uuid
+
+import aiohttp
+
+from ...utils.logging import debug_log, log
+from ...utils.network import build_worker_url, get_client_session, probe_worker
+try:
+    from ...utils.trace_logger import trace_debug, trace_info
+except ImportError:
+    def trace_debug(*_args, **_kwargs):
+        return None
+
+    def trace_info(*_args, **_kwargs):
+        return None
+
+try:
+    from ..schemas import parse_positive_int
+except ImportError:
+    def parse_positive_int(value, default):
+        try:
+            parsed = int(value)
+            return parsed if parsed > 0 else default
+        except (TypeError, ValueError):
+            return default
+
+_least_busy_rr_index = 0
+
+
+async def worker_is_active(worker):
+    """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):
+    """Ping worker's websocket endpoint to confirm it's reachable."""
+    session = await get_client_session()
+    url = build_worker_url(worker, "/distributed/worker_ws")
+    try:
+        ws = await session.ws_connect(url, heartbeat=20, timeout=3)
+        await ws.close()
+        return True
+    except 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, 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) 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,
+    prompt_obj,
+    workflow_meta,
+    client_id=None,
+    use_websocket=False,
+    trace_execution_id=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:
+        try:
+            await _dispatch_via_websocket(
+                worker_url,
+                {
+                    "prompt": prompt_obj,
+                    "workflow": workflow_meta,
+                },
+                client_id,
+            )
+            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,
+    use_websocket,
+    delegate_master,
+    trace_execution_id=None,
+    probe_concurrency=8,
+):
+    """Probe workers and return (active_workers, updated_delegate_master)."""
+    probe_limit = parse_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,
+    trace_execution_id=None,
+    probe_concurrency=8,
+    probe_timeout=3.0,
+):
+    """Select one worker by queue depth, round-robin among idle workers."""
+    if not workers:
+        return None
+
+    probe_limit = parse_positive_int(probe_concurrency, 8)
+    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

api/orchestration/media_sync.py

@@ -0,0 +1,256 @@
+import asyncio
+import hashlib
+import mimetypes
+import os
+import re
+
+import aiohttp
+
+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, target_separator):
+    """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, trace_execution_id=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()
+    try:
+        async with session.get(url, 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()
+    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},
+            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, prompt_obj, trace_execution_id=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,579 @@
+import json
+from collections import deque
+
+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, target_class):
+        cache_key = (str(start_node_id), target_class)
+        if cache_key in self._upstream_cache:
+            return self._upstream_cache[cache_key]
+
+        visited = set()
+        stack = [str(start_node_id)]
+        while stack:
+            node_id = stack.pop()
+            if node_id in visited:
+                continue
+            visited.add(node_id)
+            inputs = self.inputs_by_node.get(node_id, {})
+            for value in inputs.values():
+                if isinstance(value, list) and len(value) == 2:
+                    upstream_id = str(value[0])
+                    if self.class_by_node.get(upstream_id) == target_class:
+                        self._upstream_cache[cache_key] = True
+                        return True
+                    if upstream_id in self.inputs_by_node:
+                        stack.append(upstream_id)
+
+        self._upstream_cache[cache_key] = False
+        return False
+
+
+def _iter_prompt_nodes(prompt_obj):
+    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 _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):
+            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
+
+
+_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = {
+    "PrimitiveBoolean",
+    "PrimitiveFloat",
+    "PrimitiveInt",
+    "PrimitiveNode",
+    "PrimitiveString",
+}
+
+_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = {
+    "LoadImage",
+}
+
+_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"}
+_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"}
+
+# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the
+# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility
+# only after its connected inputs recursively prove safe.
+_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"}
+_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = {
+    "CreateList": ("inputs.",),
+}
+
+# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily.
+_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None
+
+
+def _get_delegate_master_node_class_mappings():
+    """Return ComfyUI node-class mappings when available."""
+    if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None:
+        return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS
+    try:
+        import nodes as comfy_nodes  # type: ignore
+    except Exception:  # pragma: no cover - depends on ComfyUI runtime imports
+        return {}
+    return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
+
+
+def _get_delegate_master_node_class(class_type):
+    mappings = _get_delegate_master_node_class_mappings()
+    return mappings.get(class_type) if isinstance(mappings, dict) else None
+
+
+def _normalize_delegate_master_return_type(return_type):
+    if return_type is None:
+        return ""
+    return str(return_type).strip().upper()
+
+
+def _delegate_master_type_is_safe_scalar(type_name):
+    return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES
+
+
+def _delegate_master_type_is_safe_config(type_name):
+    return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES
+
+
+def _delegate_master_output_is_safe_scalar(class_type, output_index):
+    """Return True when a registered node output is lightweight config data."""
+    if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES:
+        return True
+    node_class = _get_delegate_master_node_class(class_type)
+    return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else ()
+    try:
+        output_type = return_types[int(output_index)]
+    except (IndexError, TypeError, ValueError):
+        return False
+    return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type))
+
+
+def _get_delegate_master_input_types(class_type):
+    node_class = _get_delegate_master_node_class(class_type)
+    input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None
+    if callable(input_types):
+        try:
+            input_types = input_types()
+        except TypeError:
+            return {}
+    return input_types if isinstance(input_types, dict) else {}
+
+
+def _normalize_delegate_master_input_type(input_spec):
+    if isinstance(input_spec, (list, tuple)) and input_spec:
+        return _normalize_delegate_master_return_type(input_spec[0])
+    return _normalize_delegate_master_return_type(input_spec)
+
+
+def _delegate_master_input_is_safe_scalar(class_type, input_name):
+    """Return True when a registered downstream input expects config data."""
+    for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()):
+        if input_name.startswith(prefix):
+            return True
+    input_types = _get_delegate_master_input_types(class_type)
+    for section_name in ("required", "optional"):
+        section = input_types.get(section_name, {})
+        if isinstance(section, dict) and input_name in section:
+            input_type = _normalize_delegate_master_input_type(section[input_name])
+            return _delegate_master_type_is_safe_config(input_type)
+    return False
+
+
+def _is_delegate_master_always_retained_upstream_node(node):
+    if not isinstance(node, dict):
+        return False
+    class_type = node.get("class_type")
+    return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES
+
+
+def _is_delegate_master_retained_upstream_node(node, output_index=0):
+    """Return True for lightweight upstream nodes safe to keep on the master."""
+    if not isinstance(node, dict):
+        return False
+    class_type = node.get("class_type")
+    if not isinstance(class_type, str):
+        return False
+    return (
+        class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES
+        or class_type.startswith("Primitive")
+        or _delegate_master_output_is_safe_scalar(class_type, output_index)
+    )
+
+
+def _collect_delegate_master_retained_upstream_branch(
+    prompt_obj,
+    node_id,
+    output_index,
+    memo,
+    visiting,
+):
+    """Return safe retained branch nodes, or None when the branch is not safe."""
+    node_id = str(node_id)
+    cache_key = (node_id, output_index)
+    if cache_key in memo:
+        cached = memo[cache_key]
+        return None if cached is None else set(cached)
+    if cache_key in visiting:
+        memo[cache_key] = None
+        return None
+
+    node = prompt_obj.get(node_id)
+    if not _is_delegate_master_retained_upstream_node(node, output_index):
+        memo[cache_key] = None
+        return None
+
+    visiting.add(cache_key)
+    retained = {node_id}
+    inputs = node.get("inputs", {}) if isinstance(node, dict) else {}
+    class_type = node.get("class_type") if isinstance(node, dict) else None
+    for input_name, value in inputs.items():
+        if not (isinstance(value, list) and len(value) == 2):
+            continue
+        if not _delegate_master_input_is_safe_scalar(class_type, input_name):
+            visiting.remove(cache_key)
+            memo[cache_key] = None
+            return None
+        source_id = str(value[0])
+        branch = _collect_delegate_master_retained_upstream_branch(
+            prompt_obj,
+            source_id,
+            value[1],
+            memo,
+            visiting,
+        )
+        if branch is None:
+            visiting.remove(cache_key)
+            memo[cache_key] = None
+            return None
+        retained.update(branch)
+
+    visiting.remove(cache_key)
+    memo[cache_key] = frozenset(retained)
+    return retained
+
+
+def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
+    """Return lightweight upstream nodes needed by kept delegate-master nodes."""
+    connected = set()
+    memo = {}
+    for node_id in start_ids:
+        node = prompt_obj.get(str(node_id)) or {}
+        inputs = node.get("inputs", {})
+        class_type = node.get("class_type") if isinstance(node, dict) else None
+        for input_name, value in inputs.items():
+            if not (isinstance(value, list) and len(value) == 2):
+                continue
+            source_node = prompt_obj.get(str(value[0]))
+            if _is_delegate_master_always_retained_upstream_node(source_node):
+                connected.add(str(value[0]))
+                continue
+            if not _delegate_master_input_is_safe_scalar(class_type, input_name):
+                continue
+            branch = _collect_delegate_master_retained_upstream_branch(
+                prompt_obj,
+                value[0],
+                value[1],
+                memo,
+                set(),
+            )
+            if branch is not None:
+                connected.update(branch)
+    return connected
+
+
+def prune_prompt_for_worker(prompt_obj):
+    """Prune worker prompt to distributed nodes and their upstream dependencies."""
+    collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
+    upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
+    distributed_ids = collector_ids + upscale_ids
+    if not distributed_ids:
+        return prompt_obj
+
+    connected = _find_upstream_nodes(prompt_obj, distributed_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 for node_id in downstream)
+        if has_removed_downstream:
+            original_node = prompt_obj.get(str(dist_id), {})
+            class_type = original_node.get("class_type")
+            inputs = original_node.get("inputs", {})
+            image_connected = class_type != "DistributedCollector" or (
+                isinstance(inputs.get("images"), list)
+                and len(inputs["images"]) == 2
+            )
+            audio_connected = (
+                class_type == "DistributedCollector"
+                and isinstance(inputs.get("audio"), list)
+                and len(inputs["audio"]) == 2
+            )
+
+            if image_connected:
+                preview_id = next_id()
+                pruned_prompt[preview_id] = {
+                    "inputs": {"images": [dist_id, 0]},
+                    "class_type": "PreviewImage",
+                    "_meta": {"title": "Preview Image (auto-added)"},
+                }
+            elif audio_connected:
+                preview_id = next_id()
+                pruned_prompt[preview_id] = {
+                    "inputs": {"audio": [dist_id, 1]},
+                    "class_type": "PreviewAudio",
+                    "_meta": {"title": "Preview Audio (auto-added)"},
+                }
+
+    return pruned_prompt
+
+
+def prepare_delegate_master_prompt(prompt_obj, collector_ids):
+    """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)
+    nodes_to_keep.update(
+        _find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
+    )
+
+    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
+        original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
+        original_images = original_inputs.get("images")
+        if not (isinstance(original_images, list) and len(original_images) == 2):
+            continue
+        placeholder_id = next_id()
+        pruned_prompt[placeholder_id] = {
+            "class_type": "DistributedEmptyImage",
+            "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, prefix):
+    """Create stable per-node job IDs for distributed nodes."""
+    job_map = {}
+    distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
+        "UltimateSDUpscaleDistributed"
+    )
+    for node_id in distributed_nodes:
+        job_map[node_id] = f"{prefix}_{node_id}"
+    return job_map
+
+
+def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
+    """Configure DistributedSeed nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("DistributedSeed"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+        inputs = node.setdefault("inputs", {})
+        inputs["is_worker"] = not is_master
+        if is_master:
+            inputs["worker_id"] = ""
+        else:
+            inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
+
+
+def _override_collector_nodes(
+    prompt_copy,
+    prompt_index,
+    is_master,
+    participant_id,
+    job_id_map,
+    master_url,
+    enabled_json,
+    delegate_master,
+):
+    """Configure DistributedCollector nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("DistributedCollector"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+
+        if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
+            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_upscale_nodes(
+    prompt_copy,
+    prompt_index,
+    is_master,
+    participant_id,
+    job_id_map,
+    master_url,
+    enabled_json,
+):
+    """Configure UltimateSDUpscaleDistributed nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            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.pop("master_url", None)
+            inputs.pop("worker_id", None)
+        else:
+            inputs["master_url"] = master_url
+            inputs["worker_id"] = participant_id
+
+
+def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
+    """Configure DistributedValue nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("DistributedValue"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+        inputs = node.setdefault("inputs", {})
+        inputs["is_worker"] = not is_master
+        if is_master:
+            inputs["worker_id"] = ""
+        else:
+            inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
+
+
+def apply_participant_overrides(
+    prompt_copy,
+    participant_id,
+    enabled_worker_ids,
+    job_id_map,
+    master_url,
+    delegate_master,
+    prompt_index,
+):
+    """Return a prompt copy with hidden inputs configured for master/worker."""
+    is_master = participant_id == "master"
+    worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
+    enabled_json = json.dumps(enabled_worker_ids)
+
+    _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
+    _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
+    _override_collector_nodes(
+        prompt_copy,
+        prompt_index,
+        is_master,
+        participant_id,
+        job_id_map,
+        master_url,
+        enabled_json,
+        delegate_master,
+    )
+    _override_upscale_nodes(
+        prompt_copy,
+        prompt_index,
+        is_master,
+        participant_id,
+        job_id_map,
+        master_url,
+        enabled_json,
+    )
+
+    return prompt_copy

api/queue_orchestration.py

@@ -0,0 +1,418 @@
+import asyncio
+import time
+import uuid
+
+import server
+
+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, build_master_callback_url
+from ..utils.trace_logger import trace_debug
+from .schemas import parse_positive_float, parse_positive_int
+from .orchestration.dispatch import (
+    dispatch_worker_prompt,
+    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,
+)
+
+
+prompt_server = server.PromptServer.instance
+
+
+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."""
+    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()
+
+
+# Initialize top-level distributed queue state at module import time.
+ensure_distributed_state()
+
+
+async def _ensure_distributed_queue(job_id):
+    """Ensure a queue exists for the given distributed job ID."""
+    ensure_distributed_state()
+    async with prompt_server.distributed_jobs_lock:
+        if job_id not in prompt_server.distributed_pending_jobs:
+            prompt_server.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 = parse_positive_int(
+        settings.get("worker_probe_concurrency"),
+        ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
+    )
+    worker_prep_concurrency = parse_positive_int(
+        settings.get("worker_prep_concurrency"),
+        ORCHESTRATION_WORKER_PREP_CONCURRENCY,
+    )
+    media_sync_concurrency = parse_positive_int(
+        settings.get("media_sync_concurrency"),
+        ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
+    )
+    media_sync_timeout_seconds = parse_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,
+    config,
+    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_master_url = build_master_callback_url(
+            worker,
+            config=config,
+            prompt_server_instance=prompt_server,
+        )
+
+        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,
+            worker_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 orchestrate_distributed_execution(
+    prompt_obj,
+    workflow_meta,
+    client_id,
+    enabled_worker_ids=None,
+    delegate_master=None,
+    trace_execution_id=None,
+):
+    """Core orchestration logic for the /distributed/queue endpoint.
+
+    Returns:
+        tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
+    """
+    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=prompt_server)
+    (
+        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)
+    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}"
+        ),
+    )
+
+    # 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_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:
+            # Include master in load balancing only when master participation is enabled.
+            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":
+            # Master selected as least busy; run master workload only.
+            active_workers = []
+            delegate_master = False
+            trace_debug(
+                execution_trace_id,
+                "Load-balance selected master for execution (workers skipped).",
+            )
+        elif selected_worker is not None:
+            active_workers = [selected_worker]
+            # Worker selected as least busy; keep master orchestrator-only for this run.
+            delegate_master = True
+            trace_debug(
+                execution_trace_id,
+                f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).",
+            )
+        else:
+            trace_debug(
+                execution_trace_id,
+                "Load-balance requested but no execution candidates were available.",
+            )
+            active_workers = []
+            delegate_master = False
+
+    enabled_ids = [worker["id"] for worker in active_workers]
+
+    discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
+    job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
+
+    if not job_id_map:
+        trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
+        queue_result = await queue_prompt_payload(
+            prompt_obj,
+            workflow_meta,
+            client_id,
+            include_queue_metadata=True,
+        )
+        return (
+            queue_result["prompt_id"],
+            queue_result["number"],
+            0,
+            queue_result.get("node_errors", {}),
+        )
+
+    for job_id in job_id_map.values():
+        await _ensure_distributed_queue(job_id)
+
+    master_prompt = prompt_index.copy_prompt()
+    master_prompt = apply_participant_overrides(
+        master_prompt,
+        "master",
+        enabled_ids,
+        job_id_map,
+        master_url,
+        delegate_master,
+        prompt_index,
+    )
+
+    if delegate_master:
+        collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
+        upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
+        if upscale_nodes:
+            debug_log(
+                "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
+            )
+        elif not collector_ids:
+            debug_log(
+                "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
+            )
+        else:
+            master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
+
+    if active_workers:
+        trace_debug(
+            execution_trace_id,
+            "Active distributed workers: "
+            + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
+        )
+    worker_payloads = []
+    if active_workers:
+        worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
+        media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
+        worker_payloads = await asyncio.gather(
+            *[
+                _prepare_worker_payload(
+                    worker,
+                    prompt_index,
+                    enabled_ids,
+                    job_id_map,
+                    master_url,
+                    config,
+                    delegate_master,
+                    execution_trace_id,
+                    worker_prep_semaphore,
+                    media_sync_semaphore,
+                    media_sync_timeout_seconds,
+                )
+                for worker in active_workers
+            ]
+        )
+
+    if worker_payloads:
+        await asyncio.gather(
+            *[
+                dispatch_worker_prompt(
+                    worker,
+                    wprompt,
+                    workflow_meta,
+                    client_id,
+                    use_websocket=use_websocket,
+                    trace_execution_id=execution_trace_id,
+                )
+                for worker, wprompt in worker_payloads
+            ]
+        )
+
+    queue_result = await queue_prompt_payload(
+        master_prompt,
+        workflow_meta,
+        client_id,
+        include_queue_metadata=True,
+    )
+    prompt_id = queue_result["prompt_id"]
+    prompt_number = queue_result["number"]
+    node_errors = queue_result.get("node_errors", {})
+    trace_debug(
+        execution_trace_id,
+        f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
+    )
+    return prompt_id, prompt_number, len(worker_payloads), node_errors

api/queue_request.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+
+@dataclass(frozen=True)
+class QueueRequestPayload:
+    prompt: Dict[str, Any]
+    workflow_meta: Any
+    client_id: str
+    delegate_master: Optional[bool]
+    enabled_worker_ids: List[str]
+    auto_prepare: bool
+    trace_execution_id: Optional[str]
+
+
+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")
+
+    auto_prepare_raw = data.get("auto_prepare", True)
+    if not isinstance(auto_prepare_raw, bool):
+        raise ValueError("auto_prepare must be a boolean when provided")
+    auto_prepare = auto_prepare_raw
+
+    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")
+
+    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")
+        enabled_ids = [str(worker_id).strip() for worker_id in enabled_ids_raw if str(worker_id).strip()]
+
+    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=data.get("workflow"),
+        client_id=client_id,
+        delegate_master=delegate_master,
+        enabled_worker_ids=enabled_ids,
+        auto_prepare=auto_prepare,
+        trace_execution_id=trace_execution_id,
+    )

api/schemas.py

@@ -0,0 +1,54 @@
+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) -> bool:
+    """Return True when worker_id exists in config['workers']."""
+    worker_id_str = str(worker_id)
+    workers = (config or {}).get("workers", [])
+    return any(str(worker.get("id")) == worker_id_str for worker in workers)
+
+
+def validate_positive_int(value, field_name: str) -> str | None:
+    """Validate positive integers and return an error string when invalid."""
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError):
+        return f"Field '{field_name}' must be a positive integer."
+    if parsed <= 0:
+        return f"Field '{field_name}' must be a positive integer."
+    return None
+
+
+def parse_positive_int(value, default: int) -> int:
+    """Parse value as positive int, returning default on failure."""
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError):
+        return max(1, int(default))
+    return max(1, parsed)
+
+
+def parse_positive_float(value, default: float) -> float:
+    """Parse value as positive float, returning default on failure."""
+    try:
+        parsed = float(value)
+    except (TypeError, ValueError):
+        return max(0.0, float(default))
+    return max(0.0, parsed)

api/tunnel_routes.py

@@ -0,0 +1,51 @@
+from aiohttp import web
+import server
+
+from ..utils.cloudflare import cloudflare_tunnel_manager
+from ..utils.config import load_config
+from ..utils.logging import debug_log, log
+from ..utils.network import handle_api_error
+
+
+@server.PromptServer.instance.routes.get("/distributed/tunnel/status")
+async def tunnel_status_endpoint(request):
+    """Return Cloudflare tunnel status and last known details."""
+    try:
+        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
+        })
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+@server.PromptServer.instance.routes.post("/distributed/tunnel/start")
+async def tunnel_start_endpoint(request):
+    """Start a Cloudflare tunnel pointing at the current ComfyUI server."""
+    try:
+        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")
+        })
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+@server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
+async def tunnel_stop_endpoint(request):
+    """Stop the managed Cloudflare tunnel if running."""
+    try:
+        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")
+        })
+    except Exception as e:
+        return await handle_api_error(request, e, 500)

api/usdu_routes.py

@@ -0,0 +1,228 @@
+import asyncio
+import io
+import time
+
+from aiohttp import web
+from PIL import Image
+import server
+
+from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
+from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized
+from ..upscale.payload_parsers import _parse_tiles_from_form
+from ..utils.logging import debug_log
+from ..utils.network import handle_api_error
+
+
+@server.PromptServer.instance.routes.post("/distributed/heartbeat")
+async def heartbeat_endpoint(request):
+    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):
+                    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):
+    """Endpoint for workers to submit processed tiles in static mode."""
+    try:
+        content_length = request.headers.get('content-length')
+        if content_length and int(content_length) > 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')
+        is_last = data.get('is_last', 'False').lower() == 'true'
+
+        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()
+
+        batch_size = int(data.get('batch_size', 0))
+
+        # 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)
+                    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)
+
+                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):
+    """Endpoint for workers to submit processed images in dynamic mode."""
+    try:
+        content_length = request.headers.get('content-length')
+        if content_length and int(content_length) > 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')
+        is_last = data.get('is_last', 'False').lower() == 'true'
+
+        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:
+            image_idx = int(data.get('image_idx'))
+            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)
+                    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)
+                    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):
+    """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
+    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)
+
+                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({"image_idx": task_idx, "estimated_remaining": remaining})
+                    debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
+                    return web.json_response({
+                        "tile_idx": task_idx,
+                        "estimated_remaining": remaining,
+                        "batched_static": job_data.batched_static,
+                    })
+                except asyncio.TimeoutError:
+                    if mode == 'dynamic':
+                        return web.json_response({"image_idx": None})
+                    return web.json_response({"tile_idx": None})
+            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):
+    """Endpoint to check if a job is ready."""
+    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,708 @@
+import json
+import asyncio
+import os
+import time
+import platform
+import subprocess
+import socket
+
+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,
+    get_server_port,
+    handle_api_error,
+    normalize_host,
+    probe_worker,
+)
+from ..utils.constants import CHUNK_SIZE
+from ..workers import get_worker_manager
+from ..workers.ports import allocate_worker_ports
+from .schemas import require_fields, validate_worker_id
+from ..workers.detection import (
+    get_machine_id,
+    is_docker_environment,
+    is_runpod_environment,
+)
+try:
+    from ..utils.async_helpers import PromptValidationError, queue_prompt_payload
+except ImportError:
+    from ..utils.async_helpers import queue_prompt_payload
+
+    class PromptValidationError(RuntimeError):
+        def __init__(self, message, validation_error=None, node_errors=None):
+            super().__init__(str(message))
+            self.validation_error = validation_error if isinstance(validation_error, dict) else {}
+            self.node_errors = node_errors if isinstance(node_errors, dict) else {}
+
+
+@server.PromptServer.instance.routes.get("/distributed/worker_ws")
+async def worker_ws_endpoint(request):
+    """WebSocket endpoint for worker prompt dispatch."""
+    ws = web.WebSocketResponse(heartbeat=30)
+    await ws.prepare(request)
+
+    async for msg in ws:
+        if msg.type == aiohttp.WSMsgType.TEXT:
+            try:
+                data = json.loads(msg.data or "{}")
+            except json.JSONDecodeError:
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": None,
+                    "ok": False,
+                    "error": "Invalid JSON payload.",
+                })
+                continue
+
+            if data.get("type") != "dispatch_prompt":
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": False,
+                    "error": "Unsupported websocket message type.",
+                })
+                continue
+
+            prompt = data.get("prompt")
+            if not isinstance(prompt, dict):
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": False,
+                    "error": "Field 'prompt' must be an object.",
+                })
+                continue
+
+            try:
+                prompt_id = await queue_prompt_payload(
+                    prompt,
+                    workflow_meta=data.get("workflow"),
+                    client_id=data.get("client_id"),
+                )
+                await ws.send_json({
+                    "type": "dispatch_ack",
+                    "request_id": data.get("request_id"),
+                    "ok": True,
+                    "prompt_id": prompt_id,
+                })
+            except 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):
+    """Clear the launching flag when worker is confirmed running."""
+    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()
+        if not validate_worker_id(worker_id, config):
+            return await handle_api_error(request, f"Worker {worker_id} not found", 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():
+    """Get all network IPs, trying multiple methods."""
+    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):
+        pass
+
+    # 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):
+        pass
+
+    # Method 3: Platform-specific commands
+    try:
+        if platform.system() == "Windows":
+            # Windows ipconfig
+            result = subprocess.run(["ipconfig"], capture_output=True, text=True)
+            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(["ip", "addr"], capture_output=True, text=True)
+            except (FileNotFoundError, OSError):
+                try:
+                    result = subprocess.run(["ifconfig"], capture_output=True, text=True)
+                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):
+        pass
+
+    return ips
+
+
+def get_recommended_ip(ips):
+    """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():
+    """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(
+                        ['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():
+    """Collect network/cuda info in a worker thread to avoid blocking route handlers."""
+    cuda_device, cuda_device_count, physical_device_count = _get_cuda_info()
+    device_count = physical_device_count if physical_device_count > 0 else cuda_device_count
+    master_port = get_server_port()
+    config = load_config()
+    worker_ports = []
+    if not config.get("settings", {}).get("has_auto_populated_workers") and not config.get("workers"):
+        worker_count = device_count - (1 if cuda_device is not None and 0 <= cuda_device < device_count else 0)
+        worker_ports = allocate_worker_ports(master_port, config.get("workers", []), worker_count)
+    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": device_count,
+        "master_port": master_port,
+        "local_worker_ports": worker_ports,
+    }
+
+
+def _read_worker_log_sync(log_file, lines_to_read):
+    """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, default, minimum=1, maximum=10000):
+    """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, worker_id):
+    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):
+    """Return this instance's in-memory ComfyUI log buffer."""
+    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):
+    """Get network interfaces and recommend best IP for master."""
+    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):
+    """Get system information including machine ID for local worker detection."""
+    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):
+    """Launch a worker process from the UI."""
+    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()
+        if not validate_worker_id(worker_id, config):
+            return await handle_api_error(request, f"Worker {worker_id} not found", 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 ValueError as e:
+            return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 400)
+        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):
+    """Stop a worker process that was launched from the UI."""
+    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()
+        if not validate_worker_id(worker_id, config):
+            return await handle_api_error(request, f"Worker {worker_id} not found", 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):
+    """Get list of workers managed by this UI instance."""
+    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):
+    """Check status of all local workers (localhost/no host specified)."""
+    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):
+    """Get log content for a specific worker."""
+    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):
+    """Proxy a remote worker log request to the worker's local in-memory log endpoint."""
+    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)},
+            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)

conftest.py

@@ -0,0 +1,31 @@
+# conftest.py — project-level pytest configuration.
+#
+# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
+# (from .nodes.v3 import ...) that fail when pytest tries to import it as a
+# standalone module during Package.setup() for the root package node.
+#
+# Fix: patch Package.setup() to skip the root-package's __init__.py import.
+# All actual package context is provided by each test module via
+# importlib.util.spec_from_file_location with synthetic stub packages.
+
+from _pytest.python import Package
+
+_orig_pkg_setup = Package.setup
+
+
+def _patched_pkg_setup(self) -> None:
+    # Skip the root package setup — its __init__.py uses relative imports
+    # that require a parent package (ComfyUI's plugin loader) which is not
+    # available in the test environment.
+    if self.path == self.config.rootpath:
+        return
+    _orig_pkg_setup(self)
+
+
+Package.setup = _patched_pkg_setup
+
+collect_ignore = [
+    "__init__.py",
+    "distributed.py",
+    "distributed_upscale.py",
+]

distributed_queue_api.py

@@ -1,605 +0,0 @@
-import asyncio
-import json
-import time
-import uuid
-from collections import deque
-
-import aiohttp
-
-import execution
-import server
-
-from .utils.config import load_config
-from .utils.logging import debug_log, log
-
-
-prompt_server = server.PromptServer.instance
-
-
-def ensure_distributed_state():
-    """Ensure prompt_server has the state used by distributed queue orchestration."""
-    if not hasattr(prompt_server, "distributed_pending_jobs"):
-        prompt_server.distributed_pending_jobs = {}
-    if not hasattr(prompt_server, "distributed_jobs_lock"):
-        prompt_server.distributed_jobs_lock = asyncio.Lock()
-    if not hasattr(prompt_server, "distributed_worker_ws"):
-        prompt_server.distributed_worker_ws = {}
-
-
-async def _get_client_session():
-    """Get or create aiohttp client session (shared with distributed.py)."""
-    if not hasattr(prompt_server, "_distributed_session"):
-        prompt_server._distributed_session = aiohttp.ClientSession()
-    return prompt_server._distributed_session
-
-
-class PromptIndex:
-    """Cache prompt metadata for faster worker/master prompt preparation."""
-
-    def __init__(self, prompt_obj):
-        self._prompt_json = json.dumps(prompt_obj)
-        self.nodes_by_class = {}
-        self.class_by_node = {}
-        self.inputs_by_node = {}
-        for node_id, node in _iter_prompt_nodes(prompt_obj):
-            class_type = node.get("class_type")
-            node_id_str = str(node_id)
-            if class_type:
-                self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
-            self.class_by_node[node_id_str] = class_type
-            self.inputs_by_node[node_id_str] = node.get("inputs", {})
-        self._upstream_cache = {}
-
-    def copy_prompt(self):
-        return json.loads(self._prompt_json)
-
-    def nodes_for_class(self, class_name):
-        return self.nodes_by_class.get(class_name, [])
-
-    def has_upstream(self, start_node_id, target_class):
-        cache_key = (str(start_node_id), target_class)
-        if cache_key in self._upstream_cache:
-            return self._upstream_cache[cache_key]
-
-        visited = set()
-        stack = [str(start_node_id)]
-        while stack:
-            node_id = stack.pop()
-            if node_id in visited:
-                continue
-            visited.add(node_id)
-            inputs = self.inputs_by_node.get(node_id, {})
-            for value in inputs.values():
-                if isinstance(value, list) and len(value) == 2:
-                    upstream_id = str(value[0])
-                    if self.class_by_node.get(upstream_id) == target_class:
-                        self._upstream_cache[cache_key] = True
-                        return True
-                    if upstream_id in self.inputs_by_node:
-                        stack.append(upstream_id)
-
-        self._upstream_cache[cache_key] = False
-        return False
-
-
-def _iter_prompt_nodes(prompt_obj):
-    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 _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):
-            continue
-        max_id = max(max_id, numeric)
-
-    counter = max_id
-
-    def _next_id():
-        nonlocal counter
-        counter += 1
-        return str(counter)
-
-    return _next_id
-
-
-def _prepare_delegate_master_prompt(prompt_obj, collector_ids):
-    """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."
-                    )
-
-    next_id = _create_numeric_id_generator(pruned_prompt)
-    for collector_id in collector_ids:
-        collector_entry = pruned_prompt.get(collector_id)
-        if not collector_entry:
-            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
-
-
-async def _ensure_distributed_queue(job_id):
-    """Ensure a queue exists for the given distributed job ID."""
-    ensure_distributed_state()
-    async with prompt_server.distributed_jobs_lock:
-        if job_id not in prompt_server.distributed_pending_jobs:
-            prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
-
-
-def _generate_job_id_map(prompt_index, prefix):
-    """Create stable per-node job IDs for distributed nodes."""
-    job_map = {}
-    distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
-        "UltimateSDUpscaleDistributed"
-    )
-    for node_id in distributed_nodes:
-        job_map[node_id] = f"{prefix}_{node_id}"
-    return job_map
-
-
-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_master_url():
-    """Best-effort reconstruction of the master's public URL."""
-    cfg = load_config()
-    master_cfg = cfg.get("master", {}) or {}
-    configured_host = (master_cfg.get("host") or "").strip()
-    configured_port = master_cfg.get("port")
-    default_port = getattr(prompt_server, "port", 8188) or 8188
-    port = int(configured_port or default_port)
-
-    def _needs_https(hostname):
-        hostname = hostname.lower()
-        https_domains = (
-            ".proxy.runpod.net",
-            ".ngrok-free.app",
-            ".ngrok-free.dev",
-            ".ngrok.io",
-            ".trycloudflare.com",
-            ".cloudflare.dev",
-        )
-        return any(hostname.endswith(suffix) for suffix in https_domains)
-
-    if configured_host:
-        if configured_host.startswith(("http://", "https://")):
-            return configured_host.rstrip("/")
-
-        host = configured_host
-        scheme = "https" if _needs_https(host) or port == 443 else "http"
-        default_port_for_scheme = 443 if scheme == "https" else 80
-        # For ngrok/cloud domains without explicit port, default to the scheme's default
-        if configured_port is None and scheme == "https" and _needs_https(host):
-            port = default_port_for_scheme
-        port_part = "" if port == default_port_for_scheme else f":{port}"
-        return f"{scheme}://{host}{port_part}"
-
-    address = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
-    if address in ("0.0.0.0", "::"):
-        address = "127.0.0.1"
-    scheme = "https" if port == 443 else "http"
-    default_port_for_scheme = 443 if scheme == "https" else 80
-    port_part = "" if port == default_port_for_scheme else f":{port}"
-    return f"{scheme}://{address}{port_part}"
-
-
-def _build_worker_url(worker, endpoint=""):
-    """Construct the worker base URL with optional endpoint."""
-    host = (worker.get("host") or "").strip()
-    port = int(worker.get("port", 8188) or 8188)
-
-    if not host:
-        host = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
-
-    if host.startswith(("http://", "https://")):
-        base = host.rstrip("/")
-    else:
-        is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
-        scheme = "https" if is_cloud else "http"
-        default_port = 443 if scheme == "https" else 80
-        port_part = "" if port == default_port else f":{port}"
-        base = f"{scheme}://{host}{port_part}"
-
-    return f"{base}{endpoint}"
-
-
-async def _worker_is_active(worker):
-    """Ping worker's /prompt endpoint to confirm it's reachable."""
-    url = _build_worker_url(worker, "/prompt")
-    session = await _get_client_session()
-    try:
-        async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp:
-            return resp.status == 200
-    except Exception:
-        return False
-
-
-async def _worker_ws_is_active(worker):
-    """Ping worker's websocket endpoint to confirm it's reachable."""
-    session = await _get_client_session()
-    url = _build_worker_url(worker, "/distributed/worker_ws")
-    try:
-        ws = await session.ws_connect(url, heartbeat=20, timeout=3)
-        await ws.close()
-        return True
-    except Exception:
-        return False
-
-
-async def _get_worker_ws_connection(worker):
-    """Return a cached websocket connection to a worker (creating if needed)."""
-    ensure_distributed_state()
-    worker_id = worker.get("id")
-    existing = prompt_server.distributed_worker_ws.get(worker_id)
-    if existing and not existing.closed:
-        return existing
-
-    session = await _get_client_session()
-    url = _build_worker_url(worker, "/distributed/worker_ws")
-    ws = await session.ws_connect(url, heartbeat=20, timeout=10)
-    prompt_server.distributed_worker_ws[worker_id] = ws
-    return ws
-
-
-async def _close_worker_ws(worker_id):
-    ws = prompt_server.distributed_worker_ws.pop(worker_id, None)
-    if ws and not ws.closed:
-        await ws.close()
-
-
-async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta, client_id=None, use_websocket=False):
-    """Send the prepared prompt to a worker ComfyUI instance."""
-    url = _build_worker_url(worker, "/prompt")
-    payload = {"prompt": prompt_obj}
-    extra_data = {}
-    if workflow_meta:
-        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
-    if extra_data:
-        payload["extra_data"] = extra_data
-
-    if use_websocket:
-        request_id = uuid.uuid4().hex
-        ws_payload = {
-            "type": "dispatch_prompt",
-            "request_id": request_id,
-            "prompt": prompt_obj,
-            "workflow": workflow_meta,
-            "client_id": client_id,
-        }
-        try:
-            ws = await _get_worker_ws_connection(worker)
-            await ws.send_json(ws_payload)
-            response = await asyncio.wait_for(ws.receive(), timeout=60)
-            if response.type == aiohttp.WSMsgType.TEXT:
-                data = json.loads(response.data or "{}")
-                if (
-                    data.get("type") == "dispatch_ack"
-                    and data.get("request_id") == request_id
-                    and data.get("ok")
-                ):
-                    return
-                error_msg = data.get("error") or "Worker rejected websocket dispatch."
-                raise RuntimeError(error_msg)
-            raise RuntimeError("Unexpected websocket response from worker.")
-        except Exception as exc:
-            worker_id = worker.get("id")
-            await _close_worker_ws(worker_id)
-            log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
-
-    session = await _get_client_session()
-    async with session.post(
-        url,
-        json=payload,
-        timeout=aiohttp.ClientTimeout(total=60),
-    ) as resp:
-        resp.raise_for_status()
-
-
-async def _queue_master_prompt(prompt_obj, workflow_meta, client_id):
-    """Queue the master prompt through ComfyUI's prompt queue."""
-    payload = {"prompt": prompt_obj}
-    payload = prompt_server.trigger_on_prompt(payload)
-    prompt = payload["prompt"]
-
-    prompt_id = str(uuid.uuid4())
-    valid = await execution.validate_prompt(prompt_id, prompt, None)
-    if not valid[0]:
-        raise RuntimeError(f"Invalid prompt: {valid[1]}")
-
-    extra_data = {}
-    if workflow_meta:
-        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
-    if client_id:
-        extra_data["client_id"] = client_id
-
-    sensitive = {}
-    for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
-        if key in extra_data:
-            sensitive[key] = extra_data.pop(key)
-
-    number = getattr(prompt_server, "number", 0)
-    prompt_server.number = number + 1
-    prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
-    prompt_server.prompt_queue.put(prompt_queue_item)
-    return prompt_id
-
-
-def _apply_participant_overrides(
-    prompt_copy,
-    participant_id,
-    enabled_worker_ids,
-    job_id_map,
-    master_url,
-    delegate_master,
-    prompt_index,
-):
-    """Return a prompt copy with hidden inputs configured for master/worker."""
-    is_master = participant_id == "master"
-    worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
-    enabled_json = json.dumps(enabled_worker_ids)
-
-    # Distributed seed nodes
-    for node_id in prompt_index.nodes_for_class("DistributedSeed"):
-        node = prompt_copy.get(node_id, {})
-        inputs = node.setdefault("inputs", {})
-        inputs["is_worker"] = not is_master
-        if not is_master:
-            idx = worker_index_map.get(participant_id, 0)
-            inputs["worker_id"] = f"worker_{idx}"
-        else:
-            inputs["worker_id"] = ""
-
-    # Distributed collectors
-    for node_id in prompt_index.nodes_for_class("DistributedCollector"):
-        node = prompt_copy.get(node_id, {})
-        inputs = node.setdefault("inputs", {})
-
-        if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
-            inputs["pass_through"] = True
-            continue
-
-        unique_id = job_id_map.get(node_id, node_id)
-        inputs["multi_job_id"] = unique_id
-        inputs["is_worker"] = not is_master
-
-        if is_master:
-            inputs["enabled_worker_ids"] = enabled_json
-            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["enabled_worker_ids"] = enabled_json
-            inputs["delegate_only"] = False
-
-    # Distributed upscaler nodes
-    for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
-        node = prompt_copy.get(node_id, {})
-        inputs = node.setdefault("inputs", {})
-        unique_id = job_id_map.get(node_id, node_id)
-        inputs["multi_job_id"] = unique_id
-        inputs["is_worker"] = not is_master
-        if is_master:
-            inputs["enabled_worker_ids"] = enabled_json
-            inputs.pop("master_url", None)
-            inputs.pop("worker_id", None)
-        else:
-            inputs["master_url"] = master_url
-            inputs["worker_id"] = participant_id
-            inputs["enabled_worker_ids"] = enabled_json
-
-    return prompt_copy
-
-
-async def orchestrate_distributed_execution(
-    prompt_obj,
-    workflow_meta,
-    client_id,
-    enabled_worker_ids=None,
-    delegate_master=None,
-):
-    """Core orchestration logic for the /distributed/queue endpoint.
-
-    Returns:
-        tuple[str, int]: (prompt_id, worker_count)
-    """
-    ensure_distributed_state()
-
-    config = load_config()
-    use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
-    requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
-    workers = _resolve_enabled_workers(config, requested_ids)
-    prompt_index = PromptIndex(prompt_obj)
-
-    # Respect master delegate-only configuration
-    if delegate_master is None:
-        delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
-
-    if not workers and delegate_master:
-        debug_log("Delegate-only requested but no workers are enabled. Falling back to master execution.")
-        delegate_master = False
-
-    # Filter to active workers
-    active_workers = []
-    for worker in workers:
-        if use_websocket:
-            is_active = await _worker_ws_is_active(worker)
-        else:
-            is_active = await _worker_is_active(worker)
-        if is_active:
-            active_workers.append(worker)
-        else:
-            log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
-
-    if not active_workers and delegate_master:
-        debug_log("All workers offline while delegate-only requested; enabling master participation.")
-        delegate_master = False
-
-    enabled_ids = [worker["id"] for worker in active_workers]
-
-    discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
-    job_id_map = _generate_job_id_map(prompt_index, discovery_prefix)
-
-    if not job_id_map:
-        prompt_id = await _queue_master_prompt(prompt_obj, workflow_meta, client_id)
-        return prompt_id, 0
-
-    for job_id in job_id_map.values():
-        await _ensure_distributed_queue(job_id)
-
-    master_url = _resolve_master_url()
-    master_prompt = prompt_index.copy_prompt()
-    master_prompt = _apply_participant_overrides(
-        master_prompt,
-        "master",
-        enabled_ids,
-        job_id_map,
-        master_url,
-        delegate_master,
-        prompt_index,
-    )
-
-    if delegate_master:
-        collector_ids = _find_nodes_by_class(master_prompt, "DistributedCollector")
-        upscale_nodes = _find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
-        if upscale_nodes:
-            debug_log(
-                "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
-            )
-        elif not collector_ids:
-            debug_log(
-                "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
-            )
-        else:
-            master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids)
-
-    if active_workers:
-        debug_log(
-            "Active distributed workers: "
-            + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers)
-        )
-    worker_payloads = []
-    for worker in active_workers:
-        worker_prompt = prompt_index.copy_prompt()
-        worker_prompt = _apply_participant_overrides(
-            worker_prompt,
-            worker["id"],
-            enabled_ids,
-            job_id_map,
-            master_url,
-            delegate_master,
-            prompt_index,
-        )
-        worker_payloads.append((worker, worker_prompt))
-
-    if worker_payloads:
-        await asyncio.gather(
-            *[
-                _dispatch_worker_prompt(worker, wprompt, workflow_meta, client_id, use_websocket=use_websocket)
-                for worker, wprompt in worker_payloads
-            ]
-        )
-
-    prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id)
-    return prompt_id, len(worker_payloads)

docs/comfyui-distributed-api.md

@@ -19,9 +19,10 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
 - `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
   - Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
   - Resolves enabled/selected workers.
-  - Pings workers (`GET /prompt`) to include only reachable ones.
-  - Dispatches the workflow to workers (`POST /prompt`).
+  - Probes and dispatches workers through `/distributed/worker_ws` by default.
+  - If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
   - 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
 
@@ -48,27 +49,39 @@ Queue a workflow for distributed execution.
 {
   "prompt": { "<node_id>": { "class_type": "...", "inputs": { } } },
   "workflow": { },
-  "client_id": "optional",
+  "client_id": "external-client",
   "delegate_master": false,
-  "enabled_worker_ids": ["1", "2"]
+  "enabled_worker_ids": ["1", "2"],
+  "workers": ["1", "2"],
+  "auto_prepare": true,
+  "trace_execution_id": "exec_1700000000_ab12cd"
 }
 ```
 
 #### Fields
 
-- `prompt` (required, object)
-  - The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
+- `prompt` (required unless `workflow.prompt` is present, object)
+  - A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
+  - This is not the normal visual workflow export from the ComfyUI editor.
 - `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` (optional, string)
+- `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` (optional, array of strings)
-  - If provided, only these worker IDs will be considered.
-  - If omitted, the plugin uses workers marked as enabled in the UI config.
+- `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.
+- `auto_prepare` (optional, boolean)
+  - Kept for wire compatibility.
+  - Backend orchestration always runs with auto-prepare semantics.
+  - If top-level `prompt` is omitted, backend will attempt `workflow.prompt`.
+- `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`
 
@@ -93,7 +106,8 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
 ```json
 {
   "prompt_id": "<uuid>",
-  "worker_count": 2
+  "worker_count": 2,
+  "auto_prepare_supported": true
 }
 ```
 
@@ -110,24 +124,89 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
 
 ## Worker requirements (important)
 
-For a worker to participate, it must be reachable from the master:
+For a worker to participate, it must be reachable from the master. By default:
+
+- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
+
+If `settings.websocket_orchestration=false`:
 
 - Health check: `GET <worker-base>/prompt` must return HTTP 200.
-- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
+- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
 
 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`
+### 1) Minimal `curl` request envelope
 
 ```bash
 curl -X POST "http://127.0.0.1:8188/distributed/queue" \
@@ -135,16 +214,23 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
   -d @payload.json
 ```
 
-Where `payload.json` contains at least:
+`payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
 
 ```json
 {
   "prompt": {
-    "1": {"class_type": "KSampler", "inputs": {} }
-  }
+    "<node_id>": {
+      "class_type": "<node class>",
+      "inputs": {"<required_input>": "<value or connection>"}
+    }
+  },
+  "enabled_worker_ids": [],
+  "client_id": "external-client"
 }
 ```
 
+Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
+
 ### 2) Python (`requests`)
 
 ```python
@@ -154,6 +240,7 @@ url = "http://127.0.0.1:8188/distributed/queue"
 payload = {
     "prompt": {...},
     "workflow": {...},
+    "client_id": "external-client",
     "delegate_master": False,
     "enabled_worker_ids": ["1", "2"],
 }
@@ -171,6 +258,7 @@ 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"],
 };

docs/v3-migration.md

@@ -0,0 +1,53 @@
+# Native V3 node API
+
+The package registers one `ComfyExtension` through `comfy_entrypoint()` and
+imports the versioned `comfy_api.v0_0_2` API. Use a ComfyUI version that provides
+that API; there is no V1 registration fallback.
+
+## Compatibility
+
+- All eight node IDs, display names, categories, visible input order, defaults
+  and output order are retained. Existing execution algorithms remain in the
+  private collector, utilities and upscale modules.
+- The image/audio dividers explicitly declare ten typed outputs. Their existing
+  frontend extensions still show only the selected number of outputs. This
+  replaces the V1 `ByPassTypeTuple` indexing workaround without changing saved
+  IMAGE/AUDIO socket indices or the ten returned values.
+- Standard hidden context uses `cls.hidden`. Worker/orchestration metadata keeps
+  its existing prompt-input names and Python defaults via `accept_all_inputs`;
+  it does not become visible widgets.
+- The collector retains list-input handling. Upscale retains its always-changing
+  fingerprint and creates a private runtime helper for each V3 execution, so
+  mutable helper state is not attached to sanitized V3 class clones.
+- Routes, distributed state and the existing worker startup/shutdown hooks are
+  initialized by `runtime/bootstrap.py` during `ComfyExtension.on_load()`.
+  Worker mode still suppresses automatic worker launch. `distributed.py` is no
+  longer an entrypoint.
+
+## Verification
+
+Run ordinary unit tests from this repository:
+
+```bash
+python -m pytest tests -q -o addopts=
+```
+
+Opt into real-framework acceptance using a ComfyUI checkout and its interpreter:
+
+```bash
+COMFYUI_SOURCE_ROOT=/path/to/ComfyUI \
+  /path/to/ComfyUI/.venv/bin/python -m pytest tests -q -o addopts=
+```
+
+The acceptance subprocess uses CPU mode and the real ComfyUI loader, input
+parser, V3 class preparation, prompt validation and `PromptExecutor`. It compares
+all node schemas with the V1 fixture, checks all five bundled workflows' node
+IDs and socket/link contracts, exercises injected worker metadata, image/audio
+lists, collector aggregation and divider outputs, rejects invalid upscale enums,
+and decodes an actual preview PNG referenced by executor history.
+
+The upscale GPU/model boundary is mocked to check argument forwarding and
+per-execution helper isolation. Full checkpoint inference, browser canvas
+acceptance and multi-host HTTP transport are not exercised. No HTTP listener,
+workers or model downloads are started; preview files use a temporary scratch
+directory. The test does not install this branch into a live custom-node folder.

docs/worker-setup-guides.md

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

nodes/__init__.py

@@ -0,0 +1 @@
+"""Private execution helpers; public registration lives in nodes.v3."""

nodes/collector.py

@@ -0,0 +1,556 @@
+import torch
+import io
+import json
+import asyncio
+import time
+import base64
+
+import aiohttp
+import server as _server
+import comfy.model_management
+from comfy.utils import ProgressBar
+
+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.audio_payload import encode_audio_payload
+from ..utils.async_helpers import run_async_in_server_loop
+
+prompt_server = _server.PromptServer.instance
+
+
+class DistributedCollectorNode:
+    INPUT_IS_LIST = True
+    EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "load_balance": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Run this workflow on one least-busy participant (master included when participating).",
+                    },
+                ),
+            },
+            "optional": {
+                "images": ("IMAGE",),
+                "audio": ("AUDIO",),
+            },
+            "hidden": {
+                "multi_job_id": ("STRING", {"default": ""}),
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "master_url": ("STRING", {"default": ""}),
+                "enabled_worker_ids": ("STRING", {"default": "[]"}),
+                "worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}),
+                "worker_id": ("STRING", {"default": ""}),
+                "pass_through": ("BOOLEAN", {"default": False}),
+                "delegate_only": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "AUDIO")
+    RETURN_NAMES = ("images", "audio")
+    FUNCTION = "run"
+    CATEGORY = "image"
+    
+    @staticmethod
+    def _unwrap_list_input(value):
+        """Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST."""
+        if isinstance(value, (list, tuple)) and len(value) == 1:
+            return value[0]
+        return value
+
+    def _normalize_images_input(self, images):
+        """Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
+        if isinstance(images, (list, tuple)):
+            if not images:
+                raise ValueError("Collector received an empty image list")
+            if not all(isinstance(image, torch.Tensor) for image in images):
+                raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
+            if len(images) == 1:
+                return ensure_contiguous(images[0])
+            return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
+        return ensure_contiguous(images)
+
+    def _normalize_audio_input(self, audio):
+        """Collapse ComfyUI list AUDIO inputs into a single AUDIO payload when present."""
+        if not isinstance(audio, (list, tuple)):
+            return audio
+
+        audio_items = [item for item in audio if item is not None]
+        if not audio_items:
+            return None
+        if len(audio_items) == 1:
+            return audio_items[0]
+
+        waveforms = []
+        sample_rate = 44100
+        for item in audio_items:
+            if not isinstance(item, dict):
+                raise TypeError("Collector expected AUDIO list items to be dictionaries")
+            waveform = item.get("waveform")
+            if waveform is None or waveform.numel() == 0:
+                continue
+            waveforms.append(waveform)
+            if sample_rate == 44100:
+                sample_rate = item.get("sample_rate", 44100)
+
+        if not waveforms:
+            return None
+        return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
+
+    def run(self, images=None, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
+        if images is not None:
+            images = self._normalize_images_input(images)
+        audio = self._normalize_audio_input(audio)
+        load_balance = self._unwrap_list_input(load_balance)
+        multi_job_id = self._unwrap_list_input(multi_job_id)
+        is_worker = self._unwrap_list_input(is_worker)
+        master_url = self._unwrap_list_input(master_url)
+        enabled_worker_ids = self._unwrap_list_input(enabled_worker_ids)
+        worker_batch_size = self._unwrap_list_input(worker_batch_size)
+        worker_id = self._unwrap_list_input(worker_id)
+        pass_through = self._unwrap_list_input(pass_through)
+        delegate_only = self._unwrap_list_input(delegate_only)
+
+        remote_only_master = (
+            bool(multi_job_id)
+            and not is_worker
+            and (delegate_only or is_master_delegate_only())
+        )
+        if images is None and audio is None and not remote_only_master:
+            raise ValueError("DistributedCollector requires at least one image or audio input")
+
+        # Create empty audio if not provided
+        empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
+
+        if not multi_job_id or pass_through:
+            if 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,
+                audio,
+                load_balance,
+                multi_job_id,
+                is_worker,
+                master_url,
+                enabled_worker_ids,
+                worker_batch_size,
+                worker_id,
+                delegate_only,
+            )
+        )
+        return result
+
+    async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
+        """Send an image batch, optionally with audio, or an audio-only completion."""
+        encoded_audio = encode_audio_payload(audio)
+        session = await get_client_session()
+        url = f"{master_url}/distributed/job_complete"
+
+        payloads = []
+        batch_size = 0 if image_batch is None else image_batch.shape[0]
+        if batch_size == 0:
+            if encoded_audio is None:
+                raise ValueError("Worker completion requires image or audio data")
+            payloads.append(
+                {
+                    "job_id": str(multi_job_id),
+                    "worker_id": str(worker_id),
+                    "batch_idx": 0,
+                    "audio": encoded_audio,
+                    "is_last": True,
+                }
+            )
+        else:
+            for batch_idx in range(batch_size):
+                img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
+                byte_io = io.BytesIO()
+                img.save(byte_io, format='PNG', compress_level=0)
+                encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
+                payload = {
+                    "job_id": str(multi_job_id),
+                    "worker_id": str(worker_id),
+                    "batch_idx": int(batch_idx),
+                    "image": f"data:image/png;base64,{encoded_image}",
+                    "is_last": bool(batch_idx == batch_size - 1),
+                }
+                if payload["is_last"] and encoded_audio is not None:
+                    payload["audio"] = encoded_audio
+                payloads.append(payload)
+
+        for payload in payloads:
+            timeout_seconds = 60 if "image" in payload else 600
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    timeout=aiohttp.ClientTimeout(total=timeout_seconds),
+                ) as response:
+                    response.raise_for_status()
+            except Exception as e:
+                media_type = "image/audio" if "image" in payload else "audio-only"
+                log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
+                debug_log(f"Worker - Full error details: URL={url}")
+                raise
+
+    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,
+    ):
+        """Assemble final tensor, or return None when the job contains only audio."""
+        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))
+        if fallback_images is not None:
+            return ensure_contiguous(fallback_images)
+        return None
+
+    async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
+        if is_worker:
+            # Worker mode: send images and audio to master in a single batch
+            image_count = 0 if images is None else images.shape[0]
+            debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
+            await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
+            return (images, audio if audio is not None else self.EMPTY_AUDIO)
+        else:
+            delegate_mode = delegate_only or is_master_delegate_only()
+            # Master mode: collect images and audio from workers
+            enabled_workers_raw = json.loads(enabled_worker_ids)
+            enabled_workers = []
+            seen_enabled = set()
+            for worker_id in enabled_workers_raw:
+                worker_id_str = str(worker_id)
+                if worker_id_str in seen_enabled:
+                    continue
+                seen_enabled.add(worker_id_str)
+                enabled_workers.append(worker_id_str)
+            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)
+
+            # Create the queue before any expensive local work to avoid job_complete race.
+            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})")
+
+            if delegate_mode:
+                master_batch_size = 0
+                images_on_cpu = None
+                master_audio = None
+                debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
+            else:
+                if images is None:
+                    images_on_cpu = None
+                    master_batch_size = 0
+                else:
+                    images_on_cpu = ensure_contiguous(images.cpu())
+                    master_batch_size = images.shape[0]
+                master_audio = audio  # Keep master's audio for later
+                debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
+
+
+            # Initialize storage for collected images and audio
+            worker_images = {}  # Dict to store images by worker_id and index
+            worker_audio = {}   # Dict to store audio by worker_id
+            
+            # Collect images until all workers report they're done
+            collected_count = 0
+            workers_done = set()
+            
+            # Use unified worker timeout from config/UI with simple sliced waits
+            base_timeout = float(get_worker_timeout_seconds())
+            slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
+            last_activity = time.time()
+            
+            
+            # Get queue size before starting
+            async with prompt_server.distributed_jobs_lock:
+                q = prompt_server.distributed_pending_jobs[multi_job_id]
+                initial_size = q.qsize()
+
+            # NEW: Initialize progress bar for workers (total = num_workers)
+            p = ProgressBar(num_workers)
+
+            def mark_worker_done(done_worker_id):
+                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)
+                p.update(1)  # +1 per completed expected worker
+
+            try:
+                while len(workers_done) < num_workers:
+                    # Check for user interruption to abort collection promptly
+                    comfy.model_management.throw_exception_if_processing_interrupted()
+                    try:
+                        # Get the queue again each time to ensure we have the right reference
+                        async with prompt_server.distributed_jobs_lock:
+                            q = prompt_server.distributed_pending_jobs[multi_job_id]
+                            current_size = q.qsize()
+                        
+                        result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
+                        worker_id = result['worker_id']
+                        is_last = result.get('is_last', False)
+                        count = self._store_worker_result(worker_images, result)
+                        collected_count += count
+                        debug_log(
+                            f"Master - Got canonical result from worker {worker_id}, "
+                            f"image {result.get('image_index', 0)}, is_last={is_last}"
+                        )
+
+                        # Collect audio data if present
+                        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}")
+
+                        # Record activity and refresh timeout baseline
+                        last_activity = time.time()
+                        base_timeout = float(get_worker_timeout_seconds())
+
+                        if is_last:
+                            mark_worker_done(worker_id)
+                        
+                    except asyncio.TimeoutError:
+                        # If we still have time, continue polling; otherwise handle timeout
+                        if (time.time() - last_activity) < base_timeout:
+                            comfy.model_management.throw_exception_if_processing_interrupted()
+                            continue
+                        # Re-check for user interruption after timeout expiry
+                        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)"
+                        )
+
+                        # Probe missing workers' /prompt endpoints to check if they are actively processing
+                        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}")
+
+                        if any_busy:
+                            # Refresh last_activity and continue waiting
+                            last_activity = time.time()
+                            # Refresh base timeout in case the user changed it in UI
+                            base_timeout = float(get_worker_timeout_seconds())
+                            continue
+                        
+                        # Check queue size again with lock
+                        async with prompt_server.distributed_jobs_lock:
+                            if multi_job_id in prompt_server.distributed_pending_jobs:
+                                final_q = prompt_server.distributed_pending_jobs[multi_job_id]
+                                final_size = final_q.qsize()
+                                
+                                # Try to drain any remaining items
+                                remaining_items = []
+                                while not final_q.empty():
+                                    try:
+                                        item = final_q.get_nowait()
+                                        remaining_items.append(item)
+                                    except asyncio.QueueEmpty:
+                                        break
+                                
+                                if remaining_items:
+                                    # Process them
+                                    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)
+                            else:
+                                log(f"Master - Queue {multi_job_id} no longer exists!")
+                        break
+            except comfy.model_management.InterruptProcessingException:
+                # Cleanup queue on interruption and re-raise to abort prompt cleanly
+                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]
+                raise
+            
+            total_collected = sum(len(imgs) for imgs in worker_images.values())
+            
+            # Clean up job queue
+            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]
+
+            combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
+            try:
+                combined = self._reorder_and_combine_tensors(
+                    worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
+                )
+                if combined is None:
+                    debug_log(f"Master - Job {multi_job_id} complete with audio only")
+                else:
+                    debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
+                              f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
+
+                return (combined, combined_audio)
+            except Exception as e:
+                log(f"Master - Error combining images: {e}")
+                # Preserve collected audio even when image assembly fails.
+                return (images, combined_audio)

nodes/distributed_upscale.py

@@ -0,0 +1,270 @@
+import json
+import math
+from functools import wraps
+
+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 ..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.modes.single_gpu import SingleGpuModeMixin
+from ..upscale.modes.static import StaticModeMixin
+from ..upscale.modes.dynamic import DynamicModeMixin
+
+def sync_wrapper(async_func):
+    """Decorator to wrap async methods for synchronous execution."""
+    @wraps(async_func)
+    def sync_func(self, *args, **kwargs):
+        # 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):
+    """Parse enabled worker IDs from either JSON or list input."""
+    if isinstance(enabled_worker_ids, list):
+        return [str(worker_id) for worker_id in enabled_worker_ids]
+    if not enabled_worker_ids:
+        return []
+    if isinstance(enabled_worker_ids, str):
+        try:
+            parsed = json.loads(enabled_worker_ids)
+        except json.JSONDecodeError:
+            log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
+            return []
+        if isinstance(parsed, list):
+            return [str(wid) for wid in parsed]
+    return []
+
+class UltimateSDUpscaleDistributed(
+    DynamicModeMixin,
+    StaticModeMixin,
+    SingleGpuModeMixin,
+    ResultCollectorMixin,
+    WorkerCommsMixin,
+    JobStateMixin,
+    TileOpsMixin,
+):
+
+    """
+    Distributed version of Ultimate SD Upscale (No Upscale).
+    
+    Supports two currently selected processing modes:
+    1. Single GPU: No workers available, process everything locally
+    2. Distributed tile queue: Workers pull tiles from a shared queue
+    
+    Features:
+    - Tile-based batch handling for video/image upscaling
+    - Tiled VAE support for memory efficiency
+    - Shared work queue so faster workers can process more tiles
+    - Backward compatible with single-image workflows
+    
+    Environment Variables:
+    - COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
+    - COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
+    
+    The hidden dynamic_threshold input is retained for workflow compatibility but
+    does not affect the current mode-selection policy.
+    """
+
+    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(s):
+        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": {
+                "multi_job_id": ("STRING", {"default": ""}),
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "master_url": ("STRING", {"default": ""}),
+                "enabled_worker_ids": ("STRING", {"default": "[]"}),
+                "worker_id": ("STRING", {"default": ""}),
+                "tile_indices": ("STRING", {"default": ""}),  # Unused - kept for compatibility
+                "dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
+            },
+        }
+
+    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, 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="", is_worker=False, master_url="", enabled_worker_ids="[]", 
+            worker_id="", tile_indices="", dynamic_threshold=8):
+        """Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler."""
+        # 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, model, positive, negative, vae,
+                                          seed, steps, cfg, sampler_name, scheduler, denoise,
+                                          tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
+        
+        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, 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, enabled_worker_ids, 
+                                     dynamic_threshold)
+
+    def process_worker(self, 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):
+        """Unified worker processing - handles both static and dynamic modes."""
+        # 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 = json.loads(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, even if batch == 1.
+        if num_workers > 0 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, 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)
+        
+        # Static mode - enhanced with health monitoring and retry logic
+        return self._process_worker_static_sync(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_workers)
+
+    def process_master(self, 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, enabled_worker_ids, 
+                      dynamic_threshold):
+        """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 = json.loads(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,
+        # even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue.
+        if num_workers > 0 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, model, positive, negative, vae,
+                                         seed, steps, cfg, sampler_name, scheduler, denoise,
+                                         tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
+        
+        elif mode == "dynamic":
+            # Dynamic mode for large batches
+            return self.process_master_dynamic(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, enabled_workers)
+        
+        # Static mode - enhanced with unified job management
+        return self._process_master_static_sync(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, enabled_workers,
+                                               all_tiles, num_tiles_per_image)
+
+    def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
+        """Determines processing mode per requested policy:
+        - any workers     => prefer static (tile-based) for USDU
+        - no workers      => single_gpu
+        """
+        if num_workers == 0:
+            return "single_gpu"
+        # Default to static when distributed; master/worker may still override if special cases arise
+        return "static"
+
+# Ensure initialization before registering routes
+ensure_tile_jobs_initialized()

nodes/utilities.py

@@ -0,0 +1,354 @@
+import torch
+import json
+
+from ..utils.logging import debug_log, log
+
+
+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):
+        return {
+            "required": {
+                "seed": ("INT", {
+                    "default": 1125899906842, 
+                    "min": 0,
+                    "max": 1125899906842624,
+                    "forceInput": False  # Widget by default, can be converted to input
+                }),
+            },
+            "hidden": {
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "worker_id": ("STRING", {"default": ""}),
+            },
+        }
+    
+    RETURN_TYPES = ("INT",)
+    RETURN_NAMES = ("seed",)
+    FUNCTION = "distribute"
+    CATEGORY = "utils"
+    
+    def distribute(self, seed, is_worker=False, worker_id=""):
+        if not is_worker:
+            # Master node: pass through original values
+            debug_log(f"Distributor - Master: seed={seed}")
+            return (seed,)
+        else:
+            # Worker node: apply offset based on worker index
+            # Find worker index from enabled_worker_ids
+            try:
+                # Worker IDs are passed as "worker_0", "worker_1", etc.
+                if worker_id.startswith("worker_"):
+                    worker_index = int(worker_id.split("_")[1])
+                else:
+                    # Fallback: try to parse as direct index
+                    worker_index = int(worker_id)
+                
+                offset = worker_index + 1
+                new_seed = seed + offset
+                debug_log(f"Distributor - Worker {worker_index}: seed={seed} → {new_seed}")
+                return (new_seed,)
+            except (ValueError, IndexError) as e:
+                debug_log(f"Distributor - Error parsing worker_id '{worker_id}': {e}")
+                # 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):
+        return {
+            "required": {
+                "default_value": ("STRING", {"default": ""}),
+                "worker_values": ("STRING", {"default": "{}"}),
+            },
+            "hidden": {
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "worker_id": ("STRING", {"default": ""}),
+            },
+        }
+
+    RETURN_TYPES = (any_type,)
+    RETURN_NAMES = ("value",)
+    FUNCTION = "distribute"
+    CATEGORY = "utils"
+
+    @staticmethod
+    def _coerce(value, value_type):
+        """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, value_type):
+        """Best-effort coercion with graceful fallback to original value."""
+        try:
+            return DistributedValue._coerce(value, value_type)
+        except (TypeError, ValueError):
+            return value
+
+    def distribute(self, default_value, worker_values="{}", is_worker=False, worker_id=""):
+        values = {}
+        value_type = "STRING"
+
+        try:
+            values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
+            if not isinstance(values, dict):
+                values = {}
+        except json.JSONDecodeError as e:
+            debug_log(f"DistributedValue - Error parsing worker_values: {e}")
+            values = {}
+
+        value_type = values.get("_type", "STRING")
+        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,)
+
+        try:
+            if worker_id.startswith("worker_"):
+                idx = int(worker_id.split("_")[1])
+            else:
+                idx = int(worker_id)
+            key = str(idx + 1)  # worker_0 → key "1" (1-indexed)
+            raw = values.get(key, "")
+            if raw:
+                coerced = self._coerce(raw, value_type)
+                debug_log(f"DistributedValue - Worker {idx}: returning '{coerced}'")
+                return (coerced,)
+        except (ValueError, IndexError) as e:
+            debug_log(f"DistributedValue - Error: {e}")
+        debug_log(f"DistributedValue - Worker fallback: returning default '{coerced_default}'")
+        return (coerced_default,)
+
+class DistributedModelName:
+    @classmethod
+    def INPUT_TYPES(cls):
+        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):
+        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, unique_id, values):
+        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, unique_id=None, extra_pnginfo=None):
+        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):
+        if index > 0:
+            index = 0
+        item = super().__getitem__(index)
+        if isinstance(item, str):
+            return any_type
+        return item
+
+class ImageBatchDivider:
+    @classmethod
+    def INPUT_TYPES(s):
+        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, divide_by):
+        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 sequential segments along the time dimension."""
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "audio": ("AUDIO",),
+                "divide_by": ("INT", {
+                    "default": 2,
+                    "min": 1,
+                    "max": 10,
+                    "step": 1,
+                    "display": "number",
+                    "tooltip": "Number of sequential time segments to create"
+                }),
+            }
+        }
+
+    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, divide_by):
+        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):
+        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, width, channels):
+        import torch
+
+        shape = (0, height, width, channels)
+        tensor = torch.zeros(shape, dtype=torch.float32)
+        return (tensor,)

nodes/v3.py

@@ -0,0 +1,181 @@
+"""Native V3 schemas with the existing execution algorithms kept intact.
+
+Runtime objects are private, per-execution helpers, not registered V1 nodes.
+This avoids sharing mutable instance state through V3's sanitized class clones.
+"""
+import comfy.samplers
+from comfy_api.v0_0_2 import io
+
+from . import utilities as _utilities
+from .collector import DistributedCollectorNode as _CollectorRuntime
+from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
+
+
+class DistributedSeed(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedSeed', display_name='Distributed Seed', category='utils',
+            inputs=[io.Int.Input('seed', default=1125899906842, min=0,
+                                 max=1125899906842624, force_input=False)],
+            outputs=[io.Int.Output(display_name='seed')],
+            accept_all_inputs=True,
+        )
+
+    @classmethod
+    def execute(cls, seed, is_worker=False, worker_id=''):
+        return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id))
+
+
+class DistributedValue(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedValue', display_name='Distributed Value', category='utils',
+            inputs=[io.String.Input('default_value', default=''),
+                    io.String.Input('worker_values', default='{}')],
+            outputs=[io.AnyType.Output(display_name='value')],
+            accept_all_inputs=True,
+        )
+
+    @classmethod
+    def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''):
+        return io.NodeOutput(*_utilities.DistributedValue().distribute(
+            default_value, worker_values, is_worker, worker_id))
+
+
+class DistributedModelName(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedModelName', display_name='Distributed Model Name', category='utils',
+            inputs=[io.String.Input('text', default='')],
+            outputs=[io.AnyType.Output(display_name='output')],
+            hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, text):
+        result = _utilities.DistributedModelName().log_input(
+            text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo)
+        return io.NodeOutput(*result['result'], ui=result['ui'])
+
+
+class ImageBatchDivider(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image',
+            inputs=[io.Image.Input('images'),
+                    io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
+                                 display_mode=io.NumberDisplay.number,
+                                 tooltip='Number of parts to divide the batch into')],
+            # The existing frontend still displays only divide_by sockets.
+            outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, images, divide_by):
+        return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by))
+
+
+class AudioBatchDivider(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio',
+            inputs=[io.Audio.Input('audio'),
+                    io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
+                                 display_mode=io.NumberDisplay.number,
+                                 tooltip='Number of sequential time segments to create')],
+            outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, audio, divide_by):
+        return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by))
+
+
+class DistributedEmptyImage(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image',
+            inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1),
+                    io.Int.Input('width', default=64, min=1, max=4096, step=1),
+                    io.Int.Input('channels', default=3, min=1, max=4, step=1)],
+            outputs=[io.Image.Output()],
+        )
+
+    @classmethod
+    def execute(cls, height, width, channels):
+        return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels))
+
+
+class DistributedCollector(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedCollector', display_name='Distributed Collector', category='image',
+            inputs=[io.Boolean.Input('load_balance', default=False,
+                                     tooltip='Run this workflow on one least-busy participant (master included when participating).'),
+                    io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True)],
+            outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
+            is_input_list=True, accept_all_inputs=True,
+        )
+
+    @classmethod
+    def execute(cls, images=None, load_balance=False, audio=None, multi_job_id='',
+                is_worker=False, master_url='', enabled_worker_ids='[]', worker_batch_size=1,
+                worker_id='', pass_through=False, delegate_only=False):
+        return io.NodeOutput(*_CollectorRuntime().run(
+            images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
+            is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
+            worker_batch_size=worker_batch_size, worker_id=worker_id,
+            pass_through=pass_through, delegate_only=delegate_only))
+
+
+class UltimateSDUpscaleDistributed(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='UltimateSDUpscaleDistributed',
+            display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling',
+            inputs=[
+                io.Image.Input('upscaled_image'), io.Model.Input('model'),
+                io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'),
+                io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff),
+                io.Int.Input('steps', default=20, min=1, max=10000),
+                io.Float.Input('cfg', default=8.0, min=0.0, max=100.0),
+                io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS),
+                io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS),
+                io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01),
+                io.Int.Input('tile_width', default=512, min=64, max=2048, step=8),
+                io.Int.Input('tile_height', default=512, min=64, max=2048, step=8),
+                io.Int.Input('padding', default=32, min=0, max=256, step=8),
+                io.Int.Input('mask_blur', default=8, min=0, max=256),
+                io.Boolean.Input('force_uniform_tiles', default=True),
+                io.Boolean.Input('tiled_decode', default=False),
+            ], outputs=[io.Image.Output()], accept_all_inputs=True,
+        )
+
+    @classmethod
+    def fingerprint_inputs(cls, **kwargs):
+        return _UpscaleRuntime.IS_CHANGED(**kwargs)
+
+    @classmethod
+    def execute(cls, 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='', is_worker=False,
+                master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8):
+        return io.NodeOutput(*_UpscaleRuntime().run(
+            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, is_worker,
+            master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
+
+
+NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
+         ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed]

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,7 +1,7 @@
 [project]
 name = "ComfyUI-Distributed"
 description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
-version = "1.2.1"
+version = "1.5.0"
 license = {file = "LICENSE"}
 dependencies  = []
 
@@ -13,3 +13,7 @@ Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
 PublisherId = "robertvoy"
 DisplayName = "ComfyUI-Distributed"
 Icon = "https://raw.githubusercontent.com/robertvoy/ComfyUI-Distributed/refs/heads/main/web/distributed-logo-icon.png"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["."]

runtime/__init__.py

@@ -0,0 +1 @@
+"""Internal extension lifecycle support (not a node provider)."""

runtime/bootstrap.py

@@ -0,0 +1,35 @@
+"""Initialize routes, distributed state and the existing worker lifecycle."""
+import atexit
+import os
+
+import server
+
+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
+from ..upscale.job_store import ensure_tile_jobs_initialized
+
+_initialized = False
+
+
+def initialize():
+    """Called by ComfyExtension.on_load; initialize once per loaded package."""
+    global _initialized
+    if _initialized:
+        return
+
+    ensure_config_exists()
+    from .. import api  # noqa: F401 - registers the existing @routes.* handlers
+    from ..api.queue_orchestration import ensure_distributed_state
+
+    ensure_distributed_state(server.PromptServer.instance)
+    ensure_tile_jobs_initialized()
+
+    if not os.environ.get('COMFYUI_IS_WORKER'):
+        atexit.register(sync_cleanup)
+        delayed_auto_launch()
+        register_async_signals()
+
+    _initialized = True
+    debug_log('Loaded Distributed nodes.')
+    debug_log(f'Config file: {CONFIG_FILE}')

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/test_config_routes.py

@@ -0,0 +1,148 @@
+import copy
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+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"
+
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    class _Routes:
+        def get(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+        def post(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
+    sys.modules["server"] = server_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
+    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
+    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)
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+
+    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), patch.object(
+            config_routes, "save_config", return_value=True
+        ):
+            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_distributed_queue.py

@@ -0,0 +1,398 @@
+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
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, payload):
+        self._payload = payload
+
+    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"
+
+    # Reset package namespace to avoid stale module state across test runs.
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    # aiohttp.web stub
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    # server module stub with route decorators
+    class _Routes:
+        def get(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+        def post(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+    prompt_server_instance = types.SimpleNamespace(
+        routes=_Routes(),
+        distributed_jobs_lock=None,
+        distributed_pending_jobs={},
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
+    sys.modules["server"] = server_module
+
+    # 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
+
+    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", 7, 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
+        auto_prepare: bool
+        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,
+            auto_prepare=bool(data.get("auto_prepare", True)),
+            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)
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+    if created_torch_stub:
+        sys.modules.pop("torch", None)
+    if created_pil_stub:
+        sys.modules.pop("PIL.Image", None)
+        sys.modules.pop("PIL", None)
+
+    return module
+
+
+job_routes = _load_job_routes_module()
+
+
+class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
+    async def test_distributed_queue_happy_path_returns_prompt_metadata(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", 42, 2, {})),
+        ):
+            response = await job_routes.distributed_queue_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
+        self.assertEqual(response.payload.get("number"), 42)
+        self.assertEqual(response.payload.get("node_errors"), {})
+        self.assertTrue(response.payload.get("auto_prepare_supported"))
+
+    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()
+        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        job_routes.prompt_server.distributed_pending_jobs = {"job-1": queue}
+        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_accepts_audio_without_image(self):
+        queue = asyncio.Queue()
+        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
+        request = _FakeRequest(
+            {
+                "job_id": "audio-only-job",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "audio": self._encoded_audio_payload(),
+                "is_last": True,
+            }
+        )
+
+        with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
+            response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        decode_image.assert_not_called()
+        queued = await queue.get()
+        self.assertIsNone(queued["tensor"])
+        self.assertEqual(queued["audio"]["sample_rate"], 44100)
+
+    async def test_job_complete_rejects_payload_without_image_or_audio(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "is_last": True,
+            }
+        )
+
+        response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("image or audio", response.payload.get("message", "").lower())
+
+    async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "image": 123,
+                "audio": self._encoded_audio_payload(),
+                "is_last": True,
+            }
+        )
+
+        response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("image", response.payload.get("message", "").lower())
+
+    async def test_job_complete_returns_400_for_malformed_audio(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
+                "is_last": True,
+            }
+        )
+
+        response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("audio.shape", response.payload.get("message", "").lower())
+
+    def test_decode_audio_payload_rejects_bad_shape(self):
+        bad = {
+            "sample_rate": 44100,
+            "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,247 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _load_media_sync_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py"
+    package_name = "dist_ms_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
+
+    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
+
+    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
+
+        class _FormData:
+            def add_field(self, *args, **kwargs):
+                pass
+
+        aiohttp_module.ClientTimeout = _ClientTimeout
+        aiohttp_module.FormData = _FormData
+        sys.modules["aiohttp"] = aiohttp_module
+
+    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)
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+
+    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_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_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()

tests/api/test_usdu_routes.py

@@ -0,0 +1,304 @@
+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
+
+
+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
+
+
+class _Routes:
+    def post(self, _path):
+        def _decorator(fn):
+            return fn
+
+        return _decorator
+
+    def get(self, _path):
+        def _decorator(fn):
+            return fn
+
+        return _decorator
+
+
+def _load_usdu_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py"
+    package_name = "dist_api_usdu_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
+
+    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
+
+    prompt_server_holder = {
+        "value": types.SimpleNamespace(
+            distributed_tile_jobs_lock=asyncio.Lock(),
+            distributed_pending_tile_jobs={},
+        )
+    }
+
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
+    sys.modules["server"] = server_module
+
+    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")
+
+    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"]
+    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: []
+    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)
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+
+    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("image_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("tile_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,435 @@
+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
+
+
+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):
+        self._payload = payload
+        self.match_info = match_info or {}
+        self.query = query 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, timeout=None):
+        self.calls.append({"url": url, "params": params, "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",
+            "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"
+
+    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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    workers_pkg = types.ModuleType(f"{package_name}.workers")
+    workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
+    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 = 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 _WSMsgType:
+            TEXT = "TEXT"
+            ERROR = "ERROR"
+            CLOSED = "CLOSED"
+
+        class _WebSocketResponse:
+            def __init__(self, *args, **kwargs):
+                self.args = args
+                self.kwargs = kwargs
+
+            async def prepare(self, _request):
+                return None
+
+            async def send_json(self, _payload):
+                return None
+
+            def __aiter__(self):
+                async def _empty():
+                    if False:
+                        yield None
+                return _empty()
+
+        aiohttp_module.ClientTimeout = _ClientTimeout
+        aiohttp_module.WSMsgType = _WSMsgType
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status),
+            WebSocketResponse=_WebSocketResponse,
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    class _Routes:
+        def get(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+        def post(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
+    sys.modules["server"] = server_module
+
+    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
+    network_module.get_server_port = lambda: 8189
+    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", []))
+
+    schemas_module.require_fields = _require_fields
+    schemas_module.validate_worker_id = _validate_worker_id
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    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)
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+    if created_torch_stub:
+        sys.modules.pop("torch", None)
+
+    return module
+
+
+worker_routes = _load_worker_routes_module()
+
+
+class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
+    async def test_network_info_reports_actual_master_and_available_ports(self):
+        with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
+             patch.object(worker_routes, "get_network_ips", return_value=["127.0.0.1"]), \
+             patch.object(worker_routes, "load_config", return_value={"workers": []}), \
+             patch.object(worker_routes, "get_server_port", return_value=8189), \
+             patch.object(worker_routes, "allocate_worker_ports", return_value=[8191, 8192, 8193]) as allocate:
+            response = await worker_routes.get_network_info_endpoint(_FakeRequest())
+        self.assertEqual(response.payload["master_port"], 8189)
+        self.assertEqual(response.payload["local_worker_ports"], [8191, 8192, 8193])
+        self.assertEqual(response.payload["cuda_device_count"], 4)
+        allocate.assert_called_once_with(8189, [], 3)
+
+    async def test_network_info_does_not_reallocate_existing_configuration(self):
+        with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 4, 4)), \
+             patch.object(worker_routes, "get_network_ips", return_value=[]), \
+             patch.object(worker_routes, "load_config", return_value={"workers": [{"id": "existing", "port": 8189}]}), \
+             patch.object(worker_routes, "allocate_worker_ports") as allocate:
+            response = await worker_routes.get_network_info_endpoint(_FakeRequest())
+        self.assertEqual(response.payload["local_worker_ports"], [])
+        allocate.assert_not_called()
+
+    async def test_launch_conflict_is_a_clear_client_error(self):
+        manager = _DummyWorkerManager()
+        config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8189}]}
+        with patch.object(worker_routes, "get_worker_manager", return_value=manager), \
+             patch.object(worker_routes, "load_config", return_value=config), \
+             patch.object(manager, "launch_worker", side_effect=ValueError("Worker port 8189 conflicts with the master port 8189")):
+            response = await worker_routes.launch_worker_endpoint(_FakeRequest({"worker_id": "worker-a"}))
+        self.assertEqual(response.status, 400)
+        self.assertIn("master port 8189", response.payload["message"])
+        self.assertEqual(manager.processes, {})
+
+    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/conftest.py

@@ -0,0 +1,2 @@
+# This conftest.py marks the tests/ directory as the pytest collection root,
+# preventing pytest from traversing into the parent package's __init__.py.

tests/fixtures/v1_node_contracts.json

@@ -0,0 +1,691 @@
+{
+  "DistributedCollector": {
+    "input": {
+      "required": {
+        "load_balance": [
+          "BOOLEAN",
+          {
+            "default": false,
+            "tooltip": "Run this workflow on one least-busy participant (master included when participating)."
+          }
+        ]
+      },
+      "optional": {
+        "images": [
+          "IMAGE"
+        ],
+        "audio": [
+          "AUDIO"
+        ]
+      },
+      "hidden": {
+        "multi_job_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "is_worker": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "master_url": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "enabled_worker_ids": [
+          "STRING",
+          {
+            "default": "[]"
+          }
+        ],
+        "worker_batch_size": [
+          "INT",
+          {
+            "default": 1,
+            "min": 1,
+            "max": 1024
+          }
+        ],
+        "worker_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "pass_through": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "delegate_only": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "load_balance"
+      ],
+      "optional": [
+        "images",
+        "audio"
+      ],
+      "hidden": [
+        "multi_job_id",
+        "is_worker",
+        "master_url",
+        "enabled_worker_ids",
+        "worker_batch_size",
+        "worker_id",
+        "pass_through",
+        "delegate_only"
+      ]
+    },
+    "output": [
+      "IMAGE",
+      "AUDIO"
+    ],
+    "output_name": [
+      "images",
+      "audio"
+    ],
+    "output_is_list": [
+      false,
+      false
+    ],
+    "is_input_list": true,
+    "output_node": false,
+    "category": "image",
+    "display_name": "Distributed Collector"
+  },
+  "DistributedSeed": {
+    "input": {
+      "required": {
+        "seed": [
+          "INT",
+          {
+            "default": 1125899906842,
+            "min": 0,
+            "max": 1125899906842624,
+            "forceInput": false
+          }
+        ]
+      },
+      "hidden": {
+        "is_worker": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "worker_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "seed"
+      ],
+      "hidden": [
+        "is_worker",
+        "worker_id"
+      ]
+    },
+    "output": [
+      "INT"
+    ],
+    "output_name": [
+      "seed"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": false,
+    "category": "utils",
+    "display_name": "Distributed Seed"
+  },
+  "DistributedModelName": {
+    "input": {
+      "required": {
+        "text": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ]
+      },
+      "hidden": {
+        "unique_id": "UNIQUE_ID",
+        "extra_pnginfo": "EXTRA_PNGINFO"
+      }
+    },
+    "input_order": {
+      "required": [
+        "text"
+      ],
+      "hidden": [
+        "unique_id",
+        "extra_pnginfo"
+      ]
+    },
+    "output": [
+      "*"
+    ],
+    "output_name": [
+      "output"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": true,
+    "category": "utils",
+    "display_name": "Distributed Model Name"
+  },
+  "DistributedValue": {
+    "input": {
+      "required": {
+        "default_value": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "worker_values": [
+          "STRING",
+          {
+            "default": "{}"
+          }
+        ]
+      },
+      "hidden": {
+        "is_worker": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "worker_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "default_value",
+        "worker_values"
+      ],
+      "hidden": [
+        "is_worker",
+        "worker_id"
+      ]
+    },
+    "output": [
+      "*"
+    ],
+    "output_name": [
+      "value"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": false,
+    "category": "utils",
+    "display_name": "Distributed Value"
+  },
+  "ImageBatchDivider": {
+    "input": {
+      "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"
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "images",
+        "divide_by"
+      ]
+    },
+    "output": [
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*"
+    ],
+    "output_name": [
+      "batch_1",
+      "batch_2",
+      "batch_3",
+      "batch_4",
+      "batch_5",
+      "batch_6",
+      "batch_7",
+      "batch_8",
+      "batch_9",
+      "batch_10"
+    ],
+    "output_is_list": [
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false
+    ],
+    "is_input_list": false,
+    "output_node": true,
+    "category": "image",
+    "display_name": "Image Batch Divider"
+  },
+  "AudioBatchDivider": {
+    "input": {
+      "required": {
+        "audio": [
+          "AUDIO"
+        ],
+        "divide_by": [
+          "INT",
+          {
+            "default": 2,
+            "min": 1,
+            "max": 10,
+            "step": 1,
+            "display": "number",
+            "tooltip": "Number of sequential time segments to create"
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "audio",
+        "divide_by"
+      ]
+    },
+    "output": [
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*"
+    ],
+    "output_name": [
+      "audio_1",
+      "audio_2",
+      "audio_3",
+      "audio_4",
+      "audio_5",
+      "audio_6",
+      "audio_7",
+      "audio_8",
+      "audio_9",
+      "audio_10"
+    ],
+    "output_is_list": [
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false
+    ],
+    "is_input_list": false,
+    "output_node": true,
+    "category": "audio",
+    "display_name": "Audio Segment Divider"
+  },
+  "DistributedEmptyImage": {
+    "input": {
+      "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
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "height",
+        "width",
+        "channels"
+      ]
+    },
+    "output": [
+      "IMAGE"
+    ],
+    "output_name": [
+      "IMAGE"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": false,
+    "category": "image",
+    "display_name": "Distributed Empty Image"
+  },
+  "UltimateSDUpscaleDistributed": {
+    "input": {
+      "required": {
+        "upscaled_image": [
+          "IMAGE"
+        ],
+        "model": [
+          "MODEL"
+        ],
+        "positive": [
+          "CONDITIONING"
+        ],
+        "negative": [
+          "CONDITIONING"
+        ],
+        "vae": [
+          "VAE"
+        ],
+        "seed": [
+          "INT",
+          {
+            "default": 0,
+            "min": 0,
+            "max": 18446744073709551615
+          }
+        ],
+        "steps": [
+          "INT",
+          {
+            "default": 20,
+            "min": 1,
+            "max": 10000
+          }
+        ],
+        "cfg": [
+          "FLOAT",
+          {
+            "default": 8.0,
+            "min": 0.0,
+            "max": 100.0
+          }
+        ],
+        "sampler_name": [
+          [
+            "euler",
+            "euler_cfg_pp",
+            "euler_ancestral",
+            "euler_ancestral_cfg_pp",
+            "heun",
+            "heunpp2",
+            "exp_heun_2_x0",
+            "exp_heun_2_x0_sde",
+            "dpm_2",
+            "dpm_2_ancestral",
+            "lms",
+            "dpm_fast",
+            "dpm_adaptive",
+            "dpmpp_2s_ancestral",
+            "dpmpp_2s_ancestral_cfg_pp",
+            "dpmpp_sde",
+            "dpmpp_sde_gpu",
+            "dpmpp_2m",
+            "dpmpp_2m_cfg_pp",
+            "dpmpp_2m_sde",
+            "dpmpp_2m_sde_gpu",
+            "dpmpp_2m_sde_heun",
+            "dpmpp_2m_sde_heun_gpu",
+            "dpmpp_3m_sde",
+            "dpmpp_3m_sde_gpu",
+            "ddpm",
+            "lcm",
+            "ipndm",
+            "ipndm_v",
+            "deis",
+            "cfgpp_ud10_ab",
+            "res_multistep",
+            "res_multistep_cfg_pp",
+            "res_multistep_ancestral",
+            "res_multistep_ancestral_cfg_pp",
+            "gradient_estimation",
+            "gradient_estimation_cfg_pp",
+            "er_sde",
+            "seeds_2",
+            "seeds_3",
+            "sa_solver",
+            "sa_solver_pece",
+            "ddim",
+            "uni_pc",
+            "uni_pc_bh2"
+          ]
+        ],
+        "scheduler": [
+          [
+            "simple",
+            "sgm_uniform",
+            "karras",
+            "exponential",
+            "ddim_uniform",
+            "beta",
+            "normal",
+            "linear_quadratic",
+            "kl_optimal"
+          ]
+        ],
+        "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": {
+        "multi_job_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "is_worker": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "master_url": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "enabled_worker_ids": [
+          "STRING",
+          {
+            "default": "[]"
+          }
+        ],
+        "worker_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "tile_indices": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "dynamic_threshold": [
+          "INT",
+          {
+            "default": 8,
+            "min": 1,
+            "max": 64
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "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"
+      ],
+      "hidden": [
+        "multi_job_id",
+        "is_worker",
+        "master_url",
+        "enabled_worker_ids",
+        "worker_id",
+        "tile_indices",
+        "dynamic_threshold"
+      ]
+    },
+    "output": [
+      "IMAGE"
+    ],
+    "output_name": [
+      "IMAGE"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": false,
+    "category": "image/upscaling",
+    "display_name": "Ultimate SD Upscale Distributed (No Upscale)"
+  }
+}

tests/helpers/v3_runtime_check.py

@@ -0,0 +1,252 @@
+"""Exercise the actual loader, schemas and executor in a fresh CPU process.
+
+No HTTP listener, workers, model downloads or live installation changes.
+The original contracts were captured from 32ac027 using the same core.
+"""
+import asyncio
+import inspect
+import json
+import os
+from pathlib import Path
+import sys
+import tempfile
+from unittest.mock import patch
+
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parents[2]
+COMFY_ROOT = Path(sys.argv[1]).resolve()
+sys.path.insert(0, str(COMFY_ROOT))
+os.environ['COMFYUI_IS_WORKER'] = '1'
+import comfy.cli_args
+comfy.cli_args.args.cpu = True
+comfy.cli_args.args.disable_assets = True
+from app.assets.manager import default_asset_manager
+from comfy_api.v0_0_2 import io
+import execution
+import nodes
+import server
+import torch
+
+BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
+
+
+def normalized_input(value):
+    kind = value[0]
+    opts = dict(value[1]) if len(value) > 1 else {}
+    if kind == 'STRING':
+        # V3 serializes the same default single-line widget explicitly.
+        opts.setdefault('multiline', False)
+    if kind == 'COMBO':
+        kind = opts['options']
+        opts = {key: val for key, val in opts.items() if key != 'options'}
+    if isinstance(kind, list):
+        # Single-selection is the V1 dropdown default as well.
+        opts.setdefault('multiselect', False)
+    return [kind, opts]
+
+
+async def map_node(cls, values, extra=None):
+    """Use the actual input parser and sanitized V3 executor clones."""
+    prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {})
+    assert not missing, missing
+    # Use a worker thread for synchronous nodes: collector/upscale bridge back
+    # to PromptServer's running loop, just as ComfyUI's prompt worker does.
+    def run():
+        return asyncio.run(execution._async_map_node_over_list(
+            'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden))
+    return await asyncio.to_thread(run)
+
+
+def check_saved_workflows(mapping):
+    checked = 0
+    occurrences = 0
+    for path in sorted((ROOT / 'workflows').glob('*.json')):
+        workflow = json.loads(path.read_text())
+        local_nodes = {node['id']: node for node in workflow['nodes']}
+        for saved in workflow['nodes']:
+            if saved['type'] not in mapping:
+                continue
+            occurrences += 1
+            info = mapping[saved['type']].GET_NODE_INFO_V1()
+            inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})}
+            for socket in saved.get('inputs', []):
+                assert socket['name'] in inputs, (path.name, saved['id'], socket)
+                assert socket['type'] == inputs[socket['name']][0], (path.name, socket)
+            for index, output in enumerate(saved.get('outputs', [])):
+                assert output['type'] == info['output'][index], (path.name, index)
+            for link in workflow['links']:
+                if link[1] == saved['id']:
+                    assert link[2] < len(info['output']), (path.name, link)
+                    target = local_nodes[link[3]]['inputs'][link[4]]
+                    assert target['type'] == info['output'][link[2]], (path.name, link)
+        checked += 1
+    assert checked == 5 and occurrences == 10, (checked, occurrences)
+    print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences)
+
+
+async def check_execution(mapping, module, prompt_server, asset_manager):
+    result = await map_node(mapping['DistributedSeed'],
+                            {'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'})
+    assert result[0].result == (126,)
+    values = json.dumps({'_type': 'INT', '3': '17'})
+    result = await map_node(mapping['DistributedValue'],
+                            {'default_value': '4', 'worker_values': values,
+                             'is_worker': True, 'worker_id': 'worker_2'})
+    assert result[0].result == (17,)
+    metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}}
+    result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'},
+                            {'extra_pnginfo': metadata})
+    assert result[0].result == ('model.ckpt',)
+    assert result[0].ui == {'text': ['model.ckpt']}
+    assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']]
+    result = await map_node(mapping['DistributedEmptyImage'],
+                            {'height': 8, 'width': 8, 'channels': 3})
+    empty = result[0].result[0]
+    assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0
+    images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3)
+    for count in (1, 3, 10):
+        result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count})
+        assert len(result[0].result) == 10
+        assert torch.equal(torch.cat(result[0].result[:count]), images)
+    audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000}
+    result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3})
+    assert len(result[0].result) == 10
+    assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform'])
+    # INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars.
+    collector = mapping['DistributedCollector']
+    prepared, missing, hidden = execution.get_input_data(
+        {'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list')
+    assert not missing
+    prepared['images'] = [images[:2], images[2:]]
+    prepared['audio'] = [audio, audio]
+    result = await execution._async_map_node_over_list(
+        'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden)
+    assert torch.equal(result[0].result[0], images)
+    assert result[0].result[1]['waveform'].shape[-1] == 66
+    # Exercise real async aggregation without an HTTP endpoint or worker.
+    queue = asyncio.Queue()
+    await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True})
+    prompt_server.distributed_pending_jobs['v3-aggregate'] = queue
+    result = await map_node(collector,
+                            {'images': images[:2], 'load_balance': False,
+                             'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'})
+    assert torch.equal(result[0].result[0], images[:3])
+    assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
+
+    # Prove V3's per-execution class clones do not leak helper instance state;
+    # replace only the GPU/model boundary, not the input parser or executor.
+    upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
+    seen = []
+    def fake_upscale(self, *args):
+        assert not hasattr(self, 'acceptance_marker')
+        self.acceptance_marker = True
+        seen.append(args)
+        return (args[0],)
+    inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [],
+              'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0,
+              'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5,
+              'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0,
+              'force_uniform_tiles': True, 'tiled_decode': False,
+              'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid',
+              'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2',
+              'tile_indices': '[2,3]', 'dynamic_threshold': 9}
+    with patch.object(upscale_runtime, 'run', fake_upscale):
+        for _ in range(2):
+            result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs)
+            assert result[0].result[0] is images
+    assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in (
+        'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id',
+        'tile_indices', 'dynamic_threshold')), seen[0]
+    import math
+    assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job'))
+    assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id=''))
+
+    graph = {
+        '1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}},
+        '2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
+                                                              'pass_through': True}},
+        '3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}},
+        '4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}},
+    }
+    valid = await execution.validate_prompt('v3-graph', graph, None)
+    assert valid[0], valid
+    invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed',
+                     'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}},
+               '2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}}
+    rejected = await execution.validate_prompt('v3-invalid', invalid, None)
+    assert not rejected[0]
+    errors = [error for entry in rejected[3].values() for error in entry['errors']]
+    enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list']
+    assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors
+    import folder_paths
+    scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
+    with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp:
+        with patch.object(folder_paths, 'temp_directory', temp):
+            executor = execution.PromptExecutor(
+                prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
+            await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2])
+            assert executor.success, executor.status_messages
+            history = executor.history_result
+            record = history['outputs']['4']['images'][0]
+            preview = Path(temp) / record.get('subfolder', '') / record['filename']
+            with Image.open(preview) as image:
+                assert image.size == (8, 8) and image.mode == 'RGB'
+                assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
+    print('V3_EXECUTION_OK eight nodes; upscale GPU boundary mocked; preview artifact verified')
+
+
+async def main():
+    asset_manager = default_asset_manager()
+    prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
+    assert not (ROOT / 'distributed.py').exists(), 'obsolete root bootstrap remains'
+    assert await nodes.load_custom_node(str(ROOT)), 'ComfyUI loader rejected the pack'
+    module = sys.modules[str(ROOT).replace('.', '_x_')]
+    assert not hasattr(module, 'NODE_CLASS_MAPPINGS'), 'V1 map shadows V3 entrypoint'
+    extension = await module.comfy_entrypoint()
+    classes = await extension.get_node_list()
+    mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
+    assert len(classes) == len(mapping) == len(BASELINE) == 8
+    assert set(mapping) == set(BASELINE)
+    for node_id, cls in mapping.items():
+        assert issubclass(cls, io.ComfyNode)
+        assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
+        old = dict(BASELINE[node_id])
+        if node_id in ('ImageBatchDivider', 'AudioBatchDivider'):
+            # V1's ByPassTypeTuple advertises '*' when indexed, while its
+            # underlying tuple and existing frontend declare IMAGE/AUDIO.
+            # Native V3 declares all ten existing typed sockets explicitly.
+            assert old['output'] == ['*'] * 10
+            old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
+        new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
+        for group in ('required', 'optional'):
+            old_inputs = old['input'].get(group, {})
+            new_inputs = new['input'].get(group, {})
+            assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
+            for name, original in old_inputs.items():
+                assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
+        for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'):
+            assert old[key] == new[key], (node_id, key, old[key], new[key])
+        # Standard context lives in cls.hidden; orchestrator metadata remains
+        # accepted by its original kwarg name, without creating new widgets.
+        signature = inspect.signature(cls.execute)
+        for name, field in old['input'].get('hidden', {}).items():
+            if isinstance(field, list):
+                assert cls.GET_SCHEMA().accept_all_inputs, node_id
+                assert name in signature.parameters, (node_id, name)
+                assert signature.parameters[name].default == field[1]['default'], (node_id, name)
+            else:
+                assert name in new['input']['hidden'], (node_id, name)
+        expected_hidden = []
+        if node_id == 'DistributedModelName':
+            expected_hidden.extend(['unique_id', 'extra_pnginfo'])
+        if old['output_node']:
+            expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
+        assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
+    print('SCHEMA_PARITY_OK', len(mapping))
+    check_saved_workflows(mapping)
+    await check_execution(mapping, module, prompt_server, asset_manager)
+    print('V3_ACCEPTANCE_OK')
+
+
+asyncio.run(main())

tests/test_async_helpers.py

@@ -0,0 +1,91 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+class _PromptQueue:
+    def __init__(self):
+        self.items = []
+
+    def put(self, item):
+        self.items.append(item)
+
+
+def _load_async_helpers_module():
+    module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
+    package_name = "dist_async_helpers_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
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    execution_module = types.ModuleType("execution")
+
+    async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
+        return (True, None, ["9"], {})
+
+    execution_module.validate_prompt = _validate_prompt
+    execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
+    sys.modules["execution"] = execution_module
+
+    prompt_server = types.SimpleNamespace(
+        trigger_on_prompt=lambda payload: payload,
+        number=12,
+        prompt_queue=_PromptQueue(),
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
+    sys.modules["server"] = server_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.get_server_loop = lambda: None
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    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
+    spec.loader.exec_module(module)
+    return module, prompt_server
+
+
+async_helpers, prompt_server = _load_async_helpers_module()
+
+
+class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
+    async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
+        result = await async_helpers.queue_prompt_payload(
+            {"1": {"class_type": "Node"}},
+            workflow_meta={"id": "workflow-1"},
+            client_id="client-1",
+            include_queue_metadata=True,
+        )
+
+        self.assertIsInstance(result["prompt_id"], str)
+        self.assertTrue(result["prompt_id"])
+        self.assertEqual(result["number"], 12)
+        self.assertEqual(result["node_errors"], {})
+
+        self.assertEqual(prompt_server.number, 13)
+        self.assertEqual(len(prompt_server.prompt_queue.items), 1)
+        queued_item = prompt_server.prompt_queue.items[0]
+        self.assertEqual(queued_item[0], 12)
+        extra_data = queued_item[3]
+        self.assertEqual(extra_data["client_id"], "client-1")
+        self.assertIn("create_time", extra_data)
+        self.assertIsInstance(extra_data["create_time"], int)
+        self.assertGreater(extra_data["create_time"], 0)
+        self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/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[1] / "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/test_collector_list_inputs.py

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

tests/test_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[1] / "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
+
+    logging_module = types.ModuleType(f"{package_name}.logging")
+    logging_module.log = lambda *_args, **_kwargs: None
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.logging"] = logging_module
+
+    constants_module = types.ModuleType(f"{package_name}.constants")
+    constants_module.HEARTBEAT_TIMEOUT = 30
+    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_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/test_detection.py

@@ -0,0 +1,196 @@
+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[1] / "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
+
+    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})
+        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/test_dispatch_selection.py

@@ -0,0 +1,254 @@
+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[1] / "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)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_distributed_value.py

@@ -0,0 +1,208 @@
+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[1] / "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
+
+        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
+
+        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/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[1] / "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/test_network_helpers.py

@@ -0,0 +1,137 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _load_network_module():
+    module_path = Path(__file__).resolve().parents[1] / "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
+
+    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_ignores_stale_saved_port_and_uses_runtime_port(self):
+        cfg = {"master": {"host": "192.168.68.56", "port": 8001}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "http://192.168.68.56:8188",
+        )
+
+    def test_build_master_url_keeps_explicit_port_in_host(self):
+        cfg = {"master": {"host": "192.168.68.56:8001"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "http://192.168.68.56:8001",
+        )
+
+    def test_build_master_url_falls_back_to_server_address(self):
+        cfg = {"master": {"host": "", "port": 8001}}
+        prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "http://127.0.0.1:8190",
+        )
+
+    def test_build_master_callback_url_uses_loopback_for_local_worker(self):
+        cfg = {"master": {"host": "192.168.68.56"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
+        worker = {"id": "w1", "type": "local", "host": "localhost", "port": 8189}
+        self.assertEqual(
+            network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
+            "http://127.0.0.1:8001",
+        )
+
+    def test_build_master_callback_url_keeps_public_master_url_for_remote_worker(self):
+        cfg = {"master": {"host": "192.168.68.56"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
+        worker = {"id": "w2", "type": "remote", "host": "192.168.68.99", "port": 8189}
+        self.assertEqual(
+            network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
+            "http://192.168.68.56:8001",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_payload_parsers.py

@@ -0,0 +1,196 @@
+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[1] / "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_image_and_metadata(self):
+        tiles = pp._parse_tiles_from_form(_make_form(1))
+        self.assertEqual(len(tiles), 1)
+        tile = tiles[0]
+        self.assertIsInstance(tile["image"], PILImage.Image)
+        self.assertEqual(tile["tile_idx"], 0)
+        self.assertEqual(tile["x"], 0)
+        self.assertEqual(tile["y"], 0)
+        self.assertEqual(tile["extracted_width"], 64)
+        self.assertEqual(tile["extracted_height"], 64)
+
+    def test_multiple_tiles_preserve_count_order_and_coordinates(self):
+        tiles = pp._parse_tiles_from_form(_make_form(3))
+        self.assertEqual(len(tiles), 3)
+        for i, tile in enumerate(tiles):
+            self.assertEqual(tile["tile_idx"], i)
+        self.assertEqual(tiles[1]["x"], 64)
+        self.assertEqual(tiles[2]["x"], 128)
+
+    def test_padding_is_parsed_from_form(self):
+        tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
+        self.assertEqual(tiles[0]["padding"], 16)
+
+    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])
+
+    # --- 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/test_queue_request.py

@@ -0,0 +1,151 @@
+import importlib.util
+import unittest
+from pathlib import Path
+
+
+def _load_queue_request_module():
+    module_path = Path(__file__).resolve().parents[1] / "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"] = ["a", 2, 3]
+        payload_data["delegate_master"] = True
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertEqual(payload.enabled_worker_ids, ["a", "2", "3"])
+        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": 3}, {"name": "no-id"}]
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "3"])
+
+    def test_supports_auto_prepare_prompt_fallback(self):
+        payload_data = self._base_payload()
+        payload_data.pop("prompt", None)
+        payload_data["auto_prepare"] = True
+        payload_data["workflow"] = {
+            "prompt": {
+                "10": {"class_type": "DistributedCollector"},
+            }
+        }
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertIn("10", payload.prompt)
+        self.assertTrue(payload.auto_prepare)
+
+    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_auto_prepare_defaults_true(self):
+        payload = parse_queue_request_payload(self._base_payload())
+        self.assertTrue(payload.auto_prepare)
+
+    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_false_still_falls_back_to_workflow_prompt(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)
+        self.assertFalse(payload.auto_prepare)
+
+    def test_auto_prepare_must_be_boolean(self):
+        payload_data = self._base_payload()
+        payload_data["auto_prepare"] = "true"
+        with self.assertRaisesRegex(ValueError, "auto_prepare must be a boolean"):
+            parse_queue_request_payload(payload_data)
+
+    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_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/test_static_mode.py

@@ -0,0 +1,260 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+
+
+def _load_static_mode_module():
+    module_path = Path(__file__).resolve().parents[1] / "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)
+
+    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")
+
+    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
+
+    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_tile_from_master(self, *_args, **_kwargs):
+        self.request_calls += 1
+        return self.tile_sequence.pop(0)
+
+    async def _send_heartbeat_to_master(self, *_args, **_kwargs):
+        self.heartbeat_calls += 1
+
+    async def send_tiles_batch_to_master(
+        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/test_v3_migration.py

@@ -0,0 +1,21 @@
+"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+import pytest
+
+
+def test_native_v3_runtime():
+    comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
+    if not comfy_root:
+        pytest.skip('Set COMFYUI_SOURCE_ROOT to run native ComfyUI V3 acceptance')
+    helper = Path(__file__).parent / 'helpers' / 'v3_runtime_check.py'
+    result = subprocess.run(
+        [sys.executable, str(helper), str(Path(comfy_root).resolve())],
+        capture_output=True, text=True, timeout=120,
+        env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
+    )
+    assert result.returncode == 0, result.stdout + '\n' + result.stderr
+    assert 'V3_ACCEPTANCE_OK' in result.stdout

tests/test_worker_ports.py

@@ -0,0 +1,113 @@
+import importlib.util
+import os
+import socket
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+def load_ports():
+    path = Path(__file__).resolve().parents[1] / "workers" / "ports.py"
+    spec = importlib.util.spec_from_file_location("worker_ports_test", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class WorkerPortTests(unittest.TestCase):
+    def setUp(self):
+        self.ports = load_ports()
+
+    def test_allocation_starts_above_actual_master_and_skips_assigned_and_occupied(self):
+        workers = [
+            {"id": "existing", "host": "localhost", "port": 8190, "enabled": False},
+            {"id": "remote", "host": "remote.example", "port": 8192},
+        ]
+        with patch.object(self.ports, "is_port_available", side_effect=lambda port: port != 8191):
+            self.assertEqual(self.ports.allocate_worker_ports(8189, workers, 3), [8192, 8193, 8194])
+
+    def test_local_host_forms_reserve_ports(self):
+        for host in ["::1", "[::1]", "http://localhost/", "HTTPS://LOCALHOST", "0.0.0.0", None]:
+            with self.subTest(host=host), patch.object(self.ports, "is_port_available", return_value=True):
+                self.assertEqual(self.ports.allocate_worker_ports(8189, [{"id": "local", "host": host, "port": 8190}], 1), [8191])
+
+    def test_exhaustion_is_explicit_not_a_partial_allocation(self):
+        with patch.object(self.ports, "is_port_available", return_value=True):
+            with self.assertRaisesRegex(ValueError, "available.*ports"):
+                self.ports.allocate_worker_ports(65534, [], 2)
+
+    def test_launch_rejects_master_port_without_reassigning_configuration(self):
+        worker = {"id": "manual", "port": 8189}
+        with self.assertRaisesRegex(ValueError, "master.*8189"):
+            self.ports.validate_worker_port(worker, 8189, [])
+        self.assertEqual(worker["port"], 8189)
+
+    def test_launch_rejects_another_local_workers_reserved_port(self):
+        worker = {"id": "manual", "port": 8190}
+        other = {"id": "other", "host": None, "port": 8190, "enabled": False}
+        with self.assertRaisesRegex(ValueError, "other"):
+            self.ports.validate_worker_port(worker, 8189, [worker, other])
+
+    def test_launch_ignores_itself_and_remote_workers_with_same_port(self):
+        worker = {"id": "manual", "port": 8190}
+        other = {"id": "remote", "host": "example.com", "port": 8190}
+        with patch.object(self.ports, "is_port_available", return_value=True):
+            self.ports.validate_worker_port(worker, 8189, [worker, other])
+
+    def test_launch_rejects_an_occupied_port(self):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
+            listener.bind(("127.0.0.1", 0))
+            listener.listen()
+            port = listener.getsockname()[1]
+            self.assertFalse(self.ports.is_port_available(port))
+            with self.assertRaisesRegex(ValueError, "already in use"):
+                self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
+
+    @unittest.skipIf(os.name == "nt", "asyncio does not reuse addresses on Windows")
+    def test_recently_closed_connection_does_not_block_worker_restart(self):
+        with socket.socket() as listener:
+            listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            listener.bind(("127.0.0.1", 0))
+            listener.listen(1)
+            port = listener.getsockname()[1]
+            with socket.create_connection(("127.0.0.1", port), timeout=2) as client:
+                accepted, _ = listener.accept()
+                with accepted:
+                    accepted.settimeout(2)
+                    accepted.shutdown(socket.SHUT_WR)
+                    self.assertEqual(client.recv(1), b"")
+                    client.shutdown(socket.SHUT_WR)
+                    self.assertEqual(accepted.recv(1), b"")
+        self.assertTrue(self.ports.is_port_available(port))
+        self.ports.validate_worker_port({"id": "restart", "port": port}, 1, [])
+        with socket.socket() as restarted:
+            restarted.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            restarted.bind(("127.0.0.1", port))
+            restarted.listen(1)
+
+    def test_active_listener_with_reuseaddr_is_still_a_conflict(self):
+        with socket.socket() as listener:
+            listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            listener.bind(("127.0.0.1", 0))
+            listener.listen(1)
+            self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
+
+    @unittest.skipUnless(socket.has_ipv6, "IPv6 unavailable")
+    def test_ipv6_only_listener_is_still_a_conflict(self):
+        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as listener:
+            listener.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
+            try:
+                listener.bind(("::1", 0))
+            except OSError as exc:
+                self.skipTest(f"IPv6 loopback unavailable: {exc}")
+            listener.listen(1)
+            self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
+
+    def test_invalid_ports_fail_clearly(self):
+        for port in [0, 65536, "not-a-port", None, True, 8190.5]:
+            with self.subTest(port=port), self.assertRaisesRegex(ValueError, "port"):
+                self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_worker_process_runtime.py

@@ -0,0 +1,352 @@
+import importlib.util
+import os
+import sys
+import tempfile
+import types
+import unittest
+from argparse import Namespace
+from pathlib import Path
+from unittest.mock import patch
+
+from sqlalchemy import create_engine
+from sqlalchemy.engine import URL, make_url
+
+
+def _load_process_module(module_filename: str):
+    module_path = Path(__file__).resolve().parents[1] / "workers" / "process" / module_filename
+    package_name = "dist_proc_testpkg"
+    module_name = module_filename[:-3]
+
+    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__ = [str(module_path.parents[1])]
+    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
+
+    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")
+    process_module.get_python_executable = lambda: "/usr/bin/test-python"
+    process_module.is_process_alive = lambda _pid: False
+    process_module.terminate_process = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.process"] = process_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"workers": [], "settings": {"stop_workers_on_master_exit": False}}
+    config_module.save_config = lambda _config: None
+    sys.modules[f"{package_name}.utils.config"] = config_module
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.PROCESS_TERMINATION_TIMEOUT = 1
+    constants_module.PROCESS_WAIT_TIMEOUT = 1
+    constants_module.WORKER_CHECK_INTERVAL = 0.01
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.get_server_port = lambda: 8189
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.workers.process.{module_name}",
+        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
+
+
+root_discovery_module = _load_process_module("root_discovery.py")
+launch_builder_module = _load_process_module("launch_builder.py")
+
+
+class ComfyRootDiscoveryTests(unittest.TestCase):
+    def test_prefers_loaded_comfyui_module_path(self):
+        discovery = root_discovery_module.ComfyRootDiscovery()
+        server_module = types.SimpleNamespace(__file__="/opt/ComfyUI/server.py")
+
+        def fake_exists(path):
+            return path == "/opt/ComfyUI/main.py"
+
+        with patch.dict(sys.modules, {"server": server_module}, clear=False), \
+             patch.object(root_discovery_module.os.path, "exists", side_effect=fake_exists), \
+             patch.dict(root_discovery_module.os.environ, {}, clear=True):
+            self.assertEqual(discovery.find_comfy_root(), "/opt/ComfyUI")
+
+
+class LaunchCommandBuilderTests(unittest.TestCase):
+    def build_command(self, root, worker=None, runtime=None):
+        builder = launch_builder_module.LaunchCommandBuilder()
+        worker = worker or {"id": "worker-a", "port": 8190}
+        runtime = runtime if runtime is not None else Namespace(database_url=None)
+        with patch.object(builder, "_get_runtime_args", return_value=runtime):
+            return builder.build_launch_command(worker, str(root))
+
+    def test_database_is_unique_and_stable_by_worker_id(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            first = self.build_command(root, {"id": "worker-a", "port": 8190})
+            renamed = self.build_command(root, {
+                "id": "worker-a", "port": 9001, "name": "Renamed", "cuda_device": 3,
+            })
+            second = self.build_command(root, {"id": "worker-b", "port": 8191})
+            database = first[first.index("--database-url") + 1]
+            self.assertEqual(database, renamed[renamed.index("--database-url") + 1])
+            self.assertNotEqual(database, second[second.index("--database-url") + 1])
+            self.assertTrue(database.startswith("sqlite:///" + root.as_posix() + "/"))
+            self.assertTrue(Path(database.removeprefix("sqlite:///")).parent.is_dir())
+            self.assertNotIn("--disable-assets", first)
+
+    def test_database_uses_effective_user_or_base_directory(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            for extra_args, expected in [
+                (f'--user-directory "{root / "custom user"}"', root / "custom user"),
+                (f'--base-directory "{root / "custom base"}"', root / "custom base" / "user"),
+            ]:
+                with self.subTest(extra_args=extra_args):
+                    cmd = self.build_command(root, {
+                        "id": "worker-a", "port": 8190, "extra_args": extra_args,
+                    })
+                    database = cmd[cmd.index("--database-url") + 1]
+                    self.assertTrue(database.startswith("sqlite:///" + expected.as_posix() + "/"))
+            cmd = self.build_command(root, runtime=Namespace(
+                database_url="sqlite:///master.db", user_directory=str(root / "runtime user"),
+            ))
+            database = cmd[cmd.index("--database-url") + 1]
+            self.assertTrue(database.startswith("sqlite:///" + (root / "runtime user").as_posix() + "/"))
+            self.assertNotEqual(database, "sqlite:///master.db")
+
+    def test_relative_database_directories_use_worker_cwd(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory) / "ComfyUI"
+            root.mkdir()
+            (root / "main.py").touch()
+            launcher = Path(directory) / "launcher"
+            launcher.mkdir()
+            cases = [
+                (Namespace(database_url=None, user_directory="profiles"), "", root / "profiles"),
+                (Namespace(database_url=None, base_directory="data"), "", root / "data/user"),
+                (Namespace(database_url=None), "--user-directory=profiles", root / "profiles"),
+                (Namespace(database_url=None), "--base-directory data", root / "data/user"),
+            ]
+            for runtime, extra, expected in cases:
+                with self.subTest(runtime=runtime, extra=extra), \
+                     patch.object(os, "getcwd", return_value=str(launcher)):
+                    cmd = self.build_command(root, {
+                        "id": "worker-a", "port": 8190, "extra_args": extra,
+                    }, runtime)
+                    database = make_url(cmd[cmd.index("--database-url") + 1]).database
+                    self.assertEqual(Path(database).parent, expected / "distributed/workers")
+            self.assertFalse((launcher / "profiles").exists())
+            self.assertFalse((launcher / "data").exists())
+
+    def test_database_url_round_trips_special_characters_without_collisions(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            for name in ["user?profile", "user%3Fprofile"]:
+                with self.subTest(name=name):
+                    user = root / name
+                    runtime = Namespace(database_url=None, user_directory=str(user))
+                    # SQLAlchemy 2.0 cannot round-trip '?' in a filename. Fail
+                    # clearly on those versions rather than silently sharing a DB.
+                    sample = (user / "test.db").as_posix()
+                    serialized = URL.create("sqlite", database=sample).render_as_string()
+                    if make_url(serialized).database != sample:
+                        with self.assertRaisesRegex(ValueError, "SQLAlchemy.*database path"):
+                            self.build_command(root, runtime=runtime)
+                        self.assertFalse(user.exists())
+                        continue
+                    databases = []
+                    for worker_id in ["worker-a", "worker-b"]:
+                        cmd = self.build_command(root, {"id": worker_id, "port": 8190}, runtime)
+                        url = cmd[cmd.index("--database-url") + 1]
+                        database = make_url(url).database
+                        self.assertEqual(Path(database).parent, user / "distributed/workers")
+                        engine = create_engine(url)
+                        try:
+                            with engine.connect() as connection:
+                                actual = connection.exec_driver_sql("PRAGMA database_list").one()[2]
+                                self.assertEqual(Path(actual), Path(database))
+                        finally:
+                            engine.dispose()
+                        databases.append(database)
+                    self.assertNotEqual(*databases)
+
+    def test_explicit_database_and_disabled_assets_are_preserved(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            for extra in ["--database-url sqlite:///explicit.db", "--database-url=sqlite:///explicit.db", "--disable-assets"]:
+                with self.subTest(extra=extra):
+                    cmd = self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
+                    self.assertEqual(sum(arg.split("=")[0] == "--database-url" for arg in cmd), 0 if extra == "--disable-assets" else 1)
+                    self.assertFalse((root / "user").exists())
+
+    def test_older_comfyui_without_database_flag_keeps_existing_command(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            cmd = self.build_command(root, runtime=Namespace())
+            self.assertNotIn("--database-url", cmd)
+            self.assertFalse((root / "user").exists())
+
+    def test_worker_id_cannot_escape_database_directory(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            databases = []
+            for worker_id in ["../outside", "a/b", "a_b"]:
+                cmd = self.build_command(root, {"id": worker_id, "port": 8190})
+                path = Path(cmd[cmd.index("--database-url") + 1].removeprefix("sqlite:///"))
+                self.assertTrue(path.is_relative_to(root / "user"))
+                databases.append(path)
+            self.assertEqual(len(set(databases)), 3)
+
+    def test_extra_args_cannot_silently_override_the_configured_port(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            for extra in ["--port 8189", "--port=8189"]:
+                with self.subTest(extra=extra), self.assertRaisesRegex(ValueError, "configured worker port"):
+                    self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
+
+    def test_inherits_runtime_layout_args_for_desktop(self):
+        builder = launch_builder_module.LaunchCommandBuilder()
+        runtime_args = Namespace(
+            listen="127.0.0.1",
+            base_directory="C:/Users/test/ComfyUI",
+            temp_directory=None,
+            input_directory="C:/Users/test/ComfyUI/input",
+            output_directory="C:/Users/test/ComfyUI/output",
+            user_directory="C:/Users/test/ComfyUI/user",
+            front_end_root="C:/Program Files/ComfyUI/web_custom_versions/desktop_app",
+            extra_model_paths_config=[["C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml"]],
+            enable_manager=True,
+            disable_manager_ui=False,
+            enable_manager_legacy_ui=False,
+            windows_standalone_build=True,
+            log_stdout=True,
+            verbose="INFO",
+            enable_cors_header="*",
+        )
+        comfy_module = types.ModuleType("comfy")
+        comfy_cli_args = types.ModuleType("comfy.cli_args")
+        comfy_cli_args.args = runtime_args
+
+        worker_config = {
+            "port": 9001,
+            "extra_args": "--preview-method auto",
+        }
+
+        def fake_exists(path):
+            return path == "/desktop/ComfyUI/main.py"
+
+        with patch.dict(
+            sys.modules,
+            {"comfy": comfy_module, "comfy.cli_args": comfy_cli_args},
+            clear=False,
+        ), patch.object(launch_builder_module.os.path, "exists", side_effect=fake_exists):
+            cmd = builder.build_launch_command(worker_config, "/desktop/ComfyUI")
+
+        self.assertEqual(cmd[:2], ["/usr/bin/test-python", "/desktop/ComfyUI/main.py"])
+        self.assertIn("--listen", cmd)
+        self.assertIn("127.0.0.1", cmd)
+        self.assertIn("--base-directory", cmd)
+        self.assertIn("C:/Users/test/ComfyUI", cmd)
+        self.assertIn("--input-directory", cmd)
+        self.assertIn("--output-directory", cmd)
+        self.assertIn("--user-directory", cmd)
+        self.assertIn("--front-end-root", cmd)
+        self.assertIn("--extra-model-paths-config", cmd)
+        self.assertIn("C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml", cmd)
+        self.assertIn("--enable-manager", cmd)
+        self.assertIn("--windows-standalone-build", cmd)
+        self.assertIn("--log-stdout", cmd)
+        self.assertIn("--disable-auto-launch", cmd)
+        self.assertIn("--enable-cors-header", cmd)
+        self.assertIn("*", cmd)
+        self.assertIn("--port", cmd)
+        self.assertIn("9001", cmd)
+        self.assertNotIn("--auto-launch", cmd)
+
+
+class ProcessLaunchPortTests(unittest.TestCase):
+    def test_launch_before_master_listens_uses_configured_master_port(self):
+        lifecycle_module = _load_process_module("lifecycle.py")
+        cli_args = types.ModuleType("comfy.cli_args")
+        cli_args.args = Namespace(port=8189)
+        for missing in [AttributeError("PromptServer has no port"), None]:
+            with self.subTest(missing=missing), tempfile.TemporaryDirectory() as directory:
+                manager = types.SimpleNamespace(
+                    find_comfy_root=lambda: directory,
+                    build_launch_command=lambda _worker, _root: ["python", "main.py"],
+                    processes={}, save_processes=lambda: None,
+                )
+                with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
+                     patch.object(lifecycle_module, "get_server_port", side_effect=missing if isinstance(missing, Exception) else None, return_value=None), \
+                     patch.object(lifecycle_module, "validate_worker_port") as validate, \
+                     patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
+                    pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
+                validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
+                spawn.assert_called_once()
+                self.assertEqual(pid, 1234)
+
+    def test_pre_listen_launch_still_rejects_the_configured_master_port(self):
+        lifecycle_module = _load_process_module("lifecycle.py")
+        cli_args = types.ModuleType("comfy.cli_args")
+        cli_args.args = Namespace(port=8189)
+        manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
+        with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
+             patch.object(lifecycle_module, "get_server_port", side_effect=AttributeError("port")), \
+             patch.object(lifecycle_module.subprocess, "Popen") as spawn:
+            with self.assertRaisesRegex(ValueError, "master port 8189"):
+                lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
+        spawn.assert_not_called()
+
+    def test_conflict_is_rejected_before_building_or_spawning(self):
+        lifecycle_module = _load_process_module("lifecycle.py")
+        manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
+        with patch.object(lifecycle_module.subprocess, "Popen") as spawn:
+            with self.assertRaisesRegex(ValueError, "master port 8189"):
+                lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
+        spawn.assert_not_called()
+        self.assertEqual(manager.processes, {})
+
+    def test_available_port_reaches_existing_launch_path(self):
+        lifecycle_module = _load_process_module("lifecycle.py")
+        with tempfile.TemporaryDirectory() as directory:
+            manager = types.SimpleNamespace(
+                find_comfy_root=lambda: directory,
+                build_launch_command=lambda _worker, _root: ["python", "main.py", "--port", "8190"],
+                processes={}, save_processes=lambda: None,
+            )
+            with patch.object(lifecycle_module, "validate_worker_port") as validate, \
+                 patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
+                pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
+            validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
+            spawn.assert_called_once()
+            self.assertEqual(pid, 1234)
+            self.assertEqual(manager.processes["manual"]["pid"], 1234)
+
+
+if __name__ == "__main__":
+    unittest.main()

upscale/conditioning.py

@@ -0,0 +1,32 @@
+import copy
+
+
+def clone_control_chain(control, clone_hint=True):
+    """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, clone_hints=True):
+    """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,135 @@
+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 as _requeue_usdu
+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 _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 _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 _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 _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 _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 _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 _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
+        """Check for timed out workers and requeue their assigned images."""
+        return await _requeue_usdu(multi_job_id, batch_size)

upscale/job_store.py

@@ -0,0 +1,180 @@
+import asyncio
+import os
+import time
+from typing import List, Optional
+
+import server
+
+from ..utils.logging import debug_log
+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)))
+
+
+def ensure_tile_jobs_initialized():
+    """Ensure tile job storage is initialized on the server instance."""
+    prompt_server = server.PromptServer.instance
+    if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
+        debug_log("Initializing persistent tile job queue on server instance.")
+        prompt_server.distributed_pending_tile_jobs = {}
+        prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
+    else:
+        invalid_job_ids = [
+            job_id
+            for job_id, job_data in prompt_server.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 prompt_server.distributed_pending_tile_jobs[job_id]
+    return prompt_server
+
+
+async def _init_job_queue(
+    multi_job_id,
+    mode,
+    batch_size=None,
+    num_tiles_per_image=None,
+    all_indices=None,
+    enabled_workers=None,
+    batched_static: bool = False,
+):
+    """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 mode == 'dynamic':
+            job_data = ImageJobState(multi_job_id=multi_job_id)
+        elif mode == 'static':
+            job_data = TileJobState(multi_job_id=multi_job_id)
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
+
+        job_data.worker_status = {w: time.time() for w in enabled_workers or []}
+        job_data.assigned_to_workers = {w: [] for w in enabled_workers or []}
+
+        if mode == 'dynamic':
+            job_data.batch_size = int(batch_size or 0)
+            pending_queue = job_data.pending_images
+            for i in (all_indices or range(int(batch_size or 0))):
+                await pending_queue.put(i)
+            debug_log(f"Initialized image queue with {batch_size} pending items")
+        elif mode == 'static':
+            job_data.num_tiles_per_image = int(num_tiles_per_image or 0)
+            job_data.batch_size = int(batch_size or 0)
+            job_data.batched_static = bool(batched_static)
+            # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
+            pending_queue = job_data.pending_tasks
+            if batched_static and num_tiles_per_image is not None:
+                for i in range(num_tiles_per_image):
+                    await pending_queue.put(i)
+            else:
+                total_tiles = int(batch_size or 0) * int(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: Optional[List[int]] = None,
+):
+    """Initialize queue for dynamic mode (per-image), with collector fields."""
+    await _init_job_queue(
+        multi_job_id,
+        '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],
+):
+    """Initialize queue for static mode (batched-per-tile)."""
+    await _init_job_queue(
+        multi_job_id,
+        '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):
+    """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):
+    """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, task_id, result):
+    """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):
+    """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}")

upscale/job_timeout.py

@@ -0,0 +1,150 @@
+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 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

upscale/modes/dynamic.py

@@ -0,0 +1,313 @@
+import asyncio, torch
+from PIL import Image
+import comfy.model_management
+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 TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT
+from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job
+
+
+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_image_from_master`, `_send_full_image_to_master`, `_send_heartbeat_to_master`).
+    """
+
+    def process_master_dynamic(self, 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, enabled_workers):
+        """Dynamic mode for large batches - assigns whole images to workers dynamically, including master."""
+        # 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 = self.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(
+                self._get_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]
+                image_seed = seed
+                
+                # Pre-slice conditioning once per image (not per tile)
+                positive_sliced, negative_sliced = self._slice_conditioning(positive, 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 = self._process_and_blend_tile(
+                        tile_idx, pos, source_tensor, local_image,
+                        model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
+                        sampler_name, scheduler, denoise, tile_width, tile_height,
+                        padding, mask_blur, width, height, force_uniform_tiles,
+                        tiled_decode, batch_idx=image_idx
+                    )
+                    
+                    # Yield after each tile to minimize worker downtime
+                    run_async_in_server_loop(self._async_yield(), timeout=0.1)
+                    # Note: No per-tile drain here – that's what makes this "per-image"
+                
+                result_images[image_idx] = local_image
+                
+                # Mark as completed
+                run_async_in_server_loop(
+                    self._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(
+                    self._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(
+                    self._get_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(self._async_yield(), timeout=0.1)
+            else:
+                # Queue empty: collect any queued worker results to update progress
+                drained_count = run_async_in_server_loop(
+                    self._drain_worker_results_queue(multi_job_id),
+                    timeout=5.0
+                )
+                run_async_in_server_loop(self._async_yield(), timeout=0.1)  # Yield after drain
+                
+                # Check for timed out workers and requeue their images
+                requeued_count = run_async_in_server_loop(
+                    self._check_and_requeue_timed_out_workers(multi_job_id, batch_size),
+                    timeout=5.0
+                )
+                run_async_in_server_loop(self._async_yield(), timeout=0.1)  # Yield after requeue
+                
+                if requeued_count > 0:
+                    log(f"Requeued {requeued_count} images from timed out workers")
+                    consecutive_retries = 0  # Reset since we have work to do
+                    continue
+
+                # Now check total completed (includes newly collected)
+                completed_now = run_async_in_server_loop(
+                    self._get_total_completed_count(multi_job_id),
+                    timeout=1.0
+                )
+                
+                log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
+                
+                if completed_now >= batch_size:
+                    break
+
+                run_async_in_server_loop(self._async_yield(), timeout=0.1)  # Yield before pending check
+                
+                # Check if there are pending images in the queue (could be requeued)
+                pending_count = run_async_in_server_loop(
+                    self._get_pending_count(multi_job_id),
+                    timeout=1.0
+                )
+                
+                if pending_count > 0:
+                    consecutive_retries = 0  # Reset retries since there's work to do
+                    continue
+
+                consecutive_retries += 1
+                if consecutive_retries >= max_consecutive_retries:
+                    log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.")
+                    break  # Force exit to collection phase
+
+                debug_log("Waiting for workers")
+                # Use async sleep to allow event loop to process worker requests
+                run_async_in_server_loop(asyncio.sleep(2), timeout=3.0)
+        
+        debug_log(f"Master processed {processed_count} images locally")
+        
+        # Get all completed images to check what needs to be collected
+        all_completed = run_async_in_server_loop(
+            self._get_all_completed_images(multi_job_id),
+            timeout=5.0
+        )
+        
+        # Calculate how many we still need to collect
+        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")
+            # Use the unified worker timeout for the collection phase
+            collection_timeout = float(get_worker_timeout_seconds())
+            collected_images = run_async_in_server_loop(
+                self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count),
+                timeout=collection_timeout
+            )
+            
+            # Merge collected with already completed
+            all_completed.update(collected_images)
+        
+        # Update result images with all completed images
+        for idx, processed_img in all_completed.items():
+            if idx < batch_size:
+                result_images[idx] = processed_img
+        
+        # Convert back to tensor
+        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()
+        
+        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, 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):
+        """Worker processing in dynamic mode - processes whole images."""
+        # Round tile dimensions
+        tile_width = self.round_to_multiple(tile_width)
+        tile_height = self.round_to_multiple(tile_height)
+
+        # Get dimensions and tile grid
+        batch_size, height, width, _ = upscaled_image.shape
+        all_tiles = self.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
+
+        # Poll for job readiness to avoid races during master init
+        max_poll_attempts = 20  # ~20s at 1s sleep
+        if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts):
+            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
+            image_idx, estimated_remaining = run_async_in_server_loop(
+                self._request_image_from_master(multi_job_id, master_url, worker_id),
+                timeout=TILE_WAIT_TIMEOUT
+            )
+
+            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
+            image_seed = seed
+
+            # Pre-slice conditioning once per image (not per tile)
+            positive_sliced, negative_sliced = self._slice_conditioning(positive, 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 = self._process_and_blend_tile(
+                    tile_idx, pos, source_tensor, local_image,
+                    model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
+                    sampler_name, scheduler, denoise, tile_width, tile_height,
+                    padding, mask_blur, width, height, force_uniform_tiles,
+                    tiled_decode, batch_idx=image_idx
+                )
+                run_async_in_server_loop(
+                    self._send_heartbeat_to_master(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(
+                    self._send_full_image_to_master(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(
+                    self._send_heartbeat_to_master(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(
+                self._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,72 @@
+import math, torch
+from PIL import Image
+from ...utils.logging import debug_log, log
+from ...utils.image import tensor_to_pil, pil_to_tensor
+
+
+class SingleGpuModeMixin:
+    def process_single_gpu(self, 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):
+        """Process all tiles on a single GPU (no distribution), batching per tile like USDU."""
+        # 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
+        all_tiles = self.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}"
+        )
+
+        # 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(self.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()
+
+            # Extract batched tile
+            tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
+                source_batch, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
+            )
+
+            # Process batch
+            region = (x1, y1, x1 + ew, y1 + eh)
+            processed_batch = self.process_tiles_batch(tile_batch, model, positive, negative, vae,
+                                                       seed, steps, cfg, sampler_name, scheduler, denoise,
+                                                       tiled_decode, region, (width, height))
+
+            # Blend results back into each image using cached mask
+            tile_mask = tile_masks[tile_idx]
+            for b in range(batch_size):
+                tile_pil = tensor_to_pil(processed_batch, b)
+                # Resize back to extracted size
+                if tile_pil.size != (ew, eh):
+                    tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
+                result_images[b] = self.blend_tile(result_images[b], tile_pil, 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,570 @@
+import asyncio, time, torch
+from PIL import Image
+import comfy.model_management
+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 (
+    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,
+    _mark_task_completed, _cleanup_job, _drain_results_queue, _get_completed_count,
+)
+from ..job_models import TileJobState
+
+
+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_to_master`, `_request_tile_from_master`, `_send_heartbeat_to_master`).
+    """
+
+    def _poll_job_ready(self, multi_job_id, master_url, worker_id=None, max_attempts=JOB_POLL_MAX_ATTEMPTS):
+        """Poll master for job readiness to avoid worker/master initialization race."""
+        for attempt in range(max_attempts):
+            ready = run_async_in_server_loop(
+                self._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,
+        model,
+        positive,
+        negative,
+        vae,
+        seed,
+        steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        denoise,
+        tiled_decode,
+        width,
+        height,
+    ):
+        """Extract one tile position for the whole batch and process it."""
+        tx, ty = all_tiles[tile_id]
+        tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
+            upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
+        )
+        region = (x1, y1, x1 + ew, y1 + eh)
+        processed_batch = self.process_tiles_batch(
+            tile_batch, model, positive, negative, vae,
+            seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
+            region, (width, height)
+        )
+        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,
+    ):
+        """Send accumulated tile payloads to master and return a fresh accumulator."""
+        if not processed_tiles:
+            if is_final_flush:
+                run_async_in_server_loop(
+                    self.send_tiles_batch_to_master(
+                        [],
+                        multi_job_id,
+                        master_url,
+                        padding,
+                        worker_id,
+                        is_final_flush=True,
+                    ),
+                    timeout=TILE_SEND_TIMEOUT,
+                )
+            return processed_tiles
+        run_async_in_server_loop(
+            self.send_tiles_batch_to_master(
+                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,
+        model,
+        positive,
+        negative,
+        vae,
+        seed,
+        steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        denoise,
+        tiled_decode,
+        width,
+        height,
+    ):
+        """Process one tile_id across the batch and blend into result_images."""
+        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,
+            model,
+            positive,
+            negative,
+            vae,
+            seed,
+            steps,
+            cfg,
+            sampler_name,
+            scheduler,
+            denoise,
+            tiled_decode,
+            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):
+            tile_pil = tensor_to_pil(processed_batch, b)
+            if tile_pil.size != (ew, eh):
+                tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
+            result_images[b] = self.blend_tile(result_images[b], tile_pil, 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, 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_workers):
+        """Worker static mode processing with optional dynamic queue pulling."""
+        # Round tile dimensions
+        tile_width = self.round_to_multiple(tile_width)
+        tile_height = self.round_to_multiple(tile_height)
+        
+        # Get dimensions and calculate tiles
+        _, height, width, _ = upscaled_image.shape
+        all_tiles = self.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(self.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):
+            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
+            tile_idx, estimated_remaining, batched_static = run_async_in_server_loop(
+                self._request_tile_from_master(multi_job_id, master_url, worker_id),
+                timeout=TILE_WAIT_TIMEOUT
+            )
+
+            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,
+                model,
+                positive,
+                negative,
+                vae,
+                seed,
+                steps,
+                cfg,
+                sampler_name,
+                scheduler,
+                denoise,
+                tiled_decode,
+                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] = self.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(
+                    self._send_heartbeat_to_master(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
+                )
+
+        # 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
+        )
+        
+        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):
+        """Async helper for collection and monitoring in static mode.
+        Returns collected tasks dict. Caller should check if all tasks are complete."""
+        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 self._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 self._get_all_completed_tasks(multi_job_id)
+
+    def _process_master_static_sync(self, 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, enabled_workers,
+                                    all_tiles, num_tiles_per_image):
+        """Static mode master processing with optional dynamic queue pulling."""
+        batch_size = upscaled_image.shape[0]
+        _, height, width, _ = upscaled_image.shape
+        total_tiles = batch_size * num_tiles_per_image
+        
+        # Convert batch to PIL list for processing
+        result_images = []
+        for b in range(batch_size):
+            image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
+            result_images.append(image_pil.copy())
+        
+        # Initialize queue: pending queue holds tile ids (batched per tile)
+        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}"
+        )
+
+        # Precompute masks for all tile positions to avoid repeated Gaussian blur work during blending
+        tile_masks = []
+        for idx, (tx, ty) in enumerate(all_tiles):
+            tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
+
+        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_idx = run_async_in_server_loop(
+                self._get_next_tile_index(multi_job_id),
+                timeout=5.0
+            )
+            if tile_idx is not None:
+                consecutive_no_tile = 0
+                tile_id = tile_idx
+                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,
+                    model,
+                    positive,
+                    negative,
+                    vae,
+                    seed,
+                    steps,
+                    cfg,
+                    sampler_name,
+                    scheduler,
+                    denoise,
+                    tiled_decode,
+                    width,
+                    height,
+                )
+                log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})")
+            else:
+                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)
+        master_processed_count = processed_count
+        
+        # Continue processing any remaining tiles while collecting worker results
+        remaining_tiles = total_tiles - master_processed_count
+        if remaining_tiles > 0:
+            debug_log(f"Master waiting for {remaining_tiles} tiles from workers")
+            
+            # Collect worker results using async operations
+            try:
+                # Wait until either all tasks are collected or there are no active workers left
+                collected_tasks = run_async_in_server_loop(
+                    self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles),
+                    timeout=None
+                )
+            except comfy.model_management.InterruptProcessingException:
+                # Clean up job on interruption
+                run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
+                raise
+            
+            # Check if we need to process any remaining tasks locally after collection
+            completed_count = len(collected_tasks)
+            if completed_count < total_tiles:
+                log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures")
+                
+                # Process any remaining pending tasks (batched-per-tile)
+                while True:
+                    # Check for user interruption
+                    comfy.model_management.throw_exception_if_processing_interrupted()
+
+                    # Get next tile_id from pending queue
+                    tile_id = run_async_in_server_loop(
+                        self._get_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,
+                        model,
+                        positive,
+                        negative,
+                        vae,
+                        seed,
+                        steps,
+                        cfg,
+                        sampler_name,
+                        scheduler,
+                        denoise,
+                        tiled_decode,
+                        width,
+                        height,
+                    )
+        else:
+            # Master processed all tiles
+            collected_tasks = run_async_in_server_loop(
+                self._get_all_completed_tasks(multi_job_id),
+                timeout=5.0
+            )
+        
+        # Blend worker tiles synchronously in deterministic tile order.
+        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):
+            # Skip tiles that don't have tensor data (already processed)
+            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
+            
+            # Blend tile synchronously
+            x = tile_data.get('x', 0)
+            y = tile_data.get('y', 0)
+            # Prefer PIL image if present to avoid reconversion
+            if 'image' in tile_data:
+                tile_pil = tile_data['image']
+            else:
+                tile_tensor = tile_data['tensor']
+                tile_pil = tensor_to_pil(tile_tensor, 0)
+            orig_x, orig_y = all_tiles[tile_idx]
+            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] = self.blend_tile(result_images[batch_idx], tile_pil,
+                                                      x, y, (extracted_width, extracted_height), tile_mask, padding)
+        
+        try:
+            # Convert back to tensor
+            if batch_size == 1:
+                result_tensor = pil_to_tensor(result_images[0])
+            else:
+                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()
+            
+            log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
+            return (result_tensor,)
+        finally:
+            # Cleanup (async operation) - always execute
+            run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)

upscale/payload_parsers.py

@@ -0,0 +1,64 @@
+import io
+import json
+
+from PIL import Image
+
+
+def _parse_tiles_from_form(data):
+    """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:
+                pass
+        if 'global_idx' in meta:
+            try:
+                tile_info['global_idx'] = int(meta['global_idx'])
+            except Exception:
+                pass
+
+        tiles.append(tile_info)
+
+    return tiles

upscale/result_collector.py

@@ -0,0 +1,198 @@
+import asyncio, time
+import comfy.model_management
+import server
+from ..utils.constants import DYNAMIC_MODE_MAX_POLL_TIMEOUT, HEARTBEAT_INTERVAL
+from ..utils.logging import debug_log, log
+from ..utils.config import get_worker_timeout_seconds
+from .job_store import ensure_tile_jobs_initialized, _mark_task_completed
+from .job_timeout import _check_and_requeue_timed_out_workers
+from .job_models import BaseJobState, ImageJobState, TileJobState
+
+
+class ResultCollectorMixin:
+    """
+    Mixin for master-side result collection in USDU distributed jobs.
+
+    Expected co-mixins/attributes:
+    - JobStateMixin methods for queue/task access.
+    - `self._check_and_requeue_timed_out_workers(...)` coroutine.
+    - `self._async_yield(...)` optional helper from WorkerCommsMixin.
+    """
+
+    def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> list[str]:
+        """Log timeout elapsed seconds for each tracked worker and return worker ids."""
+        if not isinstance(job_data, BaseJobState):
+            return []
+
+        worker_status = dict(job_data.worker_status)
+        for worker_id, last_seen in worker_status.items():
+            elapsed = max(0.0, current_time - float(last_seen))
+            log(
+                "UltimateSDUpscale Master - Heartbeat timeout: "
+                f"job={multi_job_id}, worker={worker_id}, elapsed={elapsed:.1f}s"
+            )
+        return list(worker_status.keys())
+
+    async def _async_collect_results(self, multi_job_id, num_workers, mode='static', 
+                                   remaining_to_collect=None, batch_size=None):
+        """Unified async helper to collect results from workers (tiles or images)."""
+        # Get the already initialized queue
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
+                raise RuntimeError(f"Job queue not initialized for {multi_job_id}")
+            job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+            if mode == 'dynamic':
+                if not isinstance(job_data, ImageJobState):
+                    raise RuntimeError(
+                        f"Mode mismatch: expected dynamic, got {getattr(job_data, 'mode', 'unknown')}"
+                    )
+                q = job_data.queue
+                completed_images = job_data.completed_images
+                expected_count = remaining_to_collect or batch_size
+            elif mode == 'static':
+                if not isinstance(job_data, TileJobState):
+                    raise RuntimeError(
+                        f"Mode mismatch: expected static, got {getattr(job_data, 'mode', 'unknown')}"
+                    )
+                q = job_data.queue
+                expected_count = len(job_data.completed_tasks) + job_data.pending_tasks.qsize()
+            else:
+                raise RuntimeError(f"Unsupported mode: {mode}")
+        
+        item_type = "images" if mode == 'dynamic' else "tiles"
+        debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers")
+        
+        collected_results = {}
+        workers_done = set()
+        # Unify collector/upscaler wait behavior with the UI worker timeout
+        timeout = float(get_worker_timeout_seconds())
+        last_heartbeat_check = time.time()
+        wait_started_at = time.time()
+        collected_count = 0
+        
+        while len(workers_done) < num_workers:
+            # Check for user interruption
+            if comfy.model_management.processing_interrupted():
+                log("Processing interrupted by user")
+                raise comfy.model_management.InterruptProcessingException()
+                
+            # For dynamic mode with remaining_to_collect, check if we've collected enough
+            if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect:
+                break
+
+            job_data_snapshot = None
+            async with prompt_server.distributed_tile_jobs_lock:
+                current_job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+                if isinstance(current_job_data, BaseJobState):
+                    job_data_snapshot = current_job_data
+
+            try:
+                # Shorter poll for dynamic mode, but never exceed the configured timeout
+                wait_timeout = (min(DYNAMIC_MODE_MAX_POLL_TIMEOUT, timeout) if mode == 'dynamic' else timeout)
+                result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
+                worker_id = result['worker_id']
+                is_last = result.get('is_last', False)
+                
+                if mode == 'static':
+                    # Handle tiles
+                    tiles = result.get('tiles', [])
+                    debug_log(
+                        f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker "
+                        f"'{worker_id}' (is_last={is_last})"
+                    )
+
+                    for tile_data in tiles:
+                        if 'batch_idx' not in tile_data:
+                            log("UltimateSDUpscale Master - Missing batch_idx in tile data, skipping")
+                            continue
+
+                        tile_idx = tile_data['tile_idx']
+                        key = tile_data.get('global_idx', tile_idx)
+                        entry = {
+                            'tile_idx': tile_idx,
+                            'x': tile_data['x'],
+                            'y': tile_data['y'],
+                            'extracted_width': tile_data['extracted_width'],
+                            'extracted_height': tile_data['extracted_height'],
+                            'padding': tile_data['padding'],
+                            'worker_id': worker_id,
+                            'batch_idx': tile_data.get('batch_idx', 0),
+                            'global_idx': tile_data.get('global_idx', tile_idx),
+                        }
+                        if 'image' in tile_data:
+                            entry['image'] = tile_data['image']
+                        elif 'tensor' in tile_data:
+                            entry['tensor'] = tile_data['tensor']
+                        collected_results[key] = entry
+                
+                elif mode == 'dynamic':
+                    # Handle full images
+                    if 'image_idx' in result and 'image' in result:
+                        image_idx = result['image_idx']
+                        image_pil = result['image']
+                        completed_images[image_idx] = image_pil
+                        collected_results[image_idx] = image_pil
+                        collected_count += 1
+                        debug_log(f"UltimateSDUpscale Master - Received image {image_idx} from worker {worker_id}")
+                
+                if is_last:
+                    workers_done.add(worker_id)
+                    debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
+                    
+            except asyncio.TimeoutError:
+                current_time = time.time()
+                waiting_workers = self._log_worker_timeout_status(job_data_snapshot, current_time, multi_job_id)
+                if mode == 'dynamic':
+                    # Check for worker timeouts periodically
+                    if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
+                        # Use the class method to check and requeue
+                        requeued = await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size)
+                        if requeued > 0:
+                            log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers")
+                        last_heartbeat_check = current_time
+                    
+                    # Check if we've been waiting too long overall
+                    if current_time - wait_started_at > timeout:
+                        elapsed = current_time - wait_started_at
+                        log(
+                            "UltimateSDUpscale Master - Heartbeat timeout while waiting for images; "
+                            f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
+                        )
+                        break
+                else:
+                    elapsed = current_time - wait_started_at
+                    log(
+                        f"UltimateSDUpscale Master - Heartbeat timeout waiting for {item_type}; "
+                        f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
+                    )
+                    break
+        
+        debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers")
+        
+        # Clean up job queue
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                del prompt_server.distributed_pending_tile_jobs[multi_job_id]
+        
+        return collected_results if mode == 'static' else completed_images
+
+    async def _async_collect_worker_tiles(self, multi_job_id, num_workers):
+        """Async helper to collect tiles from workers."""
+        return await self._async_collect_results(multi_job_id, num_workers, mode='static')
+
+    async def _mark_image_completed(self, multi_job_id, image_idx, image_pil):
+        """Mark an image as completed in the job data."""
+        # Mark the image as completed with the image data
+        await _mark_task_completed(multi_job_id, image_idx, {'image': image_pil})
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+            if isinstance(job_data, ImageJobState):
+                job_data.completed_images[image_idx] = image_pil
+
+    async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count):
+        """Collect remaining processed images from workers."""
+        return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic', 
+                                               remaining_to_collect=remaining_to_collect, 
+                                               batch_size=batch_size)

upscale/tile_ops.py

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

upscale/worker_comms.py

@@ -0,0 +1,262 @@
+import asyncio, io, json, time
+import aiohttp
+from PIL import Image
+from ..utils.logging import debug_log, log
+from ..utils.network import get_client_session
+from ..utils.constants import TILE_SEND_TIMEOUT
+from ..utils.usdu_managment import MAX_PAYLOAD_SIZE, _send_heartbeat_to_master
+from ..utils.image import tensor_to_pil
+
+
+class WorkerCommsMixin:
+    async def _send_heartbeat_to_master(self, multi_job_id, master_url, worker_id):
+        """Proxy heartbeat helper used by worker processing mixins."""
+        await _send_heartbeat_to_master(multi_job_id, master_url, worker_id)
+
+    async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, 
+                                       padding, worker_id, is_final_flush=False):
+        """Send all processed tiles to master, chunked if large."""
+        if not processed_tiles:
+            if is_final_flush:
+                await self._send_tiles_completion_signal(multi_job_id, master_url, worker_id)
+            return  # Early exit if empty
+
+        total_tiles = len(processed_tiles)
+        debug_log(f"Worker[{worker_id[:8]}] - Preparing to send {total_tiles} tiles (size-aware chunks)")
+
+        # Prepare encoded images and sizes to enable size-aware chunking
+        encoded = []
+        for idx, tile_data in enumerate(processed_tiles):
+            img = tensor_to_pil(tile_data['tile'], 0)
+            bio = io.BytesIO()
+            # Keep compression low to balance speed and size; adjust if needed
+            img.save(bio, format='PNG', compress_level=0)
+            raw = bio.getvalue()
+            encoded.append({
+                'bytes': raw,
+                'meta': {
+                    'tile_idx': tile_data['tile_idx'],
+                    'x': tile_data['x'],
+                    'y': tile_data['y'],
+                    'extracted_width': tile_data['extracted_width'],
+                    'extracted_height': tile_data['extracted_height'],
+                    **({'batch_idx': tile_data['batch_idx']} if 'batch_idx' in tile_data else {}),
+                    **({'global_idx': tile_data['global_idx']} if 'global_idx' in tile_data else {}),
+                }
+            })
+
+        # Size-aware chunking
+        max_bytes = int(MAX_PAYLOAD_SIZE) - (1024 * 1024)  # 1MB headroom
+        i = 0
+        chunk_index = 0
+        while i < total_tiles:
+            data = aiohttp.FormData()
+            data.add_field('multi_job_id', multi_job_id)
+            data.add_field('worker_id', str(worker_id))
+            data.add_field('padding', str(padding))
+
+            metadata = []
+            used = 0
+            j = i
+            while j < total_tiles:
+                img_bytes = encoded[j]['bytes']
+                meta = encoded[j]['meta']
+                # Rough overhead for fields + JSON
+                overhead = 1024
+                if used + len(img_bytes) + overhead > max_bytes and j > i:
+                    break
+                # Accept this tile in this chunk
+                metadata.append(meta)
+                data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png')
+                used += len(img_bytes) + overhead
+                j += 1
+
+            # Ensure at least one tile per chunk
+            if j == i:
+                # Single oversized tile, send anyway
+                meta = encoded[j]['meta']
+                metadata.append(meta)
+                data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png')
+                j += 1
+
+            chunk_size = j - i
+            is_chunk_last = (j >= total_tiles)
+            data.add_field('is_last', str(bool(is_final_flush and is_chunk_last)))
+            data.add_field('batch_size', str(chunk_size))
+            data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')
+
+            # Retry logic with exponential backoff
+            max_retries = 5
+            retry_delay = 0.5
+            for attempt in range(max_retries):
+                try:
+                    session = await get_client_session()
+                    url = f"{master_url}/distributed/submit_tiles"
+                    async with session.post(url, data=data) as response:
+                        response.raise_for_status()
+                        break
+                except Exception as e:
+                    if attempt < max_retries - 1:
+                        await asyncio.sleep(retry_delay)
+                        retry_delay = min(retry_delay * 2, 5.0)
+                    else:
+                        log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}")
+                        raise
+
+            debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)")
+            chunk_index += 1
+            i = j
+
+    async def _send_tiles_completion_signal(self, multi_job_id, master_url, worker_id):
+        """Send completion signal to master in static mode when no tiles are left."""
+        data = aiohttp.FormData()
+        data.add_field('multi_job_id', multi_job_id)
+        data.add_field('worker_id', str(worker_id))
+        data.add_field('is_last', 'true')
+        data.add_field('batch_size', '0')
+
+        session = await get_client_session()
+        url = f"{master_url}/distributed/submit_tiles"
+        async with session.post(url, data=data) as response:
+            response.raise_for_status()
+            debug_log(f"Worker {worker_id} sent static completion signal")
+
+    async def _request_work_item_from_master(
+        self,
+        multi_job_id,
+        master_url,
+        worker_id,
+        endpoint="/distributed/request_image",
+    ):
+        """Request one work item from master with retry/backoff and total timeout."""
+        max_retries = 10
+        retry_delay = 0.5
+        start_time = time.monotonic()
+        url = f"{master_url}{endpoint}"
+
+        for attempt in range(max_retries):
+            if time.monotonic() - start_time > 30:
+                log(f"Total request timeout after 30s for worker {worker_id}")
+                return None
+
+            try:
+                session = await get_client_session()
+                async with session.post(url, json={
+                    'worker_id': str(worker_id),
+                    'multi_job_id': multi_job_id
+                }) as response:
+                    if response.status == 200:
+                        return await response.json()
+                    if response.status == 404:
+                        text = await response.text()
+                        debug_log(f"Job not found (404), will retry: {text}")
+                        await asyncio.sleep(1.0)
+                    else:
+                        text = await response.text()
+                        debug_log(
+                            f"Request work item failed ({response.status}) for worker {worker_id}: {text}"
+                        )
+
+            except Exception as exc:
+                if attempt < max_retries - 1:
+                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {exc}")
+                    await asyncio.sleep(retry_delay)
+                    retry_delay = min(retry_delay * 2, 5.0)
+                else:
+                    log(f"Failed to request work item after {max_retries} attempts: {exc}")
+                    raise
+
+        return None
+
+    async def _request_image_from_master(self, multi_job_id, master_url, worker_id):
+        """Request an image index to process from master in dynamic mode."""
+        data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
+        if not data:
+            return None, 0
+        image_idx = data.get('image_idx')
+        estimated_remaining = data.get('estimated_remaining', 0)
+        return image_idx, estimated_remaining
+
+    async def _request_tile_from_master(self, multi_job_id, master_url, worker_id):
+        """Request a tile index to process from master in static mode (reusing dynamic infrastructure)."""
+        data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
+        if not data:
+            return None, 0, False
+        tile_idx = data.get('tile_idx')
+        estimated_remaining = data.get('estimated_remaining', 0)
+        batched_static = data.get('batched_static', False)
+        return tile_idx, estimated_remaining, batched_static
+
+    async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, 
+                                        master_url, worker_id, is_last):
+        """Send a processed full image back to master in dynamic mode."""
+        # Serialize image to PNG
+        byte_io = io.BytesIO()
+        image_pil.save(byte_io, format='PNG', compress_level=0)
+        byte_io.seek(0)
+        
+        # Prepare form data
+        data = aiohttp.FormData()
+        data.add_field('multi_job_id', multi_job_id)
+        data.add_field('worker_id', str(worker_id))
+        data.add_field('image_idx', str(image_idx))
+        data.add_field('is_last', str(is_last))
+        data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png', 
+                      content_type='image/png')
+        
+        # Retry logic
+        max_retries = 5
+        retry_delay = 0.5
+        
+        for attempt in range(max_retries):
+            try:
+                session = await get_client_session()
+                url = f"{master_url}/distributed/submit_image"
+                
+                async with session.post(url, data=data) as response:
+                    response.raise_for_status()
+                    debug_log(f"Successfully sent image {image_idx} to master")
+                    return
+                    
+            except Exception as e:
+                if attempt < max_retries - 1:
+                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
+                    await asyncio.sleep(retry_delay)
+                    retry_delay *= 2
+                else:
+                    log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}")
+                    raise
+
+    async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id):
+        """Send completion signal to master in dynamic mode."""
+        # Send a dummy request with is_last=True
+        data = aiohttp.FormData()
+        data.add_field('multi_job_id', multi_job_id)
+        data.add_field('worker_id', str(worker_id))
+        data.add_field('is_last', 'true')
+        # No image data - just completion signal
+        
+        session = await get_client_session()
+        url = f"{master_url}/distributed/submit_image"
+        
+        async with session.post(url, data=data) as response:
+            response.raise_for_status()
+            debug_log(f"Worker {worker_id} sent completion signal")
+
+    async def _check_job_status(self, multi_job_id, master_url):
+        """Check if job is ready on the master."""
+        try:
+            session = await get_client_session()
+            url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}"
+            async with session.get(url) as response:
+                if response.status == 200:
+                    data = await response.json()
+                    return data.get('ready', False)
+                return False
+        except Exception as e:
+            debug_log(f"Job status check failed: {e}")
+            return False
+
+    async def _async_yield(self):
+        """Simple async yield to allow event loop processing."""
+        await asyncio.sleep(0)

utils/async_helpers.py

@@ -3,6 +3,10 @@ Async helper utilities for ComfyUI-Distributed.
 """
 import asyncio
 import threading
+import time
+import uuid
+import execution
+import server
 from typing import Optional, Any, Coroutine
 from .network import get_server_loop
 
@@ -47,4 +51,99 @@ def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -
     
     if error:
         raise error
-    return result
+    return result
+
+
+prompt_server = server.PromptServer.instance
+
+
+def _summarize_node_errors(node_errors: dict) -> str:
+    if not isinstance(node_errors, dict) or not node_errors:
+        return ""
+
+    parts = []
+    for node_id, entry in node_errors.items():
+        if not isinstance(entry, dict):
+            continue
+        class_type = str(entry.get("class_type") or "UnknownNode")
+        for err in entry.get("errors", []):
+            if not isinstance(err, dict):
+                continue
+            message = str(err.get("message") or "validation error")
+            details = str(err.get("details") or "").strip()
+            parts.append(
+                f"{class_type}#{node_id}: {message}{f' ({details})' if details else ''}"
+            )
+            if len(parts) >= 5:
+                return " | ".join(parts)
+    return " | ".join(parts)
+
+
+class PromptValidationError(RuntimeError):
+    """Raised when a prompt fails ComfyUI validation with structured context."""
+
+    def __init__(self, error_payload, node_errors=None):
+        payload = error_payload if isinstance(error_payload, dict) else {
+            "type": "prompt_validation_failed",
+            "message": str(error_payload),
+            "details": "",
+            "extra_info": {},
+        }
+        self.validation_error = dict(payload)
+        self.node_errors = node_errors if isinstance(node_errors, dict) else {}
+
+        if self.node_errors:
+            details = str(self.validation_error.get("details") or "").strip()
+            if not details:
+                summary = _summarize_node_errors(self.node_errors)
+                if summary:
+                    self.validation_error["details"] = summary
+
+        merged = dict(self.validation_error)
+        if self.node_errors:
+            merged["node_errors"] = self.node_errors
+        super().__init__(f"Invalid prompt: {merged}")
+
+
+async def queue_prompt_payload(
+    prompt_obj,
+    workflow_meta=None,
+    client_id=None,
+    include_queue_metadata=False,
+):
+    """Validate and queue a prompt via ComfyUI's prompt queue."""
+    payload = {"prompt": prompt_obj}
+    payload = prompt_server.trigger_on_prompt(payload)
+    prompt = payload["prompt"]
+
+    prompt_id = str(uuid.uuid4())
+    valid = await execution.validate_prompt(prompt_id, prompt, None)
+    if not valid[0]:
+        error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation"
+        node_errors = valid[3] if len(valid) > 3 else {}
+        raise PromptValidationError(error_payload, node_errors)
+
+    extra_data = {"create_time": int(time.time() * 1000)}
+    if workflow_meta:
+        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
+    if client_id:
+        extra_data["client_id"] = client_id
+
+    sensitive = {}
+    for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
+        if key in extra_data:
+            sensitive[key] = extra_data.pop(key)
+
+    number = getattr(prompt_server, "number", 0)
+    prompt_server.number = number + 1
+    prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
+    prompt_server.prompt_queue.put(prompt_queue_item)
+
+    if include_queue_metadata:
+        return {
+            "prompt_id": prompt_id,
+            "number": number,
+            "node_errors": {},
+        }
+
+    return prompt_id

utils/audio_payload.py

@@ -0,0 +1,103 @@
+import base64
+import binascii
+import os
+
+import numpy as np
+import torch
+
+from .image import ensure_contiguous
+
+
+MAX_AUDIO_PAYLOAD_BYTES = int(
+    os.environ.get("COMFYUI_MAX_AUDIO_PAYLOAD_BYTES", str(256 * 1024 * 1024))
+)
+
+
+def encode_audio_payload(audio_payload):
+    """Serialize an AUDIO dict into JSON-safe canonical envelope payload."""
+    if not isinstance(audio_payload, dict):
+        return None
+
+    waveform = audio_payload.get("waveform")
+    if waveform is None or not isinstance(waveform, torch.Tensor) or waveform.numel() == 0:
+        return None
+
+    sample_rate = audio_payload.get("sample_rate", 44100)
+    try:
+        sample_rate = int(sample_rate)
+    except (TypeError, ValueError):
+        sample_rate = 44100
+
+    waveform_cpu = waveform.detach().to(device="cpu", dtype=torch.float32).contiguous()
+    data_bytes = waveform_cpu.numpy().tobytes()
+    if len(data_bytes) > MAX_AUDIO_PAYLOAD_BYTES:
+        raise ValueError(
+            f"Audio payload too large: {len(data_bytes)} bytes exceeds {MAX_AUDIO_PAYLOAD_BYTES}."
+        )
+
+    return {
+        "sample_rate": sample_rate,
+        "shape": [int(dim) for dim in waveform_cpu.shape],
+        "dtype": "float32",
+        "data": base64.b64encode(data_bytes).decode("ascii"),
+    }
+
+
+def decode_audio_payload(audio_payload):
+    """Decode canonical envelope audio payload into an AUDIO dict."""
+    if audio_payload is None:
+        return None
+    if not isinstance(audio_payload, dict):
+        raise ValueError("Field 'audio' must be an object when provided.")
+
+    encoded = audio_payload.get("data")
+    shape = audio_payload.get("shape")
+    sample_rate = audio_payload.get("sample_rate", 44100)
+    dtype = audio_payload.get("dtype", "float32")
+
+    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 [batch, channels, samples].")
+    if dtype != "float32":
+        raise ValueError("Field 'audio.dtype' must be 'float32'.")
+
+    try:
+        shape_tuple = tuple(int(dim) for dim in shape)
+    except (TypeError, ValueError) as exc:
+        raise ValueError("Field 'audio.shape' must contain integers.") from exc
+
+    if shape_tuple[0] <= 0 or shape_tuple[1] <= 0 or shape_tuple[2] < 0:
+        raise ValueError(
+            "Field 'audio.shape' must be [batch>0, channels>0, samples>=0]."
+        )
+
+    try:
+        sample_rate = int(sample_rate)
+    except (TypeError, ValueError) as exc:
+        raise ValueError("Field 'audio.sample_rate' must be an integer.") from exc
+    if sample_rate <= 0:
+        raise ValueError("Field 'audio.sample_rate' must be positive.")
+
+    try:
+        raw = base64.b64decode(encoded, validate=True)
+    except (binascii.Error, ValueError) as exc:
+        raise ValueError("Field 'audio.data' is not valid base64.") from exc
+
+    if len(raw) > MAX_AUDIO_PAYLOAD_BYTES:
+        raise ValueError(
+            f"Field 'audio.data' too large: {len(raw)} bytes exceeds {MAX_AUDIO_PAYLOAD_BYTES}."
+        )
+
+    expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
+    if len(raw) != expected_bytes:
+        raise ValueError(
+            f"Field 'audio.data' byte size mismatch: expected {expected_bytes}, got {len(raw)}."
+        )
+
+    array = np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple)
+    waveform = torch.from_numpy(array.copy())
+    return {
+        "waveform": ensure_contiguous(waveform),
+        "sample_rate": sample_rate,
+    }

utils/cloudflare.py

@@ -1,397 +0,0 @@
-"""
-Lightweight Cloudflare Tunnel manager.
-
-Responsibilities:
-- Locate or download the cloudflared binary for the current platform.
-- Start/stop a quick tunnel pointing to the local ComfyUI server.
-- Track status and persist minimal state in the shared config.
-"""
-import asyncio
-import os
-import platform
-import re
-import shutil
-import signal
-import stat
-import subprocess
-import threading
-import time
-from urllib import request, error as urlerror
-
-from .config import load_config, save_config
-from .logging import debug_log, log
-from .network import get_server_port
-from .process import is_process_alive, terminate_process
-
-# Regex to capture the generated public URL from cloudflared output
-PUBLIC_URL_PATTERN = re.compile(r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))", re.IGNORECASE)
-
-# Default timeout (seconds) to wait for a tunnel URL before giving up
-TUNNEL_START_TIMEOUT = 25
-
-
-def _normalize_host(value):
-    if not value or not isinstance(value, str):
-        return ""
-    value = value.strip()
-    value = re.sub(r"^https?://", "", value, flags=re.IGNORECASE)
-    return value.split("/")[0]
-
-
-class CloudflareTunnelManager:
-    def __init__(self):
-        self.process = None
-        self.pid = None
-        self.public_url = None
-        self.last_error = None
-        self.log_file = None
-        self.status = "stopped"
-        self.previous_master_host = None
-
-        self._lock = asyncio.Lock()
-        self._url_event = None
-        self._reader_thread = None
-        self._loop = None
-        self._recent_logs = []
-
-        self.binary_path = None
-        self._restore_state()
-
-    # --- Binary discovery & download ---
-    @property
-    def base_dir(self):
-        return os.path.dirname(os.path.dirname(__file__))
-
-    @property
-    def bin_dir(self):
-        return os.path.join(self.base_dir, "bin")
-
-    def _get_asset_name(self):
-        system = platform.system().lower()
-        machine = platform.machine().lower()
-
-        if system == "windows":
-            if "arm" in machine:
-                return "cloudflared-windows-arm64.exe"
-            return "cloudflared-windows-amd64.exe"
-        if system == "darwin":
-            if machine in ("arm64", "aarch64"):
-                return "cloudflared-darwin-arm64"
-            return "cloudflared-darwin-amd64"
-        if system == "linux":
-            if machine in ("arm64", "aarch64"):
-                return "cloudflared-linux-arm64"
-            return "cloudflared-linux-amd64"
-
-        raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
-
-    def _download_cloudflared(self):
-        asset = self._get_asset_name()
-        url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
-        os.makedirs(self.bin_dir, exist_ok=True)
-        target_path = os.path.join(self.bin_dir, "cloudflared.exe" if asset.endswith(".exe") else "cloudflared")
-
-        debug_log(f"Downloading cloudflared from {url}")
-        try:
-            with request.urlopen(url, timeout=30) as resp:
-                with open(target_path, "wb") as f:
-                    shutil.copyfileobj(resp, f)
-        except urlerror.URLError as exc:
-            raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
-
-        # Make executable
-        st = os.stat(target_path)
-        os.chmod(target_path, st.st_mode | stat.S_IEXEC)
-        debug_log(f"Downloaded cloudflared to {target_path}")
-        return target_path
-
-    def ensure_binary(self):
-        # Allow explicit override
-        env_path = os.environ.get("CLOUDFLARED_PATH")
-        if env_path and os.path.exists(env_path):
-            self.binary_path = env_path
-            return env_path
-
-        # Check cached path
-        cached = self.binary_path
-        if cached and os.path.exists(cached):
-            return cached
-
-        # Local bin directory
-        candidate = os.path.join(self.bin_dir, "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared")
-        if os.path.exists(candidate):
-            self.binary_path = candidate
-            return candidate
-
-        # System PATH
-        path_binary = shutil.which("cloudflared")
-        if path_binary:
-            self.binary_path = path_binary
-            return path_binary
-
-        # Download on demand
-        self.binary_path = self._download_cloudflared()
-        return self.binary_path
-
-    # --- State helpers ---
-    def _persist_state(self, status=None, public_url=None, pid=None, log_file=None, previous_host=None, master_host=None):
-        cfg = load_config()
-        tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
-
-        if status is not None:
-            tunnel_cfg["status"] = status
-        if public_url is not None:
-            tunnel_cfg["public_url"] = public_url
-        if pid is not None:
-            tunnel_cfg["pid"] = pid
-        if log_file is not None:
-            tunnel_cfg["log_file"] = log_file
-        if previous_host is not None:
-            tunnel_cfg["previous_master_host"] = previous_host
-        if master_host is not None:
-            cfg.setdefault("master", {})["host"] = master_host
-
-        cfg["tunnel"] = tunnel_cfg
-        save_config(cfg)
-
-    def _restore_state(self):
-        cfg = load_config()
-        tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
-
-        self.public_url = tunnel_cfg.get("public_url") or None
-        self.previous_master_host = tunnel_cfg.get("previous_master_host")
-        self.log_file = tunnel_cfg.get("log_file")
-        pid = tunnel_cfg.get("pid")
-
-        if pid and is_process_alive(pid):
-            self.pid = pid
-            self.status = tunnel_cfg.get("status", "running")
-            debug_log(f"Detected existing cloudflared process (pid={pid})")
-        else:
-            # Clear stale info
-            self._persist_state(status="stopped", public_url="", pid=None, log_file=None)
-            self.status = "stopped"
-            self.pid = None
-
-    # --- Process output handling ---
-    def _append_log(self, line):
-        if self.log_file:
-            try:
-                with open(self.log_file, "a", encoding="utf-8", errors="replace") as f:
-                    f.write(line + "\n")
-            except Exception as exc:  # pragma: no cover - best effort logging
-                debug_log(f"Failed to write tunnel log: {exc}")
-
-        self._recent_logs.append(line)
-        # Keep last ~200 lines to avoid unbounded growth
-        if len(self._recent_logs) > 200:
-            self._recent_logs = self._recent_logs[-200:]
-
-    def _reader(self):
-        assert self.process is not None
-        loop = self._loop
-        for raw_line in iter(self.process.stdout.readline, ""):
-            line = raw_line.strip()
-            if not line:
-                continue
-
-            self._append_log(line)
-            match = PUBLIC_URL_PATTERN.search(line)
-            if match and not self.public_url:
-                self.public_url = match.group(1).rstrip("/")
-                self.status = "running"
-                if self._url_event and loop:
-                    loop.call_soon_threadsafe(self._url_event.set)
-
-            # Capture obvious error strings
-            if "error" in line.lower() and not self.last_error:
-                self.last_error = line
-
-        # Process exited
-        if self.status == "starting" and not self.public_url:
-            self.status = "error"
-            if not self.last_error:
-                self.last_error = "Cloudflare tunnel exited before becoming ready"
-            if self._url_event and loop:
-                loop.call_soon_threadsafe(self._url_event.set)
-        elif self.status == "running":
-            # Treat unexpected exit as stopped so UI can retry
-            self.status = "stopped"
-        elif self.status == "stopped":
-            # Normal stop path, nothing to do
-            pass
-
-    # --- Public API ---
-    async def start_tunnel(self):
-        async with self._lock:
-            if self.process and self.process.poll() is None:
-                return {
-                    "status": self.status,
-                    "public_url": self.public_url,
-                    "pid": self.process.pid,
-                    "log_file": self.log_file,
-                }
-
-            # If we had a stale pid, ensure it's not running
-            if self.pid and is_process_alive(self.pid):
-                debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one")
-                await self.stop_tunnel()
-
-            binary = await asyncio.to_thread(self.ensure_binary)
-            port = get_server_port()
-            self.status = "starting"
-            self.last_error = None
-            self.public_url = None
-            self._recent_logs = []
-
-            # Remember current master host so we can restore it if the tunnel stops
-            config = load_config()
-            master_host = (config.get("master") or {}).get("host") or ""
-            tunnel_cfg = config.get("tunnel") or {}
-            if tunnel_cfg.get("previous_master_host"):
-                self.previous_master_host = tunnel_cfg.get("previous_master_host")
-            else:
-                self.previous_master_host = master_host
-
-            os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True)
-            timestamp = time.strftime("%Y%m%d-%H%M%S")
-            self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log")
-
-            cmd = [
-                binary,
-                "tunnel",
-                "--no-autoupdate",
-                "--url",
-                f"http://127.0.0.1:{port}",
-            ]
-
-            debug_log(f"Starting cloudflared: {' '.join(cmd)}")
-            try:
-                self.process = subprocess.Popen(
-                    cmd,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.STDOUT,
-                    text=True,
-                    bufsize=1,
-                )
-            except FileNotFoundError:
-                self.status = "error"
-                raise RuntimeError("cloudflared binary not found")
-            except Exception as exc:
-                self.status = "error"
-                raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc
-
-            self.pid = self.process.pid
-            self._persist_state(
-                status="starting",
-                pid=self.pid,
-                log_file=self.log_file,
-                previous_host=self.previous_master_host
-            )
-
-            # Kick off background reader
-            self._loop = asyncio.get_running_loop()
-            self._url_event = asyncio.Event()
-            self._reader_thread = threading.Thread(target=self._reader, daemon=True)
-            self._reader_thread.start()
-
-            try:
-                await asyncio.wait_for(self._url_event.wait(), timeout=TUNNEL_START_TIMEOUT)
-            except asyncio.TimeoutError:
-                self.last_error = "Timed out waiting for Cloudflare to assign a URL"
-                await self.stop_tunnel()
-                raise RuntimeError(self.last_error)
-
-            if self.status != "running" or not self.public_url:
-                error_msg = self.last_error or "Cloudflare tunnel failed to start"
-                await self.stop_tunnel()
-                raise RuntimeError(error_msg)
-
-            debug_log(f"Cloudflare tunnel ready at {self.public_url}")
-            master_host = _normalize_host(self.public_url)
-            self._persist_state(
-                status="running",
-                public_url=self.public_url,
-                pid=self.pid,
-                log_file=self.log_file,
-                previous_host=self.previous_master_host or "",
-                master_host=master_host
-            )
-            return {
-                "status": self.status,
-                "public_url": self.public_url,
-                "pid": self.pid,
-                "log_file": self.log_file,
-            }
-
-    async def stop_tunnel(self):
-        async with self._lock:
-            pid = self.process.pid if self.process else self.pid
-            if not pid:
-                self._persist_state(status="stopped", public_url=None, pid=None, log_file=None)
-                self.status = "stopped"
-                return {"status": "stopped"}
-
-            debug_log(f"Stopping cloudflared (pid={pid})")
-            if self.process:
-                terminate_process(self.process, timeout=5)
-            else:
-                try:
-                    os.kill(pid, signal.SIGTERM)
-                    time.sleep(0.5)
-                except Exception as exc:  # pragma: no cover - best effort
-                    debug_log(f"Error stopping cloudflared pid {pid}: {exc}")
-
-            config = load_config()
-            tunnel_cfg = config.get("tunnel", {}) if isinstance(config.get("tunnel", {}), dict) else {}
-            self.previous_master_host = tunnel_cfg.get("previous_master_host", self.previous_master_host)
-            active_url = tunnel_cfg.get("public_url")
-            current_master_host = (config.get("master") or {}).get("host")
-            restore_host = None
-            if active_url:
-                active_host = _normalize_host(active_url)
-                current_host = _normalize_host(current_master_host)
-                if current_host == active_host:
-                    restore_host = self.previous_master_host or ""
-
-            self.status = "stopped"
-            self.public_url = None
-            self.pid = None
-            self.process = None
-            self.last_error = None
-            self._url_event = None
-            if self._reader_thread and self._reader_thread.is_alive():
-                self._reader_thread.join(timeout=1)
-            self._reader_thread = None
-            self._persist_state(
-                status="stopped",
-                public_url="",
-                pid=None,
-                log_file=self.log_file,
-                previous_host=self.previous_master_host,
-                master_host=restore_host
-            )
-            return {"status": "stopped"}
-
-    def get_status(self):
-        alive = False
-        pid = self.process.pid if self.process else self.pid
-        if pid:
-            alive = is_process_alive(pid)
-            if not alive and self.status == "running":
-                self.status = "stopped"
-
-        return {
-            "status": self.status,
-            "public_url": self.public_url,
-            "pid": pid,
-            "log_file": self.log_file,
-            "last_error": self.last_error,
-            "binary_path": self.binary_path or shutil.which("cloudflared"),
-            "recent_logs": self._recent_logs[-20:],
-            "previous_master_host": self.previous_master_host,
-        }
-
-
-# Singleton tunnel manager
-cloudflare_tunnel_manager = CloudflareTunnelManager()

utils/cloudflare/__init__.py

@@ -0,0 +1,5 @@
+from .tunnel import CloudflareTunnelManager
+
+cloudflare_tunnel_manager = CloudflareTunnelManager()
+
+__all__ = ["CloudflareTunnelManager", "cloudflare_tunnel_manager"]

utils/cloudflare/binary.py

@@ -0,0 +1,83 @@
+"""Cloudflared binary discovery and download helpers."""
+
+import os
+import platform
+import shutil
+import stat
+from urllib import error as urlerror
+from urllib import request
+
+from ..logging import debug_log
+
+
+def _get_project_root():
+    return os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+
+
+def _get_cloudflared_dir():
+    return os.path.join(_get_project_root(), "bin")
+
+
+def _get_platform_binary_name():
+    system = platform.system().lower()
+    machine = platform.machine().lower()
+
+    if system == "windows":
+        if "arm" in machine:
+            return "cloudflared-windows-arm64.exe"
+        return "cloudflared-windows-amd64.exe"
+    if system == "darwin":
+        if machine in ("arm64", "aarch64"):
+            return "cloudflared-darwin-arm64"
+        return "cloudflared-darwin-amd64"
+    if system == "linux":
+        if machine in ("arm64", "aarch64"):
+            return "cloudflared-linux-arm64"
+        return "cloudflared-linux-amd64"
+
+    raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
+
+
+def _get_binary_path(bin_dir=None):
+    bin_dir = bin_dir or _get_cloudflared_dir()
+    binary_name = "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared"
+    return os.path.join(bin_dir, binary_name)
+
+
+def _download_cloudflared():
+    asset = _get_platform_binary_name()
+    url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
+
+    bin_dir = _get_cloudflared_dir()
+    os.makedirs(bin_dir, exist_ok=True)
+    target_path = _get_binary_path(bin_dir)
+
+    debug_log(f"Downloading cloudflared from {url}")
+    try:
+        with request.urlopen(url, timeout=30) as resp:
+            with open(target_path, "wb") as f:
+                shutil.copyfileobj(resp, f)
+    except urlerror.URLError as exc:
+        raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
+
+    st = os.stat(target_path)
+    os.chmod(target_path, st.st_mode | stat.S_IEXEC)
+    debug_log(f"Downloaded cloudflared to {target_path}")
+    return target_path
+
+
+def ensure_binary() -> str:
+    """Return a usable cloudflared binary path, downloading if necessary."""
+    env_path = os.environ.get("CLOUDFLARED_PATH")
+    if env_path and os.path.exists(env_path):
+        return env_path
+
+    local_candidate = _get_binary_path()
+    if os.path.exists(local_candidate):
+        return local_candidate
+
+    path_binary = shutil.which("cloudflared")
+    if path_binary:
+        return path_binary
+
+    return _download_cloudflared()

utils/cloudflare/process_reader.py

@@ -0,0 +1,97 @@
+"""Background cloudflared process output reader."""
+
+import asyncio
+import re
+import threading
+
+from ..constants import CLOUDFLARE_LOG_BUFFER_SIZE
+from ..logging import debug_log
+
+PUBLIC_URL_PATTERN = re.compile(
+    r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))",
+    re.IGNORECASE,
+)
+class ProcessReader:
+    def __init__(self, log_file=None):
+        self._process = None
+        self._thread = None
+        self._loop = None
+        self._url_event = None
+        self._public_url = None
+        self._last_error = None
+        self._recent_logs = []
+        self._log_file = log_file
+
+    def set_log_file(self, log_file):
+        self._log_file = log_file
+
+    def _append_log(self, line):
+        if self._log_file:
+            try:
+                with open(self._log_file, "a", encoding="utf-8", errors="replace") as f:
+                    f.write(line + "\n")
+            except Exception as exc:  # pragma: no cover
+                debug_log(f"Failed to write tunnel log: {exc}")
+
+        self._recent_logs.append(line)
+        if len(self._recent_logs) > CLOUDFLARE_LOG_BUFFER_SIZE:
+            self._recent_logs = self._recent_logs[-CLOUDFLARE_LOG_BUFFER_SIZE:]
+
+    def _reader(self):
+        process = self._process
+        if process is None:
+            return
+
+        loop = self._loop
+        for raw_line in iter(process.stdout.readline, ""):
+            line = raw_line.strip()
+            if not line:
+                continue
+
+            self._append_log(line)
+            match = PUBLIC_URL_PATTERN.search(line)
+            if match and not self._public_url:
+                self._public_url = match.group(1).rstrip("/")
+                if self._url_event and loop:
+                    loop.call_soon_threadsafe(self._url_event.set)
+
+            if "error" in line.lower() and not self._last_error:
+                self._last_error = line
+
+        if self._url_event and loop:
+            if not self._last_error and not self._public_url:
+                self._last_error = "Cloudflare tunnel exited before becoming ready"
+            loop.call_soon_threadsafe(self._url_event.set)
+
+    def start(self, process, loop):
+        self._process = process
+        self._loop = loop
+        self._url_event = asyncio.Event()
+        self._public_url = None
+        self._last_error = None
+        self._recent_logs = []
+        self._thread = threading.Thread(target=self._reader, daemon=True)
+        self._thread.start()
+
+    async def wait_for_url(self, timeout):
+        if not self._url_event:
+            return None
+        await asyncio.wait_for(self._url_event.wait(), timeout=timeout)
+        return self._public_url
+
+    def stop(self):
+        if self._thread and self._thread.is_alive():
+            self._thread.join(timeout=1)
+        self._thread = None
+        self._process = None
+        self._loop = None
+        self._url_event = None
+
+    def get_url(self):
+        return self._public_url
+
+    def get_last_error(self):
+        return self._last_error
+
+    def get_recent_logs(self):
+        return list(self._recent_logs)

utils/cloudflare/state.py

@@ -0,0 +1,81 @@
+"""Cloudflare tunnel state persistence helpers."""
+
+from ..config import load_config, save_config
+from ..network import normalize_host
+
+
+def _get_tunnel_config(cfg):
+    tunnel_cfg = cfg.get("tunnel", {})
+    if isinstance(tunnel_cfg, dict):
+        return tunnel_cfg
+    return {}
+
+
+def load_tunnel_state():
+    cfg = load_config()
+    tunnel_cfg = _get_tunnel_config(cfg)
+    master_cfg = cfg.get("master", {}) if isinstance(cfg.get("master", {}), dict) else {}
+    return {
+        "status": tunnel_cfg.get("status", "stopped"),
+        "public_url": tunnel_cfg.get("public_url") or None,
+        "pid": tunnel_cfg.get("pid"),
+        "log_file": tunnel_cfg.get("log_file"),
+        "previous_master_host": tunnel_cfg.get("previous_master_host"),
+        "master_host": master_cfg.get("host"),
+    }
+
+
+def persist_tunnel_state(
+    status=None,
+    public_url=None,
+    pid=None,
+    log_file=None,
+    previous_host=None,
+    master_host=None,
+):
+    cfg = load_config()
+    tunnel_cfg = _get_tunnel_config(cfg)
+
+    if status is not None:
+        tunnel_cfg["status"] = status
+    if public_url is not None:
+        tunnel_cfg["public_url"] = public_url
+    if pid is not None:
+        tunnel_cfg["pid"] = pid
+    if log_file is not None:
+        tunnel_cfg["log_file"] = log_file
+    if previous_host is not None:
+        tunnel_cfg["previous_master_host"] = previous_host
+    if master_host is not None:
+        cfg.setdefault("master", {})["host"] = master_host
+
+    cfg["tunnel"] = tunnel_cfg
+    save_config(cfg)
+
+
+def clear_tunnel_state(log_file=None, previous_host=None, master_host=None):
+    persist_tunnel_state(
+        status="stopped",
+        public_url="",
+        pid=None,
+        log_file=log_file,
+        previous_host=previous_host,
+        master_host=master_host,
+    )
+
+
+def resolve_restore_master_host(previous_master_host):
+    """Determine whether master host should be restored after tunnel stop."""
+    cfg = load_config()
+    tunnel_cfg = _get_tunnel_config(cfg)
+    active_url = tunnel_cfg.get("public_url")
+    current_master_host = (cfg.get("master") or {}).get("host")
+
+    if not active_url:
+        return None
+
+    active_host = normalize_host(active_url)
+    current_host = normalize_host(current_master_host)
+    if current_host == active_host:
+        return previous_master_host or ""
+    return None

utils/cloudflare/tunnel.py

@@ -0,0 +1,207 @@
+"""Cloudflare tunnel lifecycle manager."""
+
+import asyncio
+import os
+import shutil
+import signal
+import subprocess
+import time
+
+from ..constants import TUNNEL_START_TIMEOUT
+from ..logging import debug_log
+from ..network import get_server_port, normalize_host
+from ..process import is_process_alive, terminate_process
+from .binary import ensure_binary
+from .process_reader import ProcessReader
+from .state import clear_tunnel_state, load_tunnel_state, persist_tunnel_state, resolve_restore_master_host
+
+
+class CloudflareTunnelManager:
+    def __init__(self):
+        self.process = None
+        self.pid = None
+        self.public_url = None
+        self.last_error = None
+        self.log_file = None
+        self.status = "stopped"
+        self.previous_master_host = None
+
+        self._lock = asyncio.Lock()
+        self._reader = ProcessReader()
+        self.binary_path = None
+
+        self._restore_state()
+
+    @property
+    def base_dir(self):
+        return os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+
+    def _restore_state(self):
+        state = load_tunnel_state()
+
+        self.public_url = state.get("public_url") or None
+        self.previous_master_host = state.get("previous_master_host")
+        self.log_file = state.get("log_file")
+        pid = state.get("pid")
+
+        if pid and is_process_alive(pid):
+            self.pid = pid
+            self.status = state.get("status") or "running"
+            debug_log(f"Detected existing cloudflared process (pid={pid})")
+        else:
+            clear_tunnel_state(log_file=self.log_file, previous_host=self.previous_master_host)
+            self.status = "stopped"
+            self.pid = None
+
+    async def start_tunnel(self):
+        async with self._lock:
+            if self.process and self.process.poll() is None:
+                return {
+                    "status": self.status,
+                    "public_url": self.public_url,
+                    "pid": self.process.pid,
+                    "log_file": self.log_file,
+                }
+
+            if self.pid and is_process_alive(self.pid):
+                debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one")
+                await self.stop_tunnel()
+
+            binary = await asyncio.to_thread(ensure_binary)
+            self.binary_path = binary
+            port = get_server_port()
+            self.status = "starting"
+            self.last_error = None
+            self.public_url = None
+
+            state = load_tunnel_state()
+            master_host = state.get("master_host") or ""
+            if state.get("previous_master_host"):
+                self.previous_master_host = state.get("previous_master_host")
+            else:
+                self.previous_master_host = master_host
+
+            os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True)
+            timestamp = time.strftime("%Y%m%d-%H%M%S")
+            self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log")
+
+            cmd = [
+                binary,
+                "tunnel",
+                "--no-autoupdate",
+                "--url",
+                f"http://127.0.0.1:{port}",
+            ]
+
+            debug_log(f"Starting cloudflared: {' '.join(cmd)}")
+            try:
+                self.process = subprocess.Popen(
+                    cmd,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.STDOUT,
+                    text=True,
+                    bufsize=1,
+                )
+            except FileNotFoundError:
+                self.status = "error"
+                raise RuntimeError("cloudflared binary not found")
+            except Exception as exc:
+                self.status = "error"
+                raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc
+
+            self.pid = self.process.pid
+            persist_tunnel_state(
+                status="starting",
+                pid=self.pid,
+                log_file=self.log_file,
+                previous_host=self.previous_master_host,
+            )
+
+            loop = asyncio.get_running_loop()
+            self._reader.set_log_file(self.log_file)
+            self._reader.start(self.process, loop)
+
+            try:
+                await self._reader.wait_for_url(timeout=TUNNEL_START_TIMEOUT)
+            except asyncio.TimeoutError:
+                self.last_error = "Timed out waiting for Cloudflare to assign a URL"
+                await self.stop_tunnel()
+                raise RuntimeError(self.last_error)
+
+            public_url = self._reader.get_url()
+            if not public_url:
+                self.last_error = self._reader.get_last_error() or "Cloudflare tunnel failed to start"
+                await self.stop_tunnel()
+                raise RuntimeError(self.last_error)
+
+            self.public_url = public_url
+            self.status = "running"
+            debug_log(f"Cloudflare tunnel ready at {self.public_url}")
+
+            persist_tunnel_state(
+                status="running",
+                public_url=self.public_url,
+                pid=self.pid,
+                log_file=self.log_file,
+                previous_host=self.previous_master_host or "",
+                master_host=normalize_host(self.public_url),
+            )
+            return {
+                "status": self.status,
+                "public_url": self.public_url,
+                "pid": self.pid,
+                "log_file": self.log_file,
+            }
+
+    async def stop_tunnel(self):
+        async with self._lock:
+            pid = self.process.pid if self.process else self.pid
+            if not pid:
+                clear_tunnel_state(log_file=self.log_file, previous_host=self.previous_master_host)
+                self.status = "stopped"
+                return {"status": "stopped"}
+
+            debug_log(f"Stopping cloudflared (pid={pid})")
+            if self.process:
+                terminate_process(self.process, timeout=5)
+            else:
+                try:
+                    os.kill(pid, signal.SIGTERM)
+                    time.sleep(0.5)
+                except Exception as exc:  # pragma: no cover
+                    debug_log(f"Error stopping cloudflared pid {pid}: {exc}")
+
+            restore_host = resolve_restore_master_host(self.previous_master_host)
+
+            self.status = "stopped"
+            self.public_url = None
+            self.pid = None
+            self.process = None
+            self.last_error = None
+            self._reader.stop()
+
+            clear_tunnel_state(
+                log_file=self.log_file,
+                previous_host=self.previous_master_host,
+                master_host=restore_host,
+            )
+            return {"status": "stopped"}
+
+    def get_status(self):
+        alive = False
+        pid = self.process.pid if self.process else self.pid
+        if pid:
+            alive = is_process_alive(pid)
+            if not alive and self.status == "running":
+                self.status = "stopped"
+
+        return {
+            "status": self.status,
+            "public_url": self.public_url,
+            "pid": pid,
+            "log_file": self.log_file,
+            "last_error": self.last_error or self._reader.get_last_error(),
+            "binary_path": self.binary_path or shutil.which("cloudflared"),
+            "recent_logs": self._reader.get_recent_logs()[-20:],
+            "previous_master_host": self.previous_master_host,
+        }

utils/config.py

@@ -1,14 +1,23 @@
 """
 Configuration management for ComfyUI-Distributed.
 """
+import asyncio
 import os
 import json
+from contextlib import asynccontextmanager
 from .logging import log
 
 # Import defaults for timeout fallbacks
 from .constants import HEARTBEAT_TIMEOUT
 
 CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json")
+_config_cache = None
+_config_mtime = 0.0
+_config_lock = asyncio.Lock()
+
+
+def _config_path():
+    return CONFIG_FILE
 
 def get_default_config():
     """Returns the default configuration dictionary. Single source of truth."""
@@ -20,7 +29,11 @@ def get_default_config():
             "auto_launch_workers": False,
             "stop_workers_on_master_exit": True,
             "master_delegate_only": False,
-            "websocket_orchestration": True
+            "websocket_orchestration": True,
+            "worker_probe_concurrency": 8,
+            "worker_prep_concurrency": 4,
+            "media_sync_concurrency": 2,
+            "media_sync_timeout_seconds": 120
         },
         "tunnel": {
             "status": "stopped",
@@ -31,34 +44,93 @@ def get_default_config():
         }
     }
 
+def _merge_with_defaults(data, defaults):
+    """Recursively merge loaded config data with default keys."""
+    if not isinstance(data, dict):
+        return defaults
+
+    merged = {}
+    for key, default_value in defaults.items():
+        loaded_value = data.get(key, default_value)
+        if isinstance(default_value, dict) and isinstance(loaded_value, dict):
+            merged[key] = _merge_with_defaults(loaded_value, default_value)
+        else:
+            merged[key] = loaded_value
+
+    # Preserve unknown keys for forward compatibility.
+    for key, value in data.items():
+        if key not in merged:
+            merged[key] = value
+
+    return merged
+
+
+def invalidate_config_cache():
+    """Invalidate in-memory config cache so next load reads from disk."""
+    global _config_cache, _config_mtime
+    _config_cache = None
+    _config_mtime = 0.0
+
+
 def load_config():
     """Loads the config, falling back to defaults if the file is missing or invalid."""
-    if os.path.exists(CONFIG_FILE):
+    global _config_cache, _config_mtime
+    path = _config_path()
+
+    try:
+        mtime = os.path.getmtime(path)
+    except OSError:
+        if _config_cache is None:
+            _config_cache = get_default_config()
+        return _config_cache
+
+    if _config_cache is None or mtime != _config_mtime:
         try:
-            with open(CONFIG_FILE, 'r') as f:
-                data = json.load(f)
-                defaults = get_default_config()
-                for key, value in defaults.items():
-                    if key not in data:
-                        data[key] = value
-                return data
+            with open(path, 'r', encoding='utf-8') as f:
+                loaded = json.load(f)
+            _config_cache = _merge_with_defaults(loaded, get_default_config())
         except Exception as e:
             log(f"Error loading config, using defaults: {e}")
-    return get_default_config()
+            _config_cache = get_default_config()
+        _config_mtime = mtime
+
+    return _config_cache
 
 def save_config(config):
     """Saves the configuration to file."""
+    tmp_path = f"{_config_path()}.tmp"
     try:
-        with open(CONFIG_FILE, 'w') as f:
+        with open(tmp_path, 'w', encoding='utf-8') as f:
             json.dump(config, f, indent=2)
+            f.flush()
+            os.fsync(f.fileno())
+        os.replace(tmp_path, _config_path())
+        invalidate_config_cache()
         return True
     except Exception as e:
+        try:
+            os.unlink(tmp_path)
+        except OSError:
+            pass
         log(f"Error saving config: {e}")
         return False
 
+
+@asynccontextmanager
+async def config_transaction():
+    """Acquire config lock, yield loaded config, and save if changed."""
+    async with _config_lock:
+        config = load_config()
+        original_snapshot = json.dumps(config, sort_keys=True)
+        yield config
+        updated_snapshot = json.dumps(config, sort_keys=True)
+        if updated_snapshot != original_snapshot:
+            if not save_config(config):
+                raise RuntimeError("Failed to save config")
+
 def ensure_config_exists():
     """Creates default config file if it doesn't exist. Used by __init__.py"""
-    if not os.path.exists(CONFIG_FILE):
+    if not os.path.exists(_config_path()):
         default_config = get_default_config()
         if save_config(default_config):
             from .logging import debug_log

utils/constants.py

@@ -13,6 +13,10 @@ PROCESS_TERMINATION_TIMEOUT = 5.0
 WORKER_CHECK_INTERVAL = 2.0
 STATUS_CHECK_INTERVAL = 5.0
 
+# Cloudflare tunnel
+TUNNEL_START_TIMEOUT = float(os.environ.get("TUNNEL_START_TIMEOUT", "25"))
+CLOUDFLARE_LOG_BUFFER_SIZE = 200
+
 # Network
 CHUNK_SIZE = 8192
 LOG_TAIL_BYTES = 65536  # 64KB
@@ -30,6 +34,7 @@ TILE_TRANSFER_TIMEOUT = 30.0
 PROCESS_WAIT_TIMEOUT = 3.0
 QUEUE_INIT_TIMEOUT = 5.0
 TILE_SEND_TIMEOUT = 60.0
+JOB_INIT_GRACE_PERIOD = 10.0
 
 # Memory operations  
 MEMORY_CLEAR_DELAY = 0.5
@@ -38,4 +43,26 @@ MEMORY_CLEAR_DELAY = 0.5
 MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20'))  # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs)
 
 # Heartbeat monitoring
+HEARTBEAT_INTERVAL = float(os.environ.get('COMFYUI_HEARTBEAT_INTERVAL', '10'))  # Heartbeat/check interval in seconds
 HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60'))  # Worker heartbeat timeout in seconds (default 60s)
+
+# USDU result collection
+DYNAMIC_MODE_MAX_POLL_TIMEOUT = 10.0
+
+# Static mode job poll loop
+JOB_POLL_INTERVAL = 1.0
+JOB_POLL_MAX_ATTEMPTS = 20
+
+# Orchestration pipeline
+ORCHESTRATION_WORKER_PROBE_CONCURRENCY = int(
+    os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PROBE_CONCURRENCY', '8')
+)
+ORCHESTRATION_WORKER_PREP_CONCURRENCY = int(
+    os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PREP_CONCURRENCY', '4')
+)
+ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = int(
+    os.environ.get('COMFYUI_ORCHESTRATION_MEDIA_SYNC_CONCURRENCY', '2')
+)
+ORCHESTRATION_MEDIA_SYNC_TIMEOUT = float(
+    os.environ.get('COMFYUI_ORCHESTRATION_MEDIA_SYNC_TIMEOUT', '120')
+)

utils/crop_model_patch.py

@@ -0,0 +1,114 @@
+from contextlib import contextmanager
+
+import torch
+
+from .logging import debug_log
+from .usdu_utils import resize_region
+
+
+@contextmanager
+def crop_model_cond(model, crop_regions, init_size, canvas_size, tile_size, latent_crop=False):
+    """Clone model and crop compatible model patches for tile-local sampling."""
+    try:
+        patched_model = model.clone()
+    except Exception:
+        # Fallback to original model when clone/patch access is unavailable.
+        yield model
+        return
+
+    patches = (
+        patched_model
+        .model_options
+        .get("transformer_options", {})
+        .get("patches", {})
+    )
+    applied_croppers = {}
+    for _module, module_patches in patches.items():
+        for patch in module_patches:
+            if id(patch) in applied_croppers:
+                continue
+            if type(patch).__name__ not in ("DiffSynthCnetPatch", "ZImageControlPatch"):
+                continue
+            try:
+                cropper = ModelPatchCropper(patch).crop(crop_regions, canvas_size, latent_crop)
+                applied_croppers[id(patch)] = cropper
+            except Exception as exc:
+                debug_log(f"crop_model_cond: patch crop skipped for {type(patch).__name__}: {exc}")
+    try:
+        yield patched_model
+    finally:
+        for cropper in applied_croppers.values():
+            del cropper
+
+
+class ModelPatchCropper:
+    """Stateful crop helper that restores model patch tensors on cleanup."""
+
+    def __init__(self, patch):
+        self.patch = patch
+        self.original_state = {
+            "image": patch.image.clone() if isinstance(patch.image, torch.Tensor) else patch.image,
+            "encoded_image": patch.encoded_image.clone() if isinstance(patch.encoded_image, torch.Tensor) else patch.encoded_image,
+            "encoded_image_size": patch.encoded_image_size,
+        }
+        self.patch_class = type(patch).__name__
+        required_attrs = (
+            "image",
+            "model_patch",
+            "vae",
+            "strength",
+            "encoded_image",
+            "encoded_image_size",
+        )
+        missing = [attr for attr in required_attrs if not hasattr(patch, attr)]
+        if missing:
+            raise AttributeError(
+                f"{self.patch_class} missing required attrs: {', '.join(missing)}"
+            )
+
+    def __del__(self):
+        self.patch.image = self.original_state["image"]
+        self.patch.encoded_image = self.original_state["encoded_image"]
+        self.patch.encoded_image_size = self.original_state["encoded_image_size"]
+
+    def crop(self, crop_regions, canvas_size, latent_crop=True):
+        patch = self.patch
+
+        if not isinstance(crop_regions, list):
+            crop_regions = [crop_regions]
+
+        image_size = (patch.image.shape[2], patch.image.shape[1])  # (W,H)
+
+        cropped_images = []
+        for crop_region in crop_regions:
+            resized_crop = resize_region(crop_region, canvas_size, image_size)
+            x1, y1, x2, y2 = resized_crop
+            cropped_image = patch.image[:, y1:y2, x1:x2, :]
+            cropped_images.append(cropped_image)
+
+        concatenated_image = torch.cat(cropped_images, dim=0)
+        patch.image = concatenated_image
+        patch.encoded_image_size = (
+            concatenated_image.shape[1],
+            concatenated_image.shape[2],
+        )
+
+        if latent_crop:
+            downscale_ratio = patch.vae.spacial_compression_encode()
+            cropped_latents = []
+            for crop_region in crop_regions:
+                resized_crop = resize_region(crop_region, canvas_size, image_size)
+                x1, y1, x2, y2 = tuple(x // downscale_ratio for x in resized_crop)
+                cropped_latent = patch.encoded_image[:, :, y1:y2, x1:x2]
+                cropped_latents.append(cropped_latent)
+            patch.encoded_image = torch.cat(cropped_latents, dim=0)
+        else:
+            patch.__init__(
+                patch.model_patch,
+                patch.vae,
+                concatenated_image,
+                patch.strength,
+                inpaint_image=patch.inpaint_image,
+                mask=patch.mask,
+            )
+        return self

utils/exceptions.py

@@ -0,0 +1,42 @@
+"""Custom exceptions for ComfyUI-Distributed."""
+
+
+class DistributedError(Exception):
+    """Base exception for all ComfyUI-Distributed errors."""
+
+
+class WorkerError(DistributedError):
+    """Error related to a specific distributed worker."""
+
+    def __init__(self, message, worker_id=None):
+        super().__init__(message)
+        self.worker_id = worker_id
+
+
+class WorkerTimeoutError(WorkerError):
+    """Worker did not respond within the expected timeout."""
+
+
+class WorkerNotAvailableError(WorkerError):
+    """Worker is unreachable or not running."""
+
+
+class JobQueueError(DistributedError):
+    """Error in distributed job queue management."""
+
+
+class TileCollectionError(DistributedError):
+    """Error collecting processed tiles from workers."""
+
+
+class ProcessError(DistributedError):
+    """Error managing a worker subprocess."""
+
+    def __init__(self, message, pid=None, worker_id=None):
+        super().__init__(message)
+        self.pid = pid
+        self.worker_id = worker_id
+
+
+class TunnelError(DistributedError):
+    """Error managing the Cloudflare tunnel."""

utils/logging.py

@@ -3,20 +3,35 @@ Shared logging utilities for ComfyUI-Distributed.
 """
 import os
 import json
+import time
 
 # Config file is in parent directory
 CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json")
 
+_debug_cache: bool | None = None
+_debug_cache_time: float = 0.0
+_DEBUG_TTL: float = 5.0
+
 def is_debug_enabled():
     """Check if debug is enabled."""
+    global _debug_cache, _debug_cache_time
+
+    now = time.monotonic()
+    if _debug_cache is not None and (now - _debug_cache_time) < _DEBUG_TTL:
+        return _debug_cache
+
+    enabled = False
     if os.path.exists(CONFIG_FILE):
         try:
             with open(CONFIG_FILE, 'r') as f:
                 config = json.load(f)
-                return config.get("settings", {}).get("debug", False)
-        except:
+                enabled = config.get("settings", {}).get("debug", False)
+        except (OSError, json.JSONDecodeError, ValueError):
             pass
-    return False
+
+    _debug_cache = enabled
+    _debug_cache_time = now
+    return enabled
 
 def debug_log(message):
     """Log debug messages only if debug is enabled in config."""
@@ -25,4 +40,4 @@ def debug_log(message):
 
 def log(message):
     """Always log important messages."""
-    print(f"[Distributed] {message}")
+    print(f"[Distributed] {message}")

utils/network.py

@@ -1,7 +1,10 @@
 """
 Network and API utilities for ComfyUI-Distributed.
 """
+import asyncio
 import aiohttp
+import re
+import server
 from aiohttp import web
 from .logging import debug_log
 
@@ -11,6 +14,11 @@ _client_session = None
 async def get_client_session():
     """Get or create a shared aiohttp client session."""
     global _client_session
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError as exc:
+        raise RuntimeError("get_client_session() requires an active asyncio event loop.") from exc
+
     if _client_session is None or _client_session.closed:
         connector = aiohttp.TCPConnector(limit=100, limit_per_host=30)
         # Don't set timeout here - set it per request
@@ -26,8 +34,14 @@ async def cleanup_client_session():
 
 async def handle_api_error(request, error, status=500):
     """Standardized error response handler."""
-    debug_log(f"API Error: {error}")
-    return web.json_response({"status": "error", "message": str(error)}, status=status)
+    if isinstance(error, list):
+        messages = [str(item) for item in error]
+        debug_log(f"API Error [{status}]: {messages}")
+        return web.json_response({"errors": messages}, status=status)
+
+    message = str(error)
+    debug_log(f"API Error [{status}]: {message}")
+    return web.json_response({"error": message}, status=status)
 
 def get_server_port():
     """Get the ComfyUI server port."""
@@ -39,3 +53,149 @@ def get_server_loop():
     import server
     return server.PromptServer.instance.loop
 
+
+def normalize_host(value):
+    if value is None:
+        return None
+    if not isinstance(value, str):
+        return value
+    host = value.strip()
+    if not host:
+        return host
+    host = re.sub(r"^https?://", "", host, flags=re.IGNORECASE)
+    return host.split("/")[0]
+
+
+def _split_host_and_port(host):
+    if not host:
+        return host, None
+
+    if host.startswith("["):
+        match = re.match(r"^(\[[^\]]+\])(?::(\d+))?$", host)
+        if match:
+            parsed_port = int(match.group(2)) if match.group(2) else None
+            return match.group(1), parsed_port
+        return host, None
+
+    if host.count(":") == 1:
+        candidate_host, candidate_port = host.rsplit(":", 1)
+        if candidate_port.isdigit():
+            return candidate_host, int(candidate_port)
+
+    return host, None
+
+
+def build_worker_url(worker, endpoint=""):
+    """Construct the worker base URL with optional endpoint."""
+    host = (worker.get("host") or "").strip()
+    port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
+
+    if not host:
+        host = getattr(server.PromptServer.instance, "address", "127.0.0.1") or "127.0.0.1"
+
+    if host.startswith(("http://", "https://")):
+        base = host.rstrip("/")
+    else:
+        is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
+        scheme = "https" if is_cloud else "http"
+        default_port = 443 if scheme == "https" else 80
+        port_part = "" if port == default_port else f":{port}"
+        base = f"{scheme}://{host}{port_part}"
+
+    return f"{base}{endpoint}"
+
+
+async def probe_worker(worker_url: str, timeout: float = 5.0) -> dict | None:
+    """GET {worker_url}/prompt. Returns parsed JSON or None on any failure."""
+    base_url = (worker_url or "").strip().rstrip("/")
+    if not base_url:
+        return None
+    probe_url = base_url if base_url.endswith("/prompt") else f"{base_url}/prompt"
+    session = await get_client_session()
+    try:
+        async with session.get(
+            probe_url,
+            timeout=aiohttp.ClientTimeout(total=float(timeout)),
+        ) as response:
+            if response.status != 200:
+                debug_log(f"[Distributed] Worker probe non-200 status: {response.status} ({probe_url})")
+                return None
+            payload = await response.json()
+            if isinstance(payload, dict):
+                return payload
+            debug_log(f"[Distributed] Worker probe returned non-object JSON: {probe_url}")
+            return None
+    except asyncio.TimeoutError:
+        debug_log(f"[Distributed] Worker probe timed out: {probe_url}")
+        return None
+    except aiohttp.ClientConnectorError:
+        debug_log(f"[Distributed] Worker unreachable: {probe_url}")
+        return None
+    except Exception as exc:
+        debug_log(f"[Distributed] Worker probe error ({probe_url}): {exc}")
+        return None
+
+
+def build_master_url(config=None, prompt_server_instance=None):
+    """Build the best public URL workers should use to reach the master."""
+    if config is None:
+        from .config import load_config
+        config = load_config()
+
+    prompt_server_instance = prompt_server_instance or server.PromptServer.instance
+    master_cfg = (config or {}).get("master", {}) or {}
+    configured_host = (master_cfg.get("host") or "").strip()
+    runtime_port = getattr(prompt_server_instance, "port", 8188) or 8188
+
+    def _needs_https(hostname):
+        hostname = hostname.lower()
+        https_domains = (
+            ".proxy.runpod.net",
+            ".ngrok-free.app",
+            ".ngrok-free.dev",
+            ".ngrok.io",
+            ".trycloudflare.com",
+            ".cloudflare.dev",
+        )
+        return any(hostname.endswith(suffix) for suffix in https_domains)
+
+    if configured_host:
+        if configured_host.startswith(("http://", "https://")):
+            return configured_host.rstrip("/")
+
+        host, explicit_port = _split_host_and_port(configured_host)
+        port = explicit_port if explicit_port is not None else int(runtime_port)
+        scheme = "https" if _needs_https(host) or port == 443 else "http"
+        default_port_for_scheme = 443 if scheme == "https" else 80
+        if explicit_port is None and scheme == "https" and _needs_https(host):
+            port = default_port_for_scheme
+        port_part = "" if port == default_port_for_scheme else f":{port}"
+        return f"{scheme}://{host}{port_part}"
+
+    address = getattr(prompt_server_instance, "address", "127.0.0.1") or "127.0.0.1"
+    if address in ("0.0.0.0", "::"):
+        address = "127.0.0.1"
+    port = int(runtime_port)
+    scheme = "https" if port == 443 else "http"
+    default_port_for_scheme = 443 if scheme == "https" else 80
+    port_part = "" if port == default_port_for_scheme else f":{port}"
+    return f"{scheme}://{address}{port_part}"
+
+
+def build_master_callback_url(worker, config=None, prompt_server_instance=None):
+    """Build the callback URL a specific worker should use to reach the master."""
+    prompt_server_instance = prompt_server_instance or server.PromptServer.instance
+
+    worker_type = str((worker or {}).get("type") or "").strip().lower()
+    worker_host = normalize_host((worker or {}).get("host"))
+    local_hosts = {"", "localhost", "127.0.0.1", "::1", "[::1]", "0.0.0.0"}
+
+    is_local_worker = worker_type == "local" or worker_host in local_hosts
+    if is_local_worker:
+        port = int(getattr(prompt_server_instance, "port", 8188) or 8188)
+        scheme = "https" if port == 443 else "http"
+        default_port_for_scheme = 443 if scheme == "https" else 80
+        port_part = "" if port == default_port_for_scheme else f":{port}"
+        return f"{scheme}://127.0.0.1{port_part}"
+
+    return build_master_url(config=config, prompt_server_instance=prompt_server_instance)

utils/trace_logger.py

@@ -0,0 +1,13 @@
+from .logging import debug_log, log
+
+
+def trace_prefix(trace_execution_id: str) -> str:
+    return f"[Distributed][exec:{trace_execution_id}]"
+
+
+def trace_debug(trace_execution_id: str, message: str) -> None:
+    debug_log(f"{trace_prefix(trace_execution_id)} {message}")
+
+
+def trace_info(trace_execution_id: str, message: str) -> None:
+    log(f"{trace_prefix(trace_execution_id)} {message}")

utils/usdu_managment.py

@@ -1,389 +1,32 @@
-import asyncio
-import time
-import json
-import copy
-import os
-import io
-from aiohttp import web, ClientTimeout
-import server
-from PIL import Image
+"""Backward-compatibility shim for USDU helpers.
 
-# Import from other utilities
-from .logging import debug_log, log
-from .network import handle_api_error, get_client_session
-# We avoid converting to tensors on the master for tiles; blending uses PIL
+Route handlers and job logic now live in:
+- upscale.job_store
+- upscale.job_timeout
+- upscale.payload_parsers
+- api.usdu_routes
+"""
 
-# 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)))
+from ..upscale.conditioning import clone_conditioning, clone_control_chain
+from ..upscale.job_store import (
+    MAX_PAYLOAD_SIZE,
+    _cleanup_job,
+    _drain_results_queue,
+    _get_completed_count,
+    _init_job_queue,
+    _mark_task_completed,
+    ensure_tile_jobs_initialized,
+    init_dynamic_job,
+    init_static_job_batched,
+)
+from ..upscale.job_timeout import _check_and_requeue_timed_out_workers
+from ..upscale.payload_parsers import _parse_tiles_from_form
+from .logging import debug_log
+from .network import get_client_session
 
-# Import HEARTBEAT_TIMEOUT from constants
-from .constants import HEARTBEAT_TIMEOUT
-from .config import load_config
-
-
-def _parse_tiles_from_form(data):
-    """Parse tiles submitted via multipart/form-data into a list of tile dicts.
-
-    Expects the following fields in the aiohttp form data:
-    - 'tiles_metadata': JSON list with per-tile metadata items containing at least
-      'tile_idx', 'x', 'y', 'extracted_width', 'extracted_height'. Optional
-      'batch_idx' and 'global_idx' are included when available.
-    - 'tile_{i}': image bytes for each tile described in tiles_metadata (PNG).
-    - 'padding': integer padding used during extraction (optional; defaults 0).
-
-    Returns: list of dicts with keys: 'image', 'tile_idx', 'x', 'y',
-    'extracted_width', 'extracted_height', and optional 'batch_idx', 'global_idx',
-    plus 'padding'.
-    """
-    try:
-        # Parse padding if present
-        padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
-    except Exception:
-        padding = 0
-
-    # Parse tiles metadata (JSON list)
-    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 = []
-    # Iterate over metadata items and corresponding uploaded files tile_0, tile_1, ...
-    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}")
-
-        # Read image bytes and decode to PIL
-        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}")
-
-        # Build tile dictionary (store PIL only; master blends via PIL)
-        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}")
-
-        # Optional fields
-        if 'batch_idx' in meta:
-            try:
-                tile_info['batch_idx'] = int(meta['batch_idx'])
-            except Exception:
-                pass
-        if 'global_idx' in meta:
-            try:
-                tile_info['global_idx'] = int(meta['global_idx'])
-            except Exception:
-                pass
-
-        tiles.append(tile_info)
-
-    return tiles
-
-
-# Unified Job Data Structure Keys
-JOB_QUEUE = 'queue'
-JOB_MODE = 'mode'
-JOB_COMPLETED_TASKS = 'completed_tasks'
-JOB_WORKER_STATUS = 'worker_status'
-JOB_ASSIGNED_TO_WORKERS = 'assigned_to_workers'
-JOB_PENDING_TASKS = 'pending_tasks'
-JOB_BATCH_SIZE = 'batch_size'  # For dynamic
-JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image'  # For static
-
-# Task Types
-TASK_TYPE_TILE = 'tile'
-TASK_TYPE_IMAGE = 'image'
-
-from typing import List, Optional
-
-async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: List[str], all_indices: Optional[List[int]] = None):
-    """Initialize queue for dynamic mode (per-image), with collector fields.
-
-    - Creates JOB_PENDING_TASKS with image indices
-    - Adds 'completed_images' dict and 'pending_images' alias used by collectors
-    """
-    await _init_job_queue(
-        multi_job_id,
-        'dynamic',
-        batch_size=batch_size,
-        all_indices=all_indices or list(range(batch_size)),
-        enabled_workers=enabled_workers,
-    )
-
-    prompt_server = ensure_tile_jobs_initialized()
-    async with prompt_server.distributed_tile_jobs_lock:
-        job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
-        job_data['completed_images'] = {}
-        job_data['pending_images'] = job_data[JOB_PENDING_TASKS]
-    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]):
-    """Initialize queue for static mode (batched-per-tile).
-
-    - Populates JOB_PENDING_TASKS with tile ids [0..num_tiles_per_image-1]
-    """
-    await _init_job_queue(
-        multi_job_id,
-        'static',
-        batch_size=batch_size,
-        num_tiles_per_image=num_tiles_per_image,
-        enabled_workers=enabled_workers,
-        batched_static=True,
-    )
-    # Initialization handled by master; avoid duplicate init logs here
-
-async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None, batched_static: bool = False):
-    """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
-
-        job_data = {
-            JOB_QUEUE: asyncio.Queue(),
-            JOB_MODE: mode,
-            JOB_COMPLETED_TASKS: {},
-            JOB_WORKER_STATUS: {w: time.time() for w in enabled_workers or []},
-            JOB_ASSIGNED_TO_WORKERS: {w: [] for w in enabled_workers or []},
-            JOB_PENDING_TASKS: asyncio.Queue(),
-        }
-
-        if mode == 'dynamic':
-            job_data[JOB_BATCH_SIZE] = batch_size
-            pending_queue = job_data[JOB_PENDING_TASKS]
-            for i in (all_indices or range(batch_size)):
-                await pending_queue.put(i)
-            debug_log(f"Initialized image queue with {batch_size} pending items")
-        elif mode == 'static':
-            job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image
-            job_data[JOB_BATCH_SIZE] = batch_size
-            job_data['batched_static'] = bool(batched_static)
-            # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
-            pending_queue = job_data[JOB_PENDING_TASKS]
-            if batched_static and num_tiles_per_image is not None:
-                for i in range(num_tiles_per_image):
-                    await pending_queue.put(i)
-            else:
-                total_tiles = batch_size * num_tiles_per_image
-                for i in range(total_tiles):
-                    await pending_queue.put(i)
-            
-            # Keep backward compatibility - if task assignments provided, still track them
-            if task_assignments and enabled_workers:
-                # task_assignments[0] is for master, 1+ are for workers
-                for i, worker_id in enumerate(enabled_workers):
-                    if i + 1 < len(task_assignments):
-                        job_data[JOB_ASSIGNED_TO_WORKERS][worker_id] = task_assignments[i + 1]
-                        debug_log(f"Worker {worker_id} pre-assigned {len(task_assignments[i + 1])} tasks (legacy mode)")
-        else:
-            raise ValueError(f"Unknown mode: {mode}")
-
-        prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
-
-# Note: legacy task distribution and queue pull helpers removed
-
-async def _drain_results_queue(multi_job_id):
-    """Drain pending results from queue and update completed_tasks. Returns count drained.
-
-    Uses non-blocking get_nowait to avoid await timeouts and reduce latency.
-    """
-    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 job_data or JOB_QUEUE not in job_data or JOB_COMPLETED_TASKS not in job_data:
-            return 0
-        q = job_data[JOB_QUEUE]
-        completed_tasks = job_data[JOB_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
-            elif 'tensor' in result and 'tile_idx' in result:  # Single tile backward compat
-                task_id = result.get('global_idx', result['tile_idx'])
-                if task_id not in completed_tasks:
-                    completed_tasks[task_id] = {
-                        'tensor': result['tensor'],
-                        'tile_idx': result['tile_idx'],
-                        'x': result['x'],
-                        'y': result['y'],
-                        'extracted_width': result['extracted_width'],
-                        'extracted_height': result['extracted_height'],
-                        'padding': result['padding'],
-                        'batch_idx': result.get('batch_idx', 0),
-                        'global_idx': task_id
-                    }
-                    collected += 1
-
-            if is_last:
-                # Track worker completion
-                if worker_id in job_data[JOB_WORKER_STATUS]:
-                    del job_data[JOB_WORKER_STATUS][worker_id]
-
-        return collected
-
-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()
-    async with prompt_server.distributed_tile_jobs_lock:
-        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-        if not job_data:
-            return 0
-
-        current_time = time.time()
-        requeued_count = 0
-        completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {})
-
-        # Allow override via config setting 'worker_timeout_seconds'
-        cfg = load_config()
-        hb_timeout = int(cfg.get('settings', {}).get('worker_timeout_seconds', HEARTBEAT_TIMEOUT))
-
-        for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()):
-            age = current_time - last_heartbeat
-            debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
-            if age > hb_timeout:
-                # Busy-only grace policy: require positive signal from worker (/prompt)
-                # We also log assignment state for diagnostics but do not grace on it alone.
-                assigned = job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, [])
-                incomplete_assigned = 0
-                try:
-                    if assigned:
-                        batched_static = bool(job_data.get('batched_static', False))
-                        if batched_static:
-                            num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1)
-                            batch_size = job_data.get(JOB_BATCH_SIZE, 1)
-                            for task_id in assigned:
-                                for b in range(batch_size):
-                                    gidx = b * num_tiles_per_image + task_id
-                                    if gidx not in completed_tasks:
-                                        incomplete_assigned += 1
-                                        break
-                        else:
-                            for task_id in assigned:
-                                if task_id not in completed_tasks:
-                                    incomplete_assigned += 1
-                    debug_log(f"Assigned diagnostics: total_assigned={len(assigned)} incomplete_assigned={incomplete_assigned}")
-                except Exception as e:
-                    debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
-
-                busy = False
-                probe_status = None
-                probe_queue = None
-                try:
-                    cfg_workers = load_config().get('workers', [])
-                    wrec = next((w for w in cfg_workers if str(w.get('id')) == str(worker)), None)
-                    if wrec:
-                        host = wrec.get('host') or 'localhost'
-                        port = int(wrec.get('port', 8188))
-                        url = f"http://{host}:{port}/prompt"
-                        debug_log(f"Probing worker {worker} at {url}")
-                        session = await get_client_session()
-                        async with session.get(url, timeout=ClientTimeout(total=2.0)) as resp:
-                            probe_status = resp.status
-                            if resp.status == 200:
-                                try:
-                                    payload = await resp.json()
-                                    probe_queue = int(payload.get('exec_info', {}).get('queue_remaining', 0))
-                                except Exception:
-                                    probe_queue = 0
-                                busy = probe_queue is not None and probe_queue > 0
-                except Exception as e:
-                    debug_log(f"Probe failed for worker {worker}: {e}")
-                finally:
-                    debug_log(f"Probe diagnostics: http_status={probe_status} queue_remaining={probe_queue}")
-
-                if busy:
-                    job_data[JOB_WORKER_STATUS][worker] = current_time
-                    debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
-                    continue
-
-                log(f"Worker {worker} timed out")
-                for task_id in job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, []):
-                    # If batched_static, task_id is a tile_idx; consider it complete only if
-                    # all corresponding global_idx entries are present in completed_tasks.
-                    batched_static = bool(job_data.get('batched_static', False))
-                    if batched_static:
-                        num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1)
-                        batch_size = job_data.get(JOB_BATCH_SIZE, 1)
-                        # Check all global indices for this tile across the batch
-                        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[JOB_PENDING_TASKS].put(task_id)
-                            requeued_count += 1
-                    else:
-                        # Legacy/global-idx mode: task_id is a global index key
-                        if task_id not in completed_tasks:
-                            await job_data[JOB_PENDING_TASKS].put(task_id)
-                            requeued_count += 1
-                if JOB_WORKER_STATUS in job_data:
-                    del job_data[JOB_WORKER_STATUS][worker]
-                if JOB_ASSIGNED_TO_WORKERS in job_data:
-                    job_data[JOB_ASSIGNED_TO_WORKERS][worker] = []
-
-        return requeued_count
-
-async def _get_completed_count(multi_job_id):
-    """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 job_data and JOB_COMPLETED_TASKS in job_data:
-            return len(job_data[JOB_COMPLETED_TASKS])
-        return 0
-
-async def _mark_task_completed(multi_job_id, task_id, result):
-    """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 job_data and JOB_COMPLETED_TASKS in job_data:
-            job_data[JOB_COMPLETED_TASKS][task_id] = result
 
 async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
-    """Send heartbeat to master."""
+    """Send heartbeat to master from worker-side processing loops."""
     try:
         data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)}
         session = await get_client_session()
@@ -393,295 +36,20 @@ async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
     except Exception as e:
         debug_log(f"Heartbeat failed: {e}")
 
-async def _cleanup_job(multi_job_id):
-    """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}")
-
-# API Endpoints (generalized)
-
-@server.PromptServer.instance.routes.post("/distributed/heartbeat")
-async def heartbeat_endpoint(request):
-    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 JOB_WORKER_STATUS in job_data:
-                    job_data[JOB_WORKER_STATUS][worker_id] = time.time()
-                    debug_log(f"Heartbeat from worker {worker_id}")
-                    return web.json_response({"status": "success"})
-                else:
-                    return await handle_api_error(request, "Worker status tracking not available", 400)
-            else:
-                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):
-    """Endpoint for workers to submit processed tiles in static mode."""
-    try:
-        content_length = request.headers.get('content-length')
-        if content_length and int(content_length) > 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')
-        is_last = data.get('is_last', 'False').lower() == 'true'
-        
-        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()
-        
-        batch_size = int(data.get('batch_size', 0))
-        tiles = []
-        
-        # 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 JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
-                        return await handle_api_error(request, "Job not configured for tile submissions", 400)
-                    if JOB_QUEUE in job_data:
-                        await job_data[JOB_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 JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
-                    return await handle_api_error(request, "Job not configured for tile submissions", 400)
-                
-                q = job_data[JOB_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"})
-            else:
-                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):
-    """Endpoint for workers to submit processed images in dynamic mode."""
-    try:
-        content_length = request.headers.get('content-length')
-        if content_length and int(content_length) > 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')
-        is_last = data.get('is_last', 'False').lower() == 'true'
-        
-        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:
-            image_idx = int(data.get('image_idx'))
-            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 JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
-                        return await handle_api_error(request, "Job not configured for image submissions", 400)
-                    if JOB_QUEUE in job_data:
-                        await job_data[JOB_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 JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
-                        return await handle_api_error(request, "Job not configured for image submissions", 400)
-                    if JOB_QUEUE in job_data:
-                        await job_data[JOB_QUEUE].put({
-                            'worker_id': worker_id,
-                            'is_last': True,
-                            'tiles': []  # For compatibility
-                        })
-                        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)
-
-# Note: Removed legacy /distributed/tile_complete endpoint. Use /distributed/submit_tiles.
-
-
-
-# Helper functions for shallow copying conditioning without duplicating models
-def clone_control_chain(control, clone_hint=True):
-    """Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
-    if control is None:
-        return None
-    new_control = copy.copy(control)  # Shallow copy (shares model)
-    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, clone_hints=True):
-    """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:
-            if new_dict['mask'] is not None:
-                new_dict['mask'] = new_dict['mask'].clone()
-        # Handle other potential fields if needed
-        if 'pooled_output' in new_dict:
-            if 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'][:]  # Shallow copy list/tuple
-        new_cond.append([new_emb, new_dict])
-    return new_cond
-
-# Helper function to ensure persistent state is initialized
-def ensure_tile_jobs_initialized():
-    """Ensure tile job storage is initialized on the server instance."""
-    prompt_server = server.PromptServer.instance
-    if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
-        debug_log("Initializing persistent tile job queue on server instance.")
-        prompt_server.distributed_pending_tile_jobs = {}
-        prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
-    else:
-        # Clean up any legacy queue structures that don't have the 'mode' field
-        # (Should be rare after fixes, but keep for safety)
-        to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
-                     if not isinstance(job_data, dict) or 'mode' not in job_data]
-        for job_id in to_remove:
-            debug_log(f"Removing legacy queue structure for job {job_id}")
-            del prompt_server.distributed_pending_tile_jobs[job_id]
-    return prompt_server
-
-# API Endpoint for tile completion
-@server.PromptServer.instance.routes.post("/distributed/request_image")
-async def request_image_endpoint(request):
-    """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
-    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, dict) or 'mode' not in job_data:
-                    return await handle_api_error(request, "Invalid job data structure", 500)
-                
-                mode = job_data['mode']
-                
-                # Handle both dynamic and static modes
-                if mode == 'dynamic' and 'pending_images' in job_data:
-                    pending_queue = job_data['pending_images']
-                elif mode == 'static' and JOB_PENDING_TASKS in job_data:
-                    pending_queue = job_data[JOB_PENDING_TASKS]
-                else:
-                    return await handle_api_error(request, "Invalid job configuration", 400)
-                
-                try:
-                    task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
-                    # Track assigned task
-                    if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']:
-                        job_data['assigned_to_workers'][worker_id].append(task_idx)
-                    # Update worker heartbeat
-                    if 'worker_status' in job_data:
-                        job_data['worker_status'][worker_id] = time.time()
-                    # Get estimated remaining count
-                    remaining = pending_queue.qsize()  # Approximate
-                    
-                    # Return appropriate response based on mode
-                    if mode == 'dynamic':
-                        debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
-                        return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
-                    else:  # static
-                        debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
-                        return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining, "batched_static": job_data.get('batched_static', False)})
-                except asyncio.TimeoutError:
-                    if mode == 'dynamic':
-                        return web.json_response({"image_idx": None})
-                    else:
-                        return web.json_response({"tile_idx": None})
-            else:
-                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):
-    """Endpoint to check if a job is ready."""
-    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(job_data and isinstance(job_data, dict) and 'queue' in job_data)
-        return web.json_response({"ready": ready})
-
 
+__all__ = [
+    "MAX_PAYLOAD_SIZE",
+    "_check_and_requeue_timed_out_workers",
+    "_cleanup_job",
+    "_drain_results_queue",
+    "_get_completed_count",
+    "_init_job_queue",
+    "_mark_task_completed",
+    "_parse_tiles_from_form",
+    "_send_heartbeat_to_master",
+    "clone_conditioning",
+    "clone_control_chain",
+    "ensure_tile_jobs_initialized",
+    "init_dynamic_job",
+    "init_static_job_batched",
+]

utils/usdu_utils.py

@@ -466,9 +466,13 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size,
     x1_px, y1_px, x2_px, y2_px = region
 
     new_latents = []
-    for t in latents:  # (B,C,H_lat_in,W_lat_in)
+    for t in latents:  # (B,C,H_lat_in,W_lat_in) or (B,C,1,H_lat_in,W_lat_in)
+        has_5d = False
+        if t.ndim == 5:
+            has_5d = True
+            t = t.squeeze(2)
         if t.ndim != 4:
-            raise ValueError(f"expected BCHW, got {t.shape}")
+            raise ValueError(f"expected BCHW or BC1HW, got {t.shape}")
 
         # 1. Resize to canvas resolution in latent units only if needed
         if t.shape[-2:] != (H_can_lat, W_can_lat):
@@ -491,6 +495,8 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size,
                                 mode="bilinear",
                                 align_corners=False)
 
+        if has_5d:
+            cropped = cropped.unsqueeze(2)
         new_latents.append(cropped)
 
     cond_dict["reference_latents"] = new_latents
@@ -508,4 +514,4 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0)
         crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
         crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
         cropped.append(n)
-    return cropped
+    return cropped

vitest.config.js

@@ -0,0 +1,8 @@
+import { defineConfig } from "vitest/config";
+
+export default defineConfig({
+  test: {
+    include: ["web/tests/**/*.test.js"],
+    environment: "node",
+  },
+});

web/apiClient.js

@@ -1,27 +1,43 @@
 import { TIMEOUTS } from './constants.js';
+import { normalizeWorkerUrl } from './urlUtils.js';
 
 export function createApiClient(baseUrl) {
-    const request = async (endpoint, options = {}, retries = TIMEOUTS.MAX_RETRIES) => {
+    const normalizedBaseUrl = normalizeWorkerUrl(baseUrl);
+
+    const request = async (
+        endpoint,
+        options = {},
+        { retries = TIMEOUTS.MAX_RETRIES, retry = true } = {},
+    ) => {
+        const maxAttempts = retry ? retries : 1;
         let lastError;
         let delay = TIMEOUTS.RETRY_DELAY; // Initial delay for exponential backoff
 
-        for (let attempt = 0; attempt < retries; attempt++) {
+        for (let attempt = 0; attempt < maxAttempts; attempt++) {
             try {
-                const response = await fetch(`${baseUrl}${endpoint}`, {
-                    headers: { 'Content-Type': 'application/json' },
-                    ...options
+                const headers = {
+                    'Content-Type': 'application/json',
+                    ...(options.headers || {}),
+                };
+                const response = await fetch(`${normalizedBaseUrl}${endpoint}`, {
+                    ...options,
+                    headers,
                 });
                 
                 if (!response.ok) {
-                    const error = await response.json().catch(() => ({ message: 'Request failed' }));
-                    throw new Error(error.message || `HTTP ${response.status}`);
+                    const error = await response.json().catch(() => ({}));
+                    const message = error.message
+                        || error.error
+                        || (Array.isArray(error.errors) ? error.errors.join('; ') : null)
+                        || `HTTP ${response.status}`;
+                    throw new Error(message);
                 }
                 
                 return await response.json();
             } catch (error) {
                 lastError = error;
-                console.log(`API Error (attempt ${attempt + 1}/${retries}): ${endpoint} - ${error.message}`);
-                if (attempt < retries - 1) {
+                console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${error.message}`);
+                if (attempt < maxAttempts - 1) {
                     await new Promise(resolve => setTimeout(resolve, delay));
                     delay *= 2; // Exponential backoff
                 }
@@ -29,6 +45,39 @@ export function createApiClient(baseUrl) {
         }
         throw lastError;
     };
+
+    const requestUrl = async (
+        url,
+        options = {},
+        { retries = TIMEOUTS.MAX_RETRIES, retry = true } = {},
+    ) => {
+        const maxAttempts = retry ? retries : 1;
+        let lastError;
+        let delay = TIMEOUTS.RETRY_DELAY;
+
+        for (let attempt = 0; attempt < maxAttempts; attempt++) {
+            try {
+                const response = await fetch(url, options);
+                if (!response.ok) {
+                    const error = await response.json().catch(() => ({}));
+                    const message = error.message
+                        || error.error
+                        || (Array.isArray(error.errors) ? error.errors.join('; ') : null)
+                        || `HTTP ${response.status}`;
+                    throw new Error(message);
+                }
+                return await response.json();
+            } catch (error) {
+                lastError = error;
+                console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${url} - ${error.message}`);
+                if (attempt < maxAttempts - 1) {
+                    await new Promise(resolve => setTimeout(resolve, delay));
+                    delay *= 2;
+                }
+            }
+        }
+        throw lastError;
+    };
     
     return {
         // Config endpoints
@@ -40,28 +89,28 @@ export function createApiClient(baseUrl) {
             return request('/distributed/config/update_worker', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId, ...data })
-            });
+            }, { retry: false });
         },
         
         async deleteWorker(workerId) {
             return request('/distributed/config/delete_worker', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId })
-            });
+            }, { retry: false });
         },
         
         async updateSetting(key, value) {
             return request('/distributed/config/update_setting', {
                 method: 'POST',
                 body: JSON.stringify({ key, value })
-            });
+            }, { retry: false });
         },
         
         async updateMaster(data) {
             return request('/distributed/config/update_master', {
                 method: 'POST',
                 body: JSON.stringify(data)
-            });
+            }, { retry: false });
         },
         
         // Worker management endpoints
@@ -69,14 +118,14 @@ export function createApiClient(baseUrl) {
             return request('/distributed/launch_worker', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId })
-            });
+            }, { retry: false });
         },
         
         async stopWorker(workerId) {
             return request('/distributed/stop_worker', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId })
-            });
+            }, { retry: false });
         },
         
         async getManagedWorkers() {
@@ -86,28 +135,89 @@ export function createApiClient(baseUrl) {
         async getWorkerLog(workerId, lines = 1000) {
             return request(`/distributed/worker_log/${workerId}?lines=${lines}`);
         },
+
+        async getRemoteWorkerLog(workerId, lines = 300) {
+            return request(`/distributed/remote_worker_log/${workerId}?lines=${lines}`);
+        },
         
         async clearLaunchingFlag(workerId) {
             return request('/distributed/worker/clear_launching', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId })
-            });
+            }, { retry: false });
         },
         
-        // Job preparation
-        async prepareJob(multiJobId) {
-            return request('/distributed/prepare_job', {
+        async queueDistributed(payload) {
+            return request('/distributed/queue', {
                 method: 'POST',
-                body: JSON.stringify({ multi_job_id: multiJobId })
-            });
+                headers: {
+                    ...(payload?.trace_execution_id
+                        ? { 'X-Idempotency-Key': payload.trace_execution_id }
+                        : {}),
+                },
+                body: JSON.stringify(payload)
+            }, { retry: false });
         },
-        
-        // Image loading
-        async loadImage(imagePath) {
-            return request('/distributed/load_image', {
+
+        async probeWorker(workerUrl, timeoutMs = TIMEOUTS.STATUS_CHECK, signal = null) {
+            const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl);
+            const controller = new AbortController();
+            const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
+            const effectiveSignal = signal
+                ? AbortSignal.any([controller.signal, signal])
+                : controller.signal;
+            try {
+                const response = await fetch(`${normalizedWorkerUrl}/prompt`, {
+                    method: 'GET',
+                    mode: 'cors',
+                    cache: 'no-store',
+                    signal: effectiveSignal,
+                });
+
+                if (!response.ok) {
+                    return { ok: false, status: response.status, queueRemaining: null };
+                }
+
+                let data;
+                try {
+                    data = await response.json();
+                } catch {
+                    return { ok: false, status: response.status, queueRemaining: null };
+                }
+
+                if (!data || typeof data !== "object" || Array.isArray(data)) {
+                    return { ok: false, status: response.status, queueRemaining: null };
+                }
+
+                const execInfo = data.exec_info;
+                if (!execInfo || typeof execInfo !== "object" || Array.isArray(execInfo)) {
+                    return { ok: false, status: response.status, queueRemaining: null };
+                }
+
+                const rawQueueRemaining = execInfo.queue_remaining;
+                const queueRemaining = Number(rawQueueRemaining);
+                if (!Number.isFinite(queueRemaining)) {
+                    return { ok: false, status: response.status, queueRemaining: null };
+                }
+
+                return {
+                    ok: true,
+                    status: response.status,
+                    queueRemaining: Math.max(0, queueRemaining),
+                };
+            } finally {
+                clearTimeout(timeoutId);
+            }
+        },
+
+        async dispatchToWorker(workerUrl, promptPayload) {
+            const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl);
+            return requestUrl(`${normalizedWorkerUrl}/prompt`, {
                 method: 'POST',
-                body: JSON.stringify({ image_path: imagePath })
-            });
+                headers: { 'Content-Type': 'application/json' },
+                mode: 'cors',
+                body: JSON.stringify(promptPayload),
+            }, { retry: false });
         },
         
         // Network info
@@ -148,14 +258,14 @@ export function createApiClient(baseUrl) {
             return request('/distributed/tunnel/start', {
                 method: 'POST',
                 body: JSON.stringify({})
-            });
+            }, { retry: false });
         },
 
         async stopTunnel() {
             return request('/distributed/tunnel/stop', {
                 method: 'POST',
                 body: JSON.stringify({})
-            });
+            }, { retry: false });
         },
 
         async getTunnelStatus() {

web/constants.js

@@ -17,6 +17,7 @@ export const BUTTON_STYLES = {
     launch: "background-color: #4a7c4a;",
     stop: "background-color: #7c4a4a;",
     log: "background-color: #685434;",
+    working: "background-color: #666;",
     clearMemory: "background-color: #555; padding: 6px 14px;",
     interrupt: "background-color: #555; padding: 6px 14px;",
 };
@@ -62,6 +63,22 @@ export const PULSE_ANIMATION_CSS = `
         animation: pulse 1.2s ease-in-out infinite;
         transform-origin: center;
     }
+
+    .worker-status--online {
+        background: var(--status-online, #3ca03c) !important;
+    }
+
+    .worker-status--offline {
+        background: var(--status-offline, #c04c4c) !important;
+    }
+
+    .worker-status--unknown {
+        background: var(--status-unknown, #888) !important;
+    }
+
+    .worker-status--processing {
+        background: var(--status-processing, #f0ad4e) !important;
+    }
     
     /* Button hover effects */
     .distributed-button:hover:not(:disabled) {
@@ -89,7 +106,9 @@ export const PULSE_ANIMATION_CSS = `
     .worker-settings.expanded {
         max-height: 500px;
         opacity: 1;
-        padding: 12px 0;
+        padding: 12px;
+        margin-top: 8px;
+        margin-bottom: 8px;
     }
 
     /* Cloudflare tunnel spinner */
@@ -114,20 +133,21 @@ export const UI_STYLES = {
     statusDot: "display: inline-block; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;",
     controlsDiv: "padding: 0 12px 12px 12px; display: flex; gap: 6px;",
     formGroup: "display: flex; flex-direction: column; gap: 5px;",
-    formLabel: "font-size: 12px; color: #ccc; font-weight: 500;",
-    formInput: "padding: 6px 10px; background: #2a2a2a; border: 1px solid #444; color: white; font-size: 12px; border-radius: 4px; transition: border-color 0.2s;",
+    formLabel: "font-size: 12px; color: var(--dist-label-text, #ccc); font-weight: 500;",
+    formInput:
+        "padding: 6px 10px; color: var(--dist-input-text, white); background: var(--dist-input-bg, transparent); font-size: 12px; transition: border-color 0.2s;",
     
     // Card styles
-    cardBase: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex;",
-    workerCard: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex; background: #2a2a2a;",
-    cardBlueprint: "border: 2px dashed #555; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);",
-    cardAdd: "border: 1px dashed #444; cursor: pointer; transition: all 0.2s ease; background: transparent;",
-    
+    cardBase: "margin-bottom: 12px; overflow: hidden; display: flex;",
+    workerCard: "margin-bottom: 12px; overflow: hidden; display: flex;",
+    cardBlueprint: "cursor: pointer; transition: all 0.2s ease;",
+    cardAdd: "cursor: pointer; transition: all 0.2s ease;",
+
     // Column styles
     columnBase: "display: flex; align-items: center; justify-content: center;",
-    checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #3a3a3a; cursor: default; background: rgba(0,0,0,0.1);",
+    checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; cursor: default;",
     contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;",
-    iconColumn: "width: 44px; flex-shrink: 0; font-size: 20px; color: #666;",
+    iconColumn: "width: 44px; flex-shrink: 0; font-size: 20px; color: var(--dist-placeholder-add-color, #666);",
     
     // Row and content styles
     infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;",
@@ -137,14 +157,16 @@ export const UI_STYLES = {
     buttonGroup: "display: flex; gap: 4px; margin-top: 10px;",
     settingsForm: "display: flex; flex-direction: column; gap: 10px;",
     checkboxGroup: "display: flex; align-items: center; gap: 8px; margin: 5px 0;",
-    formLabelClickable: "font-size: 12px; color: #ccc; cursor: pointer;",
+    formLabelClickable: "font-size: 12px; color: var(--dist-label-text, #ccc); cursor: pointer;",
     settingsToggle: "display: flex; align-items: center; gap: 6px; padding: 4px 0; cursor: pointer; user-select: none;",
     controlsWrapper: "display: flex; gap: 6px; align-items: stretch; width: 100%;",
     
     // Existing styles
-    settingsArrow: "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto; padding: 4px;",
-    infoBox: "background-color: #333; color: #999; padding: 5px 14px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1; font-weight: 500;",
-    workerSettings: "margin: 0 12px; padding: 0 12px; background: #1e1e1e; border-radius: 4px; border: 1px solid #2a2a2a;"
+    settingsArrow:
+        "font-size: 12px; color: var(--dist-settings-arrow, #888); transition: all 0.2s ease; margin-left: auto; padding: 4px;",
+    infoBox:
+        "color: var(--dist-info-box-text, #999); padding: 5px 14px; font-size: 11px; text-align: center; flex: 1; font-weight: 500;",
+    workerSettings: "margin: 0 12px; padding: 0 12px;"
 };
 
 export const TIMEOUTS = {
@@ -168,3 +190,50 @@ export const TIMEOUTS = {
     LOG_REFRESH: 2000, // log auto-refresh interval
     IMAGE_CACHE_CLEAR: 30000 // delay before clearing image cache
 };
+
+export const ENDPOINTS = {
+    // ComfyUI core
+    PROMPT: '/prompt',
+    INTERRUPT: '/interrupt',
+    UPLOAD_IMAGE: '/upload/image',
+    SYSTEM_INFO: '/system_stats',
+
+    // Distributed API
+    CONFIG: '/distributed/config',
+    UPDATE_WORKER: '/distributed/config/update_worker',
+    DELETE_WORKER: '/distributed/config/delete_worker',
+    UPDATE_SETTING: '/distributed/config/update_setting',
+    UPDATE_MASTER: '/distributed/config/update_master',
+    LAUNCH_WORKER: '/distributed/launch_worker',
+    STOP_WORKER: '/distributed/stop_worker',
+    MANAGED_WORKERS: '/distributed/managed_workers',
+    WORKER_LOG: '/distributed/worker_log',
+    REMOTE_WORKER_LOG: '/distributed/remote_worker_log',
+    LOCAL_LOG: '/distributed/local_log',
+    CLEAR_LAUNCHING: '/distributed/worker/clear_launching',
+    PREPARE_JOB: '/distributed/prepare_job',
+    LOAD_IMAGE: '/distributed/load_image',
+    NETWORK_INFO: '/distributed/network_info',
+    CHECK_FILE: '/distributed/check_file',
+    CLEAR_MEMORY: '/distributed/clear_memory',
+    SYSTEM_INFO_DIST: '/distributed/system_info',
+    TUNNEL_START: '/distributed/tunnel/start',
+    TUNNEL_STOP: '/distributed/tunnel/stop',
+    TUNNEL_STATUS: '/distributed/tunnel/status',
+};
+
+export const NODE_CLASSES = {
+    DISTRIBUTED_COLLECTOR: 'DistributedCollector',
+    DISTRIBUTED_SEED: 'DistributedSeed',
+    DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage',
+    UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed',
+    PREVIEW_IMAGE: 'PreviewImage',
+};
+
+export function generateUUID() {
+    if (crypto.randomUUID) return crypto.randomUUID();
+    return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, c => {
+        const r = Math.random() * 16 | 0;
+        return (c === 'x' ? r : (r & 0x3 | 0x8)).toString(16);
+    });
+}

web/distributed.css

@@ -0,0 +1,416 @@
+:root {
+    --btn-stop: #7c4a4a;
+    --btn-launch: #4a7c4a;
+    --btn-log: #685434;
+    --btn-working: #666;
+    --btn-success: #3a6a3a;
+    --btn-error: #6a3a3a;
+    --tunnel-enable: #665533;
+    --tunnel-disable: #7c4a4a;
+    --master-badge-fallback-bg: #243024;
+    --master-badge-fallback-text: #6bd06b;
+    --master-badge-fallback-border: #335533;
+    --master-badge-delegate-bg: #3a3a3a;
+    --master-badge-delegate-text: #ffcc66;
+    --dist-divider: #444;
+    --dist-muted-text: #888;
+    --dist-label-text: #ccc;
+    --dist-settings-arrow: #888;
+    --dist-settings-arrow-hover: #fff;
+    --dist-card-bg: #2a2a2a;
+    --dist-card-title: #f4f5f7;
+    --dist-card-subtext: #a9afb9;
+    --dist-card-placeholder-title: #aaa;
+    --dist-left-col-border: #3a3a3a;
+    --dist-left-col-bg: rgba(0, 0, 0, 0.1);
+    --dist-info-box-bg: #333;
+    --dist-info-box-text: #999;
+    --dist-input-bg: #2a2a2a;
+    --dist-input-border: #444;
+    --dist-input-text: #fff;
+    --dist-settings-bg: #1e1e1e;
+    --dist-settings-border: #2a2a2a;
+    --dist-hover-bg: #333;
+    --dist-placeholder-blueprint-border: #555;
+    --dist-placeholder-blueprint-hover-border: #777;
+    --dist-placeholder-add-border: #444;
+    --dist-placeholder-add-hover-border: #666;
+    --dist-placeholder-blueprint-bg: rgba(255, 255, 255, 0.02);
+    --dist-placeholder-blueprint-hover-bg: rgba(255, 255, 255, 0.05);
+    --dist-placeholder-add-hover-bg: rgba(255, 255, 255, 0.02);
+    --dist-placeholder-blueprint-color: #777;
+    --dist-placeholder-blueprint-hover-color: #999;
+    --dist-placeholder-add-color: #555;
+    --dist-placeholder-add-hover-color: #888;
+    --dist-log-modal-bg: #1e1e1e;
+    --dist-log-modal-border: #444;
+    --dist-log-modal-header-border: #444;
+    --dist-log-modal-title: #fff;
+    --dist-log-modal-label: #ccc;
+    --dist-log-modal-body-bg: #0d0d0d;
+    --dist-log-modal-body-text: #ddd;
+    --dist-log-modal-status: #888;
+}
+
+.distributed-panel--light {
+    --dist-divider: #c7ced9;
+    --dist-muted-text: #5b6472;
+    --dist-label-text: #2f3a4a;
+    --dist-settings-arrow: #6b7483;
+    --dist-settings-arrow-hover: #253040;
+    --dist-card-bg: #f6f8fb;
+    --dist-card-title: #1e293b;
+    --dist-card-subtext: #4b5565;
+    --dist-card-placeholder-title: #3f4a5a;
+    --dist-left-col-border: #d8deea;
+    --dist-left-col-bg: #eef2f8;
+    --dist-info-box-bg: #e8edf5;
+    --dist-info-box-text: #4f5b6c;
+    --dist-input-bg: #ffffff;
+    --dist-input-border: #b8c2d3;
+    --dist-input-text: #1f2937;
+    --dist-settings-bg: #eef2f8;
+    --dist-settings-border: #d5dcea;
+    --dist-hover-bg: #e6ebf3;
+    --dist-placeholder-blueprint-border: #aeb8c9;
+    --dist-placeholder-blueprint-hover-border: #8e9ab0;
+    --dist-placeholder-add-border: #b7c1d3;
+    --dist-placeholder-add-hover-border: #919eb4;
+    --dist-placeholder-blueprint-bg: rgba(80, 100, 140, 0.05);
+    --dist-placeholder-blueprint-hover-bg: rgba(80, 100, 140, 0.08);
+    --dist-placeholder-add-hover-bg: rgba(80, 100, 140, 0.05);
+    --dist-placeholder-blueprint-color: #5f6c82;
+    --dist-placeholder-blueprint-hover-color: #47556e;
+    --dist-placeholder-add-color: #6a7588;
+    --dist-placeholder-add-hover-color: #4c586f;
+    --master-badge-delegate-bg: #efe7ce;
+    --master-badge-delegate-text: #6a4f00;
+    --master-badge-fallback-bg: #dff1df;
+    --master-badge-fallback-text: #1f5f1f;
+    --master-badge-fallback-border: #9fc79f;
+    --dist-log-modal-bg: #f9fbff;
+    --dist-log-modal-border: #bcc7db;
+    --dist-log-modal-header-border: #c7d1e3;
+    --dist-log-modal-title: #1f2937;
+    --dist-log-modal-label: #4a5568;
+    --dist-log-modal-body-bg: #f2f5fb;
+    --dist-log-modal-body-text: #1f2937;
+    --dist-log-modal-status: #5c6779;
+}
+
+.btn--stop {
+    background-color: var(--btn-stop) !important;
+}
+
+.btn--launch {
+    background-color: var(--btn-launch) !important;
+}
+
+.btn--log {
+    background-color: var(--btn-log) !important;
+}
+
+.btn--working {
+    background-color: var(--btn-working) !important;
+}
+
+.btn--success {
+    background-color: var(--btn-success) !important;
+}
+
+.btn--error {
+    background-color: var(--btn-error) !important;
+}
+
+.master-info-badge--fallback {
+    background-color: var(--master-badge-fallback-bg) !important;
+    color: var(--master-badge-fallback-text) !important;
+    border: 1px solid var(--master-badge-fallback-border) !important;
+}
+
+.master-info-badge--delegate {
+    background-color: var(--master-badge-delegate-bg) !important;
+    color: var(--master-badge-delegate-text) !important;
+}
+
+.entity-card-content--hovered {
+    background-color: var(--dist-hover-bg) !important;
+}
+
+.placeholder-card--blueprint {
+    border-color: var(--dist-placeholder-blueprint-border) !important;
+    background-color: var(--dist-placeholder-blueprint-bg) !important;
+}
+
+.placeholder-card--blueprint.is-hovered {
+    border-color: var(--dist-placeholder-blueprint-hover-border) !important;
+    background-color: var(--dist-placeholder-blueprint-hover-bg) !important;
+}
+
+.placeholder-card--add {
+    border-color: var(--dist-placeholder-add-border) !important;
+    background-color: transparent !important;
+}
+
+.placeholder-card--add.is-hovered {
+    border-color: var(--dist-placeholder-add-hover-border) !important;
+    background-color: var(--dist-placeholder-add-hover-bg) !important;
+}
+
+.placeholder-column--blueprint {
+    color: var(--dist-placeholder-blueprint-color) !important;
+    border-right-color: var(--dist-placeholder-blueprint-border) !important;
+}
+
+.placeholder-column--blueprint.is-hovered {
+    color: var(--dist-placeholder-blueprint-hover-color) !important;
+}
+
+.placeholder-column--add {
+    color: var(--dist-placeholder-add-color) !important;
+    border-color: var(--dist-placeholder-add-border) !important;
+    border-right-color: var(--dist-placeholder-add-border) !important;
+}
+
+.placeholder-column--add.is-hovered {
+    color: var(--dist-placeholder-add-hover-color) !important;
+    border-color: var(--dist-placeholder-add-hover-border) !important;
+}
+
+.tunnel-button--enable {
+    background-color: var(--tunnel-enable) !important;
+}
+
+.tunnel-button--disable {
+    background-color: var(--tunnel-disable) !important;
+}
+
+.tunnel-status--enable {
+    background-color: var(--tunnel-enable) !important;
+}
+
+.tunnel-status--disable {
+    background-color: var(--tunnel-disable) !important;
+}
+
+/* ---- Themeable card/column/input defaults (Classic) ---- */
+
+.dist-card {
+    background: var(--dist-card-bg);
+    color: var(--dist-card-title);
+    border-radius: 6px;
+}
+
+.dist-card--blueprint {
+    border: 2px dashed var(--dist-placeholder-blueprint-border);
+    background: var(--dist-placeholder-blueprint-bg);
+}
+
+.dist-card--add {
+    border: 1px dashed var(--dist-placeholder-add-border);
+    background: transparent;
+}
+
+.dist-card__left-col {
+    border-right: 1px solid var(--dist-left-col-border);
+    background: var(--dist-left-col-bg);
+}
+
+.dist-info-box {
+    background-color: var(--dist-info-box-bg);
+    color: var(--dist-info-box-text);
+    border-radius: 4px;
+}
+
+.dist-form-input {
+    background: var(--dist-input-bg);
+    border: 1px solid var(--dist-input-border);
+    color: var(--dist-input-text);
+    border-radius: 4px;
+}
+
+.worker-settings {
+    background: var(--dist-settings-bg);
+    border: 1px solid var(--dist-settings-border);
+    border-radius: 4px;
+}
+
+.dist-worker-info__title {
+    color: var(--dist-card-title) !important;
+    font-size: 1.03em;
+    font-weight: 700;
+}
+
+.dist-worker-info__meta {
+    color: var(--dist-card-subtext) !important;
+    font-size: 0.88em;
+}
+
+.dist-worker-info__fallback {
+    color: var(--master-badge-fallback-text) !important;
+    font-size: 0.86em;
+    font-weight: 600;
+}
+
+/* ---- Nodes 2.0 Theme ---- */
+
+.distributed-panel--nodes2 .dist-card {
+    background: var(--p-surface-800, #1e1e1e);
+    border: 1px solid var(--p-surface-700, #2c2c2c);
+    border-radius: 8px;
+    box-shadow: 0 1px 3px rgba(0, 0, 0, 0.3);
+}
+
+.distributed-panel--nodes2 .dist-card--blueprint {
+    border: 2px dashed var(--p-surface-600, #404040);
+    background: rgba(255, 255, 255, 0.015);
+    box-shadow: none;
+}
+
+.distributed-panel--nodes2 .dist-card--add {
+    border: 1px dashed var(--p-surface-600, #404040);
+    background: transparent;
+    box-shadow: none;
+}
+
+.distributed-panel--nodes2 .dist-card__left-col {
+    border-right-color: var(--p-surface-700, #2c2c2c);
+    background: var(--p-surface-900, #131313);
+}
+
+.distributed-panel--nodes2 .dist-info-box {
+    background-color: var(--p-surface-700, #2c2c2c);
+    color: var(--p-text-muted-color, #9e9e9e);
+    border-radius: 6px;
+}
+
+.distributed-panel--nodes2 .dist-form-input {
+    background: var(--p-surface-800, #1e1e1e);
+    border-color: var(--p-surface-600, #404040);
+    border-radius: 6px;
+    color: var(--p-text-color, #ffffff);
+}
+
+.distributed-panel--nodes2 .worker-settings {
+    background: var(--p-surface-900, #131313);
+    border-color: var(--p-surface-700, #2c2c2c);
+    border-radius: 6px;
+}
+
+.distributed-panel--nodes2 .entity-card-content--hovered {
+    background-color: var(--p-surface-700, #2c2c2c) !important;
+}
+
+.distributed-panel--nodes2 .placeholder-card--blueprint {
+    border-color: var(--p-surface-600, #404040) !important;
+}
+
+.distributed-panel--nodes2 .placeholder-card--blueprint.is-hovered {
+    border-color: var(--p-surface-500, #555555) !important;
+}
+
+.distributed-panel--nodes2 .placeholder-card--add {
+    border-color: var(--p-surface-600, #404040) !important;
+}
+
+.distributed-panel--nodes2 .placeholder-card--add.is-hovered {
+    border-color: var(--p-surface-500, #555555) !important;
+}
+
+/* ---- End Nodes 2.0 Theme ---- */
+
+.is-hidden {
+    display: none !important;
+}
+
+.settings-arrow--expanded {
+    transform: rotate(90deg) !important;
+}
+
+.log-modal {
+    position: fixed;
+    top: 0;
+    left: 0;
+    width: 100%;
+    height: 100%;
+    background: rgba(0, 0, 0, 0.8);
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    z-index: 10000;
+}
+
+.log-modal__content {
+    background: var(--dist-log-modal-bg);
+    border-radius: 8px;
+    width: 90%;
+    max-width: 1200px;
+    height: 80%;
+    display: flex;
+    flex-direction: column;
+    border: 1px solid var(--dist-log-modal-border);
+}
+
+.log-modal__header {
+    padding: 15px 20px;
+    border-bottom: 1px solid var(--dist-log-modal-header-border);
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+}
+
+.log-modal__title {
+    margin: 0;
+    color: var(--dist-log-modal-title);
+}
+
+.log-modal__header-buttons {
+    display: flex;
+    gap: 20px;
+    align-items: center;
+}
+
+.log-modal__refresh {
+    display: flex;
+    align-items: center;
+    gap: 4px;
+}
+
+.log-modal__refresh-input {
+    cursor: pointer;
+}
+
+.log-modal__refresh-label {
+    font-size: 12px;
+    color: var(--dist-log-modal-label);
+    cursor: pointer;
+    white-space: nowrap;
+}
+
+.log-modal__close {
+    background-color: #c04c4c !important;
+    padding: 5px 10px !important;
+    font-size: 14px !important;
+    font-weight: bold !important;
+    border-radius: 6px !important;
+}
+
+.log-modal__body {
+    flex: 1;
+    overflow: auto;
+    padding: 15px;
+    font-family: 'Consolas', 'Monaco', 'Courier New', monospace;
+    font-size: 12px;
+    line-height: 1.4;
+    color: var(--dist-log-modal-body-text);
+    background: var(--dist-log-modal-body-bg);
+    white-space: pre-wrap;
+    word-wrap: break-word;
+}
+
+.log-modal__status {
+    padding: 10px 20px;
+    border-top: 1px solid var(--dist-log-modal-header-border);
+    font-size: 11px;
+    color: var(--dist-log-modal-status);
+}

web/distributedValue.js

@@ -0,0 +1,478 @@
+import { app } from "/scripts/app.js";
+import { ENDPOINTS } from "./constants.js";
+
+const NODE_CLASS = "DistributedValue";
+const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
+const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
+const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
+
+const trackedNodes = new Set();
+let workersChangedListenerAttached = false;
+
+function filterEnabledWorkers(workers) {
+    if (!Array.isArray(workers)) return [];
+    return workers.filter((worker) => Boolean(worker?.enabled));
+}
+
+async function fetchWorkers() {
+    try {
+        const resp = await fetch(ENDPOINTS.CONFIG);
+        if (!resp.ok) return [];
+        const config = await resp.json();
+        return filterEnabledWorkers(config.workers);
+    } catch {
+        return [];
+    }
+}
+
+function getRawDefaultWidget(node) {
+    return node.widgets?.find((w) => w.name === "default_value");
+}
+
+function getRawWorkerValuesWidget(node) {
+    return node.widgets?.find((w) => w.name === "worker_values");
+}
+
+function getDynamicDefaultWidget(node) {
+    return node.widgets?.find((w) => w.name === DYNAMIC_DEFAULT_WIDGET);
+}
+
+function getDynamicWorkerWidgets(node) {
+    return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX));
+}
+
+function hideWidgetForGood(widget) {
+    if (!widget) return;
+    // These are serialized backing fields, not widgets converted into sockets.
+    // Keep their native type/serializer and mutate the shared options in place
+    // so both the classic canvas and Nodes 2.0 suppress the whole widget row.
+    widget.options ??= {};
+    widget.options.hidden = true;
+
+    // Hide any attached DOM element (multiline widgets).
+    if (widget.element) widget.element.style.display = "none";
+    if (widget.inputEl) widget.inputEl.style.display = "none";
+
+    if (widget.linkedWidgets) {
+        for (const linked of widget.linkedWidgets) {
+            hideWidgetForGood(linked);
+        }
+    }
+}
+
+function hideRawWidgets(node) {
+    hideWidgetForGood(getRawDefaultWidget(node));
+    hideWidgetForGood(getRawWorkerValuesWidget(node));
+}
+
+function removeDynamicDefaultWidget(node) {
+    const idx = node.widgets?.findIndex((w) => w.name === DYNAMIC_DEFAULT_WIDGET);
+    if (idx != null && idx >= 0) {
+        node.widgets.splice(idx, 1);
+    }
+}
+
+function removeDynamicWorkerWidgets(node) {
+    if (!node.widgets) return;
+    for (let i = node.widgets.length - 1; i >= 0; i--) {
+        if (node.widgets[i].name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX)) {
+            node.widgets.splice(i, 1);
+        }
+    }
+}
+
+function readWorkerStore(node) {
+    const raw = getRawWorkerValuesWidget(node);
+    if (!raw) return {};
+    try {
+        const parsed = JSON.parse(raw.value || "{}");
+        return typeof parsed === "object" && parsed !== null ? parsed : {};
+    } catch {
+        return {};
+    }
+}
+
+function writeWorkerStore(node, store) {
+    const raw = getRawWorkerValuesWidget(node);
+    if (!raw) return;
+    raw.value = JSON.stringify(store);
+}
+
+function normalizeComboOptions(options) {
+    if (!options) return null;
+    if (Array.isArray(options)) return options;
+    if (Array.isArray(options.values)) return options.values;
+    return null;
+}
+
+function resolveGraphLink(graph, linkId) {
+    const links = graph.links || graph._links;
+    if (!links) return null;
+    const link = links[linkId] ?? (typeof links.get === "function" ? links.get(linkId) : null);
+    if (!link) return null;
+    if (Array.isArray(link)) {
+        return {
+            target_id: link[2],
+            target_slot: link[3],
+        };
+    }
+    return link;
+}
+
+function detectTargetType(node) {
+    const out = node.outputs?.[0];
+    const linkIds = out?.links || [];
+    if (!linkIds.length) {
+        return { connected: false, type: "STRING", options: null };
+    }
+
+    const graph = node.graph || app.graph;
+    if (!graph) {
+        return { connected: false, type: "STRING", options: null };
+    }
+
+    const link = resolveGraphLink(graph, linkIds[0]);
+    if (!link) {
+        return { connected: false, type: "STRING", options: null };
+    }
+
+    const targetNode = graph.getNodeById(link.target_id);
+    if (!targetNode) {
+        return { connected: false, type: "STRING", options: null };
+    }
+
+    const targetInputName = targetNode.inputs?.[link.target_slot]?.name;
+    if (!targetInputName) {
+        return { connected: false, type: "STRING", options: null };
+    }
+
+    const targetWidget = targetNode.widgets?.find((w) => w.name === targetInputName);
+    if (targetWidget) {
+        if (targetWidget.type === "combo") {
+            const comboOptions = normalizeComboOptions(targetWidget.options);
+            return { connected: true, type: "COMBO", options: comboOptions };
+        }
+        if (targetWidget.type === "number") {
+            const step = targetWidget.options?.step;
+            const precision = targetWidget.options?.precision;
+            const isInt = Number.isInteger(step) && (precision === 0 || precision == null);
+            return { connected: true, type: isInt ? "INT" : "FLOAT", options: null };
+        }
+    }
+
+    const nodeDef = targetNode.constructor?.nodeData;
+    const inputDef = nodeDef?.input?.required?.[targetInputName] || nodeDef?.input?.optional?.[targetInputName];
+    if (inputDef) {
+        const defType = inputDef[0];
+        if (Array.isArray(defType)) {
+            return { connected: true, type: "COMBO", options: defType };
+        }
+        if (defType === "INT") return { connected: true, type: "INT", options: null };
+        if (defType === "FLOAT") return { connected: true, type: "FLOAT", options: null };
+    }
+
+    return { connected: true, type: "STRING", options: null };
+}
+
+function normalizeNumber(value, fallback) {
+    const parsed = Number(value);
+    return Number.isFinite(parsed) ? parsed : fallback;
+}
+
+function getDefaultInitialValue(node, inputType, comboOptions) {
+    const rawDefault = getRawDefaultWidget(node);
+    const current = rawDefault?.value;
+
+    if (inputType === "INT") {
+        return Math.trunc(normalizeNumber(current, 0));
+    }
+    if (inputType === "FLOAT") {
+        return normalizeNumber(current, 0);
+    }
+    if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
+        const currentText = current == null ? "" : String(current);
+        return comboOptions.includes(currentText) ? currentText : comboOptions[0];
+    }
+    return current == null ? "" : String(current);
+}
+
+function setRawDefaultValue(node, value) {
+    const rawDefault = getRawDefaultWidget(node);
+    if (!rawDefault) return;
+    rawDefault.value = value;
+}
+
+function serializeWorkerStoreFromWidgets(node, inputType, comboOptions) {
+    const nextStore = { _type: inputType };
+    if (inputType === "COMBO" && Array.isArray(comboOptions)) {
+        nextStore._options = comboOptions;
+    }
+    const valuesByWorkerId = {};
+
+    for (const widget of getDynamicWorkerWidgets(node)) {
+        const key = widget.name.slice(DYNAMIC_WORKER_WIDGET_PREFIX.length);
+        if (widget.value !== "" && widget.value !== null && widget.value !== undefined) {
+            const value = String(widget.value);
+            nextStore[key] = value;
+            if (widget._dvWorkerId) {
+                valuesByWorkerId[widget._dvWorkerId] = value;
+            }
+        }
+    }
+    if (Object.keys(valuesByWorkerId).length) {
+        nextStore._by_worker_id = valuesByWorkerId;
+    }
+
+    writeWorkerStore(node, nextStore);
+}
+
+function updateWorkerStoreTypeMetadata(node, inputType, comboOptions) {
+    const store = readWorkerStore(node);
+    store._type = inputType;
+    if (inputType === "COMBO" && Array.isArray(comboOptions)) {
+        store._options = comboOptions;
+    } else {
+        delete store._options;
+    }
+    writeWorkerStore(node, store);
+}
+
+function createDynamicDefaultWidget(node, inputType, comboOptions) {
+    removeDynamicDefaultWidget(node);
+    const initial = getDefaultInitialValue(node, inputType, comboOptions);
+    let widget;
+
+    if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
+        widget = node.addWidget(
+            "combo",
+            DYNAMIC_DEFAULT_WIDGET,
+            initial,
+            (value) => {
+                widget.value = value;
+                setRawDefaultValue(node, String(value));
+            },
+            { values: comboOptions }
+        );
+    } else if (inputType === "INT") {
+        widget = node.addWidget(
+            "number",
+            DYNAMIC_DEFAULT_WIDGET,
+            initial,
+            (value) => {
+                widget.value = Math.trunc(normalizeNumber(value, 0));
+                setRawDefaultValue(node, widget.value);
+            },
+            { min: -Infinity, max: Infinity, step: 1, precision: 0 }
+        );
+    } else if (inputType === "FLOAT") {
+        widget = node.addWidget(
+            "number",
+            DYNAMIC_DEFAULT_WIDGET,
+            initial,
+            (value) => {
+                widget.value = normalizeNumber(value, 0);
+                setRawDefaultValue(node, widget.value);
+            },
+            { min: -Infinity, max: Infinity, step: 0.1, precision: 3 }
+        );
+    } else {
+        widget = node.addWidget(
+            "string",
+            DYNAMIC_DEFAULT_WIDGET,
+            initial,
+            (value) => {
+                widget.value = value ?? "";
+                setRawDefaultValue(node, widget.value);
+            },
+            {}
+        );
+    }
+
+    widget.label = "default_value";
+}
+
+function getWorkerInitialValue(store, key, workerId, inputType, comboOptions) {
+    const byWorkerId = store?._by_worker_id;
+    const saved = (byWorkerId && workerId && byWorkerId[workerId] != null)
+        ? byWorkerId[workerId]
+        : store[key];
+    if (saved == null) {
+        if (inputType === "INT" || inputType === "FLOAT") return 0;
+        if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
+            return comboOptions[0];
+        }
+        return "";
+    }
+
+    if (inputType === "INT") return Math.trunc(normalizeNumber(saved, 0));
+    if (inputType === "FLOAT") return normalizeNumber(saved, 0);
+    if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
+        const savedText = String(saved);
+        return comboOptions.includes(savedText) ? savedText : comboOptions[0];
+    }
+    return String(saved);
+}
+
+function createWorkerWidgets(node, workers, inputType, comboOptions) {
+    removeDynamicWorkerWidgets(node);
+    const store = readWorkerStore(node);
+
+    for (let i = 0; i < workers.length; i++) {
+        const key = String(i + 1);
+        const worker = workers[i];
+        const label = worker.name || worker.id || `Worker ${key}`;
+        const widgetName = `${DYNAMIC_WORKER_WIDGET_PREFIX}${key}`;
+        const initial = getWorkerInitialValue(store, key, worker.id, inputType, comboOptions);
+        let widget;
+
+        if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
+            widget = node.addWidget(
+                "combo",
+                widgetName,
+                initial,
+                (value) => {
+                    widget.value = value;
+                    serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
+                },
+                { values: comboOptions }
+            );
+        } else if (inputType === "INT") {
+            widget = node.addWidget(
+                "number",
+                widgetName,
+                initial,
+                (value) => {
+                    widget.value = Math.trunc(normalizeNumber(value, 0));
+                    serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
+                },
+                { min: -Infinity, max: Infinity, step: 1, precision: 0 }
+            );
+        } else if (inputType === "FLOAT") {
+            widget = node.addWidget(
+                "number",
+                widgetName,
+                initial,
+                (value) => {
+                    widget.value = normalizeNumber(value, 0);
+                    serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
+                },
+                { min: -Infinity, max: Infinity, step: 0.1, precision: 3 }
+            );
+        } else {
+            widget = node.addWidget(
+                "string",
+                widgetName,
+                initial,
+                (value) => {
+                    widget.value = value ?? "";
+                    serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
+                },
+                {}
+            );
+        }
+
+        widget.label = label;
+        widget._dvWorkerId = worker.id;
+    }
+
+    serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
+}
+
+function rebuildWidgets(node) {
+    hideRawWidgets(node);
+    const workers = node._dvWorkers || [];
+    const store = readWorkerStore(node);
+    const detected = detectTargetType(node);
+    const disconnected = !detected.connected;
+    const inputType = disconnected ? "STRING" : detected.type;
+    const comboOptions = disconnected ? null : detected.options;
+
+    if (disconnected) {
+        // Reset disconnected node back to the neutral default state.
+        setRawDefaultValue(node, "");
+        writeWorkerStore(node, { _type: "STRING" });
+    }
+
+    createDynamicDefaultWidget(node, inputType, comboOptions);
+    if (workers.length > 0) {
+        createWorkerWidgets(node, workers, inputType, comboOptions);
+    } else {
+        removeDynamicWorkerWidgets(node);
+        updateWorkerStoreTypeMetadata(node, inputType, comboOptions);
+    }
+
+    const size = node.computeSize();
+    size[0] = Math.max(size[0], node.size?.[0] || 0);
+    node.setSize(size);
+    if (node.setDirtyCanvas) node.setDirtyCanvas(true, true);
+}
+
+function refreshNodeWorkers(node, workers) {
+    if (!node || !node.graph) return;
+    node._dvWorkers = workers;
+    rebuildWidgets(node);
+}
+
+async function refreshTrackedNodes(workers = null) {
+    const nextWorkers = workers || (await fetchWorkers());
+    for (const node of trackedNodes) {
+        refreshNodeWorkers(node, nextWorkers);
+    }
+}
+
+function attachWorkersChangedListener() {
+    if (workersChangedListenerAttached) return;
+    if (typeof window === "undefined" || typeof window.addEventListener !== "function") return;
+
+    window.addEventListener(WORKERS_CHANGED_EVENT, (event) => {
+        const changedWorkers = filterEnabledWorkers(event?.detail?.workers);
+        if (changedWorkers.length > 0 || Array.isArray(event?.detail?.workers)) {
+            void refreshTrackedNodes(changedWorkers);
+            return;
+        }
+        void refreshTrackedNodes();
+    });
+
+    workersChangedListenerAttached = true;
+}
+
+app.registerExtension({
+    name: "Distributed.DistributedValue",
+    async nodeCreated(node) {
+        if (node.comfyClass !== NODE_CLASS) return;
+
+        try {
+            attachWorkersChangedListener();
+            trackedNodes.add(node);
+            node._dvWorkers = await fetchWorkers();
+            rebuildWidgets(node);
+
+            const originalOnConnectionsChange = node.onConnectionsChange;
+            node.onConnectionsChange = function (type, index, connected, linkInfo, ioSlot) {
+                if (originalOnConnectionsChange) {
+                    originalOnConnectionsChange.call(this, type, index, connected, linkInfo, ioSlot);
+                }
+                if (type === 2 && index === 0) {
+                    setTimeout(() => rebuildWidgets(this), 20);
+                }
+            };
+
+            const originalConfigure = node.configure;
+            node.configure = function (data) {
+                const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
+                setTimeout(() => rebuildWidgets(this), 20);
+                return result;
+            };
+
+            const originalOnRemoved = node.onRemoved;
+            node.onRemoved = function () {
+                trackedNodes.delete(this);
+                if (originalOnRemoved) {
+                    return originalOnRemoved.call(this);
+                }
+            };
+        } catch (error) {
+            console.error("Error in DistributedValue extension:", error);
+        }
+    },
+});

web/executionUtils.js

@@ -1,92 +1,32 @@
 import { api } from "../../scripts/api.js";
-import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo, findCollectorDownstreamNodes } from './workerUtils.js';
-import { TIMEOUTS } from './constants.js';
-
-/**
- * Convert paths in the API prompt to match the target platform's separator
- * @param {Object} apiPrompt - The workflow API prompt
- * @param {string} targetSeparator - The target path separator ('\\' or '/')
- * @returns {Object} The converted API prompt
- */
-function convertPathsForPlatform(apiPrompt, targetSeparator) {
-    // Validate target separator
-    if (!targetSeparator || (targetSeparator !== '/' && targetSeparator !== '\\')) {
-        console.warn('[Distributed] Invalid target separator:', targetSeparator, '- skipping path conversion');
-        return apiPrompt;
-    }
-    
-    // Regex to identify likely file paths with extensions
-    const isLikelyFilename = (value) => {
-        return value.match(/\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)(\s*\[\w+\])?$/i);
-    };
-    const isImageOrVideo = (value) => {
-        return value.match(/\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i);
-    };
-    
-    function convert(obj) {
-        if (typeof obj === 'string') {
-            // Only convert strings that look like file paths
-            if ((obj.includes('\\') || obj.includes('/')) && isLikelyFilename(obj)) {
-                const trimmed = obj.trim();
-                const hasDrive = /^[A-Za-z]:\\\\|^[A-Za-z]:\//.test(trimmed);
-                const isAbsolute = trimmed.startsWith('/') || trimmed.startsWith('\\\\');
-                const hasProtocol = /^\w+:\/\//.test(trimmed);
-
-                // For annotated relative image/video paths, keep forward slashes
-                if (!hasDrive && !isAbsolute && !hasProtocol && isImageOrVideo(trimmed)) {
-                    return trimmed.replace(/[\\\\]/g, '/');
-                }
-                // Otherwise replace any path separator with the worker's target separator
-                return trimmed.replace(/[\\\\\/]/g, targetSeparator);
-            }
-            return obj;
-        } else if (Array.isArray(obj)) {
-            return obj.map(convert);
-        } else if (typeof obj === 'object' && obj !== null) {
-            const newObj = {};
-            for (const [key, value] of Object.entries(obj)) {
-                newObj[key] = convert(value);
-            }
-            return newObj;
-        }
-        return obj;
-    }
-    
-    return convert(apiPrompt);
-}
+import { applyProbeResultToWorkerDot, findNodesByClass } from './workerUtils.js';
+import { TIMEOUTS, NODE_CLASSES, generateUUID } from './constants.js';
+import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js';
 
 export function setupInterceptor(extension) {
-    api.queuePrompt = async (number, prompt) => {
+    api.queuePrompt = async (number, prompt, ...rest) => {
         if (extension.isEnabled) {
-            const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0;
-            const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0;
-            const hasDistQueue = findNodesByClass(prompt.output, "DistributedQueue").length > 0;
+            const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
+            const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
 
             if (hasCollector || hasDistUpscale) {
                 const result = await executeParallelDistributed(extension, prompt);
                 // Immediate status check for instant feedback
-                extension.checkAllWorkerStatuses();
+                checkAllWorkerStatuses(extension);
                 // Another check after a short delay to catch state changes
-                setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
-                return result;
-            }
-
-            // DistributedQueue: route entire workflow to least-busy worker
-            if (hasDistQueue) {
-                const result = await executeQueueDistributed(extension, prompt);
-                extension.checkAllWorkerStatuses();
-                setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
+                setTimeout(() => checkAllWorkerStatuses(extension), TIMEOUTS.POST_ACTION_DELAY);
                 return result;
             }
         }
-        return extension.originalQueuePrompt(number, prompt);
+        return extension.originalQueuePrompt(number, prompt, ...rest);
     };
 }
 
 export async function executeParallelDistributed(extension, promptWrapper) {
+    const traceExecutionId = `exec_${Date.now()}_${generateUUID().slice(0, 6)}`;
     try {
-        const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution
         const enabledWorkers = extension.enabledWorkers;
+        extension.log(`[exec:${traceExecutionId}] Starting distributed execution`, "debug");
         
         // Pre-flight health check on all enabled workers
         const activeWorkers = await performPreflightCheck(extension, enabledWorkers);
@@ -136,725 +76,32 @@ export async function executeParallelDistributed(extension, promptWrapper) {
             }
         }
 
-        // Find all distributed nodes in the workflow
-        const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector");
-        const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed");
-        const allDistributedNodes = [...collectorNodes, ...upscaleNodes];
-
-        const wantsDelegate = Boolean(extension.config?.settings?.master_delegate_only);
-        let masterDelegateActive = false;
-        if (wantsDelegate) {
-            if (activeWorkers.length === 0) {
-                extension.log("Master delegate-only mode disabled: no active workers detected. Falling back to master execution.", "debug");
-            } else if (upscaleNodes.length > 0) {
-                extension.log("Master delegate-only mode is not yet supported for UltimateSDUpscaleDistributed nodes. Falling back to normal execution.", "warn");
-            } else if (collectorNodes.length === 0) {
-                extension.log("Master delegate-only mode requested but no DistributedCollector nodes found. Falling back to normal execution.", "debug");
-            } else {
-                masterDelegateActive = true;
-                extension.log("Master delegate-only mode enabled: master will orchestrate without running upstream nodes.", "debug");
-            }
-        }
-        
-        // Map original node IDs to truly unique job IDs for this specific run
-        const job_id_map = new Map(allDistributedNodes.map(node => [node.id, `${executionPrefix}_${node.id}`]));
-        
-        // Prepare a separate job queue on the backend for each unique job ID
-        const preparePromises = Array.from(job_id_map.values()).map(uniqueId => prepareDistributedJob(extension, uniqueId));
-        await Promise.all(preparePromises);
-
-        const jobs = [];
-        // Use only active workers
-        const participants = ['master', ...activeWorkers.map(w => w.id)];
-
-        for (const participantId of participants) {
-            const options = { 
-                enabled_worker_ids: activeWorkers.map(w => w.id), 
-                workflow: promptWrapper.workflow,
-                job_id_map: job_id_map, // Pass the map of unique IDs
-                delegate_master: masterDelegateActive && participantId === 'master'
-            };
-            
-            const jobApiPrompt = await prepareApiPromptForParticipant(
-                extension, promptWrapper.output, participantId, options
+        const queueResponse = await extension.api.queueDistributed({
+            prompt: promptWrapper.output,
+            workflow: promptWrapper.workflow,
+            enabled_worker_ids: activeWorkers.map((worker) => worker.id),
+            workers: activeWorkers.map((worker) => ({ id: worker.id })),
+            client_id: api.clientId,
+            delegate_master: Boolean(extension.config?.settings?.master_delegate_only),
+            auto_prepare: true,
+            trace_execution_id: traceExecutionId,
+        });
+        if (queueResponse?.prompt_id) {
+            extension.log(
+                `[exec:${traceExecutionId}] Distributed queue accepted by backend (prompt_id=${queueResponse.prompt_id}, workers=${queueResponse.worker_count ?? activeWorkers.length})`,
+                "debug"
             );
-            
-            if (participantId === 'master') {
-                jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } });
-            } else {
-                const worker = activeWorkers.find(w => w.id === participantId);
-                if (worker) {
-                    const job = {
-                        type: 'worker',
-                        worker,
-                        prompt: jobApiPrompt,
-                        workflow: promptWrapper.workflow
-                    };
-                    
-                    // Add image references if found for remote workers
-                    if (options._imageReferences) {
-                        job.imageReferences = options._imageReferences;
-                    }
-                    
-                    jobs.push(job);
-                }
-            }
+            return queueResponse;
         }
-        
-        const result = await executeJobs(extension, jobs);
-        return result;
+        throw new Error(
+            `[exec:${traceExecutionId}] Backend did not return a prompt_id for distributed queue.`
+        );
     } catch (error) {
-        extension.log("Parallel execution failed: " + error.message, "error");
-        throw error;
-    }
-}
-
-/**
- * Execute a workflow via DistributedQueue - routes to the least-busy worker.
- * The workflow runs ONLY on the selected worker, not on the master.
- */
-export async function executeQueueDistributed(extension, promptWrapper) {
-    try {
-        const enabledWorkers = extension.enabledWorkers;
-
-        if (enabledWorkers.length === 0) {
-            extension.log("DistributedQueue: No enabled workers. Falling back to local execution.", "warn");
-            return extension.originalQueuePrompt(0, promptWrapper);
-        }
-
-        // Fetch queue status from all enabled workers
-        const workerStatuses = await fetchWorkerQueueStatuses(extension, enabledWorkers);
-
-        if (workerStatuses.length === 0) {
-            extension.log("DistributedQueue: No reachable workers. Falling back to local execution.", "warn");
-            return extension.originalQueuePrompt(0, promptWrapper);
-        }
-
-        // Select the least-busy worker
-        const selectedWorker = selectLeastBusyWorker(extension, workerStatuses);
-        const worker = selectedWorker.worker;
-        const queueRemaining = selectedWorker.queueRemaining;
-
-        extension.log(`DistributedQueue: Routing to ${worker.name} (queue_remaining=${queueRemaining})`, "info");
-
-        // Prepare the prompt with skip_dispatch=True for DistributedQueue nodes
-        const promptToSend = JSON.parse(JSON.stringify(promptWrapper.output));
-        markSkipDispatch(promptToSend);
-
-        // Dispatch to the selected worker
-        const workerUrl = extension.getWorkerUrl(worker);
-        try {
-            const response = await fetch(`${workerUrl}/prompt`, {
-                method: 'POST',
-                headers: { 'Content-Type': 'application/json' },
-                mode: 'cors',
-                body: JSON.stringify({
-                    prompt: promptToSend,
-                    extra_data: { extra_pnginfo: { workflow: promptWrapper.workflow } },
-                    client_id: api.clientId
-                })
-            });
-
-            if (!response.ok) {
-                throw new Error(`Worker returned ${response.status}`);
-            }
-
-            const result = await response.json();
-            extension.log(`DistributedQueue: Dispatched to ${worker.name}, prompt_id=${result.prompt_id}`, "info");
-
-            return { prompt_id: result.prompt_id, worker_id: worker.id };
-        } catch (error) {
-            extension.log(`DistributedQueue: Failed to dispatch to ${worker.name}: ${error.message}`, "error");
-            // Fall back to local execution
-            return extension.originalQueuePrompt(0, promptWrapper);
-        }
-    } catch (error) {
-        extension.log("DistributedQueue execution failed: " + error.message, "error");
-        throw error;
-    }
-}
-
-async function fetchWorkerQueueStatuses(extension, workers) {
-    const statuses = [];
-
-    const checkPromises = workers.map(async (worker) => {
-        const url = extension.getWorkerUrl(worker, '/prompt');
-        try {
-            const response = await fetch(url, {
-                method: 'GET',
-                mode: 'cors',
-                cache: 'no-store',
-                signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
-            });
-
-            if (response.ok) {
-                const data = await response.json();
-                const queueRemaining = data.exec_info?.queue_remaining || 0;
-                return { worker, queueRemaining, online: true };
-            }
-        } catch (error) {
-            extension.log(`DistributedQueue: Worker ${worker.name} unreachable: ${error.message}`, "debug");
+        extension.log(`[exec:${traceExecutionId}] Distributed execution failed: ${error.message}`, "error");
+        if (extension.ui?.showToast) {
+            extension.ui.showToast(extension.app, "error", "Distributed Failed", error.message, 5000);
         }
         return null;
-    });
-
-    const results = await Promise.all(checkPromises);
-    return results.filter(r => r !== null);
-}
-
-function selectLeastBusyWorker(extension, statuses) {
-    // Find idle workers (queue_remaining == 0)
-    const idleWorkers = statuses.filter(s => s.queueRemaining === 0);
-
-    if (idleWorkers.length > 0) {
-        // Round-robin among idle workers
-        if (!extension._distributedQueueRRIndex) {
-            extension._distributedQueueRRIndex = 0;
-        }
-        const index = extension._distributedQueueRRIndex % idleWorkers.length;
-        extension._distributedQueueRRIndex++;
-        return idleWorkers[index];
-    }
-
-    // No idle workers - pick the one with the shortest queue
-    return statuses.reduce((min, s) => s.queueRemaining < min.queueRemaining ? s : min);
-}
-
-function markSkipDispatch(promptObj) {
-    for (const node of Object.values(promptObj)) {
-        if (node && typeof node === 'object' && node.class_type === 'DistributedQueue') {
-            node.inputs = node.inputs || {};
-            node.inputs.skip_dispatch = true;
-        }
-    }
-}
-
-export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) {
-    let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt));
-    const isMaster = participantId === 'master';
-    const delegateMaster = Boolean(options.delegate_master);
-    
-    // Find all distributed nodes once (before pruning)
-    let collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
-    const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed");
-    let allDistributedNodes = [...collectorNodes, ...upscaleNodes];
-
-    if (isMaster && delegateMaster) {
-        if (upscaleNodes.length > 0) {
-            extension.log("Delegate-only master mode does not support UltimateSDUpscaleDistributed nodes yet. Using full prompt.", "warn");
-        } else if (collectorNodes.length === 0) {
-            extension.log("Delegate-only master mode requested but no collectors found in master prompt. Using full prompt.", "debug");
-        } else {
-            jobApiPrompt = prepareMasterDelegatePrompt(extension, jobApiPrompt, collectorNodes);
-            collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
-            allDistributedNodes = [...collectorNodes, ...upscaleNodes];
-        }
-    }
-    
-    // For workers, handle platform-specific path conversion
-    if (!isMaster) {
-        const workerInfo = extension.config.workers.find(w => w.id === participantId);
-        
-        if (workerInfo && workerInfo.host) {
-            // Remote or cloud worker - needs path translation
-            try {
-                const workerUrl = extension.getWorkerUrl(workerInfo);
-                const systemInfo = await getCachedWorkerSystemInfo(workerUrl);
-                const targetSeparator = systemInfo?.platform?.path_separator;
-                
-                if (targetSeparator) {
-                    // Convert paths to match worker's platform
-                    jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator);
-                    extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug");
-                } else {
-                    extension.log(`No path separator found for worker ${participantId}, skipping path conversion`, "debug");
-                }
-            } catch (e) {
-                extension.log(`Failed to get system info for worker ${participantId}: ${e.message}`, "warn");
-                // Continue without path conversion
-            }
-        }
-        
-        // Prune the workflow to only include distributed node dependencies
-        if (allDistributedNodes.length > 0) {
-            jobApiPrompt = pruneWorkflowForWorker(extension, jobApiPrompt, allDistributedNodes);
-        }
-    }
-    
-    // Handle image references for remote workers
-    if (!isMaster && options.enabled_worker_ids) {
-        // Check if this is a remote worker
-        const workerId = participantId;
-        const workerInfo = extension.config.workers.find(w => w.id === workerId);
-        const isRemote = workerInfo && workerInfo.host;
-        
-        if (isRemote) {
-            // Find all image/video references in the pruned workflow
-            const imageReferences = findImageReferences(extension, jobApiPrompt);
-            if (imageReferences.size > 0) {
-                extension.log(`Found ${imageReferences.size} media references (images/videos) for remote worker ${workerId}`, "debug");
-                // Store image references for later processing
-                options._imageReferences = imageReferences;
-            }
-        }
-    }
-    
-    // Handle Distributed seed nodes
-    const distributorNodes = findNodesByClass(jobApiPrompt, "DistributedSeed");
-    if (distributorNodes.length > 0) {
-        extension.log(`Found ${distributorNodes.length} seed node(s)`, "debug");
-    }
-    
-    for (const seedNode of distributorNodes) {
-        const { inputs } = jobApiPrompt[seedNode.id];
-        inputs.is_worker = !isMaster;
-        if (!isMaster) {
-            const workerIndex = options.enabled_worker_ids.indexOf(participantId);
-            inputs.worker_id = `worker_${workerIndex}`;
-            extension.log(`Set seed node ${seedNode.id} for worker ${workerIndex}`, "debug");
-        }
-    }
-    
-    // Handle Distributed collector nodes (already found above)
-    for (const collector of collectorNodes) {
-        const { inputs } = jobApiPrompt[collector.id];
-        
-        // Check if this collector is downstream from a distributed upscaler
-        const hasUpstreamDistributedUpscaler = hasUpstreamNode(
-            jobApiPrompt, 
-            collector.id, 
-            'UltimateSDUpscaleDistributed'
-        );
-        
-        if (hasUpstreamDistributedUpscaler) {
-            // Set pass_through mode for this collector
-            inputs.pass_through = true;
-            extension.log(`Collector ${collector.id} set to pass-through mode (downstream from distributed upscaler)`, "debug");
-        } else {
-            // Normal collector behavior
-            // Get the unique job ID from the map created for this execution
-            const uniqueJobId = options.job_id_map ? options.job_id_map.get(collector.id) : collector.id;
-            
-            // Use the truly unique ID for this execution
-            inputs.multi_job_id = uniqueJobId;
-            inputs.is_worker = !isMaster;
-            if (isMaster) {
-                inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
-                if (delegateMaster) {
-                    inputs.delegate_only = true;
-                }
-            } else {
-                inputs.master_url = extension.getMasterUrl();
-                // Also make the worker_job_id unique to prevent potential caching issues
-                inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`;
-                inputs.worker_id = participantId;
-            }
-        }
-    }
-    
-    // Handle Ultimate SD Upscale Distributed nodes
-    for (const upscaleNode of upscaleNodes) {
-        const { inputs } = jobApiPrompt[upscaleNode.id];
-        
-        // Get the unique job ID from the map
-        const uniqueJobId = options.job_id_map ? options.job_id_map.get(upscaleNode.id) : upscaleNode.id;
-        
-        inputs.multi_job_id = uniqueJobId;
-        inputs.is_worker = !isMaster;
-        
-        if (isMaster) {
-            inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
-        } else {
-            inputs.master_url = extension.getMasterUrl();
-            inputs.worker_id = participantId;
-            // Workers also need the enabled_worker_ids to calculate tile distribution
-            inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
-        }
-    }
-    
-    return jobApiPrompt;
-}
-
-export async function prepareDistributedJob(extension, multi_job_id) {
-    try {
-        await extension.api.prepareJob(multi_job_id);
-    } catch (error) {
-        extension.log("Error preparing job: " + error.message, "error");
-        throw error;
-    }
-}
-
-function createNumericIdGenerator(promptObj) {
-    let maxId = 0;
-    for (const key of Object.keys(promptObj)) {
-        const numeric = parseInt(key, 10);
-        if (!Number.isNaN(numeric)) {
-            maxId = Math.max(maxId, numeric);
-        }
-    }
-    return () => {
-        maxId += 1;
-        return String(maxId);
-    };
-}
-
-function prepareMasterDelegatePrompt(extension, apiPrompt, collectorNodes) {
-    const collectorIds = collectorNodes.map((node) => node.id);
-    const nodesToKeep = findCollectorDownstreamNodes(apiPrompt, collectorIds);
-
-    // Ensure collectors themselves are present
-    collectorIds.forEach((id) => nodesToKeep.add(id));
-
-    const prunedPrompt = {};
-    nodesToKeep.forEach((nodeId) => {
-        const nodeData = apiPrompt[nodeId];
-        if (nodeData) {
-            prunedPrompt[nodeId] = JSON.parse(JSON.stringify(nodeData));
-        }
-    });
-
-    const prunedIds = new Set(Object.keys(prunedPrompt));
-
-    // Remove dangling references to trimmed nodes
-    for (const [nodeId, node] of Object.entries(prunedPrompt)) {
-        if (!node.inputs) continue;
-        for (const [inputName, inputValue] of Object.entries(node.inputs)) {
-            if (Array.isArray(inputValue) && inputValue.length === 2) {
-                const sourceId = String(inputValue[0]);
-                if (!prunedIds.has(sourceId)) {
-                    delete node.inputs[inputName];
-                    extension.log(`Removed upstream reference '${inputName}' from node ${nodeId} for delegate-only master prompt.`, "debug");
-                }
-            }
-        }
-    }
-
-    const nextId = createNumericIdGenerator(prunedPrompt);
-    collectorIds.forEach((collectorId) => {
-        const collectorEntry = prunedPrompt[collectorId];
-        if (!collectorEntry) return;
-
-        const placeholderId = nextId();
-        prunedPrompt[placeholderId] = {
-            class_type: "DistributedEmptyImage",
-            inputs: {
-                height: 64,
-                width: 64,
-                channels: 3
-            }
-        };
-        prunedIds.add(placeholderId);
-
-        collectorEntry.inputs = collectorEntry.inputs || {};
-        collectorEntry.inputs.images = [placeholderId, 0];
-    });
-
-    extension.log(`Prepared delegate-only master prompt with ${Object.keys(prunedPrompt).length} nodes (kept ${nodesToKeep.size}).`, "debug");
-    return prunedPrompt;
-}
-
-export async function executeJobs(extension, jobs) {
-    let masterPromptId = null;
-    
-    // Pre-load all unique images before dispatching to workers
-    const allImageReferences = new Map();
-    for (const job of jobs) {
-        if (job.type === 'worker' && job.imageReferences) {
-            for (const [filename, info] of job.imageReferences) {
-                allImageReferences.set(filename, info);
-            }
-        }
-    }
-    
-    if (allImageReferences.size > 0) {
-        extension.log(`Pre-loading ${allImageReferences.size} unique media file(s) for all workers`, "debug");
-        await loadImagesForWorker(extension, allImageReferences);
-    }
-    
-    // Now dispatch jobs in parallel
-    const promises = jobs.map(job => {
-        if (job.type === 'master') {
-            return extension.originalQueuePrompt(0, job.promptWrapper).then(result => {
-                masterPromptId = result;
-                return result;
-            });
-        } else {
-            return dispatchToWorker(extension, job.worker, job.prompt, job.workflow, job.imageReferences);
-        }
-    });
-    await Promise.all(promises);
-    
-    // Trigger immediate status check for instant feedback
-    extension.checkAllWorkerStatuses();
-    
-    return masterPromptId || { "prompt_id": "distributed-job-dispatched" };
-}
-
-async function dispatchToWorker(extension, worker, prompt, workflow, imageReferences) {
-    const workerUrl = extension.getWorkerUrl(worker);
-    
-    // Debug logging - always log to console for debugging
-    extension.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`, "debug");
-    extension.log('[Distributed] Worker URL: ' + workerUrl, "debug");
-    
-    // Handle image uploads for remote workers
-    if (imageReferences && imageReferences.size > 0) {
-        // Check if this is a local worker (same host as master)
-        const isLocalWorker = workerUrl.includes('127.0.0.1') || workerUrl.includes('localhost');
-        
-        if (isLocalWorker) {
-            extension.log(`[Distributed] Skipping image processing for local worker ${worker.name} (shares filesystem with master)`, "debug");
-        } else {
-            extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug");
-            
-            try {
-                // Load images from master
-                const images = await loadImagesForWorker(extension, imageReferences);
-                
-                // Upload images to worker
-                if (images.length > 0) {
-                    await uploadImagesToWorker(extension, workerUrl, images);
-                    extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug");
-                }
-            } catch (error) {
-                extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error");
-                // Continue with workflow execution even if image upload fails
-            }
-        }
-    }
-    
-    const promptToSend = {
-        prompt,
-        workflow,
-        client_id: api.clientId
-    };
-
-    extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug");
-
-    try {
-        if (extension.config?.settings?.websocket_orchestration) {
-            const wsUrl = buildWorkerWebSocketUrl(workerUrl);
-            await dispatchPromptViaWebSocket(extension, wsUrl, promptToSend);
-            return;
-        }
-        await fetch(`${workerUrl}/prompt`, {
-            method: 'POST',
-            headers: { 'Content-Type': 'application/json' },
-            mode: 'cors',
-            body: JSON.stringify({
-                prompt,
-                extra_data: { extra_pnginfo: { workflow } },
-                client_id: api.clientId
-            })
-        });
-    } catch (e) {
-        extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error");
-    }
-}
-
-function buildWorkerWebSocketUrl(workerUrl) {
-    const url = new URL(workerUrl);
-    url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:';
-    url.pathname = '/distributed/worker_ws';
-    url.search = '';
-    return url.toString();
-}
-
-function generateRequestId() {
-    if (crypto.randomUUID) {
-        return crypto.randomUUID();
-    }
-    return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => {
-        const r = Math.random() * 16 | 0;
-        const v = c === 'x' ? r : (r & 0x3 | 0x8);
-        return v.toString(16);
-    });
-}
-
-async function dispatchPromptViaWebSocket(extension, wsUrl, payload) {
-    const requestId = generateRequestId();
-    const message = {
-        type: 'dispatch_prompt',
-        request_id: requestId,
-        prompt: payload.prompt,
-        workflow: payload.workflow,
-        client_id: payload.client_id
-    };
-
-    await new Promise((resolve, reject) => {
-        const ws = new WebSocket(wsUrl);
-        const timeoutId = setTimeout(() => {
-            ws.close();
-            reject(new Error('Websocket dispatch timed out.'));
-        }, 60000);
-
-        ws.addEventListener('open', () => {
-            ws.send(JSON.stringify(message));
-        });
-
-        ws.addEventListener('message', (event) => {
-            try {
-                const data = JSON.parse(event.data || '{}');
-                if (data.type !== 'dispatch_ack' || data.request_id !== requestId) {
-                    return;
-                }
-                clearTimeout(timeoutId);
-                ws.close();
-                if (data.ok) {
-                    resolve();
-                } else {
-                    reject(new Error(data.error || 'Worker rejected websocket dispatch.'));
-                }
-            } catch (error) {
-                clearTimeout(timeoutId);
-                ws.close();
-                reject(error);
-            }
-        });
-
-        ws.addEventListener('error', () => {
-            clearTimeout(timeoutId);
-            ws.close();
-            reject(new Error('Websocket connection failed.'));
-        });
-    });
-    extension.log(`[Distributed] Websocket dispatch succeeded for ${wsUrl}`, "debug");
-}
-
-export async function loadImagesForWorker(extension, imageReferences) {
-    const images = [];
-    
-    // Use a cache to avoid loading the same image multiple times
-    if (!extension._imageCache) {
-        extension._imageCache = new Map();
-    }
-    
-    for (const [filename, info] of imageReferences) {
-        try {
-            // Check cache first
-            if (extension._imageCache.has(filename)) {
-                images.push(extension._imageCache.get(filename));
-                extension.log(`Using cached image: ${filename}`, "debug");
-                continue;
-            }
-            
-            // Limit cache size
-            if (extension._imageCache.size >= 10) {
-                const oldestKey = extension._imageCache.keys().next().value;
-                extension._imageCache.delete(oldestKey);
-                extension.log(`Evicted oldest cache entry: ${oldestKey} (cache limit reached)`, "debug");
-            }
-            
-            // Load image from master's filesystem via API
-            try {
-                const data = await extension.api.loadImage(filename);
-                const imageData = {
-                    name: filename,
-                    image: data.image_data,
-                    hash: data.hash  // Include hash from the response
-                };
-                images.push(imageData);
-                
-                // Cache the image for future use
-                extension._imageCache.set(filename, imageData);
-                extension.log(`Loaded and cached image: ${filename}`, "debug");
-            } catch (loadError) {
-                extension.log(`Failed to load image ${filename}: ${loadError.message}`, "error");
-                throw loadError;
-            }
-        } catch (error) {
-            extension.log(`Error loading image ${filename}: ${error.message}`, "error");
-        }
-    }
-    
-    // Clear cache after a reasonable time to avoid memory issues
-    setTimeout(() => {
-        if (extension._imageCache && extension._imageCache.size > 0) {
-            extension.log(`Clearing image cache (${extension._imageCache.size} images)`, "debug");
-            extension._imageCache.clear();
-        }
-    }, TIMEOUTS.IMAGE_CACHE_CLEAR); // Clear after 30 seconds
-    
-    return images;
-}
-
-export async function uploadImagesToWorker(extension, workerUrl, images) {
-    // Upload images to worker's ComfyUI instance
-    for (const imageData of images) {
-        // Check if file already exists with matching hash
-        if (imageData.hash) {
-            try {
-                const checkResponse = await fetch(`${workerUrl}/distributed/check_file`, {
-                    method: 'POST',
-                    headers: { 'Content-Type': 'application/json' },
-                    mode: 'cors',
-                    body: JSON.stringify({ 
-                        filename: imageData.name, 
-                        hash: imageData.hash 
-                    })
-                });
-                
-                if (checkResponse.ok) {
-                    const result = await checkResponse.json();
-                    if (result.exists && result.hash_matches) {
-                        extension.log(`File ${imageData.name} already exists on worker with matching hash, skipping upload`, "debug");
-                        continue;
-                    }
-                }
-            } catch (error) {
-                // If check fails, proceed with upload
-                extension.log(`Failed to check file existence for ${imageData.name}: ${error.message}`, "debug");
-            }
-        }
-        
-        const formData = new FormData();
-        
-        // Detect MIME type from base64 header (supports both image and video)
-        const mimeMatch = imageData.image.match(/^data:((?:image|video)\/\w+);base64,/);
-        const mimeType = mimeMatch ? mimeMatch[1] : 'image/png'; // Default to PNG if not detected
-        
-        // Convert base64 to blob
-        const base64Data = imageData.image.replace(/^data:(?:image|video)\/\w+;base64,/, '');
-        const byteCharacters = atob(base64Data);
-        const byteNumbers = new Array(byteCharacters.length);
-        for (let i = 0; i < byteCharacters.length; i++) {
-            byteNumbers[i] = byteCharacters.charCodeAt(i);
-        }
-        const byteArray = new Uint8Array(byteNumbers);
-        const blob = new Blob([byteArray], { type: mimeType });
-        
-        // Use original filename without heavy cleaning
-        let cleanName = imageData.name;
-        let subfolder = '';
-        
-        // Extract subfolder if present (handle both slash styles)
-        if (cleanName.includes('/') || cleanName.includes('\\')) {
-            const parts = cleanName.replace(/\\/g, '/').split('/');
-            subfolder = parts.slice(0, -1).join('/');
-            cleanName = parts[parts.length - 1];
-        }
-        
-        formData.append('image', blob, cleanName);
-        formData.append('type', 'input');
-        formData.append('subfolder', subfolder);
-        formData.append('overwrite', 'true');
-        
-        try {
-            const response = await fetch(`${workerUrl}/upload/image`, {
-                method: 'POST',
-                mode: 'cors',
-                body: formData
-            });
-            
-            if (!response.ok) {
-                throw new Error(`Upload failed: ${response.statusText}`);
-            }
-            
-            extension.log(`Uploaded image to worker: ${imageData.name} -> ${subfolder}/${cleanName}`, "debug");
-        } catch (error) {
-            extension.log(`Failed to upload ${imageData.name}: ${error.message}`, "error");
-            // Continue with other images
-        }
     }
 }
 
@@ -865,23 +112,18 @@ export async function performPreflightCheck(extension, workers) {
     const startTime = Date.now();
     
     const checkPromises = workers.map(async (worker) => {
-        const url = extension.getWorkerUrl(worker, '/prompt');
+        const workerUrl = getWorkerUrl(extension, worker);
         
-        extension.log(`Pre-flight checking ${worker.name} at: ${url}`, "debug");
+        extension.log(`Pre-flight checking ${worker.name} at: ${workerUrl}`, "debug");
         
         try {
-            const response = await fetch(url, {
-                method: 'GET',
-                mode: 'cors',
-                cache: 'no-store',
-                signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
-            });
+            const probeResult = await extension.api.probeWorker(workerUrl, TIMEOUTS.STATUS_CHECK);
 
-            if (response.ok) {
+            if (probeResult.ok) {
                 extension.log(`Worker ${worker.name} is active`, "debug");
                 return { worker, active: true };
             } else {
-                extension.log(`Worker ${worker.name} returned ${response.status}`, "debug");
+                extension.log(`Worker ${worker.name} returned ${probeResult.status}`, "debug");
                 return { worker, active: false };
             }
         } catch (error) {
@@ -902,13 +144,7 @@ export async function performPreflightCheck(extension, workers) {
     
     // Update UI status indicators for inactive workers
     results.filter(r => !r.active).forEach(r => {
-        const statusDot = document.getElementById(`status-${r.worker.id}`);
-        if (statusDot) {
-            // Remove pulsing animation once status is determined
-            statusDot.classList.remove('status-pulsing');
-            statusDot.style.backgroundColor = "#c04c4c"; // Red for offline
-            statusDot.title = "Offline - Cannot connect";
-        }
+        applyProbeResultToWorkerDot(r.worker.id, { ok: false });
     });
     
     return activeWorkers;

web/masterDetection.js

@@ -0,0 +1,156 @@
+import { generateUUID } from './constants.js';
+
+export async function detectMasterIP(extension) {
+    try {
+        const isRunpod = window.location.hostname.endsWith('.proxy.runpod.net');
+        if (isRunpod) {
+            extension.log("Detected Runpod environment", "info");
+        }
+
+        const data = await extension.api.getNetworkInfo();
+        extension.log("Network info: " + JSON.stringify(data), "debug");
+
+        if (data.cuda_device !== null && data.cuda_device !== undefined) {
+            extension.masterCudaDevice = data.cuda_device;
+
+            if (!extension.config.master) {
+                extension.config.master = {};
+            }
+            if (extension.config.master.cuda_device === undefined || extension.config.master.cuda_device !== data.cuda_device) {
+                extension.config.master.cuda_device = data.cuda_device;
+                try {
+                    await extension.api.updateMaster({ cuda_device: data.cuda_device });
+                    extension.log(`Stored master CUDA device: ${data.cuda_device}`, "debug");
+                } catch (error) {
+                    extension.log(`Error storing master CUDA device: ${error.message}`, "error");
+                }
+            }
+
+            extension.ui.updateMasterDisplay(extension);
+        }
+
+        if (data.cuda_device_count > 0) {
+            extension.cudaDeviceCount = data.cuda_device_count;
+            extension.log(`Detected ${extension.cudaDeviceCount} CUDA devices`, "info");
+
+            const shouldAutoPopulate =
+                !extension.config.settings.has_auto_populated_workers &&
+                (!extension.config.workers || extension.config.workers.length === 0);
+
+            extension.log(`Auto-population check: has_populated=${extension.config.settings.has_auto_populated_workers}, workers=${extension.config.workers ? extension.config.workers.length : 'null'}, should_populate=${shouldAutoPopulate}`, "debug");
+
+            if (shouldAutoPopulate) {
+                extension.log(`Auto-populating workers based on ${extension.cudaDeviceCount} CUDA devices (excluding master on CUDA ${extension.masterCudaDevice})`, "info");
+
+                const workerCount = extension.cudaDeviceCount - (
+                    Number.isInteger(extension.masterCudaDevice) &&
+                    extension.masterCudaDevice >= 0 &&
+                    extension.masterCudaDevice < extension.cudaDeviceCount ? 1 : 0
+                );
+                const workerPorts = data.local_worker_ports;
+                if (!Array.isArray(workerPorts) || workerPorts.length !== workerCount) {
+                    throw new Error("Could not allocate enough available local worker ports; check the server log");
+                }
+                const newWorkers = [];
+                let workerNum = 1;
+                let portOffset = 0;
+
+                for (let i = 0; i < extension.cudaDeviceCount; i++) {
+                    if (i === extension.masterCudaDevice) {
+                        extension.log(`Skipping CUDA ${i} (used by master)`, "debug");
+                        continue;
+                    }
+
+                    const worker = {
+                        id: generateUUID(),
+                        name: `Worker ${workerNum}`,
+                        host: isRunpod ? null : "localhost",
+                        port: workerPorts[portOffset],
+                        cuda_device: i,
+                        enabled: true,
+                        extra_args: isRunpod ? "--listen" : "",
+                    };
+                    newWorkers.push(worker);
+                    workerNum += 1;
+                    portOffset += 1;
+                }
+
+                if (newWorkers.length > 0) {
+                    extension.log(`Auto-populating ${newWorkers.length} workers`, "info");
+
+                    extension.config.workers = newWorkers;
+                    extension.config.settings.has_auto_populated_workers = true;
+
+                    for (const worker of newWorkers) {
+                        try {
+                            await extension.api.updateWorker(worker.id, worker);
+                        } catch (error) {
+                            extension.log(`Error saving worker ${worker.name}: ${error.message}`, "error");
+                        }
+                    }
+
+                    try {
+                        await extension.api.updateSetting('has_auto_populated_workers', true);
+                    } catch (error) {
+                        extension.log(`Error saving auto-population flag: ${error.message}`, "error");
+                    }
+
+                    extension.log(`Auto-populated ${newWorkers.length} workers and saved config`, "info");
+
+                    if (extension.app.extensionManager?.toast) {
+                        extension.app.extensionManager.toast.add({
+                            severity: "success",
+                            summary: "Workers Auto-populated",
+                            detail: `Automatically created ${newWorkers.length} workers based on detected CUDA devices`,
+                            life: 5000,
+                        });
+                    }
+
+                    await extension.loadConfig();
+                } else {
+                    extension.log("No additional CUDA devices available for workers (all used by master)", "debug");
+                }
+            }
+        }
+
+        if (extension.config?.master?.host) {
+            extension.log(`Master host already configured: ${extension.config.master.host}`, "debug");
+            return;
+        }
+
+        if (isRunpod) {
+            const runpodHost = window.location.hostname;
+            extension.log(`Setting Runpod master host: ${runpodHost}`, "info");
+
+            await extension.api.updateMaster({ host: runpodHost });
+
+            if (!extension.config.master) {
+                extension.config.master = {};
+            }
+            extension.config.master.host = runpodHost;
+
+            if (extension.app.extensionManager?.toast) {
+                extension.app.extensionManager.toast.add({
+                    severity: "info",
+                    summary: "Runpod Auto-Configuration",
+                    detail: `Master host set to ${runpodHost} with --listen flag for workers`,
+                    life: 5000,
+                });
+            }
+            return;
+        }
+
+        if (data.recommended_ip && data.recommended_ip !== '127.0.0.1') {
+            extension.log(`Auto-detected master IP: ${data.recommended_ip}`, "info");
+
+            await extension.api.updateMaster({ host: data.recommended_ip });
+
+            if (!extension.config.master) {
+                extension.config.master = {};
+            }
+            extension.config.master.host = data.recommended_ip;
+        }
+    } catch (error) {
+        extension.log("Error detecting master IP: " + error.message, "error");
+    }
+}