Add files via upload

2025-01-22 01:17:14 -08:00
parent 919fecfd6d
commit 70c372df82
8 changed files with 326 additions and 30 deletions
@@ -1,4 +1,4 @@
-# ComfyUI-RMBG v1.5.0
+# ComfyUI-RMBG v1.6.0
 # This custom node for ComfyUI provides functionality for background removal using various models,
 # including RMBG-2.0, INSPYRENET, and BEN. It leverages deep learning techniques
 # to process images and generate masks for background removal.
@@ -10,7 +10,7 @@
 # When using or modifying this code, please respect both the original model licenses
 # and this integration's license terms.
 #
-# Source: https://github.com/1038lab/ComfyUI-RMBG
+# Source: https://github.com/AILab-AI/ComfyUI-RMBG

 import os
 import torch
@@ -276,4 +276,4 @@ NODE_CLASS_MAPPINGS = {

 NODE_DISPLAY_NAME_MAPPINGS = {
    "ClothesSegment": "Clothes Segment (RMBG)"
-} 
+} 
@@ -0,0 +1,281 @@
+# ComfyUI-RMBG v1.6.0
+# This custom node for ComfyUI provides functionality for face parsing using Segformer model.
+# 
+# This integration script follows GPL-3.0 License.
+# When using or modifying this code, please respect both the original model licenses
+# and this integration's license terms.
+#
+# Source: https://github.com/AILab-AI/ComfyUI-RMBG
+
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+from typing import Tuple, Union
+from PIL import Image, ImageFilter
+from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
+import folder_paths
+from huggingface_hub import hf_hub_download
+import shutil
+from torchvision import transforms
+
+def pil2tensor(image: Image.Image) -> torch.Tensor:
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
+
+def tensor2pil(image: torch.Tensor) -> Image.Image:
+    return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8))
+
+def image2mask(image: Image.Image) -> torch.Tensor:
+    if isinstance(image, Image.Image):
+        image = pil2tensor(image)
+    return image.squeeze()[..., 0]
+
+def mask2image(mask: torch.Tensor) -> Image.Image:
+    if len(mask.shape) == 2:
+        mask = mask.unsqueeze(0)
+    return tensor2pil(mask)
+
+def RGB2RGBA(image: Image.Image, mask: Union[Image.Image, torch.Tensor]) -> Image.Image:
+    if isinstance(mask, torch.Tensor):
+        mask = mask2image(mask)
+    if mask.size != image.size:
+        mask = mask.resize(image.size, Image.Resampling.LANCZOS)
+    return Image.merge('RGBA', (*image.convert('RGB').split(), mask.convert('L')))
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG"))
+
+AVAILABLE_MODELS = {
+    "face_parsing": "1038lab/segformer_face"
+}
+
+class FaceSegment:
+    def __init__(self):
+        self.processor = None
+        self.model = None
+        self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "segformer_face")
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        available_classes = [
+            # "Background",  # Not a facial feature
+            "Skin", "Nose", "Eyeglasses", "Left-eye", "Right-eye",
+            "Left-eyebrow", "Right-eyebrow", "Left-ear", "Right-ear", "Mouth",
+            "Upper-lip", "Lower-lip", "Hair", "Earring", "Neck",
+            # "Hat",        # Not a facial feature
+            # "Necklace",   # Not a facial feature
+            # "Clothing"    # Not a facial feature
+        ]
+        
+        tooltips = {
+            "process_res": "Processing resolution (higher = more VRAM)",
+            "mask_blur": "Blur amount for mask edges",
+            "mask_offset": "Expand/Shrink mask boundary",
+            "background_color": "Choose background color (Alpha = transparent)",
+            "invert_output": "Invert both image and mask output",
+        }
+        
+        return {
+            "required": {
+                "images": ("IMAGE",),
+            },
+            "optional": {
+                **{cls_name: ("BOOLEAN", {"default": False}) 
+                   for cls_name in available_classes},
+                "process_res": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 32, "tooltip": tooltips["process_res"]}),
+                "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}),
+                "mask_offset": ("INT", {"default": 0, "min": -20, "max": 20, "step": 1, "tooltip": tooltips["mask_offset"]}),
+                "background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}),
+                "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "MASK")
+    RETURN_NAMES = ("images", "mask")
+    FUNCTION = "segment_face"
+    CATEGORY = "🧪AILab/🧽RMBG"
+
+    def check_model_cache(self):
+        if not os.path.exists(self.cache_dir):
+            return False, "Model directory not found"
+        
+        required_files = [
+            'config.json',
+            'model.safetensors',
+            'preprocessor_config.json'
+        ]
+        
+        missing_files = [f for f in required_files if not os.path.exists(os.path.join(self.cache_dir, f))]
+        if missing_files:
+            return False, f"Required model files missing: {', '.join(missing_files)}"
+        return True, "Model cache verified"
+
+    def clear_model(self):
+        if self.model is not None:
+            self.model.cpu()
+            del self.model
+            self.model = None
+            self.processor = None
+            torch.cuda.empty_cache()
+
+    def download_model_files(self):
+        model_id = AVAILABLE_MODELS["face_parsing"]
+        model_files = {
+            'config.json': 'config.json',
+            'model.safetensors': 'model.safetensors',
+            'preprocessor_config.json': 'preprocessor_config.json'
+        }
+        
+        os.makedirs(self.cache_dir, exist_ok=True)
+        print(f"Downloading face parsing model files...")
+        
+        try:
+            for save_name, repo_path in model_files.items():
+                print(f"Downloading {save_name}...")
+                downloaded_path = hf_hub_download(
+                    repo_id=model_id,
+                    filename=repo_path,
+                    local_dir=self.cache_dir,
+                    local_dir_use_symlinks=False
+                )
+                
+                if os.path.dirname(downloaded_path) != self.cache_dir:
+                    target_path = os.path.join(self.cache_dir, save_name)
+                    shutil.move(downloaded_path, target_path)
+            return True, "Model files downloaded successfully"
+        except Exception as e:
+            return False, f"Error downloading model files: {str(e)}"
+
+    def segment_face(self, images, process_res=512, mask_blur=0, mask_offset=0, background_color="Alpha", invert_output=False, **class_selections):
+        try:
+            # Check and download model if needed
+            cache_status, message = self.check_model_cache()
+            if not cache_status:
+                print(f"Cache check: {message}")
+                download_status, download_message = self.download_model_files()
+                if not download_status:
+                    raise RuntimeError(download_message)
+            
+            # Load model if needed
+            if self.processor is None:
+                self.processor = SegformerImageProcessor.from_pretrained(self.cache_dir)
+                self.model = AutoModelForSemanticSegmentation.from_pretrained(self.cache_dir)
+                self.model.eval()
+                for param in self.model.parameters():
+                    param.requires_grad = False
+                self.model.to(device)
+
+            # Class mapping for segmentation
+            class_map = {
+                "Background": 0, "Skin": 1, "Nose": 2, "Eyeglasses": 3,
+                "Left-eye": 4, "Right-eye": 5, "Left-eyebrow": 6, "Right-eyebrow": 7,
+                "Left-ear": 8, "Right-ear": 9, "Mouth": 10, "Upper-lip": 11,
+                "Lower-lip": 12, "Hair": 13, "Hat": 14, "Earring": 15,
+                "Necklace": 16, "Neck": 17, "Clothing": 18
+            }
+
+            # Get selected classes
+            selected_classes = [name for name, selected in class_selections.items() if selected]
+            if not selected_classes:
+                selected_classes = ["Skin", "Nose", "Eyes", "Mouth"]
+
+            # Image preprocessing
+            transform_image = transforms.Compose([
+                transforms.Resize((process_res, process_res)),
+                transforms.ToTensor(),
+            ])
+
+            batch_tensor = []
+            batch_masks = []
+            
+            for image in images:
+                orig_image = tensor2pil(image)
+                w, h = orig_image.size
+                
+                input_tensor = transform_image(orig_image)
+
+                if input_tensor.shape[0] == 4:
+                    input_tensor = input_tensor[:3]
+                
+                input_tensor = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(input_tensor)
+                
+                input_tensor = input_tensor.unsqueeze(0).to(device)
+                
+                with torch.no_grad():
+                    outputs = self.model(input_tensor)
+                    logits = outputs.logits.cpu()
+                    upsampled_logits = nn.functional.interpolate(
+                        logits,
+                        size=(h, w),
+                        mode="bilinear",
+                        align_corners=False,
+                    )
+                    pred_seg = upsampled_logits.argmax(dim=1)[0]
+
+                    # Combine selected class masks
+                    combined_mask = None
+                    for class_name in selected_classes:
+                        mask = (pred_seg == class_map[class_name]).float()
+                        if combined_mask is None:
+                            combined_mask = mask
+                        else:
+                            combined_mask = torch.clamp(combined_mask + mask, 0, 1)
+
+                    # Convert mask to PIL for processing
+                    mask_image = Image.fromarray((combined_mask.numpy() * 255).astype(np.uint8))
+
+                    if mask_blur > 0:
+                        mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur))
+
+                    if mask_offset != 0:
+                        if mask_offset > 0:
+                            mask_image = mask_image.filter(ImageFilter.MaxFilter(size=mask_offset * 2 + 1))
+                        else:
+                            mask_image = mask_image.filter(ImageFilter.MinFilter(size=-mask_offset * 2 + 1))
+
+                    if invert_output:
+                        mask_image = Image.fromarray(255 - np.array(mask_image))
+
+                    # Handle background color
+                    if background_color == "Alpha":
+                        rgba_image = RGB2RGBA(orig_image, mask_image)
+                        result_image = pil2tensor(rgba_image)
+                    else:
+                        bg_colors = {
+                            "black": (0, 0, 0),
+                            "white": (255, 255, 255),
+                            "gray": (128, 128, 128),
+                            "green": (0, 255, 0),
+                            "blue": (0, 0, 255),
+                            "red": (255, 0, 0)
+                        }
+                        
+                        rgba_image = RGB2RGBA(orig_image, mask_image)
+                        bg_image = Image.new('RGBA', orig_image.size, (*bg_colors[background_color], 255))
+                        composite_image = Image.alpha_composite(bg_image, rgba_image)
+                        result_image = pil2tensor(composite_image.convert('RGB'))
+
+                    batch_tensor.append(result_image)
+                    batch_masks.append(pil2tensor(mask_image))
+
+            # Prepare final output
+            batch_tensor = torch.cat(batch_tensor, dim=0)
+            batch_masks = torch.cat(batch_masks, dim=0)
+            
+            return (batch_tensor, batch_masks)
+
+        except Exception as e:
+            self.clear_model()
+            raise RuntimeError(f"Error in Face Parsing processing: {str(e)}")
+        finally:
+            if not self.model.training:
+                self.clear_model()
+
+NODE_CLASS_MAPPINGS = {
+    "FaceSegment": FaceSegment
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "FaceSegment": "Face Segment (RMBG)"
+} 
@@ -1,4 +1,4 @@
-# ComfyUI-RMBG v1.5.0
+# ComfyUI-RMBG v1.6.0
 # This custom node for ComfyUI provides functionality for fashion segmentation using segformer-b3-fashion model.
 # It leverages deep learning techniques to process images and generate masks for fashion items segmentation.

@@ -9,7 +9,7 @@
 # When using or modifying this code, please respect both the original model licenses
 # and this integration's license terms.
 #
-# Source: https://github.com/1038lab/ComfyUI-RMBG
+# Source: https://github.com/AILab-AI/ComfyUI-RMBG

 import os
 import torch
@@ -354,4 +354,4 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
    "FashionSegmentAccessories": "Accessories Segment (RMBG)",
    "FashionSegmentClothing": "Fashion Segment (RMBG)"
-} 
+} 
@@ -1,4 +1,4 @@
-# ComfyUI-RMBG v1.5.0
+# ComfyUI-RMBG v1.6.0
 # This custom node for ComfyUI provides functionality for background removal using various models,
 # including RMBG-2.0, INSPYRENET, and BEN. It leverages deep learning techniques
 # to process images and generate masks for background removal.
@@ -12,7 +12,7 @@
 # When using or modifying this code, please respect both the original model licenses
 # and this integration's license terms.
 #
-# Source: https://github.com/1038lab/ComfyUI-RMBG
+# Source: https://github.com/AILab-AI/ComfyUI-RMBG

 import os
 import torch
@@ -435,4 +435,4 @@ NODE_CLASS_MAPPINGS = {

 NODE_DISPLAY_NAME_MAPPINGS = {
    "RMBG": "Remove Background (RMBG)"
-} 
+} 
@@ -1,4 +1,4 @@
-# ComfyUI-RMBG v1.5.0
+# ComfyUI-RMBG v1.6.0
 # This custom node for ComfyUI provides functionality for background removal using various models,
 # including RMBG-2.0, INSPYRENET, and BEN. It leverages deep learning techniques
 # to process images and generate masks for background removal.
@@ -11,7 +11,7 @@
 # When using or modifying this code, please respect both the original model licenses
 # and this integration's license terms.
 #
-# Source: https://github.com/1038lab/ComfyUI-RMBG
+# Source: https://github.com/AILab-AI/ComfyUI-RMBG

 import os
 import sys
@@ -345,4 +345,4 @@ NODE_CLASS_MAPPINGS = {

 NODE_DISPLAY_NAME_MAPPINGS = {
    "Segment": "Segment (RMBG)"
-}
+}
@@ -1,15 +1,16 @@
-[project]
-name = "comfyui-rmbg"
-description = "A ComfyUI custom node designed for advanced image background removal utilizing multiple models, including RMBG-2.0, INSPYRENET, and BEN."
-version = "1.5.0"
-license = {file = "LICENSE"}
-dependencies = ["torch>=2.0.0", "torchvision>=0.15.0", "Pillow>=9.0.0", "numpy>=1.22.0", "huggingface-hub>=0.19.0", "tqdm>=4.65.0", "transformers>=4.35.0", "transparent-background>=1.2.4", "groundingdino-py>=0.4.0", "segment-anything>=1.0", "opencv-python>=4.7.0"]
-
-[project.urls]
-Repository = "https://github.com/1038lab/ComfyUI-RMBG"
-#  Used by Comfy Registry https://comfyregistry.org
-
-[tool.comfy]
-PublisherId = "ailab"
-DisplayName = "ComfyUI-RMBG"
-Icon = ""
+[project]
+name = "comfyui-rmbg"
+description = "A ComfyUI custom node designed for advanced image background removal utilizing multiple models, including RMBG-2.0, INSPYRENET, and BEN."
+version = "1.5.0"
+license = {file = "LICENSE"}
+
+dependencies = ["torch>=2.0.0", "torchvision>=0.15.0", "Pillow>=9.0.0", "numpy>=1.22.0", "huggingface-hub>=0.19.0", "transformers>=4.35.0", "transparent-background>=1.2.4", "tqdm>=4.65.0", "segment-anything>=1.0", "groundingdino-py>=0.4.0", "opencv-python>=4.7.0"]
+    
+[project.urls]
+Repository = "https://github.com/1038lab/ComfyUI-RMBG"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "ailab"
+DisplayName = "ComfyUI-RMBG"
+Icon = ""
@@ -1,11 +1,12 @@
+# Base dependencies
 torch>=2.0.0
 torchvision>=0.15.0
 Pillow>=9.0.0
 numpy>=1.22.0
 huggingface-hub>=0.19.0
-tqdm>=4.65.0
 transformers>=4.35.0
 transparent-background>=1.2.4
-groundingdino-py>=0.4.0
+tqdm>=4.65.0
 segment-anything>=1.0
+groundingdino-py>=0.4.0
 opencv-python>=4.7.0
@@ -1,5 +1,18 @@
 # ComfyUI-RMBG Update Log

+## v1.6.0 (2025/01/22)
+
+### New Face Segment Custom Node
+- Added a new custom node for face parsing and segmentation
+  - Support for 19 facial feature categories (Skin, Nose, Eyes, Eyebrows, etc.)
+  - Precise facial feature extraction and segmentation
+  - Multiple feature selection for combined segmentation
+  - Same parameter controls as other RMBG nodes
+  - Automatic model downloading and resource management
+  - Perfect for portrait editing and facial feature manipulation
+
+![RMBGv_1 6 0](https://github.com/user-attachments/assets/face-segment-demo)
+
 ## v1.5.0 (2025/01/05)

 ### New Fashion and accessories Segment Custom Node
@@ -12,7 +25,7 @@

 ![RMBGv_1 5 0](https://github.com/user-attachments/assets/a250c1a6-8425-4902-b902-a6e1a8bfe959)

-## v1.4.0 (2025/1/2)
+## v1.4.0 (2025/01/02)

 ### New Clothes Segment Node
 - Added intelligent clothes segmentation functionality