diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..68bc17f --- /dev/null +++ b/.gitignore @@ -0,0 +1,160 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ diff --git a/HDNodes.py b/HDNodes.py new file mode 100644 index 0000000..8b84f36 --- /dev/null +++ b/HDNodes.py @@ -0,0 +1,536 @@ +import os +import torch +import numpy as np +import math +import cv2 +import json +import copy +from PIL import Image, ImageOps +from PIL.PngImagePlugin import PngInfo +from scipy.ndimage import gaussian_filter +from skimage import exposure + + +from torchvision import transforms +from torchvision.ops import masks_to_boxes +import torchvision.transforms.functional as TF +import torch.nn.functional as torchfn +from .reactor_swapper import swap_face, get_current_faces_model, analyze_faces, get_face_single +from .reactor_logger import logger + +import folder_paths +from nodes import MAX_RESOLUTION + +import comfy +from .utils import tensor_to_pil, pil_to_tensor, tensor2pil, pil2tensor, pil2mask +from .modules import shared +from .modules.scripts import USDUMode, USDUSFMode, Script +from .modules.processing import StableDiffusionProcessing +from .modules.upscaler import UpscalerData + +# The modes available for Ultimate SD Upscale +MODES = { + "Linear": USDUMode.LINEAR, + "Chess": USDUMode.CHESS, + "None": USDUMode.NONE, +} +# The seam fix modes +SEAM_FIX_MODES = { + "None": USDUSFMode.NONE, + "Band Pass": USDUSFMode.BAND_PASS, + "Half Tile": USDUSFMode.HALF_TILE, + "Half Tile + Intersections": USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS, +} + + +def USDU_base_inputs(): + return [ + ("image", ("IMAGE",)), + # Sampling Params + ("model", ("MODEL",)), + ("positive", ("CONDITIONING",)), + ("negative", ("CONDITIONING",)), + ("vae", ("VAE",)), + ("upscale_by", ("FLOAT", {"default": 2, "min": 0.05, "max": 4, "step": 0.05})), + ("seed", ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff})), + ("steps", ("INT", {"default": 20, "min": 1, "max": 10000, "step": 1})), + ("cfg", ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0})), + ("sampler_name", (comfy.samplers.KSampler.SAMPLERS,)), + ("scheduler", (comfy.samplers.KSampler.SCHEDULERS,)), + ("denoise", ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01})), + # Upscale Params + ("upscale_model", ("UPSCALE_MODEL",)), + ("mode_type", (list(MODES.keys()),)), + ("tile_width", ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8})), + ("tile_height", ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8})), + ("mask_blur", ("INT", {"default": 8, "min": 0, "max": 64, "step": 1})), + ("tile_padding", ("INT", {"default": 32, "min": 0, "max": MAX_RESOLUTION, "step": 8})), + # Seam fix params + ("seam_fix_mode", (list(SEAM_FIX_MODES.keys()),)), + ("seam_fix_denoise", ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})), + ("seam_fix_width", ("INT", {"default": 64, "min": 0, "max": MAX_RESOLUTION, "step": 8})), + ("seam_fix_mask_blur", ("INT", {"default": 8, "min": 0, "max": 64, "step": 1})), + ("seam_fix_padding", ("INT", {"default": 16, "min": 0, "max": MAX_RESOLUTION, "step": 8})), + # Misc + ("force_uniform_tiles", ("BOOLEAN", {"default": True})), + ("tiled_decode", ("BOOLEAN", {"default": False})), + ] + + +def prepare_inputs(required: list, optional: list = None): + inputs = {} + if required: + inputs["required"] = {} + for name, type in required: + inputs["required"][name] = type + if optional: + inputs["optional"] = {} + for name, type in optional: + inputs["optional"][name] = type + return inputs + + +def remove_input(inputs: list, input_name: str): + for i, (n, _) in enumerate(inputs): + if n == input_name: + del inputs[i] + break + + +def rename_input(inputs: list, old_name: str, new_name: str): + for i, (n, t) in enumerate(inputs): + if n == old_name: + inputs[i] = (new_name, t) + break + + +class UltimateSDUpscale: + @classmethod + def INPUT_TYPES(s): + return prepare_inputs(USDU_base_inputs()) + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "upscale" + CATEGORY = "HD Nodes/upscaling" + + def upscale(self, image, model, positive, negative, vae, upscale_by, seed, + steps, cfg, sampler_name, scheduler, denoise, upscale_model, + mode_type, tile_width, tile_height, mask_blur, tile_padding, + seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode): + # + # Set up A1111 patches + # + + # Upscaler + # An object that the script works with + shared.sd_upscalers = [None] + shared.sd_upscalers[0] = UpscalerData() + # Where the actual upscaler is stored, will be used when the script upscales using the Upscaler in UpscalerData + shared.actual_upscaler = upscale_model + + # Set the batch of images + shared.batch = [tensor_to_pil(image, i) for i in range(len(image))] + + # Processing + sdprocessing = StableDiffusionProcessing( + tensor_to_pil(image), model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, upscale_by, force_uniform_tiles, tiled_decode + ) + + # + # Running the script + # + script = Script() + processed = script.run(p=sdprocessing, _=None, tile_width=tile_width, tile_height=tile_height, + mask_blur=mask_blur, padding=tile_padding, seams_fix_width=seam_fix_width, + seams_fix_denoise=seam_fix_denoise, seams_fix_padding=seam_fix_padding, + upscaler_index=0, save_upscaled_image=False, redraw_mode=MODES[mode_type], + save_seams_fix_image=False, seams_fix_mask_blur=seam_fix_mask_blur, + seams_fix_type=SEAM_FIX_MODES[seam_fix_mode], target_size_type=2, + custom_width=None, custom_height=None, custom_scale=upscale_by) + + # Return the resulting images + images = [pil_to_tensor(img) for img in shared.batch] + tensor = torch.cat(images, dim=0) + return (tensor,) + + +class UltimateSDUpscaleNoUpscale: + @classmethod + def INPUT_TYPES(s): + required = USDU_base_inputs() + remove_input(required, "upscale_model") + remove_input(required, "upscale_by") + rename_input(required, "image", "upscaled_image") + return prepare_inputs(required) + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "upscale" + CATEGORY = "HD Nodes/upscaling" + + def upscale(self, upscaled_image, model, positive, negative, vae, seed, + steps, cfg, sampler_name, scheduler, denoise, + mode_type, tile_width, tile_height, mask_blur, tile_padding, + seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode): + + shared.sd_upscalers[0] = UpscalerData() + shared.actual_upscaler = None + shared.batch = [tensor_to_pil(upscaled_image, i) for i in range(len(upscaled_image))] + sdprocessing = StableDiffusionProcessing( + tensor_to_pil(upscaled_image), model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, 1, force_uniform_tiles, tiled_decode + ) + + script = Script() + processed = script.run(p=sdprocessing, _=None, tile_width=tile_width, tile_height=tile_height, + mask_blur=mask_blur, padding=tile_padding, seams_fix_width=seam_fix_width, + seams_fix_denoise=seam_fix_denoise, seams_fix_padding=seam_fix_padding, + upscaler_index=0, save_upscaled_image=False, redraw_mode=MODES[mode_type], + save_seams_fix_image=False, seams_fix_mask_blur=seam_fix_mask_blur, + seams_fix_type=SEAM_FIX_MODES[seam_fix_mode], target_size_type=2, + custom_width=None, custom_height=None, custom_scale=1) + + images = [pil_to_tensor(img) for img in shared.batch] + tensor = torch.cat(images, dim=0) + return (tensor,) + + + +class MaskCombineOp: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask1": ("MASK",), + "mask2": ("MASK",), + "index": ("INT", {"default": 0, "min": 0, "max": 1}), + }, + } + + RETURN_TYPES = ("MASK",) + FUNCTION = "combine" + + CATEGORY = "HD Nodes" + + def combine(self, mask1, mask2, index=0): + + if (index == 0): + result = mask1 + else: + result = mask2 + + return (result,) + +class MaskCoverOp: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "destination": ("MASK",), + "source": ("MASK",), + "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + }, + } + + RETURN_TYPES = ("INT",) + FUNCTION = "cover" + + CATEGORY = "HD Nodes" + + def cover(self, destination, source, x, y): + output = destination.reshape((-1, destination.shape[-2], destination.shape[-1])).clone() + source = source.reshape((-1, source.shape[-2], source.shape[-1])) + + left, top = (x, y,) + right, bottom = (min(left + source.shape[-1], destination.shape[-1]), min(top + source.shape[-2], destination.shape[-2])) + visible_width, visible_height = (right - left, bottom - top,) + + source_portion = source[:, :visible_height, :visible_width] + destination_portion = destination[:, top:bottom, left:right] + + output[:, top:bottom, left:right] = destination_portion * source_portion + + output_ones = output.norm(1) + destination_ones = destination_portion.norm(1) + if (output_ones / destination_ones) < 0.2: + result = 1 + else: + result = 0 + + return (result,) + +class GetFaceIndex: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input_image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("INT",) + FUNCTION = "getindex" + + CATEGORY = "HD Nodes" + + def getindex(self, input_image): + target_img = tensor_to_pil(input_image) + target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) + target_faces = analyze_faces(target_img) + target_face, wrong_gender = get_face_single(target_img, target_faces) + if target_face is not None: + result = 1 + else: + result = 0 + return (result,) + + +class GetMaskArea: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "mask": ("MASK",), + "h_offset": ("INT", {"default": 100, "step":10}), + "h_cutoff": ("FLOAT", {"default": 0, "step":0.001}), + "max_width": ("INT", {"default": 1600, "min": 0, "max": MAX_RESOLUTION, "step": 100}), + "min_height": ("INT", {"default": 2400, "min": 0, "max": MAX_RESOLUTION, "step": 100}), + }, + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "MASK",) + FUNCTION = "getimage" + + CATEGORY = "HD Nodes" + + def get_mask_aabb(self, masks): + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device, dtype=torch.int) + + b = masks.shape[0] + + bounding_boxes = torch.zeros((b, 4), device=masks.device, dtype=torch.int) + is_empty = torch.zeros((b), device=masks.device, dtype=torch.bool) + for i in range(b): + mask = masks[i] + if mask.numel() == 0: + continue + if torch.max(mask != 0) == False: + is_empty[i] = True + continue + y, x = torch.where(mask) + bounding_boxes[i, 0] = torch.min(x) + bounding_boxes[i, 1] = torch.min(y) + bounding_boxes[i, 2] = torch.max(x) + bounding_boxes[i, 3] = torch.max(y) + + return bounding_boxes, is_empty + + def getimage(self, image1, image2, mask, h_offset = 100, h_cutoff=0, max_width=1600, min_height=2400): + bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0) + boxes, is_empty = self.get_mask_aabb(bounds) + + padding = 0.02 + box = boxes[0] + H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0]) + hh = int(int(H.item()) * (1.0 - h_cutoff)) + yy = int(Y.item()) - h_offset + hh = int((hh + 50) * (1 + padding)) + ww = int(int(W.item()) * (1 + 20 * padding)) + #X = int(X.item()) - 10 + #W = int(W.item()) + 20 + xx = max(int(int(X.item()) - ww / ( 1 + 20 * padding) * 10 * padding), 0) + image1_copy = copy.deepcopy(image1) + image11 = image1_copy[:,yy:yy+hh,xx:xx+ww] + image2_copy = copy.deepcopy(image2) + image22 = image2_copy[:,yy:yy+hh,xx:xx+ww] + mask = mask[:,yy:yy+hh,xx:xx+ww] + return (image11, image22, mask) + +# IMAGE LEVELS NODE + +class HD_Image_Levels: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "black_level": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 255.0, "step": 0.1}), + "mid_level": ("FLOAT", {"default": 127.5, "min": 0.0, "max": 255.0, "step": 0.1}), + "white_level": ("FLOAT", {"default": 255, "min": 0.0, "max": 255.0, "step": 0.1}), + } + } + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply_image_levels" + + CATEGORY = "HD Nodes/Image/Adjustment" + + def apply_image_levels(self, image, black_level, mid_level, white_level): + tensor_images = [] + for timg in image: + img = tensor2pil(timg) + levels = self.AdjustLevels(black_level, mid_level, white_level) + tensor_images.append(pil2tensor(levels.adjust(img))) + tensor_images = torch.cat(tensor_images, dim=0) + + # Return adjust image tensor + return (tensor_images, ) + + + class AdjustLevels: + def __init__(self, min_level, mid_level, max_level): + self.min_level = min_level + self.mid_level = mid_level + self.max_level = max_level + + def isBright(self, pil_image): + gray_img = np.array(pil_image.convert('L')) + + # 获取形状以及长宽 + img_shape = gray_img.shape + height, width = img_shape[0], img_shape[1] + size = gray_img.size + # 灰度图的直方图 + hist = cv2.calcHist([gray_img], [0], None, [256], [0, 256]) + + # 计算灰度图像素点偏离均值(128)程序 + a = 0 + ma = 0 + reduce_matrix = np.full((height, width), 128) + shift_value = gray_img - reduce_matrix + shift_sum = sum(map(sum, shift_value)) + + da = shift_sum / size + # 计算偏离128的平均偏差 + for i in range(256): + ma += (abs(i-128-da) * hist[i]) + m = abs(ma / size) + # 亮度系数 + k = abs(da) / m + # print(k) + if k[0] > 1: + # 过亮 + if da > 0: + #print("过亮") + return True + else: + #print("过暗") + return False + else: + #print("亮度正常") + return False + + def adjust(self, im): + + if not self.isBright(im): + im_arr = np.array(im) + im_arr[im_arr < self.min_level] = self.min_level + im_arr = (im_arr - self.min_level) * \ + (255 / (self.max_level - self.min_level)) + im_arr[im_arr < 0] = 0 + im_arr[im_arr > 255] = 255 + im_arr = im_arr.astype(np.uint8) + + im = Image.fromarray(im_arr) + im = ImageOps.autocontrast(im, cutoff=self.max_level) + + return im + +class SmoothEdge: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "masks": ("MASK",), + "sigma": ("FLOAT", {"default":1.5, "min":0.0, "max":12.0, "step":0.1}), + "gamma": ("INT", {"default": 20}), + }, + } + + RETURN_TYPES = ("MASK",) + + FUNCTION = "smooth" + + CATEGORY = "HD Nodes" + + def img2tensor(self, img, bgr2rgb=True, float32=True): + + if img.dtype == 'float64': + img = img.astype('float32') + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + img = torch.from_numpy(img.transpose(2, 0, 1)) + if float32: + img = img.float() + return img + + def smooth_region(self, image, tolerance): + from scipy.ndimage import gaussian_filter + image = image.convert("L") + mask_array = np.array(image) + smoothed_array = gaussian_filter(mask_array, sigma=tolerance) + threshold = np.max(smoothed_array) / 2 + smoothed_mask = np.where(smoothed_array >= threshold, 255, 0).astype(np.uint8) + smoothed_mask = exposure.adjust_gamma(smoothed_mask, gamma=20) + smoothed_image = Image.fromarray(smoothed_mask, mode="L") + return ImageOps.invert(smoothed_image.convert("RGB")) + + def smooth(self, masks, sigma=128, gamma=20): + if masks.ndim > 3: + regions = [] + for mask in masks: + mask_np = np.clip(255. * mask.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) + pil_image = Image.fromarray(mask_np, mode="L") + region_mask = self.smooth_region(pil_image, sigma) + region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1) + regions.append(region_tensor) + regions_tensor = torch.cat(regions, dim=0) + return (regions_tensor,) + else: + mask_np = np.clip(255. * masks.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) + pil_image = Image.fromarray(mask_np, mode="L") + region_mask = self.smooth_region(pil_image, sigma) + region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1) + return (region_tensor,) + +NODE_CLASS_MAPPINGS = { + "Combine HDMasks": MaskCombineOp, + "Cover HDMasks": MaskCoverOp, + "HD FaceIndex": GetFaceIndex, + "HD SmoothEdge": SmoothEdge, + "HD GetMaskArea": GetMaskArea, + "HD Image Levels": HD_Image_Levels, + "HD UltimateSDUpscale": UltimateSDUpscale, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Combine HDMasks": "Combine HDMasks", + "Cover HDMasks": "Cover HDMasks", + "HD FaceIndex": "HD FaceIndex", + "HD SmoothEdge": "HD SmoothEdge", + "HD GetMaskArea": "HD GetMaskArea", + "HD Image Levels": "HD Image Levels", + "HD UltimateSDUpscale": "HD Ultimate SD Upscale", +} diff --git a/README.md b/README.md index 447aa9b..5281152 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,7 @@ -# hd_node -comfyui节点 +# hd-nodes-comfyui +# comfyuiui节点 +## 1.判断是否背面 +## 2.更加图片色阶再判断是否提亮 +## 3.修改Ultimate SD Upscale只放大上半，提高效率 +## 4.自动裁剪图片，更加图片中人物裁剪合适大小 + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..2ab6930 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .ZHGNodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/folder_paths.py b/folder_paths.py new file mode 100644 index 0000000..4a38dee --- /dev/null +++ b/folder_paths.py @@ -0,0 +1,241 @@ +import os +import time + +supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors']) + +folder_names_and_paths = {} + +base_path = os.path.dirname(os.path.realpath(__file__)) +models_dir = os.path.join(base_path, "models") +folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions) +folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"]) + +folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions) +folder_names_and_paths["vae"] = ([os.path.join(models_dir, "vae")], supported_pt_extensions) +folder_names_and_paths["clip"] = ([os.path.join(models_dir, "clip")], supported_pt_extensions) +folder_names_and_paths["unet"] = ([os.path.join(models_dir, "unet")], supported_pt_extensions) +folder_names_and_paths["clip_vision"] = ([os.path.join(models_dir, "clip_vision")], supported_pt_extensions) +folder_names_and_paths["style_models"] = ([os.path.join(models_dir, "style_models")], supported_pt_extensions) +folder_names_and_paths["embeddings"] = ([os.path.join(models_dir, "embeddings")], supported_pt_extensions) +folder_names_and_paths["diffusers"] = ([os.path.join(models_dir, "diffusers")], ["folder"]) +folder_names_and_paths["vae_approx"] = ([os.path.join(models_dir, "vae_approx")], supported_pt_extensions) + +folder_names_and_paths["controlnet"] = ([os.path.join(models_dir, "controlnet"), os.path.join(models_dir, "t2i_adapter")], supported_pt_extensions) +folder_names_and_paths["gligen"] = ([os.path.join(models_dir, "gligen")], supported_pt_extensions) + +folder_names_and_paths["upscale_models"] = ([os.path.join(models_dir, "upscale_models")], supported_pt_extensions) + +folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], []) + +folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions) + +folder_names_and_paths["classifiers"] = ([os.path.join(models_dir, "classifiers")], {""}) + +output_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") +temp_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp") +input_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") + +filename_list_cache = {} + +if not os.path.exists(input_directory): + os.makedirs(input_directory) + +def set_output_directory(output_dir): + global output_directory + output_directory = output_dir + +def set_temp_directory(temp_dir): + global temp_directory + temp_directory = temp_dir + +def set_input_directory(input_dir): + global input_directory + input_directory = input_dir + +def get_output_directory(): + global output_directory + return output_directory + +def get_temp_directory(): + global temp_directory + return temp_directory + +def get_input_directory(): + global input_directory + return input_directory + + +#NOTE: used in http server so don't put folders that should not be accessed remotely +def get_directory_by_type(type_name): + if type_name == "output": + return get_output_directory() + if type_name == "temp": + return get_temp_directory() + if type_name == "input": + return get_input_directory() + return None + + +# determine base_dir rely on annotation if name is 'filename.ext [annotation]' format +# otherwise use default_path as base_dir +def annotated_filepath(name): + if name.endswith("[output]"): + base_dir = get_output_directory() + name = name[:-9] + elif name.endswith("[input]"): + base_dir = get_input_directory() + name = name[:-8] + elif name.endswith("[temp]"): + base_dir = get_temp_directory() + name = name[:-7] + else: + return name, None + + return name, base_dir + + +def get_annotated_filepath(name, default_dir=None): + name, base_dir = annotated_filepath(name) + + if base_dir is None: + if default_dir is not None: + base_dir = default_dir + else: + base_dir = get_input_directory() # fallback path + + return os.path.join(base_dir, name) + + +def exists_annotated_filepath(name): + name, base_dir = annotated_filepath(name) + + if base_dir is None: + base_dir = get_input_directory() # fallback path + + filepath = os.path.join(base_dir, name) + return os.path.exists(filepath) + + +def add_model_folder_path(folder_name, full_folder_path): + global folder_names_and_paths + if folder_name in folder_names_and_paths: + folder_names_and_paths[folder_name][0].append(full_folder_path) + else: + folder_names_and_paths[folder_name] = ([full_folder_path], set()) + +def get_folder_paths(folder_name): + return folder_names_and_paths[folder_name][0][:] + +def recursive_search(directory, excluded_dir_names=None): + if not os.path.isdir(directory): + return [], {} + + if excluded_dir_names is None: + excluded_dir_names = [] + + result = [] + dirs = {directory: os.path.getmtime(directory)} + for dirpath, subdirs, filenames in os.walk(directory, followlinks=True, topdown=True): + subdirs[:] = [d for d in subdirs if d not in excluded_dir_names] + for file_name in filenames: + relative_path = os.path.relpath(os.path.join(dirpath, file_name), directory) + result.append(relative_path) + for d in subdirs: + path = os.path.join(dirpath, d) + dirs[path] = os.path.getmtime(path) + return result, dirs + +def filter_files_extensions(files, extensions): + return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions or len(extensions) == 0, files))) + + + +def get_full_path(folder_name, filename): + global folder_names_and_paths + if folder_name not in folder_names_and_paths: + return None + folders = folder_names_and_paths[folder_name] + filename = os.path.relpath(os.path.join("/", filename), "/") + for x in folders[0]: + full_path = os.path.join(x, filename) + if os.path.isfile(full_path): + return full_path + + return None + +def get_filename_list_(folder_name): + global folder_names_and_paths + output_list = set() + folders = folder_names_and_paths[folder_name] + output_folders = {} + for x in folders[0]: + files, folders_all = recursive_search(x, excluded_dir_names=[".git"]) + output_list.update(filter_files_extensions(files, folders[1])) + output_folders = {**output_folders, **folders_all} + + return (sorted(list(output_list)), output_folders, time.perf_counter()) + +def cached_filename_list_(folder_name): + global filename_list_cache + global folder_names_and_paths + if folder_name not in filename_list_cache: + return None + out = filename_list_cache[folder_name] + if time.perf_counter() < (out[2] + 0.5): + return out + for x in out[1]: + time_modified = out[1][x] + folder = x + if os.path.getmtime(folder) != time_modified: + return None + + folders = folder_names_and_paths[folder_name] + for x in folders[0]: + if os.path.isdir(x): + if x not in out[1]: + return None + + return out + +def get_filename_list(folder_name): + out = cached_filename_list_(folder_name) + if out is None: + out = get_filename_list_(folder_name) + global filename_list_cache + filename_list_cache[folder_name] = out + return list(out[0]) + +def get_save_image_path(filename_prefix, output_dir, image_width=0, image_height=0): + def map_filename(filename): + prefix_len = len(os.path.basename(filename_prefix)) + prefix = filename[:prefix_len + 1] + try: + digits = int(filename[prefix_len + 1:].split('_')[0]) + except: + digits = 0 + return (digits, prefix) + + def compute_vars(input, image_width, image_height): + input = input.replace("%width%", str(image_width)) + input = input.replace("%height%", str(image_height)) + return input + + filename_prefix = compute_vars(filename_prefix, image_width, image_height) + + subfolder = os.path.dirname(os.path.normpath(filename_prefix)) + filename = os.path.basename(os.path.normpath(filename_prefix)) + + full_output_folder = os.path.join(output_dir, subfolder) + + if os.path.commonpath((output_dir, os.path.abspath(full_output_folder))) != output_dir: + print("Saving image outside the output folder is not allowed.") + return {} + + try: + counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1 + except ValueError: + counter = 1 + except FileNotFoundError: + os.makedirs(full_output_folder, exist_ok=True) + counter = 1 + return full_output_folder, filename, counter, subfolder, filename_prefix diff --git a/modules/devices.py b/modules/devices.py new file mode 100644 index 0000000..3e37b88 --- /dev/null +++ b/modules/devices.py @@ -0,0 +1,2 @@ +def torch_gc(): + pass diff --git a/modules/images.py b/modules/images.py new file mode 100644 index 0000000..502c819 --- /dev/null +++ b/modules/images.py @@ -0,0 +1,8 @@ +from PIL import Image + + +def flatten(img, bgcolor): + # Replace transparency with bgcolor + if img.mode in ("RGB"): + return img + return Image.alpha_composite(Image.new("RGBA", img.size, bgcolor), img).convert("RGB") diff --git a/modules/processing.py b/modules/processing.py new file mode 100644 index 0000000..aa8ecb7 --- /dev/null +++ b/modules/processing.py @@ -0,0 +1,163 @@ +from PIL import Image, ImageFilter +import torch +import math +from nodes import common_ksampler, VAEEncode, VAEDecode, VAEDecodeTiled +from ..utils import pil_to_tensor, tensor_to_pil, get_crop_region, expand_crop, crop_cond +from . import shared + +if (not hasattr(Image, 'Resampling')): # For older versions of Pillow + Image.Resampling = Image + + +class StableDiffusionProcessing: + + def __init__(self, init_img, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, upscale_by, uniform_tile_mode, tiled_decode): + # Variables used by the USDU script + self.init_images = [init_img] + self.image_mask = None + self.mask_blur = 0 + self.inpaint_full_res_padding = 0 + self.width = init_img.width + self.height = init_img.height + + # ComfyUI Sampler inputs + self.model = model + self.positive = positive + self.negative = negative + self.vae = vae + self.seed = seed + self.steps = steps + self.cfg = cfg + self.sampler_name = sampler_name + self.scheduler = scheduler + self.denoise = denoise + + # Variables used only by this script + self.init_size = init_img.width, init_img.height + self.upscale_by = upscale_by + self.uniform_tile_mode = uniform_tile_mode + self.tiled_decode = tiled_decode + self.vae_decoder = VAEDecode() + self.vae_encoder = VAEEncode() + self.vae_decoder_tiled = VAEDecodeTiled() + + # Other required A1111 variables for the USDU script that is currently unused in this script + self.extra_generation_params = {} + + +class Processed: + + def __init__(self, p: StableDiffusionProcessing, images: list, seed: int, info: str): + self.images = images + self.seed = seed + self.info = info + + def infotext(self, p: StableDiffusionProcessing, index): + return None + + +def fix_seed(p: StableDiffusionProcessing): + pass + + +def process_images(p: StableDiffusionProcessing) -> Processed: + # Where the main image generation happens in A1111 + + # Setup + image_mask = p.image_mask.convert('L') + init_image = p.init_images[0] + + # Locate the white region of the mask outlining the tile and add padding + crop_region = get_crop_region(image_mask, p.inpaint_full_res_padding) + + if p.uniform_tile_mode: + # Expand the crop region to match the processing size ratio and then resize it to the processing size + x1, y1, x2, y2 = crop_region + crop_width = x2 - x1 + crop_height = y2 - y1 + crop_ratio = crop_width / crop_height + p_ratio = p.width / p.height + if crop_ratio > p_ratio: + target_width = crop_width + target_height = round(crop_width / p_ratio) + else: + target_width = round(crop_height * p_ratio) + target_height = crop_height + crop_region, _ = expand_crop(crop_region, image_mask.width, image_mask.height, target_width, target_height) + tile_size = p.width, p.height + else: + # Uses the minimal size that can fit the mask, minimizes tile size but may lead to image sizes that the model is not trained on + x1, y1, x2, y2 = crop_region + crop_width = x2 - x1 + crop_height = y2 - y1 + target_width = math.ceil(crop_width / 8) * 8 + target_height = math.ceil(crop_height / 8) * 8 + crop_region, tile_size = expand_crop(crop_region, image_mask.width, + image_mask.height, target_width, target_height) + + # Blur the mask + if p.mask_blur > 0: + image_mask = image_mask.filter(ImageFilter.GaussianBlur(p.mask_blur)) + + # Crop the images to get the tiles that will be used for generation + tiles = [img.crop(crop_region) for img in shared.batch] + + # Assume the same size for all images in the batch + initial_tile_size = tiles[0].size + + # Resize if necessary + for i, tile in enumerate(tiles): + if tile.size != tile_size: + tiles[i] = tile.resize(tile_size, Image.Resampling.LANCZOS) + + # Crop conditioning + positive_cropped = crop_cond(p.positive, crop_region, p.init_size, init_image.size, tile_size) + negative_cropped = crop_cond(p.negative, crop_region, p.init_size, init_image.size, tile_size) + + # Encode the image + batched_tiles = torch.cat([pil_to_tensor(tile) for tile in tiles], dim=0) + (latent,) = p.vae_encoder.encode(p.vae, batched_tiles) + + # Generate samples + (samples,) = common_ksampler(p.model, p.seed, p.steps, p.cfg, p.sampler_name, + p.scheduler, positive_cropped, negative_cropped, latent, denoise=p.denoise) + + # Decode the sample + if not p.tiled_decode: + (decoded,) = p.vae_decoder.decode(p.vae, samples) + else: + print("[USDU] Using tiled decode") + (decoded,) = p.vae_decoder_tiled.decode(p.vae, samples, 512) # Default tile size is 512 + + # Convert the sample to a PIL image + tiles_sampled = [tensor_to_pil(decoded, i) for i in range(len(decoded))] + + for i, tile_sampled in enumerate(tiles_sampled): + init_image = shared.batch[i] + + # Resize back to the original size + if tile_sampled.size != initial_tile_size: + tile_sampled = tile_sampled.resize(initial_tile_size, Image.Resampling.LANCZOS) + + + # Put the tile into position + image_tile_only = Image.new('RGBA', init_image.size) + image_tile_only.paste(tile_sampled, 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(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 = init_image.convert('RGBA') + result.alpha_composite(image_tile_only) + + # Convert back to RGB + result = result.convert('RGB') + + shared.batch[i] = result + + processed = Processed(p, [shared.batch[0]], p.seed, None) + return processed diff --git a/modules/scripts.py b/modules/scripts.py new file mode 100644 index 0000000..3525f65 --- /dev/null +++ b/modules/scripts.py @@ -0,0 +1,473 @@ +import math +from PIL import Image, ImageDraw, ImageOps +from .processing import StableDiffusionProcessing +from .processing import Processed +from enum import Enum +from . import shared +from . import images +from . import devices +from . import processing + +class USDUMode(Enum): + LINEAR = 0 + CHESS = 1 + NONE = 2 + +class USDUSFMode(Enum): + NONE = 0 + BAND_PASS = 1 + HALF_TILE = 2 + HALF_TILE_PLUS_INTERSECTIONS = 3 + +class USDUpscaler(): + + def __init__(self, p, image, upscaler_index:int, save_redraw, save_seams_fix, tile_width, tile_height) -> None: + self.p:StableDiffusionProcessing = p + self.image:Image = image + self.scale_factor = math.ceil(max(p.width, p.height) / max(image.width, image.height)) + self.upscaler = shared.sd_upscalers[upscaler_index] + self.redraw = USDURedraw() + self.redraw.save = save_redraw + self.redraw.tile_width = tile_width if tile_width > 0 else tile_height + self.redraw.tile_height = tile_height if tile_height > 0 else tile_width + self.seams_fix = USDUSeamsFix() + self.seams_fix.save = save_seams_fix + self.seams_fix.tile_width = tile_width if tile_width > 0 else tile_height + self.seams_fix.tile_height = tile_height if tile_height > 0 else tile_width + self.initial_info = None + self.rows = math.ceil(self.p.height / 4 / self.redraw.tile_height) + self.cols = math.ceil(self.p.width / self.redraw.tile_width) + + def get_factor(self, num): + # Its just return, don't need elif + if num == 1: + return 2 + if num % 4 == 0: + return 4 + if num % 3 == 0: + return 3 + if num % 2 == 0: + return 2 + return 0 + + def get_factors(self): + scales = [] + current_scale = 1 + current_scale_factor = self.get_factor(self.scale_factor) + while current_scale_factor == 0: + self.scale_factor += 1 + current_scale_factor = self.get_factor(self.scale_factor) + while current_scale < self.scale_factor: + current_scale_factor = self.get_factor(self.scale_factor // current_scale) + scales.append(current_scale_factor) + current_scale = current_scale * current_scale_factor + if current_scale_factor == 0: + break + self.scales = enumerate(scales) + + def upscale(self): + # Log info + print(f"Canva size: {self.p.width}x{self.p.height}") + print(f"Image size: {self.image.width}x{self.image.height}") + print(f"Scale factor: {self.scale_factor}") + # Check upscaler is not empty + if self.upscaler.name == "None": + self.image = self.image.resize((self.p.width, self.p.height), resample=Image.LANCZOS) + return + # Get list with scale factors + self.get_factors() + # Upscaling image over all factors + for index, value in self.scales: + print(f"Upscaling iteration {index+1} with scale factor {value}") + self.image = self.upscaler.scaler.upscale(self.image, value, self.upscaler.data_path) + # Resize image to set values + self.image = self.image.resize((self.p.width, self.p.height), resample=Image.LANCZOS) + shared.batch = [self.image] + [img.resize((self.p.width, self.p.height), resample=Image.LANCZOS) for img in shared.batch[1:]] + + + def setup_redraw(self, redraw_mode, padding, mask_blur): + self.redraw.mode = USDUMode(redraw_mode) + self.redraw.enabled = self.redraw.mode != USDUMode.NONE + self.redraw.padding = padding + self.p.mask_blur = mask_blur + + + def setup_seams_fix(self, padding, denoise, mask_blur, width, mode): + self.seams_fix.padding = padding + self.seams_fix.denoise = denoise + self.seams_fix.mask_blur = mask_blur + self.seams_fix.width = width + self.seams_fix.mode = USDUSFMode(mode) + self.seams_fix.enabled = self.seams_fix.mode != USDUSFMode.NONE + + + def save_image(self): + if type(self.p.prompt) != list: + images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt, shared.opts.samples_format, info=self.initial_info, p=self.p) + else: + images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt[0], shared.opts.samples_format, info=self.initial_info, p=self.p) + + def calc_jobs_count(self): + redraw_job_count = (self.rows * self.cols) if self.redraw.enabled else 0 + seams_job_count = 0 + if self.seams_fix.mode == USDUSFMode.BAND_PASS: + seams_job_count = self.rows + self.cols - 2 + elif self.seams_fix.mode == USDUSFMode.HALF_TILE: + seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols + elif self.seams_fix.mode == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS: + seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols + (self.rows - 1) * (self.cols - 1) + + shared.state.job_count = redraw_job_count + seams_job_count + + def print_info(self): + print(f"Tile size: {self.redraw.tile_width}x{self.redraw.tile_height}") + print(f"Tiles amount: {self.rows * self.cols}") + print(f"Grid: {self.rows}x{self.cols}") + print(f"Redraw enabled: {self.redraw.enabled}") + print(f"Seams fix mode: {self.seams_fix.mode.name}") + + def add_extra_info(self): + self.p.extra_generation_params["Ultimate SD upscale upscaler"] = self.upscaler.name + self.p.extra_generation_params["Ultimate SD upscale tile_width"] = self.redraw.tile_width + self.p.extra_generation_params["Ultimate SD upscale tile_height"] = self.redraw.tile_height + self.p.extra_generation_params["Ultimate SD upscale mask_blur"] = self.p.mask_blur + self.p.extra_generation_params["Ultimate SD upscale padding"] = self.redraw.padding + + def process(self): + shared.state.begin() + self.calc_jobs_count() + self.result_images = [] + if self.redraw.enabled: + self.image = self.redraw.start(self.p, self.image, self.rows, self.cols) + self.initial_info = self.redraw.initial_info + self.result_images.append(self.image) + if self.redraw.save: + self.save_image() + + if self.seams_fix.enabled: + self.image = self.seams_fix.start(self.p, self.image, self.rows, self.cols) + self.initial_info = self.seams_fix.initial_info + self.result_images.append(self.image) + if self.seams_fix.save: + self.save_image() + shared.state.end() + +class USDURedraw(): + + def init_draw(self, p, width, height): + p.inpaint_full_res = True + p.inpaint_full_res_padding = self.padding + p.width = math.ceil((self.tile_width+self.padding) / 8) * 8 + p.height = math.ceil((self.tile_height+self.padding) / 8) * 8 + mask = Image.new("L", (width, height), "black") + draw = ImageDraw.Draw(mask) + return mask, draw + + def calc_rectangle(self, xi, yi): + x1 = xi * self.tile_width + y1 = yi * self.tile_height + x2 = xi * self.tile_width + self.tile_width + y2 = yi * self.tile_height + self.tile_height + + return x1, y1, x2, y2 + + def linear_process(self, p, image, rows, cols): + mask, draw = self.init_draw(p, image.width, image.height) + for yi in range(rows): + for xi in range(cols): + if shared.state.interrupted: + break + draw.rectangle(self.calc_rectangle(xi, yi), fill="white") + p.init_images = [image] + p.image_mask = mask + processed = processing.process_images(p) + draw.rectangle(self.calc_rectangle(xi, yi), fill="black") + if (len(processed.images) > 0): + image = processed.images[0] + + p.width = image.width + p.height = image.height + self.initial_info = processed.infotext(p, 0) + + return image + + def chess_process(self, p, image, rows, cols): + mask, draw = self.init_draw(p, image.width, image.height) + tiles = [] + # calc tiles colors + for yi in range(rows): + for xi in range(cols): + if shared.state.interrupted: + break + if xi == 0: + tiles.append([]) + color = xi % 2 == 0 + if yi > 0 and yi % 2 != 0: + color = not color + tiles[yi].append(color) + + for yi in range(len(tiles)): + for xi in range(len(tiles[yi])): + if shared.state.interrupted: + break + if not tiles[yi][xi]: + tiles[yi][xi] = not tiles[yi][xi] + continue + tiles[yi][xi] = not tiles[yi][xi] + draw.rectangle(self.calc_rectangle(xi, yi), fill="white") + p.init_images = [image] + p.image_mask = mask + processed = processing.process_images(p) + draw.rectangle(self.calc_rectangle(xi, yi), fill="black") + if (len(processed.images) > 0): + image = processed.images[0] + + for yi in range(len(tiles)): + for xi in range(len(tiles[yi])): + if shared.state.interrupted: + break + if not tiles[yi][xi]: + continue + draw.rectangle(self.calc_rectangle(xi, yi), fill="white") + p.init_images = [image] + p.image_mask = mask + processed = processing.process_images(p) + draw.rectangle(self.calc_rectangle(xi, yi), fill="black") + if (len(processed.images) > 0): + image = processed.images[0] + + p.width = image.width + p.height = image.height + self.initial_info = processed.infotext(p, 0) + + return image + + def start(self, p, image, rows, cols): + self.initial_info = None + if self.mode == USDUMode.LINEAR: + return self.linear_process(p, image, rows, cols) + if self.mode == USDUMode.CHESS: + return self.chess_process(p, image, rows, cols) + +class USDUSeamsFix(): + + def init_draw(self, p): + self.initial_info = None + p.width = math.ceil((self.tile_width+self.padding) / 8) * 8 + p.height = math.ceil((self.tile_height+self.padding) / 8) * 8 + + def half_tile_process(self, p, image, rows, cols): + + self.init_draw(p) + processed = None + + gradient = Image.linear_gradient("L") + row_gradient = Image.new("L", (self.tile_width, self.tile_height), "black") + row_gradient.paste(gradient.resize( + (self.tile_width, self.tile_height//2), resample=Image.BICUBIC), (0, 0)) + row_gradient.paste(gradient.rotate(180).resize( + (self.tile_width, self.tile_height//2), resample=Image.BICUBIC), + (0, self.tile_height//2)) + col_gradient = Image.new("L", (self.tile_width, self.tile_height), "black") + col_gradient.paste(gradient.rotate(90).resize( + (self.tile_width//2, self.tile_height), resample=Image.BICUBIC), (0, 0)) + col_gradient.paste(gradient.rotate(270).resize( + (self.tile_width//2, self.tile_height), resample=Image.BICUBIC), (self.tile_width//2, 0)) + + p.denoising_strength = self.denoise + p.mask_blur = self.mask_blur + + for yi in range(rows-1): + for xi in range(cols): + if shared.state.interrupted: + break + p.width = self.tile_width + p.height = self.tile_height + p.inpaint_full_res = True + p.inpaint_full_res_padding = self.padding + mask = Image.new("L", (image.width, image.height), "black") + mask.paste(row_gradient, (xi*self.tile_width, yi*self.tile_height + self.tile_height//2)) + + p.init_images = [image] + p.image_mask = mask + processed = processing.process_images(p) + if (len(processed.images) > 0): + image = processed.images[0] + + for yi in range(rows): + for xi in range(cols-1): + if shared.state.interrupted: + break + p.width = self.tile_width + p.height = self.tile_height + p.inpaint_full_res = True + p.inpaint_full_res_padding = self.padding + mask = Image.new("L", (image.width, image.height), "black") + mask.paste(col_gradient, (xi*self.tile_width+self.tile_width//2, yi*self.tile_height)) + + p.init_images = [image] + p.image_mask = mask + processed = processing.process_images(p) + if (len(processed.images) > 0): + image = processed.images[0] + + p.width = image.width + p.height = image.height + if processed is not None: + self.initial_info = processed.infotext(p, 0) + + return image + + def half_tile_process_corners(self, p, image, rows, cols): + fixed_image = self.half_tile_process(p, image, rows, cols) + processed = None + self.init_draw(p) + gradient = Image.radial_gradient("L").resize( + (self.tile_width, self.tile_height), resample=Image.BICUBIC) + gradient = ImageOps.invert(gradient) + p.denoising_strength = self.denoise + #p.mask_blur = 0 + p.mask_blur = self.mask_blur + + for yi in range(rows-1): + for xi in range(cols-1): + if shared.state.interrupted: + break + p.width = self.tile_width + p.height = self.tile_height + p.inpaint_full_res = True + p.inpaint_full_res_padding = 0 + mask = Image.new("L", (fixed_image.width, fixed_image.height), "black") + mask.paste(gradient, (xi*self.tile_width + self.tile_width//2, + yi*self.tile_height + self.tile_height//2)) + + p.init_images = [fixed_image] + p.image_mask = mask + processed = processing.process_images(p) + if (len(processed.images) > 0): + fixed_image = processed.images[0] + + p.width = fixed_image.width + p.height = fixed_image.height + if processed is not None: + self.initial_info = processed.infotext(p, 0) + + return fixed_image + + def band_pass_process(self, p, image, cols, rows): + + self.init_draw(p) + processed = None + + p.denoising_strength = self.denoise + p.mask_blur = 0 + + gradient = Image.linear_gradient("L") + mirror_gradient = Image.new("L", (256, 256), "black") + mirror_gradient.paste(gradient.resize((256, 128), resample=Image.BICUBIC), (0, 0)) + mirror_gradient.paste(gradient.rotate(180).resize((256, 128), resample=Image.BICUBIC), (0, 128)) + + row_gradient = mirror_gradient.resize((image.width, self.width), resample=Image.BICUBIC) + col_gradient = mirror_gradient.rotate(90).resize((self.width, image.height), resample=Image.BICUBIC) + + for xi in range(1, rows): + if shared.state.interrupted: + break + p.width = self.width + self.padding * 2 + p.height = image.height + p.inpaint_full_res = True + p.inpaint_full_res_padding = self.padding + mask = Image.new("L", (image.width, image.height), "black") + mask.paste(col_gradient, (xi * self.tile_width - self.width // 2, 0)) + + p.init_images = [image] + p.image_mask = mask + processed = processing.process_images(p) + if (len(processed.images) > 0): + image = processed.images[0] + for yi in range(1, cols): + if shared.state.interrupted: + break + p.width = image.width + p.height = self.width + self.padding * 2 + p.inpaint_full_res = True + p.inpaint_full_res_padding = self.padding + mask = Image.new("L", (image.width, image.height), "black") + mask.paste(row_gradient, (0, yi * self.tile_height - self.width // 2)) + + p.init_images = [image] + p.image_mask = mask + processed = processing.process_images(p) + if (len(processed.images) > 0): + image = processed.images[0] + + p.width = image.width + p.height = image.height + if processed is not None: + self.initial_info = processed.infotext(p, 0) + + return image + + def start(self, p, image, rows, cols): + if USDUSFMode(self.mode) == USDUSFMode.BAND_PASS: + return self.band_pass_process(p, image, rows, cols) + elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE: + return self.half_tile_process(p, image, rows, cols) + elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS: + return self.half_tile_process_corners(p, image, rows, cols) + else: + return image + +class Script(): + def title(self): + return "Ultimate SD upscale" + + def show(self, is_img2img): + return is_img2img + + def run(self, p, _, tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding, + upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur, + seams_fix_type, target_size_type, custom_width, custom_height, custom_scale): + + # Init + processing.fix_seed(p) + devices.torch_gc() + + p.do_not_save_grid = True + p.do_not_save_samples = True + p.inpaint_full_res = False + + p.inpainting_fill = 1 + p.n_iter = 1 + p.batch_size = 1 + + seed = p.seed + + # Init image + init_img = p.init_images[0] + if init_img == None: + return Processed(p, [], seed, "Empty image") + init_img = images.flatten(init_img, shared.opts.img2img_background_color) + + #override size + if target_size_type == 1: + p.width = custom_width + p.height = custom_height + if target_size_type == 2: + p.width = math.ceil((init_img.width * custom_scale) / 8) * 8 + p.height = math.ceil((init_img.height * custom_scale) / 8) * 8 + + # Upscaling + upscaler = USDUpscaler(p, init_img, upscaler_index, save_upscaled_image, save_seams_fix_image, tile_width, tile_height) + upscaler.upscale() + + # Drawing + upscaler.setup_redraw(redraw_mode, padding, mask_blur) + upscaler.setup_seams_fix(seams_fix_padding, seams_fix_denoise, seams_fix_mask_blur, seams_fix_width, seams_fix_type) + upscaler.print_info() + upscaler.add_extra_info() + upscaler.process() + result_images = upscaler.result_images + + return Processed(p, result_images, seed, upscaler.initial_info if upscaler.initial_info is not None else "") + diff --git a/modules/shared.py b/modules/shared.py new file mode 100644 index 0000000..9d4bdd7 --- /dev/null +++ b/modules/shared.py @@ -0,0 +1,24 @@ +class Options: + img2img_background_color = "#ffffff" # Set to white for now + + +class State: + interrupted = False + + def begin(self): + pass + + def end(self): + pass + + +opts = Options() +state = State() + +# Will only ever hold 1 upscaler +sd_upscalers = [None] +# The upscaler usable by ComfyUI nodes +actual_upscaler = None + +# Batch of images to upscale +batch = None diff --git a/modules/upscaler.py b/modules/upscaler.py new file mode 100644 index 0000000..a6556c9 --- /dev/null +++ b/modules/upscaler.py @@ -0,0 +1,30 @@ +from PIL import Image +from ..utils import tensor_to_pil, pil_to_tensor +from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel +from ..modules import shared + +if (not hasattr(Image, 'Resampling')): # For older versions of Pillow + Image.Resampling = Image + + +class Upscaler: + + def _upscale(self, img: Image, scale): + if (shared.actual_upscaler is None): + return img.resize((img.width * scale, img.height * scale), Image.Resampling.LANCZOS) + tensor = pil_to_tensor(img) + image_upscale_node = ImageUpscaleWithModel() + (upscaled,) = image_upscale_node.upscale(shared.actual_upscaler, tensor) + return tensor_to_pil(upscaled) + + def upscale(self, img: Image, scale, selected_model: str = None): + shared.batch = [self._upscale(img, scale) for img in shared.batch] + return shared.batch[0] + + +class UpscalerData: + name = "" + data_path = "" + + def __init__(self): + self.scaler = Upscaler() diff --git a/reactor_logger.py b/reactor_logger.py new file mode 100644 index 0000000..aea3d9a --- /dev/null +++ b/reactor_logger.py @@ -0,0 +1,43 @@ +import logging +import copy +import sys + +class ColoredFormatter(logging.Formatter): + COLORS = { + "DEBUG": "\033[0;36m", # CYAN + "STATUS": "\033[38;5;173m", # Calm ORANGE + "INFO": "\033[0;32m", # GREEN + "WARNING": "\033[0;33m", # YELLOW + "ERROR": "\033[0;31m", # RED + "CRITICAL": "\033[0;37;41m", # WHITE ON RED + "RESET": "\033[0m", # RESET COLOR + } + + def format(self, record): + colored_record = copy.copy(record) + levelname = colored_record.levelname + seq = self.COLORS.get(levelname, self.COLORS["RESET"]) + colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}" + return super().format(colored_record) + + +# Create a new logger +logger = logging.getLogger("ReActor") +logger.propagate = False + +# Add Custom Level +# logging.addLevelName(logging.INFO, "STATUS") +#addLoggingLevel("STATUS", logging.INFO + 5) + +# Add handler if we don't have one. +if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter( + ColoredFormatter("[%(name)s] %(asctime)s - %(levelname)s - %(message)s",datefmt="%H:%M:%S") + ) + logger.addHandler(handler) + +# Configure logger +loglevel_string = getattr(None, "reactor_loglevel", "INFO") +loglevel = getattr(logging, loglevel_string.upper(), "info") +logger.setLevel(loglevel) diff --git a/reactor_swapper.py b/reactor_swapper.py new file mode 100644 index 0000000..0c90f54 --- /dev/null +++ b/reactor_swapper.py @@ -0,0 +1,279 @@ +import copy +import os +import shutil +from dataclasses import dataclass +from typing import List, Union + +import cv2 +import numpy as np +from PIL import Image + +import insightface +from insightface.app.common import Face +try: + import torch.cuda as cuda +except: + cuda = None + +from .reactor_logger import logger +from .reactor_utils import move_path, get_image_md5hash +import folder_paths + +import warnings + +np.warnings = warnings +np.warnings.filterwarnings('ignore') + +if cuda is not None: + if cuda.is_available(): + providers = ["CUDAExecutionProvider"] + else: + providers = ["CPUExecutionProvider"] +else: + providers = ["CPUExecutionProvider"] + +models_path_old = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models") +insightface_path_old = os.path.join(models_path_old, "insightface") +insightface_models_path_old = os.path.join(insightface_path_old, "models") + +models_path = folder_paths.models_dir +insightface_path = os.path.join(models_path, "insightface") +insightface_models_path = os.path.join(insightface_path, "models") + +if os.path.exists(models_path_old): + move_path(insightface_models_path_old, insightface_models_path) + move_path(insightface_path_old, insightface_path) + move_path(models_path_old, models_path) +if os.path.exists(insightface_path) and os.path.exists(insightface_path_old): + shutil.rmtree(insightface_path_old) + shutil.rmtree(models_path_old) + + +FS_MODEL = None +CURRENT_FS_MODEL_PATH = None + +ANALYSIS_MODEL = None + +SOURCE_FACES = None +SOURCE_IMAGE_HASH = None +TARGET_FACES = None +TARGET_IMAGE_HASH = None + +def get_current_faces_model(): + global SOURCE_FACES + return SOURCE_FACES + +def getAnalysisModel(): + global ANALYSIS_MODEL + if ANALYSIS_MODEL is None: + ANALYSIS_MODEL = insightface.app.FaceAnalysis( + name="buffalo_l", providers=providers, root=insightface_path + ) + return ANALYSIS_MODEL + + +def getFaceSwapModel(model_path: str): + global FS_MODEL + global CURRENT_FS_MODEL_PATH + if CURRENT_FS_MODEL_PATH is None or CURRENT_FS_MODEL_PATH != model_path: + CURRENT_FS_MODEL_PATH = model_path + FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers) + + return FS_MODEL + + +def get_face_gender( + face, + face_index, + gender_condition, + operated: str +): + gender = [ + x.sex + for x in face + ] + gender.reverse() + face_gender = gender[face_index] + logger.status("%s Face %s: Detected Gender -%s-", operated, face_index, face_gender) + if (gender_condition == 1 and face_gender == "F") or (gender_condition == 2 and face_gender == "M"): + logger.status("OK - Detected Gender matches Condition") + try: + return sorted(face, key=lambda x: x.bbox[0])[face_index], 0 + except IndexError: + return None, 0 + else: + logger.status("WRONG - Detected Gender doesn't match Condition") + return sorted(face, key=lambda x: x.bbox[0])[face_index], 1 + + +def half_det_size(det_size): + logger.status("Trying to halve 'det_size' parameter") + return (det_size[0] // 2, det_size[1] // 2) + +def analyze_faces(img_data: np.ndarray, det_size=(640, 640)): + face_analyser = copy.deepcopy(getAnalysisModel()) + face_analyser.prepare(ctx_id=0, det_size=det_size) + return face_analyser.get(img_data) + +def get_face_single(img_data: np.ndarray, face, face_index=0, det_size=(640, 640), gender_source=0, gender_target=0): + + buffalo_path = os.path.join(insightface_models_path, "buffalo_l.zip") + if os.path.exists(buffalo_path): + os.remove(buffalo_path) + + if gender_source != 0: + if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320: + det_size_half = half_det_size(det_size) + return get_face_single(img_data, analyze_faces(img_data, det_size_half), face_index, det_size_half, gender_source, gender_target) + return get_face_gender(face,face_index,gender_source,"Source") + + if gender_target != 0: + if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320: + det_size_half = half_det_size(det_size) + return get_face_single(img_data, analyze_faces(img_data, det_size_half), face_index, det_size_half, gender_source, gender_target) + return get_face_gender(face,face_index,gender_target,"Target") + + if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320: + det_size_half = half_det_size(det_size) + return get_face_single(img_data, analyze_faces(img_data, det_size_half), face_index, det_size_half, gender_source, gender_target) + + try: + return sorted(face, key=lambda x: x.bbox[0])[face_index], 0 + except IndexError: + return None, 0 + + +def swap_face( + source_img: Union[Image.Image, None], + target_img: Image.Image, + model: Union[str, None] = None, + source_faces_index: List[int] = [0], + faces_index: List[int] = [0], + gender_source: int = 0, + gender_target: int = 0, + face_model: Union[Face, None] = None, +): + global SOURCE_FACES, SOURCE_IMAGE_HASH, TARGET_FACES, TARGET_IMAGE_HASH + result_image = target_img + + if model is not None: + + if isinstance(source_img, str): # source_img is a base64 string + import base64, io + if 'base64,' in source_img: # check if the base64 string has a data URL scheme + # split the base64 string to get the actual base64 encoded image data + base64_data = source_img.split('base64,')[-1] + # decode base64 string to bytes + img_bytes = base64.b64decode(base64_data) + else: + # if no data URL scheme, just decode + img_bytes = base64.b64decode(source_img) + + source_img = Image.open(io.BytesIO(img_bytes)) + + target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR) + + if source_img is not None: + + source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR) + + source_image_md5hash = get_image_md5hash(source_img) + + if SOURCE_IMAGE_HASH is None: + SOURCE_IMAGE_HASH = source_image_md5hash + source_image_same = False + else: + source_image_same = True if SOURCE_IMAGE_HASH == source_image_md5hash else False + if not source_image_same: + SOURCE_IMAGE_HASH = source_image_md5hash + + logger.info("Source Image MD5 Hash = %s", SOURCE_IMAGE_HASH) + logger.info("Source Image the Same? %s", source_image_same) + + if SOURCE_FACES is None or not source_image_same: + logger.status("Analyzing Source Image...") + source_faces = analyze_faces(source_img) + SOURCE_FACES = source_faces + elif source_image_same: + logger.status("Using Hashed Source Face(s) Model...") + source_faces = SOURCE_FACES + + elif face_model is not None: + + logger.status("Using Loaded Source Face Model...") + source_face_model = [face_model] + source_faces = source_face_model + + else: + logger.error("Cannot detect any Source") + + if source_faces is not None: + + target_image_md5hash = get_image_md5hash(target_img) + + if TARGET_IMAGE_HASH is None: + TARGET_IMAGE_HASH = target_image_md5hash + target_image_same = False + else: + target_image_same = True if TARGET_IMAGE_HASH == target_image_md5hash else False + if not target_image_same: + TARGET_IMAGE_HASH = target_image_md5hash + + logger.info("Target Image MD5 Hash = %s", TARGET_IMAGE_HASH) + logger.info("Target Image the Same? %s", target_image_same) + + if TARGET_FACES is None or not target_image_same: + logger.status("Analyzing Target Image...") + target_faces = analyze_faces(target_img) + TARGET_FACES = target_faces + elif target_image_same: + logger.status("Using Hashed Target Face(s) Model...") + target_faces = TARGET_FACES + + if source_img is not None: + source_face, wrong_gender = get_face_single(source_img, source_faces, face_index=source_faces_index[0], gender_source=gender_source) + else: + source_face = sorted(source_faces, key=lambda x: x.bbox[0])[source_faces_index[0]] + wrong_gender = 0 + + if len(source_faces_index) != 0 and len(source_faces_index) != 1 and len(source_faces_index) != len(faces_index): + logger.status(f'Source Faces must have no entries (default=0), one entry, or same number of entries as target faces.') + elif source_face is not None: + result = target_img + model_path = model_path = os.path.join(insightface_path, model) + face_swapper = getFaceSwapModel(model_path) + + source_face_idx = 0 + + for face_num in faces_index: + if len(source_faces_index) > 1 and source_face_idx > 0: + source_face, wrong_gender = get_face_single(source_img, source_faces, face_index=source_faces_index[source_face_idx], gender_source=gender_source) + source_face_idx += 1 + + if source_face is not None and wrong_gender == 0: + target_face, wrong_gender = get_face_single(target_img, target_faces, face_index=face_num, gender_target=gender_target) + if target_face is not None and wrong_gender == 0: + logger.status(f"Swapping...") + result = face_swapper.get(result, target_face, source_face) + elif wrong_gender == 1: + wrong_gender = 0 + if source_face_idx == len(source_faces_index): + result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) + return result_image + else: + logger.status(f"No target face found for {face_num}") + elif wrong_gender == 1: + wrong_gender = 0 + if source_face_idx == len(source_faces_index): + result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) + return result_image + else: + logger.status(f"No source face found for face number {source_face_idx}.") + + result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) + + else: + logger.status("No source face(s) in the provided Index") + else: + logger.status("No source face(s) found") + return result_image diff --git a/reactor_utils.py b/reactor_utils.py new file mode 100644 index 0000000..1a37773 --- /dev/null +++ b/reactor_utils.py @@ -0,0 +1,152 @@ +import os +from PIL import Image +import numpy as np +import torch +import cv2 +import logging +import hashlib +from insightface.app.common import Face +from safetensors.torch import save_file, safe_open + + +def tensor_to_pil(img_tensor, batch_index=0): + # Convert tensor of shape [batch_size, channels, height, width] at the batch_index to PIL Image + img_tensor = img_tensor[batch_index].unsqueeze(0) + i = 255. * img_tensor.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8).squeeze()) + return img + + +def batch_tensor_to_pil(img_tensor): + # Convert tensor of shape [batch_size, channels, height, width] to a list of PIL Images + return [tensor_to_pil(img_tensor, i) for i in range(img_tensor.shape[0])] + + +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 batched_pil_to_tensor(images): + # Takes a list of PIL images and returns a tensor of shape [batch_size, height, width, channels] + return torch.cat([pil_to_tensor(image) for image in images], dim=0) + + +def img2tensor(imgs, bgr2rgb=True, float32=True): + + def _totensor(img, bgr2rgb, float32): + if img.shape[2] == 3 and bgr2rgb: + if img.dtype == 'float64': + img = img.astype('float32') + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + img = torch.from_numpy(img.transpose(2, 0, 1)) + if float32: + img = img.float() + return img + + if isinstance(imgs, list): + return [_totensor(img, bgr2rgb, float32) for img in imgs] + else: + return _totensor(imgs, bgr2rgb, float32) + + +def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)): + + if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))): + raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}') + + if torch.is_tensor(tensor): + tensor = [tensor] + result = [] + for _tensor in tensor: + _tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max) + _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0]) + + n_dim = _tensor.dim() + if n_dim == 4: + img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy() + img_np = img_np.transpose(1, 2, 0) + if rgb2bgr: + img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) + elif n_dim == 3: + img_np = _tensor.numpy() + img_np = img_np.transpose(1, 2, 0) + if img_np.shape[2] == 1: # gray image + img_np = np.squeeze(img_np, axis=2) + else: + if rgb2bgr: + img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) + elif n_dim == 2: + img_np = _tensor.numpy() + else: + raise TypeError('Only support 4D, 3D or 2D tensor. ' f'But received with dimension: {n_dim}') + if out_type == np.uint8: + # Unlike MATLAB, numpy.unit8() WILL NOT round by default. + img_np = (img_np * 255.0).round() + img_np = img_np.astype(out_type) + result.append(img_np) + if len(result) == 1: + result = result[0] + return result + +def move_path(old_path, new_path): + if os.path.exists(old_path): + try: + models = os.listdir(old_path) + for model in models: + move_old_path = os.path.join(old_path, model) + move_new_path = os.path.join(new_path, model) + os.rename(move_old_path, move_new_path) + os.rmdir(old_path) + except Exception as e: + print(f"Error: {e}") + new_path = old_path + +def addLoggingLevel(levelName, levelNum, methodName=None): + if not methodName: + methodName = levelName.lower() + + def logForLevel(self, message, *args, **kwargs): + if self.isEnabledFor(levelNum): + self._log(levelNum, message, args, **kwargs) + + def logToRoot(message, *args, **kwargs): + logging.log(levelNum, message, *args, **kwargs) + + logging.addLevelName(levelNum, levelName) + setattr(logging, levelName, levelNum) + setattr(logging.getLoggerClass(), methodName, logForLevel) + setattr(logging, methodName, logToRoot) + +def get_image_md5hash(image: Image.Image): + md5hash = hashlib.md5(image.tobytes()) + return md5hash.hexdigest() + +def save_face_model(face: Face, filename: str) -> None: + try: + tensors = { + "bbox": torch.tensor(face["bbox"]), + "kps": torch.tensor(face["kps"]), + "det_score": torch.tensor(face["det_score"]), + "landmark_3d_68": torch.tensor(face["landmark_3d_68"]), + "pose": torch.tensor(face["pose"]), + "landmark_2d_106": torch.tensor(face["landmark_2d_106"]), + "embedding": torch.tensor(face["embedding"]), + "gender": torch.tensor(face["gender"]), + "age": torch.tensor(face["age"]), + } + save_file(tensors, filename) + print(f"Face model has been saved to '{filename}'") + except Exception as e: + print(f"Error: {e}") + +def load_face_model(filename: str): + face = {} + with safe_open(filename, framework="pt") as f: + for k in f.keys(): + face[k] = f.get_tensor(k).numpy() + return Face(face) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..d71b1e0 --- /dev/null +++ b/utils.py @@ -0,0 +1,481 @@ +import numpy as np +from PIL import Image, ImageFilter +import torch +import torch.nn.functional as F +from torchvision.transforms import GaussianBlur +import math + +if (not hasattr(Image, 'Resampling')): # For older versions of Pillow + Image.Resampling = Image + +BLUR_KERNEL_SIZE = 15 + +# PIL to Mask +def pil2mask(image): + image_np = np.array(image.convert("L")).astype(np.float32) / 255.0 + mask = torch.from_numpy(image_np) + return 1.0 - mask + +# Mask to PIL +def mask2pil(mask): + if mask.ndim > 2: + mask = mask.squeeze(0) + mask_np = mask.cpu().numpy().astype('uint8') + mask_pil = Image.fromarray(mask_np, mode="L") + return mask_pil + +# Tensor to PIL +def tensor2pil(image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +# PIL to Tensor +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +def tensor_to_pil(img_tensor, batch_index=0): + # Takes an image in a batch in the form of a tensor of shape [batch_size, channels, height, width] + # and returns an PIL Image with the corresponding mode deduced by the number of channels + + # Take the image in the batch given by batch_index + img_tensor = img_tensor[batch_index].unsqueeze(0) + i = 255. * img_tensor.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8).squeeze()) + return img + + +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