diff --git a/distributed_upscale.py b/distributed_upscale.py index 6be5de0..371050c 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -1,3 +1,15 @@ +""" +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 @@ -7,6 +19,9 @@ import aiohttp from aiohttp import web import io import server +import math +import os +import time from typing import List, Tuple # Import ComfyUI modules @@ -20,9 +35,15 @@ from .utils.network import get_server_port, get_server_loop, get_client_session, 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, MAX_BATCH + 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.""" @@ -31,6 +52,15 @@ def ensure_tile_jobs_initialized(): 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 + to_remove = [] + 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: + debug_log(f"Removing legacy queue structure for job {job_id}") + to_remove.append(job_id) + for job_id in to_remove: + 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 @@ -38,7 +68,23 @@ def ensure_tile_jobs_initialized(): class UltimateSDUpscaleDistributed: """ Distributed version of Ultimate SD Upscale (No Upscale). - Distributes tile processing across multiple workers. + + 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): @@ -67,6 +113,7 @@ class UltimateSDUpscaleDistributed: "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": ""}), @@ -75,6 +122,7 @@ class UltimateSDUpscaleDistributed: "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}), }, } @@ -106,67 +154,106 @@ class UltimateSDUpscaleDistributed: 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, multi_job_id="", is_worker=False, - master_url="", enabled_worker_ids="[]", worker_id="", tile_indices=""): + 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) + 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, multi_job_id, master_url, - worker_id, enabled_worker_ids) + 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, multi_job_id, enabled_worker_ids) + 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, multi_job_id, master_url, - worker_id, enabled_worker_ids): - """Worker processing - SYNCHRONOUS tile processing with async communication.""" + 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) + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - # Get assigned tiles for this worker - assigned_tiles = self._get_worker_tiles(all_tiles, enabled_worker_ids, worker_id) - if not assigned_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_tiles)} tiles") + 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 tile_idx in assigned_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 + # Extract tile from the specific image in the batch tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( - upscaled_image, x, y, tile_width, tile_height, padding + upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding ) - # Process tile through SD (SYNCHRONOUS - no async!) + # Process tile through SD with unique seed + image_seed = seed + batch_idx * 1000 processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, - seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise) + 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, @@ -189,17 +276,17 @@ class UltimateSDUpscaleDistributed: 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, multi_job_id, enabled_worker_ids): + 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 - _, height, width, _ = upscaled_image.shape + # 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) + 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}") @@ -208,11 +295,44 @@ class UltimateSDUpscaleDistributed: enabled_workers = json.loads(enabled_worker_ids) num_workers = len(enabled_workers) + # Synchronous queue initialization before mode determination + import threading + def sync_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] = asyncio.Queue() + debug_log(f"Synchronous queue init for {multi_job_id}") + thread = threading.Thread(target=sync_init) + thread.start() + thread.join(timeout=5.0) # Block with timeout + if thread.is_alive(): + log(f"Queue init timed out for {multi_job_id}") + raise RuntimeError("Queue init failed") + + # 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) + 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 @@ -235,33 +355,51 @@ class UltimateSDUpscaleDistributed: # loop = asyncio.new_event_loop() # asyncio.set_event_loop(loop) - # Prepare result image - image_pil = tensor_to_pil(upscaled_image, 0) - result_image = image_pil.copy() + # 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 master's tiles - master_tiles_indices = self._get_master_tiles(all_tiles, num_workers) - master_tiles = len(master_tiles_indices) + # Calculate total tiles across all images + num_tiles_per_image = len(all_tiles) + total_tiles_all_images = batch_size * num_tiles_per_image - debug_log(f"UltimateSDUpscale Master - Processing {master_tiles} tiles locally") + # 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 tile_idx in master_tiles_indices: - result_image = self._process_and_blend_tile( - tile_idx, all_tiles[tile_idx], upscaled_image, result_image, - model, positive, negative, vae, seed, steps, cfg, + 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 + padding, mask_blur, width, height, tiled_decode ) # Collect worker tiles using async - worker_tiles_expected = total_tiles - master_tiles + 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 // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) + 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 = [] @@ -274,7 +412,7 @@ class UltimateSDUpscaleDistributed: worker_start_idx += remainder - 1 if remainder > 0 else 0 worker_tile_count = tiles_per_participant - if worker_start_idx < total_tiles and worker_tile_count > 0: + 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") @@ -296,6 +434,14 @@ class UltimateSDUpscaleDistributed: 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 @@ -303,17 +449,53 @@ class UltimateSDUpscaleDistributed: # 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_image = self.blend_tile(result_image, tile_pil, + 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 - result_tensor = pil_to_tensor(result_image) + # 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) @@ -344,6 +526,13 @@ class UltimateSDUpscaleDistributed: 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) @@ -355,7 +544,7 @@ class UltimateSDUpscaleDistributed: 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): + padding, mask_blur, image_width, image_height, tiled_decode): """Process a single tile and blend it into the result image.""" x, y = tile_pos @@ -366,7 +555,7 @@ class UltimateSDUpscaleDistributed: processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise) + scheduler, denoise, tiled_decode) # Convert and blend processed_pil = tensor_to_pil(processed_tile, 0) @@ -389,6 +578,29 @@ class UltimateSDUpscaleDistributed: 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() @@ -396,7 +608,11 @@ class UltimateSDUpscaleDistributed: 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] = asyncio.Queue() + 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: @@ -409,7 +625,12 @@ class UltimateSDUpscaleDistributed: 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}") - q = prompt_server.distributed_pending_tile_jobs[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() @@ -431,9 +652,18 @@ class UltimateSDUpscaleDistributed: 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[tile_idx] = { + collected_tiles[key] = { 'tensor': tile_data['tensor'], 'tile_idx': tile_idx, 'x': tile_data['x'], @@ -441,7 +671,9 @@ class UltimateSDUpscaleDistributed: 'extracted_width': tile_data['extracted_width'], 'extracted_height': tile_data['extracted_height'], 'padding': tile_data['padding'], - 'worker_id': worker_id + '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) @@ -470,21 +702,160 @@ class UltimateSDUpscaleDistributed: 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) -> List[Tuple[int, 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.""" - 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 + 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]: @@ -524,11 +895,20 @@ class UltimateSDUpscaleDistributed: 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) -> torch.Tensor: + 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) @@ -538,14 +918,20 @@ class UltimateSDUpscaleDistributed: clean_tensor = clean_tensor.cuda() # Encode to latent - latent = VAEEncode().encode(vae, clean_tensor)[0] + 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 - image = VAEDecode().decode(vae, samples)[0] + if tiled_decode and tiled_vae_available: + image = VAEDecodeTiled().decode(vae, samples)[0] + else: + image = VAEDecode().decode(vae, samples)[0] return image @@ -572,7 +958,7 @@ class UltimateSDUpscaleDistributed: def blend_tile(self, base_image: Image.Image, tile_image: Image.Image, x: int, y: int, extracted_size: Tuple[int, int], - mask: Image.Image, padding: int) -> Image.Image: + 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 @@ -635,16 +1021,22 @@ class UltimateSDUpscaleDistributed: data.add_field('batch_size', str(chunk_size)) data.add_field('padding', str(padding)) - # Chunk metadata: Keep absolute tile_idx + # Chunk metadata: Keep absolute tile_idx and batch_idx if present metadata = [] for tile_data in chunk: - metadata.append({ + 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') @@ -696,51 +1088,493 @@ class UltimateSDUpscaleDistributed: 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): + 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 - _, height, width, _ = upscaled_image.shape + # 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) + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally") + debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally for batch of {batch_size}") - # Convert to PIL - image_pil = tensor_to_pil(upscaled_image, 0) - result_image = image_pil.copy() + # 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()) - # Mask will be created per tile for proper blending - - # Process each tile - for idx, tile_pos in enumerate(all_tiles): - result_image = self._process_and_blend_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, width, height - ) + # 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 - result_tensor = pil_to_tensor(result_image) + 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') @@ -751,10 +1585,68 @@ async def tile_complete_endpoint(request): 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)) @@ -811,7 +1703,9 @@ async def tile_complete_endpoint(request): 'y': tile_data['y'], 'extracted_width': tile_data['extracted_width'], 'extracted_height': tile_data['extracted_height'], - 'padding': padding + '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: @@ -896,7 +1790,9 @@ async def tile_complete_endpoint(request): 'y': y, 'extracted_width': extracted_width, 'extracted_height': extracted_height, - 'padding': padding + '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}") @@ -908,9 +1804,16 @@ async def tile_complete_endpoint(request): 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 prompt_server.distributed_pending_tile_jobs[multi_job_id].put({ + await q.put({ 'worker_id': worker_id, 'tiles': tiles, 'is_last': is_last @@ -919,7 +1822,7 @@ async def tile_complete_endpoint(request): else: # Put single tile (backward compat) tile_data = tiles[0] - await prompt_server.distributed_pending_tile_jobs[multi_job_id].put({ + await q.put({ 'tensor': tile_data['tensor'], 'worker_id': worker_id, 'tile_idx': tile_data['tile_idx'],