From f09fd46b2ac6213bbd502499ac7d0c71c15c9843 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Tue, 15 Jul 2025 02:40:30 -0700 Subject: [PATCH] Add files via upload --- AILab_ImageMaskTools.py | 73 ++++++++++++++------ AILab_InpaintTools.py | 146 +++++++++++++++++++++++++++++++++++++++ locales/en/nodeDefs.json | 16 +++++ locales/fr/nodeDefs.json | 16 +++++ locales/ja/nodeDefs.json | 16 +++++ locales/ko/nodeDefs.json | 16 +++++ locales/ru/nodeDefs.json | 18 ++++- locales/zh/nodeDefs.json | 16 +++++ 8 files changed, 293 insertions(+), 24 deletions(-) create mode 100644 AILab_InpaintTools.py diff --git a/AILab_ImageMaskTools.py b/AILab_ImageMaskTools.py index a4cd54f..88c3918 100644 --- a/AILab_ImageMaskTools.py +++ b/AILab_ImageMaskTools.py @@ -1,4 +1,4 @@ -# ComfyUI-RMBG v2.5.0 +# ComfyUI-RMBG v2.6.0 # # This node facilitates background removal using various models, including RMBG-2.0, INSPYRENET, BEN, BEN2, and BIREFNET-HR. # It utilizes advanced deep learning techniques to process images and generate accurate masks for background removal. @@ -31,7 +31,8 @@ # # 5. Input Nodes: # - ColorInput: A node for inputting colors in various formats. - +# +# License: GPL-3.0 # These nodes are crafted to streamline common image and mask operations within ComfyUI workflows. import os @@ -587,6 +588,8 @@ class AILab_MaskCombiner: # Image loader node class AILab_LoadImage: + upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] + @classmethod def INPUT_TYPES(cls): input_dir = folder_paths.get_input_directory() @@ -596,6 +599,7 @@ class AILab_LoadImage: "required": { "image": (sorted(files) or [""], {"image_upload": True}), "mask_channel": (["alpha", "red", "green", "blue"], {"default": "alpha", "tooltip": "Select channel to extract mask from"}), + "upscale_method": (cls.upscale_methods, {"default": "lanczos", "tooltip": "Method used for resizing the image"}), "scale_by": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 8.0, "step": 0.01, "tooltip": "Scale image by this factor (ignored if size > 0)"}), "resize_mode": (["longest_side", "shortest_side", "width", "height"], {"default": "longest_side", "tooltip": "Choose how to resize the image"}), "size": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1, "tooltip": "Target size for the selected resize mode (0 = keep original size)"}), @@ -611,14 +615,28 @@ class AILab_LoadImage: FUNCTION = "load_image" OUTPUT_NODE = False - def load_image(self, image, mask_channel="alpha", scale_by=1.0, resize_mode="longest_side", size=0, extra_pnginfo=None): + def load_image(self, image, mask_channel="alpha", upscale_method="lanczos", scale_by=1.0, resize_mode="longest_side", size=0, extra_pnginfo=None): try: image_path = folder_paths.get_annotated_filepath(image) img = Image.open(image_path) orig_width, orig_height = img.size - # Image resizing logic + resampling_map = { + "nearest-exact": Image.NEAREST, + "bilinear": Image.BILINEAR, + "area": Image.BOX, + "bicubic": Image.BICUBIC, + "lanczos": Image.LANCZOS + } + resampling = resampling_map.get(upscale_method, Image.LANCZOS) + + has_alpha = 'A' in img.getbands() + if has_alpha and mask_channel == "alpha": + original_alpha = img.getchannel('A') + + img_rgb = img.convert('RGB') + if size > 0: if resize_mode == "longest_side": if orig_width >= orig_height: @@ -627,7 +645,7 @@ class AILab_LoadImage: else: new_height = size new_width = int(orig_width * (size / orig_height)) - img = img.resize((new_width, new_height), Image.LANCZOS) + img_rgb = img_rgb.resize((new_width, new_height), resampling) elif resize_mode == "shortest_side": if orig_width <= orig_height: new_width = size @@ -635,49 +653,58 @@ class AILab_LoadImage: else: new_height = size new_width = int(orig_width * (size / orig_height)) - img = img.resize((new_width, new_height), Image.LANCZOS) + img_rgb = img_rgb.resize((new_width, new_height), resampling) elif resize_mode == "width": new_width = size new_height = int(orig_height * (size / orig_width)) - img = img.resize((new_width, new_height), Image.LANCZOS) + img_rgb = img_rgb.resize((new_width, new_height), resampling) elif resize_mode == "height": new_height = size new_width = int(orig_width * (size / orig_height)) - img = img.resize((new_width, new_height), Image.LANCZOS) + img_rgb = img_rgb.resize((new_width, new_height), resampling) elif scale_by != 1.0: new_width = int(orig_width * scale_by) new_height = int(orig_height * scale_by) - img = img.resize((new_width, new_height), Image.LANCZOS) + img_rgb = img_rgb.resize((new_width, new_height), resampling) - width, height = img.size + width, height = img_rgb.size + + mask = None + if mask_channel == "alpha" and has_alpha: + if (size > 0 or scale_by != 1.0) and 'original_alpha' in locals(): + mask_img = original_alpha.resize((width, height), resampling) + mask = np.array(mask_img).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) output_images = [] output_masks = [] - for i in ImageSequence.Iterator(img): + + for i in ImageSequence.Iterator(img_rgb): i = ImageOps.exif_transpose(i) if i.mode == 'I': i = i.point(lambda i: i * (1 / 255)) - image = i.convert("RGB") - image = np.array(image).astype(np.float32) / 255.0 + + if i.mode != 'RGB': + i = i.convert('RGB') + + image = np.array(i).astype(np.float32) / 255.0 image = torch.from_numpy(image)[None,] - if mask_channel == "alpha" and 'A' in i.getbands(): - mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 - mask = 1. - torch.from_numpy(mask) + if mask is not None: + output_masks.append(mask.unsqueeze(0)) elif mask_channel == "red" and 'R' in i.getbands(): mask = np.array(i.getchannel('R')).astype(np.float32) / 255.0 - mask = torch.from_numpy(mask) + output_masks.append(torch.from_numpy(mask).unsqueeze(0)) elif mask_channel == "green" and 'G' in i.getbands(): mask = np.array(i.getchannel('G')).astype(np.float32) / 255.0 - mask = torch.from_numpy(mask) + output_masks.append(torch.from_numpy(mask).unsqueeze(0)) elif mask_channel == "blue" and 'B' in i.getbands(): mask = np.array(i.getchannel('B')).astype(np.float32) / 255.0 - mask = torch.from_numpy(mask) + output_masks.append(torch.from_numpy(mask).unsqueeze(0)) else: - mask = torch.ones((height, width), dtype=torch.float32, device="cpu") + output_masks.append(torch.ones((1, height, width), dtype=torch.float32, device="cpu")) output_images.append(image) - output_masks.append(mask.unsqueeze(0)) if len(output_images) > 1: output_image = torch.cat(output_images, dim=0) @@ -700,7 +727,7 @@ class AILab_LoadImage: return (empty_image, empty_mask, empty_mask_image, 64, 64) @classmethod - def IS_CHANGED(cls, image, mask_channel="alpha", scale_by=1.0, resize_mode="longest_side", size=0, extra_pnginfo=None): + def IS_CHANGED(cls, image, mask_channel="alpha", upscale_method="lanczos", scale_by=1.0, resize_mode="longest_side", size=0, extra_pnginfo=None): image_path = folder_paths.get_annotated_filepath(image) m = hashlib.sha256() with open(image_path, 'rb') as f: @@ -708,7 +735,7 @@ class AILab_LoadImage: return m.digest().hex() @classmethod - def VALIDATE_INPUTS(cls, image, mask_channel="alpha", scale_by=1.0, resize_mode="longest_side", size=0, extra_pnginfo=None): + def VALIDATE_INPUTS(cls, image, mask_channel="alpha", upscale_method="lanczos", scale_by=1.0, resize_mode="longest_side", size=0, extra_pnginfo=None): if not folder_paths.exists_annotated_filepath(image): return f"Invalid image file: {image}" diff --git a/AILab_InpaintTools.py b/AILab_InpaintTools.py new file mode 100644 index 0000000..996d303 --- /dev/null +++ b/AILab_InpaintTools.py @@ -0,0 +1,146 @@ +# ComfyUI-RMBG v2.6.0 +# +# AILab Inpaint Tools +# A collection of specialized nodes for inpainting tasks in ComfyUI. +# Features a set of utilities for mask processing, latent conditioning, and inpainting workflows. +# +# 1. Inpaint Nodes: +# - AILab_ReferenceLatentMask: A node for inpainting tasks with the Flux Kontext model, using a reference latent and mask for precise region conditioning +# +# License: GPL-3.0 +# These nodes are crafted to streamline common image and mask operations within ComfyUI workflows. + +import torch +import node_helpers + +def expand_mask(mask, expand_amount): + if expand_amount == 0: + return mask + + import torch.nn.functional as F + + binary_mask = (mask > 0.5).float() + kernel_size = abs(expand_amount) * 2 + 1 + kernel_size = max(3, kernel_size) + + kernel = torch.ones(1, 1, kernel_size, kernel_size, device=mask.device) + + if expand_amount > 0: + expanded = F.conv2d( + binary_mask.reshape(-1, 1, mask.shape[-2], mask.shape[-1]), + kernel, + padding=kernel_size // 2 + ) + result = (expanded > 0).float() + else: + eroded = F.conv2d( + binary_mask.reshape(-1, 1, mask.shape[-2], mask.shape[-1]), + kernel, + padding=kernel_size // 2 + ) + result = (eroded >= kernel_size * kernel_size).float() + + if len(mask.shape) == 3: + result = result.squeeze(1) + + return result + + +def blur_mask(mask, blur_amount): + if blur_amount == 0: + return mask + + import torch.nn.functional as F + import math + + x = mask.reshape(-1, 1, mask.shape[-2], mask.shape[-1]) + kernel_size = max(3, math.ceil(blur_amount * 3) * 2 + 1) + + sigma = blur_amount + half_kernel = kernel_size // 2 + grid = torch.arange(-half_kernel, half_kernel + 1, device=mask.device).float() + + gaussian = torch.exp(-0.5 * (grid / sigma) ** 2) + gaussian = gaussian / gaussian.sum() + + gaussian_x = gaussian.view(1, 1, 1, kernel_size) + gaussian_y = gaussian.view(1, 1, kernel_size, 1) + + blurred = F.conv2d(x, gaussian_x, padding=(0, half_kernel)) + blurred = F.conv2d(blurred, gaussian_y, padding=(half_kernel, 0)) + + if len(mask.shape) == 3: + blurred = blurred.squeeze(1) + + return blurred + +class AILab_ReferenceLatentMask: + @classmethod + def INPUT_TYPES(cls): + tooltips = { + "conditioning": "Base conditioning input for inpainting task", + "latent": "Encoded latent from VAE", + "mask": "Area to inpaint (white regions)", + "expand": "Grow mask (+) or shrink mask (-)", + "blur": "Soften mask edges", + "mask_only": "Only generate content in masked area" + } + + return { + "required": { + "conditioning": ("CONDITIONING", {"tooltip": tooltips["conditioning"]}), + "latent": ("LATENT", {"tooltip": tooltips["latent"]}), + "mask": ("MASK", {"tooltip": tooltips["mask"]}), + "expand": ("INT", {"default": 5, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["expand"]}), + "blur": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 64.0, "step": 0.1, "tooltip": tooltips["blur"]}), + "mask_only": ("BOOLEAN", {"default": True, "tooltip": tooltips["mask_only"]}), + } + } + + RETURN_TYPES = ("CONDITIONING", "LATENT", "MASK") + RETURN_NAMES = ("CONDITIONING", "LATENT", "MASK") + FUNCTION = "prepare_inpaint_conditioning" + CATEGORY = "🧪AILab/🧽RMBG/🎭Inpaint" + + def add_latent_to_conditioning(self, conditioning, latent=None): + if latent is not None: + return node_helpers.conditioning_set_values( + conditioning, + {"reference_latents": [latent["samples"]]}, + append=True + ) + return conditioning + + def prepare_inpaint_conditioning(self, conditioning, latent, mask, expand=5, blur=3.0, mask_only=True): + processed_mask = mask + + if expand != 0: + processed_mask = expand_mask(processed_mask, expand) + + if blur > 0: + processed_mask = blur_mask(processed_mask, blur) + + modified_cond = node_helpers.conditioning_set_values( + conditioning, + { + "concat_latent_image": latent["samples"], + "concat_mask": processed_mask + } + ) + + final_cond = self.add_latent_to_conditioning(modified_cond, latent) + + output_latent = {"samples": latent["samples"]} + if mask_only: + output_latent["noise_mask"] = processed_mask + + return (final_cond, output_latent, processed_mask) + + +NODE_CLASS_MAPPINGS = { + "AILab_ReferenceLatentMask": AILab_ReferenceLatentMask, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AILab_ReferenceLatentMask": "Reference Latent Mask (RMBG) 🖼️🎭", +} \ No newline at end of file diff --git a/locales/en/nodeDefs.json b/locales/en/nodeDefs.json index c0931dd..2261379 100644 --- a/locales/en/nodeDefs.json +++ b/locales/en/nodeDefs.json @@ -416,5 +416,21 @@ "2": { "name": "WIDTH" }, "3": { "name": "HEIGHT" } } + }, + "AILab_ReferenceLatentMask": { + "display_name": "Kontext Reference Latent Mask (RMBG) 🎭", + "inputs": { + "conditioning": { "name": "Conditioning" }, + "latent": { "name": "Latent" }, + "mask": { "name": "Mask" }, + "expand": { "name": "Expand" }, + "blur": { "name": "Blur" }, + "mask_only": { "name": "Mask Only" } + }, + "outputs": { + "0": { "name": "CONDITIONING" }, + "1": { "name": "LATENT" }, + "2": { "name": "MASK" } + } } } \ No newline at end of file diff --git a/locales/fr/nodeDefs.json b/locales/fr/nodeDefs.json index f20f15b..17f1cae 100644 --- a/locales/fr/nodeDefs.json +++ b/locales/fr/nodeDefs.json @@ -415,5 +415,21 @@ "2": { "name": "LARGEUR" }, "3": { "name": "HAUTEUR" } } + }, + "AILab_ReferenceLatentMask": { + "display_name": "Kontext 参考潜伏遮罩 (RMBG) 🎭", + "inputs": { + "conditioning": { "name": "条件" }, + "latent": { "name": "latant" }, + "mask": { "name": "遮罩" }, + "expand": { "name": "扩展" }, + "blur": { "name": "模糊" }, + "mask_only": { "name": "仅遮罩" } + }, + "outputs": { + "0": { "name": "条件" }, + "1": { "name": "latant" }, + "2": { "name": "遮罩" } + } } } \ No newline at end of file diff --git a/locales/ja/nodeDefs.json b/locales/ja/nodeDefs.json index da19703..4864140 100644 --- a/locales/ja/nodeDefs.json +++ b/locales/ja/nodeDefs.json @@ -415,5 +415,21 @@ "2": { "name": "幅" }, "3": { "name": "高さ" } } + }, + "AILab_ReferenceLatentMask": { + "display_name": "Kontext 参考潜伏マスク (RMBG) 🎭", + "inputs": { + "conditioning": { "name": "条件" }, + "latent": { "name": "潜伏" }, + "mask": { "name": "マスク" }, + "expand": { "name": "拡張" }, + "blur": { "name": "ぼかし" }, + "mask_only": { "name": "マスクのみ" } + }, + "outputs": { + "0": { "name": "条件" }, + "1": { "name": "潜伏" }, + "2": { "name": "マスク" } + } } } diff --git a/locales/ko/nodeDefs.json b/locales/ko/nodeDefs.json index 791fb82..db46aa8 100644 --- a/locales/ko/nodeDefs.json +++ b/locales/ko/nodeDefs.json @@ -415,5 +415,21 @@ "2": { "name": "너비" }, "3": { "name": "높이" } } + }, + "AILab_ReferenceLatentMask": { + "display_name": "Kontext 참조 잠재 마스크 (RMBG) 🎭", + "inputs": { + "conditioning": { "name": "조건" }, + "latent": { "name": "잠재" }, + "mask": { "name": "마스크" }, + "expand": { "name": "확장" }, + "blur": { "name": "블러" }, + "mask_only": { "name": "마스크 전용" } + }, + "outputs": { + "0": { "name": "조건" }, + "1": { "name": "잠재" }, + "2": { "name": "마스크" } + } } } diff --git a/locales/ru/nodeDefs.json b/locales/ru/nodeDefs.json index 47935e3..978ae46 100644 --- a/locales/ru/nodeDefs.json +++ b/locales/ru/nodeDefs.json @@ -415,6 +415,22 @@ "2": { "name": "Ширина" }, "3": { "name": "Высота" } } - } + }, + "AILab_ReferenceLatentMask": { + "display_name": "Kontext Ссылка на латентную маску (RMBG) 🎭", + "inputs": { + "conditioning": { "name": "Условие" }, + "latent": { "name": "Латент" }, + "mask": { "name": "Маска" }, + "expand": { "name": "Расширение" }, + "blur": { "name": "Размытие" }, + "mask_only": { "name": "Только маска" } + }, + "outputs": { + "0": { "name": "Условие" }, + "1": { "name": "Латент" }, + "2": { "name": "Маска" } + } + } } diff --git a/locales/zh/nodeDefs.json b/locales/zh/nodeDefs.json index 66b4c4d..e6f86b8 100644 --- a/locales/zh/nodeDefs.json +++ b/locales/zh/nodeDefs.json @@ -415,5 +415,21 @@ "2": { "name": "宽度" }, "3": { "name": "高度" } } + }, + "AILab_ReferenceLatentMask": { + "display_name": "Kontext 参考潜伏遮罩 (RMBG) 🎭", + "inputs": { + "conditioning": { "name": "条件" }, + "latent": { "name": "latant" }, + "mask": { "name": "遮罩" }, + "expand": { "name": "扩展" }, + "blur": { "name": "模糊" }, + "mask_only": { "name": "仅遮罩" } + }, + "outputs": { + "0": { "name": "条件" }, + "1": { "name": "latant" }, + "2": { "name": "遮罩" } + } } } \ No newline at end of file