diff --git a/distributed_upscale.py b/distributed_upscale.py index 626d15d..c9e20cc 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -1,55 +1,21 @@ -# Import statements (original + new) -import torch -from PIL import Image, ImageFilter, ImageDraw import json -import asyncio -import aiohttp -import io import math -import time -from typing import List, Tuple from functools import wraps -# Import ComfyUI modules import comfy.samplers -import comfy.model_management -# Import shared utilities from .utils.logging import debug_log, log -from .utils.image import tensor_to_pil, pil_to_tensor -from .utils.network import get_client_session from .utils.async_helpers import run_async_in_server_loop -from .utils.config import get_worker_timeout_seconds -from .utils.constants import ( - TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT, - TILE_SEND_TIMEOUT, MAX_BATCH -) +from .utils.usdu_managment import ensure_tile_jobs_initialized -# Import for controller support -from .utils.usdu_utils import ( - crop_cond, - get_crop_region, - expand_crop, -) -from .utils.usdu_managment import ( - clone_conditioning, ensure_tile_jobs_initialized, - # Job management functions - _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 +from .upscale.tile_ops import TileOpsMixin +from .upscale.result_collector import ResultCollectorMixin +from .upscale.worker_comms import WorkerCommsMixin +from .upscale.job_state import JobStateMixin +from .upscale.modes.single_gpu import SingleGpuModeMixin +from .upscale.modes.static import StaticModeMixin +from .upscale.modes.dynamic import DynamicModeMixin - -# Note: MAX_BATCH and HEARTBEAT_TIMEOUT are imported from utils.constants -# They can be overridden via environment variables: -# - COMFYUI_MAX_BATCH (default: 20) -# - COMFYUI_HEARTBEAT_TIMEOUT (default: 90) - -# Sync wrapper decorator for async methods def sync_wrapper(async_func): """Decorator to wrap async methods for synchronous execution.""" @wraps(async_func) @@ -61,9 +27,6 @@ def sync_wrapper(async_func): ) return sync_func -# Note: tensor_to_pil and pil_to_tensor are imported from utils.image - - def _parse_enabled_worker_ids(enabled_worker_ids): """Parse enabled worker IDs from either JSON or list input.""" if isinstance(enabled_worker_ids, list): @@ -80,8 +43,16 @@ def _parse_enabled_worker_ids(enabled_worker_ids): return [str(wid) for wid in parsed] return [] +class UltimateSDUpscaleDistributed( + DynamicModeMixin, + StaticModeMixin, + SingleGpuModeMixin, + ResultCollectorMixin, + WorkerCommsMixin, + JobStateMixin, + TileOpsMixin, +): -class UltimateSDUpscaleDistributed: """ Distributed version of Ultimate SD Upscale (No Upscale). @@ -109,8 +80,6 @@ class UltimateSDUpscaleDistributed: ensure_tile_jobs_initialized() debug_log("UltimateSDUpscaleDistributed - Node initialized") - # WAN/FLOW detection removed per user request; enforcing 4n+1 for any batch > 1. - @classmethod def INPUT_TYPES(s): return { @@ -143,17 +112,16 @@ class UltimateSDUpscaleDistributed: "dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}), }, } - + RETURN_TYPES = ("IMAGE",) FUNCTION = "run" CATEGORY = "image/upscaling" - + @classmethod def IS_CHANGED(cls, **kwargs): """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, @@ -192,7 +160,7 @@ class UltimateSDUpscaleDistributed: tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, dynamic_threshold) - + def process_worker(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, @@ -232,7 +200,7 @@ class UltimateSDUpscaleDistributed: tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, master_url, worker_id, enabled_workers) - + def process_master(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, @@ -287,1666 +255,6 @@ class UltimateSDUpscaleDistributed: tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, all_tiles, num_tiles_per_image) - - # Legacy static assignment helpers removed - - 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, force_uniform_tiles, - tiled_decode, batch_idx: int = 0): - """Process a single tile and blend it into the result image.""" - x, y = tile_pos - - # Extract and process tile - tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( - upscaled_image, x, y, tile_width, tile_height, padding, force_uniform_tiles - ) - - processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, - seed, steps, cfg, sampler_name, - 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) - # 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, - x1, y1, (ew, eh), tile_mask, padding) - - return result_image - - - async def _async_collect_results(self, multi_job_id, num_workers, mode='static', - remaining_to_collect=None, batch_size=None): - """Unified async helper to collect results from workers (tiles or images).""" - # Get the already initialized queue - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - raise RuntimeError(f"Job queue not initialized for {multi_job_id}") - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, dict) or 'mode' not in job_data: - raise RuntimeError("Invalid job data structure") - if job_data['mode'] != mode: - raise RuntimeError(f"Mode mismatch: expected {mode}, got {job_data['mode']}") - q = job_data['queue'] - - # For dynamic mode, get reference to completed_images - if mode == 'dynamic': - completed_images = job_data['completed_images'] - # Calculate expected count for logging - expected_count = remaining_to_collect or batch_size - else: - # Calculate expected count from job data for static mode - expected_count = job_data.get('total_items', 0) - len(job_data.get('completed_items', set())) - - item_type = "images" if mode == 'dynamic' else "tiles" - debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers") - - collected_results = {} - workers_done = set() - # Unify collector/upscaler wait behavior with the UI worker timeout - timeout = float(get_worker_timeout_seconds()) - last_heartbeat_check = time.time() - collected_count = 0 - - while len(workers_done) < num_workers: - # Check for user interruption - if comfy.model_management.processing_interrupted(): - log("Processing interrupted by user") - raise comfy.model_management.InterruptProcessingException() - - # For dynamic mode with remaining_to_collect, check if we've collected enough - if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect: - break - - try: - # Shorter poll for dynamic mode, but never exceed the configured timeout - wait_timeout = (min(10.0, timeout) if mode == 'dynamic' else timeout) - result = await asyncio.wait_for(q.get(), timeout=wait_timeout) - worker_id = result['worker_id'] - is_last = result.get('is_last', False) - - if mode == 'static': - # Handle tiles - tiles = result.get('tiles', []) - if tiles: - # Batch mode - debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})") - - for tile_data in tiles: - # Validate required fields - if 'batch_idx' not in tile_data: - log(f"UltimateSDUpscale Master - Missing batch_idx in tile data, skipping") - continue - - tile_idx = tile_data['tile_idx'] - # Use global_idx as key if available (for batch processing) - key = tile_data.get('global_idx', tile_idx) - - # Store the full tile data including metadata; prefer PIL image if present - entry = { - 'tile_idx': tile_idx, - 'x': tile_data['x'], - 'y': tile_data['y'], - 'extracted_width': tile_data['extracted_width'], - 'extracted_height': tile_data['extracted_height'], - 'padding': tile_data['padding'], - 'worker_id': worker_id, - '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'] - collected_results[tile_idx] = result - debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})") - - elif mode == 'dynamic': - # Handle full images - if 'image_idx' in result and 'image' in result: - image_idx = result['image_idx'] - image_pil = result['image'] - completed_images[image_idx] = image_pil - collected_results[image_idx] = image_pil - collected_count += 1 - 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 - Worker {worker_id} completed") - - except asyncio.TimeoutError: - if mode == 'dynamic': - # Check for worker timeouts periodically - current_time = time.time() - if current_time - last_heartbeat_check >= 10.0: - # 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 - 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 - Overall timeout waiting for images") - break - else: - log(f"UltimateSDUpscale Master - Timeout waiting for {item_type}") - break - - 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: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - del prompt_server.distributed_pending_tile_jobs[multi_job_id] - - return collected_results if mode == 'static' else completed_images - - # Keep compatibility wrappers for existing code - async def _async_collect_worker_tiles(self, multi_job_id, num_workers): - """Async helper to collect tiles from workers.""" - return await self._async_collect_results(multi_job_id, num_workers, mode='static') - - async def _mark_image_completed(self, multi_job_id, image_idx, image_pil): - """Mark an image as completed in the job data.""" - # Mark the image as completed with the image data - await _mark_task_completed(multi_job_id, image_idx, {'image': image_pil}) - 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 'completed_images' in job_data: - job_data['completed_images'][image_idx] = image_pil - - async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count): - """Collect remaining processed images from workers.""" - return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic', - remaining_to_collect=remaining_to_collect, - batch_size=batch_size) - - def round_to_multiple(self, value: int, multiple: int = 8) -> int: - """Round value to nearest multiple.""" - return round(value / multiple) * multiple - - def calculate_tiles(self, image_width: int, image_height: int, - tile_width: int, tile_height: int, force_uniform_tiles: bool = True) -> List[Tuple[int, int]]: - """Calculate tile positions to match Ultimate SD Upscale. - - Positions are a simple grid starting at (0,0) with steps of - `tile_width` and `tile_height`, using ceil(rows/cols) to cover edges. - Uniform vs non-uniform affects only crop/resize, not positions. - """ - rows = math.ceil(image_height / tile_height) - cols = math.ceil(image_width / tile_width) - tiles: List[Tuple[int, int]] = [] - for yi in range(rows): - for xi in range(cols): - tiles.append((xi * tile_width, yi * tile_height)) - return tiles - - def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int, - tile_width: int, tile_height: int, padding: int, - force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]: - """Extract a tile region and resize to match USDU cropping logic. - - Mirrors ComfyUI_UltimateSDUpscale processing: - - Build a mask with a white rectangle at the tile rect - - Compute crop_region via get_crop_region(mask, padding) - - If force_uniform_tiles: shrink/expand to target size of - round_to_multiple(tile + padding) for each dimension - - Else: target is ceil(crop_size/8)*8 per dimension - - Extract the crop and resize to target tile_size - Returns the resized tensor and crop origin/size for blending. - """ - _, h, w, _ = image.shape - - # Create mask and compute initial padded crop region - mask = Image.new('L', (w, h), 0) - draw = ImageDraw.Draw(mask) - draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255) - x1, y1, x2, y2 = get_crop_region(mask, padding) - - # Determine target tile size (processing size) - if force_uniform_tiles: - target_w = self.round_to_multiple(tile_width + padding, 8) - target_h = self.round_to_multiple(tile_height + padding, 8) - (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) - else: - crop_w = x2 - x1 - crop_h = y2 - y1 - target_w = max(8, math.ceil(crop_w / 8) * 8) - target_h = max(8, math.ceil(crop_h / 8) * 8) - (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) - - # Actual extracted size before resizing (for blending) - extracted_width = x2 - x1 - extracted_height = y2 - y1 - - # Extract tile and resize to processing size - tile = image[:, y1:y2, x1:x2, :] - tile_pil = tensor_to_pil(tile, 0) - if tile_pil.size != (target_w, target_h): - tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS) - - tile_tensor = pil_to_tensor(tile_pil) - if image.is_cuda: - tile_tensor = tile_tensor.cuda() - - return tile_tensor, x1, y1, extracted_width, extracted_height - - def extract_batch_tile_with_padding(self, images: torch.Tensor, x: int, y: int, - tile_width: int, tile_height: int, padding: int, - force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]: - """Extract a tile region for the entire batch and resize to USDU logic. - - - Computes a single crop region from a mask at (x,y,w,h) with padding - - force_uniform_tiles controls target processing size logic - - Returns a batched tensor [B,H',W',C] and crop origin/size for blending - """ - batch, h, w, _ = images.shape - - # Create mask and compute initial padded crop region (same for all images) - mask = Image.new('L', (w, h), 0) - draw = ImageDraw.Draw(mask) - draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255) - x1, y1, x2, y2 = get_crop_region(mask, padding) - - # Determine target processing size - if force_uniform_tiles: - target_w = self.round_to_multiple(tile_width + padding, 8) - target_h = self.round_to_multiple(tile_height + padding, 8) - (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) - else: - crop_w = x2 - x1 - crop_h = y2 - y1 - target_w = max(8, math.ceil(crop_w / 8) * 8) - target_h = max(8, math.ceil(crop_h / 8) * 8) - (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) - - extracted_width = x2 - x1 - extracted_height = y2 - y1 - - # Slice batch region - tiles = images[:, y1:y2, x1:x2, :] - - # Resize each tile to target size - resized_tiles = [] - for i in range(batch): - tile_pil = tensor_to_pil(tiles, i) - if tile_pil.size != (target_w, target_h): - tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS) - resized_tiles.append(pil_to_tensor(tile_pil)) - tile_batch = torch.cat(resized_tiles, dim=0) - - if images.is_cuda: - tile_batch = tile_batch.cuda() - - return tile_batch, 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, - 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}") - - - # Import here to avoid circular dependencies - from nodes import common_ksampler, VAEEncode, VAEDecode - - # Try to import tiled VAE nodes if available - try: - from nodes import VAEEncodeTiled, VAEDecodeTiled - tiled_vae_available = True - except ImportError: - tiled_vae_available = False - if tiled_decode: - debug_log("Tiled VAE nodes not available, falling back to standard VAE") - - # Convert to PIL and back to ensure clean tensor without gradient tracking - tile_pil = tensor_to_pil(tile_tensor, 0) - clean_tensor = pil_to_tensor(tile_pil) - - # Ensure tensor is detached and doesn't require gradients - clean_tensor = clean_tensor.detach() - if hasattr(clean_tensor, 'requires_grad_'): - clean_tensor.requires_grad_(False) - - # Move to correct device - if tile_tensor.is_cuda: - clean_tensor = clean_tensor.cuda() - clean_tensor = clean_tensor.detach() # Detach again after device transfer - - # Clone conditioning per tile (shares models, clones hints for cropping) - positive_tile = clone_conditioning(positive, clone_hints=True) - negative_tile = clone_conditioning(negative, clone_hints=True) - - # Crop conditioning to tile region if provided (assumes hints at image resolution) - if region is not None and image_size is not None: - 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_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) - negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) - else: - # No region cropping needed, use cloned conditioning as-is - positive_cropped = positive_tile - negative_cropped = negative_tile - - # Encode to latent (always non-tiled, matching original node) - latent = VAEEncode().encode(vae, clean_tensor)[0] - - # Sample - samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, - positive_cropped, negative_cropped, latent, denoise=denoise)[0] - - # Decode back to image - if tiled_decode and tiled_vae_available: - image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0] - else: - image = VAEDecode().decode(vae, samples)[0] - - return image - - def process_tiles_batch(self, tile_batch: torch.Tensor, model, positive, negative, vae, - seed: int, steps: int, cfg: float, sampler_name: str, - scheduler: str, denoise: float, tiled_decode: bool, - region: Tuple[int, int, int, int], image_size: Tuple[int, int]) -> torch.Tensor: - """Process a batch of tiles together (USDU behavior). - - tile_batch: [B, H, W, C] - Returns image batch tensor [B, H, W, C] - """ - # Import locally to avoid circular deps - from nodes import common_ksampler, VAEEncode, VAEDecode - try: - from nodes import VAEEncodeTiled, VAEDecodeTiled - tiled_vae_available = True - except ImportError: - tiled_vae_available = False - - # Detach and move device - clean = tile_batch.detach() - if hasattr(clean, 'requires_grad_'): - clean.requires_grad_(False) - if tile_batch.is_cuda: - clean = clean.cuda().detach() - - # Clone/crop conditioning once for the region - positive_tile = clone_conditioning(positive, clone_hints=True) - negative_tile = clone_conditioning(negative, clone_hints=True) - - init_size = image_size - canvas_size = image_size - tile_size = (clean.shape[2], clean.shape[1]) # (W,H) - w_pad = 0 - h_pad = 0 - positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) - negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) - - # Encode -> Sample -> Decode - latent = VAEEncode().encode(vae, clean)[0] - samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, - positive_cropped, negative_cropped, latent, denoise=denoise)[0] - if tiled_decode and tiled_vae_available: - image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0] - else: - image = VAEDecode().decode(vae, samples)[0] - - return image - - def create_tile_mask(self, image_width: int, image_height: int, - x: int, y: int, tile_width: int, tile_height: int, - mask_blur: int) -> Image.Image: - """Create a mask for blending tiles - matches Ultimate SD Upscale approach. - - Creates a black image with a white rectangle at the tile position, - then applies blur to create soft edges. - """ - # Create a full-size mask matching the image dimensions - mask = Image.new('L', (image_width, image_height), 0) # Black background - - # Draw white rectangle at tile position - draw = ImageDraw.Draw(mask) - draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255) - - # Apply blur to soften edges - if mask_blur > 0: - mask = mask.filter(ImageFilter.GaussianBlur(mask_blur)) - - return mask - - def blend_tile(self, base_image: Image.Image, tile_image: Image.Image, - x: int, y: int, extracted_size: Tuple[int, int], - mask: Image.Image, padding: int) -> Image.Image: - """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 - """ - extracted_width, extracted_height = extracted_size - - # Debug logging (uncomment if needed) - # debug_log(f"[Blend] Placing tile at ({x}, {y}), size: {extracted_width}x{extracted_height}") - - # Calculate the crop region that was used for extraction - crop_region = (x, y, x + extracted_width, y + extracted_height) - - # The mask is already full-size, no need to crop - - # Resize the processed tile back to the extracted size - if tile_image.size != (extracted_width, extracted_height): - tile_resized = tile_image.resize((extracted_width, extracted_height), Image.LANCZOS) - else: - tile_resized = tile_image - - # Follow Ultimate SD Upscale blending approach: - # Put the tile into position - image_tile_only = Image.new('RGBA', base_image.size) - image_tile_only.paste(tile_resized, crop_region[:2]) - - # Add the mask as an alpha channel - # Must make a copy due to the possibility of an edge becoming black - temp = image_tile_only.copy() - temp.putalpha(mask) # Use the full image mask - image_tile_only.paste(temp, image_tile_only) - - # Add back the tile to the initial image according to the mask in the alpha channel - result = base_image.convert('RGBA') - result.alpha_composite(image_tile_only) - - # Convert back to RGB - return result.convert('RGB') - - def _slice_conditioning(self, positive, negative, batch_idx): - """Helper to slice conditioning for a specific batch index.""" - # Clone and slice conditioning properly, including ControlNet hints - 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[batch_idx:batch_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[batch_idx:batch_idx+1] - control = control.previous_controlnet - if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: - cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] - - return positive_sliced, negative_sliced - - async def _get_all_completed_tasks(self, multi_job_id): - """Helper to retrieve all completed tasks from the job data.""" - 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 JOB_COMPLETED_TASKS in job_data: - return dict(job_data[JOB_COMPLETED_TASKS]) # Return a copy - return {} - - # Note: Removed unused _process_worker_static_async to reduce redundancy - - def _process_worker_static_sync(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_workers): - """Worker static mode processing with optional dynamic queue pulling.""" - # Round tile dimensions - tile_width = self.round_to_multiple(tile_width) - tile_height = self.round_to_multiple(tile_height) - - # Get dimensions and calculate tiles - _, height, width, _ = upscaled_image.shape - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - num_tiles_per_image = len(all_tiles) - batch_size = upscaled_image.shape[0] - total_tiles = batch_size * num_tiles_per_image - - processed_tiles = [] - sliced_conditioning_cache = {} - - # Dynamic queue mode (static processing): process 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 - max_poll_attempts = 20 - 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) - else: - log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") - return (upscaled_image,) - - # Main processing loop - pull tile ids from queue - while True: - # Request a tile to process - tile_idx, estimated_remaining, batched_static = run_async_in_server_loop( - self._request_tile_from_master(multi_job_id, master_url, worker_id), - timeout=TILE_WAIT_TIMEOUT - ) - - if tile_idx is None: - debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process") - break - - # Always batched-per-tile in static mode - 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] - # Extract tile for entire batch - tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( - upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles - ) - # Process batch - region = (x1, y1, x1 + ew, y1 + eh) - processed_batch = self.process_tiles_batch( - tile_batch, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, - region, (width, height) - ) - # Queue results - for b in range(batch_size): - processed_tiles.append({ - 'tile': processed_batch[b:b+1], - 'tile_idx': tile_id, - 'x': x1, - 'y': y1, - 'extracted_width': ew, - 'extracted_height': eh, - 'padding': padding, - 'batch_idx': b, - 'global_idx': b * num_tiles_per_image + tile_id - }) - - # Send heartbeat - try: - run_async_in_server_loop( - _send_heartbeat_to_master(multi_job_id, master_url, worker_id), - timeout=5.0 - ) - except Exception as e: - debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}") - - # Send tiles in batches within loop - 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: - 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 - ) - - debug_log(f"Worker {worker_id} completed all assigned and requeued tiles") - return (upscaled_image,) - - async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total): - """Async helper for collection and monitoring in static mode. - Returns collected tasks dict. Caller should check if all tasks are complete.""" - last_progress_log = time.time() - progress_interval = 5.0 - last_heartbeat_check = time.time() - last_completed_count = 0 - - while True: - # Check for user interruption - if comfy.model_management.processing_interrupted(): - log("Processing interrupted by user") - raise comfy.model_management.InterruptProcessingException() - - # Drain any pending results - collected_count = await _drain_results_queue(multi_job_id) - - # Check and requeue timed-out workers periodically - current_time = time.time() - if current_time - last_heartbeat_check >= 10.0: - requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total) - if requeued_count > 0: - log(f"Requeued {requeued_count} tasks from timed-out workers") - last_heartbeat_check = current_time - - # Get current completion count - completed_count = await _get_completed_count(multi_job_id) - - # Progress logging - if current_time - last_progress_log >= progress_interval: - log(f"Progress: {completed_count}/{expected_total} tasks completed") - last_progress_log = current_time - - # Check if all tasks are completed - if completed_count >= expected_total: - debug_log(f"All {expected_total} tasks completed") - break - - # If no active workers remain and there are pending tasks, return for local processing - 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: - pending_queue = job_data.get(JOB_PENDING_TASKS) - active_workers = list(job_data.get(JOB_WORKER_STATUS, {}).keys()) - if pending_queue and not pending_queue.empty() and len(active_workers) == 0: - log(f"No active workers remaining with {expected_total - completed_count} tasks pending. Returning for local processing.") - break - - # Wait a bit before next check - await asyncio.sleep(0.1) - - # Get all completed tasks for return - return await self._get_all_completed_tasks(multi_job_id) - - def _process_master_static_sync(self, upscaled_image, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, - all_tiles, num_tiles_per_image): - """Static mode master processing with optional dynamic queue pulling.""" - batch_size = upscaled_image.shape[0] - _, height, width, _ = upscaled_image.shape - total_tiles = batch_size * num_tiles_per_image - - # Convert batch to PIL list for processing - result_images = [] - for b in range(batch_size): - image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) - result_images.append(image_pil.copy()) - - sliced_conditioning_cache = {} - # Initialize queue: pending queue holds tile ids (batched per tile) - log("USDU Dist: Using tile queue distribution") - run_async_in_server_loop( - init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers), - timeout=10.0 - ) - debug_log( - 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 = 2 - - while processed_count < total_tiles: - comfy.model_management.throw_exception_if_processing_interrupted() - tile_idx = run_async_in_server_loop( - self._get_next_tile_index(multi_job_id), - timeout=5.0 - ) - if tile_idx is not None: - consecutive_no_tile = 0 - processed_count += batch_size - tile_id = tile_idx - tx, ty = all_tiles[tile_id] - tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( - upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles - ) - region = (x1, y1, x1 + ew, y1 + eh) - processed_batch = self.process_tiles_batch( - tile_batch, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, - region, (width, height) - ) - 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): - tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) - result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding) - global_idx = b * num_tiles_per_image + tile_id - run_async_in_server_loop( - _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: - debug_log(f"Master processed {processed_count} tiles, moving to collection phase") - break - time.sleep(0.1) - master_processed_count = processed_count - - # Continue processing any remaining tiles while collecting worker results - remaining_tiles = total_tiles - master_processed_count - if remaining_tiles > 0: - debug_log(f"Master waiting for {remaining_tiles} tiles from workers") - - # Collect worker results using async operations - try: - # Wait until either all tasks are collected or there are no active workers left - collected_tasks = run_async_in_server_loop( - self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles), - timeout=None - ) - except comfy.model_management.InterruptProcessingException: - # Clean up job on interruption - run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) - raise - - # Check if we need to process any remaining tasks locally after collection - completed_count = len(collected_tasks) - if completed_count < total_tiles: - log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures") - - # Process any remaining pending tasks (batched-per-tile) - while True: - # Check for user interruption - comfy.model_management.throw_exception_if_processing_interrupted() - - # Get next tile_id from pending queue - tile_id = run_async_in_server_loop( - self._get_next_tile_index(multi_job_id), - timeout=5.0 - ) - - if tile_id is None: - break - - # Extract batched tile and process across available batch - tx, ty = all_tiles[tile_id] - tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( - upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles - ) - region = (x1, y1, x1 + ew, y1 + eh) - processed_batch = self.process_tiles_batch( - tile_batch, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, - region, (width, height) - ) - tile_mask = tile_masks[tile_id] - out_bs = processed_batch.shape[0] if hasattr(processed_batch, 'shape') else batch_size - for b in range(min(batch_size, out_bs)): - tile_pil = tensor_to_pil(processed_batch, b) - if tile_pil.size != (ew, eh): - tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) - result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding) - global_idx = b * num_tiles_per_image + tile_id - # Mark as completed so the collector state is consistent - run_async_in_server_loop( - _mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}), - timeout=5.0 - ) - else: - # Master processed all tiles - collected_tasks = run_async_in_server_loop( - self._get_all_completed_tasks(multi_job_id), - timeout=5.0 - ) - - # 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 and 'image' not in tile_data: - continue - - batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image) - tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image) - - if batch_idx >= batch_size: - continue - - # Blend tile synchronously - x = tile_data.get('x', 0) - y = tile_data.get('y', 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 = 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, - x, y, (extracted_width, extracted_height), tile_mask, padding) - - try: - # Convert back to tensor - if batch_size == 1: - result_tensor = pil_to_tensor(result_images[0]) - else: - result_tensors = [pil_to_tensor(img) for img in result_images] - result_tensor = torch.cat(result_tensors, dim=0) - - if upscaled_image.is_cuda: - result_tensor = result_tensor.cuda() - - log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") - return (result_tensor,) - finally: - # Cleanup (async operation) - always execute - run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) - - def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles, - model, positive, negative, vae, seed, steps, cfg, sampler_name, - scheduler, denoise, tiled_decode, tile_width, tile_height, padding, - width, height, force_uniform_tiles, sliced_conditioning_cache): - """Process a single tile.""" - # Calculate which image and tile this corresponds to - batch_idx = global_idx // num_tiles_per_image - tile_idx = global_idx % num_tiles_per_image - - # Skip if batch_idx is out of range - if batch_idx >= upscaled_image.shape[0]: - debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}") - return None - - # Get or create sliced conditioning for this batch index - if batch_idx not in sliced_conditioning_cache: - positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) - 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] - - # Extract tile from the specific image in the batch - tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( - upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding, force_uniform_tiles - ) - - # Process tile through SD with unique seed - image_seed = seed + batch_idx * 1000 - processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae, - image_seed, steps, cfg, sampler_name, - scheduler, denoise, tiled_decode, batch_idx=batch_idx, - region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height)) - - return { - 'tile': processed_tile, - 'global_idx': global_idx, - 'batch_idx': batch_idx, - 'tile_idx': tile_idx, - 'x': x1, - 'y': y1, - 'extracted_width': ew, - 'extracted_height': eh - } - - async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, - padding, worker_id): - """Send all processed tiles to master, chunked if large.""" - if not processed_tiles: - return # Early exit if empty - - total_tiles = len(processed_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 = [] - for idx, tile_data in enumerate(processed_tiles): - img = tensor_to_pil(tile_data['tile'], 0) - bio = io.BytesIO() - # Keep compression low to balance speed and size; adjust if needed - img.save(bio, format='PNG', compress_level=0) - raw = bio.getvalue() - encoded.append({ - 'bytes': raw, - 'meta': { - 'tile_idx': tile_data['tile_idx'], - 'x': tile_data['x'], - 'y': tile_data['y'], - 'extracted_width': tile_data['extracted_width'], - 'extracted_height': tile_data['extracted_height'], - **({'batch_idx': tile_data['batch_idx']} if 'batch_idx' in tile_data else {}), - **({'global_idx': tile_data['global_idx']} if 'global_idx' in tile_data else {}), - } - }) - - # Size-aware chunking - max_bytes = int(MAX_PAYLOAD_SIZE) - (1024 * 1024) # 1MB headroom - i = 0 - chunk_index = 0 - while i < total_tiles: - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('padding', str(padding)) - - metadata = [] - used = 0 - j = i - while j < total_tiles: - img_bytes = encoded[j]['bytes'] - meta = encoded[j]['meta'] - # Rough overhead for fields + JSON - overhead = 1024 - if used + len(img_bytes) + overhead > max_bytes and j > i: - break - # Accept this tile in this chunk - metadata.append(meta) - data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png') - used += len(img_bytes) + overhead - j += 1 - - # Ensure at least one tile per chunk - if j == i: - # Single oversized tile, send anyway - meta = encoded[j]['meta'] - metadata.append(meta) - data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png') - j += 1 - - chunk_size = j - i - is_chunk_last = (j >= total_tiles) - data.add_field('is_last', str(is_chunk_last)) - data.add_field('batch_size', str(chunk_size)) - data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json') - - # Retry logic with exponential backoff - max_retries = 5 - retry_delay = 0.5 - for attempt in range(max_retries): - try: - session = await get_client_session() - url = f"{master_url}/distributed/submit_tiles" - async with session.post(url, data=data) as response: - response.raise_for_status() - break - except Exception as e: - if attempt < max_retries - 1: - await asyncio.sleep(retry_delay) - retry_delay = min(retry_delay * 2, 5.0) - else: - log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}") - raise - - debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)") - chunk_index += 1 - i = j - - def process_single_gpu(self, upscaled_image, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode): - """Process all tiles on a single GPU (no distribution), batching per tile like USDU.""" - # Round tile dimensions - tile_width = self.round_to_multiple(tile_width) - tile_height = self.round_to_multiple(tile_height) - - # Get image dimensions and batch size - batch_size, height, width, _ = upscaled_image.shape - - # Calculate all tiles - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - - 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 = [] - for b in range(batch_size): - image_pil = tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB') - result_images.append(image_pil.copy()) - - # Precompute tile masks once - tile_masks = [] - for tx, ty in all_tiles: - tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)) - - # Process tiles batched across images - for tile_idx, (tx, ty) in enumerate(all_tiles): - # Extract batched tile - tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( - upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles - ) - - # Process batch - region = (x1, y1, x1 + ew, y1 + eh) - processed_batch = self.process_tiles_batch(tile_batch, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, - tiled_decode, region, (width, height)) - - # Blend results back into each image using cached mask - tile_mask = tile_masks[tile_idx] - for b in range(batch_size): - tile_pil = tensor_to_pil(processed_batch, b) - # Resize back to extracted size - if tile_pil.size != (ew, eh): - tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) - result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding) - - # Convert back to tensor - result_tensors = [pil_to_tensor(img) for img in result_images] - result_tensor = torch.cat(result_tensors, dim=0) - if upscaled_image.is_cuda: - result_tensor = result_tensor.cuda() - - return (result_tensor,) - - def process_master_dynamic(self, upscaled_image, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers): - """Dynamic mode for large batches - assigns whole images to workers dynamically, including master.""" - # 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 = self.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( - self._get_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] - 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 - ) - - # Yield after each tile to minimize worker downtime - run_async_in_server_loop(self._async_yield(), timeout=0.1) - # Note: No per-tile drain here – that's what makes this "per-image" - - result_images[image_idx] = local_image - - # Mark as completed - run_async_in_server_loop( - self._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( - self._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( - self._get_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(self._async_yield(), timeout=0.1) - else: - # Queue empty: collect any queued worker results to update progress - drained_count = run_async_in_server_loop( - self._drain_worker_results_queue(multi_job_id), - timeout=5.0 - ) - run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after drain - - # 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 - ) - run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after requeue - - 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( - self._get_total_completed_count(multi_job_id), - timeout=1.0 - ) - - log(f"USDU Dist: Images progress {completed_now}/{batch_size}") - - if completed_now >= batch_size: - break - - run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield before pending check - - # 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: - 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.") - break # Force exit to collection phase - - debug_log("Waiting for workers") - # Use async sleep to allow event loop to process worker requests - run_async_in_server_loop(asyncio.sleep(2), timeout=3.0) - - debug_log(f"Master processed {processed_count} images locally") - - # Get all completed images to check what needs to be collected - all_completed = run_async_in_server_loop( - self._get_all_completed_images(multi_job_id), - timeout=5.0 - ) - - # Calculate how many we still need to collect - 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") - # Use the unified worker timeout for the collection phase - collection_timeout = float(get_worker_timeout_seconds()) - collected_images = run_async_in_server_loop( - self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count), - timeout=collection_timeout - ) - - # Merge collected with already completed - all_completed.update(collected_images) - - # Update result images with all completed images - for idx, processed_img in all_completed.items(): - if idx < batch_size: - result_images[idx] = processed_img - - # Convert back to tensor - 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() - - debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") - log(f"Completed processing all {batch_size} images") - return (result_tensor,) - - - async def _async_yield(self): - """Simple async yield to allow event loop processing.""" - await asyncio.sleep(0) - - async def _get_next_image_index(self, multi_job_id): - """Get next image index from pending queue for master.""" - 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 'pending_images' not in job_data: - return None - try: - image_idx = await asyncio.wait_for(job_data['pending_images'].get(), timeout=1.0) - return image_idx - except asyncio.TimeoutError: - return None - - async def _get_next_tile_index(self, multi_job_id): - """Get next tile index from pending queue for master in static mode.""" - 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: - tile_idx = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=0.1) - return tile_idx - except asyncio.TimeoutError: - return None - - - async def _get_total_completed_count(self, multi_job_id): - """Get total count of all completed images (master + workers).""" - 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 'completed_images' in job_data: - return len(job_data['completed_images']) - return 0 - - async def _get_all_completed_images(self, multi_job_id): - """Get all completed 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 job_data and 'completed_images' in job_data: - 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() - 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 'queue' not in job_data or 'completed_images' not in job_data: - return 0 - q = job_data['queue'] - completed_images = job_data['completed_images'] - - collected = 0 - while not q.empty(): - try: - result = await asyncio.wait_for(q.get(), timeout=0.1) - worker_id = result['worker_id'] - is_last = result.get('is_last', False) - - if 'image_idx' in result and 'image' in result: - image_idx = result['image_idx'] - image_pil = result['image'] - if image_idx not in completed_images: - completed_images[image_idx] = image_pil - collected += 1 - debug_log(f"Drained image {image_idx} from worker {worker_id}") - - if is_last: - # Optional: track worker completion if needed - pass - except asyncio.TimeoutError: - break # No more immediately available - - if collected > 0: - debug_log(f"Drained {collected} worker images during retry") - - 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.""" - # Use the original function from usdu_managment.py - return await _check_and_requeue_timed_out_workers(multi_job_id, batch_size) - - - 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 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 - ) - - # 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"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.""" - # Enhanced retries with 404-specific delay, backoff cap, and total timeout to handle init races - max_retries = 10 - retry_delay = 0.5 - start_time = time.time() - - for attempt in range(max_retries): - # Check total timeout - if time.time() - start_time > 30: - log(f"Total request timeout after 30s for worker {worker_id}") - return None, 0 - - try: - session = await get_client_session() - url = f"{master_url}/distributed/request_image" - - async with session.post(url, json={ - 'worker_id': str(worker_id), - 'multi_job_id': multi_job_id - }) as response: - if response.status == 200: - data = await response.json() - image_idx = data.get('image_idx') - estimated_remaining = data.get('estimated_remaining', 0) - return image_idx, estimated_remaining - elif response.status == 404: - # Special handling for 404 - job not found yet - text = await response.text() - debug_log(f"Job not found (404), will retry: {text}") - await asyncio.sleep(1.0) - else: - text = await response.text() - debug_log(f"Request image failed: {response.status} - {text}") - - except Exception as e: - if attempt < max_retries - 1: - debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") - await asyncio.sleep(retry_delay) - retry_delay *= 2 - retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s - else: - log(f"Failed to request image after {max_retries} attempts: {e}") - raise - - return None, 0 - - async def _request_tile_from_master(self, multi_job_id, master_url, worker_id): - """Request a tile index to process from master in static mode (reusing dynamic infrastructure).""" - # Reuse the same retry logic as dynamic mode - max_retries = 10 - retry_delay = 0.5 - start_time = time.time() - - for attempt in range(max_retries): - # Check total timeout - if time.time() - start_time > 30: - log(f"Total request timeout after 30s for worker {worker_id}") - return None, 0 - - try: - session = await get_client_session() - url = f"{master_url}/distributed/request_image" # Same endpoint - - async with session.post(url, json={ - 'worker_id': str(worker_id), - 'multi_job_id': multi_job_id - }) as response: - if response.status == 200: - data = await response.json() - tile_idx = data.get('tile_idx') - estimated_remaining = data.get('estimated_remaining', 0) - batched_static = data.get('batched_static', False) - return tile_idx, estimated_remaining, batched_static - elif response.status == 404: - # Special handling for 404 - job not found yet - text = await response.text() - debug_log(f"Job not found (404), will retry: {text}") - await asyncio.sleep(1.0) - else: - text = await response.text() - debug_log(f"Request tile failed: {response.status} - {text}") - - except Exception as e: - if attempt < max_retries - 1: - debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") - await asyncio.sleep(retry_delay) - retry_delay *= 2 - retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s - else: - log(f"Failed to request tile after {max_retries} attempts: {e}") - raise - - return None, 0, False - - async def _check_job_status(self, multi_job_id, master_url): - """Check if job is ready on the master.""" - try: - session = await get_client_session() - url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}" - async with session.get(url) as response: - if response.status == 200: - data = await response.json() - return data.get('ready', False) - return False - except Exception as e: - debug_log(f"Job status check failed: {e}") - return False - - async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, - master_url, worker_id, is_last): - """Send a processed full image back to master in dynamic mode.""" - # Serialize image to PNG - byte_io = io.BytesIO() - image_pil.save(byte_io, format='PNG', compress_level=0) - byte_io.seek(0) - - # Prepare form data - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('image_idx', str(image_idx)) - data.add_field('is_last', str(is_last)) - data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png', - content_type='image/png') - - # Retry logic - max_retries = 5 - retry_delay = 0.5 - - for attempt in range(max_retries): - try: - session = await get_client_session() - url = f"{master_url}/distributed/submit_image" - - async with session.post(url, data=data) as response: - response.raise_for_status() - debug_log(f"Successfully sent image {image_idx} to master") - return - - except Exception as e: - if attempt < max_retries - 1: - debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") - await asyncio.sleep(retry_delay) - retry_delay *= 2 - else: - log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}") - raise - - # Note: Using _send_heartbeat_to_master from usdu_managment.py instead of duplicate - - async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id): - """Send completion signal to master in dynamic mode.""" - # Send a dummy request with is_last=True - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('is_last', 'true') - # No image data - just completion signal - - session = await get_client_session() - url = f"{master_url}/distributed/submit_image" - - async with session.post(url, data=data) as response: - 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 processing mode per requested policy: diff --git a/upscale/__init__.py b/upscale/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/upscale/job_state.py b/upscale/job_state.py new file mode 100644 index 0000000..6df741d --- /dev/null +++ b/upscale/job_state.py @@ -0,0 +1,113 @@ +import asyncio, time +import server +from ..utils.logging import debug_log +from ..utils.usdu_managment import ( + ensure_tile_jobs_initialized, + _check_and_requeue_timed_out_workers as _requeue_usdu, + JOB_COMPLETED_TASKS, JOB_PENDING_TASKS, +) + + +class JobStateMixin: + async def _get_all_completed_tasks(self, multi_job_id): + """Helper to retrieve all completed tasks from the job data.""" + 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 JOB_COMPLETED_TASKS in job_data: + return dict(job_data[JOB_COMPLETED_TASKS]) # Return a copy + return {} + + async def _get_next_image_index(self, multi_job_id): + """Get next image index from pending queue for master.""" + 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 'pending_images' not in job_data: + return None + try: + image_idx = await asyncio.wait_for(job_data['pending_images'].get(), timeout=1.0) + return image_idx + except asyncio.TimeoutError: + return None + + async def _get_next_tile_index(self, multi_job_id): + """Get next tile index from pending queue for master in static mode.""" + 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: + tile_idx = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=0.1) + return tile_idx + except asyncio.TimeoutError: + return None + + async def _get_total_completed_count(self, multi_job_id): + """Get total count of all completed images (master + workers).""" + 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 'completed_images' in job_data: + return len(job_data['completed_images']) + return 0 + + async def _get_all_completed_images(self, multi_job_id): + """Get all completed 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 job_data and 'completed_images' in job_data: + 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() + 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 'queue' not in job_data or 'completed_images' not in job_data: + return 0 + q = job_data['queue'] + completed_images = job_data['completed_images'] + + collected = 0 + while not q.empty(): + try: + result = await asyncio.wait_for(q.get(), timeout=0.1) + worker_id = result['worker_id'] + is_last = result.get('is_last', False) + + if 'image_idx' in result and 'image' in result: + image_idx = result['image_idx'] + image_pil = result['image'] + if image_idx not in completed_images: + completed_images[image_idx] = image_pil + collected += 1 + debug_log(f"Drained image {image_idx} from worker {worker_id}") + + if is_last: + # Optional: track worker completion if needed + pass + except asyncio.TimeoutError: + break # No more immediately available + + if collected > 0: + debug_log(f"Drained {collected} worker images during retry") + + 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.""" + # Use the original function from usdu_managment.py + return await _requeue_usdu(multi_job_id, batch_size) diff --git a/upscale/modes/__init__.py b/upscale/modes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/upscale/modes/dynamic.py b/upscale/modes/dynamic.py new file mode 100644 index 0000000..2a1e0ef --- /dev/null +++ b/upscale/modes/dynamic.py @@ -0,0 +1,340 @@ +import asyncio, torch +from PIL import Image +import comfy.model_management +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 TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT +from ...utils.usdu_managment import ( + ensure_tile_jobs_initialized, init_dynamic_job, + clone_conditioning, _send_heartbeat_to_master, +) + + +class DynamicModeMixin: + def process_master_dynamic(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers): + """Dynamic mode for large batches - assigns whole images to workers dynamically, including master.""" + # 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 = self.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( + self._get_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] + 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 + ) + + # Yield after each tile to minimize worker downtime + run_async_in_server_loop(self._async_yield(), timeout=0.1) + # Note: No per-tile drain here – that's what makes this "per-image" + + result_images[image_idx] = local_image + + # Mark as completed + run_async_in_server_loop( + self._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( + self._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( + self._get_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(self._async_yield(), timeout=0.1) + else: + # Queue empty: collect any queued worker results to update progress + drained_count = run_async_in_server_loop( + self._drain_worker_results_queue(multi_job_id), + timeout=5.0 + ) + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after drain + + # 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 + ) + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after requeue + + 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( + self._get_total_completed_count(multi_job_id), + timeout=1.0 + ) + + log(f"USDU Dist: Images progress {completed_now}/{batch_size}") + + if completed_now >= batch_size: + break + + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield before pending check + + # 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: + 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.") + break # Force exit to collection phase + + debug_log("Waiting for workers") + # Use async sleep to allow event loop to process worker requests + run_async_in_server_loop(asyncio.sleep(2), timeout=3.0) + + debug_log(f"Master processed {processed_count} images locally") + + # Get all completed images to check what needs to be collected + all_completed = run_async_in_server_loop( + self._get_all_completed_images(multi_job_id), + timeout=5.0 + ) + + # Calculate how many we still need to collect + 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") + # Use the unified worker timeout for the collection phase + collection_timeout = float(get_worker_timeout_seconds()) + collected_images = run_async_in_server_loop( + self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count), + timeout=collection_timeout + ) + + # Merge collected with already completed + all_completed.update(collected_images) + + # Update result images with all completed images + for idx, processed_img in all_completed.items(): + if idx < batch_size: + result_images[idx] = processed_img + + # Convert back to tensor + 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() + + 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, 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 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 + ) + + # 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"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,) diff --git a/upscale/modes/single_gpu.py b/upscale/modes/single_gpu.py new file mode 100644 index 0000000..9b96257 --- /dev/null +++ b/upscale/modes/single_gpu.py @@ -0,0 +1,67 @@ +import math, torch +from PIL import Image +from ...utils.logging import debug_log, log +from ...utils.image import tensor_to_pil, pil_to_tensor + + +class SingleGpuModeMixin: + def process_single_gpu(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode): + """Process all tiles on a single GPU (no distribution), batching per tile like USDU.""" + # Round tile dimensions + tile_width = self.round_to_multiple(tile_width) + tile_height = self.round_to_multiple(tile_height) + + # Get image dimensions and batch size + batch_size, height, width, _ = upscaled_image.shape + + # Calculate all tiles + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + + 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 = [] + for b in range(batch_size): + image_pil = tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB') + result_images.append(image_pil.copy()) + + # Precompute tile masks once + tile_masks = [] + for tx, ty in all_tiles: + tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)) + + # Process tiles batched across images + for tile_idx, (tx, ty) in enumerate(all_tiles): + # Extract batched tile + tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( + upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles + ) + + # Process batch + region = (x1, y1, x1 + ew, y1 + eh) + processed_batch = self.process_tiles_batch(tile_batch, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tiled_decode, region, (width, height)) + + # Blend results back into each image using cached mask + tile_mask = tile_masks[tile_idx] + for b in range(batch_size): + tile_pil = tensor_to_pil(processed_batch, b) + # Resize back to extracted size + if tile_pil.size != (ew, eh): + tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) + result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding) + + # Convert back to tensor + result_tensors = [pil_to_tensor(img) for img in result_images] + result_tensor = torch.cat(result_tensors, dim=0) + if upscaled_image.is_cuda: + result_tensor = result_tensor.cuda() + + return (result_tensor,) diff --git a/upscale/modes/static.py b/upscale/modes/static.py new file mode 100644 index 0000000..ae7be53 --- /dev/null +++ b/upscale/modes/static.py @@ -0,0 +1,367 @@ +import asyncio, time, torch +from PIL import Image +import comfy.model_management +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 TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH +from ...utils.usdu_managment import ( + ensure_tile_jobs_initialized, init_static_job_batched, + _mark_task_completed, _cleanup_job, _drain_results_queue, + _send_heartbeat_to_master, + JOB_PENDING_TASKS, JOB_WORKER_STATUS, +) + + +class StaticModeMixin: + def _process_worker_static_sync(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_workers): + """Worker static mode processing with optional dynamic queue pulling.""" + # Round tile dimensions + tile_width = self.round_to_multiple(tile_width) + tile_height = self.round_to_multiple(tile_height) + + # Get dimensions and calculate tiles + _, height, width, _ = upscaled_image.shape + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + num_tiles_per_image = len(all_tiles) + batch_size = upscaled_image.shape[0] + total_tiles = batch_size * num_tiles_per_image + + processed_tiles = [] + sliced_conditioning_cache = {} + + # Dynamic queue mode (static processing): process 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 + max_poll_attempts = 20 + 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) + else: + log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") + return (upscaled_image,) + + # Main processing loop - pull tile ids from queue + while True: + # Request a tile to process + tile_idx, estimated_remaining, batched_static = run_async_in_server_loop( + self._request_tile_from_master(multi_job_id, master_url, worker_id), + timeout=TILE_WAIT_TIMEOUT + ) + + if tile_idx is None: + debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process") + break + + # Always batched-per-tile in static mode + 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] + # Extract tile for entire batch + tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( + upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles + ) + # Process batch + region = (x1, y1, x1 + ew, y1 + eh) + processed_batch = self.process_tiles_batch( + tile_batch, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, + region, (width, height) + ) + # Queue results + for b in range(batch_size): + processed_tiles.append({ + 'tile': processed_batch[b:b+1], + 'tile_idx': tile_id, + 'x': x1, + 'y': y1, + 'extracted_width': ew, + 'extracted_height': eh, + 'padding': padding, + 'batch_idx': b, + 'global_idx': b * num_tiles_per_image + tile_id + }) + + # Send heartbeat + try: + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + except Exception as e: + debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}") + + # Send tiles in batches within loop + 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: + 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 + ) + + debug_log(f"Worker {worker_id} completed all assigned and requeued tiles") + return (upscaled_image,) + + async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total): + """Async helper for collection and monitoring in static mode. + Returns collected tasks dict. Caller should check if all tasks are complete.""" + last_progress_log = time.time() + progress_interval = 5.0 + last_heartbeat_check = time.time() + last_completed_count = 0 + + while True: + # Check for user interruption + if comfy.model_management.processing_interrupted(): + log("Processing interrupted by user") + raise comfy.model_management.InterruptProcessingException() + + # Drain any pending results + collected_count = await _drain_results_queue(multi_job_id) + + # Check and requeue timed-out workers periodically + current_time = time.time() + if current_time - last_heartbeat_check >= 10.0: + requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total) + if requeued_count > 0: + log(f"Requeued {requeued_count} tasks from timed-out workers") + last_heartbeat_check = current_time + + # Get current completion count + completed_count = await _get_completed_count(multi_job_id) + + # Progress logging + if current_time - last_progress_log >= progress_interval: + log(f"Progress: {completed_count}/{expected_total} tasks completed") + last_progress_log = current_time + + # Check if all tasks are completed + if completed_count >= expected_total: + debug_log(f"All {expected_total} tasks completed") + break + + # If no active workers remain and there are pending tasks, return for local processing + 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: + pending_queue = job_data.get(JOB_PENDING_TASKS) + active_workers = list(job_data.get(JOB_WORKER_STATUS, {}).keys()) + if pending_queue and not pending_queue.empty() and len(active_workers) == 0: + log(f"No active workers remaining with {expected_total - completed_count} tasks pending. Returning for local processing.") + break + + # Wait a bit before next check + await asyncio.sleep(0.1) + + # Get all completed tasks for return + return await self._get_all_completed_tasks(multi_job_id) + + def _process_master_static_sync(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, + all_tiles, num_tiles_per_image): + """Static mode master processing with optional dynamic queue pulling.""" + batch_size = upscaled_image.shape[0] + _, height, width, _ = upscaled_image.shape + total_tiles = batch_size * num_tiles_per_image + + # Convert batch to PIL list for processing + result_images = [] + for b in range(batch_size): + image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) + result_images.append(image_pil.copy()) + + sliced_conditioning_cache = {} + # Initialize queue: pending queue holds tile ids (batched per tile) + log("USDU Dist: Using tile queue distribution") + run_async_in_server_loop( + init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers), + timeout=10.0 + ) + debug_log( + 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 = 2 + + while processed_count < total_tiles: + comfy.model_management.throw_exception_if_processing_interrupted() + tile_idx = run_async_in_server_loop( + self._get_next_tile_index(multi_job_id), + timeout=5.0 + ) + if tile_idx is not None: + consecutive_no_tile = 0 + processed_count += batch_size + tile_id = tile_idx + tx, ty = all_tiles[tile_id] + tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( + upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles + ) + region = (x1, y1, x1 + ew, y1 + eh) + processed_batch = self.process_tiles_batch( + tile_batch, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, + region, (width, height) + ) + 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): + tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) + result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding) + global_idx = b * num_tiles_per_image + tile_id + run_async_in_server_loop( + _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: + debug_log(f"Master processed {processed_count} tiles, moving to collection phase") + break + time.sleep(0.1) + master_processed_count = processed_count + + # Continue processing any remaining tiles while collecting worker results + remaining_tiles = total_tiles - master_processed_count + if remaining_tiles > 0: + debug_log(f"Master waiting for {remaining_tiles} tiles from workers") + + # Collect worker results using async operations + try: + # Wait until either all tasks are collected or there are no active workers left + collected_tasks = run_async_in_server_loop( + self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles), + timeout=None + ) + except comfy.model_management.InterruptProcessingException: + # Clean up job on interruption + run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) + raise + + # Check if we need to process any remaining tasks locally after collection + completed_count = len(collected_tasks) + if completed_count < total_tiles: + log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures") + + # Process any remaining pending tasks (batched-per-tile) + while True: + # Check for user interruption + comfy.model_management.throw_exception_if_processing_interrupted() + + # Get next tile_id from pending queue + tile_id = run_async_in_server_loop( + self._get_next_tile_index(multi_job_id), + timeout=5.0 + ) + + if tile_id is None: + break + + # Extract batched tile and process across available batch + tx, ty = all_tiles[tile_id] + tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( + upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles + ) + region = (x1, y1, x1 + ew, y1 + eh) + processed_batch = self.process_tiles_batch( + tile_batch, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, + region, (width, height) + ) + tile_mask = tile_masks[tile_id] + out_bs = processed_batch.shape[0] if hasattr(processed_batch, 'shape') else batch_size + for b in range(min(batch_size, out_bs)): + tile_pil = tensor_to_pil(processed_batch, b) + if tile_pil.size != (ew, eh): + tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) + result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding) + global_idx = b * num_tiles_per_image + tile_id + # Mark as completed so the collector state is consistent + run_async_in_server_loop( + _mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}), + timeout=5.0 + ) + else: + # Master processed all tiles + collected_tasks = run_async_in_server_loop( + self._get_all_completed_tasks(multi_job_id), + timeout=5.0 + ) + + # 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 and 'image' not in tile_data: + continue + + batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image) + tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image) + + if batch_idx >= batch_size: + continue + + # Blend tile synchronously + x = tile_data.get('x', 0) + y = tile_data.get('y', 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 = 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, + x, y, (extracted_width, extracted_height), tile_mask, padding) + + try: + # Convert back to tensor + if batch_size == 1: + result_tensor = pil_to_tensor(result_images[0]) + else: + result_tensors = [pil_to_tensor(img) for img in result_images] + result_tensor = torch.cat(result_tensors, dim=0) + + if upscaled_image.is_cuda: + result_tensor = result_tensor.cuda() + + log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") + return (result_tensor,) + finally: + # Cleanup (async operation) - always execute + run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) diff --git a/upscale/result_collector.py b/upscale/result_collector.py new file mode 100644 index 0000000..b87da93 --- /dev/null +++ b/upscale/result_collector.py @@ -0,0 +1,165 @@ +import asyncio, time +import comfy.model_management +import server +from ..utils.logging import debug_log, log +from ..utils.config import get_worker_timeout_seconds +from ..utils.usdu_managment import ( + ensure_tile_jobs_initialized, _mark_task_completed, + _check_and_requeue_timed_out_workers, + JOB_COMPLETED_TASKS, JOB_WORKER_STATUS, JOB_PENDING_TASKS, +) + + +class ResultCollectorMixin: + async def _async_collect_results(self, multi_job_id, num_workers, mode='static', + remaining_to_collect=None, batch_size=None): + """Unified async helper to collect results from workers (tiles or images).""" + # Get the already initialized queue + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id not in prompt_server.distributed_pending_tile_jobs: + raise RuntimeError(f"Job queue not initialized for {multi_job_id}") + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + raise RuntimeError("Invalid job data structure") + if job_data['mode'] != mode: + raise RuntimeError(f"Mode mismatch: expected {mode}, got {job_data['mode']}") + q = job_data['queue'] + + # For dynamic mode, get reference to completed_images + if mode == 'dynamic': + completed_images = job_data['completed_images'] + # Calculate expected count for logging + expected_count = remaining_to_collect or batch_size + else: + # Calculate expected count from job data for static mode + expected_count = job_data.get('total_items', 0) - len(job_data.get('completed_items', set())) + + item_type = "images" if mode == 'dynamic' else "tiles" + debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers") + + collected_results = {} + workers_done = set() + # Unify collector/upscaler wait behavior with the UI worker timeout + timeout = float(get_worker_timeout_seconds()) + last_heartbeat_check = time.time() + collected_count = 0 + + while len(workers_done) < num_workers: + # Check for user interruption + if comfy.model_management.processing_interrupted(): + log("Processing interrupted by user") + raise comfy.model_management.InterruptProcessingException() + + # For dynamic mode with remaining_to_collect, check if we've collected enough + if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect: + break + + try: + # Shorter poll for dynamic mode, but never exceed the configured timeout + wait_timeout = (min(10.0, timeout) if mode == 'dynamic' else timeout) + result = await asyncio.wait_for(q.get(), timeout=wait_timeout) + worker_id = result['worker_id'] + is_last = result.get('is_last', False) + + if mode == 'static': + # Handle tiles + tiles = result.get('tiles', []) + if tiles: + # Batch mode + debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})") + + for tile_data in tiles: + # Validate required fields + if 'batch_idx' not in tile_data: + log(f"UltimateSDUpscale Master - Missing batch_idx in tile data, skipping") + continue + + tile_idx = tile_data['tile_idx'] + # Use global_idx as key if available (for batch processing) + key = tile_data.get('global_idx', tile_idx) + + # Store the full tile data including metadata; prefer PIL image if present + entry = { + 'tile_idx': tile_idx, + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'], + 'padding': tile_data['padding'], + 'worker_id': worker_id, + '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'] + collected_results[tile_idx] = result + debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})") + + elif mode == 'dynamic': + # Handle full images + if 'image_idx' in result and 'image' in result: + image_idx = result['image_idx'] + image_pil = result['image'] + completed_images[image_idx] = image_pil + collected_results[image_idx] = image_pil + collected_count += 1 + 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 - Worker {worker_id} completed") + + except asyncio.TimeoutError: + if mode == 'dynamic': + # Check for worker timeouts periodically + current_time = time.time() + if current_time - last_heartbeat_check >= 10.0: + # 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 - 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 - Overall timeout waiting for images") + break + else: + log(f"UltimateSDUpscale Master - Timeout waiting for {item_type}") + break + + 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: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + del prompt_server.distributed_pending_tile_jobs[multi_job_id] + + return collected_results if mode == 'static' else completed_images + + async def _async_collect_worker_tiles(self, multi_job_id, num_workers): + """Async helper to collect tiles from workers.""" + return await self._async_collect_results(multi_job_id, num_workers, mode='static') + + async def _mark_image_completed(self, multi_job_id, image_idx, image_pil): + """Mark an image as completed in the job data.""" + # Mark the image as completed with the image data + await _mark_task_completed(multi_job_id, image_idx, {'image': image_pil}) + 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 'completed_images' in job_data: + job_data['completed_images'][image_idx] = image_pil + + async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count): + """Collect remaining processed images from workers.""" + return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic', + remaining_to_collect=remaining_to_collect, + batch_size=batch_size) diff --git a/upscale/tile_ops.py b/upscale/tile_ops.py new file mode 100644 index 0000000..50acf86 --- /dev/null +++ b/upscale/tile_ops.py @@ -0,0 +1,404 @@ +import math, torch +from PIL import Image, ImageFilter, ImageDraw +from typing import List, Tuple +import comfy.samplers, comfy.model_management +from ..utils.logging import debug_log, log +from ..utils.image import tensor_to_pil, pil_to_tensor +from ..utils.usdu_utils import crop_cond, get_crop_region, expand_crop +from ..utils.usdu_managment import clone_conditioning + + +class TileOpsMixin: + def round_to_multiple(self, value: int, multiple: int = 8) -> int: + """Round value to nearest multiple.""" + return round(value / multiple) * multiple + + def calculate_tiles(self, image_width: int, image_height: int, + tile_width: int, tile_height: int, force_uniform_tiles: bool = True) -> List[Tuple[int, int]]: + """Calculate tile positions to match Ultimate SD Upscale. + + Positions are a simple grid starting at (0,0) with steps of + `tile_width` and `tile_height`, using ceil(rows/cols) to cover edges. + Uniform vs non-uniform affects only crop/resize, not positions. + """ + rows = math.ceil(image_height / tile_height) + cols = math.ceil(image_width / tile_width) + tiles: List[Tuple[int, int]] = [] + for yi in range(rows): + for xi in range(cols): + tiles.append((xi * tile_width, yi * tile_height)) + return tiles + + def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int, + tile_width: int, tile_height: int, padding: int, + force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]: + """Extract a tile region and resize to match USDU cropping logic. + + Mirrors ComfyUI_UltimateSDUpscale processing: + - Build a mask with a white rectangle at the tile rect + - Compute crop_region via get_crop_region(mask, padding) + - If force_uniform_tiles: shrink/expand to target size of + round_to_multiple(tile + padding) for each dimension + - Else: target is ceil(crop_size/8)*8 per dimension + - Extract the crop and resize to target tile_size + Returns the resized tensor and crop origin/size for blending. + """ + _, h, w, _ = image.shape + + # Create mask and compute initial padded crop region + mask = Image.new('L', (w, h), 0) + draw = ImageDraw.Draw(mask) + draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255) + x1, y1, x2, y2 = get_crop_region(mask, padding) + + # Determine target tile size (processing size) + if force_uniform_tiles: + target_w = self.round_to_multiple(tile_width + padding, 8) + target_h = self.round_to_multiple(tile_height + padding, 8) + (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) + else: + crop_w = x2 - x1 + crop_h = y2 - y1 + target_w = max(8, math.ceil(crop_w / 8) * 8) + target_h = max(8, math.ceil(crop_h / 8) * 8) + (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) + + # Actual extracted size before resizing (for blending) + extracted_width = x2 - x1 + extracted_height = y2 - y1 + + # Extract tile and resize to processing size + tile = image[:, y1:y2, x1:x2, :] + tile_pil = tensor_to_pil(tile, 0) + if tile_pil.size != (target_w, target_h): + tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS) + + tile_tensor = pil_to_tensor(tile_pil) + if image.is_cuda: + tile_tensor = tile_tensor.cuda() + + return tile_tensor, x1, y1, extracted_width, extracted_height + + def extract_batch_tile_with_padding(self, images: torch.Tensor, x: int, y: int, + tile_width: int, tile_height: int, padding: int, + force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]: + """Extract a tile region for the entire batch and resize to USDU logic. + + - Computes a single crop region from a mask at (x,y,w,h) with padding + - force_uniform_tiles controls target processing size logic + - Returns a batched tensor [B,H',W',C] and crop origin/size for blending + """ + batch, h, w, _ = images.shape + + # Create mask and compute initial padded crop region (same for all images) + mask = Image.new('L', (w, h), 0) + draw = ImageDraw.Draw(mask) + draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255) + x1, y1, x2, y2 = get_crop_region(mask, padding) + + # Determine target processing size + if force_uniform_tiles: + target_w = self.round_to_multiple(tile_width + padding, 8) + target_h = self.round_to_multiple(tile_height + padding, 8) + (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) + else: + crop_w = x2 - x1 + crop_h = y2 - y1 + target_w = max(8, math.ceil(crop_w / 8) * 8) + target_h = max(8, math.ceil(crop_h / 8) * 8) + (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) + + extracted_width = x2 - x1 + extracted_height = y2 - y1 + + # Slice batch region + tiles = images[:, y1:y2, x1:x2, :] + + # Resize each tile to target size + resized_tiles = [] + for i in range(batch): + tile_pil = tensor_to_pil(tiles, i) + if tile_pil.size != (target_w, target_h): + tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS) + resized_tiles.append(pil_to_tensor(tile_pil)) + tile_batch = torch.cat(resized_tiles, dim=0) + + if images.is_cuda: + tile_batch = tile_batch.cuda() + + return tile_batch, 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, + 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}") + + + # Import here to avoid circular dependencies + from nodes import common_ksampler, VAEEncode, VAEDecode + + # Try to import tiled VAE nodes if available + try: + from nodes import VAEEncodeTiled, VAEDecodeTiled + tiled_vae_available = True + except ImportError: + tiled_vae_available = False + if tiled_decode: + debug_log("Tiled VAE nodes not available, falling back to standard VAE") + + # Convert to PIL and back to ensure clean tensor without gradient tracking + tile_pil = tensor_to_pil(tile_tensor, 0) + clean_tensor = pil_to_tensor(tile_pil) + + # Ensure tensor is detached and doesn't require gradients + clean_tensor = clean_tensor.detach() + if hasattr(clean_tensor, 'requires_grad_'): + clean_tensor.requires_grad_(False) + + # Move to correct device + if tile_tensor.is_cuda: + clean_tensor = clean_tensor.cuda() + clean_tensor = clean_tensor.detach() # Detach again after device transfer + + # Clone conditioning per tile (shares models, clones hints for cropping) + positive_tile = clone_conditioning(positive, clone_hints=True) + negative_tile = clone_conditioning(negative, clone_hints=True) + + # Crop conditioning to tile region if provided (assumes hints at image resolution) + if region is not None and image_size is not None: + 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_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) + negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) + else: + # No region cropping needed, use cloned conditioning as-is + positive_cropped = positive_tile + negative_cropped = negative_tile + + # Encode to latent (always non-tiled, matching original node) + latent = VAEEncode().encode(vae, clean_tensor)[0] + + # Sample + samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, + positive_cropped, negative_cropped, latent, denoise=denoise)[0] + + # Decode back to image + if tiled_decode and tiled_vae_available: + image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0] + else: + image = VAEDecode().decode(vae, samples)[0] + + return image + + def process_tiles_batch(self, tile_batch: torch.Tensor, model, positive, negative, vae, + seed: int, steps: int, cfg: float, sampler_name: str, + scheduler: str, denoise: float, tiled_decode: bool, + region: Tuple[int, int, int, int], image_size: Tuple[int, int]) -> torch.Tensor: + """Process a batch of tiles together (USDU behavior). + + tile_batch: [B, H, W, C] + Returns image batch tensor [B, H, W, C] + """ + # Import locally to avoid circular deps + from nodes import common_ksampler, VAEEncode, VAEDecode + try: + from nodes import VAEEncodeTiled, VAEDecodeTiled + tiled_vae_available = True + except ImportError: + tiled_vae_available = False + + # Detach and move device + clean = tile_batch.detach() + if hasattr(clean, 'requires_grad_'): + clean.requires_grad_(False) + if tile_batch.is_cuda: + clean = clean.cuda().detach() + + # Clone/crop conditioning once for the region + positive_tile = clone_conditioning(positive, clone_hints=True) + negative_tile = clone_conditioning(negative, clone_hints=True) + + init_size = image_size + canvas_size = image_size + tile_size = (clean.shape[2], clean.shape[1]) # (W,H) + w_pad = 0 + h_pad = 0 + positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) + negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) + + # Encode -> Sample -> Decode + latent = VAEEncode().encode(vae, clean)[0] + samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, + positive_cropped, negative_cropped, latent, denoise=denoise)[0] + if tiled_decode and tiled_vae_available: + image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0] + else: + image = VAEDecode().decode(vae, samples)[0] + + return image + + def create_tile_mask(self, image_width: int, image_height: int, + x: int, y: int, tile_width: int, tile_height: int, + mask_blur: int) -> Image.Image: + """Create a mask for blending tiles - matches Ultimate SD Upscale approach. + + Creates a black image with a white rectangle at the tile position, + then applies blur to create soft edges. + """ + # Create a full-size mask matching the image dimensions + mask = Image.new('L', (image_width, image_height), 0) # Black background + + # Draw white rectangle at tile position + draw = ImageDraw.Draw(mask) + draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255) + + # Apply blur to soften edges + if mask_blur > 0: + mask = mask.filter(ImageFilter.GaussianBlur(mask_blur)) + + return mask + + def blend_tile(self, base_image: Image.Image, tile_image: Image.Image, + x: int, y: int, extracted_size: Tuple[int, int], + mask: Image.Image, padding: int) -> Image.Image: + """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 + """ + extracted_width, extracted_height = extracted_size + + # Debug logging (uncomment if needed) + # debug_log(f"[Blend] Placing tile at ({x}, {y}), size: {extracted_width}x{extracted_height}") + + # Calculate the crop region that was used for extraction + crop_region = (x, y, x + extracted_width, y + extracted_height) + + # The mask is already full-size, no need to crop + + # Resize the processed tile back to the extracted size + if tile_image.size != (extracted_width, extracted_height): + tile_resized = tile_image.resize((extracted_width, extracted_height), Image.LANCZOS) + else: + tile_resized = tile_image + + # Follow Ultimate SD Upscale blending approach: + # Put the tile into position + image_tile_only = Image.new('RGBA', base_image.size) + image_tile_only.paste(tile_resized, crop_region[:2]) + + # Add the mask as an alpha channel + # Must make a copy due to the possibility of an edge becoming black + temp = image_tile_only.copy() + temp.putalpha(mask) # Use the full image mask + image_tile_only.paste(temp, image_tile_only) + + # Add back the tile to the initial image according to the mask in the alpha channel + result = base_image.convert('RGBA') + result.alpha_composite(image_tile_only) + + # Convert back to RGB + return result.convert('RGB') + + def _slice_conditioning(self, positive, negative, batch_idx): + """Helper to slice conditioning for a specific batch index.""" + # Clone and slice conditioning properly, including ControlNet hints + 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[batch_idx:batch_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[batch_idx:batch_idx+1] + control = control.previous_controlnet + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + + return positive_sliced, negative_sliced + + 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, force_uniform_tiles, + tiled_decode, batch_idx: int = 0): + """Process a single tile and blend it into the result image.""" + x, y = tile_pos + + # Extract and process tile + tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( + upscaled_image, x, y, tile_width, tile_height, padding, force_uniform_tiles + ) + + processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, + seed, steps, cfg, sampler_name, + 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) + # 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, + x1, y1, (ew, eh), tile_mask, padding) + + return result_image + + def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive, negative, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, force_uniform_tiles, sliced_conditioning_cache): + """Process a single tile.""" + # Calculate which image and tile this corresponds to + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image + + # Skip if batch_idx is out of range + if batch_idx >= upscaled_image.shape[0]: + debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}") + return None + + # Get or create sliced conditioning for this batch index + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + 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] + + # Extract tile from the specific image in the batch + tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( + upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding, force_uniform_tiles + ) + + # Process tile through SD with unique seed + image_seed = seed + batch_idx * 1000 + processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae, + image_seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, batch_idx=batch_idx, + region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height)) + + return { + 'tile': processed_tile, + 'global_idx': global_idx, + 'batch_idx': batch_idx, + 'tile_idx': tile_idx, + 'x': x1, + 'y': y1, + 'extracted_width': ew, + 'extracted_height': eh + } diff --git a/upscale/worker_comms.py b/upscale/worker_comms.py new file mode 100644 index 0000000..430f2a2 --- /dev/null +++ b/upscale/worker_comms.py @@ -0,0 +1,271 @@ +import asyncio, io, json, time +import aiohttp +from PIL import Image +from ..utils.logging import debug_log, log +from ..utils.network import get_client_session +from ..utils.constants import TILE_SEND_TIMEOUT +from ..utils.usdu_managment import MAX_PAYLOAD_SIZE, _send_heartbeat_to_master +from ..utils.image import tensor_to_pil + + +class WorkerCommsMixin: + async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, + padding, worker_id): + """Send all processed tiles to master, chunked if large.""" + if not processed_tiles: + return # Early exit if empty + + total_tiles = len(processed_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 = [] + for idx, tile_data in enumerate(processed_tiles): + img = tensor_to_pil(tile_data['tile'], 0) + bio = io.BytesIO() + # Keep compression low to balance speed and size; adjust if needed + img.save(bio, format='PNG', compress_level=0) + raw = bio.getvalue() + encoded.append({ + 'bytes': raw, + 'meta': { + 'tile_idx': tile_data['tile_idx'], + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'], + **({'batch_idx': tile_data['batch_idx']} if 'batch_idx' in tile_data else {}), + **({'global_idx': tile_data['global_idx']} if 'global_idx' in tile_data else {}), + } + }) + + # Size-aware chunking + max_bytes = int(MAX_PAYLOAD_SIZE) - (1024 * 1024) # 1MB headroom + i = 0 + chunk_index = 0 + while i < total_tiles: + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('padding', str(padding)) + + metadata = [] + used = 0 + j = i + while j < total_tiles: + img_bytes = encoded[j]['bytes'] + meta = encoded[j]['meta'] + # Rough overhead for fields + JSON + overhead = 1024 + if used + len(img_bytes) + overhead > max_bytes and j > i: + break + # Accept this tile in this chunk + metadata.append(meta) + data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png') + used += len(img_bytes) + overhead + j += 1 + + # Ensure at least one tile per chunk + if j == i: + # Single oversized tile, send anyway + meta = encoded[j]['meta'] + metadata.append(meta) + data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png') + j += 1 + + chunk_size = j - i + is_chunk_last = (j >= total_tiles) + data.add_field('is_last', str(is_chunk_last)) + data.add_field('batch_size', str(chunk_size)) + data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json') + + # Retry logic with exponential backoff + max_retries = 5 + retry_delay = 0.5 + for attempt in range(max_retries): + try: + session = await get_client_session() + url = f"{master_url}/distributed/submit_tiles" + async with session.post(url, data=data) as response: + response.raise_for_status() + break + except Exception as e: + if attempt < max_retries - 1: + await asyncio.sleep(retry_delay) + retry_delay = min(retry_delay * 2, 5.0) + else: + log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}") + raise + + debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)") + chunk_index += 1 + i = j + + async def _request_image_from_master(self, multi_job_id, master_url, worker_id): + """Request an image index to process from master in dynamic mode.""" + # Enhanced retries with 404-specific delay, backoff cap, and total timeout to handle init races + max_retries = 10 + retry_delay = 0.5 + start_time = time.time() + + for attempt in range(max_retries): + # Check total timeout + if time.time() - start_time > 30: + log(f"Total request timeout after 30s for worker {worker_id}") + return None, 0 + + try: + session = await get_client_session() + url = f"{master_url}/distributed/request_image" + + async with session.post(url, json={ + 'worker_id': str(worker_id), + 'multi_job_id': multi_job_id + }) as response: + if response.status == 200: + data = await response.json() + image_idx = data.get('image_idx') + estimated_remaining = data.get('estimated_remaining', 0) + return image_idx, estimated_remaining + elif response.status == 404: + # Special handling for 404 - job not found yet + text = await response.text() + debug_log(f"Job not found (404), will retry: {text}") + await asyncio.sleep(1.0) + else: + text = await response.text() + debug_log(f"Request image failed: {response.status} - {text}") + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 + retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s + else: + log(f"Failed to request image after {max_retries} attempts: {e}") + raise + + return None, 0 + + async def _request_tile_from_master(self, multi_job_id, master_url, worker_id): + """Request a tile index to process from master in static mode (reusing dynamic infrastructure).""" + # Reuse the same retry logic as dynamic mode + max_retries = 10 + retry_delay = 0.5 + start_time = time.time() + + for attempt in range(max_retries): + # Check total timeout + if time.time() - start_time > 30: + log(f"Total request timeout after 30s for worker {worker_id}") + return None, 0 + + try: + session = await get_client_session() + url = f"{master_url}/distributed/request_image" # Same endpoint + + async with session.post(url, json={ + 'worker_id': str(worker_id), + 'multi_job_id': multi_job_id + }) as response: + if response.status == 200: + data = await response.json() + tile_idx = data.get('tile_idx') + estimated_remaining = data.get('estimated_remaining', 0) + batched_static = data.get('batched_static', False) + return tile_idx, estimated_remaining, batched_static + elif response.status == 404: + # Special handling for 404 - job not found yet + text = await response.text() + debug_log(f"Job not found (404), will retry: {text}") + await asyncio.sleep(1.0) + else: + text = await response.text() + debug_log(f"Request tile failed: {response.status} - {text}") + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 + retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s + else: + log(f"Failed to request tile after {max_retries} attempts: {e}") + raise + + return None, 0, False + + async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, + master_url, worker_id, is_last): + """Send a processed full image back to master in dynamic mode.""" + # Serialize image to PNG + byte_io = io.BytesIO() + image_pil.save(byte_io, format='PNG', compress_level=0) + byte_io.seek(0) + + # Prepare form data + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('image_idx', str(image_idx)) + data.add_field('is_last', str(is_last)) + data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png', + content_type='image/png') + + # Retry logic + max_retries = 5 + retry_delay = 0.5 + + for attempt in range(max_retries): + try: + session = await get_client_session() + url = f"{master_url}/distributed/submit_image" + + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"Successfully sent image {image_idx} to master") + return + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 + else: + log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}") + raise + + async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id): + """Send completion signal to master in dynamic mode.""" + # Send a dummy request with is_last=True + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', 'true') + # No image data - just completion signal + + session = await get_client_session() + url = f"{master_url}/distributed/submit_image" + + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"Worker {worker_id} sent completion signal") + + async def _check_job_status(self, multi_job_id, master_url): + """Check if job is ready on the master.""" + try: + session = await get_client_session() + url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}" + async with session.get(url) as response: + if response.status == 200: + data = await response.json() + return data.get('ready', False) + return False + except Exception as e: + debug_log(f"Job status check failed: {e}") + return False + + async def _async_yield(self): + """Simple async yield to allow event loop processing.""" + await asyncio.sleep(0)