Add files via upload
This commit is contained in:
AI Lab
@@ -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) )
|
||||
|
||||
|
||||
## 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
|
||||
|
||||
|
||||
|
||||
|
||||
## 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
|
||||
|
||||
@@ -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
|
||||
<details>
|
||||
<summary><h2>About Models</h2></summary>
|
||||
|
||||
## 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
|
||||
</details>
|
||||
|
||||
|
||||
## 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
|
||||
|
||||
+405
-234
@@ -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)"
|
||||
}
|
||||
# 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)"
|
||||
}
|
||||
+27
-20
@@ -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*"]
|
||||
[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"
|
||||
+8
-8
@@ -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
|
||||
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
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
- Improved mask processing
|
||||
- Better detail preservation
|
||||
- Enhanced edge quality
|
||||
- More accurate segmentation
|
||||
|
||||
|
||||
|
||||
- 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
|
||||
|
||||
|
||||
|
||||
- Improved mask processing
|
||||
- Better detail preservation
|
||||
- Enhanced edge quality
|
||||
- More accurate segmentation
|
||||
|
||||
|
||||
|
||||
- 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
|
||||
|
||||
Reference in New Issue
Block a user