from __future__ import annotations import asyncio, time, torch from typing import Any from PIL import Image from ...utils.logging import debug_log, log from ...utils.image import tensor_to_pil, pil_to_tensor from ...utils.async_helpers import run_async_in_server_loop from ...utils.config import get_worker_timeout_seconds from ...utils.constants import ( JOB_POLL_INTERVAL, JOB_POLL_MAX_ATTEMPTS, TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT, ) from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job from ..mode_contexts import ( DynamicModeContext, JobStateCollaborator, ResultCollectorCollaborator, TileOpsCollaborator, WorkerCommsCollaborator, ) from ..processing_args import UpscaleCoreArgs class DynamicModeMixin: """ Dynamic (per-image queue) USDU mode behaviors for master and worker roles. Expected co-mixins on `self`: - TileOpsMixin (`calculate_tiles`, `slice_conditioning`, `process_and_blend_tile`). - JobStateMixin (image queue/task completion helpers). - WorkerCommsMixin (`request_assignment`, `send_full_image`, `send_heartbeat`). """ def _build_dynamic_mode_context(self) -> DynamicModeContext: """Build explicit collaborators for dynamic mode execution.""" return DynamicModeContext( tile_ops=TileOpsCollaborator(self), job_state=JobStateCollaborator(self), worker_comms=WorkerCommsCollaborator(self), result_collector=ResultCollectorCollaborator(self), ) def _poll_job_ready_dynamic( self, multi_job_id: str, master_url: str, worker_id: str | None = None, max_attempts: int = JOB_POLL_MAX_ATTEMPTS, mode_context: DynamicModeContext | None = None, ) -> bool: """Poll master for job readiness using explicit worker-comms collaborator.""" context = mode_context or self._build_dynamic_mode_context() for attempt in range(max_attempts): ready = run_async_in_server_loop( context.worker_comms.check_job_status(multi_job_id, master_url), timeout=5.0, ) if ready: if worker_id: debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts") else: debug_log(f"Job {multi_job_id} ready after {attempt} attempts") return True time.sleep(JOB_POLL_INTERVAL) return False def _process_image_tiles( self, image_idx: int, local_image: Image.Image, single_tensor: torch.Tensor, core_args: UpscaleCoreArgs, all_tiles: list, tile_width: int, tile_height: int, padding: int, mask_blur: int, width: int, height: int, force_uniform_tiles: bool, context: DynamicModeContext, per_tile_callback: Any | None = None, ) -> Image.Image: """Process all tiles for a single image. Returns the processed PIL image.""" positive_sliced, negative_sliced = context.tile_ops.slice_conditioning( core_args.positive, core_args.negative, image_idx, ) for tile_idx, pos in enumerate(all_tiles): source_tensor = pil_to_tensor(local_image) if single_tensor.is_cuda: source_tensor = source_tensor.cuda() local_image = context.tile_ops.process_and_blend_tile( tile_idx, pos, source_tensor, local_image, core_args.model, positive_sliced, negative_sliced, core_args.vae, core_args.seed, core_args.steps, core_args.cfg, core_args.sampler_name, core_args.scheduler, core_args.denoise, tile_width, tile_height, padding, mask_blur, width, height, force_uniform_tiles, core_args.tiled_decode, batch_idx=image_idx, ) if per_tile_callback: per_tile_callback() return local_image def _handle_master_dynamic_idle_state( self, multi_job_id, batch_size, consecutive_retries, max_consecutive_retries, mode_context: DynamicModeContext | None = None, ): """Handle queue-empty state while master waits for worker progress.""" context = mode_context or self._build_dynamic_mode_context() drained_count = run_async_in_server_loop( context.job_state.drain_worker_results_queue(multi_job_id), timeout=5.0, ) run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1) requeued_count = run_async_in_server_loop( context.job_state.check_and_requeue_timed_out_workers(multi_job_id, batch_size), timeout=5.0, ) run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1) if requeued_count > 0: log(f"Requeued {requeued_count} images from timed out workers") return False, True, 0 completed_now = run_async_in_server_loop( context.job_state.total_completed_count(multi_job_id), timeout=1.0, ) log(f"USDU Dist: Images progress {completed_now}/{batch_size}") if completed_now >= batch_size: return True, False, consecutive_retries run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1) pending_count = run_async_in_server_loop( context.job_state.pending_count(multi_job_id), timeout=1.0, ) if pending_count > 0: return False, True, 0 consecutive_retries += 1 if consecutive_retries >= max_consecutive_retries: log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.") return True, False, consecutive_retries if drained_count > 0: debug_log(f"Master idle drain picked up {drained_count} images while waiting for workers") debug_log("Waiting for workers") run_async_in_server_loop(asyncio.sleep(2), timeout=3.0) return False, True, consecutive_retries def _finalize_master_dynamic_results( self, multi_job_id, batch_size, num_workers, processed_count, result_images, upscaled_image, mode_context: DynamicModeContext | None = None, ): """Collect remaining worker outputs and convert final images back to tensor.""" context = mode_context or self._build_dynamic_mode_context() all_completed = run_async_in_server_loop( context.job_state.all_completed_images(multi_job_id), timeout=5.0, ) remaining_to_collect = batch_size - len(all_completed) if remaining_to_collect > 0: debug_log(f"Waiting for {remaining_to_collect} more images from workers") collection_timeout = float(get_worker_timeout_seconds()) collected_images = run_async_in_server_loop( context.result_collector.collect_dynamic_images( multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count, ), timeout=collection_timeout, ) all_completed.update(collected_images) for idx, processed_img in all_completed.items(): if idx < batch_size: result_images[idx] = processed_img result_tensor = ( torch.cat([pil_to_tensor(img) for img in result_images], dim=0) if batch_size > 1 else pil_to_tensor(result_images[0]) ) if upscaled_image.is_cuda: result_tensor = result_tensor.cuda() return result_tensor def process_master_dynamic( self, upscaled_image: torch.Tensor, core_args: UpscaleCoreArgs, tile_width: int, tile_height: int, padding: int, mask_blur: int, force_uniform_tiles: bool, multi_job_id: str, enabled_workers: list[str], mode_context: DynamicModeContext | None = None, ) -> tuple[torch.Tensor]: """Dynamic mode for large batches - assigns whole images to workers dynamically, including master.""" context = mode_context or self._build_dynamic_mode_context() # Get batch size and dimensions batch_size, height, width, _ = upscaled_image.shape num_workers = len(enabled_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)) debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.") # Calculate tiles for processing all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) # Initialize job queue for communication try: run_async_in_server_loop( init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices), timeout=2.0 ) except Exception as 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 result_images = [tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB').copy() for b in range(batch_size)] # Process images dynamically with master participating prompt_server = ensure_tile_jobs_initialized() processed_count = 0 consecutive_retries = 0 max_consecutive_retries = 10 # Process loop - master pulls from queue and processes synchronously while processed_count < batch_size: # Try to get an image to process image_idx = run_async_in_server_loop( context.job_state.next_image_index(multi_job_id), timeout=5.0 # Short timeout to allow frequent checks ) if image_idx is not None: # Reset retry counter and process locally consecutive_retries = 0 debug_log(f"Master processing image {image_idx} dynamically") processed_count += 1 # Process locally single_tensor = upscaled_image[image_idx:image_idx+1] local_image = result_images[image_idx] local_image = self._process_image_tiles( image_idx, local_image, single_tensor, core_args, all_tiles, tile_width, tile_height, padding, mask_blur, width, height, force_uniform_tiles, context, per_tile_callback=lambda: run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1), ) result_images[image_idx] = local_image # Mark as completed run_async_in_server_loop( context.result_collector.mark_image_completed(multi_job_id, image_idx, local_image), timeout=5.0 ) # NEW: Drain after the full image is marked complete (catches workers who finished during master's processing) drained_count = run_async_in_server_loop( context.job_state.drain_worker_results_queue(multi_job_id), timeout=5.0 ) if drained_count > 0: debug_log(f"Drained {drained_count} worker images after master's image {image_idx}") # NEW: Log overall progress (includes master's image + any drained workers) completed_now = run_async_in_server_loop( context.job_state.total_completed_count(multi_job_id), timeout=1.0 ) 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(context.worker_comms.async_yield(), timeout=0.1) else: should_break, should_continue, consecutive_retries = self._handle_master_dynamic_idle_state( multi_job_id=multi_job_id, batch_size=batch_size, consecutive_retries=consecutive_retries, max_consecutive_retries=max_consecutive_retries, mode_context=context, ) if should_break: break if should_continue: continue debug_log(f"Master processed {processed_count} images locally") result_tensor = self._finalize_master_dynamic_results( multi_job_id=multi_job_id, batch_size=batch_size, num_workers=num_workers, processed_count=processed_count, result_images=result_images, upscaled_image=upscaled_image, mode_context=context, ) debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") log(f"Completed processing all {batch_size} images") return (result_tensor,) def process_worker_dynamic( self, upscaled_image: torch.Tensor, core_args: UpscaleCoreArgs, tile_width: int, tile_height: int, padding: int, mask_blur: int, force_uniform_tiles: bool, multi_job_id: str, master_url: str, worker_id: str, enabled_worker_ids: list[str] | str, dynamic_threshold: int, mode_context: DynamicModeContext | None = None, ) -> tuple[torch.Tensor]: """Worker processing in dynamic mode - processes whole images.""" context = mode_context or self._build_dynamic_mode_context() # Round tile dimensions tile_width = context.tile_ops.round_to_multiple(tile_width) tile_height = context.tile_ops.round_to_multiple(tile_height) # Get dimensions and tile grid batch_size, height, width, _ = upscaled_image.shape all_tiles = context.tile_ops.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 # Ensure the dynamic job queue exists on the master. In delegate-only # mode the master's USDU node is replaced with a lightweight collector # that cannot initialise the queue itself, so workers do it instead. # The call is idempotent — if the master already created the queue this # is a no-op. from ...utils.worker_ids import coerce_enabled_worker_ids as _coerce _enabled = _coerce(enabled_worker_ids) run_async_in_server_loop( context.worker_comms.init_dynamic_job_on_master( multi_job_id, master_url, batch_size, _enabled, ), timeout=10.0, ) # Poll for job readiness to avoid races during master init max_poll_attempts = JOB_POLL_MAX_ATTEMPTS if not self._poll_job_ready_dynamic( multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts, mode_context=context, ): 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 assignment = run_async_in_server_loop( context.worker_comms.request_assignment(multi_job_id, master_url, worker_id), timeout=TILE_WAIT_TIMEOUT, ) image_idx = assignment.task_idx if assignment.kind == "image" else None estimated_remaining = assignment.estimated_remaining 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 local_image = self._process_image_tiles( image_idx, local_image, single_tensor, core_args, all_tiles, tile_width, tile_height, padding, mask_blur, width, height, force_uniform_tiles, context, per_tile_callback=lambda: run_async_in_server_loop( context.worker_comms.send_heartbeat(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( context.worker_comms.send_full_image( 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( context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id), timeout=5.0 ) if is_last: break except Exception as e: 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( context.worker_comms.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,)