robertvoy-ComfyUI-Distributed/distributed_upscale.py

# Import statements (original + new)
import torch
from PIL import Image, ImageFilter, ImageDraw
import json
import asyncio
import aiohttp
import io
import math
import time
from typing import List, Tuple
from functools import wraps

# Import ComfyUI modules
import comfy.samplers
import comfy.model_management

# Import shared utilities
from .utils.logging import debug_log, log
from .utils.image import tensor_to_pil, pil_to_tensor
from .utils.network import get_client_session
from .utils.async_helpers import run_async_in_server_loop
from .utils.config import get_worker_timeout_seconds
from .utils.constants import (
    TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT,
    TILE_SEND_TIMEOUT, MAX_BATCH
)

# Import for controller support
from .utils.usdu_utils import (
    crop_cond,
    get_crop_region,
    expand_crop,
)
from .utils.usdu_managment import (
    clone_conditioning, ensure_tile_jobs_initialized,
    # Job management functions
    _drain_results_queue,
    _check_and_requeue_timed_out_workers, _get_completed_count, _mark_task_completed,
    _send_heartbeat_to_master, _cleanup_job,
    # Constants
    JOB_COMPLETED_TASKS, JOB_WORKER_STATUS, JOB_PENDING_TASKS,
    MAX_PAYLOAD_SIZE
)
from .utils.usdu_managment import init_dynamic_job, init_static_job_batched


# Note: MAX_BATCH and HEARTBEAT_TIMEOUT are imported from utils.constants
# They can be overridden via environment variables:
# - COMFYUI_MAX_BATCH (default: 20)
# - COMFYUI_HEARTBEAT_TIMEOUT (default: 90)

# Sync wrapper decorator for async methods
def sync_wrapper(async_func):
    """Decorator to wrap async methods for synchronous execution."""
    @wraps(async_func)
    def sync_func(self, *args, **kwargs):
        # Use run_async_in_server_loop for ComfyUI compatibility
        return run_async_in_server_loop(
            async_func(self, *args, **kwargs),
            timeout=600.0  # 10 minute timeout for long operations
        )
    return sync_func

# Note: tensor_to_pil and pil_to_tensor are imported from utils.image


def _parse_enabled_worker_ids(enabled_worker_ids):
    """Parse enabled worker IDs from either JSON or list input."""
    if isinstance(enabled_worker_ids, list):
        return [str(worker_id) for worker_id in enabled_worker_ids]
    if not enabled_worker_ids:
        return []
    if isinstance(enabled_worker_ids, str):
        try:
            parsed = json.loads(enabled_worker_ids)
        except json.JSONDecodeError:
            log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
            return []
        if isinstance(parsed, list):
            return [str(wid) for wid in parsed]
    return []


class UltimateSDUpscaleDistributed:
    """
    Distributed version of Ultimate SD Upscale (No Upscale).
    
    Supports three processing modes:
    1. Single GPU: No workers available, process everything locally
    2. Static Mode: Small batches, distributes tiles across workers (flattened)
    3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
    
    Features:
    - Multi-mode batch handling for efficient video/image upscaling
    - Tiled VAE support for memory efficiency
    - Dynamic load balancing for large batches
    - Backward compatible with single-image workflows
    
    Environment Variables:
    - COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
    - COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
    
    Threshold: dynamic_threshold input controls mode switch (default 8)
    """

    def __init__(self):
        """Initialize the node and ensure persistent storage exists."""
        # Pre-initialize the persistent storage on node creation
        ensure_tile_jobs_initialized()
        debug_log("UltimateSDUpscaleDistributed - Node initialized")

    # WAN/FLOW detection removed per user request; enforcing 4n+1 for any batch > 1.
    
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "upscaled_image": ("IMAGE",),
                "model": ("MODEL",),
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "vae": ("VAE",),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
                "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
                "denoise": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                "tile_width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
                "tile_height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
                "padding": ("INT", {"default": 32, "min": 0, "max": 256, "step": 8}),
                "mask_blur": ("INT", {"default": 8, "min": 0, "max": 256}),
                "force_uniform_tiles": ("BOOLEAN", {"default": True}),
                "tiled_decode": ("BOOLEAN", {"default": False}),
            },
            "hidden": {
                "multi_job_id": ("STRING", {"default": ""}),
                "is_worker": ("BOOLEAN", {"default": False}),
                "master_url": ("STRING", {"default": ""}),
                "enabled_worker_ids": ("STRING", {"default": "[]"}),
                "worker_id": ("STRING", {"default": ""}),
                "tile_indices": ("STRING", {"default": ""}),  # Unused - kept for compatibility
                "dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
            },
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "run"
    CATEGORY = "image/upscaling"
    
    @classmethod
    def IS_CHANGED(cls, **kwargs):
        """Force re-execution."""
        return float("nan")  # Always re-execute
    
    
    def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, 
            sampler_name, scheduler, denoise, tile_width, tile_height, padding, 
            mask_blur, force_uniform_tiles, tiled_decode,
            multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", 
            worker_id="", tile_indices="", dynamic_threshold=8):
        """Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler."""
        # Strict WAN/FLOW batching: error if batch is not 4n+1 (except allow 1)
        try:
            batch_size = int(getattr(upscaled_image, 'shape', [1])[0])
        except Exception:
            batch_size = 1
        # Enforce 4n+1 batches globally for any model when batch > 1 (master only)
        if not is_worker and batch_size != 1 and (batch_size % 4 != 1):
            raise ValueError(
                f"Batch size {batch_size} is not of the form 4n+1. "
                "This node requires batch sizes of 1 or 4n+1 (1, 5, 9, 13, ...). "
                "Please adjust the batch size."
            )
        if not multi_job_id:
            # No distributed processing, run single GPU version
            return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
                                          seed, steps, cfg, sampler_name, scheduler, denoise,
                                          tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
        
        if is_worker:
            # Worker mode: process tiles synchronously
            return self.process_worker(upscaled_image, model, positive, negative, vae,
                                      seed, steps, cfg, sampler_name, scheduler, denoise,
                                      tile_width, tile_height, padding, mask_blur,
                                      force_uniform_tiles, tiled_decode, multi_job_id, master_url,
                                      worker_id, enabled_worker_ids, dynamic_threshold)
        else:
            # Master mode: distribute and collect synchronously
            return self.process_master(upscaled_image, model, positive, negative, vae,
                                     seed, steps, cfg, sampler_name, scheduler, denoise,
                                     tile_width, tile_height, padding, mask_blur,
                                     force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, 
                                     dynamic_threshold)
    
    def process_worker(self, upscaled_image, model, positive, negative, vae,
                      seed, steps, cfg, sampler_name, scheduler, denoise,
                      tile_width, tile_height, padding, mask_blur,
                      force_uniform_tiles, tiled_decode, multi_job_id, master_url,
                      worker_id, enabled_worker_ids, dynamic_threshold):
        """Unified worker processing - handles both static and dynamic modes."""
        # Get batch size to determine mode
        batch_size = upscaled_image.shape[0]
        
        # Ensure mode consistency across master/workers via shared threshold
        # Determine mode (must match master's logic)
        enabled_workers = json.loads(enabled_worker_ids)
        num_workers = len(enabled_workers)
        # Compute number of tiles for this image to decide if tile distribution makes sense
        _, height, width, _ = upscaled_image.shape
        all_tiles = self.calculate_tiles(width, height, self.round_to_multiple(tile_width), self.round_to_multiple(tile_height), force_uniform_tiles)
        num_tiles_per_image = len(all_tiles)

        mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
        # For USDU-style processing, we want tile distribution whenever workers are available
        # and there is more than one tile to process, even if batch == 1.
        if num_workers > 0 and num_tiles_per_image > 1:
            mode = "static"
            
        debug_log(f"USDU Dist Worker - Batch size {batch_size}")
        
        if mode == "dynamic":
            return self.process_worker_dynamic(upscaled_image, model, positive, negative, vae,
                                             seed, steps, cfg, sampler_name, scheduler, denoise,
                                             tile_width, tile_height, padding, mask_blur,
                                             force_uniform_tiles, tiled_decode, multi_job_id, master_url,
                                             worker_id, enabled_worker_ids, dynamic_threshold)
        
        # Static mode - enhanced with health monitoring and retry logic
        return self._process_worker_static_sync(upscaled_image, model, positive, negative, vae,
                                               seed, steps, cfg, sampler_name, scheduler, denoise,
                                               tile_width, tile_height, padding, mask_blur,
                                               force_uniform_tiles, tiled_decode, multi_job_id, master_url,
                                               worker_id, enabled_workers)
    
    def process_master(self, upscaled_image, model, positive, negative, vae,
                      seed, steps, cfg, sampler_name, scheduler, denoise,
                      tile_width, tile_height, padding, mask_blur,
                      force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, 
                      dynamic_threshold):
        """Unified master processing with enhanced monitoring and failure handling."""
        # Round tile dimensions
        tile_width = self.round_to_multiple(tile_width)
        tile_height = self.round_to_multiple(tile_height)
        
        # Get image dimensions and batch size
        batch_size, height, width, _ = upscaled_image.shape
        
        # Calculate all tiles and grid
        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
        num_tiles_per_image = len(all_tiles)
        rows = math.ceil(height / tile_height)
        cols = math.ceil(width / tile_width)
        log(
            f"USDU Dist: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({num_tiles_per_image} tiles/image) | Batch {batch_size}"
        )
        
        # Parse enabled workers
        enabled_workers = json.loads(enabled_worker_ids)
        num_workers = len(enabled_workers)
        
        # Determine processing mode
        mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
        # Prefer tile-based static distribution when workers are available and there are multiple tiles,
        # even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue.
        if num_workers > 0 and num_tiles_per_image > 1:
            mode = "static"
        
        log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")

        if mode == "single_gpu":
            # No workers, process all tiles locally
            return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
                                         seed, steps, cfg, sampler_name, scheduler, denoise,
                                         tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
        
        elif mode == "dynamic":
            # Dynamic mode for large batches
            return self.process_master_dynamic(upscaled_image, model, positive, negative, vae,
                                             seed, steps, cfg, sampler_name, scheduler, denoise,
                                             tile_width, tile_height, padding, mask_blur,
                                             force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers)
        
        # Static mode - enhanced with unified job management
        return self._process_master_static_sync(upscaled_image, model, positive, negative, vae,
                                               seed, steps, cfg, sampler_name, scheduler, denoise,
                                               tile_width, tile_height, padding, mask_blur,
                                               force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
                                               all_tiles, num_tiles_per_image)
    
    # Legacy static assignment helpers removed
    
    def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image,
                               model, positive, negative, vae, seed, steps, cfg,
                               sampler_name, scheduler, denoise, tile_width, tile_height,
                               padding, mask_blur, image_width, image_height, force_uniform_tiles,
                               tiled_decode, batch_idx: int = 0):
        """Process a single tile and blend it into the result image."""
        x, y = tile_pos
        
        # Extract and process tile
        tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
            upscaled_image, x, y, tile_width, tile_height, padding, force_uniform_tiles
        )
        
        processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae,
                                         seed, steps, cfg, sampler_name, 
                                         scheduler, denoise, tiled_decode, batch_idx=batch_idx,
                                         region=(x1, y1, x1 + ew, y1 + eh), image_size=(image_width, image_height))
        
        # Convert and blend
        processed_pil = tensor_to_pil(processed_tile, 0)
        # Create mask for this specific tile (no cache here; only used in single-tile path)
        tile_mask = self.create_tile_mask(image_width, image_height, x, y, tile_width, tile_height, mask_blur)
        # Use extraction position and size for blending
        result_image = self.blend_tile(result_image, processed_pil, 
                                     x1, y1, (ew, eh), tile_mask, padding)
        
        return result_image
    
    
    async def _async_collect_results(self, multi_job_id, num_workers, mode='static', 
                                   remaining_to_collect=None, batch_size=None):
        """Unified async helper to collect results from workers (tiles or images)."""
        # Get the already initialized queue
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
                raise RuntimeError(f"Job queue not initialized for {multi_job_id}")
            job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
            if not isinstance(job_data, dict) or 'mode' not in job_data:
                raise RuntimeError("Invalid job data structure")
            if job_data['mode'] != mode:
                raise RuntimeError(f"Mode mismatch: expected {mode}, got {job_data['mode']}")
            q = job_data['queue']
            
            # For dynamic mode, get reference to completed_images
            if mode == 'dynamic':
                completed_images = job_data['completed_images']
                # Calculate expected count for logging
                expected_count = remaining_to_collect or batch_size
            else:
                # Calculate expected count from job data for static mode
                expected_count = job_data.get('total_items', 0) - len(job_data.get('completed_items', set()))
        
        item_type = "images" if mode == 'dynamic' else "tiles"
        debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers")
        
        collected_results = {}
        workers_done = set()
        # Unify collector/upscaler wait behavior with the UI worker timeout
        timeout = float(get_worker_timeout_seconds())
        last_heartbeat_check = time.time()
        collected_count = 0
        
        while len(workers_done) < num_workers:
            # Check for user interruption
            if comfy.model_management.processing_interrupted():
                log("Processing interrupted by user")
                raise comfy.model_management.InterruptProcessingException()
                
            # For dynamic mode with remaining_to_collect, check if we've collected enough
            if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect:
                break
                
            try:
                # Shorter poll for dynamic mode, but never exceed the configured timeout
                wait_timeout = (min(10.0, timeout) if mode == 'dynamic' else timeout)
                result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
                worker_id = result['worker_id']
                is_last = result.get('is_last', False)
                
                if mode == 'static':
                    # Handle tiles
                    tiles = result.get('tiles', [])
                    if tiles:
                        # Batch mode
                        debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})")
                        
                        for tile_data in tiles:
                            # Validate required fields
                            if 'batch_idx' not in tile_data:
                                log(f"UltimateSDUpscale Master - Missing batch_idx in tile data, skipping")
                                continue
                            
                            tile_idx = tile_data['tile_idx']
                            # Use global_idx as key if available (for batch processing)
                            key = tile_data.get('global_idx', tile_idx)
                            
                            # Store the full tile data including metadata; prefer PIL image if present
                            entry = {
                                'tile_idx': tile_idx,
                                'x': tile_data['x'],
                                'y': tile_data['y'],
                                'extracted_width': tile_data['extracted_width'],
                                'extracted_height': tile_data['extracted_height'],
                                'padding': tile_data['padding'],
                                'worker_id': worker_id,
                                'batch_idx': tile_data.get('batch_idx', 0),
                                'global_idx': tile_data.get('global_idx', tile_idx)
                            }
                            if 'image' in tile_data:
                                entry['image'] = tile_data['image']
                            elif 'tensor' in tile_data:
                                entry['tensor'] = tile_data['tensor']
                            collected_results[key] = entry
                    else:
                        # Single tile mode (backward compat)
                        tile_idx = result['tile_idx']
                        collected_results[tile_idx] = result
                        debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})")
                
                elif mode == 'dynamic':
                    # Handle full images
                    if 'image_idx' in result and 'image' in result:
                        image_idx = result['image_idx']
                        image_pil = result['image']
                        completed_images[image_idx] = image_pil
                        collected_results[image_idx] = image_pil
                        collected_count += 1
                        debug_log(f"UltimateSDUpscale Master - Received image {image_idx} from worker {worker_id}")
                
                if is_last:
                    workers_done.add(worker_id)
                    debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
                    
            except asyncio.TimeoutError:
                if mode == 'dynamic':
                    # Check for worker timeouts periodically
                    current_time = time.time()
                    if current_time - last_heartbeat_check >= 10.0:
                        # Use the class method to check and requeue
                        requeued = await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size)
                        if requeued > 0:
                            log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers")
                        last_heartbeat_check = current_time
                    
                    # Check if we've been waiting too long overall
                    if current_time - last_heartbeat_check > timeout:
                        log(f"UltimateSDUpscale Master - Overall timeout waiting for images")
                        break
                else:
                    log(f"UltimateSDUpscale Master - Timeout waiting for {item_type}")
                    break
        
        debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers")
        
        # Clean up job queue
        async with prompt_server.distributed_tile_jobs_lock:
            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
                del prompt_server.distributed_pending_tile_jobs[multi_job_id]
        
        return collected_results if mode == 'static' else completed_images
    
    # Keep compatibility wrappers for existing code
    async def _async_collect_worker_tiles(self, multi_job_id, num_workers):
        """Async helper to collect tiles from workers."""
        return await self._async_collect_results(multi_job_id, num_workers, mode='static')
    
    async def _mark_image_completed(self, multi_job_id, image_idx, image_pil):
        """Mark an image as completed in the job data."""
        # Mark the image as completed with the image data
        await _mark_task_completed(multi_job_id, image_idx, {'image': image_pil})
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
            if job_data and 'completed_images' in job_data:
                job_data['completed_images'][image_idx] = image_pil

    async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count):
        """Collect remaining processed images from workers."""
        return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic', 
                                               remaining_to_collect=remaining_to_collect, 
                                               batch_size=batch_size)
    
    def round_to_multiple(self, value: int, multiple: int = 8) -> int:
        """Round value to nearest multiple."""
        return round(value / multiple) * multiple
    
    def calculate_tiles(self, image_width: int, image_height: int,
                       tile_width: int, tile_height: int, force_uniform_tiles: bool = True) -> List[Tuple[int, int]]:
        """Calculate tile positions to match Ultimate SD Upscale.

        Positions are a simple grid starting at (0,0) with steps of
        `tile_width` and `tile_height`, using ceil(rows/cols) to cover edges.
        Uniform vs non-uniform affects only crop/resize, not positions.
        """
        rows = math.ceil(image_height / tile_height)
        cols = math.ceil(image_width / tile_width)
        tiles: List[Tuple[int, int]] = []
        for yi in range(rows):
            for xi in range(cols):
                tiles.append((xi * tile_width, yi * tile_height))
        return tiles
    
    def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int,
                                 tile_width: int, tile_height: int, padding: int,
                                 force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
        """Extract a tile region and resize to match USDU cropping logic.

        Mirrors ComfyUI_UltimateSDUpscale processing:
        - Build a mask with a white rectangle at the tile rect
        - Compute crop_region via get_crop_region(mask, padding)
        - If force_uniform_tiles: shrink/expand to target size of
          round_to_multiple(tile + padding) for each dimension
        - Else: target is ceil(crop_size/8)*8 per dimension
        - Extract the crop and resize to target tile_size
        Returns the resized tensor and crop origin/size for blending.
        """
        _, h, w, _ = image.shape

        # Create mask and compute initial padded crop region
        mask = Image.new('L', (w, h), 0)
        draw = ImageDraw.Draw(mask)
        draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
        x1, y1, x2, y2 = get_crop_region(mask, padding)

        # Determine target tile size (processing size)
        if force_uniform_tiles:
            target_w = self.round_to_multiple(tile_width + padding, 8)
            target_h = self.round_to_multiple(tile_height + padding, 8)
            (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
        else:
            crop_w = x2 - x1
            crop_h = y2 - y1
            target_w = max(8, math.ceil(crop_w / 8) * 8)
            target_h = max(8, math.ceil(crop_h / 8) * 8)
            (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)

        # Actual extracted size before resizing (for blending)
        extracted_width = x2 - x1
        extracted_height = y2 - y1

        # Extract tile and resize to processing size
        tile = image[:, y1:y2, x1:x2, :]
        tile_pil = tensor_to_pil(tile, 0)
        if tile_pil.size != (target_w, target_h):
            tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)

        tile_tensor = pil_to_tensor(tile_pil)
        if image.is_cuda:
            tile_tensor = tile_tensor.cuda()

        return tile_tensor, x1, y1, extracted_width, extracted_height

    def extract_batch_tile_with_padding(self, images: torch.Tensor, x: int, y: int,
                                        tile_width: int, tile_height: int, padding: int,
                                        force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
        """Extract a tile region for the entire batch and resize to USDU logic.

        - Computes a single crop region from a mask at (x,y,w,h) with padding
        - force_uniform_tiles controls target processing size logic
        - Returns a batched tensor [B,H',W',C] and crop origin/size for blending
        """
        batch, h, w, _ = images.shape

        # Create mask and compute initial padded crop region (same for all images)
        mask = Image.new('L', (w, h), 0)
        draw = ImageDraw.Draw(mask)
        draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
        x1, y1, x2, y2 = get_crop_region(mask, padding)

        # Determine target processing size
        if force_uniform_tiles:
            target_w = self.round_to_multiple(tile_width + padding, 8)
            target_h = self.round_to_multiple(tile_height + padding, 8)
            (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
        else:
            crop_w = x2 - x1
            crop_h = y2 - y1
            target_w = max(8, math.ceil(crop_w / 8) * 8)
            target_h = max(8, math.ceil(crop_h / 8) * 8)
            (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)

        extracted_width = x2 - x1
        extracted_height = y2 - y1

        # Slice batch region
        tiles = images[:, y1:y2, x1:x2, :]

        # Resize each tile to target size
        resized_tiles = []
        for i in range(batch):
            tile_pil = tensor_to_pil(tiles, i)
            if tile_pil.size != (target_w, target_h):
                tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)
            resized_tiles.append(pil_to_tensor(tile_pil))
        tile_batch = torch.cat(resized_tiles, dim=0)

        if images.is_cuda:
            tile_batch = tile_batch.cuda()

        return tile_batch, x1, y1, extracted_width, extracted_height
    
    def process_tile(self, tile_tensor: torch.Tensor, model, positive, negative, vae,
                     seed: int, steps: int, cfg: float, sampler_name: str, 
                     scheduler: str, denoise: float, tiled_decode: bool = False,
                     batch_idx: int = 0, region: Tuple[int, int, int, int] = None,
                     image_size: Tuple[int, int] = None) -> torch.Tensor:
        """Process a single tile through SD sampling. 
        Note: positive and negative should already be pre-sliced for the current batch_idx."""
        debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}")
        
        
        # Import here to avoid circular dependencies
        from nodes import common_ksampler, VAEEncode, VAEDecode
        
        # Try to import tiled VAE nodes if available
        try:
            from nodes import VAEEncodeTiled, VAEDecodeTiled
            tiled_vae_available = True
        except ImportError:
            tiled_vae_available = False
            if tiled_decode:
                debug_log("Tiled VAE nodes not available, falling back to standard VAE")
        
        # Convert to PIL and back to ensure clean tensor without gradient tracking
        tile_pil = tensor_to_pil(tile_tensor, 0)
        clean_tensor = pil_to_tensor(tile_pil)
        
        # Ensure tensor is detached and doesn't require gradients
        clean_tensor = clean_tensor.detach()
        if hasattr(clean_tensor, 'requires_grad_'):
            clean_tensor.requires_grad_(False)
        
        # Move to correct device
        if tile_tensor.is_cuda:
            clean_tensor = clean_tensor.cuda()
            clean_tensor = clean_tensor.detach()  # Detach again after device transfer
        
        # Clone conditioning per tile (shares models, clones hints for cropping)
        positive_tile = clone_conditioning(positive, clone_hints=True)
        negative_tile = clone_conditioning(negative, clone_hints=True)
        
        # Crop conditioning to tile region if provided (assumes hints at image resolution)
        if region is not None and image_size is not None:
            init_size = image_size  # (width, height) of full image
            canvas_size = image_size
            tile_size = (tile_tensor.shape[2], tile_tensor.shape[1])  # (width, height)
            w_pad = 0  # No extra pad needed; region already includes padding
            h_pad = 0
            positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
            negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
        else:
            # No region cropping needed, use cloned conditioning as-is
            positive_cropped = positive_tile
            negative_cropped = negative_tile
        
        # Encode to latent (always non-tiled, matching original node)
        latent = VAEEncode().encode(vae, clean_tensor)[0]
        
        # Sample
        samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, 
                                positive_cropped, negative_cropped, latent, denoise=denoise)[0]
        
        # Decode back to image
        if tiled_decode and tiled_vae_available:
            image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
        else:
            image = VAEDecode().decode(vae, samples)[0]
        
        return image

    def process_tiles_batch(self, tile_batch: torch.Tensor, model, positive, negative, vae,
                            seed: int, steps: int, cfg: float, sampler_name: str,
                            scheduler: str, denoise: float, tiled_decode: bool,
                            region: Tuple[int, int, int, int], image_size: Tuple[int, int]) -> torch.Tensor:
        """Process a batch of tiles together (USDU behavior).

        tile_batch: [B, H, W, C]
        Returns image batch tensor [B, H, W, C]
        """
        # Import locally to avoid circular deps
        from nodes import common_ksampler, VAEEncode, VAEDecode
        try:
            from nodes import VAEEncodeTiled, VAEDecodeTiled
            tiled_vae_available = True
        except ImportError:
            tiled_vae_available = False

        # Detach and move device
        clean = tile_batch.detach()
        if hasattr(clean, 'requires_grad_'):
            clean.requires_grad_(False)
        if tile_batch.is_cuda:
            clean = clean.cuda().detach()

        # Clone/crop conditioning once for the region
        positive_tile = clone_conditioning(positive, clone_hints=True)
        negative_tile = clone_conditioning(negative, clone_hints=True)

        init_size = image_size
        canvas_size = image_size
        tile_size = (clean.shape[2], clean.shape[1])  # (W,H)
        w_pad = 0
        h_pad = 0
        positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
        negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)

        # Encode -> Sample -> Decode
        latent = VAEEncode().encode(vae, clean)[0]
        samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler,
                                  positive_cropped, negative_cropped, latent, denoise=denoise)[0]
        if tiled_decode and tiled_vae_available:
            image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
        else:
            image = VAEDecode().decode(vae, samples)[0]

        return image
    
    def create_tile_mask(self, image_width: int, image_height: int,
                        x: int, y: int, tile_width: int, tile_height: int, 
                        mask_blur: int) -> Image.Image:
        """Create a mask for blending tiles - matches Ultimate SD Upscale approach.
        
        Creates a black image with a white rectangle at the tile position,
        then applies blur to create soft edges.
        """
        # Create a full-size mask matching the image dimensions
        mask = Image.new('L', (image_width, image_height), 0)  # Black background
        
        # Draw white rectangle at tile position
        draw = ImageDraw.Draw(mask)
        draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
        
        # Apply blur to soften edges
        if mask_blur > 0:
            mask = mask.filter(ImageFilter.GaussianBlur(mask_blur))
        
        return mask
    
    def blend_tile(self, base_image: Image.Image, tile_image: Image.Image,
                  x: int, y: int, extracted_size: Tuple[int, int], 
                  mask: Image.Image, padding: int) -> Image.Image:
        """Blend a processed tile back into the base image using Ultimate SD Upscale's exact approach.
        
        This follows the exact method from ComfyUI_UltimateSDUpscale/modules/processing.py
        """
        extracted_width, extracted_height = extracted_size
        
        # Debug logging (uncomment if needed)
        # debug_log(f"[Blend] Placing tile at ({x}, {y}), size: {extracted_width}x{extracted_height}")
        
        # Calculate the crop region that was used for extraction
        crop_region = (x, y, x + extracted_width, y + extracted_height)
        
        # The mask is already full-size, no need to crop
        
        # Resize the processed tile back to the extracted size
        if tile_image.size != (extracted_width, extracted_height):
            tile_resized = tile_image.resize((extracted_width, extracted_height), Image.LANCZOS)
        else:
            tile_resized = tile_image
        
        # Follow Ultimate SD Upscale blending approach:
        # Put the tile into position
        image_tile_only = Image.new('RGBA', base_image.size)
        image_tile_only.paste(tile_resized, crop_region[:2])
        
        # Add the mask as an alpha channel
        # Must make a copy due to the possibility of an edge becoming black
        temp = image_tile_only.copy()
        temp.putalpha(mask)  # Use the full image mask
        image_tile_only.paste(temp, image_tile_only)
        
        # Add back the tile to the initial image according to the mask in the alpha channel
        result = base_image.convert('RGBA')
        result.alpha_composite(image_tile_only)
        
        # Convert back to RGB
        return result.convert('RGB')
    
    def _slice_conditioning(self, positive, negative, batch_idx):
        """Helper to slice conditioning for a specific batch index."""
        # Clone and slice conditioning properly, including ControlNet hints
        positive_sliced = clone_conditioning(positive)
        negative_sliced = clone_conditioning(negative)
        
        for cond_list in [positive_sliced, negative_sliced]:
            for i in range(len(cond_list)):
                emb, cond_dict = cond_list[i]
                if emb.shape[0] > 1:
                    cond_list[i][0] = emb[batch_idx:batch_idx+1]
                if 'control' in cond_dict:
                    control = cond_dict['control']
                    while control is not None:
                        hint = control.cond_hint_original
                        if hint.shape[0] > 1:
                            control.cond_hint_original = hint[batch_idx:batch_idx+1]
                        control = control.previous_controlnet
                if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
                    cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1]
        
        return positive_sliced, negative_sliced
    
    async def _get_all_completed_tasks(self, multi_job_id):
        """Helper to retrieve all completed tasks from the job data."""
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
            if job_data and JOB_COMPLETED_TASKS in job_data:
                return dict(job_data[JOB_COMPLETED_TASKS])  # Return a copy
            return {}
    
    # Note: Removed unused _process_worker_static_async to reduce redundancy
    
    def _process_worker_static_sync(self, upscaled_image, model, positive, negative, vae,
                                    seed, steps, cfg, sampler_name, scheduler, denoise,
                                    tile_width, tile_height, padding, mask_blur,
                                    force_uniform_tiles, tiled_decode, multi_job_id, master_url,
                                    worker_id, enabled_workers):
        """Worker static mode processing with optional dynamic queue pulling."""
        # Round tile dimensions
        tile_width = self.round_to_multiple(tile_width)
        tile_height = self.round_to_multiple(tile_height)
        
        # Get dimensions and calculate tiles
        _, height, width, _ = upscaled_image.shape
        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
        num_tiles_per_image = len(all_tiles)
        batch_size = upscaled_image.shape[0]
        total_tiles = batch_size * num_tiles_per_image
        
        processed_tiles = []
        sliced_conditioning_cache = {}
        
        # Dynamic queue mode (static processing): process batched-per-tile
        log(f"USDU Dist Worker[{worker_id[:8]}]: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Tiles/image {num_tiles_per_image} | Batch {batch_size}")
        processed_count = 0

        # Poll for job readiness
        max_poll_attempts = 20
        for attempt in range(max_poll_attempts):
            ready = run_async_in_server_loop(
                self._check_job_status(multi_job_id, master_url),
                timeout=5.0
            )
            if ready:
                debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
                break
            time.sleep(1.0)
        else:
            log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
            return (upscaled_image,)

        # Main processing loop - pull tile ids from queue
        while True:
            # Request a tile to process
            tile_idx, estimated_remaining, batched_static = run_async_in_server_loop(
                self._request_tile_from_master(multi_job_id, master_url, worker_id),
                timeout=TILE_WAIT_TIMEOUT
            )

            if tile_idx is None:
                debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process")
                break

            # Always batched-per-tile in static mode
            debug_log(f"Worker[{worker_id[:8]}] - Assigned tile_id {tile_idx}")
            processed_count += batch_size
            tile_id = tile_idx
            tx, ty = all_tiles[tile_id]
            # Extract tile for entire batch
            tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
                upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
            )
            # Process batch
            region = (x1, y1, x1 + ew, y1 + eh)
            processed_batch = self.process_tiles_batch(
                tile_batch, model, positive, negative, vae,
                seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
                region, (width, height)
            )
            # Queue results
            for b in range(batch_size):
                processed_tiles.append({
                    'tile': processed_batch[b:b+1],
                    'tile_idx': tile_id,
                    'x': x1,
                    'y': y1,
                    'extracted_width': ew,
                    'extracted_height': eh,
                    'padding': padding,
                    'batch_idx': b,
                    'global_idx': b * num_tiles_per_image + tile_id
                })

            # Send heartbeat
            try:
                run_async_in_server_loop(
                    _send_heartbeat_to_master(multi_job_id, master_url, worker_id),
                    timeout=5.0
                )
            except Exception as e:
                debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}")

            # Send tiles in batches within loop
            if len(processed_tiles) >= MAX_BATCH:
                run_async_in_server_loop(
                    self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id),
                    timeout=TILE_SEND_TIMEOUT
                )
                processed_tiles = []

        # Send any remaining tiles
        if processed_tiles:
            run_async_in_server_loop(
                self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id),
                timeout=TILE_SEND_TIMEOUT
            )
        
        debug_log(f"Worker {worker_id} completed all assigned and requeued tiles")
        return (upscaled_image,)
    
    async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total):
        """Async helper for collection and monitoring in static mode.
        Returns collected tasks dict. Caller should check if all tasks are complete."""
        last_progress_log = time.time()
        progress_interval = 5.0
        last_heartbeat_check = time.time()
        last_completed_count = 0
        
        while True:
            # Check for user interruption
            if comfy.model_management.processing_interrupted():
                log("Processing interrupted by user")
                raise comfy.model_management.InterruptProcessingException()
            
            # Drain any pending results
            collected_count = await _drain_results_queue(multi_job_id)
            
            # Check and requeue timed-out workers periodically
            current_time = time.time()
            if current_time - last_heartbeat_check >= 10.0:
                requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total)
                if requeued_count > 0:
                    log(f"Requeued {requeued_count} tasks from timed-out workers")
                last_heartbeat_check = current_time
            
            # Get current completion count
            completed_count = await _get_completed_count(multi_job_id)
            
            # Progress logging
            if current_time - last_progress_log >= progress_interval:
                log(f"Progress: {completed_count}/{expected_total} tasks completed")
                last_progress_log = current_time
            
            # Check if all tasks are completed
            if completed_count >= expected_total:
                debug_log(f"All {expected_total} tasks completed")
                break
            
            # If no active workers remain and there are pending tasks, return for local processing
            prompt_server = ensure_tile_jobs_initialized()
            async with prompt_server.distributed_tile_jobs_lock:
                job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
                if job_data:
                    pending_queue = job_data.get(JOB_PENDING_TASKS)
                    active_workers = list(job_data.get(JOB_WORKER_STATUS, {}).keys())
                    if pending_queue and not pending_queue.empty() and len(active_workers) == 0:
                        log(f"No active workers remaining with {expected_total - completed_count} tasks pending. Returning for local processing.")
                        break
            
            # Wait a bit before next check
            await asyncio.sleep(0.1)
        
        # Get all completed tasks for return
        return await self._get_all_completed_tasks(multi_job_id)
    
    def _process_master_static_sync(self, upscaled_image, model, positive, negative, vae,
                                    seed, steps, cfg, sampler_name, scheduler, denoise,
                                    tile_width, tile_height, padding, mask_blur,
                                    force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
                                    all_tiles, num_tiles_per_image):
        """Static mode master processing with optional dynamic queue pulling."""
        batch_size = upscaled_image.shape[0]
        _, height, width, _ = upscaled_image.shape
        total_tiles = batch_size * num_tiles_per_image
        
        # Convert batch to PIL list for processing
        result_images = []
        for b in range(batch_size):
            image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
            result_images.append(image_pil.copy())
        
        sliced_conditioning_cache = {}
        # Initialize queue: pending queue holds tile ids (batched per tile)
        log("USDU Dist: Using tile queue distribution")
        run_async_in_server_loop(
            init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers),
            timeout=10.0
        )
        debug_log(
            f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}"
        )

        # Precompute masks for all tile positions to avoid repeated Gaussian blur work during blending
        tile_masks = []
        for idx, (tx, ty) in enumerate(all_tiles):
            tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))

        processed_count = 0
        consecutive_no_tile = 0
        max_consecutive_no_tile = 2

        while processed_count < total_tiles:
            comfy.model_management.throw_exception_if_processing_interrupted()
            tile_idx = run_async_in_server_loop(
                self._get_next_tile_index(multi_job_id),
                timeout=5.0
            )
            if tile_idx is not None:
                consecutive_no_tile = 0
                processed_count += batch_size
                tile_id = tile_idx
                tx, ty = all_tiles[tile_id]
                tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
                    upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
                )
                region = (x1, y1, x1 + ew, y1 + eh)
                processed_batch = self.process_tiles_batch(
                    tile_batch, model, positive, negative, vae,
                    seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
                    region, (width, height)
                )
                tile_mask = tile_masks[tile_id]
                for b in range(batch_size):
                    tile_pil = tensor_to_pil(processed_batch, b)
                    if tile_pil.size != (ew, eh):
                        tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
                    result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)
                    global_idx = b * num_tiles_per_image + tile_id
                    run_async_in_server_loop(
                        _mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}),
                        timeout=5.0
                    )
                log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})")
            else:
                consecutive_no_tile += 1
                if consecutive_no_tile >= max_consecutive_no_tile:
                    debug_log(f"Master processed {processed_count} tiles, moving to collection phase")
                    break
                time.sleep(0.1)
        master_processed_count = processed_count
        
        # Continue processing any remaining tiles while collecting worker results
        remaining_tiles = total_tiles - master_processed_count
        if remaining_tiles > 0:
            debug_log(f"Master waiting for {remaining_tiles} tiles from workers")
            
            # Collect worker results using async operations
            try:
                # Wait until either all tasks are collected or there are no active workers left
                collected_tasks = run_async_in_server_loop(
                    self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles),
                    timeout=None
                )
            except comfy.model_management.InterruptProcessingException:
                # Clean up job on interruption
                run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
                raise
            
            # Check if we need to process any remaining tasks locally after collection
            completed_count = len(collected_tasks)
            if completed_count < total_tiles:
                log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures")
                
                # Process any remaining pending tasks (batched-per-tile)
                while True:
                    # Check for user interruption
                    comfy.model_management.throw_exception_if_processing_interrupted()

                    # Get next tile_id from pending queue
                    tile_id = run_async_in_server_loop(
                        self._get_next_tile_index(multi_job_id),
                        timeout=5.0
                    )

                    if tile_id is None:
                        break

                    # Extract batched tile and process across available batch
                    tx, ty = all_tiles[tile_id]
                    tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
                        upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
                    )
                    region = (x1, y1, x1 + ew, y1 + eh)
                    processed_batch = self.process_tiles_batch(
                        tile_batch, model, positive, negative, vae,
                        seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
                        region, (width, height)
                    )
                    tile_mask = tile_masks[tile_id]
                    out_bs = processed_batch.shape[0] if hasattr(processed_batch, 'shape') else batch_size
                    for b in range(min(batch_size, out_bs)):
                        tile_pil = tensor_to_pil(processed_batch, b)
                        if tile_pil.size != (ew, eh):
                            tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
                        result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)
                        global_idx = b * num_tiles_per_image + tile_id
                        # Mark as completed so the collector state is consistent
                        run_async_in_server_loop(
                            _mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}),
                            timeout=5.0
                        )
        else:
            # Master processed all tiles
            collected_tasks = run_async_in_server_loop(
                self._get_all_completed_tasks(multi_job_id),
                timeout=5.0
            )
        
        # Blend worker tiles synchronously
        for global_idx, tile_data in collected_tasks.items():
            # Skip tiles that don't have tensor data (already processed)
            if 'tensor' not in tile_data and 'image' not in tile_data:
                continue
            
            batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image)
            tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image)
            
            if batch_idx >= batch_size:
                continue
            
            # Blend tile synchronously
            x = tile_data.get('x', 0)
            y = tile_data.get('y', 0)
            # Prefer PIL image if present to avoid reconversion
            if 'image' in tile_data:
                tile_pil = tile_data['image']
            else:
                tile_tensor = tile_data['tensor']
                tile_pil = tensor_to_pil(tile_tensor, 0)
            orig_x, orig_y = all_tiles[tile_idx]
            tile_mask = tile_masks[tile_idx]
            extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding)
            extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding)
            result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil,
                                                      x, y, (extracted_width, extracted_height), tile_mask, padding)
        
        try:
            # Convert back to tensor
            if batch_size == 1:
                result_tensor = pil_to_tensor(result_images[0])
            else:
                result_tensors = [pil_to_tensor(img) for img in result_images]
                result_tensor = torch.cat(result_tensors, dim=0)
            
            if upscaled_image.is_cuda:
                result_tensor = result_tensor.cuda()
            
            log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
            return (result_tensor,)
        finally:
            # Cleanup (async operation) - always execute
            run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
    
    def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles,
                                  model, positive, negative, vae, seed, steps, cfg, sampler_name,
                                  scheduler, denoise, tiled_decode, tile_width, tile_height, padding,
                                  width, height, force_uniform_tiles, sliced_conditioning_cache):
        """Process a single tile."""
        # Calculate which image and tile this corresponds to
        batch_idx = global_idx // num_tiles_per_image
        tile_idx = global_idx % num_tiles_per_image
        
        # Skip if batch_idx is out of range
        if batch_idx >= upscaled_image.shape[0]:
            debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}")
            return None
        
        # Get or create sliced conditioning for this batch index
        if batch_idx not in sliced_conditioning_cache:
            positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx)
            sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced)
        else:
            positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx]
        
        x, y = all_tiles[tile_idx]
        
        # Extract tile from the specific image in the batch
        tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
            upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding, force_uniform_tiles
        )
        
        # Process tile through SD with unique seed
        image_seed = seed + batch_idx * 1000
        processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae,
                                         image_seed, steps, cfg, sampler_name,
                                         scheduler, denoise, tiled_decode, batch_idx=batch_idx,
                                         region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height))
        
        return {
            'tile': processed_tile,
            'global_idx': global_idx,
            'batch_idx': batch_idx,
            'tile_idx': tile_idx,
            'x': x1,
            'y': y1,
            'extracted_width': ew,
            'extracted_height': eh
        }
    
    async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, 
                                       padding, worker_id):
        """Send all processed tiles to master, chunked if large."""
        if not processed_tiles:
            return  # Early exit if empty

        total_tiles = len(processed_tiles)
        debug_log(f"Worker[{worker_id[:8]}] - Preparing to send {total_tiles} tiles (size-aware chunks)")

        # Prepare encoded images and sizes to enable size-aware chunking
        encoded = []
        for idx, tile_data in enumerate(processed_tiles):
            img = tensor_to_pil(tile_data['tile'], 0)
            bio = io.BytesIO()
            # Keep compression low to balance speed and size; adjust if needed
            img.save(bio, format='PNG', compress_level=0)
            raw = bio.getvalue()
            encoded.append({
                'bytes': raw,
                'meta': {
                    'tile_idx': tile_data['tile_idx'],
                    'x': tile_data['x'],
                    'y': tile_data['y'],
                    'extracted_width': tile_data['extracted_width'],
                    'extracted_height': tile_data['extracted_height'],
                    **({'batch_idx': tile_data['batch_idx']} if 'batch_idx' in tile_data else {}),
                    **({'global_idx': tile_data['global_idx']} if 'global_idx' in tile_data else {}),
                }
            })

        # Size-aware chunking
        max_bytes = int(MAX_PAYLOAD_SIZE) - (1024 * 1024)  # 1MB headroom
        i = 0
        chunk_index = 0
        while i < total_tiles:
            data = aiohttp.FormData()
            data.add_field('multi_job_id', multi_job_id)
            data.add_field('worker_id', str(worker_id))
            data.add_field('padding', str(padding))

            metadata = []
            used = 0
            j = i
            while j < total_tiles:
                img_bytes = encoded[j]['bytes']
                meta = encoded[j]['meta']
                # Rough overhead for fields + JSON
                overhead = 1024
                if used + len(img_bytes) + overhead > max_bytes and j > i:
                    break
                # Accept this tile in this chunk
                metadata.append(meta)
                data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png')
                used += len(img_bytes) + overhead
                j += 1

            # Ensure at least one tile per chunk
            if j == i:
                # Single oversized tile, send anyway
                meta = encoded[j]['meta']
                metadata.append(meta)
                data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png')
                j += 1

            chunk_size = j - i
            is_chunk_last = (j >= total_tiles)
            data.add_field('is_last', str(is_chunk_last))
            data.add_field('batch_size', str(chunk_size))
            data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')

            # Retry logic with exponential backoff
            max_retries = 5
            retry_delay = 0.5
            for attempt in range(max_retries):
                try:
                    session = await get_client_session()
                    url = f"{master_url}/distributed/submit_tiles"
                    async with session.post(url, data=data) as response:
                        response.raise_for_status()
                        break
                except Exception as e:
                    if attempt < max_retries - 1:
                        await asyncio.sleep(retry_delay)
                        retry_delay = min(retry_delay * 2, 5.0)
                    else:
                        log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}")
                        raise

            debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)")
            chunk_index += 1
            i = j

    def process_single_gpu(self, upscaled_image, model, positive, negative, vae,
                          seed, steps, cfg, sampler_name, scheduler, denoise,
                          tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode):
        """Process all tiles on a single GPU (no distribution), batching per tile like USDU."""
        # Round tile dimensions
        tile_width = self.round_to_multiple(tile_width)
        tile_height = self.round_to_multiple(tile_height)

        # Get image dimensions and batch size
        batch_size, height, width, _ = upscaled_image.shape

        # Calculate all tiles
        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)

        rows = math.ceil(height / tile_height)
        cols = math.ceil(width / tile_width)
        log(
            f"USDU Dist: Single GPU | Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({len(all_tiles)} tiles/image) | Batch {batch_size}"
        )

        # Prepare result images list
        result_images = []
        for b in range(batch_size):
            image_pil = tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB')
            result_images.append(image_pil.copy())

        # Precompute tile masks once
        tile_masks = []
        for tx, ty in all_tiles:
            tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))

        # Process tiles batched across images
        for tile_idx, (tx, ty) in enumerate(all_tiles):
            # Extract batched tile
            tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
                upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
            )

            # Process batch
            region = (x1, y1, x1 + ew, y1 + eh)
            processed_batch = self.process_tiles_batch(tile_batch, model, positive, negative, vae,
                                                       seed, steps, cfg, sampler_name, scheduler, denoise,
                                                       tiled_decode, region, (width, height))

            # Blend results back into each image using cached mask
            tile_mask = tile_masks[tile_idx]
            for b in range(batch_size):
                tile_pil = tensor_to_pil(processed_batch, b)
                # Resize back to extracted size
                if tile_pil.size != (ew, eh):
                    tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
                result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)

        # Convert back to tensor
        result_tensors = [pil_to_tensor(img) for img in result_images]
        result_tensor = torch.cat(result_tensors, dim=0)
        if upscaled_image.is_cuda:
            result_tensor = result_tensor.cuda()

        return (result_tensor,)
    
    def process_master_dynamic(self, upscaled_image, model, positive, negative, vae,
                              seed, steps, cfg, sampler_name, scheduler, denoise,
                              tile_width, tile_height, padding, mask_blur,
                              force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers):
        """Dynamic mode for large batches - assigns whole images to workers dynamically, including master."""
        # Get batch size and dimensions
        batch_size, height, width, _ = upscaled_image.shape
        num_workers = len(enabled_workers)
        
        log(f"USDU Dist: Image queue distribution | Batch {batch_size} | Workers {num_workers} | Canvas {width}x{height} | Tile {tile_width}x{tile_height}")

        # No fixed share - all images are dynamic
        all_indices = list(range(batch_size))
        
        debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.")
        
        # Calculate tiles for processing
        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
        
        # Initialize job queue for communication
        try:
            run_async_in_server_loop(
                init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices),
                timeout=2.0
            )
        except Exception as e:
            debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}")
            raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}")
        
        # Convert batch to PIL list
        result_images = [tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB').copy() for b in range(batch_size)]
        
        # Process images dynamically with master participating
        prompt_server = ensure_tile_jobs_initialized()
        processed_count = 0
        consecutive_retries = 0
        max_consecutive_retries = 10
        
        # Process loop - master pulls from queue and processes synchronously
        while processed_count < batch_size:
            # Try to get an image to process
            image_idx = run_async_in_server_loop(
                self._get_next_image_index(multi_job_id),
                timeout=5.0  # Short timeout to allow frequent checks
            )

            if image_idx is not None:
                # Reset retry counter and process locally
                consecutive_retries = 0
                debug_log(f"Master processing image {image_idx} dynamically")
                processed_count += 1

                # Process locally
                single_tensor = upscaled_image[image_idx:image_idx+1]
                local_image = result_images[image_idx]
                image_seed = seed + image_idx * 1000
                
                # Pre-slice conditioning once per image (not per tile)
                positive_sliced = clone_conditioning(positive)
                negative_sliced = clone_conditioning(negative)
                for cond_list in [positive_sliced, negative_sliced]:
                    for i in range(len(cond_list)):
                        emb, cond_dict = cond_list[i]
                        if emb.shape[0] > 1:
                            cond_list[i][0] = emb[image_idx:image_idx+1]
                        if 'control' in cond_dict:
                            control = cond_dict['control']
                            while control is not None:
                                hint = control.cond_hint_original
                                if hint.shape[0] > 1:
                                    control.cond_hint_original = hint[image_idx:image_idx+1]
                                control = control.previous_controlnet
                        if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
                            cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1]
                
                for tile_idx, pos in enumerate(all_tiles):
                    local_image = self._process_and_blend_tile(
                        tile_idx, pos, single_tensor, local_image,
                        model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
                        sampler_name, scheduler, denoise, tile_width, tile_height,
                        padding, mask_blur, width, height, force_uniform_tiles,
                        tiled_decode, batch_idx=image_idx
                    )
                    
                    # Yield after each tile to minimize worker downtime
                    run_async_in_server_loop(self._async_yield(), timeout=0.1)
                    # Note: No per-tile drain here – that's what makes this "per-image"
                
                result_images[image_idx] = local_image
                
                # Mark as completed
                run_async_in_server_loop(
                    self._mark_image_completed(multi_job_id, image_idx, local_image),
                    timeout=5.0
                )
                
                # NEW: Drain after the full image is marked complete (catches workers who finished during master's processing)
                drained_count = run_async_in_server_loop(
                    self._drain_worker_results_queue(multi_job_id),
                    timeout=5.0
                )
                
                if drained_count > 0:
                    debug_log(f"Drained {drained_count} worker images after master's image {image_idx}")
                
                # NEW: Log overall progress (includes master's image + any drained workers)
                completed_now = run_async_in_server_loop(
                    self._get_total_completed_count(multi_job_id),
                    timeout=1.0
                )
                log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
                
                # Yield to allow workers to get new images after completing one
                run_async_in_server_loop(self._async_yield(), timeout=0.1)
            else:
                # Queue empty: collect any queued worker results to update progress
                drained_count = run_async_in_server_loop(
                    self._drain_worker_results_queue(multi_job_id),
                    timeout=5.0
                )
                run_async_in_server_loop(self._async_yield(), timeout=0.1)  # Yield after drain
                
                # Check for timed out workers and requeue their images
                requeued_count = run_async_in_server_loop(
                    self._check_and_requeue_timed_out_workers(multi_job_id, batch_size),
                    timeout=5.0
                )
                run_async_in_server_loop(self._async_yield(), timeout=0.1)  # Yield after requeue
                
                if requeued_count > 0:
                    log(f"Requeued {requeued_count} images from timed out workers")
                    consecutive_retries = 0  # Reset since we have work to do
                    continue

                # Now check total completed (includes newly collected)
                completed_now = run_async_in_server_loop(
                    self._get_total_completed_count(multi_job_id),
                    timeout=1.0
                )
                
                log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
                
                if completed_now >= batch_size:
                    break

                run_async_in_server_loop(self._async_yield(), timeout=0.1)  # Yield before pending check
                
                # Check if there are pending images in the queue (could be requeued)
                pending_count = run_async_in_server_loop(
                    self._get_pending_count(multi_job_id),
                    timeout=1.0
                )
                
                if pending_count > 0:
                    consecutive_retries = 0  # Reset retries since there's work to do
                    continue

                consecutive_retries += 1
                if consecutive_retries >= max_consecutive_retries:
                    log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.")
                    break  # Force exit to collection phase

                debug_log("Waiting for workers")
                # Use async sleep to allow event loop to process worker requests
                run_async_in_server_loop(asyncio.sleep(2), timeout=3.0)
        
        debug_log(f"Master processed {processed_count} images locally")
        
        # Get all completed images to check what needs to be collected
        all_completed = run_async_in_server_loop(
            self._get_all_completed_images(multi_job_id),
            timeout=5.0
        )
        
        # Calculate how many we still need to collect
        remaining_to_collect = batch_size - len(all_completed)
        
        if remaining_to_collect > 0:
            debug_log(f"Waiting for {remaining_to_collect} more images from workers")
            # Use the unified worker timeout for the collection phase
            collection_timeout = float(get_worker_timeout_seconds())
            collected_images = run_async_in_server_loop(
                self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count),
                timeout=collection_timeout
            )
            
            # Merge collected with already completed
            all_completed.update(collected_images)
        
        # Update result images with all completed images
        for idx, processed_img in all_completed.items():
            if idx < batch_size:
                result_images[idx] = processed_img
        
        # Convert back to tensor
        result_tensor = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) if batch_size > 1 else pil_to_tensor(result_images[0])
        if upscaled_image.is_cuda:
            result_tensor = result_tensor.cuda()
        
        debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
        log(f"Completed processing all {batch_size} images")
        return (result_tensor,)
    
    
    async def _async_yield(self):
        """Simple async yield to allow event loop processing."""
        await asyncio.sleep(0)
    
    async def _get_next_image_index(self, multi_job_id):
        """Get next image index from pending queue for master."""
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
            if not job_data or 'pending_images' not in job_data:
                return None
            try:
                image_idx = await asyncio.wait_for(job_data['pending_images'].get(), timeout=1.0)
                return image_idx
            except asyncio.TimeoutError:
                return None
    
    async def _get_next_tile_index(self, multi_job_id):
        """Get next tile index from pending queue for master in static mode."""
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
            if not job_data or JOB_PENDING_TASKS not in job_data:
                return None
            try:
                tile_idx = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=0.1)
                return tile_idx
            except asyncio.TimeoutError:
                return None
    
    
    async def _get_total_completed_count(self, multi_job_id):
        """Get total count of all completed images (master + workers)."""
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
            if job_data and 'completed_images' in job_data:
                return len(job_data['completed_images'])
            return 0
    
    async def _get_all_completed_images(self, multi_job_id):
        """Get all completed images."""
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
            if job_data and 'completed_images' in job_data:
                return job_data['completed_images'].copy()
            return {}
    
    async def _get_pending_count(self, multi_job_id):
        """Get count of pending images in the queue."""
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
            if job_data and 'pending_images' in job_data:
                return job_data['pending_images'].qsize()
            return 0
    
    async def _drain_worker_results_queue(self, multi_job_id):
        """Drain pending worker results from queue and update completed_images. Returns count of drained images."""
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
            if not job_data or 'queue' not in job_data or 'completed_images' not in job_data:
                return 0
            q = job_data['queue']
            completed_images = job_data['completed_images']

            collected = 0
            while not q.empty():
                try:
                    result = await asyncio.wait_for(q.get(), timeout=0.1)
                    worker_id = result['worker_id']
                    is_last = result.get('is_last', False)

                    if 'image_idx' in result and 'image' in result:
                        image_idx = result['image_idx']
                        image_pil = result['image']
                        if image_idx not in completed_images:
                            completed_images[image_idx] = image_pil
                            collected += 1
                            debug_log(f"Drained image {image_idx} from worker {worker_id}")

                    if is_last:
                        # Optional: track worker completion if needed
                        pass
                except asyncio.TimeoutError:
                    break  # No more immediately available

            if collected > 0:
                debug_log(f"Drained {collected} worker images during retry")
            
            return collected
    
    async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
        """Check for timed out workers and requeue their assigned images. Returns count of requeued images."""
        # Use the original function from usdu_managment.py
        return await _check_and_requeue_timed_out_workers(multi_job_id, batch_size)
    
    
    def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae,
                               seed, steps, cfg, sampler_name, scheduler, denoise,
                               tile_width, tile_height, padding, mask_blur,
                               force_uniform_tiles, tiled_decode, multi_job_id, master_url,
                               worker_id, enabled_worker_ids, dynamic_threshold):
        """Worker processing in dynamic mode - processes whole images."""
        # Round tile dimensions
        tile_width = self.round_to_multiple(tile_width)
        tile_height = self.round_to_multiple(tile_height)

        # Get dimensions and tile grid
        batch_size, height, width, _ = upscaled_image.shape
        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
        log(f"USDU Dist Worker[{worker_id[:8]}]: Processing image queue | Batch {batch_size}")

        # Keep track of processed images for is_last detection
        processed_count = 0

        # Poll for job readiness to avoid races during master init
        max_poll_attempts = 20  # ~20s at 1s sleep
        for attempt in range(max_poll_attempts):
            ready = run_async_in_server_loop(
                self._check_job_status(multi_job_id, master_url),
                timeout=5.0
            )
            if ready:
                debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
                break
            time.sleep(1.0)  # Poll every 1s
        else:
            log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
            return (upscaled_image,)

        # Loop to request and process images
        while True:
            # Request an image to process
            image_idx, estimated_remaining = run_async_in_server_loop(
                self._request_image_from_master(multi_job_id, master_url, worker_id),
                timeout=TILE_WAIT_TIMEOUT
            )

            if image_idx is None:
                debug_log(f"USDU Dist Worker - No more images to process")
                break

            debug_log(f"Worker[{worker_id[:8]}] - Assigned image {image_idx}")
            processed_count += 1

            # Determine if this should be marked as last for this worker
            is_last_for_worker = (estimated_remaining == 0)

            # Extract single image tensor
            single_tensor = upscaled_image[image_idx:image_idx+1]

            # Convert to PIL for processing
            local_image = tensor_to_pil(single_tensor, 0).copy()

            # Process all tiles for this image
            image_seed = seed + image_idx * 1000

            # Pre-slice conditioning once per image (not per tile)
            positive_sliced = clone_conditioning(positive)
            negative_sliced = clone_conditioning(negative)
            for cond_list in [positive_sliced, negative_sliced]:
                for i in range(len(cond_list)):
                        emb, cond_dict = cond_list[i]
                        if emb.shape[0] > 1:
                            cond_list[i][0] = emb[image_idx:image_idx+1]
                        if 'control' in cond_dict:
                            control = cond_dict['control']
                            while control is not None:
                                hint = control.cond_hint_original
                                if hint.shape[0] > 1:
                                    control.cond_hint_original = hint[image_idx:image_idx+1]
                                control = control.previous_controlnet
                        if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
                            cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1]

                for tile_idx, pos in enumerate(all_tiles):
                    local_image = self._process_and_blend_tile(
                        tile_idx, pos, single_tensor, local_image,
                        model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
                        sampler_name, scheduler, denoise, tile_width, tile_height,
                        padding, mask_blur, width, height, force_uniform_tiles,
                        tiled_decode, batch_idx=image_idx
                    )
                    run_async_in_server_loop(
                        _send_heartbeat_to_master(multi_job_id, master_url, worker_id),
                        timeout=5.0
                    )

                # Send processed image back to master
                try:
                    # Use the estimated remaining to determine if this is the last image
                    is_last = is_last_for_worker
                    run_async_in_server_loop(
                        self._send_full_image_to_master(local_image, image_idx, multi_job_id,
                                                        master_url, worker_id, is_last),
                        timeout=TILE_SEND_TIMEOUT
                    )
                    # Send heartbeat after processing
                    run_async_in_server_loop(
                        _send_heartbeat_to_master(multi_job_id, master_url, worker_id),
                        timeout=5.0
                    )
                    if is_last:
                        break
                except Exception as e:
                    log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
                    # Continue processing other images

        # Send final is_last signal
        debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal")
        try:
            run_async_in_server_loop(
                self._send_worker_complete_signal(multi_job_id, master_url, worker_id),
                timeout=TILE_SEND_TIMEOUT
            )
        except Exception as e:
            log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending completion signal: {e}")

        return (upscaled_image,)
    
    async def _request_image_from_master(self, multi_job_id, master_url, worker_id):
        """Request an image index to process from master in dynamic mode."""
        # Enhanced retries with 404-specific delay, backoff cap, and total timeout to handle init races
        max_retries = 10
        retry_delay = 0.5
        start_time = time.time()
        
        for attempt in range(max_retries):
            # Check total timeout
            if time.time() - start_time > 30:
                log(f"Total request timeout after 30s for worker {worker_id}")
                return None, 0
                
            try:
                session = await get_client_session()
                url = f"{master_url}/distributed/request_image"
                
                async with session.post(url, json={
                    'worker_id': str(worker_id),
                    'multi_job_id': multi_job_id
                }) as response:
                    if response.status == 200:
                        data = await response.json()
                        image_idx = data.get('image_idx')
                        estimated_remaining = data.get('estimated_remaining', 0)
                        return image_idx, estimated_remaining
                    elif response.status == 404:
                        # Special handling for 404 - job not found yet
                        text = await response.text()
                        debug_log(f"Job not found (404), will retry: {text}")
                        await asyncio.sleep(1.0)
                    else:
                        text = await response.text()
                        debug_log(f"Request image failed: {response.status} - {text}")
                        
            except Exception as e:
                if attempt < max_retries - 1:
                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
                    await asyncio.sleep(retry_delay)
                    retry_delay *= 2
                    retry_delay = min(retry_delay, 5.0)  # Cap backoff at 5s
                else:
                    log(f"Failed to request image after {max_retries} attempts: {e}")
                    raise
        
        return None, 0
    
    async def _request_tile_from_master(self, multi_job_id, master_url, worker_id):
        """Request a tile index to process from master in static mode (reusing dynamic infrastructure)."""
        # Reuse the same retry logic as dynamic mode
        max_retries = 10
        retry_delay = 0.5
        start_time = time.time()
        
        for attempt in range(max_retries):
            # Check total timeout
            if time.time() - start_time > 30:
                log(f"Total request timeout after 30s for worker {worker_id}")
                return None, 0
                
            try:
                session = await get_client_session()
                url = f"{master_url}/distributed/request_image"  # Same endpoint
                
                async with session.post(url, json={
                    'worker_id': str(worker_id),
                    'multi_job_id': multi_job_id
                }) as response:
                    if response.status == 200:
                        data = await response.json()
                        tile_idx = data.get('tile_idx')
                        estimated_remaining = data.get('estimated_remaining', 0)
                        batched_static = data.get('batched_static', False)
                        return tile_idx, estimated_remaining, batched_static
                    elif response.status == 404:
                        # Special handling for 404 - job not found yet
                        text = await response.text()
                        debug_log(f"Job not found (404), will retry: {text}")
                        await asyncio.sleep(1.0)
                    else:
                        text = await response.text()
                        debug_log(f"Request tile failed: {response.status} - {text}")
                        
            except Exception as e:
                if attempt < max_retries - 1:
                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
                    await asyncio.sleep(retry_delay)
                    retry_delay *= 2
                    retry_delay = min(retry_delay, 5.0)  # Cap backoff at 5s
                else:
                    log(f"Failed to request tile after {max_retries} attempts: {e}")
                    raise
        
        return None, 0, False
    
    async def _check_job_status(self, multi_job_id, master_url):
        """Check if job is ready on the master."""
        try:
            session = await get_client_session()
            url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}"
            async with session.get(url) as response:
                if response.status == 200:
                    data = await response.json()
                    return data.get('ready', False)
                return False
        except Exception as e:
            debug_log(f"Job status check failed: {e}")
            return False
    
    async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, 
                                        master_url, worker_id, is_last):
        """Send a processed full image back to master in dynamic mode."""
        # Serialize image to PNG
        byte_io = io.BytesIO()
        image_pil.save(byte_io, format='PNG', compress_level=0)
        byte_io.seek(0)
        
        # Prepare form data
        data = aiohttp.FormData()
        data.add_field('multi_job_id', multi_job_id)
        data.add_field('worker_id', str(worker_id))
        data.add_field('image_idx', str(image_idx))
        data.add_field('is_last', str(is_last))
        data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png', 
                      content_type='image/png')
        
        # Retry logic
        max_retries = 5
        retry_delay = 0.5
        
        for attempt in range(max_retries):
            try:
                session = await get_client_session()
                url = f"{master_url}/distributed/submit_image"
                
                async with session.post(url, data=data) as response:
                    response.raise_for_status()
                    debug_log(f"Successfully sent image {image_idx} to master")
                    return
                    
            except Exception as e:
                if attempt < max_retries - 1:
                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
                    await asyncio.sleep(retry_delay)
                    retry_delay *= 2
                else:
                    log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}")
                    raise
    
    # Note: Using _send_heartbeat_to_master from usdu_managment.py instead of duplicate
    
    async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id):
        """Send completion signal to master in dynamic mode."""
        # Send a dummy request with is_last=True
        data = aiohttp.FormData()
        data.add_field('multi_job_id', multi_job_id)
        data.add_field('worker_id', str(worker_id))
        data.add_field('is_last', 'true')
        # No image data - just completion signal
        
        session = await get_client_session()
        url = f"{master_url}/distributed/submit_image"
        
        async with session.post(url, data=data) as response:
            response.raise_for_status()
            debug_log(f"Worker {worker_id} sent completion signal")

    def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
        """Determines processing mode per requested policy:
        - any workers     => prefer static (tile-based) for USDU
        - no workers      => single_gpu
        """
        if num_workers == 0:
            return "single_gpu"
        # Default to static when distributed; master/worker may still override if special cases arise
        return "static"

# Ensure initialization before registering routes
ensure_tile_jobs_initialized()

# Node registration
NODE_CLASS_MAPPINGS = {
    "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
}