From 6029c8990676a0cdb26c75c39a46dc0743e03f04 Mon Sep 17 00:00:00 2001 From: GiusTex <112352961+GiusTex@users.noreply.github.com> Date: Sun, 17 Nov 2024 18:12:21 +0100 Subject: [PATCH] add new pad options --- nodes.py | 173 +++++++++++++++++++++++++++++++++---------------------- 1 file changed, 103 insertions(+), 70 deletions(-) diff --git a/nodes.py b/nodes.py index 1031d6b..3198bfb 100644 --- a/nodes.py +++ b/nodes.py @@ -1,12 +1,21 @@ import torch import os -from PIL import Image +from PIL import Image, ImageDraw from .utils import get_first_folder_list, tensor2pil, pil2tensor, diffuserOutpaintSamples, get_device_by_name, get_dtype_by_name, clearVram # Get the absolute path of various directories my_dir = os.path.dirname(os.path.abspath(__file__)) +def can_expand(source_width, source_height, target_width, target_height, alignment): + """Checks if the image can be expanded based on the alignment.""" + if alignment in ("Left", "Right") and source_width >= target_width: + return False + if alignment in ("Top", "Bottom") and source_height >= target_height: + return False + return True + + class PadImageForDiffusersOutpaint: _alignment_options = ["Middle", "Left", "Right", "Top", "Bottom"] @classmethod @@ -14,54 +23,73 @@ class PadImageForDiffusersOutpaint: return { "required": { "image": ("IMAGE",), - "width": ("INT", {"default": 720, "min": 320, "max": 1536, "tooltip": "The width used for the image."}), - "height": ("INT", {"default": 1280, "min": 320, "max": 1536, "tooltip": "The height used for the image."}), + "width": ("INT", {"default": 720, "tooltip": "The width used for the image."}), + "height": ("INT", {"default": 1280, "tooltip": "The height used for the image."}), "alignment": (s._alignment_options, {"tooltip": "Where the original image should be in the outpainted one"}), + "resize_image": (s._resize_option, {"tooltip": "Resize input image"}), + "custom_resize_image_percentage": ("INT", {"min": 1, "default": 50, "max": 100, "step": 1, "tooltip": "Custom resize (%)"}), + "mask_overlap_percentage": ("INT", {"min": 1, "default": 10, "max": 50, "step": 1, "tooltip": "Mask overlap (%)"}), + "overlap_left": ("BOOLEAN", {"default": True}), + "overlap_right": ("BOOLEAN", {"default": True}), + "overlap_top": ("BOOLEAN", {"default": True}), + "overlap_bottom": ("BOOLEAN", {"default": True}), }, } RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") RETURN_NAMES = ("IMAGE", "MASK", "diffuser_outpaint_cnet_image") - FUNCTION = "expand_image" + FUNCTION = "prepare_image_and_mask" CATEGORY = "DiffusersOutpaint" - def expand_image(self, image, width, height, alignment="Middle"): - - # Resize Image - def can_expand(source_width, source_height, target_width, target_height, alignment): - """Checks if the image can be expanded based on the alignment.""" - if alignment in ("Left", "Right") and source_width >= target_width: - return False - if alignment in ("Top", "Bottom") and source_height >= target_height: - return False - return True - + def prepare_image_and_mask(self, image, width, height, mask_overlap_percentage, resize_image, custom_resize_image_percentage, overlap_left, overlap_right, overlap_top, overlap_bottom, alignment="Middle"): im=tensor2pil(image) source=im.convert('RGB') + target_size = (width, height) - # Raise an error. - if source.width == width and source.height == height: - raise ValueError(f'Input image size is the same as target size, resize input image or change target size.') - + # Calculate the scaling factor to fit the image within the target size + scale_factor = min(target_size[0] / source.width, target_size[1] / source.height) + new_width = int(source.width * scale_factor) + new_height = int(source.height * scale_factor) + + # Resize the source image to fit within target size + source = source.resize((new_width, new_height), Image.LANCZOS) + # Initialize new_width and new_height new_width, new_height = source.width, source.height - # Upscale if source is smaller than target in both dimensions - if source.width < target_size[0] and source.height < target_size[1]: - scale_factor = min(target_size[0] / source.width, target_size[1] / source.height) - new_width = int(source.width * scale_factor) - new_height = int(source.height * scale_factor) - source = source.resize((new_width, new_height), Image.LANCZOS) + # Apply resize option using percentages + if resize_image == "Full": + resize_percentage = 100 + elif resize_image == "50%": + resize_percentage = 50 + elif resize_image == "33%": + resize_percentage = 33 + elif resize_image == "25%": + resize_percentage = 25 + else: # Custom + resize_percentage = custom_resize_image_percentage + + # Calculate new dimensions based on percentage + resize_factor = resize_percentage / 100 + new_width = int(source.width * resize_factor) + new_height = int(source.height * resize_factor) + + # Ensure minimum size of 64 pixels + new_width = max(new_width, 64) + new_height = max(new_height, 64) - if source.width > target_size[0] or source.height > target_size[1]: - scale_factor = min(target_size[0] / source.width, target_size[1] / source.height) - new_width = int(source.width * scale_factor) - new_height = int(source.height * scale_factor) - source = source.resize((new_width, new_height), Image.LANCZOS) + # Resize the image + source = source.resize((new_width, new_height), Image.LANCZOS) + + # Calculate the overlap in pixels based on the percentage + overlap_x = int(new_width * (mask_overlap_percentage / 100)) + overlap_y = int(new_height * (mask_overlap_percentage / 100)) + + # Ensure minimum overlap of 1 pixel + overlap_x = max(overlap_x, 1) + overlap_y = max(overlap_y, 1) - if not can_expand(source.width, source.height, target_size[0], target_size[1], alignment): - alignment = "Middle" # Calculate margins based on alignment if alignment == "Middle": margin_x = (target_size[0] - source.width) // 2 @@ -79,11 +107,17 @@ class PadImageForDiffusersOutpaint: margin_x = (target_size[0] - source.width) // 2 margin_y = target_size[1] - source.height + # Adjust margins to eliminate gaps + margin_x = max(0, min(margin_x, target_size[0] - new_width)) + margin_y = max(0, min(margin_y, target_size[1] - new_height)) + + # Create a new background image and paste the resized source image background = Image.new('RGB', target_size, (255, 255, 255)) background.paste(source, (margin_x, margin_y)) image=pil2tensor(background) #---------------------------------------------------- + # Create the mask d1, d2, d3, d4 = image.size() left, top, bottom, right = 0, 0, 0, 0 # Image @@ -92,51 +126,50 @@ class PadImageForDiffusersOutpaint: dtype=torch.float32, ) * 0.5 new_image[:, top:top + d2, left:left + d3, :] = image - #---------------------------------------------------- - # Mask coordinates - if alignment == "Middle": - margin_x = (width - new_width) // 2 - margin_y = (height - new_height) // 2 - elif alignment == "Left": - margin_x = 0 - margin_y = (height - new_height) // 2 - elif alignment == "Right": - margin_x = width - new_width - margin_y = (height - new_height) // 2 - elif alignment == "Top": - margin_x = (width - new_width) // 2 - margin_y = 0 - elif alignment == "Bottom": - margin_x = (width - new_width) // 2 - margin_y = height - new_height - - # Create mask as big as new img - mask = torch.ones( - (height, width), - dtype=torch.float32, - ) - # Create hole in mask - t = torch.zeros( - (new_height, new_width), - dtype=torch.float32 - ) - # Create holed mask - mask[margin_y:margin_y + new_height, - margin_x:margin_x + new_width - ] = t - #---------------------------------------------------- - # Prepare "cn_image" for diffusers outpaint + im=tensor2pil(new_image) pil_new_image=im.convert('RGB') - - pil_mask=tensor2pil(mask) + #---------------------------------------------------- + + # Create the mask + mask = Image.new('L', target_size, 255) + mask_draw = ImageDraw.Draw(mask) + #---------------------------------------------------- + # Calculate overlap areas + white_gaps_patch = 2 + + left_overlap = margin_x + overlap_x if overlap_left else margin_x + white_gaps_patch + right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width - white_gaps_patch + top_overlap = margin_y + overlap_y if overlap_top else margin_y + white_gaps_patch + bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height - white_gaps_patch + #---------------------------------------------------- + # Mask coordinates + if alignment == "Left": + left_overlap = margin_x + overlap_x if overlap_left else margin_x + elif alignment == "Right": + right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width + elif alignment == "Top": + top_overlap = margin_y + overlap_y if overlap_top else margin_y + elif alignment == "Bottom": + bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height + + # Draw the mask + mask_draw.rectangle([ + (left_overlap, top_overlap), + (right_overlap, bottom_overlap) + ], fill=0) + + tensor_mask=pil2tensor(mask) + #---------------------------------------------------- + if not can_expand(background.width, background.height, width, height, alignment): + alignment = "Middle" cnet_image = pil_new_image.copy() # copy background as cnet_image - cnet_image.paste(0, (0, 0), pil_mask) # paste mask over cnet_image, cropping it a bit + cnet_image.paste(0, (0, 0), mask) # paste mask over cnet_image, cropping it a bit tensor_cnet_image=pil2tensor(cnet_image) - return (new_image, mask, tensor_cnet_image,) + return (new_image, tensor_mask, tensor_cnet_image,) class LoadDiffusersOutpaintModels: