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", }