Add files via upload

2025-07-15 02:40:30 -07:00
parent c7afaf2dce
commit f09fd46b2a
8 changed files with 293 additions and 24 deletions
@@ -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}"
        
@@ -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) 🖼️🎭",
+} 
@@ -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" }
+    }
  }
 } 
@@ -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": "遮罩" }
+    }
  }
 }
@@ -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": "マスク" }
+    }
  }
 }
@@ -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": "마스크" }
+    }
  }
 } 
@@ -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": "Маска" }
+    }
+  }   
 }

@@ -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": "遮罩" }
+   }
 }
 }