866 lines
40 KiB
Python
866 lines
40 KiB
Python
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) |