diff --git a/README.md b/README.md index 6b874ff..b9cab27 100644 --- a/README.md +++ b/README.md @@ -1,39 +1,38 @@ # ComfyUI-RMBG -A ComfyUI node for removing image backgrounds using RMBG-2.0. +A ComfyUI node for removing image backgrounds with multiple models: RMBG-2.0, INSPYRENET, and BEN. $${\color{red}If\ this\ custom\ node\ helps\ you\ or\ you\ like\ my\ work,\ please\ give\ me⭐on\ this\ repo!}$$ $${\color{red}It's\ a\ greatest\ encouragement\ for\ my\ efforts!}$$ ## News +- 2024/11/29: Update Comfyui-RMBG ComfyUI Custom Node to v1.2.0 ( [update.md](https://github.com/1038lab/ComfyUI-RMBG/blob/main/update.md) ) - 2024/11/21: Update Comfyui-RMBG ComfyUI Custom Node to v1.1.0 ( [update.md](https://github.com/1038lab/ComfyUI-RMBG/blob/main/update.md) ) ![comfyui-rmbg version compare](https://github.com/user-attachments/assets/2d23cf42-ca74-49e5-a8bf-9de377bd71aa) ## Features -RMBG-2.0 is built on the innovative BiRefNet (Bilateral Reference Network) architecture, offering: -- High accuracy in complex environments -- Precise edge detection and preservation -- Excellent handling of fine details -- Support for multiple objects in a single image -- Output Comparison -- Output with background -- Batch output for video - -![RMBG_3](https://github.com/user-attachments/assets/f3ffa3c4-5a21-4c0c-a078-b4ffe681c4c4) +![RMBG Demo](https://github.com/user-attachments/assets/f3ffa3c4-5a21-4c0c-a078-b4ffe681c4c4) ## Installation -1. Clone this repository to your ComfyUI custom_nodes folder: +1. install on ComfyUI-Manager, search `Comfyui-RMBG` and install + install requirment.txt in the ComfyUI-RMBG folder + ```bash + ./ComfyUI/python_embeded/python -m pip install -r requirements.txt + ``` + +2. Clone this repository to your ComfyUI custom_nodes folder: ```bash cd ComfyUI/custom_nodes git clone https://github.com/1038lab/ComfyUI-RMBG ``` -2. RMBG Model Download: - -- The model will be automatically downloaded to `ComfyUI/models/RMBG/RMBG-2.0` when first time using the custom node. +3. Manually download the models: +- The model will be automatically downloaded to `ComfyUI/models/RMBG/` when first time using the custom node. - Manually download the RMBG-2.0 model by visiting this [link](https://huggingface.co/briaai/RMBG-2.0/tree/main), then download the files and place them in the `/ComfyUI/models/RMBG/RMBG-2.0` folder. +- Manually download the INSPYRENET models by visiting the [link](https://huggingface.co/1038lab/inspyrenet), then download the files and place them in the `/ComfyUI/models/INSPYRENET` folder. +- Manually download the BEN model by visiting the [link](https://huggingface.co/PramaLLC/BEN), then download the files and place them in the `/ComfyUI/models/BEN` folder. ## Usage ![RMBG](https://github.com/user-attachments/assets/cd0eb92e-8f2e-4ae4-95f1-899a6d83cab6) @@ -45,13 +44,17 @@ git clone https://github.com/1038lab/ComfyUI-RMBG | **Processing Resolution** | Controls the processing resolution of the input image, affecting detail and memory usage. | Choose a value between 256 and 2048, with a default of 1024. Higher resolutions provide better detail but increase memory consumption. | | **Mask Blur** | Controls the amount of blur applied to the mask edges, reducing jaggedness. | Default value is 0. Try setting it between 1 and 5 for smoother edge effects. | | **Mask Offset** | Allows for expanding or shrinking the mask boundary. Positive values expand the boundary, while negative values shrink it. | Default value is 0. Adjust based on the specific image, typically fine-tuning between -10 and 10. | +| **Background** | Choose output background color | Alpha (transparent background) Black, White, Green, Blue, Red | +| **Invert Output** | Flip mask and image output | Invert both image and mask output | | **Performance Optimization** | Properly setting options can enhance performance when processing multiple images. | If memory allows, consider increasing `process_res` and `mask_blur` values for better results, but be mindful of memory usage. | ### Basic Usage 1. Load `RMBG (Remove Background)` node from the `🧪AILab/🧽RMBG` category 2. Connect an image to the input +3. Select a model from the dropdown menu +4. select the parameters as needed (optional) 3. Get two outputs: - - IMAGE: Processed image with transparent background + - IMAGE: Processed image with transparent, black, white, green, blue, or red background - MASK: Binary mask of the foreground ### Parameters @@ -59,21 +62,41 @@ git clone https://github.com/1038lab/ComfyUI-RMBG - `process_res`: Processing resolution (512-2048, step 128) - `mask_blur`: Blur amount for the mask (0-64) - `mask_offset`: Adjust mask edges (-20 to 20) +- `background`: Choose output background color +- `invert_output`: Flip mask and image output +- `optimize`: Toggle model optimization -## About RMBG-2.0 - -RMBG-2.0 is developed by BRIA AI and uses the BiRefNet architecture which includes: - -- **Localization Module (LM)**: Generates semantic maps for primary image areas -- **Restoration Module (RM)**: Performs precise boundary restoration using: - - Original Reference: Provides general background context - - Gradient Reference: Focuses on edges and fine details +
+

About Models

 +## RMBG-2.0 +RMBG-2.0 is is developed by BRIA AI and uses the BiRefNet architecture which includes: +- High accuracy in complex environments +- Precise edge detection and preservation +- Excellent handling of fine details +- Support for multiple objects in a single image +- Output Comparison +- Output with background +- Batch output for video The model is trained on a diverse dataset of over 15,000 high-quality images, ensuring: - Balanced representation across different image types - High accuracy in various scenarios - Robust performance with complex backgrounds +## INSPYRENET +INSPYRENET is specialized in human portrait segmentation, offering: +- Fast processing speed +- Good edge detection capability +- Ideal for portrait photos and human subjects + +## BEN +BEN is robust on various image types, offering: +- Good balance between speed and accuracy +- Effective on both simple and complex scenes +- Suitable for batch processing +
+ + ## Requirements - ComfyUI - Python 3.10+ @@ -82,11 +105,15 @@ The model is trained on a diverse dataset of over 15,000 high-quality images, en - torchvision>=0.15.0 - Pillow>=9.0.0 - numpy>=1.22.0 - - transformers>=4.30.0 - - safetensors>=0.3.0 + - huggingface-hub>=0.19.0 + - tqdm>=4.65.0 + - transformers>=4.35.0 + - transparent-background>=1.2.4 ## Credits - RMBG-2.0: https://huggingface.co/briaai/RMBG-2.0 +- INSPYRENET: https://github.com/plemeri/InSPyReNet +- BEN: https://huggingface.co/PramaLLC/BEN - Created by: [1038 Lab](https://github.com/1038lab) ## License diff --git a/ailab_RMBG.py b/ailab_RMBG.py index 0815bac..82b4a15 100644 --- a/ailab_RMBG.py +++ b/ailab_RMBG.py @@ -1,234 +1,405 @@ -import os -import torch -from PIL import Image -from torchvision import transforms -from torchvision.transforms.functional import normalize -import numpy as np -import folder_paths -from transformers import AutoModelForImageSegmentation -from PIL import ImageFilter -import torch.nn.functional as F -from huggingface_hub import hf_hub_download -import shutil - -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 = { - "RMBG-2.0": "briaai/RMBG-2.0" -} - -def tensor2pil(image): - return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) - -def pil2tensor(image): - return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) - -class AILAB_RMBG: - def __init__(self): - self.model = None - self.current_model_version = None - self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "RMBG-2.0") - - @classmethod - def INPUT_TYPES(s): - tooltips = { - "sensitivity": "Adjust mask detection strength", - "process_res": "Processing resolution (higher = more VRAM)", - "mask_blur": "Blur amount for mask edges", - "mask_offset": "Expand/Shrink mask boundary", - "background": "Choose background color (Alpha = transparent background)", - "invert_output": "Invert both image and mask output", - } - - return { - "required": { - "image": ("IMAGE",), - "model_version": (list(AVAILABLE_MODELS.keys()),), - }, - "optional": { - "sensitivity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": tooltips["sensitivity"]}), - "process_res": ("INT", {"default": 1024, "min": 256, "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": (["Alpha", "black", "white", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background"]}), - "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}), - } - } - - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("image", "mask") - FUNCTION = "remove_background" - CATEGORY = "🧪AILab/🧽RMBG" - - def check_model_cache(self, model_version): - model_files_path = os.path.join(self.cache_dir) - - if not os.path.exists(self.cache_dir): - return False, "Model directory not found" - - required_files = [ - 'config.json', - 'model.safetensors', - 'birefnet.py', - 'BiRefNet_config.py' - ] - - missing_files = [f for f in required_files if not os.path.exists(os.path.join(model_files_path, f))] - - if missing_files: - return False, f"Missing model files: {', '.join(missing_files)}" - - return True, "Model cache is complete" - - def clear_model(self): - if self.model is not None: - self.model.cpu() - del self.model - self.model = None - self.current_model_version = None - torch.cuda.empty_cache() - print("Model cleared from memory") - - def download_model_files(self, model_version): - model_id = AVAILABLE_MODELS[model_version] - required_files = { - 'config.json': 'config.json', - 'model.safetensors': 'model.safetensors', - 'birefnet.py': 'birefnet.py', - 'BiRefNet_config.py': 'BiRefNet_config.py' - } - - os.makedirs(self.cache_dir, exist_ok=True) - - try: - for filename, save_name in required_files.items(): - downloaded_path = hf_hub_download( - repo_id=model_id, - filename=filename, - 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 remove_background(self, image, model_version, sensitivity=1.0, process_res=1024, - mask_blur=0, mask_offset=0, invert_output=False, background="Alpha"): - try: - cache_status, message = self.check_model_cache(model_version) - - if not cache_status: - print(f"Model cache status: {message}") - print("Downloading required model files...") - download_status, download_message = self.download_model_files(model_version) - if not download_status: - raise RuntimeError(download_message) - print("Download completed.") - - if self.current_model_version != model_version or self.model is None: - if self.model is not None: - self.clear_model() - - self.model = AutoModelForImageSegmentation.from_pretrained( - self.cache_dir, - trust_remote_code=True, - local_files_only=True - ) - torch.set_float32_matmul_precision('high') - self.model.to(device) - self.model.eval() - self.current_model_version = model_version - print(f"Loaded model version: {model_version}") - - processed_images = [] - processed_masks = [] - - bg_colors = { - "Alpha": None, - "black": (0, 0, 0), - "white": (255, 255, 255), - "green": (0, 255, 0), - "blue": (0, 0, 255), - "red": (255, 0, 0) - } - - transform_image = transforms.Compose([ - transforms.Resize((process_res, process_res)), - transforms.ToTensor(), - transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) - ]) - - for img in image: - orig_image = tensor2pil(img) - w, h = orig_image.size - - input_tensor = transform_image(orig_image).unsqueeze(0).to(device) - - with torch.no_grad(): - result = self.model(input_tensor)[-1].sigmoid().cpu() - result = result[0].squeeze() - - result = result * (1 + (1 - sensitivity)) - result = torch.clamp(result, 0, 1) - - result = F.interpolate(result.unsqueeze(0).unsqueeze(0), - size=(h, w), - mode='bilinear').squeeze() - - mask_pil = tensor2pil(result) - - if invert_output: - mask_np = np.array(mask_pil) - mask_np = 255 - mask_np - mask_pil = Image.fromarray(mask_np) - - if mask_blur > 0: - mask_pil = mask_pil.filter(ImageFilter.GaussianBlur(radius=mask_blur)) - - if mask_offset != 0: - if mask_offset > 0: - for _ in range(mask_offset): - mask_pil = mask_pil.filter(ImageFilter.MaxFilter(3)) - else: - for _ in range(-mask_offset): - mask_pil = mask_pil.filter(ImageFilter.MinFilter(3)) - - rgba_image = orig_image.copy().convert('RGBA') - rgba_image.putalpha(mask_pil) - - if background != "Alpha": - bg_color = bg_colors[background] - bg_image = Image.new('RGBA', orig_image.size, (*bg_color, 255)) - composite_image = Image.alpha_composite(bg_image, rgba_image) - processed_images.append(pil2tensor(composite_image)) - else: - processed_images.append(pil2tensor(rgba_image)) - - processed_masks.append(pil2tensor(mask_pil)) - - torch.cuda.empty_cache() - - new_ims = torch.cat(processed_images, dim=0) - new_masks = torch.cat(processed_masks, dim=0) - - return (new_ims, new_masks) - - except Exception as e: - self.clear_model() - raise RuntimeError(f"Error in RMBG processing: {str(e)}") - -NODE_CLASS_MAPPINGS = { - "AILAB_RMBG": AILAB_RMBG -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "AILAB_RMBG": "RMBG (Remove Background)" -} \ No newline at end of file +# ComfyUI-RMBG +# 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. + +# This script is under MIT License, it's completely free to use and modify. +# However, if you make changes and distribute it or include it in other code, +# please acknowledge the original source. (https://github.com/AILab-AI/ComfyUI-RMBG) +# Supporting the original authors by acknowledging their work is greatly appreciated. + +import os +import torch +from PIL import Image +from torchvision import transforms +import numpy as np +import folder_paths +from PIL import ImageFilter +import torch.nn.functional as F +from huggingface_hub import hf_hub_download +import shutil +import sys +import importlib.util +from tqdm import tqdm +from transformers import AutoModelForImageSegmentation + +device = "cuda" if torch.cuda.is_available() else "cpu" + +# Add model path +folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) + +# Model configuration +AVAILABLE_MODELS = { + "RMBG-2.0": { + "type": "rmbg", + "repo_id": "briaai/RMBG-2.0", + "files": { + "config.json": "config.json", + "model.safetensors": "model.safetensors", + "birefnet.py": "birefnet.py", + "BiRefNet_config.py": "BiRefNet_config.py" + }, + "cache_dir": "RMBG-2.0" + }, + "INSPYRENET": { + "type": "inspyrenet", + "repo_id": "1038lab/inspyrenet", + "files": { + "inspyrenet.pth": "inspyrenet.pth" + }, + "cache_dir": "INSPYRENET" + }, + "BEN": { + "type": "ben", + "repo_id": "PramaLLC/BEN", + "files": { + "model.py": "model.py", + "BEN_Base.pth": "BEN_Base.pth" + }, + "cache_dir": "BEN" + } +} + +# Utility functions +def tensor2pil(image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +def handle_model_error(message): + print(f"[RMBG ERROR] {message}") + raise RuntimeError(message) + +class BaseModelLoader: + def __init__(self): + self.model = None + self.current_model_version = None + self.base_cache_dir = os.path.join(folder_paths.models_dir, "RMBG") + + def get_cache_dir(self, model_name): + return os.path.join(self.base_cache_dir, AVAILABLE_MODELS[model_name]["cache_dir"]) + + def check_model_cache(self, model_name): + model_info = AVAILABLE_MODELS[model_name] + cache_dir = self.get_cache_dir(model_name) + + if not os.path.exists(cache_dir): + return False, "Model directory not found" + + missing_files = [] + for filename in model_info["files"].keys(): + if not os.path.exists(os.path.join(cache_dir, model_info["files"][filename])): + missing_files.append(filename) + + if missing_files: + return False, f"Missing model files: {', '.join(missing_files)}" + + return True, "Model cache verified" + + def download_model(self, model_name): + model_info = AVAILABLE_MODELS[model_name] + cache_dir = self.get_cache_dir(model_name) + + try: + os.makedirs(cache_dir, exist_ok=True) + print(f"Downloading {model_name} model files...") + + for filename in model_info["files"].keys(): + print(f"Downloading {filename}...") + hf_hub_download( + repo_id=model_info["repo_id"], + filename=filename, + local_dir=cache_dir, + local_dir_use_symlinks=False + ) + + return True, "Model files downloaded successfully" + + except Exception as e: + return False, f"Error downloading model files: {str(e)}" + + def clear_model(self): + if self.model is not None: + self.model.cpu() + del self.model + self.model = None + self.current_model_version = None + torch.cuda.empty_cache() + print("Model cleared from memory") + +class RMBGModel(BaseModelLoader): + def __init__(self): + super().__init__() + + def load_model(self, model_name): + if self.current_model_version != model_name: + self.clear_model() + + cache_dir = self.get_cache_dir(model_name) + self.model = AutoModelForImageSegmentation.from_pretrained( + cache_dir, + trust_remote_code=True, + local_files_only=True + ) + + self.model.eval() + for param in self.model.parameters(): + param.requires_grad = False + + torch.set_float32_matmul_precision('high') + self.model.to(device) + self.current_model_version = model_name + + def process_image(self, image, model_name, params): + try: + self.load_model(model_name) + + transform_image = transforms.Compose([ + transforms.Resize((params["process_res"], params["process_res"])), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + orig_image = tensor2pil(image) + w, h = orig_image.size + + input_tensor = transform_image(orig_image).unsqueeze(0).to(device) + + with torch.no_grad(): + result = self.model(input_tensor)[-1].sigmoid().cpu() + result = result[0].squeeze() + + result = result * (1 + (1 - params["sensitivity"])) + result = torch.clamp(result, 0, 1) + + result = F.interpolate(result.unsqueeze(0).unsqueeze(0), + size=(h, w), + mode='bilinear').squeeze() + + return tensor2pil(result) + + except Exception as e: + handle_model_error(f"Error in RMBG processing: {str(e)}") + +class InspyrenetModel(BaseModelLoader): + def __init__(self): + super().__init__() + + def load_model(self, model_name): + if self.current_model_version != model_name: + self.clear_model() + + try: + import transparent_background + self.model = transparent_background.Remover() + self.current_model_version = model_name + except ImportError: + try: + import pip + pip.main(['install', 'transparent_background']) + import transparent_background + self.model = transparent_background.Remover() + self.current_model_version = model_name + except Exception as e: + handle_model_error(f"Failed to install transparent_background: {str(e)}") + + def process_image(self, image, model_name, params): + try: + self.load_model(model_name) + + orig_image = tensor2pil(image) + w, h = orig_image.size + + # Resize for processing + aspect_ratio = h / w + new_w = params["process_res"] + new_h = int(params["process_res"] * aspect_ratio) + resized_image = orig_image.resize((new_w, new_h), Image.LANCZOS) + + # Process image + foreground = self.model.process(resized_image, type='rgba') + foreground = foreground.resize((w, h), Image.LANCZOS) + mask = foreground.split()[-1] + + return mask + + except Exception as e: + handle_model_error(f"Error in Inspyrenet processing: {str(e)}") + +class BENModel(BaseModelLoader): + def __init__(self): + super().__init__() + + def load_model(self, model_name): + if self.current_model_version != model_name: + self.clear_model() + + cache_dir = self.get_cache_dir(model_name) + model_path = os.path.join(cache_dir, "model.py") + module_name = f"custom_ben_model_{hash(model_path)}" + + spec = importlib.util.spec_from_file_location(module_name, model_path) + ben_module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = ben_module + spec.loader.exec_module(ben_module) + + model_weights_path = os.path.join(cache_dir, "BEN_Base.pth") + self.model = ben_module.BEN_Base() + self.model.loadcheckpoints(model_weights_path) + + self.model.eval() + for param in self.model.parameters(): + param.requires_grad = False + + torch.set_float32_matmul_precision('high') + self.model.to(device) + self.current_model_version = model_name + + def process_image(self, image, model_name, params): + try: + self.load_model(model_name) + + orig_image = tensor2pil(image) + w, h = orig_image.size + + aspect_ratio = h / w + new_w = params["process_res"] + new_h = int(params["process_res"] * aspect_ratio) + resized_image = orig_image.resize((new_w, new_h), Image.LANCZOS) + + processed_input = resized_image.convert("RGBA") + + with torch.no_grad(): + _, foreground = self.model.inference(processed_input) + + foreground = foreground.resize((w, h), Image.LANCZOS) + mask = foreground.split()[-1] + + return mask + + except Exception as e: + handle_model_error(f"Error in BEN processing: {str(e)}") + +class RMBG: + def __init__(self): + self.models = { + "RMBG-2.0": RMBGModel(), + "INSPYRENET": InspyrenetModel(), + "BEN": BENModel() + } + + @classmethod + def INPUT_TYPES(s): + tooltips = { + "image": "Input image to be processed for background removal.", + "model": "Select the background removal model to use (RMBG-2.0, INSPYRENET, BEN).", + "sensitivity": "Adjust the strength of mask detection (higher values result in more aggressive detection).", + "process_res": "Set the processing resolution (higher values require more VRAM and may increase processing time).", + "mask_blur": "Specify the amount of blur to apply to the mask edges (0 for no blur, higher values for more blur).", + "mask_offset": "Adjust the mask boundary (positive values expand the mask, negative values shrink it).", + "background": "Choose the background color for the final output (Alpha for transparent background).", + "invert_output": "Enable to invert both the image and mask output (useful for certain effects).", + "optimize": "Enable model optimization for faster processing (may affect output quality)." + } + + return { + "required": { + "image": ("IMAGE", {"tooltip": tooltips["image"]}), + "model": (list(AVAILABLE_MODELS.keys()), {"tooltip": tooltips["model"]}), + }, + "optional": { + "sensitivity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": tooltips["sensitivity"]}), + "process_res": ("INT", {"default": 1024, "min": 256, "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": (["Alpha", "black", "white", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background"]}), + "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}), + "optimize": (["default", "on"], {"default": "default", "tooltip": tooltips["optimize"]}) + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "mask") + FUNCTION = "process_image" + CATEGORY = "🧪AILab/🧽RMBG" + + def process_image(self, image, model, **params): + try: + processed_images = [] + processed_masks = [] + + bg_colors = { + "Alpha": None, + "black": (0, 0, 0), + "white": (255, 255, 255), + "green": (0, 255, 0), + "blue": (0, 0, 255), + "red": (255, 0, 0) + } + + model_instance = self.models[model] + + # Check and download model if needed + cache_status, message = model_instance.check_model_cache(model) + if not cache_status: + print(f"Cache check: {message}") + print("Downloading required model files...") + download_status, download_message = model_instance.download_model(model) + if not download_status: + handle_model_error(download_message) + print("Model files downloaded successfully") + + for img in image: + # Get mask from specific model + mask = model_instance.process_image(img, model, params) + + # Post-process mask + mask_tensor = pil2tensor(mask) + mask_tensor = mask_tensor * (1 + (1 - params["sensitivity"])) + mask_tensor = torch.clamp(mask_tensor, 0, 1) + mask = tensor2pil(mask_tensor) + + if params["mask_blur"] > 0: + mask = mask.filter(ImageFilter.GaussianBlur(radius=params["mask_blur"])) + + if params["mask_offset"] != 0: + if params["mask_offset"] > 0: + for _ in range(params["mask_offset"]): + mask = mask.filter(ImageFilter.MaxFilter(3)) + else: + for _ in range(-params["mask_offset"]): + mask = mask.filter(ImageFilter.MinFilter(3)) + + if params["invert_output"]: + mask = Image.fromarray(255 - np.array(mask)) + + # Create final image + orig_image = tensor2pil(img) + orig_rgba = orig_image.convert("RGBA") + r, g, b, _ = orig_rgba.split() + foreground = Image.merge('RGBA', (r, g, b, mask)) + + if params["background"] != "Alpha": + bg_color = bg_colors[params["background"]] + bg_image = Image.new('RGBA', orig_image.size, (*bg_color, 255)) + composite_image = Image.alpha_composite(bg_image, foreground) + processed_images.append(pil2tensor(composite_image)) + else: + processed_images.append(pil2tensor(foreground)) + + processed_masks.append(pil2tensor(mask)) + + return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0)) + + except Exception as e: + handle_model_error(f"Error in image processing: {str(e)}") + +# Node Mapping +NODE_CLASS_MAPPINGS = { + "RMBG": RMBG +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RMBG": "RMBG (RMBG-2.0, INSPYRENET, BEN)" +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 8455d50..b8b72f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,24 +1,31 @@ -[build-system] -requires = ["setuptools>=42", "wheel"] -build-backend = "setuptools.build_meta" - [project] name = "ComfyUI-RMBG" -version = "1.1.0" -description = "A ComfyUI node for removing image backgrounds using RMBG-2.0" -authors = [{ name = "AILab" }] -license = { text = "MIT" } -requires-python = ">=3.10" -dependencies = [ - "torch>=2.0.0,<3.0.0", - "torchvision>=0.15.0,<1.0.0", - "Pillow>=9.0.0,<10.0.0", - "numpy>=1.22.0,<2.0.0", - "transformers>=4.30.0,<5.0.0", - "safetensors>=0.3.0,<1.0.0", - "timm>=0.6.12,<1.0.0", - "huggingface-hub>=0.16.0,<1.0.0" +version = "1.2.0" +description = "A ComfyUI node for background removal using multiple models (RMBG-2.0/INSPYRENET/BEN)" +authors = [ + {name = "AILab", email = ""} ] +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", +] +requires-python = ">=3.8" +readme = "README.md" +license = {text = "MIT"} -[tool.setuptools.packages.find] -include = ["ComfyUI-RMBG*"] \ No newline at end of file +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +packages = ["comfyui-rmbg"] + +[project.urls] +Homepage = "https://github.com/username/ComfyUI-RMBG" +Repository = "https://github.com/username/ComfyUI-RMBG.git" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index f25b527..8399ad2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,8 @@ -torch>=2.0.0,<3.0.0 -torchvision>=0.15.0,<1.0.0 -Pillow>=9.0.0,<10.0.0 -numpy>=1.22.0,<2.0.0 -transformers>=4.30.0,<5.0.0 -safetensors>=0.3.0,<1.0.0 -timm>=0.6.12,<1.0.0 -huggingface-hub>=0.16.0,<1.0.0 \ No newline at end of file +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 \ No newline at end of file diff --git a/update.md b/update.md index 6b033a5..cef4874 100644 --- a/update.md +++ b/update.md @@ -1,58 +1,105 @@ -# ComfyUI-RMBG Update Log - -## Version 1.1.0 - -### New Features -- Added background color options - - Alpha (transparent background) - - Black, White, Green, Blue, Red - -![rmbg1 1](https://github.com/user-attachments/assets/4f7d073c-f9cc-4bdb-875c-ba51decc9d5a) - -- Improved mask processing - - Better detail preservation - - Enhanced edge quality - - More accurate segmentation - -![rmbg version compare](https://github.com/user-attachments/assets/8339aa8e-46db-4f11-aa7b-0a710f0a1711) - -- Added video batch processing - - Support for video file background removal - - Maintains original video framerate and resolution - - Multiple output format support (with Alpha channel) - - Efficient batch processing for video frames - -https://github.com/user-attachments/assets/259220d3-c148-4030-93d6-c17dd5bccee1 - -- Added model cache management - - Cache status checking - - Model memory cleanup - - Better error handling - -### Parameter Updates -- Renamed 'invert_mask' to 'invert_output' for clarity -- Added sensitivity adjustment for mask strength -- Updated tooltips for better clarity - -### Technical Improvements -- Optimized image processing pipeline -- Added proper model cache verification -- Improved memory management -- Better error handling and recovery -- Enhanced batch processing performance for videos - -### Dependencies -- Added timm>=0.6.12,<1.0.0 for model support -- Updated requirements.txt with version constraints - -### Bug Fixes -- Fixed mask detail preservation issues -- Improved mask edge quality -- Fixed memory leaks in model handling - -### Usage Notes -- The 'Alpha' background option provides transparent background -- Sensitivity parameter now controls mask strength -- Model cache is checked before each operation -- Memory is automatically cleaned when switching models -- Video processing supports various formats and maintains quality +# ComfyUI-RMBG Update Log + +## v1.2.0 (2024/11/29) + +### Major Changes +- Combined three background removal models into one unified node +- Added support for RMBG-2.0, INSPYRENET, and BEN models +- Implemented lazy loading for models (only downloads when first used) + +### Model Introduction +- RMBG-2.0 ([Homepage](https://huggingface.co/briaai/RMBG-2.0)) + - Latest version of RMBG model + - Excellent performance on complex backgrounds + - High accuracy in preserving fine details + - Best for general purpose background removal + +- INSPYRENET ([Homepage](https://github.com/plemeri/InSPyReNet)) + - Specialized in human portrait segmentation + - Fast processing speed + - Good edge detection capability + - Ideal for portrait photos and human subjects + +- BEN (Background Elimination Network) ([Homepage](https://huggingface.co/PramaLLC/BEN)) + - Robust performance on various image types + - Good balance between speed and accuracy + - Effective on both simple and complex scenes + - Suitable for batch processing + +### Features +- Unified interface for all three models +- Common parameters for all models: + - Sensitivity adjustment + - Processing resolution control + - Mask blur and offset options + - Multiple background color options + - Invert output option + - Model optimization toggle + +### Improvements +- Optimized memory usage with model clearing +- Enhanced error handling and user feedback +- Added detailed tooltips for all parameters +- Improved mask post-processing + +### Dependencies +- Updated all package dependencies to latest stable versions +- Added support for transparent-background package +- Optimized dependency management + +## Version 1.1.0 (2024/11/21) + +### New Features +- Added background color options + - Alpha (transparent background) + - Black, White, Green, Blue, Red + +![rmbg1 Demo](https://github.com/user-attachments/assets/4f7d073c-f9cc-4bdb-875c-ba51decc9d5a) + +- Improved mask processing + - Better detail preservation + - Enhanced edge quality + - More accurate segmentation + +![rmbg version compare](https://github.com/user-attachments/assets/8339aa8e-46db-4f11-aa7b-0a710f0a1711) + +- Added video batch processing + - Support for video file background removal + - Maintains original video framerate and resolution + - Multiple output format support (with Alpha channel) + - Efficient batch processing for video frames + +https://github.com/user-attachments/assets/259220d3-c148-4030-93d6-c17dd5bccee1 + +- Added model cache management + - Cache status checking + - Model memory cleanup + - Better error handling + +### Parameter Updates +- Renamed 'invert_mask' to 'invert_output' for clarity +- Added sensitivity adjustment for mask strength +- Updated tooltips for better clarity + +### Technical Improvements +- Optimized image processing pipeline +- Added proper model cache verification +- Improved memory management +- Better error handling and recovery +- Enhanced batch processing performance for videos + +### Dependencies +- Added timm>=0.6.12,<1.0.0 for model support +- Updated requirements.txt with version constraints + +### Bug Fixes +- Fixed mask detail preservation issues +- Improved mask edge quality +- Fixed memory leaks in model handling + +### Usage Notes +- The 'Alpha' background option provides transparent background +- Sensitivity parameter now controls mask strength +- Model cache is checked before each operation +- Memory is automatically cleaned when switching models +- Video processing supports various formats and maintains quality