diff --git a/AILab_ClothSegment.py b/AILab_ClothSegment.py index 9488af2..77e0073 100644 --- a/AILab_ClothSegment.py +++ b/AILab_ClothSegment.py @@ -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)" -} +} \ No newline at end of file diff --git a/AILab_FaceSegment.py b/AILab_FaceSegment.py new file mode 100644 index 0000000..483d263 --- /dev/null +++ b/AILab_FaceSegment.py @@ -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)" +} \ No newline at end of file diff --git a/AILab_FashionSegment.py b/AILab_FashionSegment.py index 500a73d..e52420d 100644 --- a/AILab_FashionSegment.py +++ b/AILab_FashionSegment.py @@ -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)" -} +} \ No newline at end of file diff --git a/AILab_RMBG.py b/AILab_RMBG.py index faddfc9..f72fca8 100644 --- a/AILab_RMBG.py +++ b/AILab_RMBG.py @@ -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)" -} +} \ No newline at end of file diff --git a/AILab_Segment.py b/AILab_Segment.py index c5bdb83..f1ceecf 100644 --- a/AILab_Segment.py +++ b/AILab_Segment.py @@ -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)" -} +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index f7007cf..7b419e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 = "" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index cf87157..0323161 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 \ No newline at end of file diff --git a/update.md b/update.md index 5912a98..5b69c4b 100644 --- a/update.md +++ b/update.md @@ -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