robertvoy-ComfyUI-Distributed/distributed_upscale.py

"""
ComfyUI Distributed Upscale Module

This module implements a distributed version of Ultimate SD Upscale with multi-mode
batch handling for efficient processing of single images and video frames.

Processing Modes:
- Single GPU: Process all tiles locally when no workers are available
- Static Mode: For small batches (≤ dynamic_threshold), flattens tiles across batch
- Dynamic Mode: For large batches, assigns whole images to workers dynamically
"""

import torch
import numpy as np
from PIL import Image, ImageFilter, ImageDraw
import json
import asyncio
import aiohttp
from aiohttp import web
import io
import server
import math
import os
import time
from typing import List, Tuple

# Import ComfyUI modules
import comfy.samplers

# Import shared utilities
from .utils.logging import debug_log, log
from .utils.config import CONFIG_FILE
from .utils.image import tensor_to_pil, pil_to_tensor
from .utils.network import get_server_port, get_server_loop, get_client_session, handle_api_error
from .utils.async_helpers import run_async_in_server_loop
from .utils.constants import (
    TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT, TILE_TRANSFER_TIMEOUT,
    QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT
)

# Make MAX_BATCH configurable
MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20'))

# Configure maximum payload size (50MB default, configurable via environment variable)
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))

# Helper function to ensure persistent state is initialized
def ensure_tile_jobs_initialized():
    """Ensure tile job storage is initialized on the server instance."""
    prompt_server = server.PromptServer.instance
    if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
        debug_log("Initializing persistent tile job queue on server instance.")
        prompt_server.distributed_pending_tile_jobs = {}
        prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
    else:
        # Clean up any legacy queue structures that don't have the 'mode' field
        # (Should be rare after fixes, but keep for safety)
        to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
                     if not isinstance(job_data, dict) or 'mode' not in job_data]
        for job_id in to_remove:
            debug_log(f"Removing legacy queue structure for job {job_id}")
            del prompt_server.distributed_pending_tile_jobs[job_id]
    return prompt_server

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

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")
    
    @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."""
        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_tiles(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_tiles(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 - determines mode and processes accordingly."""
        # 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)
        
        if batch_size == 1 or batch_size <= dynamic_threshold:
            mode = "static"
        else:
            mode = "dynamic"
            
        debug_log(f"UltimateSDUpscale Worker - Mode: {mode}, 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 - continue with existing tile-based processing
        # Round tile dimensions
        tile_width = self.round_to_multiple(tile_width)
        tile_height = self.round_to_multiple(tile_height)
        
        # Calculate all tile positions
        _, height, width, _ = upscaled_image.shape
        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
        
        # Calculate total tiles across all images
        num_tiles_per_image = len(all_tiles)
        total_tiles = batch_size * num_tiles_per_image
        
        # Get assigned global tile indices for this worker
        assigned_global_indices = self._get_worker_global_indices(total_tiles, enabled_worker_ids, worker_id)
        if not assigned_global_indices:
            return (upscaled_image,)
        
        debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(assigned_global_indices)} tiles from batch of {batch_size}")
        
        # Process tiles SYNCHRONOUSLY
        processed_tiles = []
        for global_idx in assigned_global_indices:
            # 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]}")
                continue
                
            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
            )
            
            # Process tile through SD with unique seed
            image_seed = seed + batch_idx * 1000
            processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae,
                                             image_seed + tile_idx, steps, cfg, sampler_name, 
                                             scheduler, denoise, tiled_decode)
            
            processed_tiles.append({
                'tile': processed_tile,
                'global_idx': global_idx,
                'batch_idx': batch_idx,
                'tile_idx': tile_idx,
                'x': x1,  # Use extraction position
                'y': y1,
                'extracted_width': ew,
                'extracted_height': eh
            })
        
        # Send results back using the server's event loop in a single batch
        try:
            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
            )
        except Exception as e:
            log(f"UltimateSDUpscale Worker - Error sending tiles: {e}")
        
        return (upscaled_image,)
    
    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):
        """Master processing - SYNCHRONOUS with async collection."""
        # 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)
        total_tiles = len(all_tiles)
        
        debug_log(f"UltimateSDUpscale Master - Total tiles: {total_tiles}")
        
        # Parse enabled workers
        enabled_workers = json.loads(enabled_worker_ids)
        num_workers = len(enabled_workers)
        
        # Determine processing mode
        if num_workers == 0:
            mode = "single_gpu"
        elif batch_size == 1 or batch_size <= dynamic_threshold:
            mode = "static"
        else:
            mode = "dynamic"
        
        debug_log(f"UltimateSDUpscale Master - Mode: {mode}, batch_size: {batch_size}, workers: {num_workers}")
        
        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 - continue with existing implementation
        
        # Initialize job queue FIRST before any processing
        # Initialize the queue directly in the persistent server storage
        # This ensures it's accessible from the web server's event loop
        try:
            run_async_in_server_loop(
                self._init_job_queue(multi_job_id),
                timeout=QUEUE_INIT_TIMEOUT
            )
            debug_log(f"UltimateSDUpscale Master - Queue initialization complete")
        except Exception as e:
            debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}")
            raise RuntimeError(f"Failed to initialize static mode queue: {e}")
        
        # Don't try to prepare via API since we're using local queues
        # loop = asyncio.new_event_loop()
        # asyncio.set_event_loop(loop)
        
        # 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())
        
        # Mask will be created per tile for proper blending
        
        # Calculate total tiles across all images
        num_tiles_per_image = len(all_tiles)
        total_tiles_all_images = batch_size * num_tiles_per_image
        
        # Calculate master's global tile indices 
        master_global_indices = self._get_master_global_indices(total_tiles_all_images, num_workers)
        master_tiles = len(master_global_indices)
        
        debug_log(f"UltimateSDUpscale Master - Processing {master_tiles} tiles locally from batch of {batch_size}")
        
        # Process master tiles SYNCHRONOUSLY
        for global_idx in master_global_indices:
            # 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 >= len(result_images):
                debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds result_images length {len(result_images)}")
                continue
                
            # Use unique seed per image
            image_seed = seed + batch_idx * 1000
            
            result_images[batch_idx] = self._process_and_blend_tile(
                tile_idx, all_tiles[tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx],
                model, positive, negative, vae, image_seed, steps, cfg,
                sampler_name, scheduler, denoise, tile_width, tile_height,
                padding, mask_blur, width, height, tiled_decode
            )
        
        # Collect worker tiles using async
        worker_tiles_expected = total_tiles_all_images - master_tiles
        if worker_tiles_expected > 0:
            # Calculate tile distribution for active worker calculation
            tiles_per_participant = total_tiles_all_images // (num_workers + 1)
            remainder = total_tiles_all_images % (num_workers + 1)
            
            # Calculate which workers will actually have tiles
            active_workers = []
            for i in range(num_workers):
                worker_start_idx = master_tiles + (i * tiles_per_participant)
                if i < remainder - 1:
                    worker_start_idx += i
                    worker_tile_count = tiles_per_participant + 1
                else:
                    worker_start_idx += remainder - 1 if remainder > 0 else 0
                    worker_tile_count = tiles_per_participant
                
                if worker_start_idx < total_tiles_all_images and worker_tile_count > 0:
                    active_workers.append(enabled_workers[i])
            
            debug_log(f"UltimateSDUpscale Master - Waiting for {worker_tiles_expected} tiles from {len(active_workers)} active workers")
            debug_log(f"UltimateSDUpscale Master - Active worker IDs: {active_workers}")
            debug_log(f"UltimateSDUpscale Master - All enabled worker IDs: {enabled_workers}")
            
            # Collect tiles from workers using the server's event loop
            collected_tiles = run_async_in_server_loop(
                self._async_collect_worker_tiles(multi_job_id, len(active_workers)),
                timeout=TILE_COLLECTION_TIMEOUT
            )
            
            # Blend collected tiles SYNCHRONOUSLY
            for tile_data in collected_tiles.values():
                    x = tile_data['x']
                    y = tile_data['y']
                    tile_tensor = tile_data['tensor']
                    tile_idx = tile_data['tile_idx']
                    extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding)
                    extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding)
                    
                    # Get batch index if available (for batch processing)
                    batch_idx = tile_data.get('batch_idx', 0)
                    
                    # Skip if batch_idx is out of range
                    if batch_idx >= len(result_images):
                        debug_log(f"Warning: Received tile for batch_idx {batch_idx} but only have {len(result_images)} images")
                        continue
                    
                    # Convert and blend
                    tile_pil = tensor_to_pil(tile_tensor, 0)
                    # Get the original tile position from the tile index
                    orig_x, orig_y = all_tiles[tile_idx]
                    # Create mask for this specific tile
                    tile_mask = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur)
                    # Use extraction position and size for blending
                    result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil, 
                                                 x, y, (extracted_width, extracted_height), tile_mask, padding)
        
        # Convert back to tensor
        # For batch processing, combine all result images
        if batch_size == 1:
            result_tensor = pil_to_tensor(result_images[0])
        else:
            # Convert all images and concatenate
            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()
        
        debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
        return (result_tensor,)
    
    def _get_worker_global_indices(self, total_tiles, enabled_worker_ids, worker_id):
        """Calculate which global tile indices are assigned to a specific worker in flattened mode."""
        enabled_workers = json.loads(enabled_worker_ids)
        num_workers = len(enabled_workers)
        
        debug_log(f"UltimateSDUpscale Worker - Worker ID: {worker_id}, Enabled workers: {enabled_workers}")
        
        try:
            worker_index = enabled_workers.index(worker_id)
        except ValueError:
            log(f"UltimateSDUpscale Worker - Worker {worker_id} not found in enabled list {enabled_workers}")
            return []
        
        # Calculate tile distribution
        tiles_per_participant = total_tiles // (num_workers + 1)
        remainder = total_tiles % (num_workers + 1)
        master_tiles = tiles_per_participant + (1 if remainder > 0 else 0)
        
        start_idx = master_tiles + (worker_index * tiles_per_participant)
        if worker_index < remainder - 1:
            start_idx += worker_index
            end_idx = start_idx + tiles_per_participant + 1
        else:
            start_idx += remainder - 1 if remainder > 0 else 0
            end_idx = start_idx + tiles_per_participant
        
        end_idx = min(end_idx, total_tiles)
        return list(range(start_idx, end_idx))
    
    def _get_worker_tiles(self, all_tiles, enabled_worker_ids, worker_id):
        """Calculate which tiles are assigned to a specific worker."""
        enabled_workers = json.loads(enabled_worker_ids)
        num_workers = len(enabled_workers)
        total_tiles = len(all_tiles)
        
        debug_log(f"UltimateSDUpscale Worker - Worker ID: {worker_id}, Enabled workers: {enabled_workers}")
        
        try:
            worker_index = enabled_workers.index(worker_id)
        except ValueError:
            log(f"UltimateSDUpscale Worker - Worker {worker_id} not found in enabled list {enabled_workers}")
            return []
        
        # Calculate tile distribution
        tiles_per_participant = total_tiles // (num_workers + 1)
        remainder = total_tiles % (num_workers + 1)
        master_tiles = tiles_per_participant + (1 if remainder > 0 else 0)
        
        start_idx = master_tiles + (worker_index * tiles_per_participant)
        if worker_index < remainder - 1:
            start_idx += worker_index
            end_idx = start_idx + tiles_per_participant + 1
        else:
            start_idx += remainder - 1 if remainder > 0 else 0
            end_idx = start_idx + tiles_per_participant
        
        end_idx = min(end_idx, total_tiles)
        return list(range(start_idx, end_idx))
    
    def _get_master_global_indices(self, total_tiles, num_workers):
        """Calculate which global tile indices are assigned to the master in flattened mode."""
        tiles_per_participant = total_tiles // (num_workers + 1)
        remainder = total_tiles % (num_workers + 1)
        master_tiles = tiles_per_participant + (1 if remainder > 0 else 0)
        return list(range(master_tiles))
    
    def _get_master_tiles(self, all_tiles, num_workers):
        """Calculate which tiles are assigned to the master."""
        total_tiles = len(all_tiles)
        tiles_per_participant = total_tiles // (num_workers + 1)
        remainder = total_tiles % (num_workers + 1)
        master_tiles = tiles_per_participant + (1 if remainder > 0 else 0)
        return list(range(master_tiles))
    
    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, tiled_decode):
        """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
        )
        
        processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae,
                                         seed + tile_idx, steps, cfg, sampler_name, 
                                         scheduler, denoise, tiled_decode)
        
        # Convert and blend
        processed_pil = tensor_to_pil(processed_tile, 0)
        # Create mask for this specific tile
        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 _prepare_multigpu_job(self, multi_job_id):
        """Prepare job via API endpoint to ensure server is ready."""
        session = await get_client_session()
        # Use the actual server port (get from server instance)
        port = get_server_port()
        async with session.post(f"http://localhost:{port}/distributed/prepare_job",
                              json={"multi_job_id": multi_job_id}) as response:
            if response.status != 200:
                text = await response.text()
                raise RuntimeError(f"Failed to prepare job: {response.status} - {text}")
    
    async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, worker_indices=None):
        """Initialize the job queue for dynamic mode with pending images."""
        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:
                prompt_server.distributed_pending_tile_jobs[multi_job_id] = {
                    'queue': asyncio.Queue(),
                    'mode': 'dynamic',
                    'pending_images': asyncio.Queue(),
                    'completed_images': {},
                    'completed_tiles': {},  # For uniformity
                    'batch_size': batch_size,
                    'assigned_to_workers': assigned_to_workers or {},
                    'worker_status': worker_status or {}
                }
                # Initialize pending images queue with worker indices only
                pending_images = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images']
                for i in worker_indices:
                    await pending_images.put(i)
                debug_log(f"UltimateSDUpscale Master Dynamic - Initialized job queue with {len(worker_indices)} pending images for workers")
            else:
                debug_log(f"UltimateSDUpscale Master Dynamic - Queue already exists for {multi_job_id}")
    
    async def _init_job_queue(self, multi_job_id):
        """Initialize the job queue for collecting tiles."""
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            debug_log(f"UltimateSDUpscale Master - _init_job_queue: Checking if {multi_job_id} exists in distributed_pending_tile_jobs")
            debug_log(f"UltimateSDUpscale Master - _init_job_queue: Current jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}")
            if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
                prompt_server.distributed_pending_tile_jobs[multi_job_id] = {
                    'queue': asyncio.Queue(),
                    'mode': 'static',
                    'completed_tiles': {}  # For static mode data
                }
                debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}")
                debug_log(f"UltimateSDUpscale Master - _init_job_queue: Created new queue for {multi_job_id}")
            else:
                debug_log(f"UltimateSDUpscale Master - _init_job_queue: Queue already exists for {multi_job_id}")
    
    async def _async_collect_worker_tiles(self, multi_job_id, num_workers):
        """Async helper to collect tiles from workers."""
        # 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'] != 'static':
                raise RuntimeError(f"Mode mismatch: expected static, got {job_data['mode']}")
            q = job_data['queue']
        
        collected_tiles = {}
        workers_done = set()
        timeout = TILE_WAIT_TIMEOUT
        
        debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {num_workers} workers to complete")
        
        while len(workers_done) < num_workers:
            debug_log(f"UltimateSDUpscale Master - Loop status: {len(workers_done)}/{num_workers} workers done, {len(collected_tiles)} tiles collected, workers_done set: {workers_done}")
            try:
                result = await asyncio.wait_for(q.get(), timeout=timeout)
                worker_id = result['worker_id']
                is_last = result.get('is_last', False)
                
                # Check if batch mode
                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
                        batch_idx = tile_data['batch_idx']
                        
                        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
                        collected_tiles[key] = {
                            'tensor': tile_data['tensor'],
                            '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)
                        }
                else:
                    # Single tile mode (backward compat)
                    tile_idx = result['tile_idx']
                    collected_tiles[tile_idx] = result
                    debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})")
                
                if is_last:
                    workers_done.add(worker_id)
                    debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
                    debug_log(f"UltimateSDUpscale Master - Worker '{worker_id}' marked as done. Total workers done: {len(workers_done)}/{num_workers}")
                
            except asyncio.TimeoutError:
                log(f"UltimateSDUpscale Master - Timeout waiting for tiles")
                debug_log(f"UltimateSDUpscale Master - Final status: {len(workers_done)}/{num_workers} workers done, {len(collected_tiles)} tiles collected")
                debug_log(f"UltimateSDUpscale Master - Workers that completed: {workers_done}")
                break
        
        debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_tiles)} tiles from {len(workers_done)} workers")
        
        # Clean up job queue
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
                del prompt_server.distributed_pending_tile_jobs[multi_job_id]
        
        return collected_tiles
    
    async def _async_process_dynamic(self, multi_job_id, expected_from_workers, num_workers):
        """Async helper for dynamic mode - collect processed whole images from workers.
        
        This method implements dynamic load balancing where workers request images
        as they become available, ensuring optimal utilization of GPU resources."""
        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'] != 'dynamic':
                raise RuntimeError(f"Mode mismatch: expected dynamic, got {job_data['mode']}")
            q = job_data['queue']
            completed_images = job_data['completed_images']
        
        workers_done = set()
        timeout = TILE_WAIT_TIMEOUT * 2  # Longer timeout for full images
        
        debug_log(f"UltimateSDUpscale Master Dynamic - Waiting for {expected_from_workers} images from {num_workers} workers")
        
        last_heartbeat_check = time.time()
        
        while len(completed_images) < expected_from_workers and len(workers_done) < num_workers:
            try:
                result = await asyncio.wait_for(q.get(), timeout=10.0)  # Shorter timeout for regular checks
                worker_id = result['worker_id']
                is_last = result.get('is_last', False)
                
                if 'image_idx' in result and 'image' in result:
                    image_idx = result['image_idx']
                    image_pil = result['image']
                    completed_images[image_idx] = image_pil
                    debug_log(f"UltimateSDUpscale Master Dynamic - Received image {image_idx} from worker {worker_id}")
                
                if is_last:
                    workers_done.add(worker_id)
                    debug_log(f"UltimateSDUpscale Master Dynamic - Worker {worker_id} completed")
                    
            except asyncio.TimeoutError:
                # Check for worker timeouts every 10 seconds
                current_time = time.time()
                if current_time - last_heartbeat_check >= 10.0:
                    async with prompt_server.distributed_tile_jobs_lock:
                        if multi_job_id in prompt_server.distributed_pending_tile_jobs:
                            job_data_check = prompt_server.distributed_pending_tile_jobs[multi_job_id]
                            # Check for timed out workers (60 second timeout)
                            for worker, last_heartbeat in list(job_data_check.get('worker_status', {}).items()):
                                if current_time - last_heartbeat > 60:
                                    log(f"UltimateSDUpscale Master Dynamic - Worker {worker} timed out")
                                    # Requeue assigned images from this worker
                                    for idx in job_data_check.get('assigned_to_workers', {}).get(worker, []):
                                        if idx not in completed_images:
                                            await job_data_check['pending_images'].put(idx)
                                            debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx} from timed out worker {worker}")
                                    # Mark worker as done and clean up
                                    workers_done.add(worker)
                                    if 'worker_status' in job_data_check:
                                        del job_data_check['worker_status'][worker]
                                    if 'assigned_to_workers' in job_data_check:
                                        job_data_check['assigned_to_workers'][worker] = []
                    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 Dynamic - Overall timeout waiting for images")
                    # Requeue all unfinished images
                    async with prompt_server.distributed_tile_jobs_lock:
                        if multi_job_id in prompt_server.distributed_pending_tile_jobs:
                            pending_queue = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images']
                            for idx in range(job_data_check.get('batch_size', 0)):
                                if idx not in completed_images:
                                    await pending_queue.put(idx)
                                    debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx}")
                    break
        
        debug_log(f"UltimateSDUpscale Master Dynamic - Collection complete. Got {len(completed_images)} images")
        
        # 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 completed_images
    
    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 for the image.
        
        Tiles are placed at grid positions without overlap in their placement.
        The overlap happens during extraction with padding.
        
        If force_uniform_tiles is True, tiles are distributed evenly across the image."""
        if force_uniform_tiles:
            # Calculate number of tiles needed
            num_tiles_x = math.ceil(image_width / tile_width)
            num_tiles_y = math.ceil(image_height / tile_height)
            
            # Calculate uniform spacing
            tiles = []
            
            # Handle X dimension
            if num_tiles_x == 1:
                tiles_x = [0]
            else:
                # Calculate step to distribute tiles evenly
                step_x = (image_width - tile_width) / (num_tiles_x - 1)
                # Check for excessive overlap
                if step_x < tile_width / 2:
                    debug_log(f"Warning: Uniform tiles would have excessive overlap (step_x={step_x}, tile_width={tile_width}). Falling back to non-uniform mode.")
                    force_uniform_tiles = False
                else:
                    tiles_x = [round(i * step_x) for i in range(num_tiles_x)]
                    # Ensure tiles are aligned to 8-pixel boundaries
                    tiles_x = [self.round_to_multiple(x, 8) for x in tiles_x]
            
            # Handle Y dimension (if still using uniform tiles)
            if force_uniform_tiles:
                if num_tiles_y == 1:
                    tiles_y = [0]
                else:
                    # Calculate step to distribute tiles evenly
                    step_y = (image_height - tile_height) / (num_tiles_y - 1)
                    # Check for excessive overlap
                    if step_y < tile_height / 2:
                        debug_log(f"Warning: Uniform tiles would have excessive overlap (step_y={step_y}, tile_height={tile_height}). Falling back to non-uniform mode.")
                        force_uniform_tiles = False
                    else:
                        tiles_y = [round(i * step_y) for i in range(num_tiles_y)]
                        # Ensure tiles are aligned to 8-pixel boundaries
                        tiles_y = [self.round_to_multiple(y, 8) for y in tiles_y]
        
        # If still using uniform tiles, generate positions
        if force_uniform_tiles:
            # Generate all tile positions
            for y in tiles_y:
                for x in tiles_x:
                    tiles.append((x, y))
                    
            return tiles
        
        # Fall through to non-uniform mode if needed
        if not force_uniform_tiles:
            # Original non-uniform tile calculation
            tiles = []
            for y in range(0, image_height, tile_height):
                for x in range(0, image_width, tile_width):
                    tiles.append((x, y))
            return tiles
    
    def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int,
                                 tile_width: int, tile_height: int, padding: int) -> Tuple[torch.Tensor, int, int, int, int]:
        """Extract a tile from the image with padding."""
        _, h, w, _ = image.shape
        
        # Calculate extraction bounds with padding
        x1 = max(0, x - padding)
        y1 = max(0, y - padding)
        x2 = min(w, x + tile_width + padding)
        y2 = min(h, y + tile_height + padding)
        
        # Store the actual extracted size before resizing
        extracted_width = x2 - x1
        extracted_height = y2 - y1
        
        # Debug logging (uncomment if needed)
        # debug_log(f"[Extract] Tile at ({x}, {y}) -> Extract from ({x1}, {y1}) to ({x2}, {y2}), size: {extracted_width}x{extracted_height}")
        
        # Extract tile
        tile = image[:, y1:y2, x1:x2, :]
        
        # Convert to PIL using utility function
        tile_pil = tensor_to_pil(tile, 0)
        
        # Resize to target dimensions for processing
        tile_pil = tile_pil.resize((tile_width, tile_height), Image.LANCZOS)
        
        # Convert back to tensor using utility function
        tile_tensor = pil_to_tensor(tile_pil)
        
        # Move to same device as original image
        if image.is_cuda:
            tile_tensor = tile_tensor.cuda()
        
        return tile_tensor, 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) -> torch.Tensor:
        """Process a single tile through SD sampling."""
        # 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)
        
        # Move to correct device
        if tile_tensor.is_cuda:
            clean_tensor = clean_tensor.cuda()
        
        # Encode to latent
        if tiled_decode and tiled_vae_available:
            latent = VAEEncodeTiled().encode(vae, clean_tensor)[0]
        else:
            latent = VAEEncode().encode(vae, clean_tensor)[0]
        
        # Sample
        samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, 
                                positive, negative, latent, denoise=denoise)[0]
        
        # Decode back to image
        if tiled_decode and tiled_vae_available:
            image = VAEDecodeTiled().decode(vae, samples)[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: int, 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')
    
    
    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"UltimateSDUpscale Worker - Sending {total_tiles} tiles in chunks of max {MAX_BATCH}")
        
        for start in range(0, total_tiles, MAX_BATCH):
            chunk = processed_tiles[start:start + MAX_BATCH]
            chunk_size = len(chunk)
            is_chunk_last = (start + chunk_size == total_tiles)  # True only for final chunk
            
            debug_log(f"UltimateSDUpscale Worker - Processing chunk {start//MAX_BATCH + 1}, tiles {start} to {start + chunk_size - 1}")
            
            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', str(is_chunk_last))
            data.add_field('batch_size', str(chunk_size))
            data.add_field('padding', str(padding))
            
            # Chunk metadata: Keep absolute tile_idx and batch_idx if present
            metadata = []
            for tile_data in chunk:
                meta = {
                    'tile_idx': tile_data['tile_idx'],  # Original/absolute idx
                    'x': tile_data['x'],
                    'y': tile_data['y'],
                    'extracted_width': tile_data['extracted_width'],
                    'extracted_height': tile_data['extracted_height']
                }
                # Add batch_idx and global_idx if present (for batch processing)
                if 'batch_idx' in tile_data:
                    meta['batch_idx'] = tile_data['batch_idx']
                if 'global_idx' in tile_data:
                    meta['global_idx'] = tile_data['global_idx']
                metadata.append(meta)
            
            # Add JSON metadata
            data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')
            
            # Add chunk images as 'tile_{j}' where j=0 to chunk_size-1
            for j, tile_data in enumerate(chunk):
                # Convert tensor to PIL
                img = tensor_to_pil(tile_data['tile'], 0)
                byte_io = io.BytesIO()
                img.save(byte_io, format='PNG', compress_level=0)
                byte_io.seek(0)
                
                # Add image field
                data.add_field(f'tile_{j}', byte_io, filename=f'tile_{j}.png', content_type='image/png')
            
            # Estimate payload size for logging
            payload_size = 0
            for field in data._fields:
                if hasattr(field[2], 'read'):
                    field[2].seek(0)
                    payload_size += len(field[2].read())
                    field[2].seek(0)
            debug_log(f"UltimateSDUpscale Worker - Chunk payload size: {payload_size:,} bytes")
            
            # 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/tile_complete"
                    
                    debug_log(f"UltimateSDUpscale Worker - Sending chunk of {chunk_size} tiles to {url}, attempt {attempt + 1}")
                    
                    async with session.post(url, data=data) as response:
                        response.raise_for_status()
                        debug_log(f"UltimateSDUpscale Worker - Successfully sent chunk of {chunk_size} tiles")
                        break  # Success, move to next chunk
                        
                except Exception as e:
                    if attempt < max_retries - 1:
                        debug_log(f"UltimateSDUpscale Worker - Retry {attempt + 1}/{max_retries} after error: {e}")
                        await asyncio.sleep(retry_delay)
                        retry_delay *= 2  # Exponential backoff
                    else:
                        log(f"UltimateSDUpscale Worker - Failed to send chunk after {max_retries} attempts: {e}")
                        raise

    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)."""
        # 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)
        
        debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally for batch of {batch_size}")
        
        # Convert batch to PIL list
        result_images = []
        for b in range(batch_size):
            image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
            # Ensure RGB mode for consistency
            if image_pil.mode != 'RGB':
                image_pil = image_pil.convert('RGB')
            result_images.append(image_pil.copy())
        
        # Process each image in the batch
        for batch_idx in range(batch_size):
            # Process each tile for this image
            for tile_idx, tile_pos in enumerate(all_tiles):
                # Use unique seed per image
                image_seed = seed + batch_idx * 1000
                result_images[batch_idx] = self._process_and_blend_tile(
                    tile_idx, tile_pos, upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx],
                    model, positive, negative, vae, image_seed, steps, cfg,
                    sampler_name, scheduler, denoise, tile_width, tile_height,
                    padding, mask_blur, width, height, tiled_decode
                )
        
        # 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()
        
        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."""
        # Get batch size and dimensions
        batch_size, height, width, _ = upscaled_image.shape
        num_workers = len(enabled_workers)
        
        debug_log(f"UltimateSDUpscale Master Dynamic - Processing batch of {batch_size} images with {num_workers} workers")
        
        # Calculate master's share similar to static mode
        num_participants = num_workers + 1
        images_per_participant = batch_size // num_participants
        remainder = batch_size % num_participants
        master_images_count = images_per_participant + (1 if remainder > 0 else 0)
        master_indices = list(range(master_images_count))
        worker_indices = list(range(master_images_count, batch_size))
        expected_from_workers = len(worker_indices)
        
        # Calculate tiles for processing
        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
        
        # Initialize job queue for communication - synchronous approach
        try:
            # Add worker tracking data
            assigned_to_workers = {w: [] for w in enabled_workers}
            worker_status = {w: time.time() for w in enabled_workers}
            
            # Use short timeout for async init, with threading fallback
            try:
                run_async_in_server_loop(
                    self._init_job_queue_dynamic(multi_job_id, batch_size, assigned_to_workers, worker_status, worker_indices),
                    timeout=2.0  # Short timeout
                )
            except Exception as init_error:
                debug_log(f"UltimateSDUpscale Master Dynamic - Async init failed, using threading: {init_error}")
                # Threading fallback for dynamic mode
                import threading
                def sync_dynamic_init():
                    prompt_server = ensure_tile_jobs_initialized()
                    if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
                        prompt_server.distributed_pending_tile_jobs[multi_job_id] = {
                            'queue': asyncio.Queue(),
                            'mode': 'dynamic',
                            'pending_images': asyncio.Queue(),
                            'completed_images': {},
                            'completed_tiles': {},
                            'batch_size': batch_size,
                            'assigned_to_workers': assigned_to_workers,
                            'worker_status': worker_status
                        }
                        # Initialize pending images queue with worker indices only
                        pending_images = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images']
                        for i in worker_indices:
                            pending_images.put_nowait(i)
                        debug_log(f"Synchronous dynamic init for {multi_job_id} with {len(worker_indices)} pending images for workers")
                thread = threading.Thread(target=sync_dynamic_init)
                thread.start()
                thread.join(timeout=3.0)
                if thread.is_alive():
                    raise RuntimeError("Dynamic queue init failed")
        except Exception as e:
            debug_log(f"UltimateSDUpscale Master Dynamic - Queue initialization error: {e}")
            raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}")
        
        # Convert batch to PIL list
        result_images = []
        for b in range(batch_size):
            image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
            # Ensure RGB mode for consistency
            if image_pil.mode != 'RGB':
                image_pil = image_pil.convert('RGB')
            result_images.append(image_pil.copy())
        
        # Process master's share locally
        debug_log(f"UltimateSDUpscale Master Dynamic - Processing {master_images_count} images locally")
        for idx in master_indices:
            single_tensor = upscaled_image[idx:idx+1]
            local_image = result_images[idx]
            image_seed = seed + idx * 1000
            for tile_idx, pos in enumerate(all_tiles):
                local_image = self._process_and_blend_tile(
                    tile_idx, pos, single_tensor, local_image,
                    model, positive, negative, vae, image_seed, steps, cfg,
                    sampler_name, scheduler, denoise, tile_width, tile_height,
                    padding, mask_blur, width, height, tiled_decode
                )
            result_images[idx] = local_image
        
        # Process with workers using dynamic assignment
        collected_images = run_async_in_server_loop(
            self._async_process_dynamic(multi_job_id, expected_from_workers, num_workers),
            timeout=TILE_COLLECTION_TIMEOUT * 2  # Longer timeout for full images
        )
        
        # Update result images with processed ones
        for idx, processed_img in collected_images.items():
            if idx < batch_size:
                result_images[idx] = processed_img
        
        # 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()
        
        debug_log(f"UltimateSDUpscale Master Dynamic - Job {multi_job_id} complete")
        return (result_tensor,)
    
    def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae,
                                  seed, steps, cfg, sampler_name, scheduler, denoise,
                                  tile_width, tile_height, padding, mask_blur,
                                  force_uniform_tiles, tiled_decode, multi_job_id, master_url,
                                  worker_id, enabled_worker_ids, dynamic_threshold):
            """Worker processing in dynamic mode - processes whole images."""
            # Round tile dimensions
            tile_width = self.round_to_multiple(tile_width)
            tile_height = self.round_to_multiple(tile_height)
            
            # Get dimensions
            _, height, width, _ = upscaled_image.shape
            all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
            
            debug_log(f"UltimateSDUpscale Worker Dynamic - Worker {worker_id} ready to process images")
            
            # 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"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"UltimateSDUpscale Worker Dynamic - No more images to process")
                    break
                    
                debug_log(f"UltimateSDUpscale Worker Dynamic - Mode: dynamic, assigned image {image_idx} to worker {worker_id}")
                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
                for tile_idx, pos in enumerate(all_tiles):
                    local_image = self._process_and_blend_tile(
                        tile_idx, pos, single_tensor, local_image,
                        model, positive, negative, vae, image_seed, steps, cfg,
                        sampler_name, scheduler, denoise, tile_width, tile_height,
                        padding, mask_blur, width, height, tiled_decode
                    )
                
                # Send processed image back to master
                try:
                    # Use the estimated remaining to determine if this is the last image
                    is_last = is_last_for_worker
                    run_async_in_server_loop(
                        self._send_full_image_to_master(local_image, image_idx, multi_job_id, 
                                                      master_url, worker_id, is_last),
                        timeout=TILE_SEND_TIMEOUT
                    )
                    # Send heartbeat after processing
                    run_async_in_server_loop(
                        self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
                        timeout=5.0
                    )
                    if is_last:
                        break
                except Exception as e:
                    log(f"UltimateSDUpscale Worker Dynamic - Error sending image {image_idx}: {e}")
                    # Continue processing other images
            
            # Send final is_last signal
            debug_log(f"UltimateSDUpscale Worker Dynamic - Worker {worker_id} 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"UltimateSDUpscale Worker Dynamic - 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 _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/tile_complete"
                
                async with session.post(url, data=data) as response:
                    response.raise_for_status()
                    debug_log(f"Successfully sent image {image_idx} to master")
                    return
                    
            except Exception as e:
                if attempt < max_retries - 1:
                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
                    await asyncio.sleep(retry_delay)
                    retry_delay *= 2
                else:
                    log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}")
                    raise
    
    async def _send_heartbeat_to_master(self, multi_job_id, master_url, worker_id):
        """Send heartbeat signal to master."""
        try:
            data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)}
            session = await get_client_session()
            url = f"{master_url}/distributed/heartbeat"
            
            async with session.post(url, json=data) as response:
                response.raise_for_status()
                debug_log(f"Worker {worker_id} sent heartbeat")
        except Exception as e:
            debug_log(f"Worker {worker_id} heartbeat failed: {e}")
    
    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')
        # Send empty tile data to signal completion
        data.add_field('batch_size', '0')
        
        session = await get_client_session()
        url = f"{master_url}/distributed/tile_complete"
        
        async with session.post(url, data=data) as response:
            response.raise_for_status()
            debug_log(f"Worker {worker_id} sent completion signal")


# Ensure initialization before registering routes
ensure_tile_jobs_initialized()

# API Endpoint for tile completion
@server.PromptServer.instance.routes.post("/distributed/request_image")
async def request_image_endpoint(request):
    """Endpoint for workers to request images in dynamic mode."""
    try:
        data = await request.json()
        worker_id = data.get('worker_id')
        multi_job_id = data.get('multi_job_id')
        
        if not worker_id or not multi_job_id:
            return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
        
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
                if not isinstance(job_data, dict) or 'mode' not in job_data:
                    return await handle_api_error(request, "Invalid job data structure", 500)
                if job_data['mode'] != 'dynamic':
                    return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400)
                if 'pending_images' in job_data:
                    pending_queue = job_data['pending_images']
                    try:
                        image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
                        # Track assigned image
                        if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']:
                            job_data['assigned_to_workers'][worker_id].append(image_idx)
                        # Update worker heartbeat
                        if 'worker_status' in job_data:
                            job_data['worker_status'][worker_id] = time.time()
                        # Get estimated remaining count
                        remaining = pending_queue.qsize()  # Approximate
                        debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}")
                        return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining})
                    except asyncio.TimeoutError:
                        return web.json_response({"image_idx": None})
                else:
                    return await handle_api_error(request, "Invalid dynamic mode configuration", 400)
            else:
                return await handle_api_error(request, "Job not found", 404)
    except Exception as e:
        return await handle_api_error(request, e, 500)

@server.PromptServer.instance.routes.get("/distributed/job_status")
async def job_status_endpoint(request):
    """Endpoint to check if a job is ready."""
    multi_job_id = request.query.get('multi_job_id')
    if not multi_job_id:
        return web.json_response({"ready": False})
    prompt_server = ensure_tile_jobs_initialized()
    async with prompt_server.distributed_tile_jobs_lock:
        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
        ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data)
        return web.json_response({"ready": ready})

@server.PromptServer.instance.routes.post("/distributed/heartbeat")
async def heartbeat_endpoint(request):
    """Endpoint for workers to send heartbeat signals."""
    try:
        data = await request.json()
        worker_id = data.get('worker_id')
        multi_job_id = data.get('multi_job_id')
        
        if not worker_id or not multi_job_id:
            return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
        
        prompt_server = ensure_tile_jobs_initialized()
        async with prompt_server.distributed_tile_jobs_lock:
            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
                if 'worker_status' in job_data:
                    job_data['worker_status'][worker_id] = time.time()
                    debug_log(f"UltimateSDUpscale API - Heartbeat from worker {worker_id}")
                    return web.json_response({"status": "success"})
                else:
                    return await handle_api_error(request, "Worker status tracking not available", 400)
            else:
                return await handle_api_error(request, "Job not found", 404)
    except Exception as e:
        return await handle_api_error(request, e, 500)

@server.PromptServer.instance.routes.post("/distributed/tile_complete")
async def tile_complete_endpoint(request):
    """Endpoint for receiving completed tiles from workers."""
    try:
        # Check payload size before processing
        content_length = request.headers.get('content-length')
        if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
            return await handle_api_error(request, f"Payload too large: {content_length} bytes (max: {MAX_PAYLOAD_SIZE})", 413)
        
        data = await request.post()
        multi_job_id = data.get('multi_job_id')
        worker_id = data.get('worker_id')
        is_last = data.get('is_last', 'False').lower() == 'true'
        
        if multi_job_id is None or worker_id is None:
            return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)

        prompt_server = ensure_tile_jobs_initialized()
        
        # Check if this is a full image submission (dynamic mode)
        if 'full_image' in data and 'image_idx' in data:
            image_idx = int(data.get('image_idx'))
            is_last = data.get('is_last', 'False').lower() == 'true'
            
            try:
                # Process full image
                img_data = data['full_image'].file.read()
                img = Image.open(io.BytesIO(img_data)).convert("RGB")
                
                debug_log(f"UltimateSDUpscale API - Received full image {image_idx} from worker {worker_id}")
                
                # Put into dynamic mode queue
                async with prompt_server.distributed_tile_jobs_lock:
                    if multi_job_id in prompt_server.distributed_pending_tile_jobs:
                        job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
                        if not isinstance(job_data, dict) or 'mode' not in job_data:
                            return await handle_api_error(request, "Invalid job data structure", 500)
                        if job_data['mode'] != 'dynamic':
                            return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400)
                        # Check for duplicate image submission
                        if 'completed_images' in job_data and image_idx in job_data['completed_images']:
                            log(f"UltimateSDUpscale API - Duplicate image {image_idx} from {worker_id}, ignoring")
                            return web.json_response({"status": "duplicate"})
                        if 'queue' in job_data:
                            await job_data['queue'].put({
                                'worker_id': worker_id,
                                'image_idx': image_idx,
                                'image': img,
                                'is_last': is_last
                            })
                            return web.json_response({"status": "success"})
                        else:
                            return await handle_api_error(request, "Invalid dynamic mode configuration", 400)
                    else:
                        return await handle_api_error(request, "Job not found", 404)
            except Exception as e:
                log(f"Error processing full image from worker {worker_id}: {e}")
                return await handle_api_error(request, f"Image processing error: {e}", 400)
        
        # Check for batch mode
        batch_size = int(data.get('batch_size', 0))
        tiles = []
        
        # Handle completion signals (batch_size=0 with is_last=True)
        if batch_size == 0 and is_last:
            async with prompt_server.distributed_tile_jobs_lock:
                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
                    if isinstance(job_data, dict) and 'queue' in job_data:
                        await job_data['queue'].put({
                            'worker_id': worker_id,
                            'is_last': True,
                            'tiles': []
                        })
                        debug_log(f"Received completion signal from worker {worker_id}")
                        return web.json_response({"status": "success"})
                    else:
                        return await handle_api_error(request, "Invalid job structure", 500)
                else:
                    return await handle_api_error(request, "Job not found", 404)
        
        if batch_size > 0:
            # Batch mode: Extract multiple tiles
            padding = int(data.get('padding', 32))
            debug_log(f"UltimateSDUpscale - tile_complete batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}")
            
            # Check for JSON metadata
            metadata_field = data.get('tiles_metadata')
            if metadata_field:
                # New JSON metadata format
                try:
                    # Handle different types of metadata field
                    if hasattr(metadata_field, 'file'):
                        # File-like object
                        metadata_str = metadata_field.file.read().decode('utf-8')
                    elif isinstance(metadata_field, (bytes, bytearray)):
                        # Direct bytes/bytearray
                        metadata_str = metadata_field.decode('utf-8')
                    else:
                        # String
                        metadata_str = str(metadata_field)
                    
                    metadata = json.loads(metadata_str)
                    if len(metadata) != batch_size:
                        return await handle_api_error(request, "Metadata length mismatch", 400)
                    
                    tile_data_list = []  # Temporary list to collect and sort tiles
                    
                    for i in range(batch_size):
                        tile_field = data.get(f'tile_{i}')
                        if tile_field is None:
                            log(f"Missing tile_{i} from worker {worker_id}, skipping")
                            continue
                        
                        try:
                            # Process image
                            img_data = tile_field.file.read()
                            img = Image.open(io.BytesIO(img_data)).convert("RGB")
                            img_np = np.array(img).astype(np.float32) / 255.0
                            tensor = torch.from_numpy(img_np)[None,]
                            
                            # Get metadata for this tile
                            if i < len(metadata):
                                tile_data = metadata[i]
                                tile_idx = tile_data.get('tile_idx', i)
                                
                                tile_info = {
                                    'tensor': tensor,
                                    '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': padding,
                                    'batch_idx': tile_data.get('batch_idx', 0),
                                    'global_idx': tile_data.get('global_idx', tile_idx)
                                }
                                tile_data_list.append(tile_info)
                            else:
                                log(f"Missing metadata for tile {i} from worker {worker_id}, skipping")
                                continue
                                
                        except Exception as e:
                            log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile")
                            continue
                    
                    # Sort tiles by tile_idx to ensure correct order
                    tile_data_list.sort(key=lambda x: x['tile_idx'])
                    tiles.extend(tile_data_list)
                    
                    if tile_data_list:
                        debug_log(f"Reordered {len(tile_data_list)} tiles based on tile_idx")
                except Exception as e:
                    log(f"Error processing JSON metadata from worker {worker_id}: {e}")
                    return await handle_api_error(request, f"Metadata processing error: {e}", 400)
            else:
                # Legacy format: individual fields per tile (backward compatibility)
                debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.")
                
                for i in range(batch_size):
                    tile_field = data.get(f'tile_{i}')
                    if tile_field is None:
                        log(f"Missing tile_{i} from worker {worker_id}, skipping")
                        continue
                    
                    # Get metadata for this tile
                    tile_idx = int(data.get(f'tile_{i}_idx', i))
                    x = int(data.get(f'tile_{i}_x', 0))
                    y = int(data.get(f'tile_{i}_y', 0))
                    extracted_width = int(data.get(f'tile_{i}_width', 512))
                    extracted_height = int(data.get(f'tile_{i}_height', 512))
                    
                    try:
                        # Process image
                        img_data = tile_field.file.read()
                        img = Image.open(io.BytesIO(img_data)).convert("RGB")
                        img_np = np.array(img).astype(np.float32) / 255.0
                        tensor = torch.from_numpy(img_np)[None,]
                        
                        tiles.append({
                            'tensor': tensor,
                            'tile_idx': tile_idx,
                            'x': x,
                            'y': y,
                            'extracted_width': extracted_width,
                            'extracted_height': extracted_height,
                            'padding': padding
                        })
                    except Exception as e:
                        log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile")
                        continue
        else:
            # Single tile mode (backward compatibility)
            image_file = data.get('image')
            if not image_file:
                return await handle_api_error(request, "Missing image data", 400)
                
            tile_idx = int(data.get('tile_idx', 0))
            x = int(data.get('x', 0))
            y = int(data.get('y', 0))
            extracted_width = int(data.get('extracted_width', 512))
            extracted_height = int(data.get('extracted_height', 512))
            padding = int(data.get('padding', 32))
            
            debug_log(f"UltimateSDUpscale - tile_complete single - job_id: {multi_job_id}, worker: {worker_id}, tile: {tile_idx}")
            
            try:
                # Process image
                img_data = image_file.file.read()
                img = Image.open(io.BytesIO(img_data)).convert("RGB")
                img_np = np.array(img).astype(np.float32) / 255.0
                tensor = torch.from_numpy(img_np)[None,]
                
                tiles = [{
                    'tensor': tensor,
                    'tile_idx': tile_idx,
                    'x': x,
                    'y': y,
                    'extracted_width': extracted_width,
                    'extracted_height': extracted_height,
                    'padding': padding,
                    'batch_idx': 0,  # Default for legacy single-tile
                    'global_idx': tile_idx  # Default for legacy single-tile
                }]
            except Exception as e:
                log(f"Error processing tile from worker {worker_id}: {e}")
                return await handle_api_error(request, f"Tile processing error: {e}", 400)

        # Put tiles into queue
        async with prompt_server.distributed_tile_jobs_lock:
            debug_log(f"UltimateSDUpscale - tile_complete: Checking distributed_pending_tile_jobs for job {multi_job_id}")
            debug_log(f"UltimateSDUpscale - tile_complete: Current jobs in distributed_pending_tile_jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}")
            
            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
                if not isinstance(job_data, dict) or 'mode' not in job_data:
                    return await handle_api_error(request, "Invalid job data structure", 500)
                if job_data['mode'] != 'static':
                    return await handle_api_error(request, f"Mode mismatch: expected static, got {job_data['mode']}", 400)
                
                q = job_data['queue']
                if batch_size > 0:
                    # Put batch as single item  
                    await q.put({
                        'worker_id': worker_id,
                        'tiles': tiles,
                        'is_last': is_last
                    })
                    debug_log(f"UltimateSDUpscale - Received batch of {len(tiles)} tiles for job {multi_job_id} from worker {worker_id}")
                else:
                    # Put single tile (backward compat)
                    tile_data = tiles[0]
                    await q.put({
                        'tensor': tile_data['tensor'],
                        'worker_id': worker_id,
                        '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'],
                        'padding': tile_data['padding'],
                        'is_last': is_last
                    })
                    debug_log(f"UltimateSDUpscale - Received single tile {tile_data['tile_idx']} for job {multi_job_id} from worker {worker_id}")
                    
                return web.json_response({"status": "success"})
            else:
                debug_log(f"UltimateSDUpscale - API Error: Job {multi_job_id} not found in distributed_pending_tile_jobs")
                debug_log(f"UltimateSDUpscale - API Error: Available jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}")
                return await handle_api_error(request, "Job not found or already complete", 404)
    except Exception as e:
        return await handle_api_error(request, e, 500)


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

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