import numpy as np from PIL import Image, ImageFilter import torch import torch.nn.functional as F from torchvision.transforms import GaussianBlur import math import copy import asyncio import json import io import time import os import server from aiohttp import web # Import other utilities from the same package from .logging import debug_log, log from .network import handle_api_error if (not hasattr(Image, 'Resampling')): # For older versions of Pillow Image.Resampling = Image BLUR_KERNEL_SIZE = 15 # Configure maximum payload size (50MB default, configurable via environment variable) MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) def tensor_to_pil(img_tensor, batch_index=0): # Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels] # and returns an RGB PIL Image. Assumes channels=3 return Image.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8)) def pil_to_tensor(image): # Takes a PIL image and returns a tensor of shape [1, height, width, channels] image = np.array(image).astype(np.float32) / 255.0 image = torch.from_numpy(image).unsqueeze(0) if len(image.shape) == 3: # If the image is grayscale, add a channel dimension image = image.unsqueeze(-1) return image def controlnet_hint_to_pil(tensor, batch_index=0): return tensor_to_pil(tensor.movedim(1, -1), batch_index) def pil_to_controlnet_hint(img): return pil_to_tensor(img).movedim(-1, 1) def crop_tensor(tensor, region): # Takes a tensor of shape [batch_size, height, width, channels] and crops it to the given region x1, y1, x2, y2 = region return tensor[:, y1:y2, x1:x2, :] def resize_tensor(tensor, size, mode="nearest-exact"): # Takes a tensor of shape [B, C, H, W] and resizes # it to a shape of [B, C, size[0], size[1]] using the given mode return torch.nn.functional.interpolate(tensor, size=size, mode=mode) def get_crop_region(mask, pad=0): # Takes a black and white PIL image in 'L' mode and returns the coordinates of the white rectangular mask region # Should be equivalent to the get_crop_region function from https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/masking.py coordinates = mask.getbbox() if coordinates is not None: x1, y1, x2, y2 = coordinates else: x1, y1, x2, y2 = mask.width, mask.height, 0, 0 # Apply padding x1 = max(x1 - pad, 0) y1 = max(y1 - pad, 0) x2 = min(x2 + pad, mask.width) y2 = min(y2 + pad, mask.height) return fix_crop_region((x1, y1, x2, y2), (mask.width, mask.height)) def fix_crop_region(region, image_size): # Remove the extra pixel added by the get_crop_region function image_width, image_height = image_size x1, y1, x2, y2 = region if x2 < image_width: x2 -= 1 if y2 < image_height: y2 -= 1 return x1, y1, x2, y2 def expand_crop(region, width, height, target_width, target_height): ''' Expands a crop region to a specified target size. :param region: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points of the rectangular region. Expected to have x2 > x1 and y2 > y1. :param width: The width of the image the crop region is from. :param height: The height of the image the crop region is from. :param target_width: The desired width of the crop region. :param target_height: The desired height of the crop region. ''' x1, y1, x2, y2 = region actual_width = x2 - x1 actual_height = y2 - y1 # target_width = math.ceil(actual_width / 8) * 8 # target_height = math.ceil(actual_height / 8) * 8 # Try to expand region to the right of half the difference width_diff = target_width - actual_width x2 = min(x2 + width_diff // 2, width) # Expand region to the left of the difference including the pixels that could not be expanded to the right width_diff = target_width - (x2 - x1) x1 = max(x1 - width_diff, 0) # Try the right again width_diff = target_width - (x2 - x1) x2 = min(x2 + width_diff, width) # Try to expand region to the bottom of half the difference height_diff = target_height - actual_height y2 = min(y2 + height_diff // 2, height) # Expand region to the top of the difference including the pixels that could not be expanded to the bottom height_diff = target_height - (y2 - y1) y1 = max(y1 - height_diff, 0) # Try the bottom again height_diff = target_height - (y2 - y1) y2 = min(y2 + height_diff, height) return (x1, y1, x2, y2), (target_width, target_height) def resize_region(region, init_size, resize_size): # Resize a crop so that it fits an image that was resized to the given width and height x1, y1, x2, y2 = region init_width, init_height = init_size resize_width, resize_height = resize_size x1 = math.floor(x1 * resize_width / init_width) x2 = math.ceil(x2 * resize_width / init_width) y1 = math.floor(y1 * resize_height / init_height) y2 = math.ceil(y2 * resize_height / init_height) return (x1, y1, x2, y2) def pad_image(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False): ''' Pads an image with the given number of pixels on each side and fills the padding with data from the edges. :param image: A PIL image :param left_pad: The number of pixels to pad on the left side :param right_pad: The number of pixels to pad on the right side :param top_pad: The number of pixels to pad on the top side :param bottom_pad: The number of pixels to pad on the bottom side :param blur: Whether to blur the padded edges :return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad) ''' left_edge = image.crop((0, 1, 1, image.height - 1)) right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1)) top_edge = image.crop((1, 0, image.width - 1, 1)) bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height)) new_width = image.width + left_pad + right_pad new_height = image.height + top_pad + bottom_pad padded_image = Image.new(image.mode, (new_width, new_height)) padded_image.paste(image, (left_pad, top_pad)) if fill: for i in range(left_pad): edge = left_edge.resize( (1, new_height - i * (top_pad + bottom_pad) // left_pad), resample=Image.Resampling.NEAREST) padded_image.paste(edge, (i, i * top_pad // left_pad)) for i in range(right_pad): edge = right_edge.resize( (1, new_height - i * (top_pad + bottom_pad) // right_pad), resample=Image.Resampling.NEAREST) padded_image.paste(edge, (new_width - 1 - i, i * top_pad // right_pad)) for i in range(top_pad): edge = top_edge.resize( (new_width - i * (left_pad + right_pad) // top_pad, 1), resample=Image.Resampling.NEAREST) padded_image.paste(edge, (i * left_pad // top_pad, i)) for i in range(bottom_pad): edge = bottom_edge.resize( (new_width - i * (left_pad + right_pad) // bottom_pad, 1), resample=Image.Resampling.NEAREST) padded_image.paste(edge, (i * left_pad // bottom_pad, new_height - 1 - i)) if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0): padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE)) padded_image.paste(image, (left_pad, top_pad)) return padded_image def pad_image2(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False): ''' Pads an image with the given number of pixels on each side and fills the padding with data from the edges. Faster than pad_image, but only pads with edge data in straight lines. :param image: A PIL image :param left_pad: The number of pixels to pad on the left side :param right_pad: The number of pixels to pad on the right side :param top_pad: The number of pixels to pad on the top side :param bottom_pad: The number of pixels to pad on the bottom side :param blur: Whether to blur the padded edges :return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad) ''' left_edge = image.crop((0, 1, 1, image.height - 1)) right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1)) top_edge = image.crop((1, 0, image.width - 1, 1)) bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height)) new_width = image.width + left_pad + right_pad new_height = image.height + top_pad + bottom_pad padded_image = Image.new(image.mode, (new_width, new_height)) padded_image.paste(image, (left_pad, top_pad)) if fill: if left_pad > 0: padded_image.paste(left_edge.resize((left_pad, new_height), resample=Image.Resampling.NEAREST), (0, 0)) if right_pad > 0: padded_image.paste(right_edge.resize((right_pad, new_height), resample=Image.Resampling.NEAREST), (new_width - right_pad, 0)) if top_pad > 0: padded_image.paste(top_edge.resize((new_width, top_pad), resample=Image.Resampling.NEAREST), (0, 0)) if bottom_pad > 0: padded_image.paste(bottom_edge.resize((new_width, bottom_pad), resample=Image.Resampling.NEAREST), (0, new_height - bottom_pad)) if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0): padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE)) padded_image.paste(image, (left_pad, top_pad)) return padded_image def pad_tensor(tensor, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False): ''' Pads an image tensor with the given number of pixels on each side and fills the padding with data from the edges. :param tensor: A tensor of shape [B, H, W, C] :param left_pad: The number of pixels to pad on the left side :param right_pad: The number of pixels to pad on the right side :param top_pad: The number of pixels to pad on the top side :param bottom_pad: The number of pixels to pad on the bottom side :param blur: Whether to blur the padded edges :return: A tensor of shape [B, H + top_pad + bottom_pad, W + left_pad + right_pad, C] ''' batch_size, channels, height, width = tensor.shape h_pad = left_pad + right_pad v_pad = top_pad + bottom_pad new_width = width + h_pad new_height = height + v_pad # Create empty image padded = torch.zeros((batch_size, channels, new_height, new_width), dtype=tensor.dtype) # Copy the original image into the centor of the padded tensor padded[:, :, top_pad:top_pad + height, left_pad:left_pad + width] = tensor # Duplicate the edges of the original image into the padding if top_pad > 0: padded[:, :, :top_pad, :] = padded[:, :, top_pad:top_pad + 1, :] # Top edge if bottom_pad > 0: padded[:, :, -bottom_pad:, :] = padded[:, :, -bottom_pad - 1:-bottom_pad, :] # Bottom edge if left_pad > 0: padded[:, :, :, :left_pad] = padded[:, :, :, left_pad:left_pad + 1] # Left edge if right_pad > 0: padded[:, :, :, -right_pad:] = padded[:, :, :, -right_pad - 1:-right_pad] # Right edge return padded def resize_and_pad_image(image, width, height, fill=False, blur=False): ''' Resizes an image to the given width and height and pads it to the given width and height. :param image: A PIL image :param width: The width of the resized image :param height: The height of the resized image :param fill: Whether to fill the padding with data from the edges :param blur: Whether to blur the padded edges :return: A PIL image of size (width, height) ''' width_ratio = width / image.width height_ratio = height / image.height if height_ratio > width_ratio: resize_ratio = width_ratio else: resize_ratio = height_ratio resize_width = round(image.width * resize_ratio) resize_height = round(image.height * resize_ratio) resized = image.resize((resize_width, resize_height), resample=Image.Resampling.LANCZOS) # Pad the sides of the image to get the image to the desired size that wasn't covered by the resize horizontal_pad = (width - resize_width) // 2 vertical_pad = (height - resize_height) // 2 result = pad_image2(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur) result = result.resize((width, height), resample=Image.Resampling.LANCZOS) return result, (horizontal_pad, vertical_pad) def resize_and_pad_tensor(tensor, width, height, fill=False, blur=False): ''' Resizes an image tensor to the given width and height and pads it to the given width and height. :param tensor: A tensor of shape [B, H, W, C] :param width: The width of the resized image :param height: The height of the resized image :param fill: Whether to fill the padding with data from the edges :param blur: Whether to blur the padded edges :return: A tensor of shape [B, height, width, C] ''' # Resize the image to the closest size that maintains the aspect ratio width_ratio = width / tensor.shape[3] height_ratio = height / tensor.shape[2] if height_ratio > width_ratio: resize_ratio = width_ratio else: resize_ratio = height_ratio resize_width = round(tensor.shape[3] * resize_ratio) resize_height = round(tensor.shape[2] * resize_ratio) resized = F.interpolate(tensor, size=(resize_height, resize_width), mode='nearest-exact') # Pad the sides of the image to get the image to the desired size that wasn't covered by the resize horizontal_pad = (width - resize_width) // 2 vertical_pad = (height - resize_height) // 2 result = pad_tensor(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur) result = F.interpolate(result, size=(height, width), mode='nearest-exact') return result def crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): if "control" not in cond_dict: return c = cond_dict["control"] controlnet = c.copy() cond_dict["control"] = controlnet while c is not None: # hint is shape (B, C, H, W) hint = controlnet.cond_hint_original resized_crop = resize_region(region, canvas_size, hint.shape[:-3:-1]) hint = crop_tensor(hint.movedim(1, -1), resized_crop).movedim(-1, 1) hint = resize_tensor(hint, tile_size[::-1]) controlnet.cond_hint_original = hint c = c.previous_controlnet controlnet.set_previous_controlnet(c.copy() if c is not None else None) controlnet = controlnet.previous_controlnet def region_intersection(region1, region2): """ Returns the coordinates of the intersection of two rectangular regions. :param region1: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points of the first rectangular region. Expected to have x2 > x1 and y2 > y1. :param region2: The second rectangular region with the same format as the first. :return: A tuple of the form (x1, y1, x2, y2) denoting the rectangular intersection. None if there is no intersection. """ x1, y1, x2, y2 = region1 x1_, y1_, x2_, y2_ = region2 x1 = max(x1, x1_) y1 = max(y1, y1_) x2 = min(x2, x2_) y2 = min(y2, y2_) if x1 >= x2 or y1 >= y2: return None return (x1, y1, x2, y2) def crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): if "gligen" not in cond_dict: return type, model, cond = cond_dict["gligen"] if type != "position": from warnings import warn warn(f"Unknown gligen type {type}") return cropped = [] for c in cond: emb, h, w, y, x = c # Get the coordinates of the box in the upscaled image x1 = x * 8 y1 = y * 8 x2 = x1 + w * 8 y2 = y1 + h * 8 gligen_upscaled_box = resize_region((x1, y1, x2, y2), init_size, canvas_size) # Calculate the intersection of the gligen box and the region intersection = region_intersection(gligen_upscaled_box, region) if intersection is None: continue x1, y1, x2, y2 = intersection # Offset the gligen box so that the origin is at the top left of the tile region x1 -= region[0] y1 -= region[1] x2 -= region[0] y2 -= region[1] # Add the padding x1 += w_pad y1 += h_pad x2 += w_pad y2 += h_pad # Set the new position params h = (y2 - y1) // 8 w = (x2 - x1) // 8 x = x1 // 8 y = y1 // 8 cropped.append((emb, h, w, y, x)) cond_dict["gligen"] = (type, model, cropped) def crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): if "area" not in cond_dict: return # Resize the area conditioning to the canvas size and confine it to the tile region h, w, y, x = cond_dict["area"] w, h, x, y = 8 * w, 8 * h, 8 * x, 8 * y x1, y1, x2, y2 = resize_region((x, y, x + w, y + h), init_size, canvas_size) intersection = region_intersection((x1, y1, x2, y2), region) if intersection is None: del cond_dict["area"] del cond_dict["strength"] return x1, y1, x2, y2 = intersection # Offset origin to the top left of the tile x1 -= region[0] y1 -= region[1] x2 -= region[0] y2 -= region[1] # Add the padding x1 += w_pad y1 += h_pad x2 += w_pad y2 += h_pad # Set the params for tile w, h = (x2 - x1) // 8, (y2 - y1) // 8 x, y = x1 // 8, y1 // 8 cond_dict["area"] = (h, w, y, x) def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): if "mask" not in cond_dict: return mask_tensor = cond_dict["mask"] # (B, H, W) masks = [] for i in range(mask_tensor.shape[0]): # Convert to PIL image mask = tensor_to_pil(mask_tensor, i) # W x H # Resize the mask to the canvas size mask = mask.resize(canvas_size, Image.Resampling.BICUBIC) # Crop the mask to the region mask = mask.crop(region) # Add padding mask, _ = resize_and_pad_image(mask, tile_size[0], tile_size[1], fill=True) # Resize the mask to the tile size if tile_size != mask.size: mask = mask.resize(tile_size, Image.Resampling.BICUBIC) # Convert back to tensor mask = pil_to_tensor(mask) # (1, H, W, 1) mask = mask.squeeze(-1) # (1, H, W) masks.append(mask) cond_dict["mask"] = torch.cat(masks, dim=0) # (B, H, W) def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0): cropped = [] for emb, x in cond: cond_dict = x.copy() n = [emb, cond_dict] crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) cropped.append(n) return cropped # Helper functions for shallow copying conditioning without duplicating models def clone_control_chain(control, clone_hint=True): """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" if control is None: return None new_control = copy.copy(control) # Shallow copy (shares model) if clone_hint and hasattr(control, 'cond_hint_original'): new_control.cond_hint_original = control.cond_hint_original.clone() if hasattr(control, 'previous_controlnet'): new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint) return new_control def clone_conditioning(cond_list, clone_hints=True): """Clone conditioning without duplicating ControlNet models.""" new_cond = [] for emb, cond_dict in cond_list: new_emb = emb.clone() if emb is not None else None new_dict = cond_dict.copy() if 'control' in new_dict: new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints) if 'mask' in new_dict: new_dict['mask'] = new_dict['mask'].clone() # Handle other potential fields if needed if 'pooled_output' in new_dict: new_dict['pooled_output'] = new_dict['pooled_output'].clone() if 'area' in new_dict: new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple new_cond.append([new_emb, new_dict]) return new_cond # Helper function to ensure persistent state is initialized def ensure_tile_jobs_initialized(): """Ensure tile job storage is initialized on the server instance.""" prompt_server = server.PromptServer.instance if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): debug_log("Initializing persistent tile job queue on server instance.") prompt_server.distributed_pending_tile_jobs = {} prompt_server.distributed_tile_jobs_lock = asyncio.Lock() else: # Clean up any legacy queue structures that don't have the 'mode' field # (Should be rare after fixes, but keep for safety) to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() if not isinstance(job_data, dict) or 'mode' not in job_data] for job_id in to_remove: debug_log(f"Removing legacy queue structure for job {job_id}") del prompt_server.distributed_pending_tile_jobs[job_id] return prompt_server # API Endpoint for tile completion @server.PromptServer.instance.routes.post("/distributed/request_image") async def request_image_endpoint(request): """Endpoint for workers to request images in dynamic mode.""" try: data = await request.json() worker_id = data.get('worker_id') multi_job_id = data.get('multi_job_id') if not worker_id or not multi_job_id: return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] if not isinstance(job_data, dict) or 'mode' not in job_data: return await handle_api_error(request, "Invalid job data structure", 500) if job_data['mode'] != 'dynamic': return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) if 'pending_images' in job_data: pending_queue = job_data['pending_images'] try: image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) # Track assigned image if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: job_data['assigned_to_workers'][worker_id].append(image_idx) # Update worker heartbeat if 'worker_status' in job_data: job_data['worker_status'][worker_id] = time.time() # Get estimated remaining count remaining = pending_queue.qsize() # Approximate debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}") return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining}) except asyncio.TimeoutError: return web.json_response({"image_idx": None}) else: return await handle_api_error(request, "Invalid dynamic mode configuration", 400) else: return await handle_api_error(request, "Job not found", 404) except Exception as e: return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.get("/distributed/job_status") async def job_status_endpoint(request): """Endpoint to check if a job is ready.""" multi_job_id = request.query.get('multi_job_id') if not multi_job_id: return web.json_response({"ready": False}) 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) ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data) return web.json_response({"ready": ready}) @server.PromptServer.instance.routes.post("/distributed/heartbeat") async def heartbeat_endpoint(request): """Endpoint for workers to send heartbeat signals.""" try: data = await request.json() worker_id = data.get('worker_id') multi_job_id = data.get('multi_job_id') if not worker_id or not multi_job_id: return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] if 'worker_status' in job_data: job_data['worker_status'][worker_id] = time.time() debug_log(f"UltimateSDUpscale API - Heartbeat from worker {worker_id}") return web.json_response({"status": "success"}) else: return await handle_api_error(request, "Worker status tracking not available", 400) else: return await handle_api_error(request, "Job not found", 404) except Exception as e: return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.post("/distributed/tile_complete") async def tile_complete_endpoint(request): """Endpoint for receiving completed tiles from workers.""" try: # Check payload size before processing content_length = request.headers.get('content-length') if content_length and int(content_length) > MAX_PAYLOAD_SIZE: return await handle_api_error(request, f"Payload too large: {content_length} bytes (max: {MAX_PAYLOAD_SIZE})", 413) data = await request.post() multi_job_id = data.get('multi_job_id') worker_id = data.get('worker_id') is_last = data.get('is_last', 'False').lower() == 'true' if multi_job_id is None or worker_id is None: return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) prompt_server = ensure_tile_jobs_initialized() # Check if this is a full image submission (dynamic mode) if 'full_image' in data and 'image_idx' in data: image_idx = int(data.get('image_idx')) is_last = data.get('is_last', 'False').lower() == 'true' try: # Process full image img_data = data['full_image'].file.read() img = Image.open(io.BytesIO(img_data)).convert("RGB") debug_log(f"UltimateSDUpscale API - Received full image {image_idx} from worker {worker_id}") # Put into dynamic mode queue async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] if not isinstance(job_data, dict) or 'mode' not in job_data: return await handle_api_error(request, "Invalid job data structure", 500) if job_data['mode'] != 'dynamic': return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) # Check for duplicate image submission if 'completed_images' in job_data and image_idx in job_data['completed_images']: log(f"UltimateSDUpscale API - Duplicate image {image_idx} from {worker_id}, ignoring") return web.json_response({"status": "duplicate"}) if 'queue' in job_data: await job_data['queue'].put({ 'worker_id': worker_id, 'image_idx': image_idx, 'image': img, 'is_last': is_last }) return web.json_response({"status": "success"}) else: return await handle_api_error(request, "Invalid dynamic mode configuration", 400) else: return await handle_api_error(request, "Job not found", 404) except Exception as e: log(f"Error processing full image from worker {worker_id}: {e}") return await handle_api_error(request, f"Image processing error: {e}", 400) # Check for batch mode batch_size = int(data.get('batch_size', 0)) tiles = [] # Handle completion signals (batch_size=0 with is_last=True) if batch_size == 0 and is_last: async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] if isinstance(job_data, dict) and 'queue' in job_data: await job_data['queue'].put({ 'worker_id': worker_id, 'is_last': True, 'tiles': [] }) debug_log(f"Received completion signal from worker {worker_id}") return web.json_response({"status": "success"}) else: return await handle_api_error(request, "Invalid job structure", 500) else: return await handle_api_error(request, "Job not found", 404) if batch_size > 0: # Batch mode: Extract multiple tiles padding = int(data.get('padding', 32)) debug_log(f"UltimateSDUpscale - tile_complete batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}") # Check for JSON metadata metadata_field = data.get('tiles_metadata') if metadata_field: # New JSON metadata format try: # Handle different types of metadata field if hasattr(metadata_field, 'file'): # File-like object metadata_str = metadata_field.file.read().decode('utf-8') elif isinstance(metadata_field, (bytes, bytearray)): # Direct bytes/bytearray metadata_str = metadata_field.decode('utf-8') else: # String metadata_str = str(metadata_field) metadata = json.loads(metadata_str) if len(metadata) != batch_size: return await handle_api_error(request, "Metadata length mismatch", 400) tile_data_list = [] # Temporary list to collect and sort tiles for i in range(batch_size): tile_field = data.get(f'tile_{i}') if tile_field is None: log(f"Missing tile_{i} from worker {worker_id}, skipping") continue try: # Process image img_data = tile_field.file.read() img = Image.open(io.BytesIO(img_data)).convert("RGB") img_np = np.array(img).astype(np.float32) / 255.0 tensor = torch.from_numpy(img_np)[None,] # Get metadata for this tile if i < len(metadata): tile_data = metadata[i] tile_idx = tile_data.get('tile_idx', i) tile_info = { 'tensor': tensor, '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': padding, 'batch_idx': tile_data.get('batch_idx', 0), 'global_idx': tile_data.get('global_idx', tile_idx) } tile_data_list.append(tile_info) else: log(f"Missing metadata for tile {i} from worker {worker_id}, skipping") continue except Exception as e: log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") continue # Sort tiles by tile_idx to ensure correct order tile_data_list.sort(key=lambda x: x['tile_idx']) tiles.extend(tile_data_list) if tile_data_list: debug_log(f"Reordered {len(tile_data_list)} tiles based on tile_idx") except Exception as e: log(f"Error processing JSON metadata from worker {worker_id}: {e}") return await handle_api_error(request, f"Metadata processing error: {e}", 400) else: # Legacy format: individual fields per tile (backward compatibility) debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.") for i in range(batch_size): tile_field = data.get(f'tile_{i}') if tile_field is None: log(f"Missing tile_{i} from worker {worker_id}, skipping") continue # Get metadata for this tile tile_idx = int(data.get(f'tile_{i}_idx', i)) x = int(data.get(f'tile_{i}_x', 0)) y = int(data.get(f'tile_{i}_y', 0)) extracted_width = int(data.get(f'tile_{i}_width', 512)) extracted_height = int(data.get(f'tile_{i}_height', 512)) try: # Process image img_data = tile_field.file.read() img = Image.open(io.BytesIO(img_data)).convert("RGB") img_np = np.array(img).astype(np.float32) / 255.0 tensor = torch.from_numpy(img_np)[None,] tiles.append({ 'tensor': tensor, 'tile_idx': tile_idx, 'x': x, 'y': y, 'extracted_width': extracted_width, 'extracted_height': extracted_height, 'padding': padding }) except Exception as e: log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") continue else: # Single tile mode (backward compatibility) image_file = data.get('image') if not image_file: return await handle_api_error(request, "Missing image data", 400) tile_idx = int(data.get('tile_idx', 0)) x = int(data.get('x', 0)) y = int(data.get('y', 0)) extracted_width = int(data.get('extracted_width', 512)) extracted_height = int(data.get('extracted_height', 512)) padding = int(data.get('padding', 32)) debug_log(f"UltimateSDUpscale - tile_complete single - job_id: {multi_job_id}, worker: {worker_id}, tile: {tile_idx}") try: # Process image img_data = image_file.file.read() img = Image.open(io.BytesIO(img_data)).convert("RGB") img_np = np.array(img).astype(np.float32) / 255.0 tensor = torch.from_numpy(img_np)[None,] tiles = [{ 'tensor': tensor, 'tile_idx': tile_idx, 'x': x, 'y': y, 'extracted_width': extracted_width, 'extracted_height': extracted_height, 'padding': padding, 'batch_idx': 0, # Default for legacy single-tile 'global_idx': tile_idx # Default for legacy single-tile }] except Exception as e: log(f"Error processing tile from worker {worker_id}: {e}") return await handle_api_error(request, f"Tile processing error: {e}", 400) # Put tiles into queue async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] if not isinstance(job_data, dict) or 'mode' not in job_data: return await handle_api_error(request, "Invalid job data structure", 500) if job_data['mode'] != 'static': return await handle_api_error(request, f"Mode mismatch: expected static, got {job_data['mode']}", 400) q = job_data['queue'] if batch_size > 0: # Put batch as single item await q.put({ 'worker_id': worker_id, 'tiles': tiles, 'is_last': is_last }) else: # Put single tile (backward compat) tile_data = tiles[0] await q.put({ 'tensor': tile_data['tensor'], 'worker_id': worker_id, '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'], 'padding': tile_data['padding'], 'is_last': is_last }) return web.json_response({"status": "success"}) else: return await handle_api_error(request, "Job not found or already complete", 404) except Exception as e: return await handle_api_error(request, e, 500)