diff --git a/AILab_RMBG.py b/AILab_RMBG.py index 52d7e37..fcfd71e 100644 --- a/AILab_RMBG.py +++ b/AILab_RMBG.py @@ -1,442 +1,442 @@ -# ComfyUI-RMBG v1.2.2 -# 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. - -# License Notice: -# - RMBG-2.0: Apache-2.0 License (https://huggingface.co/briaai/RMBG-2.0) -# - INSPYRENET: MIT License (https://github.com/plemeri/InSPyReNet) -# - BEN: Apache-2.0 License (https://huggingface.co/PramaLLC/BEN) -# -# 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 -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.safetensors": "inspyrenet.safetensors" - }, - "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, images, model_name, params): - try: - self.load_model(model_name) - - # Prepare batch processing - 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]) - ]) - - # Ensure input is in list format - if isinstance(images, torch.Tensor): - if len(images.shape) == 3: - images = [images] - else: - images = [img for img in images] - - # Store original image sizes - original_sizes = [tensor2pil(img).size for img in images] - - # Batch process transformations - input_tensors = [transform_image(tensor2pil(img)).unsqueeze(0) for img in images] - input_batch = torch.cat(input_tensors, dim=0).to(device) - - with torch.no_grad(): - results = self.model(input_batch)[-1].sigmoid().cpu() - masks = [] - - # Process each result and resize back to original dimensions - for i, (result, (orig_w, orig_h)) in enumerate(zip(results, original_sizes)): - result = result.squeeze() - result = result * (1 + (1 - params["sensitivity"])) - result = torch.clamp(result, 0, 1) - - # Resize back to original dimensions - result = F.interpolate(result.unsqueeze(0).unsqueeze(0), - size=(orig_h, orig_w), - mode='bilinear').squeeze() - - masks.append(tensor2pil(result)) - - return masks - - except Exception as e: - handle_model_error(f"Error in batch 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: - """ - RMBG Node: Advanced Background Removal Suite - - This node provides professional background removal capabilities using three state-of-the-art models: - - RMBG-2.0: Latest model with excellent performance on complex backgrounds - - INSPYRENET: Specialized for human portrait segmentation - - BEN: Versatile model with good balance of speed and accuracy - - Features: - - Batch processing support - - Multiple background options - - Advanced mask refinement - - High-quality edge preservation - - Memory-efficient processing - """ - 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 (Background Remover)" +# ComfyUI-RMBG v1.3.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. + +# License Notice: +# - RMBG-2.0: Apache-2.0 License (https://huggingface.co/briaai/RMBG-2.0) +# - INSPYRENET: MIT License (https://github.com/plemeri/InSPyReNet) +# - BEN: Apache-2.0 License (https://huggingface.co/PramaLLC/BEN) +# +# 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 +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.safetensors": "inspyrenet.safetensors" + }, + "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, images, model_name, params): + try: + self.load_model(model_name) + + # Prepare batch processing + 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]) + ]) + + # Ensure input is in list format + if isinstance(images, torch.Tensor): + if len(images.shape) == 3: + images = [images] + else: + images = [img for img in images] + + # Store original image sizes + original_sizes = [tensor2pil(img).size for img in images] + + # Batch process transformations + input_tensors = [transform_image(tensor2pil(img)).unsqueeze(0) for img in images] + input_batch = torch.cat(input_tensors, dim=0).to(device) + + with torch.no_grad(): + results = self.model(input_batch)[-1].sigmoid().cpu() + masks = [] + + # Process each result and resize back to original dimensions + for i, (result, (orig_w, orig_h)) in enumerate(zip(results, original_sizes)): + result = result.squeeze() + result = result * (1 + (1 - params["sensitivity"])) + result = torch.clamp(result, 0, 1) + + # Resize back to original dimensions + result = F.interpolate(result.unsqueeze(0).unsqueeze(0), + size=(orig_h, orig_w), + mode='bilinear').squeeze() + + masks.append(tensor2pil(result)) + + return masks + + except Exception as e: + handle_model_error(f"Error in batch 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: + """ + RMBG Node: Advanced Background Removal Suite + + This node provides professional background removal capabilities using three state-of-the-art models: + - RMBG-2.0: Latest model with excellent performance on complex backgrounds + - INSPYRENET: Specialized for human portrait segmentation + - BEN: Versatile model with good balance of speed and accuracy + + Features: + - Batch processing support + - Multiple background options + - Advanced mask refinement + - High-quality edge preservation + - Memory-efficient processing + """ + 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 (Background Remover)" } \ No newline at end of file diff --git a/AILab_Segment.py b/AILab_Segment.py new file mode 100644 index 0000000..d237d1c --- /dev/null +++ b/AILab_Segment.py @@ -0,0 +1,342 @@ +# ComfyUI-RMBG v1.3.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. + +# License Notice: +# - SAM: MIT License (https://github.com/facebookresearch/segment-anything) +# - GroundingDINO: MIT License (https://github.com/IDEA-Research/GroundingDINO) + +# 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 sys +import copy +import requests +from urllib.parse import urlparse + +import torch +import numpy as np +from PIL import Image +from PIL import ImageFilter +from torch.hub import download_url_to_file + +import folder_paths +import comfy.model_management +from segment_anything import sam_model_registry, SamPredictor + +SAM_MODELS = { + "sam_vit_h (2.56GB)": { + "model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_vit_h.pth", + "model_type": "vit_h" + }, + "sam_vit_l (1.25GB)": { + "model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_vit_l.pth", + "model_type": "vit_l" + }, + "sam_vit_b (375MB)": { + "model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_vit_b.pth", + "model_type": "vit_b" + } +} + +DINO_MODELS = { + "GroundingDINO_SwinT_OGC (694MB)": { + "config_url": "https://huggingface.co/1038lab/GroundingDINO/resolve/main/GroundingDINO_SwinT_OGC.cfg.py", + "model_url": "https://huggingface.co/1038lab/GroundingDINO/resolve/main/groundingdino_swint_ogc.pth", + }, + "GroundingDINO_SwinB (938MB)": { + "config_url": "https://huggingface.co/1038lab/GroundingDINO/resolve/main/GroundingDINO_SwinB.cfg.py", + "model_url": "https://huggingface.co/1038lab/GroundingDINO/resolve/main/groundingdino_swinb_cogcoor.pth" + } +} + +def normalize_array(arr): + return arr.astype(np.float32) / 255.0 + +def denormalize_array(arr): + return np.clip(255. * arr, 0, 255).astype(np.uint8) + +def create_tensor_output(image_np, masks, boxes_filt): + output_masks, output_images = [], [] + for mask in masks: + image_np_copy = copy.deepcopy(image_np) + image_np_copy[~np.any(mask, axis=0)] = np.array([0, 0, 0, 0]) + output_image, output_mask = split_image_mask( + Image.fromarray(image_np_copy)) + output_masks.append(output_mask) + output_images.append(output_image) + return (torch.cat(output_images, dim=0), torch.cat(output_masks, dim=0)) + +def split_image_mask(image): + image_rgb = image.convert("RGB") + image_rgb = np.array(image_rgb).astype(np.float32) / 255.0 + image_rgb = torch.from_numpy(image_rgb)[None,] + if 'A' in image.getbands(): + mask = np.array(image.getchannel('A')).astype(np.float32) / 255.0 + mask = torch.from_numpy(mask)[None,] + else: + mask = torch.zeros((image.height, image.width), dtype=torch.float32, device="cpu")[None,] + return (image_rgb, mask) + +def process_mask(mask_image: Image.Image, invert_output: bool = False, + mask_blur: int = 0, mask_offset: int = 0) -> Image.Image: + if invert_output: + mask_np = np.array(mask_image) + mask_image = Image.fromarray(255 - mask_np) + + if mask_blur > 0: + mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur)) + + if mask_offset != 0: + filter_type = ImageFilter.MaxFilter if mask_offset > 0 else ImageFilter.MinFilter + size = abs(mask_offset) * 2 + 1 + for _ in range(abs(mask_offset)): + mask_image = mask_image.filter(filter_type(size)) + + return mask_image + +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): + if image.mode != 'L': + image = image.convert('L') + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0) + return image.squeeze() + +def apply_background_color(image: Image.Image, mask_image: Image.Image, + background_color: str = "Alpha") -> Image.Image: + bg_colors = { + "Alpha": None, + "black": (0, 0, 0), + "white": (255, 255, 255), + "green": (0, 255, 0), + "blue": (0, 0, 255), + "red": (255, 0, 0) + } + + rgba_image = image.copy().convert('RGBA') + rgba_image.putalpha(mask_image.convert('L')) + + if background_color != "Alpha": + bg_color = bg_colors[background_color] + bg_image = Image.new('RGBA', image.size, (*bg_color, 255)) + composite_image = Image.alpha_composite(bg_image, rgba_image) + return composite_image.convert('RGB') + + return rgba_image + +class Segment: + @classmethod + def INPUT_TYPES(cls): + tooltips = { + "prompt": "Enter the object or scene you want to segment. Use tag-style or natural language for more detailed prompts.", + "threshold": "Adjust mask detection strength (higher = more strict)", + "mask_blur": "Apply Gaussian blur to mask edges (0 = disabled)", + "mask_offset": "Expand/Shrink mask boundary (positive = expand, negative = shrink)", + "background_color": "Choose background color (Alpha = transparent)", + "invert_output": "Invert the mask output", + } + return { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"default": "", "multiline": True, "placeholder": "Object to segment", "tooltip": tooltips["prompt"]}), + "sam_model": (list(SAM_MODELS.keys()),), + "dino_model": (list(DINO_MODELS.keys()),), + }, + "optional": { + "threshold": ("FLOAT", {"default": 0.35, "min": 0.05, "max": 0.95, "step": 0.01, "tooltip": tooltips["threshold"]}), + "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", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}), + "invert_output": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "segment" + CATEGORY = "🧪AILab/🧽RMBG" + + def __init__(self): + from groundingdino.datasets import transforms as T + from groundingdino.util.utils import clean_state_dict + from groundingdino.util.slconfig import SLConfig + from groundingdino.models import build_model + + self.T = T + self.clean_state_dict = clean_state_dict + self.SLConfig = SLConfig + self.build_model = build_model + + def segment(self, image, prompt, sam_model, dino_model, threshold=0.35, + mask_blur=0, mask_offset=0, background_color="Alpha", + invert_output=False): + print(f'Processing create segment for: "{prompt}"...') + + image = Image.fromarray(np.clip(255. * image[0].cpu().numpy(), 0, 255).astype(np.uint8)).convert('RGBA') + dino_model = self.load_groundingdino(dino_model) + sam_model = self.load_sam(sam_model) + boxes = self.predict_boxes(dino_model, image, prompt, threshold) + + if boxes is None or boxes.shape[0] == 0: + print(f'No objects found for: "{prompt}"') + width, height = image.size + empty_mask = torch.zeros((1, height, width), dtype=torch.uint8, device="cpu") + return (empty_mask, empty_mask) + + masks = self.generate_masks(sam_model, image, boxes) + if masks is None: + print(f'Failed to generate mask for: "{prompt}"') + width, height = image.size + empty_mask = torch.zeros((1, height, width), dtype=torch.uint8, device="cpu") + return (empty_mask, empty_mask) + + mask_image = Image.fromarray((masks[1][0].numpy() * 255).astype(np.uint8)) + mask_image = process_mask(mask_image, invert_output, mask_blur, mask_offset) + + result_image = apply_background_color(image, mask_image, background_color) + + print(f'Successfully created segment for: "{prompt}"') + return (pil2tensor(result_image), image2mask(mask_image)) + + def load_sam(self, model_name): + sam_checkpoint_path = self.get_local_filepath( + SAM_MODELS[model_name]["model_url"], "sam") + model_type = SAM_MODELS[model_name]["model_type"] + + sam = sam_model_registry[model_type](checkpoint=sam_checkpoint_path) + sam_device = comfy.model_management.get_torch_device() + sam.to(device=sam_device) + sam.eval() + return sam + + def load_groundingdino(self, model_name): + import sys + from io import StringIO + temp_stdout = StringIO() + original_stdout = sys.stdout + sys.stdout = temp_stdout + + try: + dino_model_args = self.SLConfig.fromfile( + self.get_local_filepath( + DINO_MODELS[model_name]["config_url"], + "grounding-dino" + ) + ) + dino = self.build_model(dino_model_args) + checkpoint = torch.load( + self.get_local_filepath( + DINO_MODELS[model_name]["model_url"], + "grounding-dino" + ) + ) + dino.load_state_dict(self.clean_state_dict(checkpoint['model']), strict=False) + device = comfy.model_management.get_torch_device() + dino.to(device=device) + dino.eval() + return dino + finally: + output = temp_stdout.getvalue() + sys.stdout = original_stdout + + for line in output.split('\n'): + if 'error' in line.lower(): + print(line) + + def _load_dino_image(self, image_pil): + transform = self.T.Compose([ + self.T.RandomResize([800], max_size=1333), + self.T.ToTensor(), + self.T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ]) + image, _ = transform(image_pil, None) + return image + + def _get_grounding_output(self, model, image, caption, box_threshold): + caption = caption.lower().strip() + if not caption.endswith("."): + caption = caption + "." + device = comfy.model_management.get_torch_device() + image = image.to(device) + with torch.no_grad(): + outputs = model(image[None], captions=[caption]) + logits = outputs["pred_logits"].sigmoid()[0] + boxes = outputs["pred_boxes"][0] + logits_filt = logits.clone() + boxes_filt = boxes.clone() + filt_mask = logits_filt.max(dim=1)[0] > box_threshold + logits_filt = logits_filt[filt_mask] + boxes_filt = boxes_filt[filt_mask] + return boxes_filt.cpu() + + def predict_boxes(self, model, image, prompt, threshold): + dino_image = self._load_dino_image(image.convert("RGB")) + boxes_filt = self._get_grounding_output(model, dino_image, prompt, threshold) + H, W = image.size[1], image.size[0] + for i in range(boxes_filt.size(0)): + boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H]) + boxes_filt[i][:2] -= boxes_filt[i][2:] / 2 + boxes_filt[i][2:] += boxes_filt[i][:2] + return boxes_filt + + def generate_masks(self, model, image, boxes): + if boxes.shape[0] == 0: + return None + + if not hasattr(self, 'predictor'): + self.predictor = SamPredictor(model) + + image_np = np.array(image) + image_np_rgb = image_np[..., :3] + + self.predictor.set_image(image_np_rgb) + + transformed_boxes = self.predictor.transform.apply_boxes_torch(boxes, image_np.shape[:2]) + masks, _, _ = self.predictor.predict_torch( + point_coords=None, + point_labels=None, + boxes=transformed_boxes.to(comfy.model_management.get_torch_device()), + multimask_output=False + ) + + return create_tensor_output(image_np, masks.permute(1, 0, 2, 3).cpu().numpy(), boxes) + + + def get_local_filepath(self, url, dirname, local_file_name=None): + if not local_file_name: + local_file_name = os.path.basename(urlparse(url).path) + + destination = folder_paths.get_full_path(dirname, local_file_name) + if destination: + return destination + + folder = os.path.join(folder_paths.models_dir, dirname) + os.makedirs(folder, exist_ok=True) + + destination = os.path.join(folder, local_file_name) + if not os.path.exists(destination): + try: + download_url_to_file(url, destination) + except Exception as e: + if os.path.exists(destination): + os.remove(destination) + raise Exception(f'Failed to download model from {url}: {str(e)}') + return destination + +NODE_CLASS_MAPPINGS = { + "Segment": Segment +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Segment": "Segment (RMBG)" +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 16eeb1e..504dc68 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ComfyUI-RMBG" -version = "1.2.1" +version = "1.3.0 description = "A ComfyUI custom node designed for advanced image background removal utilizing multiple models, including RMBG-2.0, INSPYRENET, and BEN." authors = [ {name = "AILab", email = "ailab@mail.com"} @@ -14,6 +14,8 @@ dependencies = [ "tqdm>=4.65.0", "transformers>=4.35.0", "transparent-background>=1.2.4", + "segment-anything>=1.0.0", + "groundingdino>=1.0.0", ] requires-python = ">=3.8" readme = "README.md" diff --git a/requirements.txt b/requirements.txt index 8399ad2..cf87157 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,4 +5,7 @@ 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 +transparent-background>=1.2.4 +groundingdino-py>=0.4.0 +segment-anything>=1.0 +opencv-python>=4.7.0 \ No newline at end of file diff --git a/update.md b/update.md index 964ab62..6910255 100644 --- a/update.md +++ b/update.md @@ -1,139 +1,193 @@ -# ComfyUI-RMBG Update Log - -## v1.2.2 (2024/12/12) -![RMBG1 2 2](https://github.com/user-attachments/assets/cb7b1ad0-a2ca-4369-9401-54957af6c636) - -### Improvements -- Changed INSPYRENET model format from .pth to .safetensors for: - - Better security - - Faster loading speed (2-3x faster) - - Improved memory efficiency - - Better cross-platform compatibility -- Simplified node display name for better UI integration - -## v1.2.1 (2024/12/02) - -### New Features -- ANPG (animated PNG), AWEBP (animated WebP) and GIF supported. - -https://github.com/user-attachments/assets/40ec0b27-4fa2-4c99-9aea-5afad9ca62a5 - -### Bug Fixes -- Fixed video processing issue - -### Performance Improvements -- Enhanced batch processing in RMBG-2.0 model -- Added support for proper batch image handling -- Improved memory efficiency by optimizing image size handling - -### Technical Details -- Added original size preservation for maintaining aspect ratios -- Implemented proper batch tensor processing -- Improved error handling and code robustness -- Performance gains: - - Single image processing: ~5-10% improvement - - Batch processing: up to 30-50% improvement (depending on batch size and GPU) - -## 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 - -## v1.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 +# ComfyUI-RMBG Update Log + +## v1.3.0 (2024/12/23) + +### New Segment (RMBG) Node +- Text-Prompted Intelligent Object Segmentation + - Use natural language prompts (e.g., "a cat", "red car") to identify and segment target objects + - Support for multiple object detection and segmentation + - Perfect for precise object extraction and recognition tasks + +### Supported Models +- SAM (Segment Anything Model) + - sam_vit_h: 2.56GB - Highest accuracy + - sam_vit_l: 1.25GB - Balanced performance + - sam_vit_b: 375MB - Lightweight option +- GroundingDINO + - SwinT: 694MB - Fast and efficient + - SwinB: 938MB - Higher precision + +### Key Features +- Intuitive Parameter Controls + - Threshold: Adjust detection precision + - Mask Blur: Smooth edges + - Mask Offset: Expand or shrink selection + - Background Options: Alpha/Black/White/Green/Blue/Red +- Automatic Model Management + - Auto-download models on first use + - Smart GPU memory handling + +### Usage Examples +1. Tag-Style Prompts + - Single object: "cat" + - Multiple objects: "cat, dog, person" + - With attributes: "red car, blue shirt" + - Format: Use commas to separate multiple objects (e.g., "a, b, c") + +2. Natural Language Prompts + - Simple sentence: "a person wearing a red jacket" + - Complex scene: "a woman in a blue dress standing next to a car" + - With location: "a cat sitting on the sofa" + - Format: Write a natural descriptive sentence + +3. Tips for Better Results + - For Tag Style: + - Separate objects with commas: "chair, table, lamp" + - Add attributes before objects: "wooden chair, glass table" + - Keep it simple and clear + - For Natural Language: + - Use complete sentences + - Include details like color, position, action + - Be as descriptive as needed + - Parameter Adjustments: + - Threshold: 0.25-0.35 for broad detection, 0.45-0.55 for precision + - Use mask blur for smoother edges + - Adjust mask offset to fine-tune selection + +## v1.2.2 (2024/12/12) +![RMBG1 2 2](https://github.com/user-attachments/assets/cb7b1ad0-a2ca-4369-9401-54957af6c636) + +### Improvements +- Changed INSPYRENET model format from .pth to .safetensors for: + - Better security + - Faster loading speed (2-3x faster) + - Improved memory efficiency + - Better cross-platform compatibility +- Simplified node display name for better UI integration + +## v1.2.1 (2024/12/02) + +### New Features +- ANPG (animated PNG), AWEBP (animated WebP) and GIF supported. + +https://github.com/user-attachments/assets/40ec0b27-4fa2-4c99-9aea-5afad9ca62a5 + +### Bug Fixes +- Fixed video processing issue + +### Performance Improvements +- Enhanced batch processing in RMBG-2.0 model +- Added support for proper batch image handling +- Improved memory efficiency by optimizing image size handling + +### Technical Details +- Added original size preservation for maintaining aspect ratios +- Implemented proper batch tensor processing +- Improved error handling and code robustness +- Performance gains: + - Single image processing: ~5-10% improvement + - Batch processing: up to 30-50% improvement (depending on batch size and GPU) + +## 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 + +## v1.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