From bf4403a7ae51d2c860893cb3b76c92cda8c65647 Mon Sep 17 00:00:00 2001 From: Level Pixel Dev Date: Wed, 4 Jun 2025 05:38:36 +0600 Subject: [PATCH] Add Inpaint Crop and Inpaint Stitch --- README.md | 2 + __init__.py | 1 + nodes/inpaint/inpaint_crop_stitch_LP.py | 988 ++++++++++++++++++++++++ 3 files changed, 991 insertions(+) create mode 100644 nodes/inpaint/inpaint_crop_stitch_LP.py diff --git a/README.md b/README.md index 992c855..56638ec 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,8 @@ Load LoRA Tag node for ComfyUI/[comfyui_lora_tag_loader](https://github.com/badj Efficiency-nodes-comfyui/[efficiency-nodes-comfyui](https://github.com/jags111/efficiency-nodes-comfyui) - Thanks for Image Overlay! +ComfyUI-Inpaint-CropAndStitch/[ComfyUI-Inpaint-CropAndStitch](https://github.com/lquesada/ComfyUI-Inpaint-CropAndStitch) - Thanks for Inpaint Crop and Inpaint Stitch! + ## License Copyright (c) 2024-present [Level Pixel](https://github.com/LevelPixel) diff --git a/__init__.py b/__init__.py index b24040f..c85a080 100644 --- a/__init__.py +++ b/__init__.py @@ -30,6 +30,7 @@ init() node_list = [ "convert.convert_LP", "image.image_utils_LP", + "inpaint.inpaint_crop_stitch_LP", "io.folder_workers_LP", "io.image_loaders_LP", "io.image_outputs_LP", diff --git a/nodes/inpaint/inpaint_crop_stitch_LP.py b/nodes/inpaint/inpaint_crop_stitch_LP.py new file mode 100644 index 0000000..6686b31 --- /dev/null +++ b/nodes/inpaint/inpaint_crop_stitch_LP.py @@ -0,0 +1,988 @@ +import comfy.utils +import math +import nodes +import numpy as np +import torch +import torchvision.transforms.functional as F +from PIL import Image +from scipy.ndimage import gaussian_filter, grey_dilation, binary_closing, binary_fill_holes + + +def rescale_i(samples, width, height, algorithm: str): + samples = samples.movedim(-1, 1) + algorithm = getattr(Image, algorithm.upper()) # i.e. Image.BICUBIC + samples_pil: Image.Image = F.to_pil_image(samples[0].cpu()).resize((width, height), algorithm) + samples = F.to_tensor(samples_pil).unsqueeze(0) + samples = samples.movedim(1, -1) + return samples + + +def rescale_m(samples, width, height, algorithm: str): + samples = samples.unsqueeze(1) + algorithm = getattr(Image, algorithm.upper()) # i.e. Image.BICUBIC + samples_pil: Image.Image = F.to_pil_image(samples[0].cpu()).resize((width, height), algorithm) + samples = F.to_tensor(samples_pil).unsqueeze(0) + samples = samples.squeeze(1) + return samples + + +def preresize_imm(image, mask, optional_context_mask, downscale_algorithm, upscale_algorithm, preresize_mode, preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height): + current_width, current_height = image.shape[2], image.shape[1] # Image size [batch, height, width, channels] + + if preresize_mode == "ensure minimum resolution": + if current_width >= preresize_min_width and current_height >= preresize_min_height: + return image, mask, optional_context_mask + + scale_factor_min_width = preresize_min_width / current_width + scale_factor_min_height = preresize_min_height / current_height + + scale_factor = max(scale_factor_min_width, scale_factor_min_height) + + target_width = int(current_width * scale_factor) + target_height = int(current_height * scale_factor) + + image = rescale_i(image, target_width, target_height, upscale_algorithm) + mask = rescale_m(mask, target_width, target_height, 'bilinear') + optional_context_mask = rescale_m(optional_context_mask, target_width, target_height, 'bilinear') + + assert target_width >= preresize_min_width and target_height >= preresize_min_height, \ + f"Internal error: After resizing, target size {target_width}x{target_height} is smaller than min size {preresize_min_width}x{preresize_min_height}" + + elif preresize_mode == "ensure minimum and maximum resolution": + if preresize_min_width <= current_width <= preresize_max_width and preresize_min_height <= current_height <= preresize_max_height: + return image, mask, optional_context_mask + + scale_factor_min_width = preresize_min_width / current_width + scale_factor_min_height = preresize_min_height / current_height + scale_factor_min = max(scale_factor_min_width, scale_factor_min_height) + + scale_factor_max_width = preresize_max_width / current_width + scale_factor_max_height = preresize_max_height / current_height + scale_factor_max = min(scale_factor_max_width, scale_factor_max_height) + + if scale_factor_min > 1 and scale_factor_max < 1: + assert False, "Cannot meet both minimum and maximum resolution requirements with aspect ratio preservation." + + if scale_factor_min > 1: # We're upscaling to meet min resolution + scale_factor = scale_factor_min + rescale_algorithm = upscale_algorithm # Use upscale algorithm for min resolution + else: # We're downscaling to meet max resolution + scale_factor = scale_factor_max + rescale_algorithm = downscale_algorithm # Use downscale algorithm for max resolution + + target_width = int(current_width * scale_factor) + target_height = int(current_height * scale_factor) + + image = rescale_i(image, target_width, target_height, rescale_algorithm) + mask = rescale_m(mask, target_width, target_height, 'nearest') # Always nearest for efficiency + optional_context_mask = rescale_m(optional_context_mask, target_width, target_height, 'nearest') # Always nearest for efficiency + + assert preresize_min_width <= target_width <= preresize_max_width, \ + f"Internal error: Target width {target_width} is outside the range {preresize_min_width} - {preresize_max_width}" + assert preresize_min_height <= target_height <= preresize_max_height, \ + f"Internal error: Target height {target_height} is outside the range {preresize_min_height} - {preresize_max_height}" + + elif preresize_mode == "ensure maximum resolution": + if current_width <= preresize_max_width and current_height <= preresize_max_height: + return image, mask, optional_context_mask + + scale_factor_max_width = preresize_max_width / current_width + scale_factor_max_height = preresize_max_height / current_height + scale_factor_max = min(scale_factor_max_width, scale_factor_max_height) + + target_width = int(current_width * scale_factor_max) + target_height = int(current_height * scale_factor_max) + + image = rescale_i(image, target_width, target_height, downscale_algorithm) + mask = rescale_m(mask, target_width, target_height, 'nearest') # Always nearest for efficiency + optional_context_mask = rescale_m(optional_context_mask, target_width, target_height, 'nearest') # Always nearest for efficiency + + assert target_width <= preresize_max_width and target_height <= preresize_max_height, \ + f"Internal error: Target size {target_width}x{target_height} is greater than max size {preresize_max_width}x{preresize_max_height}" + + return image, mask, optional_context_mask + + +def fillholes_iterative_hipass_fill_m(samples): + thresholds = [1, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1] + + mask_np = samples.squeeze(0).cpu().numpy() + + for threshold in thresholds: + thresholded_mask = mask_np >= threshold + closed_mask = binary_closing(thresholded_mask, structure=np.ones((3, 3)), border_value=1) + filled_mask = binary_fill_holes(closed_mask) + mask_np = np.maximum(mask_np, np.where(filled_mask != 0, threshold, 0)) + + final_mask = torch.from_numpy(mask_np.astype(np.float32)).unsqueeze(0) + + return final_mask + + +def hipassfilter_m(samples, threshold): + filtered_mask = samples.clone() + filtered_mask[filtered_mask < threshold] = 0 + return filtered_mask + + +def expand_m(mask, pixels): + sigma = pixels / 4 + mask_np = mask.squeeze(0).cpu().numpy() + kernel_size = math.ceil(sigma * 1.5 + 1) + kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8) + dilated_mask = grey_dilation(mask_np, footprint=kernel) + dilated_mask = dilated_mask.astype(np.float32) + dilated_mask = torch.from_numpy(dilated_mask) + dilated_mask = torch.clamp(dilated_mask, 0.0, 1.0) + return dilated_mask.unsqueeze(0) + + +def invert_m(samples): + inverted_mask = samples.clone() + inverted_mask = 1.0 - inverted_mask + return inverted_mask + + +def blur_m(samples, pixels): + mask = samples.squeeze(0) + sigma = pixels / 4 + mask_np = mask.cpu().numpy() + blurred_mask = gaussian_filter(mask_np, sigma=sigma) + blurred_mask = torch.from_numpy(blurred_mask).float() + blurred_mask = torch.clamp(blurred_mask, 0.0, 1.0) + return blurred_mask.unsqueeze(0) + + +def extend_imm(image, mask, optional_context_mask, extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor): + B, H, W, C = image.shape + + new_H = int(H * (1.0 + extend_up_factor - 1.0 + extend_down_factor - 1.0)) + new_W = int(W * (1.0 + extend_left_factor - 1.0 + extend_right_factor - 1.0)) + + assert new_H >= 0, f"Error: Trying to crop too much, height ({new_H}) must be >= 0" + assert new_W >= 0, f"Error: Trying to crop too much, width ({new_W}) must be >= 0" + + expanded_image = torch.zeros(1, new_H, new_W, C, device=image.device) + expanded_mask = torch.ones(1, new_H, new_W, device=mask.device) + expanded_optional_context_mask = torch.zeros(1, new_H, new_W, device=optional_context_mask.device) + + up_padding = int(H * (extend_up_factor - 1.0)) + down_padding = new_H - H - up_padding + left_padding = int(W * (extend_left_factor - 1.0)) + right_padding = new_W - W - left_padding + + slice_target_up = max(0, up_padding) + slice_target_down = min(new_H, up_padding + H) + slice_target_left = max(0, left_padding) + slice_target_right = min(new_W, left_padding + W) + + slice_source_up = max(0, -up_padding) + slice_source_down = min(H, new_H - up_padding) + slice_source_left = max(0, -left_padding) + slice_source_right = min(W, new_W - left_padding) + + image = image.permute(0, 3, 1, 2) # [B, H, W, C] -> [B, C, H, W] + expanded_image = expanded_image.permute(0, 3, 1, 2) # [B, H, W, C] -> [B, C, H, W] + + expanded_image[:, :, slice_target_up:slice_target_down, slice_target_left:slice_target_right] = image[:, :, slice_source_up:slice_source_down, slice_source_left:slice_source_right] + if up_padding > 0: + expanded_image[:, :, :up_padding, slice_target_left:slice_target_right] = image[:, :, 0:1, slice_source_left:slice_source_right].repeat(1, 1, up_padding, 1) + if down_padding > 0: + expanded_image[:, :, -down_padding:, slice_target_left:slice_target_right] = image[:, :, -1:, slice_source_left:slice_source_right].repeat(1, 1, down_padding, 1) + if left_padding > 0: + expanded_image[:, :, :, :left_padding] = expanded_image[:, :, :, left_padding:left_padding+1].repeat(1, 1, 1, left_padding) + if right_padding > 0: + expanded_image[:, :, :, -right_padding:] = expanded_image[:, :, :, -right_padding-1:-right_padding].repeat(1, 1, 1, right_padding) + + expanded_mask[:, slice_target_up:slice_target_down, slice_target_left:slice_target_right] = mask[:, slice_source_up:slice_source_down, slice_source_left:slice_source_right] + expanded_optional_context_mask[:, slice_target_up:slice_target_down, slice_target_left:slice_target_right] = optional_context_mask[:, slice_source_up:slice_source_down, slice_source_left:slice_source_right] + + expanded_image = expanded_image.permute(0, 2, 3, 1) # [B, C, H, W] -> [B, H, W, C] + image = image.permute(0, 2, 3, 1) # [B, C, H, W] -> [B, H, W, C] + + return expanded_image, expanded_mask, expanded_optional_context_mask + + +def findcontextarea_m(mask): + mask_squeezed = mask[0] # Now shape is [H, W] + non_zero_indices = torch.nonzero(mask_squeezed) + + H, W = mask_squeezed.shape + + if non_zero_indices.numel() == 0: + x, y = -1, -1 + w, h = -1, -1 + else: + y = torch.min(non_zero_indices[:, 0]).item() + x = torch.min(non_zero_indices[:, 1]).item() + y_max = torch.max(non_zero_indices[:, 0]).item() + x_max = torch.max(non_zero_indices[:, 1]).item() + w = x_max - x + 1 # +1 to include the max index + h = y_max - y + 1 # +1 to include the max index + + context = mask[:, y:y+h, x:x+w] + return context, x, y, w, h + + +def growcontextarea_m(context, mask, x, y, w, h, extend_factor): + img_h, img_w = mask.shape[1], mask.shape[2] + + # Compute intended growth in each direction + grow_left = int(round(w * (extend_factor-1.0) / 2.0)) + grow_right = int(round(w * (extend_factor-1.0) / 2.0)) + grow_up = int(round(h * (extend_factor-1.0) / 2.0)) + grow_down = int(round(h * (extend_factor-1.0) / 2.0)) + + # Try to grow left, but clamp at 0 + new_x = x - grow_left + if new_x < 0: + new_x = 0 + + # Try to grow up, but clamp at 0 + new_y = y - grow_up + if new_y < 0: + new_y = 0 + + # Right edge + new_x2 = x + w + grow_right + if new_x2 > img_w: + new_x2 = img_w + + # Bottom edge + new_y2 = y + h + grow_down + if new_y2 > img_h: + new_y2 = img_h + + # New width and height + new_w = new_x2 - new_x + new_h = new_y2 - new_y + + # Extract the context + new_context = mask[:, new_y:new_y+new_h, new_x:new_x+new_w] + + if new_h < 0 or new_w < 0: + new_x = 0 + new_y = 0 + new_w = mask.shape[2] + new_h = mask.shape[1] + + return new_context, new_x, new_y, new_w, new_h + + +def combinecontextmask_m(context, mask, x, y, w, h, optional_context_mask): + _, x_opt, y_opt, w_opt, h_opt = findcontextarea_m(optional_context_mask) + if x == -1: + x, y, w, h = x_opt, y_opt, w_opt, h_opt + if x_opt == -1: + x_opt, y_opt, w_opt, h_opt = x, y, w, h + if x == -1: + return torch.zeros(1, 0, 0, device=mask.device), -1, -1, -1, -1 + new_x = min(x, x_opt) + new_y = min(y, y_opt) + new_x_max = max(x + w, x_opt + w_opt) + new_y_max = max(y + h, y_opt + h_opt) + new_w = new_x_max - new_x + new_h = new_y_max - new_y + combined_context = mask[:, new_y:new_y+new_h, new_x:new_x+new_w] + return combined_context, new_x, new_y, new_w, new_h + + +def pad_to_multiple(value, multiple): + return int(math.ceil(value / multiple) * multiple) + + +def crop_magic_im(image, mask, x, y, w, h, target_w, target_h, padding, downscale_algorithm, upscale_algorithm): + image = image.clone() + mask = mask.clone() + + # Ok this is the most complex function in this node. The one that does the magic after all the preparation done by the other nodes. + # Basically this function determines the right context area that encompasses the whole context area (mask+optional_context_mask), + # that is ideally within the bounds of the original image, and that has the right aspect ratio to match target width and height. + # It may grow the image if the aspect ratio wouldn't fit in the original image. + # It keeps track of that growing to then be able to crop the image in the stitch node. + # Finally, it crops the context area and resizes it to be exactly target_w and target_h. + # It keeps track of that resize to be able to revert it in the stitch node. + + # Check for invalid inputs + if target_w <= 0 or target_h <= 0 or w == 0 or h == 0: + return image, 0, 0, image.shape[2], image.shape[1], image, mask, 0, 0, image.shape[2], image.shape[1] + + # Step 1: Pad target dimensions to be multiples of padding + if padding != 0: + target_w = pad_to_multiple(target_w, padding) + target_h = pad_to_multiple(target_h, padding) + + # Step 2: Calculate target aspect ratio + target_aspect_ratio = target_w / target_h + + # Step 3: Grow current context area to meet the target aspect ratio + B, image_h, image_w, C = image.shape + context_aspect_ratio = w / h + if context_aspect_ratio < target_aspect_ratio: + # Grow width to meet aspect ratio + new_w = int(h * target_aspect_ratio) + new_h = h + new_x = x - (new_w - w) // 2 + new_y = y + + # Adjust new_x to keep within bounds + if new_x < 0: + shift = -new_x + if new_x + new_w + shift <= image_w: + new_x += shift + else: + overflow = (new_w - image_w) // 2 + new_x = -overflow + elif new_x + new_w > image_w: + overflow = new_x + new_w - image_w + if new_x - overflow >= 0: + new_x -= overflow + else: + overflow = (new_w - image_w) // 2 + new_x = -overflow + + else: + # Grow height to meet aspect ratio + new_w = w + new_h = int(w / target_aspect_ratio) + new_x = x + new_y = y - (new_h - h) // 2 + + # Adjust new_y to keep within bounds + if new_y < 0: + shift = -new_y + if new_y + new_h + shift <= image_h: + new_y += shift + else: + overflow = (new_h - image_h) // 2 + new_y = -overflow + elif new_y + new_h > image_h: + overflow = new_y + new_h - image_h + if new_y - overflow >= 0: + new_y -= overflow + else: + overflow = (new_h - image_h) // 2 + new_y = -overflow + + # Step 4: Grow the image to accommodate the new context area + up_padding, down_padding, left_padding, right_padding = 0, 0, 0, 0 + + expanded_image_w = image_w + expanded_image_h = image_h + + # Adjust width for left overflow (x < 0) and right overflow (x + w > image_w) + if new_x < 0: + left_padding = -new_x + expanded_image_w += left_padding + if new_x + new_w > image_w: + right_padding = (new_x + new_w - image_w) + expanded_image_w += right_padding + # Adjust height for top overflow (y < 0) and bottom overflow (y + h > image_h) + if new_y < 0: + up_padding = -new_y + expanded_image_h += up_padding + if new_y + new_h > image_h: + down_padding = (new_y + new_h - image_h) + expanded_image_h += down_padding + + # Step 5: Create the new image and mask + expanded_image = torch.zeros((image.shape[0], expanded_image_h, expanded_image_w, image.shape[3]), device=image.device) + expanded_mask = torch.ones((mask.shape[0], expanded_image_h, expanded_image_w), device=mask.device) + + # Reorder the tensors to match the required dimension format for padding + image = image.permute(0, 3, 1, 2) # [B, H, W, C] -> [B, C, H, W] + expanded_image = expanded_image.permute(0, 3, 1, 2) # [B, H, W, C] -> [B, C, H, W] + + # Ensure the expanded image has enough room to hold the padded version of the original image + expanded_image[:, :, up_padding:up_padding + image_h, left_padding:left_padding + image_w] = image + + # Fill the new extended areas with the edge values of the image + if up_padding > 0: + expanded_image[:, :, :up_padding, left_padding:left_padding + image_w] = image[:, :, 0:1, left_padding:left_padding + image_w].repeat(1, 1, up_padding, 1) + if down_padding > 0: + expanded_image[:, :, -down_padding:, left_padding:left_padding + image_w] = image[:, :, -1:, left_padding:left_padding + image_w].repeat(1, 1, down_padding, 1) + if left_padding > 0: + expanded_image[:, :, up_padding:up_padding + image_h, :left_padding] = expanded_image[:, :, up_padding:up_padding + image_h, left_padding:left_padding+1].repeat(1, 1, 1, left_padding) + if right_padding > 0: + expanded_image[:, :, up_padding:up_padding + image_h, -right_padding:] = expanded_image[:, :, up_padding:up_padding + image_h, -right_padding-1:-right_padding].repeat(1, 1, 1, right_padding) + + # Reorder the tensors back to [B, H, W, C] format + expanded_image = expanded_image.permute(0, 2, 3, 1) # [B, C, H, W] -> [B, H, W, C] + image = image.permute(0, 2, 3, 1) # [B, C, H, W] -> [B, H, W, C] + + # Same for the mask + expanded_mask[:, up_padding:up_padding + image_h, left_padding:left_padding + image_w] = mask + + # Record the cto values (canvas to original) + cto_x = left_padding + cto_y = up_padding + cto_w = image_w + cto_h = image_h + + # The final expanded image and mask + canvas_image = expanded_image + canvas_mask = expanded_mask + + # Step 6: Crop the image and mask around x, y, w, h + ctc_x = new_x+left_padding + ctc_y = new_y+up_padding + ctc_w = new_w + ctc_h = new_h + + # Crop the image and mask + cropped_image = canvas_image[:, ctc_y:ctc_y + ctc_h, ctc_x:ctc_x + ctc_w] + cropped_mask = canvas_mask[:, ctc_y:ctc_y + ctc_h, ctc_x:ctc_x + ctc_w] + + # Step 7: Resize image and mask to the target width and height + # Decide which algorithm to use based on the scaling direction + if target_w > ctc_w or target_h > ctc_h: # Upscaling + cropped_image = rescale_i(cropped_image, target_w, target_h, upscale_algorithm) + cropped_mask = rescale_m(cropped_mask, target_w, target_h, upscale_algorithm) + else: # Downscaling + cropped_image = rescale_i(cropped_image, target_w, target_h, downscale_algorithm) + cropped_mask = rescale_m(cropped_mask, target_w, target_h, downscale_algorithm) + + return canvas_image, cto_x, cto_y, cto_w, cto_h, cropped_image, cropped_mask, ctc_x, ctc_y, ctc_w, ctc_h + + +def stitch_magic_im(canvas_image, inpainted_image, mask, ctc_x, ctc_y, ctc_w, ctc_h, cto_x, cto_y, cto_w, cto_h, downscale_algorithm, upscale_algorithm): + canvas_image = canvas_image.clone() + inpainted_image = inpainted_image.clone() + mask = mask.clone() + + # Resize inpainted image and mask to match the context size + _, h, w, _ = inpainted_image.shape + if ctc_w > w or ctc_h > h: # Upscaling + resized_image = rescale_i(inpainted_image, ctc_w, ctc_h, upscale_algorithm) + resized_mask = rescale_m(mask, ctc_w, ctc_h, upscale_algorithm) + else: # Downscaling + resized_image = rescale_i(inpainted_image, ctc_w, ctc_h, downscale_algorithm) + resized_mask = rescale_m(mask, ctc_w, ctc_h, downscale_algorithm) + + # Clamp mask to [0, 1] and expand to match image channels + resized_mask = resized_mask.clamp(0, 1).unsqueeze(-1) # shape: [1, H, W, 1] + + # Extract the canvas region we're about to overwrite + canvas_crop = canvas_image[:, ctc_y:ctc_y + ctc_h, ctc_x:ctc_x + ctc_w] + + # Blend: new = mask * inpainted + (1 - mask) * canvas + blended = resized_mask * resized_image + (1.0 - resized_mask) * canvas_crop + + # Paste the blended region back onto the canvas + canvas_image[:, ctc_y:ctc_y + ctc_h, ctc_x:ctc_x + ctc_w] = blended + + # Final crop to get back the original image area + output_image = canvas_image[:, cto_y:cto_y + cto_h, cto_x:cto_x + cto_w] + + return output_image + + +class InpaintCrop: + """ + ComfyUI-LevelPixel + https://github.com/LevelPixel/ComfyUI-LevelPixel + + ComfyUI-Inpaint-CropAndStitch + https://github.com/lquesada/ComfyUI-Inpaint-CropAndStitch + + This node stitches the inpainted image without altering unmasked areas. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + # Required inputs + "image": ("IMAGE",), + + # Resize algorithms + "downscale_algorithm": (["nearest", "bilinear", "bicubic", "lanczos", "box", "hamming"], {"default": "bilinear"}), + "upscale_algorithm": (["nearest", "bilinear", "bicubic", "lanczos", "box", "hamming"], {"default": "bicubic"}), + + # Mask manipulation + "mask_fill_holes": ("BOOLEAN", {"default": True, "tooltip": "Mark as masked any areas fully enclosed by mask."}), + "mask_expand_pixels": ("INT", {"default": 0, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "Expand the mask by a certain amount of pixels before processing."}), + "mask_invert": ("BOOLEAN", {"default": False,"tooltip": "Invert mask so that anything masked will be kept."}), + "mask_blend_pixels": ("INT", {"default": 32, "min": 0, "max": 64, "step": 1, "tooltip": "How many pixels to blend into the original image."}), + "mask_hipass_filter": ("FLOAT", {"default": 0.1, "min": 0, "max": 1, "step": 0.01, "tooltip": "Ignore mask values lower than this value."}), + + # Context + "context_from_mask_extend_factor": ("FLOAT", {"default": 1.2, "min": 1.0, "max": 100.0, "step": 0.01, "tooltip": "Grow the context area from the mask by a certain factor in every direction. For example, 1.5 grabs extra 50% up, down, left, and right as context."}), + "output_padding": (["0", "8", "16", "32", "64", "128", "256", "512"], {"default": "32"}), + }, + "optional": { + # Optional inputs + "mask": ("MASK",), + "optional_context_mask": ("MASK",), + + # Custom Parameters + "preresize_parameters": ("PRERESIZEDATA", ), + "extend_factor_parameters": ("EXTENDFACTORDATA", ), + "cropped_image_target_size_parameters": ("TARGETSIZEDATA", ), + } + } + + FUNCTION = "inpaint_crop" + CATEGORY = "LevelPixel/Inpaint" + DESCRIPTION = "Crops an image around a mask for inpainting, the optional context mask defines an extra area to keep for the context." + + RETURN_TYPES = ("CROPDATA", "IMAGE", "MASK") + RETURN_NAMES = ("crop_data", "cropped_image", "cropped_mask") + + + def inpaint_crop(self, image, downscale_algorithm, upscale_algorithm, mask_hipass_filter, mask_fill_holes, mask_expand_pixels, mask_invert, mask_blend_pixels, context_from_mask_extend_factor, output_padding, mask=None, optional_context_mask=None, preresize_parameters=None, extend_factor_parameters=None, cropped_image_target_size_parameters=None): + + if preresize_parameters is not None: + preresize_mode, preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height = ( + preresize_parameters["preresize_mode"], + preresize_parameters["preresize_min_width"], + preresize_parameters["preresize_min_height"], + preresize_parameters["preresize_max_width"], + preresize_parameters["preresize_max_height"] + ) + preresize = True + else: + preresize = False + preresize_mode = None + preresize_min_width = None + preresize_min_height = None + preresize_max_width = None + preresize_max_height = None + + if extend_factor_parameters is not None: + extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor = ( + extend_factor_parameters["extend_up_factor"], + extend_factor_parameters["extend_down_factor"], + extend_factor_parameters["extend_left_factor"], + extend_factor_parameters["extend_right_factor"] + ) + extend_for_outpainting = True + else: + extend_for_outpainting = False + extend_up_factor = 1.0 + extend_down_factor = 1.0 + extend_left_factor = 1.0 + extend_right_factor = 1.0 + + if cropped_image_target_size_parameters is not None: + target_width, target_height, aspect_ratio_limit = ( + cropped_image_target_size_parameters["target_width"], + cropped_image_target_size_parameters["target_height"], + cropped_image_target_size_parameters["aspect_ratio_limit"] + ) + output_resize_to_target_size = True + else: + output_resize_to_target_size = False + target_width = None + target_height = None + aspect_ratio_limit = None + + image = image.clone() + if mask is not None: + mask = mask.clone() + if optional_context_mask is not None: + optional_context_mask = optional_context_mask.clone() + + output_padding = int(output_padding) + + # Check that some parameters make sense + if preresize and preresize_mode == "ensure minimum and maximum resolution": + assert preresize_max_width >= preresize_min_width, "Preresize maximum width must be greater than or equal to minimum width" + assert preresize_max_height >= preresize_min_height, "Preresize maximum height must be greater than or equal to minimum height" + + if image.shape[0] > 1: + assert output_resize_to_target_size, "cropped_image_target_size_parameters must be enabled when input is a batch of images, given all images in the batch output have to be the same size" + assert aspect_ratio_limit == 1, "aspect_ratio_limit must be equal to 1 when processing a batch of images" + + # When a LoadImage node passes a mask without user editing, it may be the wrong shape. + # Detect and fix that to avoid shape mismatch errors. + if mask is not None and (image.shape[0] == 1 or mask.shape[0] == 1 or mask.shape[0] == image.shape[0]): + if mask.shape[1] != image.shape[1] or mask.shape[2] != image.shape[2]: + if torch.count_nonzero(mask) == 0: + mask = torch.zeros((mask.shape[0], image.shape[1], image.shape[2]), device=image.device, dtype=image.dtype) + + if optional_context_mask is not None and (image.shape[0] == 1 or optional_context_mask.shape[0] == 1 or optional_context_mask.shape[0] == image.shape[0]): + if optional_context_mask.shape[1] != image.shape[1] or optional_context_mask.shape[2] != image.shape[2]: + if torch.count_nonzero(optional_context_mask) == 0: + optional_context_mask = torch.zeros((optional_context_mask.shape[0], image.shape[1], image.shape[2]), device=image.device, dtype=image.dtype) + + # If no mask is provided, create one with the shape of the image + if mask is None: + mask = torch.zeros_like(image[:, :, :, 0]) + + # If there is only one image for many masks, replicate it for all masks + if mask.shape[0] > 1 and image.shape[0] == 1: + assert image.dim() == 4, f"Expected 4D BHWC image tensor, got {image.shape}" + image = image.expand(mask.shape[0], -1, -1, -1).clone() + + # If there is only one mask for many images, replicate it for all images + if image.shape[0] > 1 and mask.shape[0] == 1: + assert mask.dim() == 3, f"Expected 3D BHW mask tensor, got {mask.shape}" + mask = mask.expand(image.shape[0], -1, -1).clone() + + # If no optional_context_mask is provided, create one with the shape of the image + if optional_context_mask is None: + optional_context_mask = torch.zeros_like(image[:, :, :, 0]) + + # If there is only one optional_context_mask for many images, replicate it for all images + if image.shape[0] > 1 and optional_context_mask.shape[0] == 1: + assert optional_context_mask.dim() == 3, f"Expected 3D BHW optional_context_mask tensor, got {optional_context_mask.shape}" + optional_context_mask = optional_context_mask.expand(image.shape[0], -1, -1).clone() + + # Validate data + assert image.ndimension() == 4, f"Expected 4 dimensions for image, got {image.ndimension()}" + assert mask.ndimension() == 3, f"Expected 3 dimensions for mask, got {mask.ndimension()}" + assert optional_context_mask.ndimension() == 3, f"Expected 3 dimensions for optional_context_mask, got {optional_context_mask.ndimension()}" + assert mask.shape[1:] == image.shape[1:3], f"Mask dimensions do not match image dimensions. Expected {image.shape[1:3]}, got {mask.shape[1:]}" + assert optional_context_mask.shape[1:] == image.shape[1:3], f"optional_context_mask dimensions do not match image dimensions. Expected {image.shape[1:3]}, got {optional_context_mask.shape[1:]}" + assert mask.shape[0] == image.shape[0], f"Mask batch does not match image batch. Expected {image.shape[0]}, got {mask.shape[0]}" + assert optional_context_mask.shape[0] == image.shape[0], f"Optional context mask batch does not match image batch. Expected {image.shape[0]}, got {optional_context_mask.shape[0]}" + + # Run for each image separately + result_stitcher = { + 'downscale_algorithm': downscale_algorithm, + 'upscale_algorithm': upscale_algorithm, + 'blend_pixels': mask_blend_pixels, + 'canvas_to_orig_x': [], + 'canvas_to_orig_y': [], + 'canvas_to_orig_w': [], + 'canvas_to_orig_h': [], + 'canvas_image': [], + 'cropped_to_canvas_x': [], + 'cropped_to_canvas_y': [], + 'cropped_to_canvas_w': [], + 'cropped_to_canvas_h': [], + 'cropped_mask_for_blend': [], + } + + result_image = [] + result_mask = [] + + batch_size = image.shape[0] + for b in range(batch_size): + one_image = image[b].unsqueeze(0) + one_mask = mask[b].unsqueeze(0) + one_optional_context_mask = optional_context_mask[b].unsqueeze(0) + + outputs = self.inpaint_crop_single_image( + one_image, downscale_algorithm, upscale_algorithm, preresize, preresize_mode, + preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height, + extend_for_outpainting, extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor, + mask_hipass_filter, mask_fill_holes, mask_expand_pixels, mask_invert, mask_blend_pixels, + context_from_mask_extend_factor, output_resize_to_target_size, target_width, target_height, + output_padding, one_mask, one_optional_context_mask) + + stitcher, cropped_image, cropped_mask = outputs[:3] + for key in ['canvas_to_orig_x', 'canvas_to_orig_y', 'canvas_to_orig_w', 'canvas_to_orig_h', 'canvas_image', 'cropped_to_canvas_x', 'cropped_to_canvas_y', 'cropped_to_canvas_w', 'cropped_to_canvas_h', 'cropped_mask_for_blend']: + result_stitcher[key].append(stitcher[key]) + + cropped_image = cropped_image.clone().squeeze(0) + result_image.append(cropped_image) + cropped_mask = cropped_mask.clone().squeeze(0) + result_mask.append(cropped_mask) + + result_image = torch.stack(result_image, dim=0) + result_mask = torch.stack(result_mask, dim=0) + + return result_stitcher, result_image, result_mask, + + + def inpaint_crop_single_image(self, image, downscale_algorithm, upscale_algorithm, preresize, preresize_mode, preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height, extend_for_outpainting, extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor, mask_hipass_filter, mask_fill_holes, mask_expand_pixels, mask_invert, mask_blend_pixels, context_from_mask_extend_factor, output_resize_to_target_size, output_target_width, output_target_height, output_padding, mask, optional_context_mask): + if preresize: + image, mask, optional_context_mask = preresize_imm(image, mask, optional_context_mask, downscale_algorithm, upscale_algorithm, preresize_mode, preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height) + + if mask_fill_holes: + mask = fillholes_iterative_hipass_fill_m(mask) + + if mask_expand_pixels > 0: + mask = expand_m(mask, mask_expand_pixels) + + if mask_invert: + mask = invert_m(mask) + + if mask_blend_pixels > 0: + mask = expand_m(mask, mask_blend_pixels) + mask = blur_m(mask, mask_blend_pixels*0.5) + + if mask_hipass_filter >= 0.01: + mask = hipassfilter_m(mask, mask_hipass_filter) + optional_context_mask = hipassfilter_m(optional_context_mask, mask_hipass_filter) + + if extend_for_outpainting: + image, mask, optional_context_mask = extend_imm(image, mask, optional_context_mask, extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor) + + context, x, y, w, h = findcontextarea_m(mask) + # If no mask, mask everything for some inpainting. + if x == -1 or w == -1 or h == -1 or y == -1: + x, y, w, h = 0, 0, image.shape[2], image.shape[1] + context = mask[:, y:y+h, x:x+w] + + if context_from_mask_extend_factor >= 1.01: + context, x, y, w, h = growcontextarea_m(context, mask, x, y, w, h, context_from_mask_extend_factor) + # If no mask, mask everything for some inpainting. + if x == -1 or w == -1 or h == -1 or y == -1: + x, y, w, h = 0, 0, image.shape[2], image.shape[1] + context = mask[:, y:y+h, x:x+w] + + context, x, y, w, h = combinecontextmask_m(context, mask, x, y, w, h, optional_context_mask) + # If no mask, mask everything for some inpainting. + if x == -1 or w == -1 or h == -1 or y == -1: + x, y, w, h = 0, 0, image.shape[2], image.shape[1] + context = mask[:, y:y+h, x:x+w] + + if not output_resize_to_target_size: + canvas_image, cto_x, cto_y, cto_w, cto_h, cropped_image, cropped_mask, ctc_x, ctc_y, ctc_w, ctc_h = crop_magic_im(image, mask, x, y, w, h, w, h, output_padding, downscale_algorithm, upscale_algorithm) + else: # if output_resize_to_target_size: + canvas_image, cto_x, cto_y, cto_w, cto_h, cropped_image, cropped_mask, ctc_x, ctc_y, ctc_w, ctc_h = crop_magic_im(image, mask, x, y, w, h, output_target_width, output_target_height, output_padding, downscale_algorithm, upscale_algorithm) + + # For blending, grow the mask even further and make it blurrier. + cropped_mask_blend = cropped_mask.clone() + if mask_blend_pixels > 0: + cropped_mask_blend = blur_m(cropped_mask_blend, mask_blend_pixels*0.5) + + stitcher = { + 'canvas_to_orig_x': cto_x, + 'canvas_to_orig_y': cto_y, + 'canvas_to_orig_w': cto_w, + 'canvas_to_orig_h': cto_h, + 'canvas_image': canvas_image, + 'cropped_to_canvas_x': ctc_x, + 'cropped_to_canvas_y': ctc_y, + 'cropped_to_canvas_w': ctc_w, + 'cropped_to_canvas_h': ctc_h, + 'cropped_mask_for_blend': cropped_mask_blend, + } + + return stitcher, cropped_image, cropped_mask + + + +class InpaintStitch: + """ + ComfyUI-LevelPixel + https://github.com/LevelPixel/ComfyUI-LevelPixel + + ComfyUI-Inpaint-CropAndStitch + https://github.com/lquesada/ComfyUI-Inpaint-CropAndStitch + + This node stitches the inpainted image without altering unmasked areas. + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "crop_data": ("CROPDATA",), + "inpainted_image": ("IMAGE",), + } + } + + CATEGORY = "LevelPixel/Inpaint" + DESCRIPTION = "Stitches an image cropped with Inpaint Crop back into the original image" + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "inpaint_stitch" + + + def inpaint_stitch(self, crop_data, inpainted_image): + inpainted_image = inpainted_image.clone() + results = [] + + batch_size = inpainted_image.shape[0] + assert len(crop_data['cropped_to_canvas_x']) == batch_size or len(crop_data['cropped_to_canvas_x']) == 1, "Stitch batch size doesn't match image batch size" + override = False + if len(crop_data['cropped_to_canvas_x']) != batch_size and len(crop_data['cropped_to_canvas_x']) == 1: + override = True + for b in range(batch_size): + one_image = inpainted_image[b] + one_stitcher = {} + for key in ['downscale_algorithm', 'upscale_algorithm', 'blend_pixels']: + one_stitcher[key] = crop_data[key] + for key in ['canvas_to_orig_x', 'canvas_to_orig_y', 'canvas_to_orig_w', 'canvas_to_orig_h', 'canvas_image', 'cropped_to_canvas_x', 'cropped_to_canvas_y', 'cropped_to_canvas_w', 'cropped_to_canvas_h', 'cropped_mask_for_blend']: + if override: # One stitcher for many images, always read 0. + one_stitcher[key] = crop_data[key][0] + else: + one_stitcher[key] = crop_data[key][b] + one_image = one_image.unsqueeze(0) + one_image, = self.inpaint_stitch_single_image(one_stitcher, one_image) + one_image = one_image.squeeze(0) + one_image = one_image.clone() + results.append(one_image) + + result_batch = torch.stack(results, dim=0) + + return (result_batch,) + + def inpaint_stitch_single_image(self, stitcher, inpainted_image): + downscale_algorithm = stitcher['downscale_algorithm'] + upscale_algorithm = stitcher['upscale_algorithm'] + canvas_image = stitcher['canvas_image'] + + ctc_x = stitcher['cropped_to_canvas_x'] + ctc_y = stitcher['cropped_to_canvas_y'] + ctc_w = stitcher['cropped_to_canvas_w'] + ctc_h = stitcher['cropped_to_canvas_h'] + + cto_x = stitcher['canvas_to_orig_x'] + cto_y = stitcher['canvas_to_orig_y'] + cto_w = stitcher['canvas_to_orig_w'] + cto_h = stitcher['canvas_to_orig_h'] + + mask = stitcher['cropped_mask_for_blend'] # shape: [1, H, W] + + output_image = stitch_magic_im(canvas_image, inpainted_image, mask, ctc_x, ctc_y, ctc_w, ctc_h, cto_x, cto_y, cto_w, cto_h, downscale_algorithm, upscale_algorithm) + + return (output_image,) + + +def compute_target_size(width, height, target_resolution, aspect_ratio_limit=1.7): + ratio = max(1 / aspect_ratio_limit, min(width / height, aspect_ratio_limit)) + height_new = target_resolution * 2 / (ratio + 1) + width_new = ratio * height_new + target_size = { + "target_height": int(round(height_new)), + "target_width": int(round(width_new)), + } + + return target_size + +def calculate_target_size(image, mask, target_resolution, aspect_ratio_limit=1.7): + B, H, W, C = image.shape + mask = mask.round() + + for b in range(B): + rows = torch.any(mask[min(b, mask.shape[0]-1)] > 0, dim=1) + cols = torch.any(mask[min(b, mask.shape[0]-1)] > 0, dim=0) + + row_indices = torch.where(rows)[0] + col_indices = torch.where(cols)[0] + + if row_indices.numel() == 0 or col_indices.numel() == 0: + width, height = W, H + else: + y_min, y_max = row_indices[[0, -1]] + x_min, x_max = col_indices[[0, -1]] + width = (x_max - x_min + 1).item() + height = (y_max - y_min + 1).item() + + return compute_target_size(width, height, target_resolution, aspect_ratio_limit) + +class CalculateTargetSize: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", ), + "mask": ("MASK", ), + "target_size": ("INT", {"default": 1024, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 1}), + "aspect_ratio_limit": ("FLOAT", {"default": 1.7, "min": 1, "max": 100, "step": 0.01}), + }, + } + + RETURN_TYPES = ("INT", "INT",) + RETURN_NAMES = ("height", "width",) + FUNCTION = "calculate_target_size" + CATEGORY = "LevelPixel/Inpaint" + def calculate_target_size(self, image, mask, target_size=1024, aspect_ratio_limit=2): + target_size = calculate_target_size(image, mask, target_size, aspect_ratio_limit) + target_height, target_width = (target_size["target_height"], target_size["target_width"]) + + return (target_height, target_width, ) + +class CroppedImageTargetSizeParameters: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", ), + "mask": ("MASK", ), + "target_size": ("INT", {"default": 1024, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 1}), + "aspect_ratio_limit": ("FLOAT", {"default": 1.7, "min": 1, "max": 100, "step": 0.01}), + }, + } + + RETURN_TYPES = ("TARGETSIZEDATA",) + RETURN_NAMES = ("target_size",) + FUNCTION = "cropped_image_target_size_parameters" + CATEGORY = "LevelPixel/Inpaint" + def cropped_image_target_size_parameters(self, image, mask, target_size=1024, aspect_ratio_limit=2): + target_size = calculate_target_size(image, mask, target_size, aspect_ratio_limit) + target_size["aspect_ratio_limit"] = aspect_ratio_limit + print(target_size) + return (target_size, ) + +class PreresizeParameters: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "preresize_mode": (["ensure minimum resolution", "ensure maximum resolution", "ensure minimum and maximum resolution"], {"default": "ensure minimum resolution"}), + "preresize_min_width": ("INT", {"default": 1024, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1}), + "preresize_min_height": ("INT", {"default": 1024, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1}), + "preresize_max_width": ("INT", {"default": nodes.MAX_RESOLUTION, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1}), + "preresize_max_height": ("INT", {"default": nodes.MAX_RESOLUTION, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1}), + }, + } + + RETURN_TYPES = ("PRERESIZEDATA",) + RETURN_NAMES = ("preresize_data",) + FUNCTION = "preresize_parameters" + CATEGORY = "LevelPixel/Inpaint" + def preresize_parameters(self, preresize_mode, preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height): + preresize_data = { + "preresize_mode": preresize_mode, + "preresize_min_width": preresize_min_width, + "preresize_min_height": preresize_min_height, + "preresize_max_width": preresize_max_width, + "preresize_max_height": preresize_max_height + } + return (preresize_data, ) + +class ExtendFactorParameters: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "extend_up_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}), + "extend_down_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}), + "extend_left_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}), + "extend_right_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ("EXTENDFACTORDATA",) + RETURN_NAMES = ("extend_factor_data",) + FUNCTION = "extend_factor_parameters" + CATEGORY = "LevelPixel/Inpaint" + def extend_factor_parameters(self, extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor): + extend_factor_data = { + "extend_up_factor": extend_up_factor, + "extend_down_factor": extend_down_factor, + "extend_left_factor": extend_left_factor, + "extend_right_factor": extend_right_factor + } + return (extend_factor_data, ) + + +NODE_CLASS_MAPPINGS = { + "InpaintCrop|LP": InpaintCrop, + "InpaintStitch|LP": InpaintStitch, + "CroppedImageTargetSizeParameters|LP": CroppedImageTargetSizeParameters, + "ExtendFactorParameters|LP": ExtendFactorParameters, + "PreresizeParameters|LP": PreresizeParameters, + "CalculateTargetSize|LP": CalculateTargetSize +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "InpaintCrop|LP": "Inpaint Crop [LP]", + "InpaintStitch|LP": "Inpaint Stitch [LP]", + "CroppedImageTargetSizeParameters|LP": "Cropped Image Target Size Parameters [LP]", + "ExtendFactorParameters|LP": "Extend Factor Parameters [LP]", + "PreresizeParameters|LP": "Preresize Parameters [LP]", + "CalculateTargetSize|LP": "Calculate Target Size [LP]" +} \ No newline at end of file