From 684b2766199c682eff8856c6a3fa64bb80a58530 Mon Sep 17 00:00:00 2001
From: Robert Wojciechowski <hello@robertvoy.com>
Date: Tue, 3 Mar 2026 20:50:40 +0000
Subject: [PATCH] refactor distributed modules and reorganize tests for quality
 improvements

---
 api/config_routes.py                          |  18 +-
 api/job_routes.py                             |  47 +-
 api/orchestration/dispatch.py                 |  59 +-
 api/orchestration/media_sync.py               |  14 +-
 api/orchestration/prompt_transform.py         |  40 +-
 api/queue_orchestration.py                    | 378 +++++++-----
 api/tunnel_routes.py                          |   8 +-
 api/usdu_routes.py                            |  20 +-
 api/worker_routes.py                          |  82 ++-
 nodes/branch.py                               |  11 +-
 nodes/branch_collector.py                     | 247 +++-----
 nodes/collector.py                            | 556 ++++++++++--------
 nodes/context_kwargs.py                       |  13 +
 nodes/distributed_upscale.py                  | 217 +++++--
 nodes/hidden_inputs.py                        |  31 +
 nodes/list_collector.py                       | 200 +++----
 nodes/list_splitter.py                        |  66 ++-
 nodes/queue_wait.py                           |  46 ++
 nodes/runtime_helpers.py                      |  16 +
 nodes/utilities.py                            |  68 ++-
 ...istributed_queue.py => test_job_routes.py} |   0
 tests/api/test_queue_orchestration.py         | 177 ++++++
 tests/api/test_tunnel_routes.py               | 113 ++++
 tests/api/test_usdu_routes.py                 |   1 +
 tests/api/test_worker_routes.py               |   2 +-
 tests/unit/__init__.py                        |   1 +
 tests/unit/nodes/__init__.py                  |   1 +
 tests/{ => unit/nodes}/test_batch_dividers.py |   2 +-
 tests/{ => unit/nodes}/test_branch.py         |   2 +-
 .../nodes/test_branch_collector.py}           |  78 ++-
 tests/unit/nodes/test_collector.py            | 228 +++++++
 tests/unit/nodes/test_distributed_upscale.py  | 239 ++++++++
 .../nodes}/test_distributed_value.py          |   2 +-
 tests/{ => unit/nodes}/test_list_collector.py |  55 +-
 tests/{ => unit/nodes}/test_list_splitter.py  |   2 +-
 tests/unit/system/__init__.py                 |   1 +
 tests/{ => unit/system}/test_async_helpers.py |   2 +-
 tests/{ => unit/system}/test_config.py        |   5 +-
 tests/{ => unit/system}/test_detection.py     |   5 +-
 .../system/test_dispatch.py}                  |   2 +-
 tests/unit/system/test_distributed.py         | 152 +++++
 tests/unit/system/test_exceptions.py          |  40 ++
 tests/unit/system/test_lifecycle.py           | 228 +++++++
 tests/unit/system/test_logging.py             |  66 +++
 .../system/test_network.py}                   |   8 +-
 .../system}/test_prompt_transform.py          |   2 +-
 tests/{ => unit/system}/test_queue_request.py |   2 +-
 tests/unit/system/test_worker_monitor.py      | 154 +++++
 tests/unit/upscale/__init__.py                |   1 +
 tests/unit/upscale/test_conditioning.py       |  53 ++
 tests/unit/upscale/test_dynamic.py            | 234 ++++++++
 tests/unit/upscale/test_image.py              |  39 ++
 tests/unit/upscale/test_job_state.py          | 131 +++++
 tests/unit/upscale/test_job_store.py          | 140 +++++
 tests/{ => unit/upscale}/test_job_timeout.py  |   2 +-
 .../upscale}/test_payload_parsers.py          |   2 +-
 tests/unit/upscale/test_result_collector.py   | 287 +++++++++
 tests/unit/upscale/test_single_gpu.py         | 206 +++++++
 .../upscale/test_static.py}                   |  52 +-
 tests/unit/upscale/test_tile_ops.py           | 179 ++++++
 tests/unit/upscale/test_tile_processing.py    | 115 ++++
 tests/unit/upscale/test_usdu_utils.py         | 128 ++++
 upscale/conditioning.py                       |   8 +-
 upscale/job_state.py                          |   4 +-
 upscale/job_store.py                          | 127 ++--
 upscale/job_timeout.py                        |   5 +
 upscale/modes/dynamic.py                      | 277 ++++++---
 upscale/modes/single_gpu.py                   |  82 ++-
 upscale/modes/static.py                       | 503 ++++++++++------
 upscale/payload_parsers.py                    |  16 +-
 upscale/processing_args.py                    |  19 +
 upscale/result_collector.py                   |  55 +-
 upscale/tile_processing.py                    |  56 ++
 upscale/worker_comms.py                       |  23 +-
 utils/__init__.py                             |   2 -
 utils/async_helpers.py                        |   6 +-
 utils/audio_payload.py                        |   5 +-
 utils/cloudflare/binary.py                    |  20 +-
 utils/cloudflare/process_reader.py            |  80 ++-
 utils/cloudflare/state.py                     |  72 ++-
 utils/cloudflare/tunnel.py                    | 104 +++-
 utils/config.py                               |  73 ++-
 utils/constants.py                            |   2 +
 utils/crop_model_patch.py                     |  20 +-
 utils/image.py                                |  39 ++
 utils/logging.py                              |  51 +-
 utils/network.py                              |  69 ++-
 utils/process.py                              | 114 +++-
 utils/usdu_managment.py                       |  35 +-
 utils/usdu_utils.py                           | 149 ++++-
 utils/worker_ids.py                           |  20 +
 workers/__init__.py                           |  16 +-
 workers/detection.py                          |  18 +-
 workers/process/__init__.py                   |  12 +-
 workers/process/launch_builder.py             |   5 +-
 workers/process/lifecycle.py                  |  70 ++-
 workers/process/persistence.py                |   4 +-
 workers/process/root_discovery.py             |   2 +-
 workers/process_manager.py                    |   5 +-
 workers/startup.py                            |  36 +-
 workers/worker_monitor.py                     | 120 ++--
 101 files changed, 6017 insertions(+), 1583 deletions(-)
 create mode 100644 nodes/context_kwargs.py
 create mode 100644 nodes/hidden_inputs.py
 create mode 100644 nodes/queue_wait.py
 create mode 100644 nodes/runtime_helpers.py
 rename tests/api/{test_distributed_queue.py => test_job_routes.py} (100%)
 create mode 100644 tests/api/test_queue_orchestration.py
 create mode 100644 tests/api/test_tunnel_routes.py
 create mode 100644 tests/unit/__init__.py
 create mode 100644 tests/unit/nodes/__init__.py
 rename tests/{ => unit/nodes}/test_batch_dividers.py (98%)
 rename tests/{ => unit/nodes}/test_branch.py (97%)
 rename tests/{test_join.py => unit/nodes/test_branch_collector.py} (79%)
 create mode 100644 tests/unit/nodes/test_collector.py
 create mode 100644 tests/unit/nodes/test_distributed_upscale.py
 rename tests/{ => unit/nodes}/test_distributed_value.py (99%)
 rename tests/{ => unit/nodes}/test_list_collector.py (80%)
 rename tests/{ => unit/nodes}/test_list_splitter.py (99%)
 create mode 100644 tests/unit/system/__init__.py
 rename tests/{ => unit/system}/test_async_helpers.py (98%)
 rename tests/{ => unit/system}/test_config.py (98%)
 rename tests/{ => unit/system}/test_detection.py (96%)
 rename tests/{test_dispatch_selection.py => unit/system/test_dispatch.py} (99%)
 create mode 100644 tests/unit/system/test_distributed.py
 create mode 100644 tests/unit/system/test_exceptions.py
 create mode 100644 tests/unit/system/test_lifecycle.py
 create mode 100644 tests/unit/system/test_logging.py
 rename tests/{test_network_helpers.py => unit/system/test_network.py} (91%)
 rename tests/{ => unit/system}/test_prompt_transform.py (99%)
 rename tests/{ => unit/system}/test_queue_request.py (98%)
 create mode 100644 tests/unit/system/test_worker_monitor.py
 create mode 100644 tests/unit/upscale/__init__.py
 create mode 100644 tests/unit/upscale/test_conditioning.py
 create mode 100644 tests/unit/upscale/test_dynamic.py
 create mode 100644 tests/unit/upscale/test_image.py
 create mode 100644 tests/unit/upscale/test_job_state.py
 create mode 100644 tests/unit/upscale/test_job_store.py
 rename tests/{ => unit/upscale}/test_job_timeout.py (99%)
 rename tests/{ => unit/upscale}/test_payload_parsers.py (99%)
 create mode 100644 tests/unit/upscale/test_result_collector.py
 create mode 100644 tests/unit/upscale/test_single_gpu.py
 rename tests/{test_static_mode.py => unit/upscale/test_static.py} (80%)
 create mode 100644 tests/unit/upscale/test_tile_ops.py
 create mode 100644 tests/unit/upscale/test_tile_processing.py
 create mode 100644 tests/unit/upscale/test_usdu_utils.py
 create mode 100644 upscale/processing_args.py
 create mode 100644 upscale/tile_processing.py
 create mode 100644 utils/worker_ids.py

diff --git a/api/config_routes.py b/api/config_routes.py
index 02d1afb..677245c 100644
--- a/api/config_routes.py
+++ b/api/config_routes.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 from contextlib import asynccontextmanager
 
@@ -26,7 +28,7 @@ from ..utils.logging import debug_log, log
 from ..utils.network import handle_api_error, normalize_host
 
 
-def _positive_int(value):
+def _positive_int(value: int) -> bool:
     return value > 0
 
 
@@ -87,13 +89,13 @@ def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
 
 
 @server.PromptServer.instance.routes.get("/distributed/config")
-async def get_config_endpoint(request):
+async def get_config_endpoint(request: web.Request) -> web.StreamResponse:
     config = load_config()
     return web.json_response(config)
 
 
 @server.PromptServer.instance.routes.post("/distributed/config")
-async def update_config_endpoint(request):
+async def update_config_endpoint(request: web.Request) -> web.StreamResponse:
     """Bulk config update with schema validation."""
     try:
         data = await request.json()
@@ -145,7 +147,7 @@ async def update_config_endpoint(request):
 
 
 @server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
-async def queue_status_endpoint(request):
+async def queue_status_endpoint(request: web.Request) -> web.StreamResponse:
     """Check if a job queue is initialized."""
     try:
         job_id = request.match_info['job_id']
@@ -163,7 +165,7 @@ async def queue_status_endpoint(request):
         return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_worker")
-async def update_worker_endpoint(request):
+async def update_worker_endpoint(request: web.Request) -> web.StreamResponse:
     try:
         data = await request.json()
         worker_id = data.get("worker_id")
@@ -207,7 +209,7 @@ async def update_worker_endpoint(request):
         return await handle_api_error(request, e, 400)
 
 @server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
-async def delete_worker_endpoint(request):
+async def delete_worker_endpoint(request: web.Request) -> web.StreamResponse:
     try:
         data = await request.json()
         worker_id = data.get("worker_id")
@@ -239,7 +241,7 @@ async def delete_worker_endpoint(request):
         return await handle_api_error(request, e, 400)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_setting")
-async def update_setting_endpoint(request):
+async def update_setting_endpoint(request: web.Request) -> web.StreamResponse:
     """Updates a specific key in the settings object."""
     try:
         data = await request.json()
@@ -262,7 +264,7 @@ async def update_setting_endpoint(request):
         return await handle_api_error(request, e, 400)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_master")
-async def update_master_endpoint(request):
+async def update_master_endpoint(request: web.Request) -> web.StreamResponse:
     """Updates master configuration."""
     try:
         data = await request.json()
diff --git a/api/job_routes.py b/api/job_routes.py
index 3843101..78bd753 100644
--- a/api/job_routes.py
+++ b/api/job_routes.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 import asyncio
 import io
@@ -5,6 +7,7 @@ import os
 import base64
 import binascii
 import time
+from typing import Any
 
 from aiohttp import web
 import server
@@ -20,8 +23,12 @@ try:
 except ImportError:
     from .queue_orchestration import orchestrate_distributed_execution
 
-    def ensure_distributed_state():
-        return None
+    def ensure_distributed_state() -> None:
+        prompt_server_instance = server.PromptServer.instance
+        if not hasattr(prompt_server_instance, "distributed_jobs_lock"):
+            prompt_server_instance.distributed_jobs_lock = asyncio.Lock()
+        if not hasattr(prompt_server_instance, "distributed_pending_jobs"):
+            prompt_server_instance.distributed_pending_jobs = {}
 from .queue_request import parse_queue_request_payload
 
 prompt_server = server.PromptServer.instance
@@ -30,7 +37,7 @@ prompt_server = server.PromptServer.instance
 # { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool }
 
 
-def _decode_image_sync(image_path):
+def _decode_image_sync(image_path: str) -> dict[str, Any]:
     """Decode image/video file and compute hash in a threadpool worker."""
     import base64
     import hashlib
@@ -76,7 +83,7 @@ def _decode_image_sync(image_path):
     }
 
 
-def _check_file_sync(filename, expected_hash):
+def _check_file_sync(filename: str, expected_hash: str) -> dict[str, Any]:
     """Check file presence and hash in a threadpool worker."""
     import hashlib
     import folder_paths
@@ -101,7 +108,7 @@ def _check_file_sync(filename, expected_hash):
     }
 
 
-def _decode_canonical_png_tensor(image_payload):
+def _decode_canonical_png_tensor(image_payload: str) -> torch.Tensor:
     """Decode canonical base64 PNG payload into a contiguous IMAGE tensor."""
     if not isinstance(image_payload, str) or not image_payload.strip():
         raise ValueError("Field 'image' must be a non-empty base64 PNG string.")
@@ -132,7 +139,7 @@ def _decode_canonical_png_tensor(image_payload):
         raise ValueError(f"Failed to decode PNG image payload: {exc}") from exc
 
 
-def _decode_audio_payload(audio_payload):
+def _decode_audio_payload(audio_payload: dict[str, Any]) -> dict[str, Any]:
     """Decode canonical audio payload into an AUDIO dict."""
     from ..utils.audio_payload import decode_audio_payload
 
@@ -140,7 +147,7 @@ def _decode_audio_payload(audio_payload):
 
 
 @server.PromptServer.instance.routes.post("/distributed/prepare_job")
-async def prepare_job_endpoint(request):
+async def prepare_job_endpoint(request: web.Request) -> web.StreamResponse:
     try:
         data = await request.json()
         multi_job_id = data.get('multi_job_id')
@@ -158,7 +165,7 @@ async def prepare_job_endpoint(request):
         return await handle_api_error(request, e)
 
 @server.PromptServer.instance.routes.post("/distributed/clear_memory")
-async def clear_memory_endpoint(request):
+async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
     debug_log("Received request to clear VRAM.")
     try:
         # Use ComfyUI's prompt server queue system like the /free endpoint does
@@ -204,7 +211,7 @@ async def clear_memory_endpoint(request):
 
 
 @server.PromptServer.instance.routes.post("/distributed/queue")
-async def distributed_queue_endpoint(request):
+async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse:
     """Queue a distributed workflow, mirroring the UI orchestration pipeline."""
     try:
         raw_payload = await request.json()
@@ -234,7 +241,7 @@ async def distributed_queue_endpoint(request):
         return await handle_api_error(request, exc, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/load_image")
-async def load_image_endpoint(request):
+async def load_image_endpoint(request: web.Request) -> web.StreamResponse:
     """Load an image or video file and return it as base64 data with hash."""
     try:
         data = await request.json()
@@ -251,7 +258,7 @@ async def load_image_endpoint(request):
         return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/check_file")
-async def check_file_endpoint(request):
+async def check_file_endpoint(request: web.Request) -> web.StreamResponse:
     """Check if a file exists and matches the given hash."""
     try:
         data = await request.json()
@@ -269,7 +276,7 @@ async def check_file_endpoint(request):
 
 
 @server.PromptServer.instance.routes.post("/distributed/job_complete")
-async def job_complete_endpoint(request):
+async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
     try:
         data = await request.json()
     except Exception as exc:
@@ -314,14 +321,16 @@ async def job_complete_endpoint(request):
             async with prompt_server.distributed_jobs_lock:
                 pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
                 if pending is not None:
+                    queue_item = {
+                        "tensor": tensor,
+                        "worker_id": worker_id,
+                        "image_index": int(batch_idx),
+                        "is_last": is_last,
+                    }
+                    if decoded_audio is not None:
+                        queue_item["audio"] = decoded_audio
                     await pending.put(
-                        {
-                            "tensor": tensor,
-                            "worker_id": worker_id,
-                            "image_index": int(batch_idx),
-                            "is_last": is_last,
-                            "audio": decoded_audio,
-                        }
+                        queue_item
                     )
                     queue_size = pending.qsize()
                     break
diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py
index 48125b5..34c0f16 100644
--- a/api/orchestration/dispatch.py
+++ b/api/orchestration/dispatch.py
@@ -1,6 +1,7 @@
 import asyncio
 import json
 import uuid
+from typing import Any
 
 import aiohttp
 
@@ -9,11 +10,11 @@ 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_fallback(_execution_id: str, message: str) -> None:
+        debug_log(message)
 
-    def trace_info(*_args, **_kwargs):
-        return None
+    trace_debug = _trace_fallback
+    trace_info = _trace_fallback
 
 try:
     from ..schemas import parse_positive_int
@@ -28,13 +29,13 @@ except ImportError:
 _least_busy_rr_index = 0
 
 
-async def worker_is_active(worker):
+async def worker_is_active(worker: dict[str, Any]) -> bool:
     """Ping worker's /prompt endpoint to confirm it's reachable."""
     url = build_worker_url(worker)
     return await probe_worker(url, timeout=3.0) is not None
 
 
-async def worker_ws_is_active(worker):
+async def worker_ws_is_active(worker: dict[str, Any]) -> bool:
     """Ping worker's websocket endpoint to confirm it's reachable."""
     session = await get_client_session()
     url = build_worker_url(worker, "/distributed/worker_ws")
@@ -96,13 +97,13 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
 
 
 async def dispatch_worker_prompt(
-    worker,
-    prompt_obj,
-    workflow_meta,
-    client_id=None,
-    use_websocket=False,
-    trace_execution_id=None,
-):
+    worker: dict[str, Any],
+    prompt_obj: dict[str, Any],
+    workflow_meta: dict[str, Any] | None,
+    client_id: str | None = None,
+    use_websocket: bool = False,
+    trace_execution_id: str | None = None,
+) -> None:
     """Send the prepared prompt to a worker ComfyUI instance."""
     worker_url = build_worker_url(worker)
     url = build_worker_url(worker, "/prompt")
@@ -142,12 +143,12 @@ async def dispatch_worker_prompt(
 
 
 async def select_active_workers(
-    workers,
-    use_websocket,
-    delegate_master,
-    trace_execution_id=None,
-    probe_concurrency=8,
-):
+    workers: list[dict[str, Any]],
+    use_websocket: bool,
+    delegate_master: bool,
+    trace_execution_id: str | None = None,
+    probe_concurrency: int = 8,
+) -> tuple[list[dict[str, Any]], bool]:
     """Probe workers and return (active_workers, updated_delegate_master)."""
     probe_limit = parse_positive_int(probe_concurrency, 8)
     probe_semaphore = asyncio.Semaphore(probe_limit)
@@ -223,11 +224,11 @@ def _select_idle_round_robin(statuses):
 
 
 async def select_least_busy_worker(
-    workers,
-    trace_execution_id=None,
-    probe_concurrency=8,
-    probe_timeout=3.0,
-):
+    workers: list[dict[str, Any]],
+    trace_execution_id: str | None = None,
+    probe_concurrency: int = 8,
+    probe_timeout: float = 3.0,
+) -> dict[str, Any] | None:
     """Select one worker by queue depth, round-robin among idle workers."""
     if not workers:
         return None
@@ -269,11 +270,11 @@ async def select_least_busy_worker(
 
 
 async def rank_workers_by_load(
-    workers,
-    trace_execution_id=None,
-    probe_concurrency=8,
-    probe_timeout=3.0,
-):
+    workers: list[dict[str, Any]],
+    trace_execution_id: str | None = None,
+    probe_concurrency: int = 8,
+    probe_timeout: float = 3.0,
+) -> list[dict[str, Any]]:
     """Return all workers sorted by queue depth (ascending), preserving unreachable workers at the end."""
     if not workers:
         return []
diff --git a/api/orchestration/media_sync.py b/api/orchestration/media_sync.py
index 84e51fe..fc61c6e 100644
--- a/api/orchestration/media_sync.py
+++ b/api/orchestration/media_sync.py
@@ -3,6 +3,7 @@ import hashlib
 import mimetypes
 import os
 import re
+from typing import Any
 
 import aiohttp
 
@@ -33,7 +34,7 @@ def _normalize_media_reference(value):
     return None
 
 
-def convert_paths_for_platform(obj, target_separator):
+def convert_paths_for_platform(obj: Any, target_separator: str) -> Any:
     """Recursively normalize likely file paths for the worker platform separator."""
     if target_separator not in ("/", "\\"):
         return obj
@@ -124,7 +125,10 @@ def _load_media_file_sync(filename):
     return file_bytes, file_hash, mime_type
 
 
-async def fetch_worker_path_separator(worker, trace_execution_id=None):
+async def fetch_worker_path_separator(
+    worker: dict[str, Any],
+    trace_execution_id: str | None = None,
+) -> str | None:
     """Best-effort fetch of a worker's path separator from /distributed/system_info."""
     url = build_worker_url(worker, "/distributed/system_info")
     session = await get_client_session()
@@ -193,7 +197,11 @@ async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_
     return True, worker_path
 
 
-async def sync_worker_media(worker, prompt_obj, trace_execution_id=None):
+async def sync_worker_media(
+    worker: dict[str, Any],
+    prompt_obj: dict[str, Any],
+    trace_execution_id: str | None = None,
+) -> None:
     """Sync referenced media files from master to a remote worker before dispatch."""
     media_refs = _find_media_references(prompt_obj)
     if not media_refs:
diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py
index e2064b7..0111585 100644
--- a/api/orchestration/prompt_transform.py
+++ b/api/orchestration/prompt_transform.py
@@ -1,5 +1,6 @@
 import json
 from collections import deque
+from typing import Any
 
 from ...utils.logging import debug_log
 
@@ -27,7 +28,7 @@ class PromptIndex:
     def nodes_for_class(self, class_name):
         return self.nodes_by_class.get(class_name, [])
 
-    def has_upstream(self, start_node_id, target_class):
+    def has_upstream(self, start_node_id: str, target_class: str) -> bool:
         cache_key = (str(start_node_id), target_class)
         if cache_key in self._upstream_cache:
             return self._upstream_cache[cache_key]
@@ -107,6 +108,9 @@ def _find_downstream_of_output_slot(prompt_obj, node_id, slot_index):
             try:
                 input_slot = int(value[1])
             except (TypeError, ValueError):
+                debug_log(
+                    f"Prompt transform: skipping malformed slot reference for node {candidate_id}: {value}"
+                )
                 continue
             if input_slot == slot:
                 direct_consumers.add(str(candidate_id))
@@ -124,6 +128,7 @@ def _create_numeric_id_generator(prompt_obj):
         try:
             numeric = int(node_id)
         except (TypeError, ValueError):
+            debug_log(f"Prompt transform: ignoring non-numeric node id while generating ids: {node_id}")
             continue
         max_id = max(max_id, numeric)
 
@@ -267,7 +272,12 @@ def _prune_worker_downstream_of_branch_collectors(prompt_obj):
     return prompt_obj
 
 
-def prune_prompt_for_branch_worker(prompt_obj, branch_node_id, assigned_branch, num_branches):
+def prune_prompt_for_branch_worker(
+    prompt_obj: dict[str, Any],
+    branch_node_id: str,
+    assigned_branch: int | list[int] | tuple[int, ...] | set[int],
+    num_branches: int,
+) -> dict[str, Any]:
     """Prune non-assigned branch paths while keeping shared downstream nodes."""
     branch_id = str(branch_node_id)
 
@@ -277,6 +287,7 @@ def prune_prompt_for_branch_worker(prompt_obj, branch_node_id, assigned_branch,
             try:
                 idx = int(value)
             except (TypeError, ValueError):
+                debug_log(f"Prompt transform: invalid assigned branch slot ignored: {value}")
                 continue
             if idx >= 0:
                 assigned_slots.add(idx)
@@ -318,7 +329,7 @@ def prune_prompt_for_branch_worker(prompt_obj, branch_node_id, assigned_branch,
     return prompt_obj
 
 
-def prune_prompt_for_worker(prompt_obj):
+def prune_prompt_for_worker(prompt_obj: dict[str, Any]) -> dict[str, Any]:
     """Prune worker prompt to distributed nodes and their upstream dependencies."""
     collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
     list_collector_ids = find_nodes_by_class(prompt_obj, "DistributedListCollector")
@@ -364,7 +375,10 @@ def prune_prompt_for_worker(prompt_obj):
     return pruned_prompt
 
 
-def prepare_delegate_master_prompt(prompt_obj, collector_ids):
+def prepare_delegate_master_prompt(
+    prompt_obj: dict[str, Any],
+    collector_ids: list[str],
+) -> dict[str, Any]:
     """Prune master prompt so it only executes post-collector nodes in delegate mode."""
     downstream = _find_downstream_nodes(prompt_obj, collector_ids)
     nodes_to_keep = set(collector_ids)
@@ -420,7 +434,7 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
     return pruned_prompt
 
 
-def generate_job_id_map(prompt_index, prefix):
+def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str]:
     """Create stable per-node job IDs for distributed nodes."""
     job_map = {}
     distributed_nodes = (
@@ -716,14 +730,14 @@ def _override_branch_collector_nodes(
 
 
 def apply_participant_overrides(
-    prompt_copy,
-    participant_id,
-    enabled_worker_ids,
-    job_id_map,
-    master_url,
-    delegate_master,
-    prompt_index,
-):
+    prompt_copy: dict[str, Any],
+    participant_id: str,
+    enabled_worker_ids: list[str],
+    job_id_map: dict[str, str],
+    master_url: str,
+    delegate_master: bool,
+    prompt_index: PromptIndex,
+) -> dict[str, Any]:
     """Return a prompt copy with hidden inputs configured for master/worker."""
     is_master = participant_id == "master"
     worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py
index 33db94c..52c2ec0 100644
--- a/api/queue_orchestration.py
+++ b/api/queue_orchestration.py
@@ -1,6 +1,7 @@
 import asyncio
 import time
 import uuid
+from typing import Any
 
 import server
 
@@ -192,14 +193,208 @@ async def _prepare_worker_payload(
         return worker, worker_prompt
 
 
-async def orchestrate_distributed_execution(
-    prompt_obj,
+async def _select_execution_workers(
+    workers,
+    use_websocket,
+    delegate_master,
+    load_balance_requested,
+    has_branch_nodes,
+    master_url,
+    execution_trace_id,
+    worker_probe_concurrency,
+):
+    active_workers, delegate_master = await select_active_workers(
+        workers,
+        use_websocket,
+        delegate_master,
+        trace_execution_id=execution_trace_id,
+        probe_concurrency=worker_probe_concurrency,
+    )
+
+    if load_balance_requested:
+        candidate_workers = list(active_workers)
+        if not delegate_master:
+            candidate_workers.append(
+                {
+                    "id": "master",
+                    "name": "Master",
+                    "host": master_url,
+                    "type": "local",
+                }
+            )
+
+        selected_worker = None
+        if candidate_workers:
+            selected_worker = await select_least_busy_worker(
+                candidate_workers,
+                trace_execution_id=execution_trace_id,
+                probe_concurrency=worker_probe_concurrency,
+            )
+        if selected_worker is None and candidate_workers:
+            trace_debug(
+                execution_trace_id,
+                "Load-balance selection probe failed; using first available candidate.",
+            )
+            selected_worker = candidate_workers[0]
+
+        if selected_worker is not None and str(selected_worker.get("id")) == "master":
+            active_workers = []
+            delegate_master = False
+            trace_debug(
+                execution_trace_id,
+                "Load-balance selected master for execution (workers skipped).",
+            )
+        elif selected_worker is not None:
+            active_workers = [selected_worker]
+            delegate_master = True
+            trace_debug(
+                execution_trace_id,
+                f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).",
+            )
+        else:
+            trace_debug(
+                execution_trace_id,
+                "Load-balance requested but no execution candidates were available.",
+            )
+            active_workers = []
+            delegate_master = False
+
+    if has_branch_nodes and len(active_workers) > 1:
+        active_workers = await rank_workers_by_load(
+            active_workers,
+            trace_execution_id=execution_trace_id,
+            probe_concurrency=worker_probe_concurrency,
+        )
+
+    return active_workers, delegate_master
+
+
+def _distributed_queue_nodes(prompt_index):
+    return (
+        prompt_index.nodes_for_class("DistributedCollector")
+        + prompt_index.nodes_for_class("DistributedListCollector")
+        + prompt_index.nodes_for_class("DistributedBranchCollector")
+        + prompt_index.nodes_for_class("DistributedBranch")
+        + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
+    )
+
+
+async def _ensure_job_queues(prompt_index, job_id_map):
+    for node_id in _distributed_queue_nodes(prompt_index):
+        job_id = job_id_map.get(node_id)
+        if job_id:
+            await _ensure_distributed_queue(job_id)
+
+
+def _build_master_prompt(
+    prompt_index,
+    enabled_ids,
+    job_id_map,
+    master_url,
+    delegate_master,
+):
+    master_prompt = prompt_index.copy_prompt()
+    master_prompt = apply_participant_overrides(
+        master_prompt,
+        "master",
+        enabled_ids,
+        job_id_map,
+        master_url,
+        delegate_master,
+        prompt_index,
+    )
+
+    if not delegate_master:
+        return master_prompt
+
+    collector_ids = (
+        find_nodes_by_class(master_prompt, "DistributedCollector")
+        + find_nodes_by_class(master_prompt, "DistributedListCollector")
+        + find_nodes_by_class(master_prompt, "DistributedBranchCollector")
+    )
+    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."
+        )
+        return master_prompt
+    if not collector_ids:
+        debug_log(
+            "Delegate-only master mode requested but no collector/branch-collector nodes found in master prompt. Running full prompt on master."
+        )
+        return master_prompt
+    return prepare_delegate_master_prompt(master_prompt, collector_ids)
+
+
+async def _prepare_worker_payloads_for_dispatch(
+    active_workers,
+    prompt_index,
+    enabled_ids,
+    job_id_map,
+    master_url,
+    delegate_master,
+    execution_trace_id,
+    worker_prep_concurrency,
+    media_sync_concurrency,
+    media_sync_timeout_seconds,
+):
+    if not active_workers:
+        return []
+
+    worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
+    media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
+    return await asyncio.gather(
+        *[
+            _prepare_worker_payload(
+                worker,
+                prompt_index,
+                enabled_ids,
+                job_id_map,
+                master_url,
+                delegate_master,
+                execution_trace_id,
+                worker_prep_semaphore,
+                media_sync_semaphore,
+                media_sync_timeout_seconds,
+            )
+            for worker in active_workers
+        ]
+    )
+
+
+async def _dispatch_worker_payloads(
+    worker_payloads,
     workflow_meta,
     client_id,
-    enabled_worker_ids=None,
-    delegate_master=None,
-    trace_execution_id=None,
+    use_websocket,
+    execution_trace_id,
 ):
+    if not worker_payloads:
+        return
+
+    await asyncio.gather(
+        *[
+            dispatch_worker_prompt(
+                worker,
+                wprompt,
+                workflow_meta,
+                client_id,
+                use_websocket=use_websocket,
+                trace_execution_id=execution_trace_id,
+            )
+            for worker, wprompt in worker_payloads
+        ]
+    )
+
+
+async def orchestrate_distributed_execution(
+    prompt_obj: dict[str, Any],
+    workflow_meta: dict[str, Any] | None,
+    client_id: str | None,
+    enabled_worker_ids: list[str] | set[str] | None = None,
+    delegate_master: bool | None = None,
+    trace_execution_id: str | None = None,
+) -> tuple[str, int]:
     """Core orchestration logic for the /distributed/queue endpoint.
 
     Returns:
@@ -247,72 +442,17 @@ async def orchestrate_distributed_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,
+    active_workers, delegate_master = await _select_execution_workers(
+        workers=workers,
+        use_websocket=use_websocket,
+        delegate_master=delegate_master,
+        load_balance_requested=load_balance_requested,
+        has_branch_nodes=has_branch_nodes,
+        master_url=master_url,
+        execution_trace_id=execution_trace_id,
+        worker_probe_concurrency=worker_probe_concurrency,
     )
 
-    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
-
-    if has_branch_nodes and len(active_workers) > 1:
-        active_workers = await rank_workers_by_load(
-            active_workers,
-            trace_execution_id=execution_trace_id,
-            probe_concurrency=worker_probe_concurrency,
-        )
-
     enabled_ids = [worker["id"] for worker in active_workers]
 
     discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
@@ -323,89 +463,41 @@ async def orchestrate_distributed_execution(
         prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
         return prompt_id, 0
 
-    queue_node_ids = (
-        prompt_index.nodes_for_class("DistributedCollector")
-        + prompt_index.nodes_for_class("DistributedListCollector")
-        + prompt_index.nodes_for_class("DistributedBranchCollector")
-        + prompt_index.nodes_for_class("DistributedBranch")
-        + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
-    )
-    for node_id in queue_node_ids:
-        job_id = job_id_map.get(node_id)
-        if job_id:
-            await _ensure_distributed_queue(job_id)
+    await _ensure_job_queues(prompt_index, job_id_map)
 
-    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,
+    master_prompt = _build_master_prompt(
+        prompt_index=prompt_index,
+        enabled_ids=enabled_ids,
+        job_id_map=job_id_map,
+        master_url=master_url,
+        delegate_master=delegate_master,
     )
 
-    if delegate_master:
-        collector_ids = (
-            find_nodes_by_class(master_prompt, "DistributedCollector")
-            + find_nodes_by_class(master_prompt, "DistributedListCollector")
-            + find_nodes_by_class(master_prompt, "DistributedBranchCollector")
-        )
-        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 collector/branch-collector nodes 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,
-                    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
-            ]
-        )
+    worker_payloads = await _prepare_worker_payloads_for_dispatch(
+        active_workers=active_workers,
+        prompt_index=prompt_index,
+        enabled_ids=enabled_ids,
+        job_id_map=job_id_map,
+        master_url=master_url,
+        delegate_master=delegate_master,
+        execution_trace_id=execution_trace_id,
+        worker_prep_concurrency=worker_prep_concurrency,
+        media_sync_concurrency=media_sync_concurrency,
+        media_sync_timeout_seconds=media_sync_timeout_seconds,
+    )
+    await _dispatch_worker_payloads(
+        worker_payloads=worker_payloads,
+        workflow_meta=workflow_meta,
+        client_id=client_id,
+        use_websocket=use_websocket,
+        execution_trace_id=execution_trace_id,
+    )
 
     prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
     trace_debug(
diff --git a/api/tunnel_routes.py b/api/tunnel_routes.py
index 4fdd89d..89158a4 100644
--- a/api/tunnel_routes.py
+++ b/api/tunnel_routes.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from aiohttp import web
 import server
 
@@ -8,7 +10,7 @@ from ..utils.network import handle_api_error
 
 
 @server.PromptServer.instance.routes.get("/distributed/tunnel/status")
-async def tunnel_status_endpoint(request):
+async def tunnel_status_endpoint(request: web.Request) -> web.StreamResponse:
     """Return Cloudflare tunnel status and last known details."""
     try:
         status = cloudflare_tunnel_manager.get_status()
@@ -23,7 +25,7 @@ async def tunnel_status_endpoint(request):
         return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/tunnel/start")
-async def tunnel_start_endpoint(request):
+async def tunnel_start_endpoint(request: web.Request) -> web.StreamResponse:
     """Start a Cloudflare tunnel pointing at the current ComfyUI server."""
     try:
         result = await cloudflare_tunnel_manager.start_tunnel()
@@ -37,7 +39,7 @@ async def tunnel_start_endpoint(request):
         return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
-async def tunnel_stop_endpoint(request):
+async def tunnel_stop_endpoint(request: web.Request) -> web.StreamResponse:
     """Stop the managed Cloudflare tunnel if running."""
     try:
         result = await cloudflare_tunnel_manager.stop_tunnel()
diff --git a/api/usdu_routes.py b/api/usdu_routes.py
index 04ca0d4..abd06c5 100644
--- a/api/usdu_routes.py
+++ b/api/usdu_routes.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import io
 import time
@@ -8,13 +10,13 @@ import server
 
 from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
 from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized, init_dynamic_job
-from ..upscale.payload_parsers import _parse_tiles_from_form
+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):
+async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse:
     try:
         data = await request.json()
         worker_id = data.get('worker_id')
@@ -38,7 +40,7 @@ async def heartbeat_endpoint(request):
 
 
 @server.PromptServer.instance.routes.post("/distributed/submit_tiles")
-async def submit_tiles_endpoint(request):
+async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
     """Endpoint for workers to submit processed tiles in static mode."""
     try:
         content_length = request.headers.get('content-length')
@@ -73,7 +75,7 @@ async def submit_tiles_endpoint(request):
                     return web.json_response({"status": "success"})
 
         try:
-            tiles = _parse_tiles_from_form(data)
+            tiles = parse_tiles_from_form(data)
         except ValueError as e:
             return await handle_api_error(request, str(e), 400)
 
@@ -106,7 +108,7 @@ async def submit_tiles_endpoint(request):
 
 
 @server.PromptServer.instance.routes.post("/distributed/submit_image")
-async def submit_image_endpoint(request):
+async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
     """Endpoint for workers to submit processed images in dynamic mode."""
     try:
         content_length = request.headers.get('content-length')
@@ -166,7 +168,7 @@ async def submit_image_endpoint(request):
 
 
 @server.PromptServer.instance.routes.post("/distributed/request_image")
-async def request_image_endpoint(request):
+async def request_image_endpoint(request: web.Request) -> web.StreamResponse:
     """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
     try:
         data = await request.json()
@@ -216,7 +218,7 @@ async def request_image_endpoint(request):
 
 
 @server.PromptServer.instance.routes.post("/distributed/init_list_queue")
-async def init_list_queue_endpoint(request):
+async def init_list_queue_endpoint(request: web.Request) -> web.StreamResponse:
     """Initialize a shared queue of list indices for dynamic DistributedListSplitter mode."""
     try:
         data = await request.json()
@@ -255,7 +257,7 @@ async def init_list_queue_endpoint(request):
 
 
 @server.PromptServer.instance.routes.post("/distributed/request_list_item")
-async def request_list_item_endpoint(request):
+async def request_list_item_endpoint(request: web.Request) -> web.StreamResponse:
     """Allow workers to pull one list item index for dynamic list distribution."""
     try:
         data = await request.json()
@@ -290,7 +292,7 @@ async def request_list_item_endpoint(request):
 
 
 @server.PromptServer.instance.routes.get("/distributed/job_status")
-async def job_status_endpoint(request):
+async def job_status_endpoint(request: web.Request) -> web.StreamResponse:
     """Endpoint to check if a job is ready."""
     multi_job_id = request.query.get('multi_job_id')
     if not multi_job_id:
diff --git a/api/worker_routes.py b/api/worker_routes.py
index 336eec7..7d32d46 100644
--- a/api/worker_routes.py
+++ b/api/worker_routes.py
@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 import json
 import asyncio
 import os
 import time
 import platform
-import subprocess
+import subprocess  # nosec B404 - commands are fixed and never shell-expanded
 import socket
+from typing import Any
 
 import torch
 import aiohttp
@@ -41,7 +44,7 @@ except ImportError:
 
 
 @server.PromptServer.instance.routes.get("/distributed/worker_ws")
-async def worker_ws_endpoint(request):
+async def worker_ws_endpoint(request: web.Request) -> web.WebSocketResponse:
     """WebSocket endpoint for worker prompt dispatch."""
     ws = web.WebSocketResponse(heartbeat=30)
     await ws.prepare(request)
@@ -113,7 +116,7 @@ async def worker_ws_endpoint(request):
 
 
 @server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
-async def clear_launching_state(request):
+async def clear_launching_state(request: web.Request) -> web.StreamResponse:
     """Clear the launching flag when worker is confirmed running."""
     try:
         wm = get_worker_manager()
@@ -139,8 +142,9 @@ async def clear_launching_state(request):
         return await handle_api_error(request, e, 500)
 
 
-def get_network_ips():
+def get_network_ips() -> list[str]:
     """Get all network IPs, trying multiple methods."""
+    command_timeout = 5.0
     ips = []
     hostname = socket.gethostname()
 
@@ -151,8 +155,8 @@ def get_network_ips():
             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
+    except (socket.gaierror, OSError) as exc:
+        debug_log(f"get_network_ips: getaddrinfo failed for hostname {hostname}: {exc}")
 
     # Method 2: Try to connect to external server and get local IP
     try:
@@ -162,14 +166,19 @@ def get_network_ips():
         s.close()
         if local_ip not in ips:
             ips.append(local_ip)
-    except (OSError, socket.error):
-        pass
+    except (OSError, socket.error) as exc:
+        debug_log(f"get_network_ips: UDP local IP probe failed: {exc}")
 
     # Method 3: Platform-specific commands
     try:
         if platform.system() == "Windows":
             # Windows ipconfig
-            result = subprocess.run(["ipconfig"], capture_output=True, text=True)
+            result = subprocess.run(  # nosec B603 - static command, no user input
+                ["ipconfig"],
+                capture_output=True,
+                text=True,
+                timeout=command_timeout,
+            )
             lines = result.stdout.split('\n')
             for i, line in enumerate(lines):
                 if 'IPv4' in line and i + 1 < len(lines):
@@ -179,10 +188,20 @@ def get_network_ips():
         else:
             # Unix/Linux/Mac ifconfig or ip addr
             try:
-                result = subprocess.run(["ip", "addr"], capture_output=True, text=True)
+                result = subprocess.run(  # nosec B603 - static command, no user input
+                    ["ip", "addr"],
+                    capture_output=True,
+                    text=True,
+                    timeout=command_timeout,
+                )
             except (FileNotFoundError, OSError):
                 try:
-                    result = subprocess.run(["ifconfig"], capture_output=True, text=True)
+                    result = subprocess.run(  # nosec B603 - static command, no user input
+                        ["ifconfig"],
+                        capture_output=True,
+                        text=True,
+                        timeout=command_timeout,
+                    )
                 except (FileNotFoundError, OSError):
                     result = None
 
@@ -193,13 +212,13 @@ def get_network_ips():
                     ip = match.group(1)
                     if ip and ip not in ips:
                         ips.append(ip)
-    except (OSError, subprocess.SubprocessError):
-        pass
+    except (OSError, subprocess.SubprocessError) as exc:
+        debug_log(f"get_network_ips: platform command probe failed: {exc}")
 
     return ips
 
 
-def get_recommended_ip(ips):
+def get_recommended_ip(ips: list[str]) -> str | None:
     """Choose the best IP for master-worker communication."""
     # Priority order:
     # 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x)
@@ -234,7 +253,7 @@ def get_recommended_ip(ips):
         return None
 
 
-def _get_cuda_info():
+def _get_cuda_info() -> tuple[int | None, int, int]:
     """Detect CUDA device index and total physical GPU count.
 
     Returns (cuda_device, cuda_device_count, physical_device_count).
@@ -250,7 +269,7 @@ def _get_cuda_info():
             if visible_devices:
                 cuda_device = visible_devices[0]
                 try:
-                    result = subprocess.run(
+                    result = subprocess.run(  # nosec B603 - static command, no user input
                         ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
                         capture_output=True,
                         text=True,
@@ -274,7 +293,7 @@ def _get_cuda_info():
         return None, 0, 0
 
 
-def _collect_network_info_sync():
+def _collect_network_info_sync() -> dict[str, Any]:
     """Collect network/cuda info in a worker thread to avoid blocking route handlers."""
     cuda_device, cuda_device_count, physical_device_count = _get_cuda_info()
     hostname = socket.gethostname()
@@ -289,7 +308,7 @@ def _collect_network_info_sync():
     }
 
 
-def _read_worker_log_sync(log_file, lines_to_read):
+def _read_worker_log_sync(log_file: str, lines_to_read: int) -> dict[str, Any]:
     """Read worker log content from disk in a threadpool worker."""
     file_size = os.path.getsize(log_file)
 
@@ -325,7 +344,12 @@ def _read_worker_log_sync(log_file, lines_to_read):
     }
 
 
-def _parse_positive_int_query(value, default, minimum=1, maximum=10000):
+def _parse_positive_int_query(
+    value: Any,
+    default: int,
+    minimum: int = 1,
+    maximum: int | None = 10000,
+) -> int:
     """Parse bounded positive integer query params with sane fallback."""
     try:
         parsed = int(value)
@@ -337,7 +361,7 @@ def _parse_positive_int_query(value, default, minimum=1, maximum=10000):
     return parsed
 
 
-def _find_worker_by_id(config, worker_id):
+def _find_worker_by_id(config: dict[str, Any], worker_id: str) -> dict[str, Any] | None:
     worker_id_str = str(worker_id).strip()
     for worker in config.get("workers", []):
         if str(worker.get("id")).strip() == worker_id_str:
@@ -346,7 +370,7 @@ def _find_worker_by_id(config, worker_id):
 
 
 @server.PromptServer.instance.routes.get("/distributed/local_log")
-async def get_local_log_endpoint(request):
+async def get_local_log_endpoint(request: web.Request) -> web.StreamResponse:
     """Return this instance's in-memory ComfyUI log buffer."""
     try:
         from app.logger import get_logs
@@ -391,7 +415,7 @@ async def get_local_log_endpoint(request):
 
 
 @server.PromptServer.instance.routes.get("/distributed/network_info")
-async def get_network_info_endpoint(request):
+async def get_network_info_endpoint(request: web.Request) -> web.StreamResponse:
     """Get network interfaces and recommend best IP for master."""
     try:
         loop = asyncio.get_running_loop()
@@ -406,7 +430,7 @@ async def get_network_info_endpoint(request):
         return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.get("/distributed/system_info")
-async def get_system_info_endpoint(request):
+async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse:
     """Get system information including machine ID for local worker detection."""
     try:
         import socket
@@ -430,7 +454,7 @@ async def get_system_info_endpoint(request):
         return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/launch_worker")
-async def launch_worker_endpoint(request):
+async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
     """Launch a worker process from the UI."""
     try:
         wm = get_worker_manager()
@@ -491,7 +515,7 @@ async def launch_worker_endpoint(request):
 
 
 @server.PromptServer.instance.routes.post("/distributed/stop_worker")
-async def stop_worker_endpoint(request):
+async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse:
     """Stop a worker process that was launched from the UI."""
     try:
         wm = get_worker_manager()
@@ -521,7 +545,7 @@ async def stop_worker_endpoint(request):
 
 
 @server.PromptServer.instance.routes.get("/distributed/managed_workers")
-async def get_managed_workers_endpoint(request):
+async def get_managed_workers_endpoint(request: web.Request) -> web.StreamResponse:
     """Get list of workers managed by this UI instance."""
     try:
         managed = get_worker_manager().get_managed_workers()
@@ -534,7 +558,7 @@ async def get_managed_workers_endpoint(request):
 
 
 @server.PromptServer.instance.routes.get("/distributed/local-worker-status")
-async def get_local_worker_status_endpoint(request):
+async def get_local_worker_status_endpoint(request: web.Request) -> web.StreamResponse:
     """Check status of all local workers (localhost/no host specified)."""
     try:
         config = load_config()
@@ -604,7 +628,7 @@ async def get_local_worker_status_endpoint(request):
 
 
 @server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}")
-async def get_worker_log_endpoint(request):
+async def get_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
     """Get log content for a specific worker."""
     try:
         wm = get_worker_manager()
@@ -647,7 +671,7 @@ async def get_worker_log_endpoint(request):
 
 
 @server.PromptServer.instance.routes.get("/distributed/remote_worker_log/{worker_id}")
-async def get_remote_worker_log_endpoint(request):
+async def get_remote_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
     """Proxy a remote worker log request to the worker's local in-memory log endpoint."""
     try:
         worker_id = str(request.match_info["worker_id"]).strip()
diff --git a/nodes/branch.py b/nodes/branch.py
index b1aa7bc..87267ca 100644
--- a/nodes/branch.py
+++ b/nodes/branch.py
@@ -1,9 +1,11 @@
+from typing import Any
+
 from .utilities import any_type
 
 
 class DistributedBranch:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls: type["DistributedBranch"]) -> dict[str, Any]:
         return {
             "required": {
                 "input": (any_type,),
@@ -22,5 +24,10 @@ class DistributedBranch:
     FUNCTION = "branch"
     CATEGORY = "utils"
 
-    def branch(self, input, num_branches=2, is_worker=False, worker_id="", assigned_branch=-1, multi_job_id=""):
+    def branch(
+        self,
+        input: Any,
+        num_branches: int = 2,
+        **_kwargs: Any,
+    ) -> tuple[Any, ...]:
         return tuple([input] * 10)
diff --git a/nodes/branch_collector.py b/nodes/branch_collector.py
index 4cabd91..9b0f5f9 100644
--- a/nodes/branch_collector.py
+++ b/nodes/branch_collector.py
@@ -1,42 +1,42 @@
 import asyncio
-import base64
-import io
-import json
-import time
+from dataclasses import dataclass
+from typing import Any
 
 import aiohttp
 import torch
 
 from ..utils.async_helpers import run_async_in_server_loop
-from ..utils.config import get_worker_timeout_seconds, is_master_delegate_only
-from ..utils.constants import HEARTBEAT_INTERVAL
-from ..utils.image import ensure_contiguous, tensor_to_pil
+from ..utils.config import is_master_delegate_only
+from ..utils.image import encode_tensor_png_data_url, ensure_contiguous
 from ..utils.logging import log
 from ..utils.network import get_client_session
+from ..utils.worker_ids import parse_enabled_worker_ids
+from .context_kwargs import parse_distributed_hidden_context
+from .hidden_inputs import build_distributed_hidden_inputs
+from .queue_wait import collect_worker_queue_results
+from .runtime_helpers import (
+    get_prompt_server_instance as _get_prompt_server_instance,
+    throw_if_processing_interrupted as _throw_if_processing_interrupted,
+)
 from .utilities import any_type
 
 
 MAX_BRANCH_OUTPUTS = 10
 
 
-def _get_prompt_server_instance():
-    import server as _server
-
-    return _server.PromptServer.instance
-
-
-def _throw_if_processing_interrupted():
-    try:
-        import comfy.model_management as model_management
-
-        model_management.throw_exception_if_processing_interrupted()
-    except Exception:
-        return
-
+@dataclass(frozen=True)
+class BranchRunContext:
+    multi_job_id: str = ""
+    is_worker: bool = False
+    master_url: str = ""
+    enabled_worker_ids: str = "[]"
+    worker_id: str = ""
+    assigned_branch: int = -1
+    delegate_only: bool = False
 
 class DistributedBranchCollector:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls: type["DistributedBranchCollector"]) -> dict[str, Any]:
         optional_inputs = {
             f"branch_{idx + 1}": (any_type,)
             for idx in range(MAX_BRANCH_OUTPUTS)
@@ -46,15 +46,10 @@ class DistributedBranchCollector:
                 "num_branches": ("INT", {"default": 2, "min": 2, "max": MAX_BRANCH_OUTPUTS, "step": 1}),
             },
             "optional": optional_inputs,
-            "hidden": {
-                "multi_job_id": ("STRING", {"default": ""}),
-                "is_worker": ("BOOLEAN", {"default": False}),
-                "master_url": ("STRING", {"default": ""}),
-                "enabled_worker_ids": ("STRING", {"default": "[]"}),
-                "worker_id": ("STRING", {"default": ""}),
-                "assigned_branch": ("INT", {"default": -1, "min": -1, "max": MAX_BRANCH_OUTPUTS - 1}),
-                "delegate_only": ("BOOLEAN", {"default": False}),
-            },
+            "hidden": build_distributed_hidden_inputs(
+                include_assigned_branch=True,
+                assigned_branch_max=MAX_BRANCH_OUTPUTS - 1,
+            ),
         }
 
     RETURN_TYPES = tuple([any_type] * MAX_BRANCH_OUTPUTS)
@@ -68,6 +63,18 @@ class DistributedBranchCollector:
             values.append(kwargs.get(f"branch_{idx + 1}", None))
         return values
 
+    def _build_run_context(self, **kwargs: Any) -> BranchRunContext:
+        try:
+            assigned_branch = int(kwargs.get("assigned_branch", -1))
+        except (TypeError, ValueError):
+            assigned_branch = -1
+
+        common_context = parse_distributed_hidden_context(kwargs)
+        return BranchRunContext(
+            assigned_branch=assigned_branch,
+            **common_context,
+        )
+
     def _resolve_local_branch_value(self, branch_values, assigned_branch):
         try:
             assigned_idx = int(assigned_branch)
@@ -108,71 +115,23 @@ class DistributedBranchCollector:
 
     def run(
         self,
-        num_branches=2,
-        multi_job_id="",
-        is_worker=False,
-        master_url="",
-        enabled_worker_ids="[]",
-        worker_id="",
-        assigned_branch=-1,
-        delegate_only=False,
-        branch_1=None,
-        branch_2=None,
-        branch_3=None,
-        branch_4=None,
-        branch_5=None,
-        branch_6=None,
-        branch_7=None,
-        branch_8=None,
-        branch_9=None,
-        branch_10=None,
-    ):
-        branch_values = [
-            branch_1,
-            branch_2,
-            branch_3,
-            branch_4,
-            branch_5,
-            branch_6,
-            branch_7,
-            branch_8,
-            branch_9,
-            branch_10,
-        ]
+        num_branches: int = 2,
+        **kwargs: Any,
+    ) -> tuple[Any, ...]:
+        branch_values = self._branch_values(**kwargs)
+        context = self._build_run_context(**kwargs)
 
-        if not multi_job_id:
+        if not context.multi_job_id:
             return self._build_outputs(branch_values, num_branches)
 
         return run_async_in_server_loop(
             self.execute(
                 branch_values,
                 num_branches=num_branches,
-                multi_job_id=multi_job_id,
-                is_worker=is_worker,
-                master_url=master_url,
-                enabled_worker_ids=enabled_worker_ids,
-                worker_id=worker_id,
-                assigned_branch=assigned_branch,
-                delegate_only=delegate_only,
+                context=context,
             )
         )
 
-    def _parse_enabled_workers(self, enabled_worker_ids):
-        try:
-            raw = json.loads(enabled_worker_ids)
-        except Exception:
-            raw = []
-
-        workers = []
-        seen = set()
-        for worker_id in raw if isinstance(raw, list) else []:
-            worker_id_str = str(worker_id)
-            if worker_id_str in seen:
-                continue
-            seen.add(worker_id_str)
-            workers.append(worker_id_str)
-        return workers
-
     def _participant_branch_slots(self, enabled_workers, delegate_mode, num_branches):
         participants = list(enabled_workers) if delegate_mode else ["master"] + list(enabled_workers)
         participant_count = len(participants)
@@ -218,16 +177,11 @@ class DistributedBranchCollector:
                 "DistributedBranchCollector tensor result must be IMAGE-like with shape [B,H,W,C] or [H,W,C]."
             )
 
-        img = tensor_to_pil(image_batch, 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(branch_idx),
-            "image": f"data:image/png;base64,{encoded_image}",
+            "image": encode_tensor_png_data_url(image_batch, 0),
             "is_last": True,
         }
 
@@ -240,18 +194,17 @@ class DistributedBranchCollector:
         ) as response:
             response.raise_for_status()
 
-    async def execute(
+    async def execute(self, *args: Any, **kwargs: Any) -> tuple[Any, ...]:
+        """Compatibility wrapper with a uniform execute() signature across collectors."""
+        return await self._execute_branch(*args, **kwargs)
+
+    async def _execute_branch(
         self,
-        branch_values,
-        num_branches=2,
-        multi_job_id="",
-        is_worker=False,
-        master_url="",
-        enabled_worker_ids="[]",
-        worker_id="",
-        assigned_branch=-1,
-        delegate_only=False,
-    ):
+        branch_values: list[Any],
+        num_branches: int = 2,
+        context: BranchRunContext | None = None,
+    ) -> tuple[Any, ...]:
+        run_context = context or BranchRunContext()
         try:
             branch_count = int(num_branches)
         except (TypeError, ValueError):
@@ -263,27 +216,27 @@ class DistributedBranchCollector:
             if branch_values[idx] is not None:
                 outputs[idx] = branch_values[idx]
 
-        local_value, local_branch_idx = self._resolve_local_branch_value(branch_values, assigned_branch)
+        local_value, local_branch_idx = self._resolve_local_branch_value(branch_values, run_context.assigned_branch)
         if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:
             outputs[local_branch_idx] = local_value
 
-        if is_worker:
+        if run_context.is_worker:
             if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:
                 try:
                     await self._send_branch_result_to_master(
                         local_value,
                         local_branch_idx,
-                        multi_job_id,
-                        master_url,
-                        worker_id,
+                        run_context.multi_job_id,
+                        run_context.master_url,
+                        run_context.worker_id,
                     )
                     outputs[local_branch_idx] = local_value
                 except Exception as exc:
                     log(f"Worker - DistributedBranchCollector failed to send branch result: {exc}")
             return tuple(outputs)
 
-        delegate_mode = bool(delegate_only or is_master_delegate_only())
-        enabled_workers = self._parse_enabled_workers(enabled_worker_ids)
+        delegate_mode = bool(run_context.delegate_only or is_master_delegate_only())
+        enabled_workers = parse_enabled_worker_ids(run_context.enabled_worker_ids)
         slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, branch_count)
         expected_worker_ids = self._expected_worker_ids_for_branches(enabled_workers, delegate_mode, branch_count)
         expected_workers = set(expected_worker_ids)
@@ -293,56 +246,42 @@ class DistributedBranchCollector:
 
         prompt_server = _get_prompt_server_instance()
         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()
+            if run_context.multi_job_id not in prompt_server.distributed_pending_jobs:
+                prompt_server.distributed_pending_jobs[run_context.multi_job_id] = asyncio.Queue()
 
-        workers_done = set()
-        base_timeout = float(get_worker_timeout_seconds())
-        slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
-        last_activity = time.time()
+        workers_done: set[str] = set()
+
+        def _handle_queue_result(result: dict[str, Any]) -> None:
+            image_index = result.get("image_index")
+            tensor = result.get("tensor")
+            try:
+                slot_idx = int(image_index)
+            except (TypeError, ValueError):
+                slot_idx = -1
+
+            if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and tensor is not None:
+                if isinstance(tensor, torch.Tensor):
+                    if tensor.is_cuda:
+                        tensor = tensor.cpu()
+                    tensor = ensure_contiguous(tensor)
+                outputs[slot_idx] = tensor
 
         try:
-            while len(workers_done) < len(expected_workers):
-                _throw_if_processing_interrupted()
-                try:
-                    async with prompt_server.distributed_jobs_lock:
-                        queue = prompt_server.distributed_pending_jobs[multi_job_id]
-                    result = await asyncio.wait_for(queue.get(), timeout=slice_timeout)
-                except asyncio.TimeoutError:
-                    if (time.time() - last_activity) < base_timeout:
-                        continue
-                    missing_workers = sorted(expected_workers - workers_done)
-                    log(
-                        "Master - DistributedBranchCollector heartbeat timeout. "
-                        f"Still waiting for workers: {missing_workers}"
-                    )
-                    break
-
-                worker_id_value = str(result.get("worker_id", ""))
-                image_index = result.get("image_index")
-                tensor = result.get("tensor")
-                is_last = bool(result.get("is_last", False))
-
-                try:
-                    slot_idx = int(image_index)
-                except (TypeError, ValueError):
-                    slot_idx = -1
-
-                if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and tensor is not None:
-                    if isinstance(tensor, torch.Tensor):
-                        if tensor.is_cuda:
-                            tensor = tensor.cpu()
-                        tensor = ensure_contiguous(tensor)
-                    outputs[slot_idx] = tensor
-
-                last_activity = time.time()
-                base_timeout = float(get_worker_timeout_seconds())
-                if is_last and worker_id_value in expected_workers:
-                    workers_done.add(worker_id_value)
+            workers_done = await collect_worker_queue_results(
+                prompt_server=prompt_server,
+                multi_job_id=run_context.multi_job_id,
+                expected_workers=expected_workers,
+                on_result=_handle_queue_result,
+                timeout_log_prefix=(
+                    "Master - DistributedBranchCollector heartbeat timeout. "
+                    "Still waiting for workers: "
+                ),
+                throw_if_interrupted=_throw_if_processing_interrupted,
+            )
 
         finally:
             async with prompt_server.distributed_jobs_lock:
-                prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
+                prompt_server.distributed_pending_jobs.pop(run_context.multi_job_id, None)
 
         missing_workers = sorted(expected_workers - workers_done)
         if missing_workers:
diff --git a/nodes/collector.py b/nodes/collector.py
index f32a8ca..bbf43c1 100644
--- a/nodes/collector.py
+++ b/nodes/collector.py
@@ -1,13 +1,15 @@
-import torch
-import io
-import json
 import asyncio
-import time
 import base64
+import io
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
 
 import aiohttp
 import server as _server
 import comfy.model_management
+import torch
 from comfy.utils import ProgressBar
 
 from ..utils.logging import debug_log, log
@@ -15,17 +17,32 @@ from ..utils.config import get_worker_timeout_seconds, load_config, is_master_de
 from ..utils.constants import HEARTBEAT_INTERVAL
 from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
 from ..utils.network import build_worker_url, get_client_session, probe_worker
+from ..utils.worker_ids import parse_enabled_worker_ids
 from ..utils.audio_payload import encode_audio_payload
 from ..utils.async_helpers import run_async_in_server_loop
+from .context_kwargs import parse_distributed_hidden_context
+from .hidden_inputs import build_distributed_hidden_inputs
 
 prompt_server = _server.PromptServer.instance
 
 
+@dataclass(frozen=True)
+class CollectorRunContext:
+    multi_job_id: str = ""
+    is_worker: bool = False
+    master_url: str = ""
+    enabled_worker_ids: str = "[]"
+    worker_batch_size: int = 1
+    worker_id: str = ""
+    pass_through: bool = False
+    delegate_only: bool = False
+
+
 class DistributedCollectorNode:
     EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
 
     @classmethod
-    def INPUT_TYPES(s):
+    def INPUT_TYPES(cls: type["DistributedCollectorNode"]) -> dict[str, Any]:
         return {
             "required": {
                 "images": ("IMAGE",),
@@ -38,50 +55,65 @@ class DistributedCollectorNode:
                 ),
             },
             "optional": { "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}),
-            },
+            "hidden": build_distributed_hidden_inputs(
+                include_worker_batch_size=True,
+                include_pass_through=True,
+            ),
         }
 
     RETURN_TYPES = ("IMAGE", "AUDIO")
     RETURN_NAMES = ("images", "audio")
     FUNCTION = "run"
     CATEGORY = "image"
-    
-    def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
+
+    def _build_run_context(self, **kwargs: Any) -> CollectorRunContext:
+        try:
+            worker_batch_size = int(kwargs.get("worker_batch_size", 1))
+        except (TypeError, ValueError):
+            worker_batch_size = 1
+
+        common_context = parse_distributed_hidden_context(kwargs)
+        return CollectorRunContext(
+            worker_batch_size=max(worker_batch_size, 1),
+            pass_through=bool(kwargs.get("pass_through", False)),
+            **common_context,
+        )
+
+    def run(
+        self,
+        images: torch.Tensor,
+        load_balance: bool = False,
+        audio: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        context = self._build_run_context(**kwargs)
         # Create empty audio if not provided
         empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
 
-        if not multi_job_id or pass_through:
-            if pass_through:
+        if not context.multi_job_id or context.pass_through:
+            if context.pass_through:
                 debug_log("Collector: pass-through mode enabled, returning images unchanged")
             return (images, audio if audio is not None else empty_audio)
 
         # Use async helper to run in server loop
         result = run_async_in_server_loop(
             self.execute(
-                images,
-                audio,
-                load_balance,
-                multi_job_id,
-                is_worker,
-                master_url,
-                enabled_worker_ids,
-                worker_batch_size,
-                worker_id,
-                delegate_only,
+                images=images,
+                audio=audio,
+                load_balance=load_balance,
+                context=context,
             )
         )
         return result
 
-    async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
+    async def send_batch_to_master(
+        self,
+        image_batch: torch.Tensor,
+        audio: dict[str, Any] | None,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+    ) -> None:
         """Send image batch to master via canonical JSON envelopes."""
         batch_size = image_batch.shape[0]
         if batch_size == 0:
@@ -235,235 +267,263 @@ class DistributedCollectorNode:
         else:
             raise ValueError("No image data collected from master or workers")
 
-    async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
-        if is_worker:
-            # Worker mode: send images and audio to master in a single batch
-            debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master")
-            await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
-            return (images, audio if audio is not None else self.EMPTY_AUDIO)
-        else:
-            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")
+    async def _ensure_pending_queue(self, multi_job_id):
+        async with prompt_server.distributed_jobs_lock:
+            if multi_job_id not in prompt_server.distributed_pending_jobs:
+                prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
+                debug_log(f"Master - Initialized queue early for job {multi_job_id}")
             else:
-                images_on_cpu = images.cpu()
-                master_batch_size = images.shape[0]
-                master_audio = audio  # Keep master's audio for later
-                debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
+                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})")
 
-                # Ensure master images are contiguous
-                images_on_cpu = ensure_contiguous(images_on_cpu)
+    async def _cleanup_pending_queue(self, multi_job_id):
+        async with prompt_server.distributed_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_jobs:
+                del prompt_server.distributed_pending_jobs[multi_job_id]
 
+    def _build_master_inputs(self, images, audio, delegate_mode, multi_job_id, num_workers):
+        if delegate_mode:
+            debug_log(
+                f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers"
+            )
+            return 0, None, None
 
-            # 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()
+        images_on_cpu = ensure_contiguous(images.cpu())
+        master_batch_size = images.shape[0]
+        debug_log(
+            f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers..."
+        )
+        return master_batch_size, images_on_cpu, audio
 
-            # 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:
+    async def _probe_missing_workers_busy(self, missing_workers):
+        any_busy = False
+        try:
+            cfg = load_config()
+            cfg_workers = cfg.get('workers', [])
+            for wid in list(missing_workers):
+                wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
+                if not wrec:
+                    debug_log(f"Collector probe: worker {wid} not found in config")
+                    continue
+                worker_url = build_worker_url(wrec)
+                try:
+                    payload = await probe_worker(worker_url, timeout=2.0)
+                    queue_remaining = None
+                    if payload is not None:
+                        queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
                     debug_log(
-                        f"Master - Ignoring completion from unexpected worker {done_worker_id} for job {multi_job_id}"
+                        "Collector probe: worker "
+                        f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
                     )
-                    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"
-                            )
+                    if payload is not None and queue_remaining and queue_remaining > 0:
+                        any_busy = True
                         log(
-                            f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} "
-                            f"(elapsed={elapsed:.1f}s)"
+                            f"Master - Probe grace: worker {wid} appears busy "
+                            f"(queue_remaining={queue_remaining}). Continuing to wait."
                         )
-
-                        # 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]
+                except Exception as e:
+                    debug_log(f"Collector probe failed for worker {wid}: {e}")
+        except Exception as e:
+            debug_log(f"Collector probe setup error: {e}")
+        return any_busy
 
-            try:
-                combined = self._reorder_and_combine_tensors(
-                    worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
+    async def _drain_remaining_queue_items(
+        self,
+        multi_job_id: str,
+        worker_images: dict[str, dict[int, torch.Tensor]],
+        mark_worker_done: Callable[[str], None],
+    ) -> int:
+        collected_count = 0
+        async with prompt_server.distributed_jobs_lock:
+            if multi_job_id not in prompt_server.distributed_pending_jobs:
+                log(f"Master - Queue {multi_job_id} no longer exists!")
+                return collected_count
+
+            final_q = prompt_server.distributed_pending_jobs[multi_job_id]
+            remaining_items = []
+            while not final_q.empty():
+                try:
+                    remaining_items.append(final_q.get_nowait())
+                except asyncio.QueueEmpty:
+                    break
+
+        for item in remaining_items:
+            worker_id = item['worker_id']
+            is_last = item.get('is_last', False)
+            collected_count += self._store_worker_result(worker_images, item)
+            if is_last:
+                mark_worker_done(worker_id)
+        return collected_count
+
+    async def _collect_worker_results(
+        self,
+        multi_job_id: str,
+        enabled_workers: list[str],
+        expected_workers: set[str],
+        worker_images: dict[str, dict[int, torch.Tensor]],
+        worker_audio: dict[str, dict[str, Any]],
+    ) -> int:
+        num_workers = len(expected_workers)
+        workers_done = set()
+        collected_count = 0
+        base_timeout = float(get_worker_timeout_seconds())
+        slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
+        last_activity = time.time()
+        progress = ProgressBar(num_workers)
+
+        def mark_worker_done(done_worker_id: str) -> None:
+            done_worker_id = str(done_worker_id)
+            if done_worker_id not in expected_workers:
+                debug_log(
+                    f"Master - Ignoring completion from unexpected worker {done_worker_id} for job {multi_job_id}"
                 )
-                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
+            if done_worker_id in workers_done:
+                debug_log(
+                    f"Master - Ignoring duplicate completion from worker {done_worker_id} for job {multi_job_id}"
+                )
+                return
+            workers_done.add(done_worker_id)
+            progress.update(1)
 
-                # Combine audio from master and workers
-                combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
+        while len(workers_done) < num_workers:
+            comfy.model_management.throw_exception_if_processing_interrupted()
+            try:
+                async with prompt_server.distributed_jobs_lock:
+                    q = prompt_server.distributed_pending_jobs[multi_job_id]
 
-                return (combined, combined_audio)
-            except Exception as e:
-                log(f"Master - Error combining images: {e}")
-                # Return just the master images as fallback
-                return (images, audio if audio is not None else self.EMPTY_AUDIO)
+                result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
+                worker_id = result['worker_id']
+                is_last = result.get('is_last', False)
+                collected_count += self._store_worker_result(worker_images, result)
+                debug_log(
+                    f"Master - Got canonical result from worker {worker_id}, "
+                    f"image {result.get('image_index', 0)}, is_last={is_last}"
+                )
+
+                result_audio = result.get('audio')
+                if result_audio is not None:
+                    worker_audio[worker_id] = result_audio
+                    debug_log(f"Master - Got audio from worker {worker_id}")
+
+                last_activity = time.time()
+                base_timeout = float(get_worker_timeout_seconds())
+                if is_last:
+                    mark_worker_done(worker_id)
+            except asyncio.TimeoutError:
+                if (time.time() - last_activity) < base_timeout:
+                    comfy.model_management.throw_exception_if_processing_interrupted()
+                    continue
+
+                comfy.model_management.throw_exception_if_processing_interrupted()
+                missing_workers = set(str(w) for w in enabled_workers) - workers_done
+                elapsed = time.time() - last_activity
+                for missing_worker_id in sorted(missing_workers):
+                    log(
+                        "Master - Heartbeat timeout: "
+                        f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
+                    )
+                log(
+                    f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} "
+                    f"(elapsed={elapsed:.1f}s)"
+                )
+
+                if await self._probe_missing_workers_busy(missing_workers):
+                    last_activity = time.time()
+                    base_timeout = float(get_worker_timeout_seconds())
+                    continue
+
+                collected_count += await self._drain_remaining_queue_items(
+                    multi_job_id,
+                    worker_images,
+                    mark_worker_done,
+                )
+                break
+
+        return collected_count
+
+    async def _execute_master(
+        self,
+        images,
+        audio,
+        multi_job_id,
+        enabled_worker_ids,
+        delegate_only,
+    ):
+        delegate_mode = delegate_only or is_master_delegate_only()
+        enabled_workers = parse_enabled_worker_ids(enabled_worker_ids)
+        expected_workers = set(enabled_workers)
+        num_workers = len(expected_workers)
+        if num_workers == 0:
+            return (images, audio if audio is not None else self.EMPTY_AUDIO)
+
+        await self._ensure_pending_queue(multi_job_id)
+        master_batch_size, images_on_cpu, master_audio = self._build_master_inputs(
+            images, audio, delegate_mode, multi_job_id, num_workers
+        )
+
+        worker_images = {}
+        worker_audio = {}
+        try:
+            await self._collect_worker_results(
+                multi_job_id=multi_job_id,
+                enabled_workers=enabled_workers,
+                expected_workers=expected_workers,
+                worker_images=worker_images,
+                worker_audio=worker_audio,
+            )
+        except comfy.model_management.InterruptProcessingException:
+            await self._cleanup_pending_queue(multi_job_id)
+            raise
+
+        await self._cleanup_pending_queue(multi_job_id)
+        try:
+            combined = self._reorder_and_combine_tensors(
+                worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
+            )
+            debug_log(
+                f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
+                f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})"
+            )
+            combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
+            return (combined, combined_audio)
+        except Exception as e:
+            log(f"Master - Error combining images: {e}")
+            return (images, audio if audio is not None else self.EMPTY_AUDIO)
+
+    async def execute(self, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, dict[str, Any]]:
+        """Compatibility wrapper with a uniform execute() signature across collectors."""
+        return await self._execute_collector(*args, **kwargs)
+
+    async def _execute_collector(
+        self,
+        images: torch.Tensor,
+        audio: dict[str, Any] | None,
+        load_balance: bool = False,
+        context: CollectorRunContext | None = None,
+    ) -> tuple[torch.Tensor, dict[str, Any]]:
+        run_context = context or CollectorRunContext()
+        _ = load_balance
+        _ = run_context.worker_batch_size
+        if run_context.is_worker:
+            debug_log(
+                "Worker - Job "
+                f"{run_context.multi_job_id} complete. Sending {images.shape[0]} image(s) to master"
+            )
+            await self.send_batch_to_master(
+                images,
+                audio,
+                run_context.multi_job_id,
+                run_context.master_url,
+                run_context.worker_id,
+            )
+            return (images, audio if audio is not None else self.EMPTY_AUDIO)
+
+        return await self._execute_master(
+            images=images,
+            audio=audio,
+            multi_job_id=run_context.multi_job_id,
+            enabled_worker_ids=run_context.enabled_worker_ids,
+            delegate_only=run_context.delegate_only,
+        )
diff --git a/nodes/context_kwargs.py b/nodes/context_kwargs.py
new file mode 100644
index 0000000..5eb3804
--- /dev/null
+++ b/nodes/context_kwargs.py
@@ -0,0 +1,13 @@
+from typing import Any
+
+
+def parse_distributed_hidden_context(kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Parse common distributed hidden inputs from node kwargs."""
+    return {
+        "multi_job_id": str(kwargs.get("multi_job_id", "")),
+        "is_worker": bool(kwargs.get("is_worker", False)),
+        "master_url": str(kwargs.get("master_url", "")),
+        "enabled_worker_ids": str(kwargs.get("enabled_worker_ids", "[]")),
+        "worker_id": str(kwargs.get("worker_id", "")),
+        "delegate_only": bool(kwargs.get("delegate_only", False)),
+    }
diff --git a/nodes/distributed_upscale.py b/nodes/distributed_upscale.py
index 2e92e44..c71503b 100644
--- a/nodes/distributed_upscale.py
+++ b/nodes/distributed_upscale.py
@@ -1,25 +1,28 @@
 import json
 import math
 from functools import wraps
+from typing import Any, Callable
 
 import comfy.samplers
 
 from ..utils.logging import debug_log, log
 from ..utils.async_helpers import run_async_in_server_loop
 from ..upscale.job_store import ensure_tile_jobs_initialized
+from .hidden_inputs import build_distributed_hidden_inputs
 
 from ..upscale.tile_ops import TileOpsMixin
 from ..upscale.result_collector import ResultCollectorMixin
 from ..upscale.worker_comms import WorkerCommsMixin
 from ..upscale.job_state import JobStateMixin
+from ..upscale.processing_args import UpscaleCoreArgs
 from ..upscale.modes.single_gpu import SingleGpuModeMixin
 from ..upscale.modes.static import StaticModeMixin
 from ..upscale.modes.dynamic import DynamicModeMixin
 
-def sync_wrapper(async_func):
+def sync_wrapper(async_func: Callable[..., Any]) -> Callable[..., Any]:
     """Decorator to wrap async methods for synchronous execution."""
     @wraps(async_func)
-    def sync_func(self, *args, **kwargs):
+    def sync_func(self, *args: Any, **kwargs: Any) -> Any:
         # Use run_async_in_server_loop for ComfyUI compatibility
         return run_async_in_server_loop(
             async_func(self, *args, **kwargs),
@@ -81,7 +84,14 @@ class UltimateSDUpscaleDistributed(
         debug_log("UltimateSDUpscaleDistributed - Node initialized")
 
     @classmethod
-    def INPUT_TYPES(s):
+    def INPUT_TYPES(cls: type["UltimateSDUpscaleDistributed"]) -> dict[str, Any]:
+        hidden_inputs = build_distributed_hidden_inputs()
+        hidden_inputs.update(
+            {
+                "tile_indices": ("STRING", {"default": ""}),  # Unused - kept for compatibility
+                "dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
+            }
+        )
         return {
             "required": {
                 "upscaled_image": ("IMAGE",),
@@ -102,15 +112,7 @@ class UltimateSDUpscaleDistributed(
                 "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}),
-            },
+            "hidden": hidden_inputs,
         }
 
     RETURN_TYPES = ("IMAGE",)
@@ -122,12 +124,47 @@ class UltimateSDUpscaleDistributed(
         """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):
+    def run(
+        self,
+        upscaled_image: Any,
+        model: Any,
+        positive: Any,
+        negative: Any,
+        vae: Any,
+        seed: int,
+        steps: int,
+        cfg: float,
+        sampler_name: str,
+        scheduler: str,
+        denoise: float,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        tiled_decode: bool,
+        multi_job_id: str = "",
+        is_worker: bool = False,
+        master_url: str = "",
+        enabled_worker_ids: str = "[]",
+        worker_id: str = "",
+        tile_indices: str = "",
+        dynamic_threshold: int = 8,
+    ) -> tuple[Any, ...]:
         """Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler."""
+        core_args = UpscaleCoreArgs(
+            model=model,
+            positive=positive,
+            negative=negative,
+            vae=vae,
+            seed=seed,
+            steps=steps,
+            cfg=cfg,
+            sampler_name=sampler_name,
+            scheduler=scheduler,
+            denoise=denoise,
+            tiled_decode=tiled_decode,
+        )
         # Strict WAN/FLOW batching: error if batch is not 4n+1 (except allow 1)
         try:
             batch_size = int(getattr(upscaled_image, 'shape', [1])[0])
@@ -142,9 +179,15 @@ class UltimateSDUpscaleDistributed(
             )
         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)
+            return self.process_single_gpu(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+            )
         
         if is_worker:
             # Worker mode: process tiles synchronously
@@ -155,24 +198,64 @@ class UltimateSDUpscaleDistributed(
                                       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)
+            return self.process_master(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+                multi_job_id=multi_job_id,
+                enabled_worker_ids=enabled_worker_ids,
+                dynamic_threshold=dynamic_threshold,
+            )
 
-    def process_worker(self, upscaled_image, 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):
+    def process_worker(
+        self,
+        upscaled_image: Any,
+        model: Any,
+        positive: Any,
+        negative: Any,
+        vae: Any,
+        seed: int,
+        steps: int,
+        cfg: float,
+        sampler_name: str,
+        scheduler: str,
+        denoise: float,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        tiled_decode: bool,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+        enabled_worker_ids: str,
+        dynamic_threshold: int,
+    ) -> tuple[Any, ...]:
         """Unified worker processing - handles both static and dynamic modes."""
+        core_args = UpscaleCoreArgs(
+            model=model,
+            positive=positive,
+            negative=negative,
+            vae=vae,
+            seed=seed,
+            steps=steps,
+            cfg=cfg,
+            sampler_name=sampler_name,
+            scheduler=scheduler,
+            denoise=denoise,
+            tiled_decode=tiled_decode,
+        )
         # Get batch size to determine mode
         batch_size = upscaled_image.shape[0]
         
         # Ensure mode consistency across master/workers via shared threshold
         # Determine mode (must match master's logic)
-        enabled_workers = json.loads(enabled_worker_ids)
+        enabled_workers = _parse_enabled_worker_ids(enabled_worker_ids)
         num_workers = len(enabled_workers)
         # Compute number of tiles for this image to decide if tile distribution makes sense
         _, height, width, _ = upscaled_image.shape
@@ -188,11 +271,20 @@ class UltimateSDUpscaleDistributed(
         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)
+            return self.process_worker_dynamic(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+                multi_job_id=multi_job_id,
+                master_url=master_url,
+                worker_id=worker_id,
+                enabled_worker_ids=enabled_worker_ids,
+                dynamic_threshold=dynamic_threshold,
+            )
         
         # Static mode - enhanced with health monitoring and retry logic
         return self._process_worker_static_sync(upscaled_image, model, positive, negative, vae,
@@ -201,11 +293,19 @@ class UltimateSDUpscaleDistributed(
                                                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):
+    def process_master(
+        self,
+        upscaled_image: Any,
+        core_args: UpscaleCoreArgs,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        multi_job_id: str,
+        enabled_worker_ids: str,
+        dynamic_threshold: int,
+    ) -> tuple[Any, ...]:
         """Unified master processing with enhanced monitoring and failure handling."""
         # Round tile dimensions
         tile_width = self.round_to_multiple(tile_width)
@@ -224,7 +324,7 @@ class UltimateSDUpscaleDistributed(
         )
         
         # Parse enabled workers
-        enabled_workers = json.loads(enabled_worker_ids)
+        enabled_workers = _parse_enabled_worker_ids(enabled_worker_ids)
         num_workers = len(enabled_workers)
         
         # Determine processing mode
@@ -238,22 +338,35 @@ class UltimateSDUpscaleDistributed(
 
         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)
+            return self.process_single_gpu(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+            )
         
         elif mode == "dynamic":
             # Dynamic mode for large batches
-            return self.process_master_dynamic(upscaled_image, 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)
+            return self.process_master_dynamic(
+                upscaled_image=upscaled_image,
+                core_args=core_args,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                mask_blur=mask_blur,
+                force_uniform_tiles=force_uniform_tiles,
+                multi_job_id=multi_job_id,
+                enabled_workers=enabled_workers,
+            )
         
         # Static mode - enhanced with unified job management
-        return self._process_master_static_sync(upscaled_image, model, positive, negative, vae,
-                                               seed, steps, cfg, sampler_name, scheduler, denoise,
+        return self._process_master_static_sync(upscaled_image, core_args.model, core_args.positive, core_args.negative, core_args.vae,
+                                               core_args.seed, core_args.steps, core_args.cfg, core_args.sampler_name, core_args.scheduler, core_args.denoise,
                                                tile_width, tile_height, padding, mask_blur,
-                                               force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
+                                               force_uniform_tiles, core_args.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:
diff --git a/nodes/hidden_inputs.py b/nodes/hidden_inputs.py
new file mode 100644
index 0000000..8128805
--- /dev/null
+++ b/nodes/hidden_inputs.py
@@ -0,0 +1,31 @@
+from typing import Any
+
+
+def build_distributed_hidden_inputs(
+    *,
+    include_assigned_branch: bool = False,
+    assigned_branch_max: int = 9,
+    include_worker_batch_size: bool = False,
+    include_pass_through: bool = False,
+) -> dict[str, tuple[str, dict[str, Any]]]:
+    hidden_inputs: dict[str, tuple[str, dict[str, Any]]] = {
+        "multi_job_id": ("STRING", {"default": ""}),
+        "is_worker": ("BOOLEAN", {"default": False}),
+        "master_url": ("STRING", {"default": ""}),
+        "enabled_worker_ids": ("STRING", {"default": "[]"}),
+        "worker_id": ("STRING", {"default": ""}),
+        "delegate_only": ("BOOLEAN", {"default": False}),
+    }
+    if include_assigned_branch:
+        hidden_inputs["assigned_branch"] = (
+            "INT",
+            {"default": -1, "min": -1, "max": assigned_branch_max},
+        )
+    if include_worker_batch_size:
+        hidden_inputs["worker_batch_size"] = (
+            "INT",
+            {"default": 1, "min": 1, "max": 1024},
+        )
+    if include_pass_through:
+        hidden_inputs["pass_through"] = ("BOOLEAN", {"default": False})
+    return hidden_inputs
diff --git a/nodes/list_collector.py b/nodes/list_collector.py
index 031cc2b..4a275e8 100644
--- a/nodes/list_collector.py
+++ b/nodes/list_collector.py
@@ -1,34 +1,33 @@
 import asyncio
-import base64
-import io
-import json
-import time
+from dataclasses import dataclass
+from typing import Any
 
 import aiohttp
 import torch
 
 from ..utils.async_helpers import run_async_in_server_loop
-from ..utils.config import get_worker_timeout_seconds, is_master_delegate_only
-from ..utils.constants import HEARTBEAT_INTERVAL
-from ..utils.image import ensure_contiguous, tensor_to_pil
+from ..utils.config import is_master_delegate_only
+from ..utils.image import encode_tensor_png_data_url, ensure_contiguous
 from ..utils.logging import debug_log, log
 from ..utils.network import get_client_session
+from ..utils.worker_ids import parse_enabled_worker_ids
+from .context_kwargs import parse_distributed_hidden_context
+from .hidden_inputs import build_distributed_hidden_inputs
+from .queue_wait import collect_worker_queue_results
+from .runtime_helpers import (
+    get_prompt_server_instance as _get_prompt_server_instance,
+    throw_if_processing_interrupted as _throw_if_processing_interrupted,
+)
 
 
-def _get_prompt_server_instance():
-    import server as _server
-
-    return _server.PromptServer.instance
-
-
-def _throw_if_processing_interrupted():
-    try:
-        import comfy.model_management as model_management
-
-        model_management.throw_exception_if_processing_interrupted()
-    except Exception:
-        # Tests and non-Comfy contexts may not have interruption support.
-        return
+@dataclass(frozen=True)
+class ListCollectorRunContext:
+    multi_job_id: str = ""
+    is_worker: bool = False
+    master_url: str = ""
+    enabled_worker_ids: str = "[]"
+    worker_id: str = ""
+    delegate_only: bool = False
 
 
 class DistributedListCollector:
@@ -36,19 +35,12 @@ class DistributedListCollector:
     OUTPUT_IS_LIST = (True,)
 
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls: type["DistributedListCollector"]) -> dict[str, Any]:
         return {
             "required": {
                 "images": ("IMAGE",),
             },
-            "hidden": {
-                "multi_job_id": ("STRING", {"default": ""}),
-                "is_worker": ("BOOLEAN", {"default": False}),
-                "master_url": ("STRING", {"default": ""}),
-                "enabled_worker_ids": ("STRING", {"default": "[]"}),
-                "worker_id": ("STRING", {"default": ""}),
-                "delegate_only": ("BOOLEAN", {"default": False}),
-            },
+            "hidden": build_distributed_hidden_inputs(),
         }
 
     RETURN_TYPES = ("IMAGE",)
@@ -56,40 +48,40 @@ class DistributedListCollector:
     FUNCTION = "run"
     CATEGORY = "image"
 
+    def _build_run_context(self, **kwargs: Any) -> ListCollectorRunContext:
+        return ListCollectorRunContext(**parse_distributed_hidden_context(kwargs))
+
     def run(
         self,
-        images,
-        multi_job_id="",
-        is_worker=False,
-        master_url="",
-        enabled_worker_ids="[]",
-        worker_id="",
-        delegate_only=False,
-    ):
+        images: list[Any] | torch.Tensor | None,
+        **kwargs: Any,
+    ) -> tuple[list[Any]]:
+        context = self._build_run_context(**kwargs)
         image_list = self._normalize_image_list(images)
-        if not multi_job_id:
+        if not context.multi_job_id:
             return (image_list,)
 
         return run_async_in_server_loop(
             self.execute(
                 image_list,
-                multi_job_id=multi_job_id,
-                is_worker=is_worker,
-                master_url=master_url,
-                enabled_worker_ids=enabled_worker_ids,
-                worker_id=worker_id,
-                delegate_only=delegate_only,
+                context=context,
             )
         )
 
-    def _normalize_image_list(self, images):
+    def _normalize_image_list(self, images: list[Any] | torch.Tensor | None) -> list[Any]:
         if images is None:
             return []
         if isinstance(images, list):
             return list(images)
         return [images]
 
-    async def send_list_to_master(self, image_list, multi_job_id, master_url, worker_id):
+    async def send_list_to_master(
+        self,
+        image_list: list[Any],
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+    ) -> None:
         if not image_list:
             return
 
@@ -111,15 +103,11 @@ class DistributedListCollector:
         url = f"{master_url}/distributed/job_complete"
 
         for send_pos, (image_index, image_batch) in enumerate(sendable):
-            img = tensor_to_pil(image_batch, 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(image_index),
-                "image": f"data:image/png;base64,{encoded_image}",
+                "image": encode_tensor_png_data_url(image_batch, 0),
                 "is_last": bool(send_pos == len(sendable) - 1),
             }
             try:
@@ -145,22 +133,6 @@ class DistributedListCollector:
         worker_images[worker_id][int(image_index)] = tensor
         return 1
 
-    def _parse_enabled_workers(self, enabled_worker_ids):
-        try:
-            raw = json.loads(enabled_worker_ids)
-        except Exception:
-            raw = []
-
-        workers = []
-        seen = set()
-        for worker_id in raw if isinstance(raw, list) else []:
-            worker_id_str = str(worker_id)
-            if worker_id_str in seen:
-                continue
-            seen.add(worker_id_str)
-            workers.append(worker_id_str)
-        return workers
-
     def _reorder_and_combine_list(self, worker_images, worker_order, master_images, delegate_mode):
         ordered = []
         if not delegate_mode:
@@ -192,25 +164,33 @@ class DistributedListCollector:
 
         return combined
 
-    async def execute(
+    async def execute(self, *args: Any, **kwargs: Any) -> tuple[list[Any]]:
+        """Compatibility wrapper with a uniform execute() signature across collectors."""
+        return await self._execute_list(*args, **kwargs)
+
+    async def _execute_list(
         self,
-        images,
-        multi_job_id="",
-        is_worker=False,
-        master_url="",
-        enabled_worker_ids="[]",
-        worker_id="",
-        delegate_only=False,
-    ):
+        images: list[Any] | torch.Tensor | None,
+        context: ListCollectorRunContext | None = None,
+    ) -> tuple[list[Any]]:
+        run_context = context or ListCollectorRunContext()
         image_list = self._normalize_image_list(images)
 
-        if is_worker:
-            debug_log(f"Worker - Job {multi_job_id} complete. Sending {len(image_list)} image list item(s) to master")
-            await self.send_list_to_master(image_list, multi_job_id, master_url, worker_id)
+        if run_context.is_worker:
+            debug_log(
+                "Worker - Job "
+                f"{run_context.multi_job_id} complete. Sending {len(image_list)} image list item(s) to master"
+            )
+            await self.send_list_to_master(
+                image_list,
+                run_context.multi_job_id,
+                run_context.master_url,
+                run_context.worker_id,
+            )
             return (image_list,)
 
-        delegate_mode = bool(delegate_only or is_master_delegate_only())
-        enabled_workers = self._parse_enabled_workers(enabled_worker_ids)
+        delegate_mode = bool(run_context.delegate_only or is_master_delegate_only())
+        enabled_workers = parse_enabled_worker_ids(run_context.enabled_worker_ids)
         expected_workers = set(enabled_workers)
         if not expected_workers:
             return (image_list,)
@@ -218,49 +198,41 @@ class DistributedListCollector:
         prompt_server = _get_prompt_server_instance()
 
         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 list collector job {multi_job_id}")
+            if run_context.multi_job_id not in prompt_server.distributed_pending_jobs:
+                prompt_server.distributed_pending_jobs[run_context.multi_job_id] = asyncio.Queue()
+                debug_log(
+                    "Master - Initialized queue early for list collector job "
+                    f"{run_context.multi_job_id}"
+                )
 
         master_images = [] if delegate_mode else image_list
 
         worker_images = {}
-        workers_done = set()
-        base_timeout = float(get_worker_timeout_seconds())
-        slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
-        last_activity = time.time()
+        workers_done: set[str] = set()
+
+        def _handle_queue_result(result: dict[str, Any]) -> None:
+            self._store_worker_result(worker_images, result)
 
         try:
-            while len(workers_done) < len(expected_workers):
-                _throw_if_processing_interrupted()
-                try:
-                    async with prompt_server.distributed_jobs_lock:
-                        queue = prompt_server.distributed_pending_jobs[multi_job_id]
-                    result = await asyncio.wait_for(queue.get(), timeout=slice_timeout)
-                except asyncio.TimeoutError:
-                    if (time.time() - last_activity) < base_timeout:
-                        continue
-                    missing_workers = sorted(expected_workers - workers_done)
-                    log(
-                        "Master - List collector heartbeat timeout. "
-                        f"Still waiting for workers: {missing_workers}"
-                    )
-                    break
-
-                worker_id_value = str(result.get("worker_id", ""))
-                is_last = bool(result.get("is_last", False))
-                self._store_worker_result(worker_images, result)
-                last_activity = time.time()
-                base_timeout = float(get_worker_timeout_seconds())
-                if is_last and worker_id_value in expected_workers:
-                    workers_done.add(worker_id_value)
+            workers_done = await collect_worker_queue_results(
+                prompt_server=prompt_server,
+                multi_job_id=run_context.multi_job_id,
+                expected_workers=expected_workers,
+                on_result=_handle_queue_result,
+                timeout_log_prefix=(
+                    "Master - List collector heartbeat timeout. "
+                    "Still waiting for workers: "
+                ),
+                throw_if_interrupted=_throw_if_processing_interrupted,
+            )
 
         finally:
             async with prompt_server.distributed_jobs_lock:
-                prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
+                prompt_server.distributed_pending_jobs.pop(run_context.multi_job_id, None)
 
         combined = self._reorder_and_combine_list(worker_images, enabled_workers, master_images, delegate_mode)
         debug_log(
-            f"Master - List collector job {multi_job_id} complete. Combined {len(combined)} image list item(s)."
+            "Master - List collector job "
+            f"{run_context.multi_job_id} complete. Combined {len(combined)} image list item(s)."
         )
         return (combined,)
diff --git a/nodes/list_splitter.py b/nodes/list_splitter.py
index e5c15c6..7b3cf95 100644
--- a/nodes/list_splitter.py
+++ b/nodes/list_splitter.py
@@ -1,4 +1,6 @@
 import asyncio
+from dataclasses import dataclass
+from typing import Any
 
 import aiohttp
 
@@ -10,12 +12,22 @@ from ..utils.network import get_client_session
 from .utilities import _chunk_bounds
 
 
+@dataclass(frozen=True)
+class ListSplitterRunContext:
+    participant_index: int = 0
+    total_participants: int = 1
+    multi_job_id: str = ""
+    is_worker: bool = False
+    master_url: str = ""
+    worker_id: str = ""
+
+
 class DistributedListSplitter:
     INPUT_IS_LIST = True
     OUTPUT_IS_LIST = (True,)
 
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls: type["DistributedListSplitter"]) -> dict[str, Any]:
         return {
             "required": {
                 "images": ("IMAGE",),
@@ -36,31 +48,50 @@ class DistributedListSplitter:
     FUNCTION = "split"
     CATEGORY = "image"
 
+    def _build_run_context(self, **kwargs: Any) -> ListSplitterRunContext:
+        try:
+            participant_index = int(kwargs.get("participant_index", 0))
+        except (TypeError, ValueError):
+            participant_index = 0
+        try:
+            total_participants = int(kwargs.get("total_participants", 1))
+        except (TypeError, ValueError):
+            total_participants = 1
+
+        return ListSplitterRunContext(
+            participant_index=participant_index,
+            total_participants=max(total_participants, 1),
+            multi_job_id=str(kwargs.get("multi_job_id", "")),
+            is_worker=bool(kwargs.get("is_worker", False)),
+            master_url=str(kwargs.get("master_url", "")),
+            worker_id=str(kwargs.get("worker_id", "")),
+        )
+
     def split(
         self,
-        images,
-        mode="static",
-        participant_index=0,
-        total_participants=1,
-        multi_job_id="",
-        is_worker=False,
-        master_url="",
-        worker_id="",
-    ):
+        images: list[Any] | Any | None,
+        mode: str = "static",
+        **kwargs: Any,
+    ) -> tuple[list[Any]]:
+        context = self._build_run_context(**kwargs)
         if images is None:
             return ([],)
 
         image_list = images if isinstance(images, list) else [images]
-        if mode != "dynamic" or not multi_job_id:
-            return self._split_static(image_list, participant_index, total_participants)
+        if mode != "dynamic" or not context.multi_job_id:
+            return self._split_static(
+                image_list,
+                context.participant_index,
+                context.total_participants,
+            )
 
         return run_async_in_server_loop(
             self._split_dynamic(
                 image_list,
-                multi_job_id,
-                bool(is_worker),
-                master_url,
-                worker_id,
+                context.multi_job_id,
+                context.is_worker,
+                context.master_url,
+                context.worker_id,
             )
         )
 
@@ -161,7 +192,8 @@ class DistributedListSplitter:
 
             try:
                 pulled.append(int(item_idx))
-            except (TypeError, ValueError):
+            except (TypeError, ValueError) as exc:
+                debug_log(f"DistributedListSplitter dynamic mode: invalid item index '{item_idx}' ignored: {exc}")
                 continue
 
         return pulled
diff --git a/nodes/queue_wait.py b/nodes/queue_wait.py
new file mode 100644
index 0000000..b4ae006
--- /dev/null
+++ b/nodes/queue_wait.py
@@ -0,0 +1,46 @@
+import asyncio
+import time
+from collections.abc import Callable
+from typing import Any
+
+from ..utils.config import get_worker_timeout_seconds
+from ..utils.constants import HEARTBEAT_INTERVAL
+from ..utils.logging import log
+
+
+async def collect_worker_queue_results(
+    *,
+    prompt_server: Any,
+    multi_job_id: str,
+    expected_workers: set[str],
+    on_result: Callable[[dict[str, Any]], None],
+    timeout_log_prefix: str,
+    throw_if_interrupted: Callable[[], None],
+) -> set[str]:
+    workers_done: set[str] = set()
+    base_timeout = float(get_worker_timeout_seconds())
+    slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
+    last_activity = time.time()
+
+    while len(workers_done) < len(expected_workers):
+        throw_if_interrupted()
+        try:
+            async with prompt_server.distributed_jobs_lock:
+                queue = prompt_server.distributed_pending_jobs[multi_job_id]
+            result = await asyncio.wait_for(queue.get(), timeout=slice_timeout)
+        except asyncio.TimeoutError:
+            if (time.time() - last_activity) < base_timeout:
+                continue
+            missing_workers = sorted(expected_workers - workers_done)
+            log(f"{timeout_log_prefix}{missing_workers}")
+            break
+
+        on_result(result)
+        worker_id_value = str(result.get("worker_id", ""))
+        is_last = bool(result.get("is_last", False))
+        last_activity = time.time()
+        base_timeout = float(get_worker_timeout_seconds())
+        if is_last and worker_id_value in expected_workers:
+            workers_done.add(worker_id_value)
+
+    return workers_done
diff --git a/nodes/runtime_helpers.py b/nodes/runtime_helpers.py
new file mode 100644
index 0000000..94aebee
--- /dev/null
+++ b/nodes/runtime_helpers.py
@@ -0,0 +1,16 @@
+from typing import Any
+
+
+def get_prompt_server_instance() -> Any:
+    import server as _server
+
+    return _server.PromptServer.instance
+
+
+def throw_if_processing_interrupted() -> None:
+    try:
+        import comfy.model_management as model_management
+
+        model_management.throw_exception_if_processing_interrupted()
+    except Exception:
+        return
diff --git a/nodes/utilities.py b/nodes/utilities.py
index 3fda3de..bc43bfc 100644
--- a/nodes/utilities.py
+++ b/nodes/utilities.py
@@ -1,5 +1,6 @@
 import torch
 import json
+from typing import Any
 
 from ..utils.logging import debug_log, log
 
@@ -28,7 +29,7 @@ class DistributedSeed:
     """
     
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls: type["DistributedSeed"]) -> dict[str, Any]:
         return {
             "required": {
                 "seed": ("INT", {
@@ -49,7 +50,12 @@ class DistributedSeed:
     FUNCTION = "distribute"
     CATEGORY = "utils"
     
-    def distribute(self, seed, is_worker=False, worker_id=""):
+    def distribute(
+        self,
+        seed: int,
+        is_worker: bool = False,
+        worker_id: str = "",
+    ) -> tuple[int]:
         if not is_worker:
             # Master node: pass through original values
             debug_log(f"Distributor - Master: seed={seed}")
@@ -92,7 +98,7 @@ class DistributedValue:
     """
 
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls: type["DistributedValue"]) -> dict[str, Any]:
         return {
             "required": {
                 "default_value": ("STRING", {"default": ""}),
@@ -110,7 +116,7 @@ class DistributedValue:
     CATEGORY = "utils"
 
     @staticmethod
-    def _coerce(value, value_type):
+    def _coerce(value: Any, value_type: str) -> Any:
         """Convert a string value to the requested type."""
         if value_type == "INT":
             return int(float(value))
@@ -119,14 +125,31 @@ class DistributedValue:
         return value  # STRING and COMBO stay as strings
 
     @staticmethod
-    def _coerce_safe(value, value_type):
+    def _coerce_safe(value: Any, value_type: str) -> Any:
         """Best-effort coercion with graceful fallback to original value."""
         try:
             return DistributedValue._coerce(value, value_type)
         except (TypeError, ValueError):
             return value
 
-    def distribute(self, default_value, worker_values="{}", is_worker=False, worker_id=""):
+    @staticmethod
+    def _infer_value_type(value: Any) -> str:
+        """Infer coercion type from the provided default value."""
+        if isinstance(value, bool):
+            return "STRING"
+        if isinstance(value, int):
+            return "INT"
+        if isinstance(value, float):
+            return "FLOAT"
+        return "STRING"
+
+    def distribute(
+        self,
+        default_value: Any,
+        worker_values: str | dict[str, Any] = "{}",
+        is_worker: bool = False,
+        worker_id: str = "",
+    ) -> tuple[Any]:
         values = {}
         value_type = "STRING"
 
@@ -138,7 +161,11 @@ class DistributedValue:
             debug_log(f"DistributedValue - Error parsing worker_values: {e}")
             values = {}
 
-        value_type = values.get("_type", "STRING")
+        inferred_type = self._infer_value_type(default_value)
+        value_type = values.get("_type", inferred_type)
+        if value_type not in {"STRING", "COMBO", "INT", "FLOAT"}:
+            value_type = inferred_type
+        values["_type"] = value_type
         coerced_default = self._coerce_safe(default_value, value_type)
 
         if not is_worker:
@@ -163,7 +190,7 @@ class DistributedValue:
 
 class DistributedModelName:
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls: type["DistributedModelName"]) -> dict[str, Any]:
         return {
             "required": {
                 "text": ("STRING", {"default": ""}),
@@ -180,7 +207,7 @@ class DistributedModelName:
     OUTPUT_NODE = True
     CATEGORY = "utils"
 
-    def _stringify(self, value):
+    def _stringify(self, value: Any) -> str:
         if isinstance(value, str):
             return value
         if isinstance(value, (int, float, bool)):
@@ -190,7 +217,7 @@ class DistributedModelName:
         except Exception:
             return str(value)
 
-    def _update_workflow(self, extra_pnginfo, unique_id, values):
+    def _update_workflow(self, extra_pnginfo: Any, unique_id: Any, values: list[str]) -> None:
         if not extra_pnginfo:
             return
         info = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo
@@ -208,7 +235,12 @@ class DistributedModelName:
         if node:
             node["widgets_values"] = [values]
 
-    def log_input(self, text, unique_id=None, extra_pnginfo=None):
+    def log_input(
+        self,
+        text: Any,
+        unique_id: Any = None,
+        extra_pnginfo: Any = None,
+    ) -> dict[str, Any]:
         values = []
         if isinstance(text, list):
             for val in text:
@@ -224,7 +256,7 @@ class DistributedModelName:
         return {"ui": {"text": values}, "result": (values,)}
 
 class ByPassTypeTuple(tuple):
-    def __getitem__(self, index):
+    def __getitem__(self, index: int) -> Any:
         if index > 0:
             index = 0
         item = super().__getitem__(index)
@@ -234,7 +266,7 @@ class ByPassTypeTuple(tuple):
 
 class ImageBatchDivider:
     @classmethod
-    def INPUT_TYPES(s):
+    def INPUT_TYPES(cls: type["ImageBatchDivider"]) -> dict[str, Any]:
         return {
             "required": {
                 "images": ("IMAGE",),
@@ -255,7 +287,7 @@ class ImageBatchDivider:
     OUTPUT_NODE = True
     CATEGORY = "image"
     
-    def divide_batch(self, images, divide_by):
+    def divide_batch(self, images: torch.Tensor, divide_by: int) -> tuple[torch.Tensor, ...]:
         total_splits = max(1, min(int(divide_by), 10))
         total_frames = images.shape[0]
         empty_tensor = images[:0]
@@ -272,7 +304,7 @@ class AudioBatchDivider:
     """Divides an audio waveform into multiple parts along the time/samples dimension."""
 
     @classmethod
-    def INPUT_TYPES(s):
+    def INPUT_TYPES(cls: type["AudioBatchDivider"]) -> dict[str, Any]:
         return {
             "required": {
                 "audio": ("AUDIO",),
@@ -293,7 +325,7 @@ class AudioBatchDivider:
     OUTPUT_NODE = True
     CATEGORY = "audio"
 
-    def divide_audio(self, audio, divide_by):
+    def divide_audio(self, audio: dict[str, Any], divide_by: int) -> tuple[dict[str, Any], ...]:
         import torch
 
         waveform = audio.get("waveform")
@@ -333,7 +365,7 @@ class DistributedEmptyImage:
     """Produces an empty IMAGE batch used when the master delegates all work."""
 
     @classmethod
-    def INPUT_TYPES(cls):
+    def INPUT_TYPES(cls: type["DistributedEmptyImage"]) -> dict[str, Any]:
         return {
             "required": {
                 "height": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}),
@@ -346,7 +378,7 @@ class DistributedEmptyImage:
     FUNCTION = "create"
     CATEGORY = "image"
 
-    def create(self, height, width, channels):
+    def create(self, height: int, width: int, channels: int) -> tuple[torch.Tensor]:
         import torch
 
         shape = (0, height, width, channels)
diff --git a/tests/api/test_distributed_queue.py b/tests/api/test_job_routes.py
similarity index 100%
rename from tests/api/test_distributed_queue.py
rename to tests/api/test_job_routes.py
diff --git a/tests/api/test_queue_orchestration.py b/tests/api/test_queue_orchestration.py
new file mode 100644
index 0000000..d6977ef
--- /dev/null
+++ b/tests/api/test_queue_orchestration.py
@@ -0,0 +1,177 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = []
+    sys.modules[f"{package_name}.api"] = api_pkg
+
+    orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration")
+    orchestration_pkg.__path__ = []
+    sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    prompt_server_instance = types.SimpleNamespace(address="127.0.0.1", port=8188)
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
+    sys.modules["server"] = server_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+
+    async def _queue_prompt_payload(prompt_obj, workflow_meta, client_id):
+        _ = (prompt_obj, workflow_meta, client_id)
+        return "prompt-id"
+
+    async_helpers_module.queue_prompt_payload = _queue_prompt_payload
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"settings": {}, "workers": []}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 2
+    constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 120.0
+    constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8
+    constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.build_master_url = lambda **_kwargs: "http://127.0.0.1:8188"
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    trace_logger_module = types.ModuleType(f"{package_name}.utils.trace_logger")
+    trace_logger_module.trace_debug = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.trace_logger"] = trace_logger_module
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+    schemas_module.parse_positive_int = lambda value, default: int(value) if str(value).isdigit() and int(value) > 0 else default
+    schemas_module.parse_positive_float = (
+        lambda value, default: float(value) if isinstance(value, (int, float, str)) and float(value) > 0 else default
+    )
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    dispatch_module = types.ModuleType(f"{package_name}.api.orchestration.dispatch")
+    dispatch_module.dispatch_worker_prompt = lambda *args, **kwargs: None
+    dispatch_module.rank_workers_by_load = lambda workers, **kwargs: workers
+    dispatch_module.select_active_workers = lambda workers, *_args, **_kwargs: (workers, False)
+    dispatch_module.select_least_busy_worker = lambda workers, **_kwargs: workers[0] if workers else None
+    sys.modules[f"{package_name}.api.orchestration.dispatch"] = dispatch_module
+
+    media_sync_module = types.ModuleType(f"{package_name}.api.orchestration.media_sync")
+    media_sync_module.convert_paths_for_platform = lambda prompt, _sep: prompt
+    media_sync_module.fetch_worker_path_separator = lambda *_args, **_kwargs: "/"
+    media_sync_module.sync_worker_media = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.api.orchestration.media_sync"] = media_sync_module
+
+    prompt_transform_module = types.ModuleType(f"{package_name}.api.orchestration.prompt_transform")
+    prompt_transform_module.PromptIndex = object
+    prompt_transform_module.apply_participant_overrides = lambda prompt, *_args, **_kwargs: prompt
+    prompt_transform_module.find_nodes_by_class = lambda _prompt, _class_name: []
+    prompt_transform_module.generate_job_id_map = lambda _index, _prefix: {}
+    prompt_transform_module.prepare_delegate_master_prompt = lambda prompt, _ids: prompt
+    prompt_transform_module.prune_prompt_for_worker = lambda prompt: prompt
+    sys.modules[f"{package_name}.api.orchestration.prompt_transform"] = prompt_transform_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[2] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_queue_orchestration_module():
+    package_name = "dist_queue_orch_testpkg"
+    _bootstrap_package(package_name)
+    return _load_module(package_name, "api/queue_orchestration.py", "api.queue_orchestration")
+
+
+queue_orchestration = _load_queue_orchestration_module()
+
+
+class QueueOrchestrationHelpersTests(unittest.TestCase):
+    def test_resolve_enabled_workers_filters_disabled_and_requested(self):
+        config = {
+            "workers": [
+                {"id": "1", "enabled": True, "host": "a", "port": 8188},
+                {"id": "2", "enabled": False, "host": "b", "port": 8189},
+                {"id": "3", "enabled": True, "host": "c", "port": "bad-port"},
+            ]
+        }
+
+        enabled_only = queue_orchestration._resolve_enabled_workers(config)
+        requested = queue_orchestration._resolve_enabled_workers(config, requested_ids={"2", "3"})
+
+        self.assertEqual([worker["id"] for worker in enabled_only], ["1", "3"])
+        self.assertEqual([worker["id"] for worker in requested], ["2", "3"])
+        self.assertEqual(requested[1]["port"], 8188)
+
+    def test_resolve_orchestration_limits_uses_defaults_for_invalid_values(self):
+        config = {
+            "settings": {
+                "worker_probe_concurrency": "x",
+                "worker_prep_concurrency": "3",
+                "media_sync_concurrency": "-1",
+                "media_sync_timeout_seconds": "42.5",
+            }
+        }
+
+        limits = queue_orchestration._resolve_orchestration_limits(config)
+
+        self.assertEqual(limits[0], 8)   # default
+        self.assertEqual(limits[1], 3)   # parsed
+        self.assertEqual(limits[2], 2)   # default
+        self.assertEqual(limits[3], 42.5)
+
+    def test_is_load_balance_enabled_accepts_common_forms(self):
+        self.assertTrue(queue_orchestration._is_load_balance_enabled(True))
+        self.assertTrue(queue_orchestration._is_load_balance_enabled(1))
+        self.assertTrue(queue_orchestration._is_load_balance_enabled(" yes "))
+        self.assertFalse(queue_orchestration._is_load_balance_enabled(0))
+        self.assertFalse(queue_orchestration._is_load_balance_enabled("off"))
+
+    def test_prompt_requests_load_balance_detects_collector_flag(self):
+        prompt_index = types.SimpleNamespace(
+            nodes_for_class=lambda class_name: ["10"] if class_name == "DistributedCollector" else [],
+            inputs_by_node={"10": {"load_balance": "true"}},
+        )
+        self.assertTrue(queue_orchestration._prompt_requests_load_balance(prompt_index))
+
+    def test_ensure_distributed_state_initializes_attrs(self):
+        prompt_server = types.SimpleNamespace()
+        queue_orchestration.ensure_distributed_state(prompt_server)
+        self.assertTrue(hasattr(prompt_server, "distributed_pending_jobs"))
+        self.assertTrue(hasattr(prompt_server, "distributed_jobs_lock"))
+        self.assertIsInstance(prompt_server.distributed_jobs_lock, asyncio.Lock)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/api/test_tunnel_routes.py b/tests/api/test_tunnel_routes.py
new file mode 100644
index 0000000..1767909
--- /dev/null
+++ b/tests/api/test_tunnel_routes.py
@@ -0,0 +1,113 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+def _load_tunnel_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "tunnel_routes.py"
+    package_name = "dist_tunnel_routes_testpkg"
+
+    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
+
+    class _Routes:
+        def get(self, _path):
+            return lambda fn: fn
+
+        def post(self, _path):
+            return lambda fn: fn
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
+    sys.modules["server"] = server_module
+
+    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
+
+    cloudflare_module = types.ModuleType(f"{package_name}.utils.cloudflare")
+
+    class _TunnelManager:
+        def get_status(self):
+            return {"status": "running"}
+
+        async def start_tunnel(self):
+            return {"status": "running", "public_url": "https://example.trycloudflare.com"}
+
+        async def stop_tunnel(self):
+            return {"status": "stopped"}
+
+    cloudflare_module.cloudflare_tunnel_manager = _TunnelManager()
+    sys.modules[f"{package_name}.utils.cloudflare"] = cloudflare_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"master": {"host": "https://master.example.com"}}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.tunnel_routes", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+tunnel_routes = _load_tunnel_routes_module()
+
+
+class TunnelRoutesTests(unittest.IsolatedAsyncioTestCase):
+    async def test_status_endpoint_returns_tunnel_and_master_host(self):
+        response = await tunnel_routes.tunnel_status_endpoint(request=None)
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload["status"], "success")
+        self.assertEqual(response.payload["tunnel"]["status"], "running")
+        self.assertEqual(response.payload["master_host"], "https://master.example.com")
+
+    async def test_start_and_stop_endpoints_return_success_payload(self):
+        start_response = await tunnel_routes.tunnel_start_endpoint(request=None)
+        stop_response = await tunnel_routes.tunnel_stop_endpoint(request=None)
+
+        self.assertEqual(start_response.status, 200)
+        self.assertEqual(start_response.payload["tunnel"]["status"], "running")
+        self.assertEqual(stop_response.status, 200)
+        self.assertEqual(stop_response.payload["tunnel"]["status"], "stopped")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/api/test_usdu_routes.py b/tests/api/test_usdu_routes.py
index 1d2c3aa..e74a2a6 100644
--- a/tests/api/test_usdu_routes.py
+++ b/tests/api/test_usdu_routes.py
@@ -165,6 +165,7 @@ def _load_usdu_routes_module():
     sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
 
     parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers")
+    parsers_module.parse_tiles_from_form = lambda _data: []
     parsers_module._parse_tiles_from_form = lambda _data: []
     sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
 
diff --git a/tests/api/test_worker_routes.py b/tests/api/test_worker_routes.py
index 9a4c562..c016d1d 100644
--- a/tests/api/test_worker_routes.py
+++ b/tests/api/test_worker_routes.py
@@ -62,7 +62,7 @@ class _DummyWorkerManager:
         worker_id = str(worker["id"])
         self.processes[worker_id] = {
             "pid": 12345,
-            "log_file": f"/tmp/distributed_worker_{worker_id}.log",
+            "log_file": f"/tmp/distributed_worker_{worker_id}.log",  # nosec B108 - deterministic fake test path
             "process": None,
         }
         return 12345
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py
new file mode 100644
index 0000000..e0310a0
--- /dev/null
+++ b/tests/unit/__init__.py
@@ -0,0 +1 @@
+"""Unit tests."""
diff --git a/tests/unit/nodes/__init__.py b/tests/unit/nodes/__init__.py
new file mode 100644
index 0000000..e0310a0
--- /dev/null
+++ b/tests/unit/nodes/__init__.py
@@ -0,0 +1 @@
+"""Unit tests."""
diff --git a/tests/test_batch_dividers.py b/tests/unit/nodes/test_batch_dividers.py
similarity index 98%
rename from tests/test_batch_dividers.py
rename to tests/unit/nodes/test_batch_dividers.py
index 8db4a1d..01f9033 100644
--- a/tests/test_batch_dividers.py
+++ b/tests/unit/nodes/test_batch_dividers.py
@@ -8,7 +8,7 @@ import torch
 
 
 def _load_utilities_module():
-    module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
+    module_path = Path(__file__).resolve().parents[3] / "nodes" / "utilities.py"
     package_name = "dist_divider_testpkg"
 
     for mod_name in list(sys.modules):
diff --git a/tests/test_branch.py b/tests/unit/nodes/test_branch.py
similarity index 97%
rename from tests/test_branch.py
rename to tests/unit/nodes/test_branch.py
index d1f97e6..5e1e5f4 100644
--- a/tests/test_branch.py
+++ b/tests/unit/nodes/test_branch.py
@@ -29,7 +29,7 @@ def _bootstrap_package(package_name):
 
 
 def _load_module(package_name, module_rel_path, module_name):
-    module_path = Path(__file__).resolve().parents[1] / module_rel_path
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
     spec = importlib.util.spec_from_file_location(
         f"{package_name}.{module_name}",
         module_path,
diff --git a/tests/test_join.py b/tests/unit/nodes/test_branch_collector.py
similarity index 79%
rename from tests/test_join.py
rename to tests/unit/nodes/test_branch_collector.py
index f537b10..651ead8 100644
--- a/tests/test_join.py
+++ b/tests/unit/nodes/test_branch_collector.py
@@ -1,5 +1,7 @@
 import asyncio
+import base64
 import importlib.util
+import io
 import json
 import sys
 import types
@@ -51,8 +53,16 @@ def _bootstrap_package(package_name):
         array = (tensor.numpy() * 255.0).astype("uint8")
         return Image.fromarray(array)
 
+    def encode_tensor_png_data_url(image_batch, batch_index=0):
+        image = tensor_to_pil(image_batch, batch_index)
+        byte_io = io.BytesIO()
+        image.save(byte_io, format="PNG", compress_level=0)
+        encoded = base64.b64encode(byte_io.getvalue()).decode("utf-8")
+        return f"data:image/png;base64,{encoded}"
+
     image_module.ensure_contiguous = ensure_contiguous
     image_module.tensor_to_pil = tensor_to_pil
+    image_module.encode_tensor_png_data_url = encode_tensor_png_data_url
     sys.modules[f"{package_name}.utils.image"] = image_module
 
     network_module = types.ModuleType(f"{package_name}.utils.network")
@@ -69,7 +79,7 @@ def _bootstrap_package(package_name):
 
 
 def _load_module(package_name, module_rel_path, module_name):
-    module_path = Path(__file__).resolve().parents[1] / module_rel_path
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
     spec = importlib.util.spec_from_file_location(
         f"{package_name}.{module_name}",
         module_path,
@@ -85,12 +95,30 @@ def _load_branch_collector_module():
     package_name = "dist_join_testpkg"
     _bootstrap_package(package_name)
     _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
+    _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
+    _load_module(package_name, "nodes/runtime_helpers.py", "nodes.runtime_helpers")
+    _load_module(package_name, "nodes/queue_wait.py", "nodes.queue_wait")
+    _load_module(package_name, "utils/worker_ids.py", "utils.worker_ids")
     return _load_module(package_name, "nodes/branch_collector.py", "nodes.branch_collector")
 
 
 branch_collector_module = _load_branch_collector_module()
 
 
+def _context(**overrides):
+    defaults = {
+        "multi_job_id": "",
+        "is_worker": False,
+        "master_url": "",
+        "enabled_worker_ids": "[]",
+        "worker_id": "",
+        "assigned_branch": -1,
+        "delegate_only": False,
+    }
+    defaults.update(overrides)
+    return branch_collector_module.BranchRunContext(**defaults)
+
+
 class _FakeResponseCtx:
     def __init__(self, recorder, payload, url):
         self._recorder = recorder
@@ -138,11 +166,13 @@ class DistributedBranchCollectorTests(unittest.IsolatedAsyncioTestCase):
         outputs = await self.node.execute(
             [None, tensor, None, None, None, None, None, None, None, None],
             num_branches=2,
-            multi_job_id="join-job",
-            is_worker=True,
-            master_url="http://master.local:8188",
-            worker_id="worker-a",
-            assigned_branch=1,
+            context=_context(
+                multi_job_id="join-job",
+                is_worker=True,
+                master_url="http://master.local:8188",
+                worker_id="worker-a",
+                assigned_branch=1,
+            ),
         )
 
         self.assertEqual(len(fake_session.calls), 1)
@@ -167,10 +197,11 @@ class DistributedBranchCollectorTests(unittest.IsolatedAsyncioTestCase):
         outputs = await self.node.execute(
             [master_tensor, None, None, None, None, None, None, None, None, None],
             num_branches=3,
-            multi_job_id="join-job",
-            is_worker=False,
-            enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
-            assigned_branch=0,
+            context=_context(
+                multi_job_id="join-job",
+                enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
+                assigned_branch=0,
+            ),
         )
 
         self.assertTrue(torch.equal(outputs[0], master_tensor))
@@ -190,10 +221,11 @@ class DistributedBranchCollectorTests(unittest.IsolatedAsyncioTestCase):
         outputs = await self.node.execute(
             [master_tensor, None, None, None, None, None, None, None, None, None],
             num_branches=2,
-            multi_job_id="join-job",
-            is_worker=False,
-            enabled_worker_ids=json.dumps(["worker-a"]),
-            assigned_branch=0,
+            context=_context(
+                multi_job_id="join-job",
+                enabled_worker_ids=json.dumps(["worker-a"]),
+                assigned_branch=0,
+            ),
         )
 
         self.assertIsNone(outputs[2])
@@ -211,10 +243,11 @@ class DistributedBranchCollectorTests(unittest.IsolatedAsyncioTestCase):
         outputs = await self.node.execute(
             [master_tensor, None, None, None, None, None, None, None, None, None],
             num_branches=3,
-            multi_job_id="join-job",
-            is_worker=False,
-            enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
-            assigned_branch=0,
+            context=_context(
+                multi_job_id="join-job",
+                enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
+                assigned_branch=0,
+            ),
         )
 
         self.assertTrue(torch.equal(outputs[0], master_tensor))
@@ -235,10 +268,11 @@ class DistributedBranchCollectorTests(unittest.IsolatedAsyncioTestCase):
             # Local value arrives on slot 1 even though participant is assigned slot 0.
             [None, local_tensor, None, None, None, None, None, None, None, None],
             num_branches=3,
-            multi_job_id="join-job",
-            is_worker=False,
-            enabled_worker_ids=json.dumps([]),
-            assigned_branch=0,
+            context=_context(
+                multi_job_id="join-job",
+                enabled_worker_ids=json.dumps([]),
+                assigned_branch=0,
+            ),
         )
 
         self.assertTrue(torch.equal(outputs[0], local_tensor))
diff --git a/tests/unit/nodes/test_collector.py b/tests/unit/nodes/test_collector.py
new file mode 100644
index 0000000..6e053d6
--- /dev/null
+++ b/tests/unit/nodes/test_collector.py
@@ -0,0 +1,228 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
+    nodes_pkg.__path__ = []
+    sys.modules[f"{package_name}.nodes"] = nodes_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.get_worker_timeout_seconds = lambda: 0.2
+    config_module.load_config = lambda: {"workers": []}
+    config_module.is_master_delegate_only = lambda: False
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.HEARTBEAT_INTERVAL = 1.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+
+    def ensure_contiguous(tensor):
+        return tensor.contiguous()
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    def pil_to_tensor(image):
+        array = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(array).unsqueeze(0)
+
+    image_module.ensure_contiguous = ensure_contiguous
+    image_module.tensor_to_pil = tensor_to_pil
+    image_module.pil_to_tensor = pil_to_tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.build_worker_url = lambda _worker: "http://worker.local:8188"
+    network_module.get_client_session = lambda: None
+
+    async def probe_worker(_url, timeout=2.0):
+        _ = timeout
+        return None
+
+    network_module.probe_worker = probe_worker
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
+    audio_payload_module.encode_audio_payload = lambda payload: payload
+    sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    if "aiohttp" not in sys.modules:
+        try:
+            import aiohttp as _aiohttp  # noqa: F401
+        except Exception:
+            aiohttp_module = types.ModuleType("aiohttp")
+
+            class _ClientTimeout:
+                def __init__(self, total=None):
+                    self.total = total
+
+            aiohttp_module.ClientTimeout = _ClientTimeout
+            sys.modules["aiohttp"] = aiohttp_module
+
+    class _FakePromptServer:
+        def __init__(self):
+            self.distributed_pending_jobs = {}
+            self.distributed_jobs_lock = asyncio.Lock()
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=_FakePromptServer())
+    sys.modules["server"] = server_module
+
+    comfy_module = types.ModuleType("comfy")
+    model_mgmt = types.ModuleType("comfy.model_management")
+
+    class _InterruptProcessingException(Exception):
+        pass
+
+    model_mgmt.throw_exception_if_processing_interrupted = lambda: None
+    model_mgmt.InterruptProcessingException = _InterruptProcessingException
+    comfy_module.model_management = model_mgmt
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.model_management"] = model_mgmt
+
+    comfy_utils_module = types.ModuleType("comfy.utils")
+
+    class _ProgressBar:
+        def __init__(self, _total):
+            self.total = _total
+
+        def update(self, _step):
+            return None
+
+    comfy_utils_module.ProgressBar = _ProgressBar
+    sys.modules["comfy.utils"] = comfy_utils_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_collector_module():
+    package_name = "dist_collector_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
+    _load_module(package_name, "utils/worker_ids.py", "utils.worker_ids")
+    return _load_module(package_name, "nodes/collector.py", "nodes.collector")
+
+
+collector_module = _load_collector_module()
+
+
+class DistributedCollectorTests(unittest.TestCase):
+    def setUp(self):
+        self.node = collector_module.DistributedCollectorNode()
+
+    def test_store_worker_result_tracks_by_worker_and_index(self):
+        worker_images = {}
+        stored = self.node._store_worker_result(
+            worker_images,
+            {
+                "worker_id": "worker-a",
+                "image_index": 2,
+                "tensor": torch.full((1, 1, 1, 1), 0.25),
+            },
+        )
+
+        self.assertEqual(stored, 1)
+        self.assertIn("worker-a", worker_images)
+        self.assertIn(2, worker_images["worker-a"])
+
+    def test_combine_audio_honors_master_then_worker_order(self):
+        master_audio = {"waveform": torch.ones((1, 2, 1)), "sample_rate": 48000}
+        worker_audio = {
+            "worker-b": {"waveform": torch.full((1, 2, 1), 3.0), "sample_rate": 48000},
+            "worker-a": {"waveform": torch.full((1, 2, 2), 2.0), "sample_rate": 48000},
+        }
+        empty_audio = {"waveform": torch.zeros((1, 2, 1)), "sample_rate": 44100}
+
+        combined = self.node._combine_audio(
+            master_audio=master_audio,
+            worker_audio=worker_audio,
+            empty_audio=empty_audio,
+            worker_order=["worker-a", "worker-b"],
+        )
+
+        expected = torch.cat(
+            [
+                master_audio["waveform"],
+                worker_audio["worker-a"]["waveform"],
+                worker_audio["worker-b"]["waveform"],
+            ],
+            dim=-1,
+        )
+        self.assertEqual(combined["sample_rate"], 48000)
+        self.assertTrue(torch.equal(combined["waveform"], expected))
+
+    def test_reorder_and_combine_tensors_uses_enabled_worker_priority(self):
+        worker_images = {
+            "worker-b": {
+                1: torch.full((1, 1, 1, 1), 4.0),
+                0: torch.full((1, 1, 1, 1), 3.0),
+            },
+            "worker-a": {
+                0: torch.full((1, 1, 1, 1), 2.0),
+            },
+        }
+        master_images = torch.full((1, 1, 1, 1), 1.0)
+
+        combined = self.node._reorder_and_combine_tensors(
+            worker_images=worker_images,
+            worker_order=["worker-a", "worker-b"],
+            master_batch_size=1,
+            images_on_cpu=master_images,
+            delegate_mode=False,
+            fallback_images=None,
+        )
+
+        self.assertEqual(combined.shape, (4, 1, 1, 1))
+        self.assertEqual(combined[0, 0, 0, 0].item(), 1.0)
+        self.assertEqual(combined[1, 0, 0, 0].item(), 2.0)
+        self.assertEqual(combined[2, 0, 0, 0].item(), 3.0)
+        self.assertEqual(combined[3, 0, 0, 0].item(), 4.0)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/nodes/test_distributed_upscale.py b/tests/unit/nodes/test_distributed_upscale.py
new file mode 100644
index 0000000..5c055b5
--- /dev/null
+++ b/tests/unit/nodes/test_distributed_upscale.py
@@ -0,0 +1,239 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
+    nodes_pkg.__path__ = []
+    sys.modules[f"{package_name}.nodes"] = nodes_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
+    modes_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.run_async_in_server_loop = lambda coro, timeout=None: (_ for _ in ()).throw(
+        RuntimeError(f"run_async_in_server_loop should not be used in these tests (timeout={timeout})")
+    )
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: None
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    tile_ops_module = types.ModuleType(f"{package_name}.upscale.tile_ops")
+
+    class _TileOpsMixin:
+        def round_to_multiple(self, value):
+            return int(value)
+
+        def calculate_tiles(self, width, height, tile_width, tile_height, force_uniform_tiles):
+            _ = force_uniform_tiles
+            tiles = []
+            for y in range(0, int(height), max(1, int(tile_height))):
+                for x in range(0, int(width), max(1, int(tile_width))):
+                    tiles.append((x, y))
+            return tiles or [(0, 0)]
+
+    tile_ops_module.TileOpsMixin = _TileOpsMixin
+    sys.modules[f"{package_name}.upscale.tile_ops"] = tile_ops_module
+
+    result_collector_module = types.ModuleType(f"{package_name}.upscale.result_collector")
+    result_collector_module.ResultCollectorMixin = type("ResultCollectorMixin", (), {})
+    sys.modules[f"{package_name}.upscale.result_collector"] = result_collector_module
+
+    worker_comms_module = types.ModuleType(f"{package_name}.upscale.worker_comms")
+    worker_comms_module.WorkerCommsMixin = type("WorkerCommsMixin", (), {})
+    sys.modules[f"{package_name}.upscale.worker_comms"] = worker_comms_module
+
+    job_state_module = types.ModuleType(f"{package_name}.upscale.job_state")
+    job_state_module.JobStateMixin = type("JobStateMixin", (), {})
+    sys.modules[f"{package_name}.upscale.job_state"] = job_state_module
+
+    single_gpu_module = types.ModuleType(f"{package_name}.upscale.modes.single_gpu")
+
+    class _SingleGpuModeMixin:
+        def process_single_gpu(self, *_args, **_kwargs):
+            return ("single_gpu_result",)
+
+    single_gpu_module.SingleGpuModeMixin = _SingleGpuModeMixin
+    sys.modules[f"{package_name}.upscale.modes.single_gpu"] = single_gpu_module
+
+    static_mode_module = types.ModuleType(f"{package_name}.upscale.modes.static")
+
+    class _StaticModeMixin:
+        def _process_worker_static_sync(self, *_args, **_kwargs):
+            return ("worker_static_result",)
+
+        def _process_master_static_sync(self, *_args, **_kwargs):
+            return ("master_static_result",)
+
+    static_mode_module.StaticModeMixin = _StaticModeMixin
+    sys.modules[f"{package_name}.upscale.modes.static"] = static_mode_module
+
+    dynamic_mode_module = types.ModuleType(f"{package_name}.upscale.modes.dynamic")
+
+    class _DynamicModeMixin:
+        def process_worker_dynamic(self, *_args, **_kwargs):
+            return ("worker_dynamic_result",)
+
+        def process_master_dynamic(self, *_args, **_kwargs):
+            return ("master_dynamic_result",)
+
+    dynamic_mode_module.DynamicModeMixin = _DynamicModeMixin
+    sys.modules[f"{package_name}.upscale.modes.dynamic"] = dynamic_mode_module
+
+    comfy_module = types.ModuleType("comfy")
+    samplers_module = types.ModuleType("comfy.samplers")
+
+    class _KSampler:
+        SAMPLERS = ("euler",)
+        SCHEDULERS = ("normal",)
+
+    samplers_module.KSampler = _KSampler
+    comfy_module.samplers = samplers_module
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.samplers"] = samplers_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_distributed_upscale_module():
+    package_name = "dist_distributed_upscale_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
+    _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args")
+    return _load_module(package_name, "nodes/distributed_upscale.py", "nodes.distributed_upscale")
+
+
+upscale_module = _load_distributed_upscale_module()
+
+
+class DistributedUpscaleTests(unittest.TestCase):
+    def setUp(self):
+        self.node = upscale_module.UltimateSDUpscaleDistributed()
+        self.base_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
+        self.common_kwargs = {
+            "model": object(),
+            "positive": object(),
+            "negative": object(),
+            "vae": object(),
+            "seed": 1,
+            "steps": 10,
+            "cfg": 7.5,
+            "sampler_name": "euler",
+            "scheduler": "normal",
+            "denoise": 0.5,
+            "tile_width": 32,
+            "tile_height": 32,
+            "padding": 16,
+            "mask_blur": 4,
+            "force_uniform_tiles": True,
+            "tiled_decode": False,
+        }
+
+    def test_parse_enabled_worker_ids_supports_json_and_lists(self):
+        self.assertEqual(
+            upscale_module._parse_enabled_worker_ids(["worker-a", 2]),
+            ["worker-a", "2"],
+        )
+        self.assertEqual(
+            upscale_module._parse_enabled_worker_ids('["worker-a","worker-b"]'),
+            ["worker-a", "worker-b"],
+        )
+        self.assertEqual(upscale_module._parse_enabled_worker_ids("invalid-json"), [])
+        self.assertEqual(upscale_module._parse_enabled_worker_ids(None), [])
+
+    def test_determine_processing_mode(self):
+        self.assertEqual(self.node._determine_processing_mode(batch_size=1, num_workers=0, dynamic_threshold=8), "single_gpu")
+        self.assertEqual(self.node._determine_processing_mode(batch_size=9, num_workers=2, dynamic_threshold=8), "static")
+
+    def test_run_raises_for_non_4n_plus_1_master_batches(self):
+        bad_batch = torch.zeros((2, 64, 64, 3), dtype=torch.float32)
+        with self.assertRaises(ValueError):
+            self.node.run(
+                upscaled_image=bad_batch,
+                multi_job_id="job-1",
+                is_worker=False,
+                **self.common_kwargs,
+            )
+
+    def test_run_dispatches_single_gpu_when_not_distributed(self):
+        self.node.process_single_gpu = lambda *_args, **_kwargs: ("single",)
+        result = self.node.run(
+            upscaled_image=self.base_image,
+            multi_job_id="",
+            is_worker=False,
+            **self.common_kwargs,
+        )
+        self.assertEqual(result, ("single",))
+
+    def test_run_dispatches_worker_or_master_paths(self):
+        self.node.process_worker = lambda *_args, **_kwargs: ("worker",)
+        self.node.process_master = lambda *_args, **_kwargs: ("master",)
+
+        worker_result = self.node.run(
+            upscaled_image=self.base_image,
+            multi_job_id="job-2",
+            is_worker=True,
+            master_url="http://master.local:8188",
+            enabled_worker_ids='["worker-a"]',
+            worker_id="worker-a",
+            tile_indices="",
+            dynamic_threshold=8,
+            **self.common_kwargs,
+        )
+        master_result = self.node.run(
+            upscaled_image=self.base_image,
+            multi_job_id="job-2",
+            is_worker=False,
+            enabled_worker_ids='["worker-a"]',
+            worker_id="",
+            tile_indices="",
+            dynamic_threshold=8,
+            **self.common_kwargs,
+        )
+
+        self.assertEqual(worker_result, ("worker",))
+        self.assertEqual(master_result, ("master",))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_distributed_value.py b/tests/unit/nodes/test_distributed_value.py
similarity index 99%
rename from tests/test_distributed_value.py
rename to tests/unit/nodes/test_distributed_value.py
index d2551c6..53809ff 100644
--- a/tests/test_distributed_value.py
+++ b/tests/unit/nodes/test_distributed_value.py
@@ -13,7 +13,7 @@ class DistributedValueTests(unittest.TestCase):
         from pathlib import Path
         from unittest.mock import MagicMock
 
-        module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
+        module_path = Path(__file__).resolve().parents[3] / "nodes" / "utilities.py"
         pkg_name = "dv_test_pkg"
 
         for mod_name in list(sys.modules):
diff --git a/tests/test_list_collector.py b/tests/unit/nodes/test_list_collector.py
similarity index 80%
rename from tests/test_list_collector.py
rename to tests/unit/nodes/test_list_collector.py
index 6772583..19843ef 100644
--- a/tests/test_list_collector.py
+++ b/tests/unit/nodes/test_list_collector.py
@@ -1,5 +1,7 @@
 import asyncio
+import base64
 import importlib.util
+import io
 import json
 import sys
 import types
@@ -51,8 +53,16 @@ def _bootstrap_package(package_name):
         array = (tensor.numpy() * 255.0).astype("uint8")
         return Image.fromarray(array)
 
+    def encode_tensor_png_data_url(image_batch, batch_index=0):
+        image = tensor_to_pil(image_batch, batch_index)
+        byte_io = io.BytesIO()
+        image.save(byte_io, format="PNG", compress_level=0)
+        encoded = base64.b64encode(byte_io.getvalue()).decode("utf-8")
+        return f"data:image/png;base64,{encoded}"
+
     image_module.ensure_contiguous = ensure_contiguous
     image_module.tensor_to_pil = tensor_to_pil
+    image_module.encode_tensor_png_data_url = encode_tensor_png_data_url
     sys.modules[f"{package_name}.utils.image"] = image_module
 
     network_module = types.ModuleType(f"{package_name}.utils.network")
@@ -69,7 +79,7 @@ def _bootstrap_package(package_name):
 
 
 def _load_module(package_name, module_rel_path, module_name):
-    module_path = Path(__file__).resolve().parents[1] / module_rel_path
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
     spec = importlib.util.spec_from_file_location(
         f"{package_name}.{module_name}",
         module_path,
@@ -84,12 +94,29 @@ def _load_module(package_name, module_rel_path, module_name):
 def _load_list_collector_module():
     package_name = "dist_list_collector_testpkg"
     _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
+    _load_module(package_name, "nodes/runtime_helpers.py", "nodes.runtime_helpers")
+    _load_module(package_name, "nodes/queue_wait.py", "nodes.queue_wait")
+    _load_module(package_name, "utils/worker_ids.py", "utils.worker_ids")
     return _load_module(package_name, "nodes/list_collector.py", "nodes.list_collector")
 
 
 collector_module = _load_list_collector_module()
 
 
+def _context(**overrides):
+    defaults = {
+        "multi_job_id": "",
+        "is_worker": False,
+        "master_url": "",
+        "enabled_worker_ids": "[]",
+        "worker_id": "",
+        "delegate_only": False,
+    }
+    defaults.update(overrides)
+    return collector_module.ListCollectorRunContext(**defaults)
+
+
 class _FakeResponseCtx:
     def __init__(self, recorder, payload, url):
         self._recorder = recorder
@@ -140,10 +167,12 @@ class DistributedListCollectorTests(unittest.IsolatedAsyncioTestCase):
 
         result = await self.node.execute(
             images,
-            multi_job_id="job-list",
-            is_worker=True,
-            master_url="http://master.local:8188",
-            worker_id="worker-a",
+            context=_context(
+                multi_job_id="job-list",
+                is_worker=True,
+                master_url="http://master.local:8188",
+                worker_id="worker-a",
+            ),
         )
 
         self.assertEqual(len(result[0]), 2)
@@ -169,9 +198,10 @@ class DistributedListCollectorTests(unittest.IsolatedAsyncioTestCase):
 
         result = await self.node.execute(
             master_images,
-            multi_job_id="job-list",
-            is_worker=False,
-            enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
+            context=_context(
+                multi_job_id="job-list",
+                enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
+            ),
         )
 
         self.assertEqual(len(result[0]), 3)
@@ -192,10 +222,11 @@ class DistributedListCollectorTests(unittest.IsolatedAsyncioTestCase):
 
         result = await self.node.execute(
             master_images,
-            multi_job_id="job-list",
-            is_worker=False,
-            enabled_worker_ids=json.dumps(["worker-a"]),
-            delegate_only=True,
+            context=_context(
+                multi_job_id="job-list",
+                enabled_worker_ids=json.dumps(["worker-a"]),
+                delegate_only=True,
+            ),
         )
 
         self.assertEqual(len(result[0]), 1)
diff --git a/tests/test_list_splitter.py b/tests/unit/nodes/test_list_splitter.py
similarity index 99%
rename from tests/test_list_splitter.py
rename to tests/unit/nodes/test_list_splitter.py
index 0fb69ce..849447f 100644
--- a/tests/test_list_splitter.py
+++ b/tests/unit/nodes/test_list_splitter.py
@@ -99,7 +99,7 @@ def _bootstrap_package(package_name):
 
 
 def _load_module(package_name, module_rel_path, module_name):
-    module_path = Path(__file__).resolve().parents[1] / module_rel_path
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
     spec = importlib.util.spec_from_file_location(
         f"{package_name}.{module_name}",
         module_path,
diff --git a/tests/unit/system/__init__.py b/tests/unit/system/__init__.py
new file mode 100644
index 0000000..e0310a0
--- /dev/null
+++ b/tests/unit/system/__init__.py
@@ -0,0 +1 @@
+"""Unit tests."""
diff --git a/tests/test_async_helpers.py b/tests/unit/system/test_async_helpers.py
similarity index 98%
rename from tests/test_async_helpers.py
rename to tests/unit/system/test_async_helpers.py
index 4499ce9..afa281f 100644
--- a/tests/test_async_helpers.py
+++ b/tests/unit/system/test_async_helpers.py
@@ -61,7 +61,7 @@ def _bootstrap_package(package_name):
 def _load_async_helpers_module():
     package_name = "dist_async_helpers_testpkg"
     fake_prompt_server = _bootstrap_package(package_name)
-    module_path = Path(__file__).resolve().parents[1] / "utils/async_helpers.py"
+    module_path = Path(__file__).resolve().parents[3] / "utils/async_helpers.py"
     spec = importlib.util.spec_from_file_location(
         f"{package_name}.utils.async_helpers",
         module_path,
diff --git a/tests/test_config.py b/tests/unit/system/test_config.py
similarity index 98%
rename from tests/test_config.py
rename to tests/unit/system/test_config.py
index f5eaa0e..fa05043 100644
--- a/tests/test_config.py
+++ b/tests/unit/system/test_config.py
@@ -10,7 +10,7 @@ from unittest.mock import patch
 
 
 def _load_config_module():
-    module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
+    module_path = Path(__file__).resolve().parents[3] / "utils" / "config.py"
     package_name = "dist_cfg_testpkg"
 
     for mod_name in list(sys.modules):
@@ -27,6 +27,7 @@ def _load_config_module():
     sys.modules[f"{package_name}.logging"] = logging_module
 
     constants_module = types.ModuleType(f"{package_name}.constants")
+    constants_module.GPU_CONFIG_FILE = "gpu_config.json"
     constants_module.HEARTBEAT_TIMEOUT = 30
     sys.modules[f"{package_name}.constants"] = constants_module
 
@@ -213,7 +214,7 @@ class SaveConfigTests(unittest.TestCase):
                 data = config.get_default_config()
                 config.save_config(data)
                 # Cache is now None; load_config should re-read
-                self.assertIsNone(config._config_cache)
+                self.assertIsNone(config._config_state().cache)
 
     def test_written_file_is_valid_json(self):
         data = config.get_default_config()
diff --git a/tests/test_detection.py b/tests/unit/system/test_detection.py
similarity index 96%
rename from tests/test_detection.py
rename to tests/unit/system/test_detection.py
index ba6cd0e..7159855 100644
--- a/tests/test_detection.py
+++ b/tests/unit/system/test_detection.py
@@ -7,7 +7,7 @@ from unittest.mock import patch
 
 
 def _load_detection_module():
-    module_path = Path(__file__).resolve().parents[1] / "workers" / "detection.py"
+    module_path = Path(__file__).resolve().parents[3] / "workers" / "detection.py"
     package_name = "dist_det_testpkg"
 
     for mod_name in list(sys.modules):
@@ -32,6 +32,7 @@ def _load_detection_module():
 
     network_module = types.ModuleType(f"{package_name}.utils.network")
     network_module.normalize_host = lambda value: value
+    network_module.build_worker_url = lambda worker, endpoint="": f"http://{worker.get('host', 'localhost')}{endpoint}"
 
     async def _fake_session():
         raise RuntimeError("network calls not used in these tests")
@@ -156,7 +157,7 @@ class IsLocalWorkerTests(unittest.IsolatedAsyncioTestCase):
         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})
+        result = await detection.is_local_worker({"host": "0.0.0.0", "port": 8188})  # nosec B104 - explicit wildcard host test case
         self.assertTrue(result)
 
     async def test_true_when_type_is_local(self):
diff --git a/tests/test_dispatch_selection.py b/tests/unit/system/test_dispatch.py
similarity index 99%
rename from tests/test_dispatch_selection.py
rename to tests/unit/system/test_dispatch.py
index 73c432e..9e04b29 100644
--- a/tests/test_dispatch_selection.py
+++ b/tests/unit/system/test_dispatch.py
@@ -8,7 +8,7 @@ from unittest.mock import patch
 
 
 def _load_dispatch_module():
-    module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
+    module_path = Path(__file__).resolve().parents[3] / "api" / "orchestration" / "dispatch.py"
 
     package_name = "dist_dispatch_testpkg"
     root_pkg = types.ModuleType(package_name)
diff --git a/tests/unit/system/test_distributed.py b/tests/unit/system/test_distributed.py
new file mode 100644
index 0000000..49f2ce5
--- /dev/null
+++ b/tests/unit/system/test_distributed.py
@@ -0,0 +1,152 @@
+import importlib.util
+import os
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _load_distributed_module(is_worker):
+    module_path = Path(__file__).resolve().parents[3] / "distributed.py"
+    package_name = "dist_distributed_entry_testpkg"
+    calls = {
+        "ensure_config_exists": 0,
+        "ensure_distributed_state": 0,
+        "ensure_tile_jobs_initialized": 0,
+        "delayed_auto_launch": 0,
+        "register_async_signals": 0,
+        "sync_cleanup_registered": 0,
+    }
+
+    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
+
+    # Absolute dependency used by module.
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace())
+    sys.modules["server"] = server_module
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.ensure_config_exists = lambda: calls.__setitem__("ensure_config_exists", calls["ensure_config_exists"] + 1)
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _cleanup_client_session():
+        return None
+
+    network_module.cleanup_client_session = _cleanup_client_session
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    workers_pkg = types.ModuleType(f"{package_name}.workers")
+    workers_pkg.__path__ = []
+    workers_pkg.get_worker_manager = lambda: None
+    sys.modules[f"{package_name}.workers"] = workers_pkg
+
+    startup_module = types.ModuleType(f"{package_name}.workers.startup")
+    startup_module.delayed_auto_launch = lambda: calls.__setitem__("delayed_auto_launch", calls["delayed_auto_launch"] + 1)
+    startup_module.register_async_signals = lambda: calls.__setitem__("register_async_signals", calls["register_async_signals"] + 1)
+    startup_module.sync_cleanup = lambda: None
+    sys.modules[f"{package_name}.workers.startup"] = startup_module
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: calls.__setitem__(
+        "ensure_tile_jobs_initialized", calls["ensure_tile_jobs_initialized"] + 1
+    )
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    nodes_module = types.ModuleType(f"{package_name}.nodes")
+    nodes_module.NODE_CLASS_MAPPINGS = {}
+    nodes_module.NODE_DISPLAY_NAME_MAPPINGS = {}
+    nodes_module.ImageBatchDivider = object
+    nodes_module.DistributedCollectorNode = object
+    nodes_module.DistributedSeed = object
+    nodes_module.DistributedModelName = object
+    nodes_module.DistributedValue = object
+    nodes_module.AudioBatchDivider = object
+    nodes_module.DistributedEmptyImage = object
+    nodes_module.DistributedListSplitter = object
+    nodes_module.DistributedListCollector = object
+    nodes_module.DistributedBranch = object
+    nodes_module.DistributedBranchCollector = object
+    nodes_module.AnyType = object
+    nodes_module.ByPassTypeTuple = tuple
+    nodes_module.any_type = "*"
+    sys.modules[f"{package_name}.nodes"] = nodes_module
+
+    api_module = types.ModuleType(f"{package_name}.api")
+    sys.modules[f"{package_name}.api"] = api_module
+
+    queue_orch_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
+    queue_orch_module.ensure_distributed_state = lambda _ps: calls.__setitem__(
+        "ensure_distributed_state", calls["ensure_distributed_state"] + 1
+    )
+    sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orch_module
+
+    import atexit
+
+    original_register = atexit.register
+    atexit.register = lambda fn: calls.__setitem__("sync_cleanup_registered", calls["sync_cleanup_registered"] + 1) or fn
+
+    original_env = os.environ.get("COMFYUI_IS_WORKER")
+    try:
+        if is_worker:
+            os.environ["COMFYUI_IS_WORKER"] = "1"
+        else:
+            os.environ.pop("COMFYUI_IS_WORKER", None)
+
+        spec = importlib.util.spec_from_file_location(f"{package_name}.distributed", module_path)
+        module = importlib.util.module_from_spec(spec)
+        assert spec is not None and spec.loader is not None
+        sys.modules[spec.name] = module
+        spec.loader.exec_module(module)
+    finally:
+        atexit.register = original_register
+        if original_env is None:
+            os.environ.pop("COMFYUI_IS_WORKER", None)
+        else:
+            os.environ["COMFYUI_IS_WORKER"] = original_env
+
+    return calls
+
+
+class DistributedEntryTests(unittest.TestCase):
+    def test_import_initializes_core_state_for_worker(self):
+        calls = _load_distributed_module(is_worker=True)
+
+        self.assertEqual(calls["ensure_config_exists"], 1)
+        self.assertEqual(calls["ensure_distributed_state"], 1)
+        self.assertEqual(calls["ensure_tile_jobs_initialized"], 1)
+        self.assertEqual(calls["delayed_auto_launch"], 0)
+        self.assertEqual(calls["register_async_signals"], 0)
+
+    def test_import_triggers_master_startup_hooks(self):
+        calls = _load_distributed_module(is_worker=False)
+
+        self.assertEqual(calls["ensure_config_exists"], 1)
+        self.assertEqual(calls["ensure_distributed_state"], 1)
+        self.assertEqual(calls["ensure_tile_jobs_initialized"], 1)
+        self.assertEqual(calls["delayed_auto_launch"], 1)
+        self.assertEqual(calls["register_async_signals"], 1)
+        self.assertGreaterEqual(calls["sync_cleanup_registered"], 2)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/system/test_exceptions.py b/tests/unit/system/test_exceptions.py
new file mode 100644
index 0000000..9658e67
--- /dev/null
+++ b/tests/unit/system/test_exceptions.py
@@ -0,0 +1,40 @@
+import importlib.util
+import unittest
+from pathlib import Path
+
+
+def _load_exceptions_module():
+    module_path = Path(__file__).resolve().parents[3] / "utils" / "exceptions.py"
+    spec = importlib.util.spec_from_file_location("dist_test_exceptions", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+exceptions = _load_exceptions_module()
+
+
+class ExceptionsModuleTests(unittest.TestCase):
+    def test_worker_error_tracks_worker_id(self):
+        err = exceptions.WorkerError("worker failed", worker_id="w-1")
+        self.assertIsInstance(err, exceptions.DistributedError)
+        self.assertEqual(str(err), "worker failed")
+        self.assertEqual(err.worker_id, "w-1")
+
+    def test_process_error_tracks_pid_and_worker_id(self):
+        err = exceptions.ProcessError("process failed", pid=1001, worker_id="worker-x")
+        self.assertIsInstance(err, exceptions.DistributedError)
+        self.assertEqual(err.pid, 1001)
+        self.assertEqual(err.worker_id, "worker-x")
+
+    def test_specialized_errors_inherit_expected_base_types(self):
+        self.assertTrue(issubclass(exceptions.WorkerTimeoutError, exceptions.WorkerError))
+        self.assertTrue(issubclass(exceptions.WorkerNotAvailableError, exceptions.WorkerError))
+        self.assertTrue(issubclass(exceptions.JobQueueError, exceptions.DistributedError))
+        self.assertTrue(issubclass(exceptions.TileCollectionError, exceptions.DistributedError))
+        self.assertTrue(issubclass(exceptions.TunnelError, exceptions.DistributedError))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/system/test_lifecycle.py b/tests/unit/system/test_lifecycle.py
new file mode 100644
index 0000000..e652e48
--- /dev/null
+++ b/tests/unit/system/test_lifecycle.py
@@ -0,0 +1,228 @@
+import importlib.util
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    workers_pkg = types.ModuleType(f"{package_name}.workers")
+    workers_pkg.__path__ = []
+    sys.modules[f"{package_name}.workers"] = workers_pkg
+
+    process_pkg = types.ModuleType(f"{package_name}.workers.process")
+    process_pkg.__path__ = []
+    sys.modules[f"{package_name}.workers.process"] = process_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_lifecycle_module():
+    package_name = "dist_lifecycle_testpkg"
+    _bootstrap_package(package_name)
+
+    state = {
+        "config": {"settings": {"stop_workers_on_master_exit": True}, "managed_processes": {}},
+        "saved_configs": [],
+        "launch_calls": [],
+        "terminated": [],
+        "alive_pids": set(),
+    }
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: state["config"]
+    config_module.save_config = lambda cfg: state["saved_configs"].append(dict(cfg))
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.PROCESS_TERMINATION_TIMEOUT = 5.0
+    constants_module.PROCESS_WAIT_TIMEOUT = 1.0
+    constants_module.WORKER_CHECK_INTERVAL = 0.01
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    process_module = types.ModuleType(f"{package_name}.utils.process")
+
+    class _FakeProcess:
+        def __init__(self, pid=1234, poll_result=None):
+            self.pid = pid
+            self._poll_result = poll_result
+            self.returncode = poll_result
+
+        def poll(self):
+            return self._poll_result
+
+    def _launch_process_with_timeout(command, timeout_seconds=10.0, **kwargs):
+        state["launch_calls"].append((list(command), float(timeout_seconds), dict(kwargs)))
+        return _FakeProcess(pid=4321, poll_result=None)
+
+    def _terminate_process(process, timeout=5):
+        state["terminated"].append((process.pid, timeout))
+        process._poll_result = 0
+        process.returncode = 0
+
+    process_module.launch_process_with_timeout = _launch_process_with_timeout
+    process_module.terminate_process = _terminate_process
+    process_module.is_process_alive = lambda pid: int(pid) in state["alive_pids"]
+    process_module.get_python_executable = lambda: "/usr/bin/python3"
+    process_module._FakeProcess = _FakeProcess
+    sys.modules[f"{package_name}.utils.process"] = process_module
+
+    lifecycle_module = _load_module(package_name, "workers/process/lifecycle.py", "workers.process.lifecycle")
+    return lifecycle_module, state, process_module
+
+
+lifecycle_module, test_state, process_stub_module = _load_lifecycle_module()
+
+
+class _FakeManager:
+    def __init__(self, comfy_root):
+        self._comfy_root = comfy_root
+        self.processes = {}
+        self.save_calls = 0
+
+    def find_comfy_root(self):
+        return self._comfy_root
+
+    def build_launch_command(self, worker_config, _comfy_root):
+        return ["python", "main.py", "--port", str(worker_config["port"])]
+
+    def save_processes(self):
+        self.save_calls += 1
+
+
+class ProcessLifecycleTests(unittest.TestCase):
+    def setUp(self):
+        test_state["config"] = {"settings": {"stop_workers_on_master_exit": True}, "managed_processes": {}}
+        test_state["saved_configs"].clear()
+        test_state["launch_calls"].clear()
+        test_state["terminated"].clear()
+        test_state["alive_pids"].clear()
+
+    def test_launch_worker_uses_monitor_wrapper_when_enabled(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+
+            pid = lifecycle.launch_worker(
+                {
+                    "id": "1",
+                    "name": "Worker A",
+                    "port": 8189,
+                    "cuda_device": 2,
+                }
+            )
+
+            self.assertEqual(pid, 4321)
+            self.assertEqual(manager.save_calls, 1)
+            self.assertIn("1", manager.processes)
+            self.assertTrue(manager.processes["1"]["is_monitor"])
+            self.assertIn("logs/workers", manager.processes["1"]["log_file"])
+            self.assertEqual(len(test_state["launch_calls"]), 1)
+            launched_command, timeout_seconds, launch_kwargs = test_state["launch_calls"][0]
+            self.assertEqual(timeout_seconds, 1.0)
+            self.assertEqual(launched_command[0], "/usr/bin/python3")
+            self.assertTrue(launched_command[1].endswith("workers/worker_monitor.py"))
+            self.assertEqual(launch_kwargs["env"]["CUDA_VISIBLE_DEVICES"], "2")
+
+    def test_launch_worker_skips_monitor_when_disabled(self):
+        test_state["config"] = {"settings": {"stop_workers_on_master_exit": False}, "managed_processes": {}}
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+
+            lifecycle.launch_worker(
+                {
+                    "id": "2",
+                    "name": "Worker-B",
+                    "port": 8190,
+                    "cuda_device": 0,
+                }
+            )
+
+            launched_command, _, _ = test_state["launch_calls"][0]
+            self.assertEqual(launched_command[0], "python")
+            self.assertFalse(manager.processes["2"]["is_monitor"])
+
+    def test_stop_worker_uses_fallback_terminate_when_tree_kill_fails(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+            running_process = process_stub_module._FakeProcess(pid=88, poll_result=None)
+            manager.processes["7"] = {
+                "pid": 88,
+                "process": running_process,
+                "started_at": 0.0,
+                "config": {},
+            }
+            lifecycle._kill_process_tree = lambda _pid: False
+
+            ok, message = lifecycle.stop_worker("7")
+
+            self.assertTrue(ok)
+            self.assertIn("fallback", message.lower())
+            self.assertEqual(test_state["terminated"], [(88, 5.0)])
+            self.assertNotIn("7", manager.processes)
+
+    def test_stop_worker_without_process_uses_tree_kill(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+            manager.processes["9"] = {
+                "pid": 99,
+                "process": None,
+                "started_at": 0.0,
+                "config": {},
+            }
+            lifecycle._kill_process_tree = lambda _pid: True
+
+            ok, message = lifecycle.stop_worker("9")
+
+            self.assertTrue(ok)
+            self.assertEqual(message, "Worker stopped")
+            self.assertNotIn("9", manager.processes)
+
+    def test_check_worker_process_falls_back_to_pid_probe(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = _FakeManager(tmpdir)
+            lifecycle = lifecycle_module.ProcessLifecycle(manager)
+            lifecycle._is_process_running = lambda pid: int(pid) == 42
+
+            running, from_subprocess = lifecycle._check_worker_process(
+                "any",
+                {"pid": 42, "process": None},
+            )
+            self.assertTrue(running)
+            self.assertFalse(from_subprocess)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/system/test_logging.py b/tests/unit/system/test_logging.py
new file mode 100644
index 0000000..587fb03
--- /dev/null
+++ b/tests/unit/system/test_logging.py
@@ -0,0 +1,66 @@
+import io
+import json
+import os
+import tempfile
+import unittest
+from contextlib import redirect_stdout
+
+from utils import logging as logging_utils
+
+
+class LoggingUtilsTests(unittest.TestCase):
+    def setUp(self):
+        self.original_config_file = logging_utils.GPU_CONFIG_FILE
+        self.state = logging_utils._debug_state()
+        self.original_cache = self.state.enabled
+        self.original_cache_time = self.state.updated_at
+
+    def tearDown(self):
+        logging_utils.GPU_CONFIG_FILE = self.original_config_file
+        self.state.enabled = self.original_cache
+        self.state.updated_at = self.original_cache_time
+
+    def _reset_cache(self):
+        self.state.enabled = None
+        self.state.updated_at = 0.0
+
+    def test_is_debug_enabled_reads_config_flag(self):
+        with tempfile.NamedTemporaryFile("w", delete=False) as temp_file:
+            json.dump({"settings": {"debug": True}}, temp_file)
+            temp_path = temp_file.name
+        self.addCleanup(lambda: os.path.exists(temp_path) and os.unlink(temp_path))
+
+        logging_utils.GPU_CONFIG_FILE = temp_path
+        self._reset_cache()
+
+        self.assertTrue(logging_utils.is_debug_enabled())
+
+    def test_is_debug_enabled_handles_invalid_json(self):
+        with tempfile.NamedTemporaryFile("w", delete=False) as temp_file:
+            temp_file.write("{invalid-json")
+            temp_path = temp_file.name
+        self.addCleanup(lambda: os.path.exists(temp_path) and os.unlink(temp_path))
+
+        logging_utils.GPU_CONFIG_FILE = temp_path
+        self._reset_cache()
+
+        self.assertFalse(logging_utils.is_debug_enabled())
+
+    def test_debug_log_emits_only_when_enabled(self):
+        self.state.enabled = True
+        self.state.updated_at = 1e18
+        output = io.StringIO()
+        with redirect_stdout(output):
+            logging_utils.debug_log("hello")
+        self.assertIn("[Distributed] hello", output.getvalue())
+
+        self.state.enabled = False
+        self.state.updated_at = 1e18
+        output = io.StringIO()
+        with redirect_stdout(output):
+            logging_utils.debug_log("hidden")
+        self.assertEqual(output.getvalue(), "")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_network_helpers.py b/tests/unit/system/test_network.py
similarity index 91%
rename from tests/test_network_helpers.py
rename to tests/unit/system/test_network.py
index 749986e..af4a1bd 100644
--- a/tests/test_network_helpers.py
+++ b/tests/unit/system/test_network.py
@@ -6,7 +6,7 @@ from pathlib import Path
 
 
 def _load_network_module():
-    module_path = Path(__file__).resolve().parents[1] / "utils" / "network.py"
+    module_path = Path(__file__).resolve().parents[3] / "utils" / "network.py"
 
     package_name = "dist_utils_testpkg"
     package_module = types.ModuleType(package_name)
@@ -17,6 +17,10 @@ def _load_network_module():
     logging_module.debug_log = lambda *_args, **_kwargs: None
     sys.modules[f"{package_name}.logging"] = logging_module
 
+    exceptions_module = types.ModuleType(f"{package_name}.exceptions")
+    exceptions_module.DistributedError = type("DistributedError", (Exception,), {})
+    sys.modules[f"{package_name}.exceptions"] = exceptions_module
+
     server_module = types.ModuleType("server")
     server_module.PromptServer = types.SimpleNamespace(
         instance=types.SimpleNamespace(address="127.0.0.1", port=8188, loop=None)
@@ -92,7 +96,7 @@ class NetworkHelpersTests(unittest.TestCase):
 
     def test_build_master_url_falls_back_to_server_address(self):
         cfg = {"master": {"host": ""}}
-        prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
+        prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)  # nosec B104 - wildcard bind normalization test
         self.assertEqual(
             network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
             "http://127.0.0.1:8190",
diff --git a/tests/test_prompt_transform.py b/tests/unit/system/test_prompt_transform.py
similarity index 99%
rename from tests/test_prompt_transform.py
rename to tests/unit/system/test_prompt_transform.py
index 9b8ea82..ee92728 100644
--- a/tests/test_prompt_transform.py
+++ b/tests/unit/system/test_prompt_transform.py
@@ -7,7 +7,7 @@ from pathlib import Path
 
 
 def _load_prompt_transform_module():
-    module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py"
+    module_path = Path(__file__).resolve().parents[3] / "api" / "orchestration" / "prompt_transform.py"
     package_name = "dist_pt_testpkg"
 
     for mod_name in list(sys.modules):
diff --git a/tests/test_queue_request.py b/tests/unit/system/test_queue_request.py
similarity index 98%
rename from tests/test_queue_request.py
rename to tests/unit/system/test_queue_request.py
index 6c150de..156388a 100644
--- a/tests/test_queue_request.py
+++ b/tests/unit/system/test_queue_request.py
@@ -4,7 +4,7 @@ from pathlib import Path
 
 
 def _load_queue_request_module():
-    module_path = Path(__file__).resolve().parents[1] / "api" / "queue_request.py"
+    module_path = Path(__file__).resolve().parents[3] / "api" / "queue_request.py"
     spec = importlib.util.spec_from_file_location("queue_request", module_path)
     module = importlib.util.module_from_spec(spec)
     assert spec is not None and spec.loader is not None
diff --git a/tests/unit/system/test_worker_monitor.py b/tests/unit/system/test_worker_monitor.py
new file mode 100644
index 0000000..ef80880
--- /dev/null
+++ b/tests/unit/system/test_worker_monitor.py
@@ -0,0 +1,154 @@
+import importlib.util
+import os
+import tempfile
+import types
+import unittest
+from pathlib import Path
+
+
+def _load_worker_monitor_module():
+    module_path = Path(__file__).resolve().parents[3] / "workers" / "worker_monitor.py"
+    spec = importlib.util.spec_from_file_location("dist_test_worker_monitor", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+worker_monitor = _load_worker_monitor_module()
+
+
+class _FakeWorkerProcess:
+    def __init__(self, pid=9001, poll_sequence=None):
+        self.pid = pid
+        self._poll_sequence = list(poll_sequence or [None])
+        self._poll_calls = 0
+        self.returncode = None
+        self.terminate_calls = 0
+        self.kill_calls = 0
+        self.wait_calls = 0
+
+    def poll(self):
+        if self._poll_calls < len(self._poll_sequence):
+            value = self._poll_sequence[self._poll_calls]
+        else:
+            value = self._poll_sequence[-1]
+        self._poll_calls += 1
+        if value is not None:
+            self.returncode = value
+        return value
+
+    def terminate(self):
+        self.terminate_calls += 1
+        self.returncode = 0
+
+    def kill(self):
+        self.kill_calls += 1
+        self.returncode = -9
+
+    def wait(self, timeout=None):
+        _ = timeout
+        self.wait_calls += 1
+        if self.returncode is None:
+            self.returncode = 0
+        return self.returncode
+
+
+class WorkerMonitorTests(unittest.TestCase):
+    def setUp(self):
+        self._orig_launch = worker_monitor.launch_process_with_timeout
+        self._orig_is_alive = worker_monitor.is_process_alive
+        self._orig_terminate = getattr(worker_monitor, "terminate_process", None)
+        self._orig_sleep = worker_monitor.time.sleep
+        self._orig_signal = worker_monitor.signal.signal
+        self._env_backup = dict(os.environ)
+
+        worker_monitor.time.sleep = lambda _seconds: None
+        worker_monitor.signal.signal = lambda *_args, **_kwargs: None
+
+    def tearDown(self):
+        worker_monitor.launch_process_with_timeout = self._orig_launch
+        worker_monitor.is_process_alive = self._orig_is_alive
+        if self._orig_terminate is not None:
+            worker_monitor.terminate_process = self._orig_terminate
+        worker_monitor.time.sleep = self._orig_sleep
+        worker_monitor.signal.signal = self._orig_signal
+        os.environ.clear()
+        os.environ.update(self._env_backup)
+
+    def test_main_validates_required_env_and_args(self):
+        os.environ.pop("COMFYUI_MASTER_PID", None)
+        self.assertEqual(worker_monitor.main(["python", "main.py"]), 1)
+
+        os.environ["COMFYUI_MASTER_PID"] = "not-an-int"
+        self.assertEqual(worker_monitor.main(["python", "main.py"]), 1)
+
+        os.environ["COMFYUI_MASTER_PID"] = "101"
+        self.assertEqual(worker_monitor.main([]), 1)
+
+    def test_main_delegates_to_monitor_and_run(self):
+        calls = []
+
+        def _monitor(master_pid, command):
+            calls.append((master_pid, list(command)))
+            return 7
+
+        os.environ["COMFYUI_MASTER_PID"] = "202"
+        original_monitor = worker_monitor.monitor_and_run
+        worker_monitor.monitor_and_run = _monitor
+        try:
+            exit_code = worker_monitor.main(["python", "worker.py"])
+        finally:
+            worker_monitor.monitor_and_run = original_monitor
+
+        self.assertEqual(exit_code, 7)
+        self.assertEqual(calls, [(202, ["python", "worker.py"])])
+
+    def test_monitor_and_run_returns_worker_exit_code(self):
+        fake_process = _FakeWorkerProcess(poll_sequence=[None, 5])
+        worker_monitor.launch_process_with_timeout = lambda *_args, **_kwargs: fake_process
+        worker_monitor.is_process_alive = lambda _pid: True
+        worker_monitor.terminate_process = lambda *_args, **_kwargs: None
+
+        exit_code = worker_monitor.monitor_and_run(100, ["python", "worker.py"])
+
+        self.assertEqual(exit_code, 5)
+
+    def test_monitor_and_run_terminates_worker_when_master_dies(self):
+        fake_process = _FakeWorkerProcess(poll_sequence=[None, None])
+        term_calls = []
+
+        worker_monitor.launch_process_with_timeout = lambda *_args, **_kwargs: fake_process
+        worker_monitor.is_process_alive = lambda _pid: False
+
+        def _terminate(process, timeout):
+            term_calls.append((process.pid, timeout))
+            process.returncode = 0
+
+        worker_monitor.terminate_process = _terminate
+
+        exit_code = worker_monitor.monitor_and_run(333, ["python", "worker.py"])
+
+        self.assertEqual(exit_code, 0)
+        self.assertEqual(term_calls, [(9001, worker_monitor.PROCESS_TERMINATION_TIMEOUT)])
+
+    def test_monitor_and_run_writes_pid_file_when_requested(self):
+        fake_process = _FakeWorkerProcess(poll_sequence=[0])
+        worker_monitor.launch_process_with_timeout = lambda *_args, **_kwargs: fake_process
+        worker_monitor.is_process_alive = lambda _pid: True
+        worker_monitor.terminate_process = lambda *_args, **_kwargs: None
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            pid_file = Path(tmpdir) / "worker.pid"
+            os.environ["WORKER_PID_FILE"] = str(pid_file)
+            exit_code = worker_monitor.monitor_and_run(444, ["python", "worker.py"])
+
+            self.assertEqual(exit_code, 0)
+            contents = pid_file.read_text(encoding="utf-8")
+            monitor_pid, worker_pid = contents.split(",")
+            self.assertTrue(monitor_pid.isdigit())
+            self.assertEqual(worker_pid, "9001")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/upscale/__init__.py b/tests/unit/upscale/__init__.py
new file mode 100644
index 0000000..e0310a0
--- /dev/null
+++ b/tests/unit/upscale/__init__.py
@@ -0,0 +1 @@
+"""Unit tests."""
diff --git a/tests/unit/upscale/test_conditioning.py b/tests/unit/upscale/test_conditioning.py
new file mode 100644
index 0000000..2dcf41d
--- /dev/null
+++ b/tests/unit/upscale/test_conditioning.py
@@ -0,0 +1,53 @@
+import unittest
+
+import torch
+
+from upscale.conditioning import clone_conditioning, clone_control_chain
+
+
+class _FakeControl:
+    def __init__(self, hint, previous=None):
+        self.cond_hint_original = hint
+        self.previous_controlnet = previous
+
+
+class ConditioningUtilsTests(unittest.TestCase):
+    def test_clone_control_chain_clones_hints_and_links(self):
+        tail = _FakeControl(torch.ones((1, 2)))
+        head = _FakeControl(torch.zeros((1, 2)), previous=tail)
+
+        cloned = clone_control_chain(head, clone_hint=True)
+
+        self.assertIsNot(cloned, head)
+        self.assertIsNot(cloned.previous_controlnet, tail)
+        self.assertTrue(torch.equal(cloned.cond_hint_original, head.cond_hint_original))
+        self.assertTrue(torch.equal(cloned.previous_controlnet.cond_hint_original, tail.cond_hint_original))
+        self.assertIsNot(cloned.cond_hint_original, head.cond_hint_original)
+
+    def test_clone_conditioning_clones_emb_and_known_fields(self):
+        control = _FakeControl(torch.ones((1, 3)))
+        cond = [
+            [
+                torch.tensor([[1.0, 2.0, 3.0]]),
+                {
+                    "control": control,
+                    "mask": torch.ones((1, 1, 1)),
+                    "pooled_output": torch.zeros((1, 4)),
+                    "area": [0, 1, 2, 3],
+                },
+            ]
+        ]
+
+        cloned = clone_conditioning(cond, clone_hints=True)
+
+        self.assertTrue(torch.equal(cloned[0][0], cond[0][0]))
+        self.assertIsNot(cloned[0][0], cond[0][0])
+        self.assertIsNot(cloned[0][1]["control"], cond[0][1]["control"])
+        self.assertIsNot(cloned[0][1]["mask"], cond[0][1]["mask"])
+        self.assertIsNot(cloned[0][1]["pooled_output"], cond[0][1]["pooled_output"])
+        self.assertEqual(cloned[0][1]["area"], cond[0][1]["area"])
+        self.assertIsNot(cloned[0][1]["area"], cond[0][1]["area"])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/upscale/test_dynamic.py b/tests/unit/upscale/test_dynamic.py
new file mode 100644
index 0000000..e4686f7
--- /dev/null
+++ b/tests/unit/upscale/test_dynamic.py
@@ -0,0 +1,234 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+try:
+    import torch
+except ModuleNotFoundError:  # pragma: no cover - optional dependency in CI/runtime
+    torch = None
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
+    modes_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    def pil_to_tensor(image):
+        if torch is None:
+            raise RuntimeError("torch is required for this test helper")
+        array = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(array).unsqueeze(0)
+
+    image_module.tensor_to_pil = tensor_to_pil
+    image_module.pil_to_tensor = pil_to_tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.run_async_in_server_loop = lambda _coro, timeout=None: (_ for _ in ()).throw(
+        RuntimeError(f"run_async_in_server_loop unexpectedly called with timeout={timeout}")
+    )
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.get_worker_timeout_seconds = lambda: 5.0
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.TILE_WAIT_TIMEOUT = 5.0
+    constants_module.TILE_SEND_TIMEOUT = 5.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    if torch is None:
+        torch_module = types.ModuleType("torch")
+        torch_module.cat = lambda *_args, **_kwargs: None
+        sys.modules["torch"] = torch_module
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: None
+
+    async def init_dynamic_job(*_args, **_kwargs):
+        return None
+
+    job_store_module.init_dynamic_job = init_dynamic_job
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    comfy_module = types.ModuleType("comfy")
+    model_mgmt = types.ModuleType("comfy.model_management")
+    model_mgmt.throw_exception_if_processing_interrupted = lambda: None
+    comfy_module.model_management = model_mgmt
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.model_management"] = model_mgmt
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_dynamic_mode_module():
+    package_name = "dist_dynamic_mode_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args")
+    return _load_module(package_name, "upscale/modes/dynamic.py", "upscale.modes.dynamic")
+
+
+dynamic_module = _load_dynamic_mode_module()
+
+
+def _run_sync(value, timeout=None):
+    _ = timeout
+    if asyncio.iscoroutine(value):
+        return asyncio.run(value)
+    return value
+
+
+dynamic_module.run_async_in_server_loop = _run_sync
+
+
+class _DummyDynamicNode(dynamic_module.DynamicModeMixin):
+    def __init__(self, *, poll_ready=True, assignments=None):
+        self._poll_ready = poll_ready
+        self._assignments = list(assignments or [])
+        self.sent_images = []
+        self.heartbeat_count = 0
+        self.completion_sent = False
+
+    def round_to_multiple(self, value):
+        return int(value)
+
+    def calculate_tiles(self, _width, _height, _tile_width, _tile_height, _force_uniform_tiles):
+        return [(0, 0)]
+
+    def _poll_job_ready(self, *_args, **_kwargs):
+        return self._poll_ready
+
+    async def _request_image_from_master(self, *_args, **_kwargs):
+        if self._assignments:
+            return self._assignments.pop(0)
+        return (None, 0)
+
+    def _slice_conditioning(self, positive, negative, _image_idx):
+        return positive, negative
+
+    def _process_and_blend_tile(self, _tile_idx, _pos, _source_tensor, local_image, *_args, **_kwargs):
+        return local_image
+
+    async def _send_heartbeat_to_master(self, *_args, **_kwargs):
+        self.heartbeat_count += 1
+
+    async def _send_full_image_to_master(self, local_image, image_idx, _multi_job_id, _master_url, _worker_id, is_last):
+        self.sent_images.append((image_idx, is_last, local_image.size))
+
+    async def _send_worker_complete_signal(self, *_args, **_kwargs):
+        self.completion_sent = True
+
+
+@unittest.skipIf(torch is None, "torch is not installed")
+class DynamicModeTests(unittest.TestCase):
+    def setUp(self):
+        self.image_batch = torch.zeros((1, 8, 8, 3), dtype=torch.float32)
+        self.core_args = dynamic_module.UpscaleCoreArgs(
+            model=None,
+            positive=None,
+            negative=None,
+            vae=None,
+            seed=1,
+            steps=10,
+            cfg=7.0,
+            sampler_name="euler",
+            scheduler="normal",
+            denoise=0.5,
+            tiled_decode=False,
+        )
+
+    def test_worker_dynamic_returns_early_if_job_not_ready(self):
+        node = _DummyDynamicNode(poll_ready=False)
+
+        result = node.process_worker_dynamic(
+            upscaled_image=self.image_batch,
+            core_args=self.core_args,
+            tile_width=8,
+            tile_height=8,
+            padding=4,
+            mask_blur=2,
+            force_uniform_tiles=True,
+            multi_job_id="job-1",
+            master_url="http://master.local:8188",
+            worker_id="worker-a",
+            enabled_worker_ids='["worker-a"]',
+            dynamic_threshold=8,
+        )
+
+        self.assertEqual(result[0].shape, self.image_batch.shape)
+        self.assertEqual(node.sent_images, [])
+        self.assertFalse(node.completion_sent)
+
+    def test_worker_dynamic_processes_assigned_image_and_sends_completion(self):
+        node = _DummyDynamicNode(poll_ready=True, assignments=[(0, 0)])
+
+        result = node.process_worker_dynamic(
+            upscaled_image=self.image_batch,
+            core_args=self.core_args,
+            tile_width=8,
+            tile_height=8,
+            padding=4,
+            mask_blur=2,
+            force_uniform_tiles=True,
+            multi_job_id="job-2",
+            master_url="http://master.local:8188",
+            worker_id="worker-a",
+            enabled_worker_ids='["worker-a"]',
+            dynamic_threshold=8,
+        )
+
+        self.assertEqual(result[0].shape, self.image_batch.shape)
+        self.assertEqual(len(node.sent_images), 1)
+        self.assertEqual(node.sent_images[0][0], 0)
+        self.assertTrue(node.sent_images[0][1])
+        self.assertTrue(node.completion_sent)
+        self.assertGreaterEqual(node.heartbeat_count, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/upscale/test_image.py b/tests/unit/upscale/test_image.py
new file mode 100644
index 0000000..62ec441
--- /dev/null
+++ b/tests/unit/upscale/test_image.py
@@ -0,0 +1,39 @@
+import unittest
+
+import torch
+from PIL import Image
+
+from utils.image import ensure_contiguous, pil_to_tensor, tensor_to_pil
+
+
+class ImageUtilsTests(unittest.TestCase):
+    def test_tensor_to_pil_converts_batch_item(self):
+        tensor = torch.tensor(
+            [[[[0.0, 0.5, 1.0], [1.0, 0.0, 0.0]]]],
+            dtype=torch.float32,
+        )
+        image = tensor_to_pil(tensor, batch_index=0)
+
+        self.assertEqual(image.size, (2, 1))
+        self.assertEqual(image.getpixel((0, 0)), (0, 127, 255))
+        self.assertEqual(image.getpixel((1, 0)), (255, 0, 0))
+
+    def test_pil_to_tensor_adds_channel_for_grayscale(self):
+        grayscale = Image.new("L", (3, 2), color=128)
+        tensor = pil_to_tensor(grayscale)
+
+        self.assertEqual(tuple(tensor.shape), (1, 2, 3, 1))
+        self.assertAlmostEqual(float(tensor[0, 0, 0, 0]), 128 / 255.0, places=6)
+
+    def test_ensure_contiguous_makes_non_contiguous_tensor_contiguous(self):
+        base = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)
+        non_contiguous = base.transpose(1, 2)
+        self.assertFalse(non_contiguous.is_contiguous())
+
+        contiguous = ensure_contiguous(non_contiguous)
+        self.assertTrue(contiguous.is_contiguous())
+        self.assertTrue(torch.equal(contiguous, non_contiguous))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/upscale/test_job_state.py b/tests/unit/upscale/test_job_state.py
new file mode 100644
index 0000000..a5f32a8
--- /dev/null
+++ b/tests/unit/upscale/test_job_state.py
@@ -0,0 +1,131 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_job_state_modules():
+    package_name = "dist_job_state_testpkg"
+    _bootstrap_package(package_name)
+    job_models = _load_module(package_name, "upscale/job_models.py", "upscale.job_models")
+
+    prompt_server = types.SimpleNamespace(
+        distributed_pending_tile_jobs={},
+        distributed_tile_jobs_lock=asyncio.Lock(),
+    )
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    call_log = {"args": None}
+    job_timeout_module = types.ModuleType(f"{package_name}.upscale.job_timeout")
+
+    async def _requeue(multi_job_id, batch_size):
+        call_log["args"] = (multi_job_id, batch_size)
+        return 3
+
+    job_timeout_module._check_and_requeue_timed_out_workers = _requeue
+    sys.modules[f"{package_name}.upscale.job_timeout"] = job_timeout_module
+
+    job_state = _load_module(package_name, "upscale/job_state.py", "upscale.job_state")
+    return job_models, job_state, prompt_server, call_log
+
+
+job_models_module, job_state_module, prompt_server, timeout_call_log = _load_job_state_modules()
+
+
+class _Node(job_state_module.JobStateMixin):
+    pass
+
+
+class JobStateMixinTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        prompt_server.distributed_pending_tile_jobs = {}
+        prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
+        self.node = _Node()
+
+    async def test_get_all_completed_tasks_supports_tile_and_image_modes(self):
+        tile_state = job_models_module.TileJobState("tile-job")
+        tile_state.completed_tasks[4] = "tile-4"
+        image_state = job_models_module.ImageJobState("image-job")
+        image_state.completed_images[2] = "image-2"
+        prompt_server.distributed_pending_tile_jobs = {
+            "tile-job": tile_state,
+            "image-job": image_state,
+        }
+
+        self.assertEqual(await self.node._get_all_completed_tasks("tile-job"), {4: "tile-4"})
+        self.assertEqual(await self.node._get_all_completed_tasks("image-job"), {2: "image-2"})
+
+    async def test_next_index_and_pending_count_from_queues(self):
+        image_state = job_models_module.ImageJobState("image-job")
+        await image_state.pending_images.put(7)
+        tile_state = job_models_module.TileJobState("tile-job")
+        await tile_state.pending_tasks.put(9)
+        prompt_server.distributed_pending_tile_jobs = {
+            "image-job": image_state,
+            "tile-job": tile_state,
+        }
+
+        self.assertEqual(await self.node._get_next_image_index("image-job"), 7)
+        self.assertEqual(await self.node._get_next_tile_index("tile-job"), 9)
+        self.assertEqual(await self.node._get_pending_count("image-job"), 0)
+        self.assertEqual(await self.node._get_pending_count("tile-job"), 0)
+
+    async def test_drain_worker_results_queue_collects_completed_images(self):
+        image_state = job_models_module.ImageJobState("image-job")
+        await image_state.queue.put({"worker_id": "w1", "image_idx": 1, "image": "img1"})
+        await image_state.queue.put({"worker_id": "w2", "image_idx": 2, "image": "img2"})
+        prompt_server.distributed_pending_tile_jobs = {"image-job": image_state}
+
+        drained = await self.node._drain_worker_results_queue("image-job")
+
+        self.assertEqual(drained, 2)
+        self.assertEqual(image_state.completed_images[1], "img1")
+        self.assertEqual(image_state.completed_images[2], "img2")
+
+    async def test_check_and_requeue_delegates_to_timeout_helper(self):
+        count = await self.node._check_and_requeue_timed_out_workers("job-abc", 5)
+        self.assertEqual(count, 3)
+        self.assertEqual(timeout_call_log["args"], ("job-abc", 5))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/upscale/test_job_store.py b/tests/unit/upscale/test_job_store.py
new file mode 100644
index 0000000..338cb65
--- /dev/null
+++ b/tests/unit/upscale/test_job_store.py
@@ -0,0 +1,140 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace())
+    sys.modules["server"] = server_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_job_store_modules():
+    package_name = "dist_job_store_testpkg"
+    _bootstrap_package(package_name)
+    job_models = _load_module(package_name, "upscale/job_models.py", "upscale.job_models")
+    job_store = _load_module(package_name, "upscale/job_store.py", "upscale.job_store")
+    return job_models, job_store
+
+
+job_models_module, job_store_module = _load_job_store_modules()
+
+
+class JobStoreTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.prompt_server = job_store_module.server.PromptServer.instance
+        self.prompt_server.distributed_pending_tile_jobs = {}
+        self.prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
+
+    async def test_ensure_tile_jobs_initialized_prunes_invalid_entries(self):
+        self.prompt_server.distributed_pending_tile_jobs["valid"] = job_models_module.TileJobState("job-valid")
+        self.prompt_server.distributed_pending_tile_jobs["invalid"] = {"not": "state"}
+
+        job_store_module.ensure_tile_jobs_initialized()
+
+        self.assertIn("valid", self.prompt_server.distributed_pending_tile_jobs)
+        self.assertNotIn("invalid", self.prompt_server.distributed_pending_tile_jobs)
+
+    async def test_init_dynamic_job_populates_pending_images(self):
+        await job_store_module.init_dynamic_job(
+            multi_job_id="job-dynamic",
+            batch_size=3,
+            enabled_workers=["worker-a"],
+            all_indices=[2, 0, 1],
+        )
+
+        job_data = self.prompt_server.distributed_pending_tile_jobs["job-dynamic"]
+        self.assertIsInstance(job_data, job_models_module.ImageJobState)
+        self.assertEqual(job_data.batch_size, 3)
+        self.assertIn("worker-a", job_data.worker_status)
+
+        pulled = [job_data.pending_images.get_nowait() for _ in range(3)]
+        self.assertEqual(pulled, [2, 0, 1])
+
+    async def test_init_static_job_batched_queues_tile_ids(self):
+        await job_store_module.init_static_job_batched(
+            multi_job_id="job-static",
+            batch_size=2,
+            num_tiles_per_image=4,
+            enabled_workers=["worker-a", "worker-b"],
+        )
+
+        job_data = self.prompt_server.distributed_pending_tile_jobs["job-static"]
+        self.assertIsInstance(job_data, job_models_module.TileJobState)
+        self.assertTrue(job_data.batched_static)
+
+        pulled = [job_data.pending_tasks.get_nowait() for _ in range(4)]
+        self.assertEqual(pulled, [0, 1, 2, 3])
+
+    async def test_drain_results_queue_collects_images_tiles_and_completion(self):
+        job_data = job_models_module.TileJobState("job-drain")
+        job_data.worker_status = {"worker-a": 1.0, "worker-b": 1.0}
+        self.prompt_server.distributed_pending_tile_jobs["job-drain"] = job_data
+
+        await job_data.queue.put(
+            {
+                "worker_id": "worker-a",
+                "is_last": True,
+                "image_idx": 5,
+                "image": "img-5",
+            }
+        )
+        await job_data.queue.put(
+            {
+                "worker_id": "worker-b",
+                "is_last": False,
+                "tiles": [
+                    {"tile_idx": 0, "global_idx": 0, "tensor": "tile-0"},
+                    {"tile_idx": 1, "global_idx": 1, "tensor": "tile-1"},
+                ],
+            }
+        )
+
+        drained = await job_store_module._drain_results_queue("job-drain")
+
+        self.assertEqual(drained, 3)
+        self.assertEqual(job_data.completed_tasks[5], "img-5")
+        self.assertEqual(job_data.completed_tasks[0]["tensor"], "tile-0")
+        self.assertEqual(job_data.completed_tasks[1]["tensor"], "tile-1")
+        self.assertNotIn("worker-a", job_data.worker_status)
+        self.assertIn("worker-b", job_data.worker_status)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_job_timeout.py b/tests/unit/upscale/test_job_timeout.py
similarity index 99%
rename from tests/test_job_timeout.py
rename to tests/unit/upscale/test_job_timeout.py
index 451a0a3..a57aa90 100644
--- a/tests/test_job_timeout.py
+++ b/tests/unit/upscale/test_job_timeout.py
@@ -9,7 +9,7 @@ from pathlib import Path
 
 
 def _load_job_timeout_module():
-    module_path = Path(__file__).resolve().parents[1] / "upscale" / "job_timeout.py"
+    module_path = Path(__file__).resolve().parents[3] / "upscale" / "job_timeout.py"
     package_name = "dist_job_timeout_testpkg"
 
     for mod_name in list(sys.modules):
diff --git a/tests/test_payload_parsers.py b/tests/unit/upscale/test_payload_parsers.py
similarity index 99%
rename from tests/test_payload_parsers.py
rename to tests/unit/upscale/test_payload_parsers.py
index 8043a0b..e793b7c 100644
--- a/tests/test_payload_parsers.py
+++ b/tests/unit/upscale/test_payload_parsers.py
@@ -15,7 +15,7 @@ except ImportError:
 
 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"
+    module_path = Path(__file__).resolve().parents[3] / "upscale" / "payload_parsers.py"
     spec = importlib.util.spec_from_file_location("upscale_payload_parsers", module_path)
     module = importlib.util.module_from_spec(spec)
     assert spec is not None and spec.loader is not None
diff --git a/tests/unit/upscale/test_result_collector.py b/tests/unit/upscale/test_result_collector.py
new file mode 100644
index 0000000..f42666e
--- /dev/null
+++ b/tests/unit/upscale/test_result_collector.py
@@ -0,0 +1,287 @@
+import asyncio
+import importlib.util
+import sys
+import time
+import types
+import unittest
+from pathlib import Path
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.get_worker_timeout_seconds = lambda: 0.05
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.DYNAMIC_MODE_MAX_POLL_TIMEOUT = 10.0
+    constants_module.HEARTBEAT_INTERVAL = 0.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    comfy_module = types.ModuleType("comfy")
+    model_mgmt = types.ModuleType("comfy.model_management")
+
+    class _InterruptProcessingException(Exception):
+        pass
+
+    model_mgmt.processing_interrupted = lambda: False
+    model_mgmt.InterruptProcessingException = _InterruptProcessingException
+    comfy_module.model_management = model_mgmt
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.model_management"] = model_mgmt
+
+    prompt_server = types.SimpleNamespace(
+        distributed_pending_tile_jobs={},
+        distributed_tile_jobs_lock=asyncio.Lock(),
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
+    sys.modules["server"] = server_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_result_collector_modules():
+    package_name = "dist_result_collector_testpkg"
+    _bootstrap_package(package_name)
+    job_models = _load_module(package_name, "upscale/job_models.py", "upscale.job_models")
+    job_store = _load_module(package_name, "upscale/job_store.py", "upscale.job_store")
+
+    job_timeout = types.ModuleType(f"{package_name}.upscale.job_timeout")
+
+    async def _noop_requeue(_multi_job_id, _batch_size):
+        return 0
+
+    job_timeout._check_and_requeue_timed_out_workers = _noop_requeue
+    job_timeout.check_and_requeue_timed_out_workers = _noop_requeue
+    sys.modules[f"{package_name}.upscale.job_timeout"] = job_timeout
+
+    result_collector = _load_module(
+        package_name,
+        "upscale/result_collector.py",
+        "upscale.result_collector",
+    )
+    return job_models, job_store, result_collector
+
+
+job_models_module, job_store_module, result_collector_module = _load_result_collector_modules()
+
+
+class _Node(result_collector_module.ResultCollectorMixin):
+    def __init__(self):
+        self.requeue_calls = []
+
+    async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
+        self.requeue_calls.append((multi_job_id, batch_size))
+        return 1
+
+
+class _WrapperNode(result_collector_module.ResultCollectorMixin):
+    def __init__(self):
+        self.calls = []
+
+    async def _async_collect_results(
+        self,
+        multi_job_id,
+        num_workers,
+        mode="static",
+        remaining_to_collect=None,
+        batch_size=None,
+    ):
+        self.calls.append(
+            (multi_job_id, num_workers, mode, remaining_to_collect, batch_size)
+        )
+        return {"mode": mode}
+
+
+class ResultCollectorMixinTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.prompt_server = job_store_module.server.PromptServer.instance
+        self.prompt_server.distributed_pending_tile_jobs = {}
+        self.prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
+        self.node = _Node()
+
+    async def test_static_collection_collects_tiles_and_cleans_up_job(self):
+        job_data = job_models_module.TileJobState("job-static")
+        self.prompt_server.distributed_pending_tile_jobs["job-static"] = job_data
+        await job_data.queue.put(
+            {
+                "worker_id": "worker-a",
+                "is_last": True,
+                "tiles": [
+                    {"tile_idx": 3},
+                    {
+                        "tile_idx": 2,
+                        "x": 10,
+                        "y": 20,
+                        "extracted_width": 64,
+                        "extracted_height": 64,
+                        "padding": 8,
+                        "batch_idx": 1,
+                        "global_idx": 9,
+                        "tensor": "tile-9",
+                    },
+                ],
+            }
+        )
+
+        collected = await self.node._async_collect_results(
+            "job-static",
+            num_workers=1,
+            mode="static",
+        )
+
+        self.assertEqual(list(collected.keys()), [9])
+        self.assertEqual(collected[9]["tensor"], "tile-9")
+        self.assertEqual(collected[9]["worker_id"], "worker-a")
+        self.assertNotIn("job-static", self.prompt_server.distributed_pending_tile_jobs)
+
+    async def test_dynamic_collection_stops_at_remaining_target_and_cleans_up(self):
+        job_data = job_models_module.ImageJobState("job-dynamic")
+        self.prompt_server.distributed_pending_tile_jobs["job-dynamic"] = job_data
+        await job_data.queue.put(
+            {
+                "worker_id": "worker-a",
+                "is_last": False,
+                "image_idx": 2,
+                "image": "img-2",
+            }
+        )
+
+        completed = await self.node._async_collect_results(
+            "job-dynamic",
+            num_workers=2,
+            mode="dynamic",
+            remaining_to_collect=1,
+            batch_size=3,
+        )
+
+        self.assertEqual(completed[2], "img-2")
+        self.assertNotIn("job-dynamic", self.prompt_server.distributed_pending_tile_jobs)
+
+    async def test_dynamic_timeout_calls_requeue_helper(self):
+        job_data = job_models_module.ImageJobState("job-timeout")
+        job_data.worker_status = {"worker-a": time.time() - 5}
+        self.prompt_server.distributed_pending_tile_jobs["job-timeout"] = job_data
+
+        old_timeout = result_collector_module.get_worker_timeout_seconds
+        result_collector_module.get_worker_timeout_seconds = lambda: 0.01
+        try:
+            completed = await self.node._async_collect_results(
+                "job-timeout",
+                num_workers=1,
+                mode="dynamic",
+                batch_size=4,
+            )
+        finally:
+            result_collector_module.get_worker_timeout_seconds = old_timeout
+
+        self.assertEqual(completed, {})
+        self.assertEqual(self.node.requeue_calls, [("job-timeout", 4)])
+        self.assertNotIn("job-timeout", self.prompt_server.distributed_pending_tile_jobs)
+
+    async def test_log_worker_timeout_status_handles_non_job_state(self):
+        worker_ids = self.node._log_worker_timeout_status(
+            {"worker_status": {"worker-a": time.time()}},
+            current_time=time.time(),
+            multi_job_id="job-invalid",
+        )
+        self.assertEqual(worker_ids, [])
+
+    async def test_log_worker_timeout_status_returns_worker_ids_for_job_state(self):
+        job_data = job_models_module.ImageJobState("job-state")
+        job_data.worker_status = {"worker-a": time.time() - 1}
+
+        worker_ids = self.node._log_worker_timeout_status(
+            job_data,
+            current_time=time.time(),
+            multi_job_id="job-state",
+        )
+        self.assertEqual(worker_ids, ["worker-a"])
+
+    async def test_mode_mismatch_and_unsupported_mode_raise(self):
+        self.prompt_server.distributed_pending_tile_jobs["job-mismatch"] = (
+            job_models_module.TileJobState("job-mismatch")
+        )
+        with self.assertRaises(RuntimeError):
+            await self.node._async_collect_results(
+                "job-mismatch",
+                num_workers=1,
+                mode="dynamic",
+                batch_size=1,
+            )
+
+        self.prompt_server.distributed_pending_tile_jobs["job-unsupported"] = (
+            job_models_module.TileJobState("job-unsupported")
+        )
+        with self.assertRaises(RuntimeError):
+            await self.node._async_collect_results(
+                "job-unsupported",
+                num_workers=1,
+                mode="unsupported",
+            )
+
+    async def test_mark_image_completed_updates_job_state(self):
+        job_data = job_models_module.ImageJobState("job-mark")
+        self.prompt_server.distributed_pending_tile_jobs["job-mark"] = job_data
+
+        await self.node._mark_image_completed("job-mark", 7, "img-7")
+
+        self.assertEqual(job_data.completed_images[7], "img-7")
+
+    async def test_wrapper_methods_delegate_to_async_collect_results(self):
+        node = _WrapperNode()
+
+        tile_result = await node._async_collect_worker_tiles("job-wrap", 3)
+        image_result = await node._async_collect_dynamic_images(
+            "job-wrap",
+            remaining_to_collect=2,
+            num_workers=4,
+            batch_size=8,
+            master_processed_count=0,
+        )
+
+        self.assertEqual(tile_result["mode"], "static")
+        self.assertEqual(image_result["mode"], "dynamic")
+        self.assertEqual(
+            node.calls,
+            [
+                ("job-wrap", 3, "static", None, None),
+                ("job-wrap", 4, "dynamic", 2, 8),
+            ],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/upscale/test_single_gpu.py b/tests/unit/upscale/test_single_gpu.py
new file mode 100644
index 0000000..6a2985b
--- /dev/null
+++ b/tests/unit/upscale/test_single_gpu.py
@@ -0,0 +1,206 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+try:
+    import torch
+except ModuleNotFoundError:  # pragma: no cover - optional dependency in CI/runtime
+    torch = None
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
+    modes_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    def pil_to_tensor(image):
+        if torch is None:
+            raise RuntimeError("torch is required for this test helper")
+        array = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(array).unsqueeze(0)
+
+    def blend_processed_batch_item(
+        result_images,
+        processed_batch,
+        batch_index,
+        blend_fn,
+        x1,
+        y1,
+        ew,
+        eh,
+        tile_mask,
+        padding,
+    ):
+        tile_pil = tensor_to_pil(processed_batch, batch_index)
+        if tile_pil.size != (ew, eh):
+            tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
+        result_images[batch_index] = blend_fn(
+            result_images[batch_index],
+            tile_pil,
+            x1,
+            y1,
+            (ew, eh),
+            tile_mask,
+            padding,
+        )
+
+    image_module.tensor_to_pil = tensor_to_pil
+    image_module.pil_to_tensor = pil_to_tensor
+    image_module.blend_processed_batch_item = blend_processed_batch_item
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    if torch is None:
+        torch_module = types.ModuleType("torch")
+        torch_module.cat = lambda *_args, **_kwargs: None
+        sys.modules["torch"] = torch_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_single_gpu_mode_module():
+    package_name = "dist_single_gpu_mode_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args")
+    _load_module(package_name, "upscale/tile_processing.py", "upscale.tile_processing")
+    return _load_module(package_name, "upscale/modes/single_gpu.py", "upscale.modes.single_gpu")
+
+
+single_gpu_module = _load_single_gpu_mode_module()
+
+
+class _DummySingleGpuNode(single_gpu_module.SingleGpuModeMixin):
+    def __init__(self):
+        self.round_calls = []
+        self.process_batch_calls = 0
+
+    def round_to_multiple(self, value):
+        self.round_calls.append(int(value))
+        return int(value)
+
+    def calculate_tiles(self, _width, _height, _tile_width, _tile_height, _force_uniform_tiles):
+        return [(0, 0)]
+
+    def create_tile_mask(self, _width, _height, _tx, _ty, _tile_width, _tile_height, _mask_blur):
+        return None
+
+    def extract_batch_tile_with_padding(self, source_batch, _tx, _ty, _tile_width, _tile_height, _padding, _force_uniform_tiles):
+        _, h, w, _ = source_batch.shape
+        return source_batch, 0, 0, w, h
+
+    def process_tiles_batch(
+        self,
+        tile_batch,
+        _model,
+        _positive,
+        _negative,
+        _vae,
+        _seed,
+        _steps,
+        _cfg,
+        _sampler_name,
+        _scheduler,
+        _denoise,
+        _tiled_decode,
+        _region,
+        _canvas_shape,
+    ):
+        self.process_batch_calls += 1
+        return tile_batch
+
+    def blend_tile(self, _base_image, tile_pil, _x1, _y1, _size, _tile_mask, _padding):
+        return tile_pil
+
+
+@unittest.skipIf(torch is None, "torch is not installed")
+class SingleGpuModeTests(unittest.TestCase):
+    def setUp(self):
+        self.node = _DummySingleGpuNode()
+        self.images = torch.rand((2, 8, 8, 3), dtype=torch.float32)
+        self.core_args = single_gpu_module.UpscaleCoreArgs(
+            model=None,
+            positive=None,
+            negative=None,
+            vae=None,
+            seed=1,
+            steps=10,
+            cfg=7.5,
+            sampler_name="euler",
+            scheduler="normal",
+            denoise=0.5,
+            tiled_decode=False,
+        )
+
+    def test_process_single_gpu_preserves_batch_shape(self):
+        result = self.node.process_single_gpu(
+            upscaled_image=self.images,
+            core_args=self.core_args,
+            tile_width=8,
+            tile_height=8,
+            padding=4,
+            mask_blur=2,
+            force_uniform_tiles=True,
+        )[0]
+
+        self.assertEqual(result.shape, self.images.shape)
+        self.assertTrue(torch.allclose(result, self.images, atol=(1.5 / 255.0)))
+
+    def test_process_single_gpu_invokes_rounding_and_batch_processing(self):
+        self.node.process_single_gpu(
+            upscaled_image=self.images,
+            core_args=self.core_args,
+            tile_width=16,
+            tile_height=24,
+            padding=4,
+            mask_blur=2,
+            force_uniform_tiles=True,
+        )
+
+        self.assertEqual(self.node.round_calls[:2], [16, 24])
+        self.assertEqual(self.node.process_batch_calls, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_static_mode.py b/tests/unit/upscale/test_static.py
similarity index 80%
rename from tests/test_static_mode.py
rename to tests/unit/upscale/test_static.py
index a824493..13d274a 100644
--- a/tests/test_static_mode.py
+++ b/tests/unit/upscale/test_static.py
@@ -9,7 +9,7 @@ import torch
 
 
 def _load_static_mode_module():
-    module_path = Path(__file__).resolve().parents[1] / "upscale" / "modes" / "static.py"
+    module_path = Path(__file__).resolve().parents[3] / "upscale" / "modes" / "static.py"
     package_name = "dist_static_mode_testpkg"
 
     for mod_name in list(sys.modules):
@@ -65,8 +65,34 @@ def _load_static_mode_module():
         arr = np.array(image).astype(np.float32) / 255.0
         return torch.from_numpy(arr).unsqueeze(0)
 
+    def _blend_processed_batch_item(
+        result_images,
+        processed_batch,
+        batch_index,
+        blend_fn,
+        x1,
+        y1,
+        ew,
+        eh,
+        tile_mask,
+        padding,
+    ):
+        tile_pil = _tensor_to_pil(processed_batch, batch_index)
+        if tile_pil.size != (ew, eh):
+            tile_pil = tile_pil.resize((ew, eh), PILImage.LANCZOS)
+        result_images[batch_index] = blend_fn(
+            result_images[batch_index],
+            tile_pil,
+            x1,
+            y1,
+            (ew, eh),
+            tile_mask,
+            padding,
+        )
+
     image_module.tensor_to_pil = _tensor_to_pil
     image_module.pil_to_tensor = _pil_to_tensor
+    image_module.blend_processed_batch_item = _blend_processed_batch_item
     sys.modules[f"{package_name}.utils.image"] = image_module
 
     async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
@@ -102,6 +128,10 @@ def _load_static_mode_module():
         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
     job_store_module._mark_task_completed = _noop
     job_store_module._cleanup_job = _noop
     job_store_module._drain_results_queue = _noop
@@ -116,6 +146,26 @@ def _load_static_mode_module():
     job_models_module.TileJobState = _TileJobState
     sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
 
+    processing_args_path = Path(__file__).resolve().parents[3] / "upscale" / "processing_args.py"
+    processing_args_spec = importlib.util.spec_from_file_location(
+        f"{package_name}.upscale.processing_args",
+        processing_args_path,
+    )
+    processing_args_module = importlib.util.module_from_spec(processing_args_spec)
+    assert processing_args_spec is not None and processing_args_spec.loader is not None
+    sys.modules[processing_args_spec.name] = processing_args_module
+    processing_args_spec.loader.exec_module(processing_args_module)
+
+    tile_processing_path = Path(__file__).resolve().parents[3] / "upscale" / "tile_processing.py"
+    tile_processing_spec = importlib.util.spec_from_file_location(
+        f"{package_name}.upscale.tile_processing",
+        tile_processing_path,
+    )
+    tile_processing_module = importlib.util.module_from_spec(tile_processing_spec)
+    assert tile_processing_spec is not None and tile_processing_spec.loader is not None
+    sys.modules[tile_processing_spec.name] = tile_processing_module
+    tile_processing_spec.loader.exec_module(tile_processing_module)
+
     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
diff --git a/tests/unit/upscale/test_tile_ops.py b/tests/unit/upscale/test_tile_ops.py
new file mode 100644
index 0000000..ec22ce7
--- /dev/null
+++ b/tests/unit/upscale/test_tile_ops.py
@@ -0,0 +1,179 @@
+import copy
+import importlib.util
+import sys
+import types
+import unittest
+from contextlib import nullcontext
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    def pil_to_tensor(image):
+        array = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(array).unsqueeze(0)
+
+    image_module.tensor_to_pil = tensor_to_pil
+    image_module.pil_to_tensor = pil_to_tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    usdu_utils_module = types.ModuleType(f"{package_name}.utils.usdu_utils")
+
+    def crop_cond(cond, *_args, **_kwargs):
+        return cond
+
+    def get_crop_region(mask, padding):
+        bbox = mask.getbbox()
+        if bbox is None:
+            return (0, 0, mask.size[0], mask.size[1])
+        x1, y1, x2, y2 = bbox
+        return (
+            max(0, x1 - int(padding)),
+            max(0, y1 - int(padding)),
+            min(mask.size[0], x2 + int(padding)),
+            min(mask.size[1], y2 + int(padding)),
+        )
+
+    def expand_crop(region, width, height, target_w, target_h):
+        x1, y1, x2, y2 = region
+        crop_w = x2 - x1
+        crop_h = y2 - y1
+        if crop_w >= target_w and crop_h >= target_h:
+            return (region, (crop_w, crop_h))
+        new_x2 = min(width, x1 + max(crop_w, target_w))
+        new_y2 = min(height, y1 + max(crop_h, target_h))
+        return ((x1, y1, new_x2, new_y2), (new_x2 - x1, new_y2 - y1))
+
+    usdu_utils_module.crop_cond = crop_cond
+    usdu_utils_module.get_crop_region = get_crop_region
+    usdu_utils_module.expand_crop = expand_crop
+    sys.modules[f"{package_name}.utils.usdu_utils"] = usdu_utils_module
+
+    crop_model_patch_module = types.ModuleType(f"{package_name}.utils.crop_model_patch")
+    crop_model_patch_module.crop_model_cond = lambda model, *_args, **_kwargs: nullcontext(model)
+    sys.modules[f"{package_name}.utils.crop_model_patch"] = crop_model_patch_module
+
+    conditioning_module = types.ModuleType(f"{package_name}.upscale.conditioning")
+    conditioning_module.clone_conditioning = lambda cond, clone_hints=False: copy.deepcopy(cond)
+    sys.modules[f"{package_name}.upscale.conditioning"] = conditioning_module
+
+    comfy_module = types.ModuleType("comfy")
+    comfy_samplers = types.ModuleType("comfy.samplers")
+    comfy_model_management = types.ModuleType("comfy.model_management")
+    comfy_module.samplers = comfy_samplers
+    comfy_module.model_management = comfy_model_management
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.samplers"] = comfy_samplers
+    sys.modules["comfy.model_management"] = comfy_model_management
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[3] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_tile_ops_module():
+    package_name = "dist_tile_ops_testpkg"
+    _bootstrap_package(package_name)
+    return _load_module(package_name, "upscale/tile_ops.py", "upscale.tile_ops")
+
+
+tile_ops_module = _load_tile_ops_module()
+
+
+class _DummyTileOps(tile_ops_module.TileOpsMixin):
+    pass
+
+
+class _FakeControl:
+    def __init__(self, cond_hint_original, previous_controlnet=None):
+        self.cond_hint_original = cond_hint_original
+        self.previous_controlnet = previous_controlnet
+
+
+class TileOpsTests(unittest.TestCase):
+    def setUp(self):
+        self.node = _DummyTileOps()
+
+    def test_round_and_calculate_tiles(self):
+        self.assertEqual(self.node.round_to_multiple(13, 8), 16)
+        tiles = self.node.calculate_tiles(10, 9, 4, 5, force_uniform_tiles=True)
+        self.assertEqual(len(tiles), 6)
+        self.assertEqual(tiles[0], (0, 0))
+        self.assertEqual(tiles[-1], (8, 5))
+
+    def test_create_tile_mask_and_blend_tile(self):
+        base = Image.new("RGB", (8, 8), (0, 0, 0))
+        tile = Image.new("RGB", (4, 4), (255, 255, 255))
+        mask = self.node.create_tile_mask(8, 8, 2, 2, 4, 4, mask_blur=0)
+
+        self.assertEqual(mask.getpixel((0, 0)), 0)
+        self.assertEqual(mask.getpixel((3, 3)), 255)
+
+        blended = self.node.blend_tile(
+            base_image=base,
+            tile_image=tile,
+            x=2,
+            y=2,
+            extracted_size=(4, 4),
+            mask=mask,
+            padding=0,
+        )
+        self.assertEqual(blended.size, (8, 8))
+        self.assertEqual(blended.getpixel((3, 3)), (255, 255, 255))
+
+    def test_slice_conditioning_selects_batch_index(self):
+        control = _FakeControl(cond_hint_original=torch.ones((2, 1, 1)))
+        positive = [[torch.arange(6, dtype=torch.float32).reshape(2, 3), {"control": control, "mask": torch.ones((2, 1, 1))}]]
+        negative = [[torch.arange(6, dtype=torch.float32).reshape(2, 3), {"mask": torch.zeros((2, 1, 1))}]]
+
+        pos_sliced, neg_sliced = self.node._slice_conditioning(positive, negative, batch_idx=1)
+
+        self.assertEqual(tuple(pos_sliced[0][0].shape), (1, 3))
+        self.assertEqual(tuple(neg_sliced[0][0].shape), (1, 3))
+        self.assertEqual(tuple(pos_sliced[0][1]["mask"].shape), (1, 1, 1))
+        self.assertEqual(tuple(neg_sliced[0][1]["mask"].shape), (1, 1, 1))
+        self.assertEqual(tuple(pos_sliced[0][1]["control"].cond_hint_original.shape), (1, 1, 1))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/upscale/test_tile_processing.py b/tests/unit/upscale/test_tile_processing.py
new file mode 100644
index 0000000..b2462ee
--- /dev/null
+++ b/tests/unit/upscale/test_tile_processing.py
@@ -0,0 +1,115 @@
+import unittest
+
+from upscale.processing_args import UpscaleCoreArgs
+from upscale.tile_processing import TileBatchArgs, extract_and_process_tile_batch
+
+
+class _DummyTileProcessor:
+    def __init__(self):
+        self.extract_calls = []
+        self.process_calls = []
+
+    def extract_batch_tile_with_padding(
+        self,
+        upscaled_image,
+        tx,
+        ty,
+        tile_width,
+        tile_height,
+        padding,
+        force_uniform_tiles,
+    ):
+        self.extract_calls.append(
+            {
+                "upscaled_image": upscaled_image,
+                "tx": tx,
+                "ty": ty,
+                "tile_width": tile_width,
+                "tile_height": tile_height,
+                "padding": padding,
+                "force_uniform_tiles": force_uniform_tiles,
+            }
+        )
+        return "tile-batch", 10, 20, 30, 40
+
+    def process_tiles_batch(
+        self,
+        tile_batch,
+        model,
+        positive,
+        negative,
+        vae,
+        seed,
+        steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        denoise,
+        tiled_decode,
+        region,
+        canvas_shape,
+    ):
+        self.process_calls.append(
+            {
+                "tile_batch": tile_batch,
+                "model": model,
+                "positive": positive,
+                "negative": negative,
+                "vae": vae,
+                "seed": seed,
+                "steps": steps,
+                "cfg": cfg,
+                "sampler_name": sampler_name,
+                "scheduler": scheduler,
+                "denoise": denoise,
+                "tiled_decode": tiled_decode,
+                "region": region,
+                "canvas_shape": canvas_shape,
+            }
+        )
+        return "processed-batch"
+
+
+class TileProcessingTests(unittest.TestCase):
+    def test_extract_and_process_tile_batch_forwards_args_and_computes_region(self):
+        node = _DummyTileProcessor()
+        tile_args = TileBatchArgs(
+            core=UpscaleCoreArgs(
+                model="model",
+                positive="positive",
+                negative="negative",
+                vae="vae",
+                seed=42,
+                steps=20,
+                cfg=7.5,
+                sampler_name="euler",
+                scheduler="normal",
+                denoise=0.4,
+                tiled_decode=False,
+            ),
+            tile_width=256,
+            tile_height=192,
+            padding=32,
+            force_uniform_tiles=True,
+            width=1024,
+            height=768,
+        )
+        result = extract_and_process_tile_batch(
+            node=node,
+            upscaled_image="upscaled-image",
+            tx=3,
+            ty=4,
+            args=tile_args,
+        )
+
+        self.assertEqual(result, ("processed-batch", 10, 20, 30, 40))
+        self.assertEqual(len(node.extract_calls), 1)
+        self.assertEqual(len(node.process_calls), 1)
+        self.assertEqual(node.extract_calls[0]["tx"], 3)
+        self.assertEqual(node.extract_calls[0]["ty"], 4)
+        self.assertEqual(node.process_calls[0]["region"], (10, 20, 40, 60))
+        self.assertEqual(node.process_calls[0]["canvas_shape"], (1024, 768))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/upscale/test_usdu_utils.py b/tests/unit/upscale/test_usdu_utils.py
new file mode 100644
index 0000000..90b1e2b
--- /dev/null
+++ b/tests/unit/upscale/test_usdu_utils.py
@@ -0,0 +1,128 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+from PIL import Image
+
+
+def _ensure_torch_stubs():
+    if "torch" not in sys.modules:
+        torch_module = types.ModuleType("torch")
+
+        class _Tensor:
+            pass
+
+        torch_module.Tensor = _Tensor
+        torch_module.float32 = object()
+        torch_module.zeros = lambda *_args, **_kwargs: None
+        torch_module.from_numpy = lambda array: array
+        torch_module.cat = lambda *_args, **_kwargs: None
+
+        nn_module = types.ModuleType("torch.nn")
+        functional_module = types.ModuleType("torch.nn.functional")
+        functional_module.interpolate = lambda tensor, **_kwargs: tensor
+        nn_module.functional = functional_module
+
+        torch_module.nn = nn_module
+
+        sys.modules["torch"] = torch_module
+        sys.modules["torch.nn"] = nn_module
+        sys.modules["torch.nn.functional"] = functional_module
+
+    if "torchvision" not in sys.modules:
+        torchvision_module = types.ModuleType("torchvision")
+        transforms_module = types.ModuleType("torchvision.transforms")
+
+        class _GaussianBlur:
+            def __init__(self, *_args, **_kwargs):
+                pass
+
+            def __call__(self, value):
+                return value
+
+        transforms_module.GaussianBlur = _GaussianBlur
+        torchvision_module.transforms = transforms_module
+
+        sys.modules["torchvision"] = torchvision_module
+        sys.modules["torchvision.transforms"] = transforms_module
+
+
+def _load_usdu_utils_module():
+    _ensure_torch_stubs()
+    module_path = Path(__file__).resolve().parents[3] / "utils" / "usdu_utils.py"
+    spec = importlib.util.spec_from_file_location("dist_test_usdu_utils", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+usdu_utils = _load_usdu_utils_module()
+
+
+class UsduUtilsGeometryTests(unittest.TestCase):
+    def test_fix_crop_region_decrements_upper_bounds_when_interior(self):
+        self.assertEqual(
+            usdu_utils.fix_crop_region((2, 3, 9, 9), (10, 10)),
+            (2, 3, 8, 8),
+        )
+        self.assertEqual(
+            usdu_utils.fix_crop_region((1, 1, 10, 10), (10, 10)),
+            (1, 1, 10, 10),
+        )
+
+    def test_get_crop_region_with_padding_and_empty_mask(self):
+        mask = Image.new("L", (10, 10), 0)
+        for x in range(2, 5):
+            for y in range(3, 7):
+                mask.putpixel((x, y), 255)
+
+        region = usdu_utils.get_crop_region(mask, pad=1)
+        self.assertEqual(region, (1, 2, 5, 7))
+
+        empty_region = usdu_utils.get_crop_region(Image.new("L", (8, 6), 0), pad=2)
+        self.assertEqual(empty_region, (6, 4, 1, 1))
+
+    def test_expand_crop_and_resize_region(self):
+        expanded_region, expanded_size = usdu_utils.expand_crop(
+            region=(2, 2, 6, 6),
+            width=12,
+            height=12,
+            target_width=8,
+            target_height=10,
+        )
+        self.assertEqual(expanded_region, (0, 0, 8, 10))
+        self.assertEqual(expanded_size, (8, 10))
+
+        resized = usdu_utils.resize_region(
+            region=(2, 2, 6, 10),
+            init_size=(8, 16),
+            resize_size=(16, 32),
+        )
+        self.assertEqual(resized, (4, 4, 12, 20))
+
+    def test_region_intersection(self):
+        self.assertEqual(
+            usdu_utils.region_intersection((0, 0, 5, 5), (3, 2, 8, 7)),
+            (3, 2, 5, 5),
+        )
+        self.assertIsNone(usdu_utils.region_intersection((0, 0, 1, 1), (2, 2, 4, 4)))
+
+    def test_pad_image2_preserves_expected_output_size(self):
+        image = Image.new("RGB", (4, 3), "black")
+        padded = usdu_utils.pad_image2(
+            image,
+            left_pad=2,
+            right_pad=1,
+            top_pad=3,
+            bottom_pad=2,
+            fill=False,
+            blur=False,
+        )
+        self.assertEqual(padded.size, (7, 8))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/upscale/conditioning.py b/upscale/conditioning.py
index 3e776bc..0dcbf09 100644
--- a/upscale/conditioning.py
+++ b/upscale/conditioning.py
@@ -1,7 +1,8 @@
 import copy
+from typing import Any
 
 
-def clone_control_chain(control, clone_hint=True):
+def clone_control_chain(control: Any, clone_hint: bool = True) -> Any:
     """Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
     if control is None:
         return None
@@ -14,7 +15,10 @@ def clone_control_chain(control, clone_hint=True):
     return new_control
 
 
-def clone_conditioning(cond_list, clone_hints=True):
+def clone_conditioning(
+    cond_list: list[tuple[Any, dict[str, Any]]] | list[list[Any]],
+    clone_hints: bool = True,
+) -> list[list[Any]]:
     """Clone conditioning without duplicating ControlNet models."""
     new_cond = []
     for emb, cond_dict in cond_list:
diff --git a/upscale/job_state.py b/upscale/job_state.py
index 51ed1f9..451ccb8 100644
--- a/upscale/job_state.py
+++ b/upscale/job_state.py
@@ -2,7 +2,7 @@ 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_timeout import check_and_requeue_timed_out_workers
 from .job_models import ImageJobState, TileJobState
 
 
@@ -132,4 +132,4 @@ class JobStateMixin:
 
     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)
+        return await check_and_requeue_timed_out_workers(multi_job_id, batch_size)
diff --git a/upscale/job_store.py b/upscale/job_store.py
index c027a33..87d2ff9 100644
--- a/upscale/job_store.py
+++ b/upscale/job_store.py
@@ -1,7 +1,8 @@
 import asyncio
 import os
 import time
-from typing import List, Optional
+from dataclasses import dataclass
+from typing import Any
 
 import server
 
@@ -12,7 +13,17 @@ from .job_models import BaseJobState, ImageJobState, TileJobState
 MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
 
 
-def ensure_tile_jobs_initialized():
+@dataclass(frozen=True)
+class JobQueueInitConfig:
+    mode: str
+    batch_size: int = 0
+    num_tiles_per_image: int = 0
+    all_indices: list[int] | None = None
+    enabled_workers: list[str] | None = None
+    batched_static: bool = False
+
+
+def ensure_tile_jobs_initialized() -> Any:
     """Ensure tile job storage is initialized on the server instance."""
     prompt_server = server.PromptServer.instance
     if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
@@ -32,14 +43,9 @@ def ensure_tile_jobs_initialized():
 
 
 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,
-):
+    multi_job_id: str,
+    config: JobQueueInitConfig,
+) -> None:
     """Unified initialization for job queues in static and dynamic modes."""
     prompt_server = ensure_tile_jobs_initialized()
     async with prompt_server.distributed_tile_jobs_lock:
@@ -47,33 +53,34 @@ async def _init_job_queue(
             debug_log(f"Queue already exists for {multi_job_id}")
             return
 
-        if mode == 'dynamic':
+        if config.mode == 'dynamic':
             job_data = ImageJobState(multi_job_id=multi_job_id)
-        elif mode == 'static':
+        elif config.mode == 'static':
             job_data = TileJobState(multi_job_id=multi_job_id)
         else:
-            raise ValueError(f"Unknown mode: {mode}")
+            raise ValueError(f"Unknown mode: {config.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 []}
+        job_data.worker_status = {w: time.time() for w in config.enabled_workers or []}
+        job_data.assigned_to_workers = {w: [] for w in config.enabled_workers or []}
 
-        if mode == 'dynamic':
-            job_data.batch_size = int(batch_size or 0)
+        if config.mode == 'dynamic':
+            job_data.batch_size = int(config.batch_size or 0)
             pending_queue = job_data.pending_images
-            for i in (all_indices or range(int(batch_size or 0))):
+            indices = config.all_indices or list(range(int(config.batch_size or 0)))
+            for i in indices:
                 await pending_queue.put(i)
-            debug_log(f"Initialized image queue with {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)
+            debug_log(f"Initialized image queue with {config.batch_size} pending items")
+        elif config.mode == 'static':
+            job_data.num_tiles_per_image = int(config.num_tiles_per_image or 0)
+            job_data.batch_size = int(config.batch_size or 0)
+            job_data.batched_static = bool(config.batched_static)
             # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
             pending_queue = job_data.pending_tasks
-            if batched_static and num_tiles_per_image is not None:
-                for i in range(num_tiles_per_image):
+            if config.batched_static and config.num_tiles_per_image > 0:
+                for i in range(config.num_tiles_per_image):
                     await pending_queue.put(i)
             else:
-                total_tiles = int(batch_size or 0) * int(num_tiles_per_image or 0)
+                total_tiles = int(config.batch_size or 0) * int(config.num_tiles_per_image or 0)
                 for i in range(total_tiles):
                     await pending_queue.put(i)
 
@@ -83,16 +90,18 @@ async def _init_job_queue(
 async def init_dynamic_job(
     multi_job_id: str,
     batch_size: int,
-    enabled_workers: List[str],
-    all_indices: Optional[List[int]] = None,
-):
+    enabled_workers: list[str],
+    all_indices: list[int] | None = None,
+) -> None:
     """Initialize queue for dynamic mode (per-image), with collector fields."""
     await _init_job_queue(
         multi_job_id,
-        'dynamic',
-        batch_size=batch_size,
-        all_indices=all_indices or list(range(batch_size)),
-        enabled_workers=enabled_workers,
+        JobQueueInitConfig(
+            mode='dynamic',
+            batch_size=batch_size,
+            all_indices=all_indices or list(range(batch_size)),
+            enabled_workers=enabled_workers,
+        ),
     )
     debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
 
@@ -101,20 +110,22 @@ async def init_static_job_batched(
     multi_job_id: str,
     batch_size: int,
     num_tiles_per_image: int,
-    enabled_workers: List[str],
-):
+    enabled_workers: list[str],
+) -> None:
     """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,
+        JobQueueInitConfig(
+            mode='static',
+            batch_size=batch_size,
+            num_tiles_per_image=num_tiles_per_image,
+            enabled_workers=enabled_workers,
+            batched_static=True,
+        ),
     )
 
 
-async def _drain_results_queue(multi_job_id):
+async def _drain_results_queue(multi_job_id: str) -> int:
     """Drain pending results from queue and update completed_tasks. Returns count drained."""
     prompt_server = ensure_tile_jobs_initialized()
     async with prompt_server.distributed_tile_jobs_lock:
@@ -152,7 +163,7 @@ async def _drain_results_queue(multi_job_id):
         return collected
 
 
-async def _get_completed_count(multi_job_id):
+async def _get_completed_count(multi_job_id: str) -> int:
     """Get count of completed tasks."""
     prompt_server = ensure_tile_jobs_initialized()
     async with prompt_server.distributed_tile_jobs_lock:
@@ -162,7 +173,7 @@ async def _get_completed_count(multi_job_id):
         return 0
 
 
-async def _mark_task_completed(multi_job_id, task_id, result):
+async def _mark_task_completed(multi_job_id: str, task_id: int, result: Any) -> None:
     """Mark a task as completed."""
     prompt_server = ensure_tile_jobs_initialized()
     async with prompt_server.distributed_tile_jobs_lock:
@@ -171,10 +182,38 @@ async def _mark_task_completed(multi_job_id, task_id, result):
             job_data.completed_tasks[task_id] = result
 
 
-async def _cleanup_job(multi_job_id):
+async def _cleanup_job(multi_job_id: str) -> None:
     """Cleanup the job data."""
     prompt_server = ensure_tile_jobs_initialized()
     async with prompt_server.distributed_tile_jobs_lock:
         if multi_job_id in prompt_server.distributed_pending_tile_jobs:
             del prompt_server.distributed_pending_tile_jobs[multi_job_id]
             debug_log(f"Cleaned up job {multi_job_id}")
+
+
+async def init_job_queue(
+    multi_job_id: str,
+    config: JobQueueInitConfig,
+) -> None:
+    """Public wrapper for unified job queue initialization."""
+    await _init_job_queue(multi_job_id, config)
+
+
+async def drain_results_queue(multi_job_id: str) -> int:
+    """Public wrapper for draining queued worker results."""
+    return await _drain_results_queue(multi_job_id)
+
+
+async def get_completed_count(multi_job_id: str) -> int:
+    """Public wrapper for counting completed tasks in a job."""
+    return await _get_completed_count(multi_job_id)
+
+
+async def mark_task_completed(multi_job_id: str, task_id: int, result: Any) -> None:
+    """Public wrapper for marking a task complete."""
+    await _mark_task_completed(multi_job_id, task_id, result)
+
+
+async def cleanup_job(multi_job_id: str) -> None:
+    """Public wrapper for deleting finished job state."""
+    await _cleanup_job(multi_job_id)
diff --git a/upscale/job_timeout.py b/upscale/job_timeout.py
index 37baa9c..0bc4d5d 100644
--- a/upscale/job_timeout.py
+++ b/upscale/job_timeout.py
@@ -148,3 +148,8 @@ async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
                 job_data.assigned_to_workers[worker] = []
 
         return requeued_count
+
+
+async def check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
+    """Public wrapper for worker-timeout requeue checks."""
+    return await _check_and_requeue_timed_out_workers(multi_job_id, total_tasks)
diff --git a/upscale/modes/dynamic.py b/upscale/modes/dynamic.py
index f22292f..dc2998a 100644
--- a/upscale/modes/dynamic.py
+++ b/upscale/modes/dynamic.py
@@ -1,12 +1,15 @@
+from __future__ import annotations
+
 import asyncio, torch
+from typing import Any
 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
+from ..processing_args import UpscaleCoreArgs
 
 
 class DynamicModeMixin:
@@ -19,10 +22,113 @@ class DynamicModeMixin:
     - 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):
+    def _handle_master_dynamic_idle_state(
+        self,
+        multi_job_id,
+        batch_size,
+        consecutive_retries,
+        max_consecutive_retries,
+    ):
+        """Handle queue-empty state while master waits for worker 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)
+
+        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)
+
+        if requeued_count > 0:
+            log(f"Requeued {requeued_count} images from timed out workers")
+            return False, True, 0
+
+        completed_now = run_async_in_server_loop(
+            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:
+            return True, False, consecutive_retries
+
+        run_async_in_server_loop(self._async_yield(), timeout=0.1)
+
+        pending_count = run_async_in_server_loop(
+            self._get_pending_count(multi_job_id),
+            timeout=1.0,
+        )
+        if pending_count > 0:
+            return False, True, 0
+
+        consecutive_retries += 1
+        if consecutive_retries >= max_consecutive_retries:
+            log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.")
+            return True, False, consecutive_retries
+
+        if drained_count > 0:
+            debug_log(f"Master idle drain picked up {drained_count} images while waiting for workers")
+        debug_log("Waiting for workers")
+        run_async_in_server_loop(asyncio.sleep(2), timeout=3.0)
+        return False, True, consecutive_retries
+
+    def _finalize_master_dynamic_results(
+        self,
+        multi_job_id,
+        batch_size,
+        num_workers,
+        processed_count,
+        result_images,
+        upscaled_image,
+    ):
+        """Collect remaining worker outputs and convert final images back to tensor."""
+        all_completed = run_async_in_server_loop(
+            self._get_all_completed_images(multi_job_id),
+            timeout=5.0,
+        )
+        remaining_to_collect = batch_size - len(all_completed)
+        if remaining_to_collect > 0:
+            debug_log(f"Waiting for {remaining_to_collect} more images from workers")
+            collection_timeout = float(get_worker_timeout_seconds())
+            collected_images = run_async_in_server_loop(
+                self._async_collect_dynamic_images(
+                    multi_job_id,
+                    remaining_to_collect,
+                    num_workers,
+                    batch_size,
+                    processed_count,
+                ),
+                timeout=collection_timeout,
+            )
+            all_completed.update(collected_images)
+
+        for idx, processed_img in all_completed.items():
+            if idx < batch_size:
+                result_images[idx] = processed_img
+
+        result_tensor = (
+            torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
+            if batch_size > 1
+            else pil_to_tensor(result_images[0])
+        )
+        if upscaled_image.is_cuda:
+            result_tensor = result_tensor.cuda()
+        return result_tensor
+
+    def process_master_dynamic(
+        self,
+        upscaled_image: torch.Tensor,
+        core_args: UpscaleCoreArgs,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        multi_job_id: str,
+        enabled_workers: list[str],
+    ) -> tuple[torch.Tensor]:
         """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
@@ -74,10 +180,14 @@ class DynamicModeMixin:
                 # Process locally
                 single_tensor = upscaled_image[image_idx:image_idx+1]
                 local_image = result_images[image_idx]
-                image_seed = seed
+                image_seed = core_args.seed
                 
                 # Pre-slice conditioning once per image (not per tile)
-                positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx)
+                positive_sliced, negative_sliced = self._slice_conditioning(
+                    core_args.positive,
+                    core_args.negative,
+                    image_idx,
+                )
                 
                 for tile_idx, pos in enumerate(all_tiles):
                     source_tensor = pil_to_tensor(local_image)
@@ -85,10 +195,21 @@ class DynamicModeMixin:
                         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,
+                        core_args.model,
+                        positive_sliced,
+                        negative_sliced,
+                        core_args.vae,
+                        image_seed,
+                        core_args.steps,
+                        core_args.cfg,
+                        core_args.sampler_name,
+                        core_args.scheduler,
+                        core_args.denoise,
+                        tile_width,
+                        tile_height,
                         padding, mask_blur, width, height, force_uniform_tiles,
-                        tiled_decode, batch_idx=image_idx
+                        core_args.tiled_decode,
+                        batch_idx=image_idx,
                     )
                     
                     # Yield after each tile to minimize worker downtime
@@ -122,99 +243,46 @@ class DynamicModeMixin:
                 # 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
+                should_break, should_continue, consecutive_retries = self._handle_master_dynamic_idle_state(
+                    multi_job_id=multi_job_id,
+                    batch_size=batch_size,
+                    consecutive_retries=consecutive_retries,
+                    max_consecutive_retries=max_consecutive_retries,
                 )
-                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:
+                if should_break:
                     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
+                if should_continue:
                     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
+        result_tensor = self._finalize_master_dynamic_results(
+            multi_job_id=multi_job_id,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            processed_count=processed_count,
+            result_images=result_images,
+            upscaled_image=upscaled_image,
         )
         
-        # 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):
+    def process_worker_dynamic(
+        self,
+        upscaled_image: torch.Tensor,
+        core_args: UpscaleCoreArgs,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+        enabled_worker_ids: list[str] | str,
+        dynamic_threshold: int,
+    ) -> tuple[torch.Tensor]:
         """Worker processing in dynamic mode - processes whole images."""
         # Round tile dimensions
         tile_width = self.round_to_multiple(tile_width)
@@ -259,10 +327,14 @@ class DynamicModeMixin:
             local_image = tensor_to_pil(single_tensor, 0).copy()
 
             # Process all tiles for this image
-            image_seed = seed
+            image_seed = core_args.seed
 
             # Pre-slice conditioning once per image (not per tile)
-            positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx)
+            positive_sliced, negative_sliced = self._slice_conditioning(
+                core_args.positive,
+                core_args.negative,
+                image_idx,
+            )
 
             for tile_idx, pos in enumerate(all_tiles):
                 source_tensor = pil_to_tensor(local_image)
@@ -270,10 +342,21 @@ class DynamicModeMixin:
                     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,
+                    core_args.model,
+                    positive_sliced,
+                    negative_sliced,
+                    core_args.vae,
+                    image_seed,
+                    core_args.steps,
+                    core_args.cfg,
+                    core_args.sampler_name,
+                    core_args.scheduler,
+                    core_args.denoise,
+                    tile_width,
+                    tile_height,
                     padding, mask_blur, width, height, force_uniform_tiles,
-                    tiled_decode, batch_idx=image_idx
+                    core_args.tiled_decode,
+                    batch_idx=image_idx,
                 )
                 run_async_in_server_loop(
                     self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
diff --git a/upscale/modes/single_gpu.py b/upscale/modes/single_gpu.py
index 5b86595..b6b10b6 100644
--- a/upscale/modes/single_gpu.py
+++ b/upscale/modes/single_gpu.py
@@ -1,29 +1,59 @@
+from __future__ import annotations
+
 import math, torch
-from PIL import Image
-from ...utils.logging import debug_log, log
-from ...utils.image import tensor_to_pil, pil_to_tensor
+from typing import Any, Protocol, cast
+from ...utils.logging import log
+from ...utils.image import blend_processed_batch_item, pil_to_tensor, tensor_to_pil
+from ..processing_args import UpscaleCoreArgs
+from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch
+
+
+class _SingleGpuOps(Protocol):
+    def round_to_multiple(self, value: int) -> int: ...
+    def calculate_tiles(self, width: int, height: int, tile_width: int, tile_height: int, force_uniform_tiles: bool): ...
+    def create_tile_mask(self, width: int, height: int, tx: int, ty: int, tile_width: int, tile_height: int, mask_blur: int): ...
+    def blend_tile(self, base_image: Any, tile_pil: Any, x1: int, y1: int, size: tuple[int, int], tile_mask: Any, padding: int): ...
+    def extract_batch_tile_with_padding(self, source_batch: Any, tx: int, ty: int, tile_width: int, tile_height: int, padding: int, force_uniform_tiles: bool): ...
+    def process_tiles_batch(self, *args: Any, **kwargs: Any): ...
 
 
 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):
+    def process_single_gpu(
+        self,
+        upscaled_image: torch.Tensor,
+        core_args: UpscaleCoreArgs,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        force_uniform_tiles: bool,
+    ) -> tuple[torch.Tensor]:
         """Process all tiles on a single GPU (no distribution), batching per tile like USDU."""
+        ops = cast(_SingleGpuOps, self)
         # Round tile dimensions
-        tile_width = self.round_to_multiple(tile_width)
-        tile_height = self.round_to_multiple(tile_height)
+        tile_width = ops.round_to_multiple(tile_width)
+        tile_height = ops.round_to_multiple(tile_height)
 
         # Get image dimensions and batch size
         batch_size, height, width, _ = upscaled_image.shape
 
         # Calculate all tiles
-        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+        all_tiles = ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
 
         rows = math.ceil(height / tile_height)
         cols = math.ceil(width / tile_width)
         log(
             f"USDU Dist: Single GPU | Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({len(all_tiles)} tiles/image) | Batch {batch_size}"
         )
+        tile_batch_args = TileBatchArgs(
+            core=core_args,
+            tile_width=tile_width,
+            tile_height=tile_height,
+            padding=padding,
+            force_uniform_tiles=force_uniform_tiles,
+            width=width,
+            height=height,
+        )
 
         # Prepare result images list
         result_images = []
@@ -34,7 +64,7 @@ class SingleGpuModeMixin:
         # 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))
+            tile_masks.append(ops.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
 
         # Process tiles batched across images
         for tile_idx, (tx, ty) in enumerate(all_tiles):
@@ -43,25 +73,29 @@ class SingleGpuModeMixin:
             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
+            processed_batch, x1, y1, ew, eh = extract_and_process_tile_batch(
+                node=ops,
+                upscaled_image=source_batch,
+                tx=tx,
+                ty=ty,
+                args=tile_batch_args,
             )
 
-            # 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)
+                blend_processed_batch_item(
+                    result_images,
+                    processed_batch,
+                    b,
+                    ops.blend_tile,
+                    x1,
+                    y1,
+                    ew,
+                    eh,
+                    tile_mask,
+                    padding,
+                )
 
         # Convert back to tensor
         result_tensors = [pil_to_tensor(img) for img in result_images]
diff --git a/upscale/modes/static.py b/upscale/modes/static.py
index abd4d67..354d5d1 100644
--- a/upscale/modes/static.py
+++ b/upscale/modes/static.py
@@ -2,7 +2,7 @@ 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.image import blend_processed_batch_item, pil_to_tensor, tensor_to_pil
 from ...utils.async_helpers import run_async_in_server_loop
 from ...utils.config import get_worker_timeout_seconds
 from ...utils.constants import (
@@ -15,9 +15,11 @@ from ...utils.constants import (
 )
 from ..job_store import (
     ensure_tile_jobs_initialized, init_static_job_batched,
-    _mark_task_completed, _cleanup_job, _drain_results_queue, _get_completed_count,
+    cleanup_job, drain_results_queue, get_completed_count, mark_task_completed,
 )
 from ..job_models import TileJobState
+from ..processing_args import UpscaleCoreArgs
+from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch
 
 
 class StaticModeMixin:
@@ -71,14 +73,33 @@ class StaticModeMixin:
     ):
         """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
+        tile_batch_args = TileBatchArgs(
+            core=UpscaleCoreArgs(
+                model=model,
+                positive=positive,
+                negative=negative,
+                vae=vae,
+                seed=seed,
+                steps=steps,
+                cfg=cfg,
+                sampler_name=sampler_name,
+                scheduler=scheduler,
+                denoise=denoise,
+                tiled_decode=tiled_decode,
+            ),
+            tile_width=tile_width,
+            tile_height=tile_height,
+            padding=padding,
+            force_uniform_tiles=force_uniform_tiles,
+            width=width,
+            height=height,
         )
-        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)
+        processed_batch, x1, y1, ew, eh = extract_and_process_tile_batch(
+            node=self,
+            upscaled_image=upscaled_image,
+            tx=tx,
+            ty=ty,
+            args=tile_batch_args,
         )
         return processed_batch, x1, y1, ew, eh
 
@@ -177,13 +198,21 @@ class StaticModeMixin:
         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)
+            blend_processed_batch_item(
+                result_images,
+                processed_batch,
+                b,
+                self.blend_tile,
+                x1,
+                y1,
+                ew,
+                eh,
+                tile_mask,
+                padding,
+            )
             global_idx = b * num_tiles_per_image + tile_id
             run_async_in_server_loop(
-                _mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}),
+                mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}),
                 timeout=5.0
             )
         return processed_items
@@ -328,7 +357,7 @@ class StaticModeMixin:
                 raise comfy.model_management.InterruptProcessingException()
             
             # Drain any pending results
-            collected_count = await _drain_results_queue(multi_job_id)
+            collected_count = await drain_results_queue(multi_job_id)
             
             # Check and requeue timed-out workers periodically
             current_time = time.time()
@@ -339,7 +368,7 @@ class StaticModeMixin:
                 last_heartbeat_check = current_time
             
             # Get current completion count
-            completed_count = await _get_completed_count(multi_job_id)
+            completed_count = await get_completed_count(multi_job_id)
             
             # Progress logging
             if current_time - last_progress_log >= progress_interval:
@@ -368,157 +397,198 @@ class StaticModeMixin:
         # 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]
+    def _init_master_static_state(
+        self,
+        upscaled_image,
+        multi_job_id,
+        batch_size,
+        num_tiles_per_image,
+        enabled_workers,
+        all_tiles,
+        tile_width,
+        tile_height,
+        mask_blur,
+    ):
         _, 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)
+        result_images = [tensor_to_pil(upscaled_image[b:b + 1], 0).copy() for b in range(batch_size)]
+
         log("USDU Dist: Using tile queue distribution")
         run_async_in_server_loop(
             init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers),
-            timeout=10.0
-        )
-        debug_log(
-            f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}"
+            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))
+        tile_masks = [
+            self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)
+            for tx, ty in all_tiles
+        ]
+        return result_images, tile_masks, width, height
 
+    def _process_master_static_initial_tiles(
+        self,
+        multi_job_id,
+        total_tiles,
+        all_tiles,
+        upscaled_image,
+        result_images,
+        tile_masks,
+        batch_size,
+        num_tiles_per_image,
+        tile_width,
+        tile_height,
+        padding,
+        force_uniform_tiles,
+        model,
+        positive,
+        negative,
+        vae,
+        seed,
+        steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        denoise,
+        tiled_decode,
+        width,
+        height,
+    ):
         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:
+            tile_id = run_async_in_server_loop(self._get_next_tile_index(multi_job_id), timeout=5.0)
+            if tile_id is None:
                 consecutive_no_tile += 1
                 if consecutive_no_tile >= max_consecutive_no_tile:
                     debug_log(f"Master processed {processed_count} tiles, moving to collection phase")
                     break
                 time.sleep(0.1)
-        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()
+                continue
 
-                    # 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
+            consecutive_no_tile = 0
+            processed_count += self._master_process_one_tile(
+                tile_id,
+                all_tiles,
+                upscaled_image,
+                result_images,
+                tile_masks,
+                multi_job_id,
+                batch_size,
+                num_tiles_per_image,
+                tile_width,
+                tile_height,
+                padding,
+                force_uniform_tiles,
+                model,
+                positive,
+                negative,
+                vae,
+                seed,
+                steps,
+                cfg,
+                sampler_name,
+                scheduler,
+                denoise,
+                tiled_decode,
+                width,
+                height,
             )
-        
-        # Blend worker tiles synchronously in deterministic tile order.
+            log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})")
+        return processed_count
+
+    def _collect_remaining_static_tiles(
+        self,
+        multi_job_id,
+        total_tiles,
+        master_processed_count,
+        all_tiles,
+        upscaled_image,
+        result_images,
+        tile_masks,
+        batch_size,
+        num_tiles_per_image,
+        tile_width,
+        tile_height,
+        padding,
+        force_uniform_tiles,
+        model,
+        positive,
+        negative,
+        vae,
+        seed,
+        steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        denoise,
+        tiled_decode,
+        width,
+        height,
+    ):
+        remaining_tiles = total_tiles - master_processed_count
+        if remaining_tiles <= 0:
+            return run_async_in_server_loop(self._get_all_completed_tasks(multi_job_id), timeout=5.0)
+
+        debug_log(f"Master waiting for {remaining_tiles} tiles from workers")
+        collected_tasks = run_async_in_server_loop(
+            self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles),
+            timeout=None,
+        )
+
+        completed_count = len(collected_tasks)
+        if completed_count >= total_tiles:
+            return collected_tasks
+
+        log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures")
+        while True:
+            comfy.model_management.throw_exception_if_processing_interrupted()
+            tile_id = run_async_in_server_loop(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,
+            )
+        return collected_tasks
+
+    def _blend_static_collected_tiles(
+        self,
+        collected_tasks,
+        result_images,
+        all_tiles,
+        tile_masks,
+        num_tiles_per_image,
+        batch_size,
+        tile_width,
+        tile_height,
+        padding,
+    ):
         def _sort_key(item):
             global_idx, tile_data = item
             batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image)
@@ -526,45 +596,130 @@ class StaticModeMixin:
             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_pil = tile_data['image'] if 'image' in tile_data else tensor_to_pil(tile_data['tensor'], 0)
             tile_mask = tile_masks[tile_idx]
             extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding)
             extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding)
-            result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil,
-                                                      x, y, (extracted_width, extracted_height), tile_mask, padding)
-        
+            result_images[batch_idx] = self.blend_tile(
+                result_images[batch_idx],
+                tile_pil,
+                x,
+                y,
+                (extracted_width, extracted_height),
+                tile_mask,
+                padding,
+            )
+
+    def _result_images_to_tensor(self, result_images, batch_size, upscaled_image):
+        if batch_size == 1:
+            result_tensor = pil_to_tensor(result_images[0])
+        else:
+            result_tensor = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
+        if upscaled_image.is_cuda:
+            result_tensor = result_tensor.cuda()
+        return result_tensor
+
+    def _process_master_static_sync(self, upscaled_image, 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]
+        total_tiles = batch_size * num_tiles_per_image
+        result_images, tile_masks, width, height = self._init_master_static_state(
+            upscaled_image=upscaled_image,
+            multi_job_id=multi_job_id,
+            batch_size=batch_size,
+            num_tiles_per_image=num_tiles_per_image,
+            enabled_workers=enabled_workers,
+            all_tiles=all_tiles,
+            tile_width=tile_width,
+            tile_height=tile_height,
+            mask_blur=mask_blur,
+        )
+
         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()
-            
+            master_processed_count = self._process_master_static_initial_tiles(
+                multi_job_id=multi_job_id,
+                total_tiles=total_tiles,
+                all_tiles=all_tiles,
+                upscaled_image=upscaled_image,
+                result_images=result_images,
+                tile_masks=tile_masks,
+                batch_size=batch_size,
+                num_tiles_per_image=num_tiles_per_image,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                force_uniform_tiles=force_uniform_tiles,
+                model=model,
+                positive=positive,
+                negative=negative,
+                vae=vae,
+                seed=seed,
+                steps=steps,
+                cfg=cfg,
+                sampler_name=sampler_name,
+                scheduler=scheduler,
+                denoise=denoise,
+                tiled_decode=tiled_decode,
+                width=width,
+                height=height,
+            )
+
+            collected_tasks = self._collect_remaining_static_tiles(
+                multi_job_id=multi_job_id,
+                total_tiles=total_tiles,
+                master_processed_count=master_processed_count,
+                all_tiles=all_tiles,
+                upscaled_image=upscaled_image,
+                result_images=result_images,
+                tile_masks=tile_masks,
+                batch_size=batch_size,
+                num_tiles_per_image=num_tiles_per_image,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+                force_uniform_tiles=force_uniform_tiles,
+                model=model,
+                positive=positive,
+                negative=negative,
+                vae=vae,
+                seed=seed,
+                steps=steps,
+                cfg=cfg,
+                sampler_name=sampler_name,
+                scheduler=scheduler,
+                denoise=denoise,
+                tiled_decode=tiled_decode,
+                width=width,
+                height=height,
+            )
+
+            self._blend_static_collected_tiles(
+                collected_tasks=collected_tasks,
+                result_images=result_images,
+                all_tiles=all_tiles,
+                tile_masks=tile_masks,
+                num_tiles_per_image=num_tiles_per_image,
+                batch_size=batch_size,
+                tile_width=tile_width,
+                tile_height=tile_height,
+                padding=padding,
+            )
+
+            result_tensor = self._result_images_to_tensor(result_images, batch_size, upscaled_image)
             log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
             return (result_tensor,)
         finally:
-            # Cleanup (async operation) - always execute
-            run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
+            run_async_in_server_loop(cleanup_job(multi_job_id), timeout=5.0)
diff --git a/upscale/payload_parsers.py b/upscale/payload_parsers.py
index 6dc64e4..37baadb 100644
--- a/upscale/payload_parsers.py
+++ b/upscale/payload_parsers.py
@@ -1,10 +1,12 @@
 import io
 import json
+from collections.abc import Mapping
+from typing import Any
 
 from PIL import Image
 
 
-def _parse_tiles_from_form(data):
+def parse_tiles_from_form(data: Mapping[str, Any]) -> list[dict[str, Any]]:
     """Parse tiles submitted via multipart/form-data into a list of tile dicts."""
     try:
         padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
@@ -51,14 +53,18 @@ def _parse_tiles_from_form(data):
         if 'batch_idx' in meta:
             try:
                 tile_info['batch_idx'] = int(meta['batch_idx'])
-            except Exception:
-                pass
+            except Exception as exc:
+                raise ValueError(f"Invalid batch_idx for tile {i}: {meta.get('batch_idx')} ({exc})")
         if 'global_idx' in meta:
             try:
                 tile_info['global_idx'] = int(meta['global_idx'])
-            except Exception:
-                pass
+            except Exception as exc:
+                raise ValueError(f"Invalid global_idx for tile {i}: {meta.get('global_idx')} ({exc})")
 
         tiles.append(tile_info)
 
     return tiles
+
+
+# Backward compatibility for existing imports.
+_parse_tiles_from_form = parse_tiles_from_form
diff --git a/upscale/processing_args.py b/upscale/processing_args.py
new file mode 100644
index 0000000..6256398
--- /dev/null
+++ b/upscale/processing_args.py
@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class UpscaleCoreArgs:
+    """Shared denoise/sampling arguments for USDU tile processing."""
+
+    model: Any
+    positive: Any
+    negative: Any
+    vae: Any
+    seed: int
+    steps: int
+    cfg: float
+    sampler_name: str
+    scheduler: str
+    denoise: float
+    tiled_decode: bool
diff --git a/upscale/result_collector.py b/upscale/result_collector.py
index 5ceb79c..a0517bd 100644
--- a/upscale/result_collector.py
+++ b/upscale/result_collector.py
@@ -4,8 +4,8 @@ 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_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
 
 
@@ -33,6 +33,10 @@ class ResultCollectorMixin:
             )
         return list(worker_status.keys())
 
+    async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
+        """Default timeout requeue hook; override in host mixins when needed."""
+        return await check_and_requeue_timed_out_workers(multi_job_id, batch_size)
+
     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)."""
@@ -110,21 +114,13 @@ class ResultCollectorMixin:
                         tile_idx = tile_data['tile_idx']
                         key = tile_data.get('global_idx', tile_idx)
                         entry = {
+                            **tile_data,
                             '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),
+                            'global_idx': key,
                         }
-                        if 'image' in tile_data:
-                            entry['image'] = tile_data['image']
-                        elif 'tensor' in tile_data:
-                            entry['tensor'] = tile_data['tensor']
-                        collected_results[key] = entry
+                        _ = entry['tile_idx'], entry['worker_id'], entry['global_idx']
+                        collected_results[entry['global_idx']] = entry
                 
                 elif mode == 'dynamic':
                     # Handle full images
@@ -142,12 +138,21 @@ class ResultCollectorMixin:
                     
             except asyncio.TimeoutError:
                 current_time = time.time()
-                waiting_workers = self._log_worker_timeout_status(job_data_snapshot, current_time, multi_job_id)
+                waiting_workers = type(self)._log_worker_timeout_status(
+                    self,
+                    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)
+                        # Use the class hook (can be overridden by host mixins/tests).
+                        requeued = await type(self)._check_and_requeue_timed_out_workers(
+                            self,
+                            multi_job_id,
+                            batch_size,
+                        )
                         if requeued > 0:
                             log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers")
                         last_heartbeat_check = current_time
@@ -179,12 +184,12 @@ class ResultCollectorMixin:
 
     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')
+        return await type(self)._async_collect_results(self, 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})
+        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)
@@ -193,6 +198,12 @@ class ResultCollectorMixin:
 
     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)
+        _ = master_processed_count
+        return await type(self)._async_collect_results(
+            self,
+            multi_job_id,
+            num_workers,
+            mode='dynamic',
+            remaining_to_collect=remaining_to_collect,
+            batch_size=batch_size,
+        )
diff --git a/upscale/tile_processing.py b/upscale/tile_processing.py
new file mode 100644
index 0000000..765ebea
--- /dev/null
+++ b/upscale/tile_processing.py
@@ -0,0 +1,56 @@
+from dataclasses import dataclass
+from typing import Any
+
+from .processing_args import UpscaleCoreArgs
+
+
+@dataclass(frozen=True)
+class TileBatchArgs:
+    """Tile/canvas parameters layered on top of shared core processing args."""
+
+    core: UpscaleCoreArgs
+    tile_width: int
+    tile_height: int
+    padding: int
+    force_uniform_tiles: bool
+    width: int
+    height: int
+
+
+def extract_and_process_tile_batch(
+    *,
+    node: Any,
+    upscaled_image: Any,
+    tx: int,
+    ty: int,
+    args: TileBatchArgs,
+) -> tuple[Any, int, int, int, int]:
+    """Extract one tile position for the whole batch and process it."""
+    tile_batch, x1, y1, ew, eh = node.extract_batch_tile_with_padding(
+        upscaled_image,
+        tx,
+        ty,
+        args.tile_width,
+        args.tile_height,
+        args.padding,
+        args.force_uniform_tiles,
+    )
+    region = (x1, y1, x1 + ew, y1 + eh)
+    core = args.core
+    processed_batch = node.process_tiles_batch(
+        tile_batch,
+        core.model,
+        core.positive,
+        core.negative,
+        core.vae,
+        core.seed,
+        core.steps,
+        core.cfg,
+        core.sampler_name,
+        core.scheduler,
+        core.denoise,
+        core.tiled_decode,
+        region,
+        (args.width, args.height),
+    )
+    return processed_batch, x1, y1, ew, eh
diff --git a/upscale/worker_comms.py b/upscale/worker_comms.py
index ddd0004..06b11c2 100644
--- a/upscale/worker_comms.py
+++ b/upscale/worker_comms.py
@@ -1,20 +1,33 @@
 import asyncio, io, json, time
+from typing import Any
 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.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):
+    async def _send_heartbeat_to_master(
+        self,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+    ) -> None:
         """Proxy heartbeat helper used by worker processing mixins."""
-        await _send_heartbeat_to_master(multi_job_id, master_url, worker_id)
+        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):
+    async def send_tiles_batch_to_master(
+        self,
+        processed_tiles: list[dict[str, Any]],
+        multi_job_id: str,
+        master_url: str,
+        padding: int,
+        worker_id: str,
+        is_final_flush: bool = False,
+    ) -> None:
         """Send all processed tiles to master, chunked if large."""
         if not processed_tiles:
             if is_final_flush:
diff --git a/utils/__init__.py b/utils/__init__.py
index 51b62aa..a2635dc 100644
--- a/utils/__init__.py
+++ b/utils/__init__.py
@@ -1,5 +1,3 @@
 """
 Utility modules for ComfyUI-Distributed extension.
 """
-
-# Make utils importable as a package
\ No newline at end of file
diff --git a/utils/async_helpers.py b/utils/async_helpers.py
index 36718bd..3a36879 100644
--- a/utils/async_helpers.py
+++ b/utils/async_helpers.py
@@ -105,7 +105,11 @@ class PromptValidationError(RuntimeError):
         super().__init__(f"Invalid prompt: {merged}")
 
 
-async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
+async def queue_prompt_payload(
+    prompt_obj: dict[str, Any],
+    workflow_meta: dict[str, Any] | None = None,
+    client_id: str | None = None,
+) -> str:
     """Validate and queue a prompt via ComfyUI's prompt queue."""
     payload = {"prompt": prompt_obj}
     payload = prompt_server.trigger_on_prompt(payload)
diff --git a/utils/audio_payload.py b/utils/audio_payload.py
index ff70225..76c5e78 100644
--- a/utils/audio_payload.py
+++ b/utils/audio_payload.py
@@ -1,6 +1,7 @@
 import base64
 import binascii
 import os
+from typing import Any
 
 import numpy as np
 import torch
@@ -13,7 +14,7 @@ MAX_AUDIO_PAYLOAD_BYTES = int(
 )
 
 
-def encode_audio_payload(audio_payload):
+def encode_audio_payload(audio_payload: dict[str, Any] | None) -> dict[str, Any] | None:
     """Serialize an AUDIO dict into JSON-safe canonical envelope payload."""
     if not isinstance(audio_payload, dict):
         return None
@@ -43,7 +44,7 @@ def encode_audio_payload(audio_payload):
     }
 
 
-def decode_audio_payload(audio_payload):
+def decode_audio_payload(audio_payload: dict[str, Any] | None) -> dict[str, Any] | None:
     """Decode canonical envelope audio payload into an AUDIO dict."""
     if audio_payload is None:
         return None
diff --git a/utils/cloudflare/binary.py b/utils/cloudflare/binary.py
index 4966345..e6eaed0 100644
--- a/utils/cloudflare/binary.py
+++ b/utils/cloudflare/binary.py
@@ -6,6 +6,7 @@ import shutil
 import stat
 from urllib import error as urlerror
 from urllib import request
+from urllib.parse import urlsplit
 
 from ..logging import debug_log
 
@@ -44,9 +45,24 @@ def _get_binary_path(bin_dir=None):
     return os.path.join(bin_dir, binary_name)
 
 
+def _get_release_download_base() -> str:
+    """Resolve cloudflared release download base URL from environment/config."""
+    configured = (os.environ.get("CLOUDFLARED_RELEASE_BASE") or "").strip()
+    if configured:
+        return configured.rstrip("/")
+    return "github.com/cloudflare/cloudflared/releases/latest/download"
+
+
 def _download_cloudflared():
     asset = _get_platform_binary_name()
-    url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
+    base = _get_release_download_base()
+    if "://" not in base:
+        scheme = (os.environ.get("CLOUDFLARED_RELEASE_SCHEME") or "https").strip() or "https"
+        base = f"{scheme}://{base}"
+    url = f"{base}/{asset}"
+    parsed = urlsplit(url)
+    if parsed.scheme not in {"http", "https"}:
+        raise RuntimeError(f"Unsupported cloudflared download URL scheme: {parsed.scheme}")
 
     bin_dir = _get_cloudflared_dir()
     os.makedirs(bin_dir, exist_ok=True)
@@ -54,7 +70,7 @@ def _download_cloudflared():
 
     debug_log(f"Downloading cloudflared from {url}")
     try:
-        with request.urlopen(url, timeout=30) as resp:
+        with request.urlopen(url, timeout=30) as resp:  # nosec B310 - scheme is allowlisted above
             with open(target_path, "wb") as f:
                 shutil.copyfileobj(resp, f)
     except urlerror.URLError as exc:
diff --git a/utils/cloudflare/process_reader.py b/utils/cloudflare/process_reader.py
index cf6de9f..e476a36 100644
--- a/utils/cloudflare/process_reader.py
+++ b/utils/cloudflare/process_reader.py
@@ -3,6 +3,7 @@
 import asyncio
 import re
 import threading
+from typing import Any
 
 from ..constants import CLOUDFLARE_LOG_BUFFER_SIZE
 from ..logging import debug_log
@@ -15,6 +16,7 @@ class ProcessReader:
     def __init__(self, log_file=None):
         self._process = None
         self._thread = None
+        self._task = None
         self._loop = None
         self._url_event = None
         self._public_url = None
@@ -37,6 +39,23 @@ class ProcessReader:
         if len(self._recent_logs) > CLOUDFLARE_LOG_BUFFER_SIZE:
             self._recent_logs = self._recent_logs[-CLOUDFLARE_LOG_BUFFER_SIZE:]
 
+    @staticmethod
+    def _normalize_line(raw_line):
+        if isinstance(raw_line, bytes):
+            return raw_line.decode("utf-8", errors="replace").strip()
+        return str(raw_line).strip()
+
+    def _process_log_line(self, line):
+        self._append_log(line)
+        match = PUBLIC_URL_PATTERN.search(line)
+        if match and not self._public_url:
+            self._public_url = match.group(1).rstrip("/")
+            return True
+
+        if "error" in line.lower() and not self._last_error:
+            self._last_error = line
+        return False
+
     def _reader(self):
         process = self._process
         if process is None:
@@ -44,34 +63,63 @@ class ProcessReader:
 
         loop = self._loop
         for raw_line in iter(process.stdout.readline, ""):
-            line = raw_line.strip()
+            line = self._normalize_line(raw_line)
             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
+            found_url = self._process_log_line(line)
+            if found_url and self._url_event and loop:
+                loop.call_soon_threadsafe(self._url_event.set)
 
         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):
+    async def _async_reader(self):
+        process = self._process
+        if process is None or process.stdout is None:
+            if self._url_event:
+                self._url_event.set()
+            return
+
+        while True:
+            raw_line = await process.stdout.readline()
+            if not raw_line:
+                break
+
+            line = self._normalize_line(raw_line)
+            if not line:
+                continue
+
+            found_url = self._process_log_line(line)
+            if found_url and self._url_event:
+                self._url_event.set()
+
+        if self._url_event:
+            if not self._last_error and not self._public_url:
+                self._last_error = "Cloudflare tunnel exited before becoming ready"
+            self._url_event.set()
+
+    def start(self, process: Any, loop: asyncio.AbstractEventLoop) -> None:
         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()
+        self._thread = None
+        self._task = None
+
+        stdout = getattr(process, "stdout", None)
+        if stdout is None:
+            return
+
+        if asyncio.iscoroutinefunction(getattr(stdout, "readline", None)):
+            self._task = loop.create_task(self._async_reader())
+        else:
+            self._thread = threading.Thread(target=self._reader, daemon=True)
+            self._thread.start()
 
     async def wait_for_url(self, timeout):
         if not self._url_event:
@@ -79,7 +127,11 @@ class ProcessReader:
         await asyncio.wait_for(self._url_event.wait(), timeout=timeout)
         return self._public_url
 
-    def stop(self):
+    def stop(self) -> None:
+        if self._task and not self._task.done():
+            self._task.cancel()
+        self._task = None
+
         if self._thread and self._thread.is_alive():
             self._thread.join(timeout=1)
         self._thread = None
diff --git a/utils/cloudflare/state.py b/utils/cloudflare/state.py
index 4ec08b5..dc377b2 100644
--- a/utils/cloudflare/state.py
+++ b/utils/cloudflare/state.py
@@ -1,17 +1,30 @@
 """Cloudflare tunnel state persistence helpers."""
 
+from dataclasses import dataclass
+from typing import Any
+
 from ..config import load_config, save_config
 from ..network import normalize_host
 
 
-def _get_tunnel_config(cfg):
+@dataclass(frozen=True)
+class TunnelStateUpdate:
+    status: str | None = None
+    public_url: str | None = None
+    pid: int | None = None
+    log_file: str | None = None
+    previous_host: str | None = None
+    master_host: str | None = None
+
+
+def _get_tunnel_config(cfg: dict[str, Any]) -> dict[str, Any]:
     tunnel_cfg = cfg.get("tunnel", {})
     if isinstance(tunnel_cfg, dict):
         return tunnel_cfg
     return {}
 
 
-def load_tunnel_state():
+def load_tunnel_state() -> dict[str, Any]:
     cfg = load_config()
     tunnel_cfg = _get_tunnel_config(cfg)
     master_cfg = cfg.get("master", {}) if isinstance(cfg.get("master", {}), dict) else {}
@@ -25,46 +38,45 @@ def load_tunnel_state():
     }
 
 
-def persist_tunnel_state(
-    status=None,
-    public_url=None,
-    pid=None,
-    log_file=None,
-    previous_host=None,
-    master_host=None,
-):
+def persist_tunnel_state(update: TunnelStateUpdate) -> 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
+    if update.status is not None:
+        tunnel_cfg["status"] = update.status
+    if update.public_url is not None:
+        tunnel_cfg["public_url"] = update.public_url
+    if update.pid is not None:
+        tunnel_cfg["pid"] = update.pid
+    if update.log_file is not None:
+        tunnel_cfg["log_file"] = update.log_file
+    if update.previous_host is not None:
+        tunnel_cfg["previous_master_host"] = update.previous_host
+    if update.master_host is not None:
+        cfg.setdefault("master", {})["host"] = update.master_host
 
     cfg["tunnel"] = tunnel_cfg
     save_config(cfg)
 
 
-def clear_tunnel_state(log_file=None, previous_host=None, master_host=None):
+def clear_tunnel_state(
+    log_file: str | None = None,
+    previous_host: str | None = None,
+    master_host: str | None = None,
+) -> None:
     persist_tunnel_state(
-        status="stopped",
-        public_url="",
-        pid=None,
-        log_file=log_file,
-        previous_host=previous_host,
-        master_host=master_host,
+        TunnelStateUpdate(
+            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):
+def resolve_restore_master_host(previous_master_host: str | None) -> str | None:
     """Determine whether master host should be restored after tunnel stop."""
     cfg = load_config()
     tunnel_cfg = _get_tunnel_config(cfg)
diff --git a/utils/cloudflare/tunnel.py b/utils/cloudflare/tunnel.py
index 936819d..a11bfa5 100644
--- a/utils/cloudflare/tunnel.py
+++ b/utils/cloudflare/tunnel.py
@@ -4,16 +4,22 @@ import asyncio
 import os
 import shutil
 import signal
-import subprocess
 import time
+from typing import Any
 
-from ..constants import TUNNEL_START_TIMEOUT
+from ..constants import PROCESS_TERMINATION_TIMEOUT, 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
+from .state import (
+    TunnelStateUpdate,
+    clear_tunnel_state,
+    load_tunnel_state,
+    persist_tunnel_state,
+    resolve_restore_master_host,
+)
 
 
 class CloudflareTunnelManager:
@@ -36,6 +42,12 @@ class CloudflareTunnelManager:
     def base_dir(self):
         return os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
 
+    @staticmethod
+    def _build_local_tunnel_url(port: int) -> str:
+        scheme = (os.environ.get("CLOUDFLARE_LOCAL_BIND_SCHEME") or "http").strip() or "http"
+        host = (os.environ.get("CLOUDFLARE_LOCAL_BIND_HOST") or "127.0.0.1").strip() or "127.0.0.1"
+        return f"{scheme}://{host}:{port}"
+
     def _restore_state(self):
         state = load_tunnel_state()
 
@@ -53,9 +65,47 @@ class CloudflareTunnelManager:
             self.status = "stopped"
             self.pid = None
 
-    async def start_tunnel(self):
+    @staticmethod
+    def _process_is_running(process):
+        poll = getattr(process, "poll", None)
+        if callable(poll):
+            return poll() is None
+        return getattr(process, "returncode", None) is None
+
+    @staticmethod
+    async def _launch_process_with_timeout(command, timeout_seconds, **create_kwargs):
+        """Launch subprocess with bounded startup wait."""
+        try:
+            return await asyncio.wait_for(
+                asyncio.create_subprocess_exec(*command, **create_kwargs),
+                timeout=float(timeout_seconds),
+            )
+        except asyncio.TimeoutError as exc:
+            raise TimeoutError(
+                f"Timed out launching cloudflared after {timeout_seconds} seconds"
+            ) from exc
+
+    @staticmethod
+    async def _terminate_process_with_timeout(process, timeout_seconds):
+        """Terminate process object with timeout for both sync and async process types."""
+        poll = getattr(process, "poll", None)
+        if callable(poll):
+            terminate_process(process, timeout=timeout_seconds)
+            return
+
+        if getattr(process, "returncode", None) is not None:
+            return
+
+        process.terminate()
+        try:
+            await asyncio.wait_for(process.wait(), timeout=float(timeout_seconds))
+        except asyncio.TimeoutError:
+            process.kill()
+            await process.wait()
+
+    async def start_tunnel(self) -> dict[str, Any]:
         async with self._lock:
-            if self.process and self.process.poll() is None:
+            if self.process and self._process_is_running(self.process):
                 return {
                     "status": self.status,
                     "public_url": self.public_url,
@@ -90,17 +140,16 @@ class CloudflareTunnelManager:
                 "tunnel",
                 "--no-autoupdate",
                 "--url",
-                f"http://127.0.0.1:{port}",
+                self._build_local_tunnel_url(port),
             ]
 
             debug_log(f"Starting cloudflared: {' '.join(cmd)}")
             try:
-                self.process = subprocess.Popen(
+                self.process = await self._launch_process_with_timeout(
                     cmd,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.STDOUT,
-                    text=True,
-                    bufsize=1,
+                    timeout_seconds=10.0,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.STDOUT,
                 )
             except FileNotFoundError:
                 self.status = "error"
@@ -111,10 +160,12 @@ class CloudflareTunnelManager:
 
             self.pid = self.process.pid
             persist_tunnel_state(
-                status="starting",
-                pid=self.pid,
-                log_file=self.log_file,
-                previous_host=self.previous_master_host,
+                TunnelStateUpdate(
+                    status="starting",
+                    pid=self.pid,
+                    log_file=self.log_file,
+                    previous_host=self.previous_master_host,
+                )
             )
 
             loop = asyncio.get_running_loop()
@@ -139,12 +190,14 @@ class CloudflareTunnelManager:
             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),
+                TunnelStateUpdate(
+                    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,
@@ -153,7 +206,7 @@ class CloudflareTunnelManager:
                 "log_file": self.log_file,
             }
 
-    async def stop_tunnel(self):
+    async def stop_tunnel(self) -> dict[str, Any]:
         async with self._lock:
             pid = self.process.pid if self.process else self.pid
             if not pid:
@@ -163,7 +216,10 @@ class CloudflareTunnelManager:
 
             debug_log(f"Stopping cloudflared (pid={pid})")
             if self.process:
-                terminate_process(self.process, timeout=5)
+                await self._terminate_process_with_timeout(
+                    self.process,
+                    timeout_seconds=PROCESS_TERMINATION_TIMEOUT,
+                )
             else:
                 try:
                     os.kill(pid, signal.SIGTERM)
@@ -187,7 +243,7 @@ class CloudflareTunnelManager:
             )
             return {"status": "stopped"}
 
-    def get_status(self):
+    def get_status(self) -> dict[str, Any]:
         alive = False
         pid = self.process.pid if self.process else self.pid
         if pid:
diff --git a/utils/config.py b/utils/config.py
index c580415..b6b4ecb 100644
--- a/utils/config.py
+++ b/utils/config.py
@@ -4,22 +4,39 @@ Configuration management for ComfyUI-Distributed.
 import asyncio
 import os
 import json
+from dataclasses import dataclass
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
+from functools import lru_cache
+from typing import Any
 from .logging import log
 
 # Import defaults for timeout fallbacks
-from .constants import HEARTBEAT_TIMEOUT
+from .constants import GPU_CONFIG_FILE, 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()
+CONFIG_FILE = GPU_CONFIG_FILE
 
 
-def _config_path():
+@dataclass
+class _ConfigState:
+    cache: dict[str, Any] | None = None
+    mtime: float = 0.0
+
+
+@lru_cache(maxsize=1)
+def _config_state() -> _ConfigState:
+    return _ConfigState()
+
+
+@lru_cache(maxsize=1)
+def _config_lock() -> asyncio.Lock:
+    return asyncio.Lock()
+
+
+def _config_path() -> str:
     return CONFIG_FILE
 
-def get_default_config():
+def get_default_config() -> dict[str, Any]:
     """Returns the default configuration dictionary. Single source of truth."""
     return {
         "master": {"host": ""},
@@ -44,7 +61,7 @@ def get_default_config():
         }
     }
 
-def _merge_with_defaults(data, defaults):
+def _merge_with_defaults(data: dict[str, Any], defaults: dict[str, Any]) -> dict[str, Any]:
     """Recursively merge loaded config data with default keys."""
     if not isinstance(data, dict):
         return defaults
@@ -65,38 +82,38 @@ def _merge_with_defaults(data, defaults):
     return merged
 
 
-def invalidate_config_cache():
+def invalidate_config_cache() -> None:
     """Invalidate in-memory config cache so next load reads from disk."""
-    global _config_cache, _config_mtime
-    _config_cache = None
-    _config_mtime = 0.0
+    state = _config_state()
+    state.cache = None
+    state.mtime = 0.0
 
 
-def load_config():
+def load_config() -> dict[str, Any]:
     """Loads the config, falling back to defaults if the file is missing or invalid."""
-    global _config_cache, _config_mtime
+    state = _config_state()
     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 state.cache is None:
+            state.cache = get_default_config()
+        return state.cache
 
-    if _config_cache is None or mtime != _config_mtime:
+    if state.cache is None or mtime != state.mtime:
         try:
             with open(path, 'r', encoding='utf-8') as f:
                 loaded = json.load(f)
-            _config_cache = _merge_with_defaults(loaded, get_default_config())
+            state.cache = _merge_with_defaults(loaded, get_default_config())
         except Exception as e:
             log(f"Error loading config, using defaults: {e}")
-            _config_cache = get_default_config()
-        _config_mtime = mtime
+            state.cache = get_default_config()
+        state.mtime = mtime
 
-    return _config_cache
+    return state.cache
 
-def save_config(config):
+def save_config(config: dict[str, Any]) -> bool:
     """Saves the configuration to file."""
     tmp_path = f"{_config_path()}.tmp"
     try:
@@ -110,16 +127,16 @@ def save_config(config):
     except Exception as e:
         try:
             os.unlink(tmp_path)
-        except OSError:
-            pass
+        except OSError as cleanup_error:
+            log(f"Error removing temporary config file '{tmp_path}': {cleanup_error}")
         log(f"Error saving config: {e}")
         return False
 
 
 @asynccontextmanager
-async def config_transaction():
+async def config_transaction() -> AsyncIterator[dict[str, Any]]:
     """Acquire config lock, yield loaded config, and save if changed."""
-    async with _config_lock:
+    async with _config_lock():
         config = load_config()
         original_snapshot = json.dumps(config, sort_keys=True)
         yield config
@@ -128,7 +145,7 @@ async def config_transaction():
             if not save_config(config):
                 raise RuntimeError("Failed to save config")
 
-def ensure_config_exists():
+def ensure_config_exists() -> None:
     """Creates default config file if it doesn't exist. Used by __init__.py"""
     if not os.path.exists(_config_path()):
         default_config = get_default_config()
diff --git a/utils/constants.py b/utils/constants.py
index caca4f9..f25191e 100644
--- a/utils/constants.py
+++ b/utils/constants.py
@@ -23,6 +23,8 @@ LOG_TAIL_BYTES = 65536  # 64KB
 
 # File paths
 WORKER_LOG_PATTERN = "distributed_worker_*.log"
+PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))
+GPU_CONFIG_FILE = os.path.join(PROJECT_ROOT, "gpu_config.json")
 
 # Worker management
 WORKER_STARTUP_DELAY = 2.0
diff --git a/utils/crop_model_patch.py b/utils/crop_model_patch.py
index 17ba2a4..fd24bbd 100644
--- a/utils/crop_model_patch.py
+++ b/utils/crop_model_patch.py
@@ -1,13 +1,24 @@
 from contextlib import contextmanager
+from typing import Any, Iterator, TypeAlias
 
 import torch
 
 from .logging import debug_log
 from .usdu_utils import resize_region
 
+CropRegion: TypeAlias = tuple[int, int, int, int]
+CropRegions: TypeAlias = list[CropRegion] | CropRegion
+
 
 @contextmanager
-def crop_model_cond(model, crop_regions, init_size, canvas_size, tile_size, latent_crop=False):
+def crop_model_cond(
+    model: Any,
+    crop_regions: CropRegions,
+    init_size: tuple[int, int],
+    canvas_size: tuple[int, int],
+    tile_size: tuple[int, int],
+    latent_crop: bool = False,
+) -> Iterator[Any]:
     """Clone model and crop compatible model patches for tile-local sampling."""
     try:
         patched_model = model.clone()
@@ -71,7 +82,12 @@ class ModelPatchCropper:
         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):
+    def crop(
+        self,
+        crop_regions: CropRegions,
+        canvas_size: tuple[int, int],
+        latent_crop: bool = True,
+    ) -> "ModelPatchCropper":
         patch = self.patch
 
         if not isinstance(crop_regions, list):
diff --git a/utils/image.py b/utils/image.py
index c6786df..efd60e9 100644
--- a/utils/image.py
+++ b/utils/image.py
@@ -1,6 +1,10 @@
 """
 Image and tensor conversion utilities for ComfyUI-Distributed.
 """
+import base64
+import io
+from typing import Any
+
 import torch
 import numpy as np
 from PIL import Image
@@ -23,3 +27,38 @@ def ensure_contiguous(tensor):
         return tensor.contiguous()
     return tensor
 
+
+def encode_tensor_png_data_url(image_batch, batch_index=0):
+    """Encode one batched tensor image as a PNG data URL."""
+    image = tensor_to_pil(image_batch, batch_index)
+    byte_io = io.BytesIO()
+    image.save(byte_io, format="PNG", compress_level=0)
+    encoded = base64.b64encode(byte_io.getvalue()).decode("utf-8")
+    return f"data:image/png;base64,{encoded}"
+
+
+def blend_processed_batch_item(
+    result_images: list[Any],
+    processed_batch: Any,
+    batch_index: int,
+    blend_fn: Any,
+    x1: int,
+    y1: int,
+    ew: int,
+    eh: int,
+    tile_mask: Any,
+    padding: int,
+) -> None:
+    """Blend one processed batch item back into result_images in-place."""
+    tile_pil = tensor_to_pil(processed_batch, batch_index)
+    if tile_pil.size != (ew, eh):
+        tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
+    result_images[batch_index] = blend_fn(
+        result_images[batch_index],
+        tile_pil,
+        x1,
+        y1,
+        (ew, eh),
+        tile_mask,
+        padding,
+    )
diff --git a/utils/logging.py b/utils/logging.py
index f4d272c..905470e 100644
--- a/utils/logging.py
+++ b/utils/logging.py
@@ -4,40 +4,53 @@ Shared logging utilities for ComfyUI-Distributed.
 import os
 import json
 import time
+from dataclasses import dataclass
+from functools import lru_cache
+from .constants import GPU_CONFIG_FILE
 
-# 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():
+
+@dataclass
+class _DebugState:
+    enabled: bool | None = None
+    updated_at: float = 0.0
+
+
+@lru_cache(maxsize=1)
+def _debug_state() -> _DebugState:
+    return _DebugState()
+
+def is_debug_enabled() -> bool:
     """Check if debug is enabled."""
-    global _debug_cache, _debug_cache_time
+    state = _debug_state()
 
     now = time.monotonic()
-    if _debug_cache is not None and (now - _debug_cache_time) < _DEBUG_TTL:
-        return _debug_cache
+    if state.enabled is not None and (now - state.updated_at) < _DEBUG_TTL:
+        return state.enabled
 
-    enabled = False
-    if os.path.exists(CONFIG_FILE):
+    enabled = state.enabled if state.enabled is not None else False
+    if os.path.exists(GPU_CONFIG_FILE):
         try:
-            with open(CONFIG_FILE, 'r') as f:
+            with open(GPU_CONFIG_FILE, "r", encoding="utf-8") as f:
                 config = json.load(f)
-                enabled = config.get("settings", {}).get("debug", False)
-        except (OSError, json.JSONDecodeError, ValueError):
-            pass
+                enabled = bool(config.get("settings", {}).get("debug", False))
+        except (OSError, json.JSONDecodeError, ValueError) as exc:
+            # Fall back to the previous cached value to avoid flapping behavior.
+            print(
+                "[Distributed] Failed to read debug flag from config; "
+                f"using cached value {enabled}: {exc}"
+                )
 
-    _debug_cache = enabled
-    _debug_cache_time = now
+    state.enabled = enabled
+    state.updated_at = now
     return enabled
 
-def debug_log(message):
+def debug_log(message: str) -> None:
     """Log debug messages only if debug is enabled in config."""
     if is_debug_enabled():
         print(f"[Distributed] {message}")
 
-def log(message):
+def log(message: str) -> None:
     """Always log important messages."""
     print(f"[Distributed] {message}")
diff --git a/utils/network.py b/utils/network.py
index ce5c7a7..9317f48 100644
--- a/utils/network.py
+++ b/utils/network.py
@@ -1,60 +1,80 @@
+from __future__ import annotations
+
 """
 Network and API utilities for ComfyUI-Distributed.
 """
 import asyncio
 import aiohttp
-import re
 import server
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any
 from aiohttp import web
+from urllib.parse import urlsplit
+from .exceptions import DistributedError
 from .logging import debug_log
 
-# Shared session for connection pooling
-_client_session = None
+@dataclass
+class _NetworkState:
+    client_session: aiohttp.ClientSession | None = None
 
-async def get_client_session():
+
+@lru_cache(maxsize=1)
+def _network_state() -> _NetworkState:
+    return _NetworkState()
+
+async def get_client_session() -> aiohttp.ClientSession:
     """Get or create a shared aiohttp client session."""
-    global _client_session
+    state = _network_state()
     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:
+    if state.client_session is None or state.client_session.closed:
         connector = aiohttp.TCPConnector(limit=100, limit_per_host=30)
         # Don't set timeout here - set it per request
-        _client_session = aiohttp.ClientSession(connector=connector)
-    return _client_session
+        state.client_session = aiohttp.ClientSession(connector=connector)
+    return state.client_session
 
-async def cleanup_client_session():
+async def cleanup_client_session() -> None:
     """Clean up the shared client session."""
-    global _client_session
-    if _client_session and not _client_session.closed:
-        await _client_session.close()
-        _client_session = None
+    state = _network_state()
+    if state.client_session and not state.client_session.closed:
+        await state.client_session.close()
+        state.client_session = None
 
-async def handle_api_error(request, error, status=500):
+async def handle_api_error(request: web.Request, error: Any, status: int = 500) -> web.StreamResponse:
     """Standardized error response handler."""
+    _ = request
     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)
 
+    if isinstance(error, DistributedError):
+        debug_log(f"API Error [{status}] ({error.__class__.__name__}): {error}")
+        return web.json_response(
+            {"error": str(error), "error_type": error.__class__.__name__},
+            status=status,
+        )
+
     message = str(error)
     debug_log(f"API Error [{status}]: {message}")
     return web.json_response({"error": message}, status=status)
 
-def get_server_port():
+def get_server_port() -> int:
     """Get the ComfyUI server port."""
     import server
     return server.PromptServer.instance.port
 
-def get_server_loop():
+def get_server_loop() -> asyncio.AbstractEventLoop:
     """Get the ComfyUI server event loop."""
     import server
     return server.PromptServer.instance.loop
 
 
-def normalize_host(value):
+def normalize_host(value: Any) -> Any:
     if value is None:
         return None
     if not isinstance(value, str):
@@ -62,11 +82,13 @@ def normalize_host(value):
     host = value.strip()
     if not host:
         return host
-    host = re.sub(r"^https?://", "", host, flags=re.IGNORECASE)
+    if "://" in host:
+        parsed = urlsplit(host)
+        host = parsed.netloc or parsed.path
     return host.split("/")[0]
 
 
-def build_worker_url(worker, endpoint=""):
+def build_worker_url(worker: dict[str, Any], endpoint: str = "") -> str:
     """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)
@@ -86,7 +108,7 @@ def build_worker_url(worker, endpoint=""):
     return f"{base}{endpoint}"
 
 
-async def probe_worker(worker_url: str, timeout: float = 5.0) -> dict | None:
+async def probe_worker(worker_url: str, timeout: float = 5.0) -> dict[str, Any] | None:
     """GET {worker_url}/prompt. Returns parsed JSON or None on any failure."""
     base_url = (worker_url or "").strip().rstrip("/")
     if not base_url:
@@ -117,7 +139,10 @@ async def probe_worker(worker_url: str, timeout: float = 5.0) -> dict | None:
         return None
 
 
-def build_master_url(config=None, prompt_server_instance=None):
+def build_master_url(
+    config: dict[str, Any] | None = None,
+    prompt_server_instance: Any | None = None,
+) -> str:
     """Build the best public URL workers should use to reach the master."""
     if config is None:
         from .config import load_config
@@ -158,7 +183,7 @@ def build_master_url(config=None, prompt_server_instance=None):
         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", "::"):
+    if address in ("0.0.0.0", "::"):  # nosec B104 - normalization of wildcard bind host
         address = "127.0.0.1"
     scheme = "https" if port == 443 else "http"
     default_port_for_scheme = 443 if scheme == "https" else 80
diff --git a/utils/process.py b/utils/process.py
index 85887b9..42bf24d 100644
--- a/utils/process.py
+++ b/utils/process.py
@@ -1,18 +1,121 @@
 """
 Process management utilities for ComfyUI-Distributed.
 """
+import asyncio
 import os
-import subprocess
 import platform
 import signal
+import subprocess
+import threading
+from collections.abc import Sequence
+from typing import Any
 
-def is_process_alive(pid):
+
+def _run_coroutine_blocking(coro):
+    """Run a coroutine in blocking contexts, even if an event loop is active."""
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return asyncio.run(coro)
+
+    outcome = {}
+
+    def _runner():
+        try:
+            outcome["value"] = asyncio.run(coro)
+        except Exception as exc:  # pragma: no cover - thin bridge
+            outcome["error"] = exc
+
+    thread = threading.Thread(target=_runner, daemon=True)
+    thread.start()
+    thread.join()
+
+    error = outcome.get("error")
+    if error is not None:
+        raise error
+    return outcome.get("value")
+
+
+class AsyncProcessAdapter:
+    """Sync-friendly adapter around asyncio subprocess.Process."""
+
+    def __init__(self, process):
+        self._process = process
+        self.pid = process.pid
+
+    @property
+    def returncode(self):
+        return self._process.returncode
+
+    @property
+    def stdout(self):
+        return self._process.stdout
+
+    def poll(self):
+        return self._process.returncode
+
+    def terminate(self):
+        self._process.terminate()
+
+    def kill(self):
+        self._process.kill()
+
+    def wait(self, timeout=None):
+        async def _wait():
+            if timeout is None:
+                return await self._process.wait()
+            return await asyncio.wait_for(self._process.wait(), timeout=float(timeout))
+
+        return _run_coroutine_blocking(_wait())
+
+
+def _normalize_stdio_value(value):
+    if value is subprocess.PIPE:
+        return asyncio.subprocess.PIPE
+    if value is subprocess.STDOUT:
+        return asyncio.subprocess.STDOUT
+    if value is subprocess.DEVNULL:
+        return asyncio.subprocess.DEVNULL
+    return value
+
+
+def launch_process_with_timeout(
+    command: Sequence[str | bytes | os.PathLike[str]],
+    timeout_seconds: float = 10.0,
+    **create_kwargs: Any,
+) -> AsyncProcessAdapter:
+    """Launch a child process with startup timeout and sync-friendly lifecycle methods."""
+    if isinstance(command, (str, bytes)):
+        raise TypeError("Command must be a sequence, not a string")
+
+    normalized_kwargs = dict(create_kwargs)
+    normalized_kwargs["stdout"] = _normalize_stdio_value(normalized_kwargs.get("stdout"))
+    normalized_kwargs["stderr"] = _normalize_stdio_value(normalized_kwargs.get("stderr"))
+
+    async def _launch():
+        return await asyncio.wait_for(
+            asyncio.create_subprocess_exec(*command, **normalized_kwargs),
+            timeout=float(timeout_seconds),
+        )
+
+    try:
+        process = _run_coroutine_blocking(_launch())
+    except asyncio.TimeoutError as exc:
+        raise TimeoutError(f"Timed out launching process after {timeout_seconds} seconds") from exc
+
+    return AsyncProcessAdapter(process)
+
+def is_process_alive(pid: int) -> bool:
     """Check if a process with given PID is still alive."""
     try:
         if platform.system() == "Windows":
             # Windows: use tasklist
-            result = subprocess.run(['tasklist', '/FI', f'PID eq {pid}'], 
-                                  capture_output=True, text=True)
+            result = subprocess.run(
+                ['tasklist', '/FI', f'PID eq {pid}'],
+                capture_output=True,
+                text=True,
+                timeout=5.0,
+            )
             return str(pid) in result.stdout
         else:
             # Unix: send signal 0
@@ -27,7 +130,7 @@ def terminate_process(process, timeout=5):
         process.terminate()
         try:
             process.wait(timeout=timeout)
-        except subprocess.TimeoutExpired:
+        except (subprocess.TimeoutExpired, TimeoutError):
             process.kill()
             process.wait()
 
@@ -35,4 +138,3 @@ def get_python_executable():
     """Get the Python executable path."""
     import sys
     return sys.executable
-
diff --git a/utils/usdu_managment.py b/utils/usdu_managment.py
index a4109f3..757efb5 100644
--- a/utils/usdu_managment.py
+++ b/utils/usdu_managment.py
@@ -10,22 +10,22 @@ Route handlers and job logic now live in:
 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,
+    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 ..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
 
 
-async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
+async def send_heartbeat_to_master(multi_job_id: str, master_url: str, worker_id: str) -> None:
     """Send heartbeat to master from worker-side processing loops."""
     try:
         data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)}
@@ -37,8 +37,27 @@ async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
         debug_log(f"Heartbeat failed: {e}")
 
 
+# Backward-compatible aliases for legacy imports.
+_check_and_requeue_timed_out_workers = check_and_requeue_timed_out_workers
+_cleanup_job = cleanup_job
+_drain_results_queue = drain_results_queue
+_get_completed_count = get_completed_count
+_init_job_queue = init_job_queue
+_mark_task_completed = mark_task_completed
+_parse_tiles_from_form = parse_tiles_from_form
+_send_heartbeat_to_master = send_heartbeat_to_master
+
+
 __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",
     "_check_and_requeue_timed_out_workers",
     "_cleanup_job",
     "_drain_results_queue",
diff --git a/utils/usdu_utils.py b/utils/usdu_utils.py
index 39b1cd3..4fd594f 100644
--- a/utils/usdu_utils.py
+++ b/utils/usdu_utils.py
@@ -4,6 +4,7 @@ import torch
 import torch.nn.functional as F
 from torchvision.transforms import GaussianBlur
 import math
+from typing import Any
 
 if (not hasattr(Image, 'Resampling')):  # For older versions of Pillow
     Image.Resampling = Image
@@ -11,13 +12,13 @@ if (not hasattr(Image, 'Resampling')):  # For older versions of Pillow
 BLUR_KERNEL_SIZE = 15
 
 
-def tensor_to_pil(img_tensor, batch_index=0):
+def tensor_to_pil(img_tensor: torch.Tensor, batch_index: int = 0) -> Image.Image:
     # Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels]
     # and returns an RGB PIL Image. Assumes channels=3
     return Image.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8))
 
 
-def pil_to_tensor(image):
+def pil_to_tensor(image: Image.Image) -> torch.Tensor:
     # Takes a PIL image and returns a tensor of shape [1, height, width, channels]
     image = np.array(image).astype(np.float32) / 255.0
     image = torch.from_numpy(image).unsqueeze(0)
@@ -26,27 +27,31 @@ def pil_to_tensor(image):
     return image
 
 
-def controlnet_hint_to_pil(tensor, batch_index=0):
+def controlnet_hint_to_pil(tensor: torch.Tensor, batch_index: int = 0) -> Image.Image:
     return tensor_to_pil(tensor.movedim(1, -1), batch_index)
 
 
-def pil_to_controlnet_hint(img):
+def pil_to_controlnet_hint(img: Image.Image) -> torch.Tensor:
     return pil_to_tensor(img).movedim(-1, 1)
 
 
-def crop_tensor(tensor, region):
+def crop_tensor(tensor: torch.Tensor, region: tuple[int, int, int, int]) -> torch.Tensor:
     # Takes a tensor of shape [batch_size, height, width, channels] and crops it to the given region
     x1, y1, x2, y2 = region
     return tensor[:, y1:y2, x1:x2, :]
 
 
-def resize_tensor(tensor, size, mode="nearest-exact"):
+def resize_tensor(
+    tensor: torch.Tensor,
+    size: tuple[int, int],
+    mode: str = "nearest-exact",
+) -> torch.Tensor:
     # Takes a tensor of shape [B, C, H, W] and resizes
     # it to a shape of [B, C, size[0], size[1]] using the given mode
     return torch.nn.functional.interpolate(tensor, size=size, mode=mode)
 
 
-def get_crop_region(mask, pad=0):
+def get_crop_region(mask: Image.Image, pad: int = 0) -> tuple[int, int, int, int]:
     # Takes a black and white PIL image in 'L' mode and returns the coordinates of the white rectangular mask region
     # Should be equivalent to the get_crop_region function from https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/masking.py
     coordinates = mask.getbbox()
@@ -62,7 +67,10 @@ def get_crop_region(mask, pad=0):
     return fix_crop_region((x1, y1, x2, y2), (mask.width, mask.height))
 
 
-def fix_crop_region(region, image_size):
+def fix_crop_region(
+    region: tuple[int, int, int, int],
+    image_size: tuple[int, int],
+) -> tuple[int, int, int, int]:
     # Remove the extra pixel added by the get_crop_region function
     image_width, image_height = image_size
     x1, y1, x2, y2 = region
@@ -73,7 +81,13 @@ def fix_crop_region(region, image_size):
     return x1, y1, x2, y2
 
 
-def expand_crop(region, width, height, target_width, target_height):
+def expand_crop(
+    region: tuple[int, int, int, int],
+    width: int,
+    height: int,
+    target_width: int,
+    target_height: int,
+) -> tuple[tuple[int, int, int, int], tuple[int, int]]:
     '''
     Expands a crop region to a specified target size.
     :param region: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points
@@ -112,7 +126,11 @@ def expand_crop(region, width, height, target_width, target_height):
     return (x1, y1, x2, y2), (target_width, target_height)
 
 
-def resize_region(region, init_size, resize_size):
+def resize_region(
+    region: tuple[int, int, int, int],
+    init_size: tuple[int, int],
+    resize_size: tuple[int, int],
+) -> tuple[int, int, int, int]:
     # Resize a crop so that it fits an image that was resized to the given width and height
     x1, y1, x2, y2 = region
     init_width, init_height = init_size
@@ -124,7 +142,15 @@ def resize_region(region, init_size, resize_size):
     return (x1, y1, x2, y2)
 
 
-def pad_image(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False):
+def pad_image(
+    image: Image.Image,
+    left_pad: int,
+    right_pad: int,
+    top_pad: int,
+    bottom_pad: int,
+    fill: bool = False,
+    blur: bool = False,
+) -> Image.Image:
     '''
     Pads an image with the given number of pixels on each side and fills the padding with data from the edges.
     :param image: A PIL image
@@ -166,7 +192,15 @@ def pad_image(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=
     return padded_image
 
 
-def pad_image2(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False):
+def pad_image2(
+    image: Image.Image,
+    left_pad: int,
+    right_pad: int,
+    top_pad: int,
+    bottom_pad: int,
+    fill: bool = False,
+    blur: bool = False,
+) -> Image.Image:
     '''
     Pads an image with the given number of pixels on each side and fills the padding with data from the edges. 
     Faster than pad_image, but only pads with edge data in straight lines.
@@ -203,7 +237,15 @@ def pad_image2(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur
     return padded_image
 
 
-def pad_tensor(tensor, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False):
+def pad_tensor(
+    tensor: torch.Tensor,
+    left_pad: int,
+    right_pad: int,
+    top_pad: int,
+    bottom_pad: int,
+    fill: bool = False,
+    blur: bool = False,
+) -> torch.Tensor:
     '''
     Pads an image tensor with the given number of pixels on each side and fills the padding with data from the edges.
     :param tensor: A tensor of shape [B, H, W, C]
@@ -239,7 +281,13 @@ def pad_tensor(tensor, left_pad, right_pad, top_pad, bottom_pad, fill=False, blu
     return padded
 
 
-def resize_and_pad_image(image, width, height, fill=False, blur=False):
+def resize_and_pad_image(
+    image: Image.Image,
+    width: int,
+    height: int,
+    fill: bool = False,
+    blur: bool = False,
+) -> tuple[Image.Image, tuple[int, int]]:
     '''
     Resizes an image to the given width and height and pads it to the given width and height.
     :param image: A PIL image
@@ -266,7 +314,13 @@ def resize_and_pad_image(image, width, height, fill=False, blur=False):
     return result, (horizontal_pad, vertical_pad)
 
 
-def resize_and_pad_tensor(tensor, width, height, fill=False, blur=False):
+def resize_and_pad_tensor(
+    tensor: torch.Tensor,
+    width: int,
+    height: int,
+    fill: bool = False,
+    blur: bool = False,
+) -> torch.Tensor:
     '''
     Resizes an image tensor to the given width and height and pads it to the given width and height.
     :param tensor: A tensor of shape [B, H, W, C]
@@ -294,7 +348,15 @@ def resize_and_pad_tensor(tensor, width, height, fill=False, blur=False):
     return result
 
 
-def crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
+def crop_controlnet(
+    cond_dict: dict[str, Any],
+    region: tuple[int, int, int, int],
+    init_size: tuple[int, int],
+    canvas_size: tuple[int, int],
+    tile_size: tuple[int, int],
+    w_pad: int,
+    h_pad: int,
+) -> None:
     if "control" not in cond_dict:
         return
     c = cond_dict["control"]
@@ -312,7 +374,10 @@ def crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad,
         controlnet = controlnet.previous_controlnet
 
 
-def region_intersection(region1, region2):
+def region_intersection(
+    region1: tuple[int, int, int, int],
+    region2: tuple[int, int, int, int],
+) -> tuple[int, int, int, int] | None:
     """
     Returns the coordinates of the intersection of two rectangular regions.
     :param region1: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points 
@@ -332,7 +397,15 @@ def region_intersection(region1, region2):
     return (x1, y1, x2, y2)
 
 
-def crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
+def crop_gligen(
+    cond_dict: dict[str, Any],
+    region: tuple[int, int, int, int],
+    init_size: tuple[int, int],
+    canvas_size: tuple[int, int],
+    tile_size: tuple[int, int],
+    w_pad: int,
+    h_pad: int,
+) -> None:
     if "gligen" not in cond_dict:
         return
     type, model, cond = cond_dict["gligen"]
@@ -378,7 +451,15 @@ def crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_p
     cond_dict["gligen"] = (type, model, cropped)
 
 
-def crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
+def crop_area(
+    cond_dict: dict[str, Any],
+    region: tuple[int, int, int, int],
+    init_size: tuple[int, int],
+    canvas_size: tuple[int, int],
+    tile_size: tuple[int, int],
+    w_pad: int,
+    h_pad: int,
+) -> None:
     if "area" not in cond_dict:
         return
 
@@ -412,7 +493,15 @@ def crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad
     cond_dict["area"] = (h, w, y, x)
 
 
-def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
+def crop_mask(
+    cond_dict: dict[str, Any],
+    region: tuple[int, int, int, int],
+    init_size: tuple[int, int],
+    canvas_size: tuple[int, int],
+    tile_size: tuple[int, int],
+    w_pad: int,
+    h_pad: int,
+) -> None:
     if "mask" not in cond_dict:
         return
     mask_tensor = cond_dict["mask"]  # (B, H, W)
@@ -442,7 +531,15 @@ def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad
     cond_dict["mask"] = torch.cat(masks, dim=0)  # (B, H, W)
 
 # Added Flux-Kontext Support crop_reference_latents by TBG ETUR
-def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
+def crop_reference_latents(
+    cond_dict: dict[str, Any],
+    region: tuple[int, int, int, int],
+    init_size: tuple[int, int],
+    canvas_size: tuple[int, int],
+    tile_size: tuple[int, int],
+    w_pad: int,
+    h_pad: int,
+) -> None:
     """
     1. Resize each latent to `canvas_size` in latent units.
     2. Crop the rectangle `region` (pixel coordinates).
@@ -503,7 +600,15 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size,
 
 
 
-def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0):
+def crop_cond(
+    cond: list[list[Any]],
+    region: tuple[int, int, int, int],
+    init_size: tuple[int, int],
+    canvas_size: tuple[int, int],
+    tile_size: tuple[int, int],
+    w_pad: int = 0,
+    h_pad: int = 0,
+) -> list[list[Any]]:
     cropped = []
     for emb, x in cond:
         cond_dict = x.copy()
diff --git a/utils/worker_ids.py b/utils/worker_ids.py
new file mode 100644
index 0000000..f5549f7
--- /dev/null
+++ b/utils/worker_ids.py
@@ -0,0 +1,20 @@
+import json
+from typing import Any
+
+
+def parse_enabled_worker_ids(enabled_worker_ids: str | list[Any] | None) -> list[str]:
+    """Parse and de-duplicate enabled worker IDs from a JSON list payload."""
+    try:
+        raw = json.loads(enabled_worker_ids) if isinstance(enabled_worker_ids, str) else enabled_worker_ids
+    except Exception:
+        raw = []
+
+    workers = []
+    seen = set()
+    for worker_id in raw if isinstance(raw, list) else []:
+        worker_id_str = str(worker_id)
+        if worker_id_str in seen:
+            continue
+        seen.add(worker_id_str)
+        workers.append(worker_id_str)
+    return workers
diff --git a/workers/__init__.py b/workers/__init__.py
index 4d80f2c..79f5b3e 100644
--- a/workers/__init__.py
+++ b/workers/__init__.py
@@ -1,11 +1,15 @@
 from .process_manager import WorkerProcessManager
+from functools import lru_cache
 
-_worker_manager: WorkerProcessManager | None = None
+# Keep monitor module importable via workers package for subprocess integration.
+from . import worker_monitor as worker_monitor  # noqa: F401
 
 
+@lru_cache(maxsize=1)
+def _worker_manager() -> WorkerProcessManager:
+    manager = WorkerProcessManager()
+    manager.queues = {}
+    return manager
+
 def get_worker_manager() -> WorkerProcessManager:
-    global _worker_manager
-    if _worker_manager is None:
-        _worker_manager = WorkerProcessManager()
-        _worker_manager.queues = {}
-    return _worker_manager
+    return _worker_manager()
diff --git a/workers/detection.py b/workers/detection.py
index ec98731..94acf45 100644
--- a/workers/detection.py
+++ b/workers/detection.py
@@ -1,17 +1,18 @@
 import os
 import platform
 import uuid
+from typing import Any, Mapping
 
 import aiohttp
 
-from ..utils.network import normalize_host, get_client_session
+from ..utils.network import build_worker_url, get_client_session, normalize_host
 from ..utils.logging import debug_log
 
 
-async def is_local_worker(worker_config):
+async def is_local_worker(worker_config: Mapping[str, Any]) -> bool:
     """Check if a worker is running on the same machine as the master."""
     host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
-    if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local':
+    if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local':  # nosec B104 - intentional local/wildcard handling
         return True
     
     # For cloud workers, check if on same physical host
@@ -20,19 +21,18 @@ async def is_local_worker(worker_config):
     
     return False
 
-async def is_same_physical_host(worker_config):
+async def is_same_physical_host(worker_config: Mapping[str, Any]) -> bool:
     """Compare machine IDs to determine if worker is on same physical host."""
     try:
         # Get master machine ID
         master_machine_id = get_machine_id()
         
         # Fetch worker's machine ID via API
-        host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
-        port = worker_config.get('port', 8188)
-        
+        system_info_url = build_worker_url(dict(worker_config), "/distributed/system_info")
+
         session = await get_client_session()
         async with session.get(
-            f"http://{host}:{port}/distributed/system_info",
+            system_info_url,
             timeout=aiohttp.ClientTimeout(total=5)
         ) as resp:
             if resp.status == 200:
@@ -46,7 +46,7 @@ async def is_same_physical_host(worker_config):
         debug_log(f"Error checking same physical host: {e}")
         return False
 
-def get_machine_id():
+def get_machine_id() -> str:
     """Get a unique identifier for this machine."""
     # Try multiple methods to get a stable machine ID
     try:
diff --git a/workers/process/__init__.py b/workers/process/__init__.py
index 90b040e..b685b91 100644
--- a/workers/process/__init__.py
+++ b/workers/process/__init__.py
@@ -1,11 +1 @@
-from .launch_builder import LaunchCommandBuilder
-from .lifecycle import ProcessLifecycle
-from .persistence import ProcessPersistence
-from .root_discovery import ComfyRootDiscovery
-
-__all__ = [
-    "ComfyRootDiscovery",
-    "LaunchCommandBuilder",
-    "ProcessLifecycle",
-    "ProcessPersistence",
-]
+"""Process subsystem package for worker lifecycle management."""
diff --git a/workers/process/launch_builder.py b/workers/process/launch_builder.py
index 9c03e4a..48c21d5 100644
--- a/workers/process/launch_builder.py
+++ b/workers/process/launch_builder.py
@@ -2,6 +2,7 @@ import glob
 import os
 import shlex
 import shutil
+from typing import Any
 
 from ...utils.logging import debug_log
 from ...utils.process import get_python_executable
@@ -10,7 +11,7 @@ from ...utils.process import get_python_executable
 class LaunchCommandBuilder:
     """Build command-lines for launching worker ComfyUI processes."""
 
-    def _find_windows_terminal(self):
+    def _find_windows_terminal(self) -> str | None:
         """Find Windows Terminal executable."""
         possible_paths = [
             os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\WindowsApps\wt.exe"),
@@ -31,7 +32,7 @@ class LaunchCommandBuilder:
             return wt_path
         return None
 
-    def build_launch_command(self, worker_config, comfy_root):
+    def build_launch_command(self, worker_config: dict[str, Any], comfy_root: str) -> list[str]:
         """Build the command to launch a worker."""
         main_py = os.path.join(comfy_root, "main.py")
 
diff --git a/workers/process/lifecycle.py b/workers/process/lifecycle.py
index b5a04e7..2f20ca6 100644
--- a/workers/process/lifecycle.py
+++ b/workers/process/lifecycle.py
@@ -3,11 +3,17 @@ import platform
 import signal
 import subprocess
 import time
+from typing import Any
 
 from ...utils.config import load_config, save_config
 from ...utils.constants import PROCESS_TERMINATION_TIMEOUT, PROCESS_WAIT_TIMEOUT, WORKER_CHECK_INTERVAL
 from ...utils.logging import debug_log, log
-from ...utils.process import get_python_executable, is_process_alive, terminate_process
+from ...utils.process import (
+    get_python_executable,
+    is_process_alive,
+    launch_process_with_timeout,
+    terminate_process,
+)
 
 try:
     import psutil
@@ -24,7 +30,7 @@ class ProcessLifecycle:
     def __init__(self, manager):
         self._manager = manager
 
-    def launch_worker(self, worker_config, show_window=False):
+    def launch_worker(self, worker_config: dict[str, Any], show_window: bool = False) -> int:
         """Launch a worker process with logging."""
         _ = show_window  # Kept for API compatibility.
         comfy_root = self._manager.find_comfy_root()
@@ -74,25 +80,22 @@ class ProcessLifecycle:
             else:
                 monitored_cmd = cmd
 
+            popen_kwargs = {
+                "env": env,
+                "cwd": cwd,
+                "stdout": log_handle,
+                "stderr": subprocess.STDOUT,
+            }
             if platform.system() == "Windows":
-                create_no_window = 0x08000000
-                process = subprocess.Popen(
-                    monitored_cmd,
-                    env=env,
-                    cwd=cwd,
-                    stdout=log_handle,
-                    stderr=subprocess.STDOUT,
-                    creationflags=create_no_window,
-                )
+                popen_kwargs["creationflags"] = 0x08000000
             else:
-                process = subprocess.Popen(
-                    monitored_cmd,
-                    env=env,
-                    cwd=cwd,
-                    stdout=log_handle,
-                    stderr=subprocess.STDOUT,
-                    start_new_session=True,
-                )
+                popen_kwargs["start_new_session"] = True
+
+            process = launch_process_with_timeout(
+                monitored_cmd,
+                timeout_seconds=PROCESS_WAIT_TIMEOUT,
+                **popen_kwargs,
+            )
 
         worker_id = str(worker_config["id"])
         self._manager.processes[worker_id] = {
@@ -114,7 +117,7 @@ class ProcessLifecycle:
         debug_log(f"Log file: {log_file}")
         return process.pid
 
-    def stop_worker(self, worker_id):
+    def stop_worker(self, worker_id: str | int) -> tuple[bool, str]:
         """Stop a worker process."""
         worker_id = str(worker_id)
         if worker_id not in self._manager.processes:
@@ -161,7 +164,7 @@ class ProcessLifecycle:
             log(f"[Distributed] Exception during stop: {exc}")
             return False, f"Error stopping worker: {str(exc)}"
 
-    def get_managed_workers(self):
+    def get_managed_workers(self) -> dict[str, dict[str, Any]]:
         """Get list of workers managed by this process."""
         managed = {}
         for worker_id, proc_info in list(self._manager.processes.items()):
@@ -178,7 +181,7 @@ class ProcessLifecycle:
 
         return managed
 
-    def cleanup_all(self):
+    def cleanup_all(self) -> None:
         """Stop all managed workers (called on shutdown)."""
         for worker_id in list(self._manager.processes.keys()):
             try:
@@ -190,11 +193,11 @@ class ProcessLifecycle:
         config["managed_processes"] = {}
         save_config(config)
 
-    def _is_process_running(self, pid):
+    def _is_process_running(self, pid: int) -> bool:
         """Check if a process with given PID is running."""
         return is_process_alive(pid)
 
-    def _check_worker_process(self, worker_id, proc_info):
+    def _check_worker_process(self, worker_id: str, proc_info: dict[str, Any]) -> tuple[bool, bool]:
         """Check if a worker process is still running and return status."""
         _ = worker_id  # Signature retained for compatibility with existing callers.
         process = proc_info.get("process")
@@ -222,7 +225,7 @@ class ProcessLifecycle:
                         debug_log(f"Terminating child {child.pid}")
                         child.terminate()
                     except psutil.NoSuchProcess:
-                        pass
+                        debug_log(f"Child process {child.pid} already exited before terminate")
 
                 _, alive = psutil.wait_procs(children, timeout=PROCESS_WAIT_TIMEOUT)
                 for child in alive:
@@ -230,7 +233,7 @@ class ProcessLifecycle:
                         debug_log(f"Force killing child {child.pid}")
                         child.kill()
                     except psutil.NoSuchProcess:
-                        pass
+                        debug_log(f"Child process {child.pid} already exited before kill")
 
                 try:
                     debug_log(f"Terminating parent {pid}")
@@ -255,6 +258,7 @@ class ProcessLifecycle:
                     ["wmic", "process", "where", f"ParentProcessId={pid}", "get", "ProcessId"],
                     capture_output=True,
                     text=True,
+                    timeout=PROCESS_WAIT_TIMEOUT,
                 )
                 if result.returncode == 0:
                     lines = result.stdout.strip().split("\n")[1:]
@@ -266,6 +270,7 @@ class ProcessLifecycle:
                                 ["taskkill", "/F", "/PID", child_pid],
                                 capture_output=True,
                                 check=False,
+                                timeout=PROCESS_WAIT_TIMEOUT,
                             )
                         except (FileNotFoundError, OSError) as exc:
                             debug_log(f"[Distributed] Warning: taskkill failed for PID {child_pid}: {exc}")
@@ -274,6 +279,7 @@ class ProcessLifecycle:
                     ["taskkill", "/F", "/PID", str(pid), "/T"],
                     capture_output=True,
                     text=True,
+                    timeout=PROCESS_WAIT_TIMEOUT,
                 )
                 debug_log(f"[Distributed] Taskkill result: {result.stdout.strip()}")
                 return result.returncode == 0
@@ -282,9 +288,17 @@ class ProcessLifecycle:
                 return False
 
         try:
-            subprocess.run(["pkill", "-TERM", "-P", str(pid)], check=False)
+            subprocess.run(
+                ["pkill", "-TERM", "-P", str(pid)],
+                check=False,
+                timeout=PROCESS_WAIT_TIMEOUT,
+            )
             time.sleep(WORKER_CHECK_INTERVAL)
-            subprocess.run(["pkill", "-KILL", "-P", str(pid)], check=False)
+            subprocess.run(
+                ["pkill", "-KILL", "-P", str(pid)],
+                check=False,
+                timeout=PROCESS_WAIT_TIMEOUT,
+            )
             os.kill(pid, signal.SIGKILL)
             return True
         except Exception as exc:
diff --git a/workers/process/persistence.py b/workers/process/persistence.py
index 19869b8..a0e800a 100644
--- a/workers/process/persistence.py
+++ b/workers/process/persistence.py
@@ -8,7 +8,7 @@ class ProcessPersistence:
     def __init__(self, manager):
         self._manager = manager
 
-    def load_processes(self):
+    def load_processes(self) -> None:
         """Load persisted process information from config."""
         config = load_config()
         managed_processes = config.get("managed_processes", {})
@@ -27,7 +27,7 @@ class ProcessPersistence:
             elif pid:
                 debug_log(f"[Distributed] Worker {worker_id} (PID: {pid}) is no longer running")
 
-    def save_processes(self):
+    def save_processes(self) -> None:
         """Save process information to config."""
         config = load_config()
         managed_processes = {}
diff --git a/workers/process/root_discovery.py b/workers/process/root_discovery.py
index f8a3201..51d2504 100644
--- a/workers/process/root_discovery.py
+++ b/workers/process/root_discovery.py
@@ -6,7 +6,7 @@ from ...utils.logging import debug_log, log
 class ComfyRootDiscovery:
     """Resolve the ComfyUI root directory across local and container layouts."""
 
-    def find_comfy_root(self):
+    def find_comfy_root(self) -> str:
         # Start from current file location.
         current_dir = os.path.dirname(os.path.abspath(__file__))
         potential_root = os.path.dirname(os.path.dirname(current_dir))
diff --git a/workers/process_manager.py b/workers/process_manager.py
index 823cb5b..0f170f4 100644
--- a/workers/process_manager.py
+++ b/workers/process_manager.py
@@ -1,4 +1,7 @@
-from .process import ComfyRootDiscovery, LaunchCommandBuilder, ProcessLifecycle, ProcessPersistence
+from .process.launch_builder import LaunchCommandBuilder
+from .process.lifecycle import ProcessLifecycle
+from .process.persistence import ProcessPersistence
+from .process.root_discovery import ComfyRootDiscovery
 
 
 class WorkerProcessManager:
diff --git a/workers/startup.py b/workers/startup.py
index d42d226..9b43e5d 100644
--- a/workers/startup.py
+++ b/workers/startup.py
@@ -3,7 +3,6 @@ import threading
 import time
 import atexit
 import signal
-import sys
 import platform
 
 import server
@@ -16,7 +15,7 @@ from ..utils.constants import WORKER_STARTUP_DELAY
 from . import get_worker_manager
 
 
-def auto_launch_workers():
+def auto_launch_workers() -> None:
     """Launch enabled workers if auto_launch_workers is set to true."""
     wm = get_worker_manager()
     try:
@@ -76,7 +75,7 @@ def auto_launch_workers():
         log(f"Error during auto-launch: {e}")
 
 # Schedule auto-launch after a short delay to ensure server is ready
-def delayed_auto_launch():
+def delayed_auto_launch() -> None:
     """Delay auto-launch to ensure server is fully initialized."""
     import threading
     timer = threading.Timer(WORKER_STARTUP_DELAY, auto_launch_workers)
@@ -84,9 +83,11 @@ def delayed_auto_launch():
     timer.start()
 
 # Async cleanup function for proper shutdown
-async def async_cleanup_and_exit(signum=None):
+async def async_cleanup_and_exit(signum: int | None = None) -> None:
     """Async-friendly cleanup and exit."""
+    _ = signum
     wm = get_worker_manager()
+    cleanup_error = None
     try:
         config = load_config()
         if config.get('settings', {}).get('stop_workers_on_master_exit', True):
@@ -101,28 +102,27 @@ async def async_cleanup_and_exit(signum=None):
             log(f"[Distributed] Warning: Cloudflare tunnel did not stop cleanly during shutdown: {tunnel_error}")
     except Exception as e:
         print(f"[Distributed] Error during cleanup: {e}")
+        cleanup_error = e
     
-    # On Windows, we need to exit differently
-    if platform.system() == "Windows":
-        # Force exit on Windows
-        sys.exit(0)
-    else:
-        # On Unix, stop the event loop gracefully
-        loop = asyncio.get_running_loop()
-        loop.stop()
+    # Stop the event loop gracefully on all platforms.
+    loop = asyncio.get_running_loop()
+    loop.stop()
+    if cleanup_error is not None:
+        raise cleanup_error
 
-def register_async_signals():
+def register_async_signals() -> None:
     """Register async signal handlers for graceful shutdown."""
     wm = get_worker_manager()
     if platform.system() == "Windows":
         # Windows doesn't support add_signal_handler, use traditional signal handling
-        def signal_handler(signum, frame):
+        def signal_handler(signum: int, _frame: object) -> None:
             # Schedule the async cleanup in the event loop
             loop = server.PromptServer.instance.loop
             if loop and loop.is_running():
                 asyncio.run_coroutine_threadsafe(async_cleanup_and_exit(signum), loop)
             else:
                 # Fallback to sync cleanup if loop isn't running
+                cleanup_error = None
                 try:
                     config = load_config()
                     if config.get('settings', {}).get('stop_workers_on_master_exit', True):
@@ -133,7 +133,10 @@ def register_async_signals():
                         wm.save_processes()
                 except Exception as e:
                     print(f"[Distributed] Error during cleanup: {e}")
-                sys.exit(0)
+                    cleanup_error = e
+                if cleanup_error is not None:
+                    raise RuntimeError("Cleanup failed during Windows fallback signal handling") from cleanup_error
+                raise SystemExit(0)
         
         signal.signal(signal.SIGINT, signal_handler)
         signal.signal(signal.SIGTERM, signal_handler)
@@ -147,7 +150,7 @@ def register_async_signals():
         # SIGHUP is Unix-only
         loop.add_signal_handler(signal.SIGHUP, lambda: asyncio.create_task(async_cleanup_and_exit(signal.SIGHUP)))
 
-def sync_cleanup():
+def sync_cleanup() -> None:
     """Synchronous wrapper for atexit."""
     wm = get_worker_manager()
     try:
@@ -172,3 +175,4 @@ def sync_cleanup():
             log(f"[Distributed] Warning: Cloudflare tunnel did not stop cleanly during sync cleanup: {tunnel_error}")
     except Exception as e:
         print(f"[Distributed] Error during cleanup: {e}")
+        raise
diff --git a/workers/worker_monitor.py b/workers/worker_monitor.py
index 07e95a1..3caa93e 100644
--- a/workers/worker_monitor.py
+++ b/workers/worker_monitor.py
@@ -9,40 +9,56 @@ import time
 import subprocess
 import platform
 import signal
-
-# Add package root to path so this script works when launched by file path.
-NODE_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-if NODE_ROOT not in sys.path:
-    sys.path.insert(0, NODE_ROOT)
+from typing import Any
 
 try:
-    from utils.process import is_process_alive, terminate_process
+    from utils.process import is_process_alive, launch_process_with_timeout, terminate_process
     from utils.constants import WORKER_CHECK_INTERVAL, PROCESS_TERMINATION_TIMEOUT
 except ImportError:
     # Fallback if running from different context
-    def is_process_alive(pid):
-        """Check if a process with given PID is still alive."""
+    def is_process_alive(pid: int | str) -> bool:
+        """Best-effort fallback process liveness probe."""
         try:
-            if platform.system() == "Windows":
-                # Windows: use tasklist
-                result = subprocess.run(['tasklist', '/FI', f'PID eq {pid}'], 
-                                      capture_output=True, text=True)
-                return str(pid) in result.stdout
-            else:
-                # Unix: send signal 0
-                os.kill(pid, 0)
-                return True
-        except (OSError, subprocess.SubprocessError):
+            pid = int(pid)
+        except (TypeError, ValueError):
+            return False
+
+        if os.name == "nt":
+            try:
+                output = subprocess.check_output(
+                    ["tasklist", "/FI", f"PID eq {pid}"],
+                    stderr=subprocess.DEVNULL,
+                    text=True,
+                    timeout=5.0,
+                )
+            except Exception:
+                return False
+            return str(pid) in output
+
+        try:
+            os.kill(pid, 0)
+            return True
+        except OSError:
             return False
     
+    def launch_process_with_timeout(
+        command: list[str],
+        timeout_seconds: float = 10.0,
+        **kwargs: Any,
+    ):
+        _ = kwargs
+        raise RuntimeError(
+            "launch_process_with_timeout unavailable: utils.process import failed"
+        )
+
     WORKER_CHECK_INTERVAL = 2.0
     PROCESS_TERMINATION_TIMEOUT = 5.0
 
-def monitor_and_run(master_pid, command):
-    """Run command and monitor master process."""
+def monitor_and_run(master_pid: int, command: list[str]) -> int:
+    """Run command and monitor master process; return process exit code."""
     # Start the actual worker process
     print(f"[Distributed] Launching worker command: {' '.join(command)}")
-    worker_process = subprocess.Popen(command)
+    worker_process = launch_process_with_timeout(command, timeout_seconds=10.0)
     
     print(f"[Distributed] Started worker PID: {worker_process.pid}")
     print(f"[Distributed] Monitoring master PID: {master_pid}")
@@ -56,11 +72,24 @@ def monitor_and_run(master_pid, command):
                 f.write(f"{monitor_pid},{worker_process.pid}")
             print(f"[Distributed] Wrote PID info to {pid_info_file}")
         except Exception as e:
-            print(f"[Distributed] Could not write PID file: {e}")
+            raise RuntimeError(f"Could not write PID file '{pid_info_file}': {e}") from e
     
     # Define cleanup function
-    def cleanup_worker(signum=None, frame=None):
+    should_stop = False
+    monitor_exit_code = 0
+
+    def cleanup_worker(
+        signum: int | None = None,
+        frame: Any | None = None,
+        exit_code: int = 0,
+    ) -> None:
         """Clean up worker process when monitor is terminated."""
+        nonlocal should_stop
+        nonlocal monitor_exit_code
+        _ = frame
+        should_stop = True
+        monitor_exit_code = int(exit_code)
+
         if signum:
             print(f"\n[Distributed] Received signal {signum}, terminating worker...")
         else:
@@ -80,53 +109,70 @@ def monitor_and_run(master_pid, command):
                     worker_process.wait()
         
         print("[Distributed] Worker terminated.")
-        sys.exit(0)
+
+    def _signal_handler(signum, frame):
+        cleanup_worker(signum=signum, frame=frame, exit_code=0)
     
     # Register signal handlers for graceful shutdown
-    signal.signal(signal.SIGTERM, cleanup_worker)
-    signal.signal(signal.SIGINT, cleanup_worker)
+    signal.signal(signal.SIGTERM, _signal_handler)
+    signal.signal(signal.SIGINT, _signal_handler)
     if platform.system() != "Windows":
-        signal.signal(signal.SIGHUP, cleanup_worker)
+        signal.signal(signal.SIGHUP, _signal_handler)
     
     # Monitor loop
     check_interval = WORKER_CHECK_INTERVAL
     
     try:
         while True:
+            if should_stop:
+                break
+
             # Check if worker is still running
             if worker_process.poll() is not None:
                 print(f"[Distributed] Worker process exited with code: {worker_process.returncode}")
-                sys.exit(worker_process.returncode)
+                monitor_exit_code = int(worker_process.returncode or 0)
+                break
             
             # Check if master is still alive
             if not is_process_alive(master_pid):
                 print(f"[Distributed] Master process {master_pid} is no longer running. Terminating worker...")
-                cleanup_worker()
+                cleanup_worker(exit_code=0)
+                break
             
             time.sleep(check_interval)
             
     except KeyboardInterrupt:
-        cleanup_worker()
+        cleanup_worker(exit_code=0)
+
+    return monitor_exit_code
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI entrypoint for worker monitor script."""
+    args = list(sys.argv[1:] if argv is None else argv)
 
-if __name__ == "__main__":
     # Get master PID from environment
     master_pid = os.environ.get('COMFYUI_MASTER_PID')
     if not master_pid:
         print("[Distributed] Error: COMFYUI_MASTER_PID not set")
-        sys.exit(1)
+        return 1
     
     try:
         master_pid = int(master_pid)
     except ValueError:
         print(f"[Distributed] Error: Invalid master PID: {master_pid}")
-        sys.exit(1)
+        return 1
     
     # Get the actual command to run (all remaining arguments)
-    if len(sys.argv) < 2:
+    if not args:
         print("[Distributed] Error: No command specified")
-        sys.exit(1)
+        return 1
     
-    command = sys.argv[1:]
+    command = args
     
     # Start monitoring
-    monitor_and_run(master_pid, command)
+    return monitor_and_run(master_pid, command)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())