From 1cb193de32e89e4abd42b560daf270ef36ce0c13 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 21:32:33 +1000 Subject: [PATCH] optimisations --- distributed_upscale.py | 369 +++++++++++++++++++++++++++-------------- 1 file changed, 244 insertions(+), 125 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index d22fb5a..2908d68 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -1,15 +1,4 @@ -""" -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 statements (original + new) import torch import numpy as np from PIL import Image, ImageFilter, ImageDraw @@ -38,6 +27,12 @@ from .utils.constants import ( QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT ) +# Import for controller support +from .utils.usdu_utils import crop_cond + +# Additional import for deepcopy +import copy + # Make MAX_BATCH configurable MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) @@ -134,6 +129,7 @@ class UltimateSDUpscaleDistributed: """Force re-execution.""" return float("nan") # Always re-execute + def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id="", is_worker=False, @@ -173,10 +169,7 @@ class UltimateSDUpscaleDistributed: 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" + mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) debug_log(f"UltimateSDUpscale Worker - Mode: {mode}, batch_size: {batch_size}") @@ -201,7 +194,7 @@ class UltimateSDUpscaleDistributed: 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) + assigned_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id) if not assigned_global_indices: return (upscaled_image,) @@ -209,6 +202,9 @@ class UltimateSDUpscaleDistributed: # Process tiles SYNCHRONOUSLY processed_tiles = [] + # Track which batch indices we've already pre-sliced conditioning for + sliced_conditioning_cache = {} + for global_idx in assigned_global_indices: # Calculate which image and tile this corresponds to batch_idx = global_idx // num_tiles_per_image @@ -218,6 +214,22 @@ class UltimateSDUpscaleDistributed: if batch_idx >= upscaled_image.shape[0]: debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}") continue + + # Pre-slice conditioning once per image (cache it) + if batch_idx not in sliced_conditioning_cache: + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[batch_idx:batch_idx+1] + # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] x, y = all_tiles[tile_idx] @@ -228,9 +240,10 @@ class UltimateSDUpscaleDistributed: # Process tile through SD with unique seed image_seed = seed + batch_idx * 1000 - processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, + processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae, image_seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise, tiled_decode) + scheduler, denoise, tiled_decode, batch_idx=batch_idx, + region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height)) processed_tiles.append({ 'tile': processed_tile, @@ -278,12 +291,7 @@ class UltimateSDUpscaleDistributed: num_workers = len(enabled_workers) # Determine processing mode - if num_workers == 0: - mode = "single_gpu" - elif batch_size == 1 or batch_size <= dynamic_threshold: - mode = "static" - else: - mode = "dynamic" + mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) debug_log(f"UltimateSDUpscale Master - Mode: {mode}, batch_size: {batch_size}, workers: {num_workers}") @@ -315,10 +323,6 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}") raise RuntimeError(f"Failed to initialize static mode queue: {e}") - # Don't try to prepare via API since we're using local queues - # loop = asyncio.new_event_loop() - # asyncio.set_event_loop(loop) - # Convert batch to PIL list for processing result_images = [] for b in range(batch_size): @@ -338,6 +342,9 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Processing {master_tiles} tiles locally from batch of {batch_size}") # Process master tiles SYNCHRONOUSLY + # Track which batch indices we've already pre-sliced conditioning for + sliced_conditioning_cache = {} + for global_idx in master_global_indices: # Calculate which image and tile this corresponds to batch_idx = global_idx // num_tiles_per_image @@ -347,15 +354,31 @@ class UltimateSDUpscaleDistributed: if batch_idx >= len(result_images): debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds result_images length {len(result_images)}") continue + + # Pre-slice conditioning once per image (cache it) + if batch_idx not in sliced_conditioning_cache: + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[batch_idx:batch_idx+1] + # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] # 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, + 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 + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx ) # Collect worker tiles using async @@ -431,9 +454,8 @@ class UltimateSDUpscaleDistributed: 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): + def _get_worker_global_indices(self, total_tiles, enabled_workers, 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}") @@ -444,71 +466,20 @@ class UltimateSDUpscaleDistributed: log(f"UltimateSDUpscale Worker - Worker {worker_id} not found in enabled list {enabled_workers}") return [] - # Calculate tile distribution - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - - start_idx = master_tiles + (worker_index * tiles_per_participant) - if worker_index < remainder - 1: - start_idx += worker_index - end_idx = start_idx + tiles_per_participant + 1 - else: - start_idx += remainder - 1 if remainder > 0 else 0 - end_idx = start_idx + tiles_per_participant - - end_idx = min(end_idx, total_tiles) - return list(range(start_idx, end_idx)) - - def _get_worker_tiles(self, all_tiles, enabled_worker_ids, worker_id): - """Calculate which tiles are assigned to a specific worker.""" - enabled_workers = json.loads(enabled_worker_ids) - num_workers = len(enabled_workers) - total_tiles = len(all_tiles) - - debug_log(f"UltimateSDUpscale Worker - Worker ID: {worker_id}, Enabled workers: {enabled_workers}") - - try: - worker_index = enabled_workers.index(worker_id) - except ValueError: - log(f"UltimateSDUpscale Worker - Worker {worker_id} not found in enabled list {enabled_workers}") - return [] - - # Calculate tile distribution - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - - start_idx = master_tiles + (worker_index * tiles_per_participant) - if worker_index < remainder - 1: - start_idx += worker_index - end_idx = start_idx + tiles_per_participant + 1 - else: - start_idx += remainder - 1 if remainder > 0 else 0 - end_idx = start_idx + tiles_per_participant - - end_idx = min(end_idx, total_tiles) - return list(range(start_idx, end_idx)) + global_indices = list(range(total_tiles)) + all_assignments = self._distribute_items(global_indices, num_workers + 1) + return all_assignments[worker_index + 1] # +1 because 0 is master def _get_master_global_indices(self, total_tiles, num_workers): """Calculate which global tile indices are assigned to the master in flattened mode.""" - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - return list(range(master_tiles)) - - def _get_master_tiles(self, all_tiles, num_workers): - """Calculate which tiles are assigned to the master.""" - total_tiles = len(all_tiles) - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - return list(range(master_tiles)) + global_indices = list(range(total_tiles)) + all_assignments = self._distribute_items(global_indices, num_workers + 1) + return all_assignments[0] def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, image_width, image_height, tiled_decode): + padding, mask_blur, image_width, image_height, tiled_decode, batch_idx: int = 0): """Process a single tile and blend it into the result image.""" x, y = tile_pos @@ -519,7 +490,8 @@ class UltimateSDUpscaleDistributed: processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise, tiled_decode) + scheduler, denoise, tiled_decode, batch_idx=batch_idx, + region=(x1, y1, x1 + ew, y1 + eh), image_size=(image_width, image_height)) # Convert and blend processed_pil = tensor_to_pil(processed_tile, 0) @@ -531,17 +503,6 @@ class UltimateSDUpscaleDistributed: return result_image - async def _prepare_multigpu_job(self, multi_job_id): - """Prepare job via API endpoint to ensure server is ready.""" - session = await get_client_session() - # Use the actual server port (get from server instance) - port = get_server_port() - async with session.post(f"http://localhost:{port}/distributed/prepare_job", - json={"multi_job_id": multi_job_id}) as response: - if response.status != 200: - text = await response.text() - raise RuntimeError(f"Failed to prepare job: {response.status} - {text}") - async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, all_indices=None): """Initialize the job queue for dynamic mode with pending images.""" prompt_server = ensure_tile_jobs_initialized() @@ -864,9 +825,31 @@ class UltimateSDUpscaleDistributed: return tile_tensor, x1, y1, extracted_width, extracted_height def process_tile(self, tile_tensor: torch.Tensor, model, positive, negative, vae, - seed: int, steps: int, cfg: float, sampler_name: str, - scheduler: str, denoise: float, tiled_decode: bool = False) -> torch.Tensor: - """Process a single tile through SD sampling.""" + seed: int, steps: int, cfg: float, sampler_name: str, + scheduler: str, denoise: float, tiled_decode: bool = False, + batch_idx: int = 0, region: Tuple[int, int, int, int] = None, + image_size: Tuple[int, int] = None) -> torch.Tensor: + """Process a single tile through SD sampling. + Note: positive and negative should already be pre-sliced for the current batch_idx.""" + debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}") + + # Log conditioning details + debug_log(f"[process_tile] Positive conditioning items: {len(positive)}") + debug_log(f"[process_tile] Negative conditioning items: {len(negative)}") + + # Check for ControlNet in conditioning + for i, (emb, cond_dict) in enumerate(positive): + if 'control' in cond_dict: + control = cond_dict['control'] + debug_log(f"[process_tile] Found ControlNet in positive[{i}]") + debug_log(f"[process_tile] Control id: {id(control)}") + if hasattr(control, 'control_model'): + debug_log(f"[process_tile] control_model id: {id(control.control_model)}") + if hasattr(control, 'cond_hint_original'): + debug_log(f"[process_tile] cond_hint_original shape: {control.cond_hint_original.shape}") + debug_log(f"[process_tile] Expected to use hint at index {batch_idx} from batch") + debug_log(f"[process_tile] Embedding shape: {emb.shape}") + # Import here to avoid circular dependencies from nodes import common_ksampler, VAEEncode, VAEDecode @@ -887,15 +870,38 @@ class UltimateSDUpscaleDistributed: if tile_tensor.is_cuda: clean_tensor = clean_tensor.cuda() + # Crop conditioning to tile region if provided (assumes hints at image resolution) + if region is not None and image_size is not None: + positive_cropped = copy.deepcopy(positive) + negative_cropped = copy.deepcopy(negative) + init_size = image_size # (width, height) of full image + canvas_size = image_size + tile_size = (tile_tensor.shape[2], tile_tensor.shape[1]) # (width, height) + w_pad = 0 # No extra pad needed; region already includes padding + h_pad = 0 + positive_cropped = crop_cond(positive_cropped, region, init_size, canvas_size, tile_size, w_pad, h_pad) + negative_cropped = crop_cond(negative_cropped, region, init_size, canvas_size, tile_size, w_pad, h_pad) + else: + # No region cropping needed, use conditioning as-is + positive_cropped = positive + negative_cropped = negative + # Encode to latent if tiled_decode and tiled_vae_available: latent = VAEEncodeTiled().encode(vae, clean_tensor)[0] else: latent = VAEEncode().encode(vae, clean_tensor)[0] + # Log latent shape + if 'samples' in latent: + debug_log(f"[process_tile] Latent shape: {latent['samples'].shape}") + # Sample + debug_log(f"[process_tile] About to call common_ksampler with batch_idx={batch_idx}") + debug_log(f"[process_tile] Model id: {id(model)}") samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, - positive, negative, latent, denoise=denoise)[0] + positive_cropped, negative_cropped, latent, denoise=denoise)[0] + debug_log(f"[process_tile] common_ksampler completed") # Decode back to image if tiled_decode and tiled_vae_available: @@ -928,7 +934,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: int, padding: int) -> Image.Image: + mask: Image.Image, padding: int) -> Image.Image: """Blend a processed tile back into the base image using Ultimate SD Upscale's exact approach. This follows the exact method from ComfyUI_UltimateSDUpscale/modules/processing.py @@ -1083,15 +1089,28 @@ class UltimateSDUpscaleDistributed: # Process each image in the batch for batch_idx in range(batch_size): - # Process each tile for this image + debug_log(f"[process_single_gpu] Processing batch_idx {batch_idx}") + # Pre-slice conditioning once per image (not per tile) + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[batch_idx:batch_idx+1] + # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + + # Process each tile for this image with pre-sliced conditioning 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, + 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 + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx ) # Convert back to tensor @@ -1160,25 +1179,30 @@ class UltimateSDUpscaleDistributed: consecutive_retries = 0 debug_log(f"Master processing image {image_idx} dynamically") processed_count += 1 - - # Get total completed count for accurate progress - total_completed = run_async_in_server_loop( - self._get_total_completed_count(multi_job_id), - timeout=1.0 - ) - log(f"{total_completed}/{batch_size} completed images") # Process locally single_tensor = upscaled_image[image_idx:image_idx+1] local_image = result_images[image_idx] image_seed = seed + image_idx * 1000 + # Pre-slice conditioning once per image (not per tile) + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(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] + # Don't slice ControlNet - crop_cond will handle spatial cropping + 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, negative, vae, image_seed, steps, cfg, + 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 + padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx ) result_images[image_idx] = local_image @@ -1194,6 +1218,16 @@ class UltimateSDUpscaleDistributed: self._drain_worker_results_queue(multi_job_id), timeout=5.0 ) + + # Check for timed out workers and requeue their images + requeued_count = run_async_in_server_loop( + self._check_and_requeue_timed_out_workers(multi_job_id, batch_size), + timeout=5.0 + ) + if requeued_count > 0: + log(f"Requeued {requeued_count} images from timed out workers") + consecutive_retries = 0 # Reset since we have work to do + continue # Now check total completed (includes newly collected) completed_now = run_async_in_server_loop( @@ -1201,9 +1235,6 @@ class UltimateSDUpscaleDistributed: timeout=1.0 ) - # Show accurate progress if we drained any new results - if drained_count > 0: - log(f"{completed_now}/{batch_size} completed images") debug_log(f"Queue empty, completed count: {completed_now}/{batch_size}") @@ -1211,6 +1242,17 @@ class UltimateSDUpscaleDistributed: debug_log("All images completed, exiting master loop") break + # Check if there are pending images in the queue (could be requeued) + pending_count = run_async_in_server_loop( + self._get_pending_count(multi_job_id), + timeout=1.0 + ) + + if pending_count > 0: + debug_log(f"Found {pending_count} pending images (possibly requeued), continuing to process") + consecutive_retries = 0 # Reset retries since there's work to do + continue + consecutive_retries += 1 if consecutive_retries >= max_consecutive_retries: log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.") @@ -1298,6 +1340,15 @@ class UltimateSDUpscaleDistributed: return job_data['completed_images'].copy() return {} + async def _get_pending_count(self, multi_job_id): + """Get count of pending images in the queue.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'pending_images' in job_data: + return job_data['pending_images'].qsize() + return 0 + async def _drain_worker_results_queue(self, multi_job_id): """Drain pending worker results from queue and update completed_images. Returns count of drained images.""" prompt_server = ensure_tile_jobs_initialized() @@ -1334,6 +1385,36 @@ class UltimateSDUpscaleDistributed: return collected + async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size): + """Check for timed out workers and requeue their assigned images. Returns count of requeued images.""" + 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: + return 0 + + current_time = time.time() + requeued_count = 0 + completed_images = job_data.get('completed_images', {}) + + # Check for timed out workers (60 second timeout) + for worker, last_heartbeat in list(job_data.get('worker_status', {}).items()): + if current_time - last_heartbeat > 60: + log(f"Worker {worker} timed out") + # Requeue assigned images from this worker + for idx in job_data.get('assigned_to_workers', {}).get(worker, []): + if idx not in completed_images: + await job_data['pending_images'].put(idx) + requeued_count += 1 + debug_log(f"Requeued image {idx} from timed out worker {worker}") + # Clean up worker tracking + if 'worker_status' in job_data: + del job_data['worker_status'][worker] + if 'assigned_to_workers' in job_data: + job_data['assigned_to_workers'][worker] = [] + + return requeued_count + def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, @@ -1395,12 +1476,25 @@ class UltimateSDUpscaleDistributed: # Process all tiles for this image image_seed = seed + image_idx * 1000 + + # Pre-slice conditioning once per image (not per tile) + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(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] + # Don't slice ControlNet - crop_cond will handle spatial cropping + 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, negative, vae, image_seed, steps, cfg, + 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 + padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx ) # Send processed image back to master @@ -1566,6 +1660,31 @@ class UltimateSDUpscaleDistributed: response.raise_for_status() debug_log(f"Worker {worker_id} sent completion signal") + def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str: + """Determines the processing mode based on batch size and worker count.""" + if num_workers == 0: + return "single_gpu" + if batch_size <= dynamic_threshold: + return "static" + return "dynamic" + + def _distribute_items(self, items: list, num_participants: int) -> List[List[any]]: + """Distributes a list of items among N participants.""" + if num_participants == 0: + return [items] # All items for the single participant (master) + + 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 # Ensure initialization before registering routes ensure_tile_jobs_initialized()