diff --git a/distributed_upscale.py b/distributed_upscale.py index a32a7f5..660e8e0 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -33,13 +33,14 @@ from .utils.usdu_utils import ( from .utils.usdu_managment import ( clone_conditioning, ensure_tile_jobs_initialized, # Job management functions - _init_job_queue, _drain_results_queue, + _drain_results_queue, _check_and_requeue_timed_out_workers, _get_completed_count, _mark_task_completed, _send_heartbeat_to_master, _cleanup_job, # Constants JOB_COMPLETED_TASKS, JOB_WORKER_STATUS, JOB_PENDING_TASKS, MAX_PAYLOAD_SIZE ) +from .utils.usdu_managment import init_dynamic_job, init_static_job_batched # Note: MAX_BATCH and HEARTBEAT_TIMEOUT are imported from utils.constants @@ -172,13 +173,18 @@ class UltimateSDUpscaleDistributed: # Determine mode (must match master's logic) enabled_workers = json.loads(enabled_worker_ids) num_workers = len(enabled_workers) - + # Compute number of tiles for this image to decide if tile distribution makes sense + _, height, width, _ = upscaled_image.shape + all_tiles = self.calculate_tiles(width, height, self.round_to_multiple(tile_width), self.round_to_multiple(tile_height), force_uniform_tiles) + num_tiles_per_image = len(all_tiles) + mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) - if batch_size > 1 and num_workers > 0: - debug_log("UltimateSDUpscale Worker - Forcing static mode for USDU-style batched-per-tile processing") + # For USDU-style processing, we want tile distribution whenever workers are available + # and there is more than one tile to process, even if batch == 1. + if num_workers > 0 and num_tiles_per_image > 1: mode = "static" - debug_log(f"UltimateSDUpscale Worker - Mode: {mode}, batch_size: {batch_size}") + debug_log(f"USDU Dist Worker - Batch size {batch_size}") if mode == "dynamic": return self.process_worker_dynamic(upscaled_image, model, positive, negative, vae, @@ -207,11 +213,14 @@ class UltimateSDUpscaleDistributed: # Get image dimensions and batch size batch_size, height, width, _ = upscaled_image.shape - # Calculate all tiles + # Calculate all tiles and grid all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) num_tiles_per_image = len(all_tiles) - - debug_log(f"UltimateSDUpscale Master - Tiles per image: {num_tiles_per_image}") + rows = math.ceil(height / tile_height) + cols = math.ceil(width / tile_width) + log( + f"USDU Dist: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({num_tiles_per_image} tiles/image) | Batch {batch_size}" + ) # Parse enabled workers enabled_workers = json.loads(enabled_worker_ids) @@ -219,11 +228,12 @@ class UltimateSDUpscaleDistributed: # Determine processing mode mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) - if batch_size > 1 and num_workers > 0 and num_tiles_per_image > 1: - debug_log("UltimateSDUpscale Master - Forcing static mode for USDU-style batched-per-tile processing") + # Prefer tile-based static distribution when workers are available and there are multiple tiles, + # even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue. + if num_workers > 0 and num_tiles_per_image > 1: mode = "static" - debug_log(f"UltimateSDUpscale Master - Mode: {mode}, batch_size: {batch_size}, workers: {num_workers}") + log(f"USDU Dist: Workers {num_workers}") if mode == "single_gpu": # No workers, process all tiles locally @@ -275,61 +285,6 @@ class UltimateSDUpscaleDistributed: return result_image - async def _init_job_queue_unified(self, multi_job_id, mode='static', batch_size=None, - total_tiles=None, enabled_workers=None, num_tiles_per_image=None, - assigned_to_workers=None, all_indices=None, task_assignments=None, - batched_static: bool = True): - """Unified job queue initialization for both static and dynamic modes.""" - if mode == 'dynamic': - # Dynamic mode initialization - await _init_job_queue(multi_job_id, 'dynamic', batch_size=batch_size, - all_indices=all_indices, - enabled_workers=list(assigned_to_workers.keys()) if assigned_to_workers else []) - # Add additional fields for backward compatibility - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - job_data['completed_images'] = {} - job_data['completed_tiles'] = {} # For uniformity - # Fix: The endpoint expects 'pending_images', not 'pending_tasks' - job_data['pending_images'] = job_data[JOB_PENDING_TASKS] - debug_log(f"UltimateSDUpscale Master Dynamic - Initialized job queue with {batch_size} pending images for all participants") - else: - # Static mode initialization - # If num_tiles_per_image not provided, calculate it - if num_tiles_per_image is None and total_tiles is not None and batch_size is not None: - num_tiles_per_image = total_tiles // batch_size if batch_size > 0 else total_tiles - - await _init_job_queue(multi_job_id, 'static', - batch_size=batch_size, - num_tiles_per_image=num_tiles_per_image, - all_indices=None, - enabled_workers=enabled_workers, - task_assignments=task_assignments, - batched_static=batched_static) - - # Add completed_tiles and total_items for backward compatibility - 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] - job_data['completed_tiles'] = {} - if total_tiles is not None: - job_data['total_items'] = total_tiles - debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}") - - # Keep compatibility wrappers - async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, all_indices=None): - """Initialize the job queue for dynamic mode with pending images.""" - await self._init_job_queue_unified(multi_job_id, mode='dynamic', batch_size=batch_size, - assigned_to_workers=assigned_to_workers, all_indices=all_indices) - - async def _init_job_queue(self, multi_job_id, total_tiles, enabled_workers, batch_size=1, num_tiles_per_image=None, task_assignments=None, batched_static: bool = True): - """Initialize the job queue for collecting tiles.""" - await self._init_job_queue_unified(multi_job_id, mode='static', batch_size=batch_size, - total_tiles=total_tiles, enabled_workers=enabled_workers, - num_tiles_per_image=num_tiles_per_image, task_assignments=task_assignments, - batched_static=batched_static) async def _async_collect_results(self, multi_job_id, num_workers, mode='static', remaining_to_collect=None, batch_size=None): @@ -355,9 +310,8 @@ class UltimateSDUpscaleDistributed: # Calculate expected count from job data for static mode expected_count = job_data.get('total_items', 0) - len(job_data.get('completed_items', set())) - mode_str = "Dynamic" if mode == 'dynamic' else "Static" item_type = "images" if mode == 'dynamic' else "tiles" - debug_log(f"UltimateSDUpscale Master {mode_str} - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers") + debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers") collected_results = {} workers_done = set() @@ -426,11 +380,11 @@ class UltimateSDUpscaleDistributed: completed_images[image_idx] = image_pil collected_results[image_idx] = image_pil collected_count += 1 - debug_log(f"UltimateSDUpscale Master Dynamic - Received image {image_idx} from worker {worker_id}") + debug_log(f"UltimateSDUpscale Master - Received image {image_idx} from worker {worker_id}") if is_last: workers_done.add(worker_id) - debug_log(f"UltimateSDUpscale Master {mode_str} - Worker {worker_id} completed") + debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed") except asyncio.TimeoutError: if mode == 'dynamic': @@ -440,18 +394,18 @@ class UltimateSDUpscaleDistributed: # Use the class method to check and requeue requeued = await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size) if requeued > 0: - log(f"UltimateSDUpscale Master Dynamic - Requeued {requeued} images from timed out workers") + log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers") last_heartbeat_check = current_time # Check if we've been waiting too long overall if current_time - last_heartbeat_check > timeout: - log(f"UltimateSDUpscale Master Dynamic - Overall timeout waiting for images") + log(f"UltimateSDUpscale Master - Overall timeout waiting for images") break else: log(f"UltimateSDUpscale Master - Timeout waiting for {item_type}") break - debug_log(f"UltimateSDUpscale Master {mode_str} - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers") + debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers") # Clean up job queue async with prompt_server.distributed_tile_jobs_lock: @@ -833,7 +787,7 @@ class UltimateSDUpscaleDistributed: sliced_conditioning_cache = {} # Dynamic queue mode (static processing): process batched-per-tile - debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} ready to process tiles dynamically (batched per tile)") + log(f"USDU Dist Worker[{worker_id[:8]}]: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Tiles/image {num_tiles_per_image} | Batch {batch_size}") processed_count = 0 # Poll for job readiness @@ -844,7 +798,7 @@ class UltimateSDUpscaleDistributed: timeout=5.0 ) if ready: - debug_log(f"Job {multi_job_id} ready after {attempt} attempts") + debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts") break time.sleep(1.0) else: @@ -860,11 +814,11 @@ class UltimateSDUpscaleDistributed: ) if tile_idx is None: - debug_log(f"UltimateSDUpscale Worker - No more tiles to process") + debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process") break # Always batched-per-tile in static mode - debug_log(f"UltimateSDUpscale Worker - Assigned tile_id {tile_idx} to worker {worker_id}") + debug_log(f"Worker[{worker_id[:8]}] - Assigned tile_id {tile_idx}") processed_count += batch_size tile_id = tile_idx tx, ty = all_tiles[tile_id] @@ -900,15 +854,15 @@ class UltimateSDUpscaleDistributed: timeout=5.0 ) except Exception as e: - debug_log(f"Heartbeat failed: {e}") + debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}") # Send tiles in batches - if len(processed_tiles) >= MAX_BATCH: - run_async_in_server_loop( - self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), - timeout=TILE_SEND_TIMEOUT - ) - processed_tiles = [] + if len(processed_tiles) >= MAX_BATCH: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] # Send any remaining tiles if processed_tiles: @@ -1004,13 +958,14 @@ class UltimateSDUpscaleDistributed: sliced_conditioning_cache = {} # Initialize queue: pending queue holds tile ids (batched per tile) - debug_log(f"UltimateSDUpscale Master - Using dynamic tile distribution") + log("USDU Dist: Using tile queue distribution") run_async_in_server_loop( - self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image, batched_static=True), + init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers), timeout=10.0 ) - debug_log(f"UltimateSDUpscale Master - Initialized batched static job queue with {num_tiles_per_image} tile ids for dynamic distribution") - debug_log(f"UltimateSDUpscale Master - Processing tiles dynamically from queue") + debug_log( + f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}" + ) processed_count = 0 consecutive_no_tile = 0 @@ -1047,6 +1002,7 @@ class UltimateSDUpscaleDistributed: _mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}), timeout=5.0 ) + log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})") else: consecutive_no_tile += 1 if consecutive_no_tile >= max_consecutive_no_tile: @@ -1111,7 +1067,7 @@ class UltimateSDUpscaleDistributed: local_tile_idx, all_tiles[local_tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx + padding, mask_blur, width, height, force_uniform_tiles, tiled_decode, batch_idx=batch_idx ) # Mark as completed and store in collected_tasks @@ -1222,7 +1178,7 @@ class UltimateSDUpscaleDistributed: return # Early exit if empty total_tiles = len(processed_tiles) - debug_log(f"UltimateSDUpscale Worker - Preparing to send {total_tiles} tiles") + debug_log(f"Worker[{worker_id[:8]}] - Preparing to send {total_tiles} tiles (size-aware chunks)") # Prepare encoded images and sizes to enable size-aware chunking encoded = [] @@ -1303,7 +1259,7 @@ class UltimateSDUpscaleDistributed: log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}") raise - debug_log(f"UltimateSDUpscale Worker - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)") + debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)") chunk_index += 1 i = j @@ -1321,7 +1277,11 @@ class UltimateSDUpscaleDistributed: # 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}, batched per tile") + rows = math.ceil(height / tile_height) + cols = math.ceil(width / tile_width) + log( + f"USDU Dist: Single GPU | Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({len(all_tiles)} tiles/image) | Batch {batch_size}" + ) # Prepare result images list result_images = [] @@ -1368,27 +1328,24 @@ class UltimateSDUpscaleDistributed: 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") - + log(f"USDU Dist: Image queue distribution | Batch {batch_size} | Workers {num_workers} | Canvas {width}x{height} | Tile {tile_width}x{tile_height}") + # No fixed share - all images are dynamic all_indices = list(range(batch_size)) - log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.") + debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.") # Calculate tiles for processing all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) # Initialize job queue for communication try: - assigned_to_workers = {w: [] for w in enabled_workers} - worker_status = {w: time.time() for w in enabled_workers} - run_async_in_server_loop( - self._init_job_queue_dynamic(multi_job_id, batch_size, assigned_to_workers, worker_status, all_indices), + init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices), timeout=2.0 ) except Exception as e: - debug_log(f"UltimateSDUpscale Master Dynamic - Queue initialization error: {e}") + debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}") raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}") # Convert batch to PIL list @@ -1472,7 +1429,7 @@ class UltimateSDUpscaleDistributed: self._get_total_completed_count(multi_job_id), timeout=1.0 ) - log(f"Progress: {completed_now}/{batch_size} images completed") + log(f"USDU Dist: Images progress {completed_now}/{batch_size}") # Yield to allow workers to get new images after completing one run_async_in_server_loop(self._async_yield(), timeout=0.1) @@ -1502,7 +1459,7 @@ class UltimateSDUpscaleDistributed: timeout=1.0 ) - log(f"Progress: {completed_now}/{batch_size} images completed") + log(f"USDU Dist: Images progress {completed_now}/{batch_size}") if completed_now >= batch_size: break @@ -1559,7 +1516,7 @@ class UltimateSDUpscaleDistributed: if upscaled_image.is_cuda: result_tensor = result_tensor.cuda() - debug_log(f"UltimateSDUpscale Master Dynamic - Job {multi_job_id} complete") + debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") log(f"Completed processing all {batch_size} images") return (result_tensor,) @@ -1665,128 +1622,127 @@ class UltimateSDUpscaleDistributed: 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), + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold): + """Worker processing in dynamic mode - processes whole images.""" + # Round tile dimensions + tile_width = self.round_to_multiple(tile_width) + tile_height = self.round_to_multiple(tile_height) + + # Get dimensions and tile grid + batch_size, height, width, _ = upscaled_image.shape + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + log(f"USDU Dist Worker[{worker_id[:8]}]: Processing image queue | Batch {batch_size}") + + # Keep track of processed images for is_last detection + processed_count = 0 + + # Poll for job readiness to avoid races during master init + max_poll_attempts = 20 # ~20s at 1s sleep + for attempt in range(max_poll_attempts): + ready = run_async_in_server_loop( + self._check_job_status(multi_job_id, master_url), + timeout=5.0 + ) + if ready: + debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts") + break + time.sleep(1.0) # Poll every 1s + else: + log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") + return (upscaled_image,) + + # Loop to request and process images + while True: + # Request an image to process + image_idx, estimated_remaining = run_async_in_server_loop( + self._request_image_from_master(multi_job_id, master_url, worker_id), + timeout=TILE_WAIT_TIMEOUT + ) + + if image_idx is None: + debug_log(f"USDU Dist Worker - No more images to process") + break + + debug_log(f"Worker[{worker_id[:8]}] - Assigned image {image_idx}") + processed_count += 1 + + # Determine if this should be marked as last for this worker + is_last_for_worker = (estimated_remaining == 0) + + # Extract single image tensor + single_tensor = upscaled_image[image_idx:image_idx+1] + + # Convert to PIL for processing + local_image = tensor_to_pil(single_tensor, 0).copy() + + # Process all tiles for this image + image_seed = seed + image_idx * 1000 + + # Pre-slice conditioning once per image (not per tile) + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[image_idx:image_idx+1] + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[image_idx:image_idx+1] + control = control.previous_controlnet + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] + + for tile_idx, pos in enumerate(all_tiles): + local_image = self._process_and_blend_tile( + tile_idx, pos, single_tensor, local_image, + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, force_uniform_tiles, + tiled_decode, batch_idx=image_idx + ) + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), timeout=5.0 ) - 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 - - # Pre-slice conditioning once per image (not per tile) - positive_sliced = clone_conditioning(positive) - negative_sliced = clone_conditioning(negative) - for cond_list in [positive_sliced, negative_sliced]: - for i in range(len(cond_list)): - emb, cond_dict = cond_list[i] - if emb.shape[0] > 1: - cond_list[i][0] = emb[image_idx:image_idx+1] - if 'control' in cond_dict: - control = cond_dict['control'] - while control is not None: - hint = control.cond_hint_original - if hint.shape[0] > 1: - control.cond_hint_original = hint[image_idx:image_idx+1] - control = control.previous_controlnet - if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: - cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] - - for tile_idx, pos in enumerate(all_tiles): - local_image = self._process_and_blend_tile( - tile_idx, pos, single_tensor, local_image, - model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, - sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, force_uniform_tiles, - tiled_decode, batch_idx=image_idx - ) - run_async_in_server_loop( - _send_heartbeat_to_master(multi_job_id, master_url, worker_id), - timeout=5.0 - ) - - # Send processed image back to master - try: - # Use the estimated remaining to determine if this is the last image - is_last = is_last_for_worker - run_async_in_server_loop( - self._send_full_image_to_master(local_image, image_idx, multi_job_id, - master_url, worker_id, is_last), - timeout=TILE_SEND_TIMEOUT - ) - # Send heartbeat after processing - run_async_in_server_loop( - _send_heartbeat_to_master(multi_job_id, master_url, worker_id), - timeout=5.0 - ) - if is_last: - break - except Exception as e: - log(f"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") + + # 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_worker_complete_signal(multi_job_id, master_url, worker_id), + self._send_full_image_to_master(local_image, image_idx, multi_job_id, + master_url, worker_id, is_last), timeout=TILE_SEND_TIMEOUT ) + # Send heartbeat after processing + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + if is_last: + break except Exception as e: - log(f"UltimateSDUpscale Worker Dynamic - Error sending completion signal: {e}") - - return (upscaled_image,) + log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}") + # Continue processing other images + + # Send final is_last signal + debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal") + try: + run_async_in_server_loop( + self._send_worker_complete_signal(multi_job_id, master_url, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + except Exception as e: + log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending completion signal: {e}") + + return (upscaled_image,) async def _request_image_from_master(self, multi_job_id, master_url, worker_id): """Request an image index to process from master in dynamic mode.""" @@ -1957,15 +1913,13 @@ class UltimateSDUpscaleDistributed: def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str: """Determines processing mode per requested policy: - - batch_size > 1 => static (tile-based) - - batch_size == 1 => dynamic (per-image) + - any workers => prefer static (tile-based) for USDU - no workers => single_gpu """ if num_workers == 0: return "single_gpu" - return "static" if batch_size > 1 else "dynamic" - - # Note: Using _distribute_tasks from usdu_managment.py instead of duplicate _distribute_items + # Default to static when distributed; master/worker may still override if special cases arise + return "static" # Ensure initialization before registering routes ensure_tile_jobs_initialized() diff --git a/utils/usdu_managment.py b/utils/usdu_managment.py index 97d30e4..0a97ff3 100644 --- a/utils/usdu_managment.py +++ b/utils/usdu_managment.py @@ -7,13 +7,11 @@ import io from aiohttp import web import server from PIL import Image -import numpy as np -import torch # Import from other utilities from .logging import debug_log, log from .network import handle_api_error, get_client_session -from .image import tensor_to_pil +from .image import pil_to_tensor # 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))) @@ -22,6 +20,87 @@ MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 from .constants import HEARTBEAT_TIMEOUT +def _parse_tiles_from_form(data): + """Parse tiles submitted via multipart/form-data into a list of tile dicts. + + Expects the following fields in the aiohttp form data: + - 'tiles_metadata': JSON list with per-tile metadata items containing at least + 'tile_idx', 'x', 'y', 'extracted_width', 'extracted_height'. Optional + 'batch_idx' and 'global_idx' are included when available. + - 'tile_{i}': image bytes for each tile described in tiles_metadata (PNG). + - 'padding': integer padding used during extraction (optional; defaults 0). + + Returns: list of dicts with keys: 'tensor', 'tile_idx', 'x', 'y', + 'extracted_width', 'extracted_height', and optional 'batch_idx', 'global_idx', + plus 'padding'. + """ + try: + # Parse padding if present + padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0 + except Exception: + padding = 0 + + # Parse tiles metadata (JSON list) + meta_raw = data.get('tiles_metadata') + if meta_raw is None: + raise ValueError("Missing tiles_metadata") + + try: + metadata = json.loads(meta_raw) + except Exception as e: + raise ValueError(f"Invalid tiles_metadata JSON: {e}") + + if not isinstance(metadata, list): + raise ValueError("tiles_metadata must be a list") + + tiles = [] + # Iterate over metadata items and corresponding uploaded files tile_0, tile_1, ... + for i, meta in enumerate(metadata): + file_field = data.get(f'tile_{i}') + if file_field is None or not hasattr(file_field, 'file'): + raise ValueError(f"Missing tile data for index {i}") + + # Read image bytes and decode to PIL + raw = file_field.file.read() + try: + img = Image.open(io.BytesIO(raw)).convert("RGB") + except Exception as e: + raise ValueError(f"Invalid image data for tile {i}: {e}") + + # Convert to tensor [1, H, W, C] + tile_tensor = pil_to_tensor(img) + + # Build tile dictionary + try: + tile_info = { + 'tensor': tile_tensor, + 'tile_idx': int(meta.get('tile_idx', i)), + 'x': int(meta.get('x', 0)), + 'y': int(meta.get('y', 0)), + 'extracted_width': int(meta.get('extracted_width', img.width)), + 'extracted_height': int(meta.get('extracted_height', img.height)), + 'padding': int(padding), + } + except Exception as e: + raise ValueError(f"Invalid metadata values for tile {i}: {e}") + + # Optional fields + if 'batch_idx' in meta: + try: + tile_info['batch_idx'] = int(meta['batch_idx']) + except Exception: + pass + if 'global_idx' in meta: + try: + tile_info['global_idx'] = int(meta['global_idx']) + except Exception: + pass + + tiles.append(tile_info) + + return tiles + + # Unified Job Data Structure Keys JOB_QUEUE = 'queue' JOB_MODE = 'mode' @@ -36,6 +115,45 @@ JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static TASK_TYPE_TILE = 'tile' TASK_TYPE_IMAGE = 'image' +from typing import List, Optional + +async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: List[str], all_indices: Optional[List[int]] = None): + """Initialize queue for dynamic mode (per-image), with collector fields. + + - Creates JOB_PENDING_TASKS with image indices + - Adds 'completed_images' dict and 'pending_images' alias used by collectors + """ + await _init_job_queue( + multi_job_id, + 'dynamic', + batch_size=batch_size, + all_indices=all_indices or list(range(batch_size)), + enabled_workers=enabled_workers, + ) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + job_data['completed_images'] = {} + job_data['pending_images'] = job_data[JOB_PENDING_TASKS] + debug_log(f"Dynamic job {multi_job_id} initialized with {batch_size} images") + + +async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_per_image: int, enabled_workers: List[str]): + """Initialize queue for static mode (batched-per-tile). + + - Populates JOB_PENDING_TASKS with tile ids [0..num_tiles_per_image-1] + """ + await _init_job_queue( + multi_job_id, + 'static', + batch_size=batch_size, + num_tiles_per_image=num_tiles_per_image, + enabled_workers=enabled_workers, + batched_static=True, + ) + debug_log(f"Static job {multi_job_id} initialized with {num_tiles_per_image} tile ids for batch {batch_size}") + async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None, batched_static: bool = False): """Unified initialization for job queues in static and dynamic modes.""" prompt_server = ensure_tile_jobs_initialized() @@ -87,34 +205,7 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data -def _distribute_tasks(items: list, num_participants: int) -> list[list[any]]: - """Distribute a list of items among N participants (master + workers).""" - if num_participants <= 1: - return [items] - - items_per_participant = len(items) // num_participants - remainder = len(items) % num_participants - assignments = [] - start_idx = 0 - for i in range(num_participants): - count = items_per_participant + (1 if i < remainder else 0) - end_idx = start_idx + count - assignments.append(items[start_idx:end_idx]) - start_idx = end_idx - return assignments - -async def _get_next_task(multi_job_id): - """Get next task from pending queue (generalized for tiles/images).""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if not job_data or JOB_PENDING_TASKS not in job_data: - return None - try: - task_id = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0) - return task_id - except asyncio.TimeoutError: - return None +# Note: legacy task distribution and queue pull helpers removed async def _drain_results_queue(multi_job_id): """Drain pending results from queue and update completed_tasks. Returns count drained.""" @@ -331,108 +422,10 @@ async def submit_tiles_endpoint(request): debug_log(f"Received completion signal from worker {worker_id}") return web.json_response({"status": "success"}) - # Handle batch tiles with metadata - if batch_size > 0: - padding = int(data.get('padding', 32)) - metadata_field = data.get('tiles_metadata') - if metadata_field: - if hasattr(metadata_field, 'file'): - metadata_str = metadata_field.file.read().decode('utf-8') - elif isinstance(metadata_field, (bytes, bytearray)): - metadata_str = metadata_field.decode('utf-8') - else: - 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 = [] - - for i in range(batch_size): - tile_field = data.get(f'tile_{i}') - if tile_field is None: - continue - - 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,] - - if i < len(metadata): - tile_meta = metadata[i] - tile_idx = tile_meta.get('tile_idx', i) - tile_info = { - 'tensor': tensor, - 'tile_idx': tile_idx, - 'x': tile_meta['x'], - 'y': tile_meta['y'], - 'extracted_width': tile_meta['extracted_width'], - 'extracted_height': tile_meta['extracted_height'], - 'padding': padding, - 'batch_idx': tile_meta.get('batch_idx', 0), - 'global_idx': tile_meta.get('global_idx', tile_idx) - } - tile_data_list.append(tile_info) - - tile_data_list.sort(key=lambda x: x['tile_idx']) - tiles.extend(tile_data_list) - else: - # Legacy format - for i in range(batch_size): - tile_field = data.get(f'tile_{i}') - if tile_field is None: - continue - - 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)) - - 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 - }) - else: - # Single tile legacy - 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)) - - 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, - 'global_idx': tile_idx - }] + try: + tiles = _parse_tiles_from_form(data) + except ValueError as e: + return await handle_api_error(request, str(e), 400) # Submit tiles to queue async with prompt_server.distributed_tile_jobs_lock: @@ -448,6 +441,7 @@ async def submit_tiles_endpoint(request): 'tiles': tiles, 'is_last': is_last }) + debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})") else: await q.put({ 'worker_id': worker_id, @@ -523,188 +517,7 @@ async def submit_image_endpoint(request): except Exception as e: return await handle_api_error(request, e, 500) -# Keep legacy endpoint for backward compatibility -@server.PromptServer.instance.routes.post("/distributed/tile_complete") -async def tile_complete_endpoint(request): - try: - content_length = request.headers.get('content-length') - if content_length and int(content_length) > MAX_PAYLOAD_SIZE: - return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) - - data = await request.post() - multi_job_id = data.get('multi_job_id') - worker_id = data.get('worker_id') - is_last = data.get('is_last', 'False').lower() == 'true' - - if multi_job_id is None or worker_id is None: - return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - - if 'full_image' in data and 'image_idx' in data: - image_idx = int(data.get('image_idx')) - img_data = data['full_image'].file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - - debug_log(f"Received full image {image_idx} from worker {worker_id}") - - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic': - return await handle_api_error(request, "Mode mismatch for image submission", 400) - if JOB_QUEUE in job_data: - await job_data[JOB_QUEUE].put({ - 'worker_id': worker_id, - 'image_idx': image_idx, - 'image': img, - 'is_last': is_last - }) - return web.json_response({"status": "success"}) - - batch_size = int(data.get('batch_size', 0)) - tiles = [] - - if batch_size == 0 and is_last: - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if JOB_QUEUE in job_data: - await job_data[JOB_QUEUE].put({ - 'worker_id': worker_id, - 'is_last': True, - 'tiles': [] - }) - debug_log(f"Received completion signal from worker {worker_id}") - return web.json_response({"status": "success"}) - - if batch_size > 0: - padding = int(data.get('padding', 32)) - metadata_field = data.get('tiles_metadata') - if metadata_field: - if hasattr(metadata_field, 'file'): - metadata_str = metadata_field.file.read().decode('utf-8') - elif isinstance(metadata_field, (bytes, bytearray)): - metadata_str = metadata_field.decode('utf-8') - else: - 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 = [] - - for i in range(batch_size): - tile_field = data.get(f'tile_{i}') - if tile_field is None: - continue - - 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,] - - if i < len(metadata): - tile_meta = metadata[i] - tile_idx = tile_meta.get('tile_idx', i) - tile_info = { - 'tensor': tensor, - 'tile_idx': tile_idx, - 'x': tile_meta['x'], - 'y': tile_meta['y'], - 'extracted_width': tile_meta['extracted_width'], - 'extracted_height': tile_meta['extracted_height'], - 'padding': padding, - 'batch_idx': tile_meta.get('batch_idx', 0), - 'global_idx': tile_meta.get('global_idx', tile_idx) - } - tile_data_list.append(tile_info) - - tile_data_list.sort(key=lambda x: x['tile_idx']) - tiles.extend(tile_data_list) - else: - # Legacy - for i in range(batch_size): - tile_field = data.get(f'tile_{i}') - if tile_field is None: - continue - - 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)) - - 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 - }) - else: - # Single tile legacy - 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)) - - 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, - 'global_idx': tile_idx - }] - - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if JOB_MODE in job_data and job_data[JOB_MODE] != 'static': - return await handle_api_error(request, "Mode mismatch for tile submission", 400) - - q = job_data[JOB_QUEUE] - if batch_size > 0 or len(tiles) > 0: - await q.put({ - 'worker_id': worker_id, - 'tiles': tiles, - 'is_last': is_last - }) - else: - await q.put({ - 'worker_id': worker_id, - 'is_last': True, - 'tiles': [] - }) - - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - return await handle_api_error(request, e, 500) +# Note: Removed legacy /distributed/tile_complete endpoint. Use /distributed/submit_tiles.