diff --git a/distributed_upscale.py b/distributed_upscale.py index 660e8e0..130dae5 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -277,7 +277,7 @@ class UltimateSDUpscaleDistributed: # Convert and blend processed_pil = tensor_to_pil(processed_tile, 0) - # Create mask for this specific tile + # Create mask for this specific tile (no cache here; only used in single-tile path) tile_mask = self.create_tile_mask(image_width, image_height, x, y, tile_width, tile_height, mask_blur) # Use extraction position and size for blending result_image = self.blend_tile(result_image, processed_pil, @@ -353,9 +353,8 @@ class UltimateSDUpscaleDistributed: # 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_results[key] = { - 'tensor': tile_data['tensor'], + # Store the full tile data including metadata; prefer PIL image if present + entry = { 'tile_idx': tile_idx, 'x': tile_data['x'], 'y': tile_data['y'], @@ -366,6 +365,11 @@ class UltimateSDUpscaleDistributed: 'batch_idx': tile_data.get('batch_idx', 0), 'global_idx': tile_data.get('global_idx', tile_idx) } + if 'image' in tile_data: + entry['image'] = tile_data['image'] + elif 'tensor' in tile_data: + entry['tensor'] = tile_data['tensor'] + collected_results[key] = entry else: # Single tile mode (backward compat) tile_idx = result['tile_idx'] @@ -935,7 +939,7 @@ class UltimateSDUpscaleDistributed: break # Wait a bit before next check - await asyncio.sleep(0.5) + await asyncio.sleep(0.1) # Get all completed tasks for return return await self._get_all_completed_tasks(multi_job_id) @@ -967,9 +971,14 @@ class UltimateSDUpscaleDistributed: f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}" ) + # Precompute masks for all tile positions to avoid repeated Gaussian blur work during blending + tile_masks = [] + for idx, (tx, ty) in enumerate(all_tiles): + tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)) + processed_count = 0 consecutive_no_tile = 0 - max_consecutive_no_tile = 5 + max_consecutive_no_tile = 2 while processed_count < total_tiles: comfy.model_management.throw_exception_if_processing_interrupted() @@ -991,7 +1000,7 @@ class UltimateSDUpscaleDistributed: seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, region, (width, height) ) - tile_mask = self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur) + tile_mask = tile_masks[tile_id] for b in range(batch_size): tile_pil = tensor_to_pil(processed_batch, b) if tile_pil.size != (ew, eh): @@ -1087,7 +1096,7 @@ class UltimateSDUpscaleDistributed: # Blend worker tiles synchronously for global_idx, tile_data in collected_tasks.items(): # Skip tiles that don't have tensor data (already processed) - if 'tensor' not in tile_data: + if 'tensor' not in tile_data and 'image' not in tile_data: continue batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image) @@ -1099,10 +1108,14 @@ class UltimateSDUpscaleDistributed: # Blend tile synchronously x = tile_data.get('x', 0) y = tile_data.get('y', 0) - tile_tensor = tile_data['tensor'] - tile_pil = tensor_to_pil(tile_tensor, 0) + # Prefer PIL image if present to avoid reconversion + if 'image' in tile_data: + tile_pil = tile_data['image'] + else: + tile_tensor = tile_data['tensor'] + tile_pil = tensor_to_pil(tile_tensor, 0) orig_x, orig_y = all_tiles[tile_idx] - tile_mask = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur) + tile_mask = tile_masks[tile_idx] extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding) extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding) result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil, @@ -1546,7 +1559,7 @@ class UltimateSDUpscaleDistributed: if not job_data or JOB_PENDING_TASKS not in job_data: return None try: - tile_idx = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0) + tile_idx = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=0.1) return tile_idx except asyncio.TimeoutError: return None diff --git a/utils/usdu_managment.py b/utils/usdu_managment.py index 0a97ff3..f6f00ed 100644 --- a/utils/usdu_managment.py +++ b/utils/usdu_managment.py @@ -67,13 +67,14 @@ def _parse_tiles_from_form(data): except Exception as e: raise ValueError(f"Invalid image data for tile {i}: {e}") - # Convert to tensor [1, H, W, C] + # Convert to tensor [1, H, W, C] (keep PIL too for faster blending later) tile_tensor = pil_to_tensor(img) # Build tile dictionary try: tile_info = { 'tensor': tile_tensor, + 'image': img, # keep PIL to avoid reconversion later 'tile_idx': int(meta.get('tile_idx', i)), 'x': int(meta.get('x', 0)), 'y': int(meta.get('y', 0)), @@ -136,7 +137,7 @@ async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: 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") + debug_log(f"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]): @@ -152,7 +153,7 @@ async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_ 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}") + debug_log(f"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.""" @@ -174,9 +175,9 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima if mode == 'dynamic': job_data[JOB_BATCH_SIZE] = batch_size pending_queue = job_data[JOB_PENDING_TASKS] - for i in all_indices or range(batch_size): + for i in (all_indices or range(batch_size)): await pending_queue.put(i) - debug_log(f"Initialized dynamic queue with {batch_size} pending images") + debug_log(f"Initialized image queue with {batch_size} pending items") elif mode == 'static': job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image job_data[JOB_BATCH_SIZE] = batch_size @@ -186,12 +187,12 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima if batched_static and num_tiles_per_image is not None: for i in range(num_tiles_per_image): await pending_queue.put(i) - debug_log(f"Initialized static queue (batched) with {num_tiles_per_image} tile ids for dynamic distribution across batch of {batch_size}") + debug_log(f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}") else: total_tiles = batch_size * num_tiles_per_image for i in range(total_tiles): await pending_queue.put(i) - debug_log(f"Initialized static queue with {total_tiles} pending tiles for dynamic distribution") + debug_log(f"Initialized tile queue with {total_tiles} pending tasks") # Keep backward compatibility - if task assignments provided, still track them if task_assignments and enabled_workers: @@ -412,7 +413,7 @@ async def submit_tiles_endpoint(request): 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: expected static mode", 400) + return await handle_api_error(request, "Job not configured for tile submissions", 400) if JOB_QUEUE in job_data: await job_data[JOB_QUEUE].put({ 'worker_id': worker_id, @@ -432,7 +433,7 @@ async def submit_tiles_endpoint(request): 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: expected static mode", 400) + return await handle_api_error(request, "Job not configured for tile submissions", 400) q = job_data[JOB_QUEUE] if batch_size > 0 or len(tiles) > 0: @@ -485,7 +486,7 @@ async def submit_image_endpoint(request): 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: expected dynamic mode", 400) + return await handle_api_error(request, "Job not configured for image submissions", 400) if JOB_QUEUE in job_data: await job_data[JOB_QUEUE].put({ 'worker_id': worker_id, @@ -501,7 +502,7 @@ async def submit_image_endpoint(request): 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: expected dynamic mode", 400) + return await handle_api_error(request, "Job not configured for image submissions", 400) if JOB_QUEUE in job_data: await job_data[JOB_QUEUE].put({ 'worker_id': worker_id, @@ -599,7 +600,7 @@ async def request_image_endpoint(request): elif mode == 'static' and JOB_PENDING_TASKS in job_data: pending_queue = job_data[JOB_PENDING_TASKS] else: - return await handle_api_error(request, f"Invalid {mode} mode configuration", 400) + return await handle_api_error(request, "Invalid job configuration", 400) try: task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)