nodes/1038lab-ComfyUI-RMBG
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
AI Lab
2026-01-01 19:18:04 -08:00
committed by GitHub
parent 95d5e5c94f
commit d6fe942a1e
14 changed files with 4571 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/AILab_BiRefNet.py
+484
View File
@@ -0,0 +1,484 @@
# ComfyUI-RMBG
# This custom node for ComfyUI provides functionality for background removal using BiRefNet models.
#
# Model License Notice:
# - BiRefNet Models: Apache-2.0 License (https://huggingface.co/ZhengPeng7)
#
# 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, ImageFilter
from torchvision import transforms
import numpy as np
import folder_paths
from huggingface_hub import hf_hub_download
import sys
import importlib.util
from safetensors.torch import load_file
import cv2
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
MODEL_CONFIG = {
"BiRefNet-general": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet.py": "birefnet.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet-general.safetensors": "BiRefNet-general.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "General purpose model with balanced performance",
"default_res": 1024,
"max_res": 2048,
"min_res": 512
},
"BiRefNet_512x512": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet.py": "birefnet.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet_512x512.safetensors": "BiRefNet_512x512.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "Optimized for 512x512 resolution, faster processing",
"default_res": 512,
"max_res": 1024,
"min_res": 256,
"force_res": True
},
"BiRefNet-HR": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet.py": "birefnet.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet-HR.safetensors": "BiRefNet-HR.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "High resolution general purpose model",
"default_res": 2048,
"max_res": 2560,
"min_res": 1024
},
"BiRefNet-portrait": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet.py": "birefnet.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet-portrait.safetensors": "BiRefNet-portrait.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "Optimized for portrait/human matting",
"default_res": 1024,
"max_res": 2048,
"min_res": 512
},
"BiRefNet-matting": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet.py": "birefnet.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet-matting.safetensors": "BiRefNet-matting.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "General purpose matting model",
"default_res": 1024,
"max_res": 2048,
"min_res": 512
},
"BiRefNet-HR-matting": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet.py": "birefnet.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet-HR-matting.safetensors": "BiRefNet-HR-matting.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "High resolution matting model",
"default_res": 2048,
"max_res": 2560,
"min_res": 1024
},
"BiRefNet_lite": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet_lite.py": "birefnet_lite.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet_lite.safetensors": "BiRefNet_lite.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "Lightweight version for faster processing",
"default_res": 1024,
"max_res": 2048,
"min_res": 512
},
"BiRefNet_lite-2K": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet_lite.py": "birefnet_lite.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet_lite-2K.safetensors": "BiRefNet_lite-2K.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "Lightweight version optimized for 2K resolution",
"default_res": 2048,
"max_res": 2560,
"min_res": 1024
},
"BiRefNet_dynamic": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet.py": "birefnet.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet_dynamic.safetensors": "BiRefNet_dynamic.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "Dynamic model for high-resolution dichotomous image segmentation",
"default_res": 1024,
"max_res": 2048,
"min_res": 512
},
"BiRefNet_lite-matting": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet_lite.py": "birefnet_lite.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet_lite-matting.safetensors": "BiRefNet_lite-matting.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "Lightweight matting model for general purpose",
"default_res": 1024,
"max_res": 2048,
"min_res": 512
},
"BiRefNet_toonout": {
"repo_id": "1038lab/BiRefNet",
"files": {
"birefnet.py": "birefnet.py",
"BiRefNet_config.py": "BiRefNet_config.py",
"BiRefNet_toonout.safetensors": "BiRefNet_toonout.safetensors",
"config.json": "config.json"
},
"cache_dir": "BiRefNet",
"description": "A model to get a toon style outline from an image.",
"default_res": 1024,
"max_res": 2048,
"min_res": 512
}
}
# 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"[BiRefNet ERROR] {message}")
raise RuntimeError(message)
def refine_foreground(image_bchw, masks_b1hw):
b, c, h, w = image_bchw.shape
if b != masks_b1hw.shape[0]:
raise ValueError("images and masks must have the same batch size")
image_np = image_bchw.cpu().numpy()
mask_np = masks_b1hw.cpu().numpy()
refined_fg = []
for i in range(b):
mask = mask_np[i, 0]
thresh = 0.45
mask_binary = (mask > thresh).astype(np.float32)
edge_blur = cv2.GaussianBlur(mask_binary, (3, 3), 0)
transition_mask = np.logical_and(mask > 0.05, mask < 0.95)
alpha = 0.85
mask_refined = np.where(transition_mask,
alpha * mask + (1-alpha) * edge_blur,
mask_binary)
edge_region = np.logical_and(mask > 0.2, mask < 0.8)
mask_refined = np.where(edge_region,
mask_refined * 0.98,
mask_refined)
result = []
for c in range(image_np.shape[1]):
channel = image_np[i, c]
refined = channel * mask_refined
result.append(refined)
refined_fg.append(np.stack(result))
return torch.from_numpy(np.stack(refined_fg))
class BiRefNetModel:
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, MODEL_CONFIG[model_name]["cache_dir"])
def check_model_cache(self, 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_CONFIG[model_name]["files"].keys():
if not os.path.exists(os.path.join(cache_dir, 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):
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_CONFIG[model_name]["files"].keys():
print(f"Downloading {filename}...")
hf_hub_download(
repo_id=MODEL_CONFIG[model_name]["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")
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_filename = [k for k in MODEL_CONFIG[model_name]["files"].keys() if k.endswith('.py') and k != "BiRefNet_config.py"][0]
model_path = os.path.join(cache_dir, model_filename)
config_path = os.path.join(cache_dir, "BiRefNet_config.py")
weights_filename = [k for k in MODEL_CONFIG[model_name]["files"].keys() if k.endswith('.safetensors')][0]
weights_path = os.path.join(cache_dir, weights_filename)
try:
# Fix relative imports in model file
with open(model_path, 'r', encoding='utf-8') as f:
model_content = f.read()
model_content = model_content.replace("from .BiRefNet_config", "from BiRefNet_config")
with open(model_path, 'w', encoding='utf-8') as f:
f.write(model_content)
# Load config and model dynamically
spec = importlib.util.spec_from_file_location("BiRefNet_config", config_path)
config_module = importlib.util.module_from_spec(spec)
sys.modules["BiRefNet_config"] = config_module
spec.loader.exec_module(config_module)
spec = importlib.util.spec_from_file_location("birefnet", model_path)
model_module = importlib.util.module_from_spec(spec)
sys.modules["birefnet"] = model_module
spec.loader.exec_module(model_module)
# Initialize model
self.model = model_module.BiRefNet(config_module.BiRefNetConfig())
# Load weights
state_dict = load_file(weights_path)
self.model.load_state_dict(state_dict)
self.model.eval()
self.model.half()
torch.set_float32_matmul_precision('high')
self.model.to(device)
self.current_model_version = model_name
except Exception as e:
handle_model_error(f"Error loading BiRefNet model: {str(e)}")
def process_image(self, image, params):
try:
transform_image = transforms.Compose([
transforms.Resize((params["process_res"], params["process_res"]),
interpolation=transforms.InterpolationMode.BICUBIC),
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).half()
with torch.no_grad():
preds = self.model(input_tensor)
pred = preds[-1].sigmoid().cpu()
pred = pred[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
mask = pred_pil.resize((w, h), Image.BICUBIC)
return mask
except Exception as e:
handle_model_error(f"Error in BiRefNet processing: {str(e)}")
class BiRefNetRMBG:
def __init__(self):
self.model = BiRefNetModel()
@classmethod
def INPUT_TYPES(s):
tooltips = {
"image": "Input image to be processed for background removal.",
"model": "Select the BiRefNet model variant to use.",
"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).",
"invert_output": "Enable to invert both the image and mask output (useful for certain effects).",
"refine_foreground": "Use Fast Foreground Colour Estimation to optimize transparent background",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)"
}
return {
"required": {
"image": ("IMAGE", {"tooltip": tooltips["image"]}),
"model": (list(MODEL_CONFIG.keys()), {"tooltip": tooltips["model"]}),
},
"optional": {
"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"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"refine_foreground": ("BOOLEAN", {"default": False, "tooltip": tooltips["refine_foreground"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "process_image"
CATEGORY = "🧪AILab/🧽RMBG"
def process_image(self, image, model, **params):
try:
model_config = MODEL_CONFIG[model]
process_res = model_config.get("default_res", 1024)
if model_config.get("force_res", False):
base_res = 512
process_res = ((process_res + base_res - 1) // base_res) * base_res
else:
process_res = process_res // 32 * 32
print(f"Using {model} model with {process_res} resolution")
params["process_res"] = process_res
processed_images = []
processed_masks = []
cache_status, message = self.model.check_model_cache(model)
if not cache_status:
print(f"Cache check: {message}")
print("Downloading required model files...")
download_status, download_message = self.model.download_model(model)
if not download_status:
handle_model_error(download_message)
print("Model files downloaded successfully")
self.model.load_model(model)
for img in image:
mask = self.model.process_image(img, params)
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))
img_tensor = torch.from_numpy(np.array(tensor2pil(img))).permute(2, 0, 1).unsqueeze(0) / 255.0
mask_tensor = torch.from_numpy(np.array(mask)).unsqueeze(0).unsqueeze(0) / 255.0
if params.get("refine_foreground", False):
refined_fg = refine_foreground(img_tensor, mask_tensor)
refined_fg = tensor2pil(refined_fg[0].permute(1, 2, 0))
orig_image = tensor2pil(img)
r, g, b = refined_fg.split()
foreground = Image.merge('RGBA', (r, g, b, mask))
else:
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":
processed_images.append(pil2tensor(foreground))
else:
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
background_color = params.get("background_color", "#222222")
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, foreground)
processed_images.append(pil2tensor(composite_image.convert("RGB")))
processed_masks.append(pil2tensor(mask))
mask_images = []
for mask_tensor in processed_masks:
mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
mask_images.append(mask_image)
mask_image_output = torch.cat(mask_images, dim=0)
return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0), mask_image_output)
except Exception as e:
handle_model_error(f"Error in image processing: {str(e)}")
# Node Mapping
NODE_CLASS_MAPPINGS = {
"BiRefNetRMBG": BiRefNetRMBG
}
NODE_DISPLAY_NAME_MAPPINGS = {
"BiRefNetRMBG": "BiRefNet (RMBG)"
}
py/AILab_BodySegment.py
+251
View File
@@ -0,0 +1,251 @@
# 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 integration script follows GPL-3.0 License.
# When using or modifying this code, please respect both the original model licenses
# and this integration's license terms.
#
# Source: https://github.com/AILab-AI/ComfyUI-RMBG
import os
import torch
import torch.nn as nn
import numpy as np
from typing import Tuple, Union
from PIL import Image, ImageFilter
import onnxruntime
import folder_paths
from huggingface_hub import hf_hub_download
import shutil
from torchvision import transforms
def pil2tensor(image: Image.Image) -> torch.Tensor:
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
def tensor2pil(image: torch.Tensor) -> Image.Image:
return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8))
def image2mask(image: Image.Image) -> torch.Tensor:
if isinstance(image, Image.Image):
image = pil2tensor(image)
return image.squeeze()[..., 0]
def mask2image(mask: torch.Tensor) -> Image.Image:
if len(mask.shape) == 2:
mask = mask.unsqueeze(0)
return tensor2pil(mask)
def RGB2RGBA(image: Image.Image, mask: Union[Image.Image, torch.Tensor]) -> Image.Image:
if isinstance(mask, torch.Tensor):
mask = mask2image(mask)
if mask.size != image.size:
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
return Image.merge('RGBA', (*image.convert('RGB').split(), mask.convert('L')))
device = "cuda" if torch.cuda.is_available() else "cpu"
folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG"))
class BodySegment:
def __init__(self):
self.model = None
self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "body_segment")
self.model_file = "deeplabv3p-resnet50-human.onnx"
@classmethod
def INPUT_TYPES(cls):
available_classes = [
"Hair", "Glasses", "Top-clothes", "Bottom-clothes",
"Torso-skin", "Face", "Left-arm", "Right-arm",
"Left-leg", "Right-leg", "Left-foot", "Right-foot"
]
tooltips = {
"process_res": "Processing resolution (fixed at 512x512)",
"mask_blur": "Blur amount for mask edges",
"mask_offset": "Expand/Shrink mask boundary",
"invert_output": "Invert both image and mask output",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)"
}
return {
"required": {
"images": ("IMAGE",),
},
"optional": {
**{cls_name: ("BOOLEAN", {"default": False})
for cls_name in available_classes},
"mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}),
"mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
},
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "segment_body"
CATEGORY = "🧪AILab/🧽RMBG"
def check_model_cache(self):
model_path = os.path.join(self.cache_dir, self.model_file)
if not os.path.exists(model_path):
return False, "Model file not found"
return True, "Model cache verified"
def clear_model(self):
if self.model is not None:
del self.model
self.model = None
def download_model_files(self):
model_id = "Metal3d/deeplabv3p-resnet50-human"
os.makedirs(self.cache_dir, exist_ok=True)
print("Downloading body segmentation model...")
try:
downloaded_path = hf_hub_download(
repo_id=model_id,
filename=self.model_file,
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, self.model_file)
shutil.move(downloaded_path, target_path)
return True, "Model file downloaded successfully"
except Exception as e:
return False, f"Error downloading model file: {str(e)}"
def segment_body(self, images, mask_blur=0, mask_offset=0, background="Alpha", background_color="#222222", invert_output=False, **class_selections):
try:
# Check and download model if needed
cache_status, message = self.check_model_cache()
if not cache_status:
print(f"Cache check: {message}")
download_status, download_message = self.download_model_files()
if not download_status:
raise RuntimeError(download_message)
# Load model if needed
if self.model is None:
self.model = onnxruntime.InferenceSession(
os.path.join(self.cache_dir, self.model_file)
)
# Class mapping
class_map = {
"Hair": 2, "Glasses": 4, "Top-clothes": 5,
"Bottom-clothes": 9, "Torso-skin": 10, "Face": 13,
"Left-arm": 14, "Right-arm": 15, "Left-leg": 16,
"Right-leg": 17, "Left-foot": 18, "Right-foot": 19
}
# Get selected classes
selected_classes = [name for name, selected in class_selections.items() if selected]
if not selected_classes:
selected_classes = ["Face", "Hair", "Top-clothes", "Bottom-clothes"]
batch_tensor = []
batch_masks = []
for image in images:
orig_image = tensor2pil(image)
w, h = orig_image.size
# Resize to 512x512 (model requirement)
input_image = orig_image.resize((512, 512))
input_array = np.array(input_image).astype(np.float32) / 127.5 - 1
# Add batch dimension
input_array = np.expand_dims(input_array, axis=0)
# Run inference
input_name = self.model.get_inputs()[0].name
output_name = self.model.get_outputs()[0].name
result = self.model.run([output_name], {input_name: input_array})
# Process results
result = np.array(result[0])
pred_seg = result.argmax(axis=3).squeeze(0)
# Combine selected class masks
combined_mask = np.zeros_like(pred_seg, dtype=np.float32)
for class_name in selected_classes:
mask = (pred_seg == class_map[class_name]).astype(np.float32)
combined_mask = np.clip(combined_mask + mask, 0, 1)
# Convert to PIL and resize back to original size
mask_image = Image.fromarray((combined_mask * 255).astype(np.uint8))
mask_image = mask_image.resize((w, h), Image.Resampling.LANCZOS)
if mask_blur > 0:
mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur))
if mask_offset != 0:
if mask_offset > 0:
mask_image = mask_image.filter(ImageFilter.MaxFilter(size=mask_offset * 2 + 1))
else:
mask_image = mask_image.filter(ImageFilter.MinFilter(size=-mask_offset * 2 + 1))
if invert_output:
mask_image = Image.fromarray(255 - np.array(mask_image))
# Handle background color
if background == "Alpha":
rgba_image = RGB2RGBA(orig_image, mask_image)
result_image = pil2tensor(rgba_image)
else:
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba_image = RGB2RGBA(orig_image, mask_image)
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
result_image = pil2tensor(composite_image.convert('RGB'))
batch_tensor.append(result_image)
batch_masks.append(pil2tensor(mask_image))
# Create mask image for visualization
mask_images = []
for mask_tensor in batch_masks:
# Convert mask to RGB image format for visualization
mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
mask_images.append(mask_image)
mask_image_output = torch.cat(mask_images, dim=0)
# Prepare final output
batch_tensor = torch.cat(batch_tensor, dim=0)
batch_masks = torch.cat(batch_masks, dim=0)
return (batch_tensor, batch_masks, mask_image_output)
except Exception as e:
self.clear_model()
raise RuntimeError(f"Error in Body Segmentation processing: {str(e)}")
finally:
self.clear_model()
NODE_CLASS_MAPPINGS = {
"BodySegment": BodySegment
}
NODE_DISPLAY_NAME_MAPPINGS = {
"BodySegment": "Body Segment (RMBG)"
}
py/AILab_ClothSegment.py
+288
View File
@@ -0,0 +1,288 @@
# 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 integration script follows GPL-3.0 License.
# When using or modifying this code, please respect both the original model licenses
# and this integration's license terms.
#
# Source: https://github.com/AILab-AI/ComfyUI-RMBG
import os
import torch
import torch.nn as nn
import numpy as np
from typing import Tuple, Union
from PIL import Image, ImageFilter
from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
import folder_paths
from huggingface_hub import hf_hub_download
import shutil
from torchvision import transforms
def pil2tensor(image: Image.Image) -> torch.Tensor:
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
def tensor2pil(image: torch.Tensor) -> Image.Image:
return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8))
def image2mask(image: Image.Image) -> torch.Tensor:
if isinstance(image, Image.Image):
image = pil2tensor(image)
return image.squeeze()[..., 0]
def mask2image(mask: torch.Tensor) -> Image.Image:
if len(mask.shape) == 2:
mask = mask.unsqueeze(0)
return tensor2pil(mask)
def RGB2RGBA(image: Image.Image, mask: Union[Image.Image, torch.Tensor]) -> Image.Image:
if isinstance(mask, torch.Tensor):
mask = mask2image(mask)
if mask.size != image.size:
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
return Image.merge('RGBA', (*image.convert('RGB').split(), mask.convert('L')))
device = "cuda" if torch.cuda.is_available() else "cpu"
folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG"))
AVAILABLE_MODELS = {
"segformer_b2_clothes": "1038lab/segformer_clothes"
}
class ClothesSegment:
def __init__(self):
self.processor = None
self.model = None
self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "segformer_clothes")
@classmethod
def INPUT_TYPES(cls):
available_classes = ["Hat", "Hair", "Face", "Sunglasses", "Upper-clothes", "Skirt", "Dress", "Belt", "Pants", "Left-arm", "Right-arm", "Left-leg", "Right-leg", "Bag", "Scarf", "Left-shoe", "Right-shoe","Background"]
tooltips = {
"process_res": "Processing resolution (higher = more VRAM)",
"mask_blur": "Blur amount for mask edges",
"mask_offset": "Expand/Shrink mask boundary",
"invert_output": "Invert both image and mask output",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)"
}
return {
"required": {
"images": ("IMAGE",),
},
"optional": {
**{cls_name: ("BOOLEAN", {"default": False})
for cls_name in available_classes},
"process_res": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 32, "tooltip": tooltips["process_res"]}),
"mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}),
"mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
},
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "segment_clothes"
CATEGORY = "🧪AILab/🧽RMBG"
def check_model_cache(self):
if not os.path.exists(self.cache_dir):
return False, "Model directory not found"
required_files = [
'config.json',
'model.safetensors',
'preprocessor_config.json'
]
missing_files = [f for f in required_files if not os.path.exists(os.path.join(self.cache_dir, f))]
if missing_files:
return False, f"Required model files missing: {', '.join(missing_files)}"
return True, "Model cache verified"
def clear_model(self):
if self.model is not None:
self.model.cpu()
del self.model
self.model = None
self.processor = None
torch.cuda.empty_cache()
def download_model_files(self):
model_id = AVAILABLE_MODELS["segformer_b2_clothes"]
model_files = {
'config.json': 'config.json',
'model.safetensors': 'model.safetensors',
'preprocessor_config.json': 'preprocessor_config.json'
}
os.makedirs(self.cache_dir, exist_ok=True)
print(f"Downloading Clothes Segformer model files...")
try:
for save_name, repo_path in model_files.items():
print(f"Downloading {save_name}...")
downloaded_path = hf_hub_download(
repo_id=model_id,
filename=repo_path,
local_dir=self.cache_dir,
local_dir_use_symlinks=False
)
if os.path.dirname(downloaded_path) != self.cache_dir:
target_path = os.path.join(self.cache_dir, save_name)
shutil.move(downloaded_path, target_path)
return True, "Model files downloaded successfully"
except Exception as e:
return False, f"Error downloading model files: {str(e)}"
def segment_clothes(self, images, process_res=1024, mask_blur=0, mask_offset=0, background="Alpha", background_color="#222222", invert_output=False, **class_selections):
try:
# Check and download model if needed
cache_status, message = self.check_model_cache()
if not cache_status:
print(f"Cache check: {message}")
download_status, download_message = self.download_model_files()
if not download_status:
raise RuntimeError(download_message)
# Load model if needed
if self.processor is None:
self.processor = SegformerImageProcessor.from_pretrained(self.cache_dir)
self.model = AutoModelForSemanticSegmentation.from_pretrained(self.cache_dir)
self.model.eval()
for param in self.model.parameters():
param.requires_grad = False
self.model.to(device)
# Class mapping for segmentation
class_map = {
"Background": 0, "Hat": 1, "Hair": 2, "Sunglasses": 3,
"Upper-clothes": 4, "Skirt": 5, "Pants": 6, "Dress": 7,
"Belt": 8, "Left-shoe": 9, "Right-shoe": 10, "Face": 11,
"Left-leg": 12, "Right-leg": 13, "Left-arm": 14, "Right-arm": 15,
"Bag": 16, "Scarf": 17
}
# Get selected classes
selected_classes = [name for name, selected in class_selections.items() if selected]
if not selected_classes:
selected_classes = ["Upper-clothes"]
# Image preprocessing
transform_image = transforms.Compose([
transforms.Resize((process_res, process_res)),
transforms.ToTensor(),
])
batch_tensor = []
batch_masks = []
for image in images:
orig_image = tensor2pil(image)
w, h = orig_image.size
input_tensor = transform_image(orig_image)
if input_tensor.shape[0] == 4:
input_tensor = input_tensor[:3]
input_tensor = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(input_tensor)
input_tensor = input_tensor.unsqueeze(0).to(device)
with torch.no_grad():
outputs = self.model(input_tensor)
logits = outputs.logits.cpu()
upsampled_logits = nn.functional.interpolate(
logits,
size=(h, w),
mode="bilinear",
align_corners=False,
)
pred_seg = upsampled_logits.argmax(dim=1)[0]
# Combine selected class masks
combined_mask = None
for class_name in selected_classes:
mask = (pred_seg == class_map[class_name]).float()
if combined_mask is None:
combined_mask = mask
else:
combined_mask = torch.clamp(combined_mask + mask, 0, 1)
# Convert mask to PIL for processing
mask_image = Image.fromarray((combined_mask.numpy() * 255).astype(np.uint8))
if mask_blur > 0:
mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur))
if mask_offset != 0:
if mask_offset > 0:
mask_image = mask_image.filter(ImageFilter.MaxFilter(size=mask_offset * 2 + 1))
else:
mask_image = mask_image.filter(ImageFilter.MinFilter(size=-mask_offset * 2 + 1))
if invert_output:
mask_image = Image.fromarray(255 - np.array(mask_image))
# Handle background color
if background == "Alpha":
rgba_image = RGB2RGBA(orig_image, mask_image)
result_image = pil2tensor(rgba_image)
else:
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba_image = RGB2RGBA(orig_image, mask_image)
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
result_image = pil2tensor(composite_image.convert('RGB'))
batch_tensor.append(result_image)
batch_masks.append(pil2tensor(mask_image))
# Create mask image for visualization
mask_images = []
for mask_tensor in batch_masks:
# Convert mask to RGB image format for visualization
mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
mask_images.append(mask_image)
mask_image_output = torch.cat(mask_images, dim=0)
# Prepare final output
batch_tensor = torch.cat(batch_tensor, dim=0)
batch_masks = torch.cat(batch_masks, dim=0)
return (batch_tensor, batch_masks, mask_image_output)
except Exception as e:
self.clear_model()
raise RuntimeError(f"Error in Clothes Segformer processing: {str(e)}")
finally:
if not getattr(self.model, "training", False):
self.clear_model()
NODE_CLASS_MAPPINGS = {
"ClothesSegment": ClothesSegment
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ClothesSegment": "Clothes Segment (RMBG)"
}
py/AILab_FaceSegment.py
+293
View File
@@ -0,0 +1,293 @@
# ComfyUI-RMBG
# This custom node for ComfyUI provides functionality for face parsing using Segformer model.
#
# This integration script follows GPL-3.0 License.
# When using or modifying this code, please respect both the original model licenses
# and this integration's license terms.
#
# Source: https://github.com/AILab-AI/ComfyUI-RMBG
import os
import torch
import torch.nn as nn
import numpy as np
from typing import Tuple, Union
from PIL import Image, ImageFilter
from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
import folder_paths
from huggingface_hub import hf_hub_download
import shutil
from torchvision import transforms
def pil2tensor(image: Image.Image) -> torch.Tensor:
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
def tensor2pil(image: torch.Tensor) -> Image.Image:
return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8))
def image2mask(image: Image.Image) -> torch.Tensor:
if isinstance(image, Image.Image):
image = pil2tensor(image)
return image.squeeze()[..., 0]
def mask2image(mask: torch.Tensor) -> Image.Image:
if len(mask.shape) == 2:
mask = mask.unsqueeze(0)
return tensor2pil(mask)
def RGB2RGBA(image: Image.Image, mask: Union[Image.Image, torch.Tensor]) -> Image.Image:
if isinstance(mask, torch.Tensor):
mask = mask2image(mask)
if mask.size != image.size:
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
return Image.merge('RGBA', (*image.convert('RGB').split(), mask.convert('L')))
device = "cuda" if torch.cuda.is_available() else "cpu"
folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG"))
AVAILABLE_MODELS = {
"face_parsing": "1038lab/segformer_face"
}
class FaceSegment:
def __init__(self):
self.processor = None
self.model = None
self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "segformer_face")
@classmethod
def INPUT_TYPES(cls):
available_classes = [
"Skin", "Nose", "Eyeglasses", "Left-eye", "Right-eye",
"Left-eyebrow", "Right-eyebrow", "Left-ear", "Right-ear", "Mouth",
"Upper-lip", "Lower-lip", "Hair", "Earring", "Neck",
]
tooltips = {
"process_res": "Processing resolution (higher = more VRAM)",
"mask_blur": "Blur amount for mask edges",
"mask_offset": "Expand/Shrink mask boundary",
"invert_output": "Invert both image and mask output",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)"
}
return {
"required": {
"images": ("IMAGE",),
},
"optional": {
**{cls_name: ("BOOLEAN", {"default": False})
for cls_name in available_classes},
"process_res": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 32, "tooltip": tooltips["process_res"]}),
"mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}),
"mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
},
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "segment_face"
CATEGORY = "🧪AILab/🧽RMBG"
def check_model_cache(self):
if not os.path.exists(self.cache_dir):
return False, "Model directory not found"
required_files = [
'config.json',
'model.safetensors',
'preprocessor_config.json'
]
missing_files = [f for f in required_files if not os.path.exists(os.path.join(self.cache_dir, f))]
if missing_files:
return False, f"Required model files missing: {', '.join(missing_files)}"
return True, "Model cache verified"
def clear_model(self):
if self.model is not None:
self.model.cpu()
del self.model
self.model = None
self.processor = None
torch.cuda.empty_cache()
def download_model_files(self):
model_id = AVAILABLE_MODELS["face_parsing"]
model_files = {
'config.json': 'config.json',
'model.safetensors': 'model.safetensors',
'preprocessor_config.json': 'preprocessor_config.json'
}
os.makedirs(self.cache_dir, exist_ok=True)
print(f"Downloading face parsing model files...")
try:
for save_name, repo_path in model_files.items():
print(f"Downloading {save_name}...")
downloaded_path = hf_hub_download(
repo_id=model_id,
filename=repo_path,
local_dir=self.cache_dir,
local_dir_use_symlinks=False
)
if os.path.dirname(downloaded_path) != self.cache_dir:
target_path = os.path.join(self.cache_dir, save_name)
shutil.move(downloaded_path, target_path)
return True, "Model files downloaded successfully"
except Exception as e:
return False, f"Error downloading model files: {str(e)}"
def segment_face(self, images, process_res=512, mask_blur=0, mask_offset=0, background="Alpha", background_color="#222222", invert_output=False, **class_selections):
try:
# Check and download model if needed
cache_status, message = self.check_model_cache()
if not cache_status:
print(f"Cache check: {message}")
download_status, download_message = self.download_model_files()
if not download_status:
raise RuntimeError(download_message)
# Load model if needed
if self.processor is None:
self.processor = SegformerImageProcessor.from_pretrained(self.cache_dir)
self.model = AutoModelForSemanticSegmentation.from_pretrained(self.cache_dir)
self.model.eval()
for param in self.model.parameters():
param.requires_grad = False
self.model.to(device)
# Class mapping for segmentation
class_map = {
"Background": 0, "Skin": 1, "Nose": 2, "Eyeglasses": 3,
"Left-eye": 4, "Right-eye": 5, "Left-eyebrow": 6, "Right-eyebrow": 7,
"Left-ear": 8, "Right-ear": 9, "Mouth": 10, "Upper-lip": 11,
"Lower-lip": 12, "Hair": 13, "Hat": 14, "Earring": 15,
"Necklace": 16, "Neck": 17, "Clothing": 18
}
# Get selected classes
selected_classes = [name for name, selected in class_selections.items() if selected]
if not selected_classes:
selected_classes = ["Skin", "Nose", "Left-eye", "Right-eye", "Mouth"]
# Validate selected classes
invalid_classes = [cls for cls in selected_classes if cls not in class_map]
if invalid_classes:
raise ValueError(f"Invalid class selections: {', '.join(invalid_classes)}. Valid classes are: {', '.join(class_map.keys())}")
# Image preprocessing
transform_image = transforms.Compose([
transforms.Resize((process_res, process_res)),
transforms.ToTensor(),
])
batch_tensor = []
batch_masks = []
for image in images:
orig_image = tensor2pil(image)
w, h = orig_image.size
input_tensor = transform_image(orig_image)
if input_tensor.shape[0] == 4:
input_tensor = input_tensor[:3]
input_tensor = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(input_tensor)
input_tensor = input_tensor.unsqueeze(0).to(device)
with torch.no_grad():
outputs = self.model(input_tensor)
logits = outputs.logits.cpu()
upsampled_logits = nn.functional.interpolate(
logits,
size=(h, w),
mode="bilinear",
align_corners=False,
)
pred_seg = upsampled_logits.argmax(dim=1)[0]
# Combine selected class masks
combined_mask = None
for class_name in selected_classes:
mask = (pred_seg == class_map[class_name]).float()
if combined_mask is None:
combined_mask = mask
else:
combined_mask = torch.clamp(combined_mask + mask, 0, 1)
# Convert mask to PIL for processing
mask_image = Image.fromarray((combined_mask.numpy() * 255).astype(np.uint8))
if mask_blur > 0:
mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur))
if mask_offset != 0:
if mask_offset > 0:
mask_image = mask_image.filter(ImageFilter.MaxFilter(size=mask_offset * 2 + 1))
else:
mask_image = mask_image.filter(ImageFilter.MinFilter(size=-mask_offset * 2 + 1))
if invert_output:
mask_image = Image.fromarray(255 - np.array(mask_image))
# Handle background color
if background == "Alpha":
rgba_image = RGB2RGBA(orig_image, mask_image)
result_image = pil2tensor(rgba_image)
else:
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba_image = RGB2RGBA(orig_image, mask_image)
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
result_image = pil2tensor(composite_image.convert('RGB'))
batch_tensor.append(result_image)
batch_masks.append(pil2tensor(mask_image))
# Create mask image for visualization
mask_images = []
for mask_tensor in batch_masks:
# Convert mask to RGB image format for visualization
mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
mask_images.append(mask_image)
mask_image_output = torch.cat(mask_images, dim=0)
# Prepare final output
batch_tensor = torch.cat(batch_tensor, dim=0)
batch_masks = torch.cat(batch_masks, dim=0)
return (batch_tensor, batch_masks, mask_image_output)
except Exception as e:
self.clear_model()
raise RuntimeError(f"Error in Face Parsing processing: {str(e)}")
finally:
if self.model is not None and not self.model.training:
self.clear_model()
NODE_CLASS_MAPPINGS = {
"FaceSegment": FaceSegment
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FaceSegment": "Face Segment (RMBG)"
}
py/AILab_FashionSegment.py
+378
View File
@@ -0,0 +1,378 @@
# ComfyUI-RMBG
# This custom node for ComfyUI provides functionality for fashion segmentation using segformer-b3-fashion model.
# It leverages deep learning techniques to process images and generate masks for fashion items segmentation.
#
# This integration script follows GPL-3.0 License.
# When using or modifying this code, please respect both the original model licenses
# and this integration's license terms.
#
# Source: https://github.com/AILab-AI/ComfyUI-RMBG
import os
import torch
import torch.nn as nn
import numpy as np
from typing import Tuple, Union
from PIL import Image, ImageFilter
from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
import folder_paths
from huggingface_hub import hf_hub_download
import shutil
from torchvision import transforms
def pil2tensor(image: Image.Image) -> torch.Tensor:
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
def tensor2pil(image: torch.Tensor) -> Image.Image:
return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8))
def RGB2RGBA(image: Image.Image, mask: Union[Image.Image, torch.Tensor]) -> Image.Image:
if isinstance(mask, torch.Tensor):
mask = mask2image(mask)
if mask.size != image.size:
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
return Image.merge('RGBA', (*image.convert('RGB').split(), mask.convert('L')))
def mask2image(mask: torch.Tensor) -> Image.Image:
if len(mask.shape) == 2:
mask = mask.unsqueeze(0)
return tensor2pil(mask)
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 = {
"segformer_fashion": "1038lab/segformer_fashion"
}
class FashionSegmentAccessories:
@classmethod
def INPUT_TYPES(cls):
accessories_classes = [
# Head accessories
"hat",
"glasses",
"headband, head covering, hair accessory",
# Neck and upper body accessories
"scarf",
"tie",
# Hand accessories
"glove",
"watch",
# Waist accessories
"belt",
# Leg accessories
"leg warmer",
# Other accessories
"bag, wallet",
"umbrella"
]
details_classes = [
# Upper body details
"collar",
"lapel",
"neckline",
"epaulette",
"pocket",
# Decorative details
"buckle",
"zipper",
"applique",
"bow",
"flower",
"bead",
"fringe",
"ribbon",
"rivet",
"ruffle",
"sequin",
"tassel"
]
return {
"required": {},
"optional": {
**{cls_name: ("BOOLEAN", {"default": False})
for cls_name in accessories_classes + details_classes},
},
}
RETURN_TYPES = ("ACCESSORIES_OPTIONS",)
RETURN_NAMES = ("accessories_options",)
FUNCTION = "get_options"
CATEGORY = "🧪AILab/🧽RMBG"
def get_options(self, **class_selections):
selected = [name for name, selected in class_selections.items() if selected]
return (selected,)
class FashionSegmentClothing:
def __init__(self):
self.processor = None
self.model = None
self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "segformer_fashion")
self.class_map = {
"Unlabelled": 0, "shirt, blouse": 1, "top, t-shirt, sweatshirt": 2, "sweater": 3,
"cardigan": 4, "jacket": 5, "vest": 6, "pants": 7, "shorts": 8, "skirt": 9, "coat": 10,
"dress": 11, "jumpsuit": 12, "cape": 13,
"glasses": 14, "hat": 15, "headband, head covering, hair accessory": 16, "tie": 17, "glove": 18,
"watch": 19, "belt": 20, "leg warmer": 21, "tights, stockings": 22,
"sock": 23, "shoe": 24, "bag, wallet": 25, "scarf": 26, "umbrella": 27,
"hood": 28, "collar": 29, "lapel": 30, "epaulette": 31, "sleeve": 32,
"pocket": 33, "neckline": 34, "buckle": 35, "zipper": 36, "applique": 37,
"bead": 38, "bow": 39, "flower": 40, "fringe": 41, "ribbon": 42,
"rivet": 43, "ruffle": 44, "sequin": 45, "tassel": 46
}
@classmethod
def INPUT_TYPES(cls):
clothing_classes = [
# Upper body
"coat",
"jacket",
"cardigan",
"vest",
"sweater",
"hood",
"shirt, blouse",
"top, t-shirt, sweatshirt",
"sleeve",
# Full body
"dress",
"jumpsuit",
"cape",
# Lower body
"pants",
"shorts",
"skirt",
# Socks and shoes
"tights, stockings",
"sock",
"shoe"
]
tooltips = {
"accessories_options": "Select the accessories to be segmented",
"process_res": "Processing resolution (higher = more VRAM)",
"mask_blur": "Blur amount for mask edges",
"mask_offset": "Expand/Shrink mask boundary",
"invert_output": "Invert both image and mask output",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)",
}
return {
"required": {
"images": ("IMAGE",),
},
"optional": {
"accessories_options": ("ACCESSORIES_OPTIONS",),
**{cls_name: ("BOOLEAN", {"default": False,})
for cls_name in clothing_classes},
"process_res": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 32, "tooltip": tooltips["process_res"]}),
"mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}),
"mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
},
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "segment_fashion"
CATEGORY = "🧪AILab/🧽RMBG"
def check_model_cache(self):
if not os.path.exists(self.cache_dir):
return False, "Model directory not found"
required_files = [
'config.json',
'model.safetensors',
'preprocessor_config.json'
]
missing_files = [f for f in required_files if not os.path.exists(os.path.join(self.cache_dir, f))]
if missing_files:
return False, f"Missing required model files: {', '.join(missing_files)}"
return True, "Model cache verified"
def clear_model(self):
if self.model is not None:
self.model.cpu()
del self.model
self.model = None
self.processor = None
torch.cuda.empty_cache()
def download_model_files(self):
model_id = AVAILABLE_MODELS["segformer_fashion"]
model_files = {
'config.json': 'config.json',
'model.safetensors': 'model.safetensors',
'preprocessor_config.json': 'preprocessor_config.json'
}
os.makedirs(self.cache_dir, exist_ok=True)
print(f"Downloading fashion segmentation model files...")
try:
for save_name, repo_path in model_files.items():
print(f"Downloading {save_name}...")
downloaded_path = hf_hub_download(
repo_id=model_id,
filename=repo_path,
local_dir=self.cache_dir,
local_dir_use_symlinks=False
)
if os.path.dirname(downloaded_path) != self.cache_dir:
target_path = os.path.join(self.cache_dir, save_name)
shutil.move(downloaded_path, target_path)
return True, "Model files downloaded successfully"
except Exception as e:
return False, f"Error downloading model files: {str(e)}"
def segment_fashion(self, images, accessories_options=None, process_res=512, mask_blur=0, mask_offset=0,
background="Alpha", background_color="#222222", invert_output=False, **class_selections):
if accessories_options is None:
accessories_options = []
try:
# Check and download model
cache_status, message = self.check_model_cache()
if not cache_status:
print(f"Cache check: {message}")
download_status, download_message = self.download_model_files()
if not download_status:
raise RuntimeError(download_message)
# Load model if needed
if self.processor is None:
self.processor = SegformerImageProcessor.from_pretrained(self.cache_dir)
self.model = AutoModelForSemanticSegmentation.from_pretrained(self.cache_dir)
self.model.eval()
for param in self.model.parameters():
param.requires_grad = False
self.model.to(device)
# Get all selected classes
selected_classes = []
# Add clothing selections
selected_classes.extend([name for name, selected in class_selections.items() if selected])
# Add accessories selections
selected_classes.extend(accessories_options)
if not selected_classes:
selected_classes = ["shirt, blouse"]
transform_image = transforms.Compose([
transforms.Resize((process_res, process_res)),
transforms.ToTensor(),
])
batch_tensor = []
batch_masks = []
for image in images:
orig_image = tensor2pil(image)
w, h = orig_image.size
inputs = self.processor(images=orig_image, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = self.model(**inputs)
logits = outputs.logits.cpu()
upsampled_logits = nn.functional.interpolate(
logits,
size=(h, w),
mode="bilinear",
align_corners=False,
)
pred_seg = upsampled_logits.argmax(dim=1)[0]
# Merge masks for selected classes
combined_mask = None
for class_name in selected_classes:
mask = (pred_seg == self.class_map[class_name]).float()
if combined_mask is None:
combined_mask = mask
else:
combined_mask = torch.clamp(combined_mask + mask, 0, 1)
# Convert mask to PIL for processing
mask_image = Image.fromarray((combined_mask.numpy() * 255).astype(np.uint8))
if mask_blur > 0:
mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur))
if mask_offset != 0:
if mask_offset > 0:
mask_image = mask_image.filter(ImageFilter.MaxFilter(size=mask_offset * 2 + 1))
else:
mask_image = mask_image.filter(ImageFilter.MinFilter(size=-mask_offset * 2 + 1))
if invert_output:
mask_image = Image.fromarray(255 - np.array(mask_image))
# Process background color
if background == "Alpha":
rgba_image = RGB2RGBA(orig_image, mask_image)
result_image = pil2tensor(rgba_image)
else:
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba_image = RGB2RGBA(orig_image, mask_image)
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
result_image = pil2tensor(composite_image.convert('RGB'))
batch_tensor.append(result_image)
batch_masks.append(pil2tensor(mask_image))
# Create mask image for visualization
mask_images = []
for mask_tensor in batch_masks:
# Convert mask to RGB image format for visualization
mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
mask_images.append(mask_image)
mask_image_output = torch.cat(mask_images, dim=0)
batch_tensor = torch.cat(batch_tensor, dim=0)
batch_masks = torch.cat(batch_masks, dim=0)
return (batch_tensor, batch_masks, mask_image_output)
except Exception as e:
self.clear_model()
raise RuntimeError(f"Error in fashion segmentation: {str(e)}")
finally:
if self.model is not None and not self.model.training:
self.clear_model()
def __del__(self):
self.clear_model()
NODE_CLASS_MAPPINGS = {
"FashionSegmentAccessories": FashionSegmentAccessories,
"FashionSegmentClothing": FashionSegmentClothing
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FashionSegmentAccessories": "Accessories Segment (RMBG)",
"FashionSegmentClothing": "Fashion Segment (RMBG)"
}
py/AILab_InpaintTools.py
+146
View File
@@ -0,0 +1,146 @@
# ComfyUI-RMBG
#
# AILab Inpaint Tools
# A collection of specialized nodes for inpainting tasks in ComfyUI.
# Features a set of utilities for mask processing, latent conditioning, and inpainting workflows.
#
# 1. Inpaint Nodes:
# - AILab_ReferenceLatentMask: A node for inpainting tasks with the Flux Kontext model, using a reference latent and mask for precise region conditioning
#
# License: GPL-3.0
# These nodes are crafted to streamline common image and mask operations within ComfyUI workflows.
import torch
import node_helpers
def expand_mask(mask, expand_amount):
if expand_amount == 0:
return mask
import torch.nn.functional as F
binary_mask = (mask > 0.5).float()
kernel_size = abs(expand_amount) * 2 + 1
kernel_size = max(3, kernel_size)
kernel = torch.ones(1, 1, kernel_size, kernel_size, device=mask.device)
if expand_amount > 0:
expanded = F.conv2d(
binary_mask.reshape(-1, 1, mask.shape[-2], mask.shape[-1]),
kernel,
padding=kernel_size // 2
)
result = (expanded > 0).float()
else:
eroded = F.conv2d(
binary_mask.reshape(-1, 1, mask.shape[-2], mask.shape[-1]),
kernel,
padding=kernel_size // 2
)
result = (eroded >= kernel_size * kernel_size).float()
if len(mask.shape) == 3:
result = result.squeeze(1)
return result
def blur_mask(mask, blur_amount):
if blur_amount == 0:
return mask
import torch.nn.functional as F
import math
x = mask.reshape(-1, 1, mask.shape[-2], mask.shape[-1])
kernel_size = max(3, math.ceil(blur_amount * 3) * 2 + 1)
sigma = blur_amount
half_kernel = kernel_size // 2
grid = torch.arange(-half_kernel, half_kernel + 1, device=mask.device).float()
gaussian = torch.exp(-0.5 * (grid / sigma) ** 2)
gaussian = gaussian / gaussian.sum()
gaussian_x = gaussian.view(1, 1, 1, kernel_size)
gaussian_y = gaussian.view(1, 1, kernel_size, 1)
blurred = F.conv2d(x, gaussian_x, padding=(0, half_kernel))
blurred = F.conv2d(blurred, gaussian_y, padding=(half_kernel, 0))
if len(mask.shape) == 3:
blurred = blurred.squeeze(1)
return blurred
class AILab_ReferenceLatentMask:
@classmethod
def INPUT_TYPES(cls):
tooltips = {
"conditioning": "Base conditioning input for inpainting task",
"latent": "Encoded latent from VAE",
"mask": "Area to inpaint (white regions)",
"expand": "Grow mask (+) or shrink mask (-)",
"blur": "Soften mask edges",
"mask_only": "Only generate content in masked area"
}
return {
"required": {
"conditioning": ("CONDITIONING", {"tooltip": tooltips["conditioning"]}),
"latent": ("LATENT", {"tooltip": tooltips["latent"]}),
"mask": ("MASK", {"tooltip": tooltips["mask"]}),
"expand": ("INT", {"default": 5, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["expand"]}),
"blur": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 64.0, "step": 0.1, "tooltip": tooltips["blur"]}),
"mask_only": ("BOOLEAN", {"default": True, "tooltip": tooltips["mask_only"]}),
}
}
RETURN_TYPES = ("CONDITIONING", "LATENT", "MASK")
RETURN_NAMES = ("CONDITIONING", "LATENT", "MASK")
FUNCTION = "prepare_inpaint_conditioning"
CATEGORY = "🧪AILab/🎭Inpaint"
def add_latent_to_conditioning(self, conditioning, latent=None):
if latent is not None:
return node_helpers.conditioning_set_values(
conditioning,
{"reference_latents": [latent["samples"]]},
append=True
)
return conditioning
def prepare_inpaint_conditioning(self, conditioning, latent, mask, expand=5, blur=3.0, mask_only=True):
processed_mask = mask
if expand != 0:
processed_mask = expand_mask(processed_mask, expand)
if blur > 0:
processed_mask = blur_mask(processed_mask, blur)
modified_cond = node_helpers.conditioning_set_values(
conditioning,
{
"concat_latent_image": latent["samples"],
"concat_mask": processed_mask
}
)
final_cond = self.add_latent_to_conditioning(modified_cond, latent)
output_latent = {"samples": latent["samples"]}
if mask_only:
output_latent["noise_mask"] = processed_mask
return (final_cond, output_latent, processed_mask)
NODE_CLASS_MAPPINGS = {
"AILab_ReferenceLatentMask": AILab_ReferenceLatentMask,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"AILab_ReferenceLatentMask": "Reference Latent Mask (RMBG) 🖼️🎭",
}
py/AILab_LamaRemover.py
+191
View File
@@ -0,0 +1,191 @@
# ComfyUI-RMBG
# This custom node for ComfyUI provides functionality for Object removal using Big-Lama model.
#
# reference from https://github.com/advimman/lama
#
# This integration script follows GPL-3.0 License.
# When using or modifying this code, please respect both the original model licenses
# and this integration's license terms.
#
# Source: https://github.com/AILab-AI/ComfyUI-RMBG
import os
import torch
import numpy as np
from PIL import Image, ImageOps, ImageFilter
import folder_paths
from comfy.model_management import get_torch_device
from torchvision import transforms
from huggingface_hub import hf_hub_download
import shutil
import gc
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 pil2comfy(image):
img_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
if len(img_tensor.shape) == 3:
img_tensor = img_tensor.unsqueeze(0)
return img_tensor
def pad_image(image, is_mask=False):
w, h = image.size
if w % 8 != 0:
w = w + (8 - w % 8)
if h % 8 != 0:
h = h + (8 - h % 8)
fill_color = 0 if is_mask else None
padded = Image.new(image.mode, (w, h), color=fill_color)
padded.paste(image, (0, 0))
return padded
def cropimage(image, w, h):
return image.crop((0, 0, w, h))
DEVICE = get_torch_device()
folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG"))
class AILab_LamaRemover:
@classmethod
def INPUT_TYPES(s):
tooltips = {
"images": "Input images to be processed",
"masks": "Masks defining areas to be removed (white=remove)",
"removal_strength": "Strength of the removal effect (higher values increase the effect area)",
"edge_smoothness": "Controls edge smoothness (higher values create smoother transitions)"
}
return {
"required": {
"images": ("IMAGE", {"tooltip": tooltips["images"]}),
"masks": ("MASK", {"tooltip": tooltips["masks"]}),
"removal_strength": ("INT", {"default": 230, "min": 0, "max": 255, "step": 1, "display": "slider", "tooltip": tooltips["removal_strength"]}),
"edge_smoothness": ("INT", {"default": 8, "min": 0, "max": 20, "step": 1, "display": "slider", "tooltip": tooltips["edge_smoothness"]}),
},
}
CATEGORY = "🧪AILab/🧽RMBG"
RETURN_NAMES = ("images",)
RETURN_TYPES = ("IMAGE",)
FUNCTION = "remove_object"
def __init__(self):
self.model = None
self.device = DEVICE
self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "Lama")
self.model_path = os.path.join(self.cache_dir, "big-lama.pt")
self.to_pil = transforms.ToPILImage()
def load_model(self):
if self.model is not None:
return
if not os.path.exists(self.model_path):
self.download_model()
try:
self.model = torch.jit.load(self.model_path, map_location=self.device)
except Exception as e:
print(f"Can't use comfy device: {str(e)}")
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = torch.jit.load(self.model_path, map_location=self.device)
self.model.eval()
self.model.to(self.device)
def download_model(self):
print("Downloading Big-Lama model...")
os.makedirs(self.cache_dir, exist_ok=True)
try:
downloaded_path = hf_hub_download(
repo_id="1038lab/Lama",
filename="big-lama.pt",
local_dir=self.cache_dir,
local_dir_use_symlinks=False
)
if os.path.dirname(downloaded_path) != self.cache_dir:
shutil.move(downloaded_path, self.model_path)
print("Big-Lama model downloaded successfully")
except Exception as e:
raise RuntimeError(f"Error downloading Big-Lama model: {str(e)}")
def process_with_model(self, img_tensor, mask_tensor):
with torch.inference_mode():
img_tensor = img_tensor.to(self.device)
mask_tensor = mask_tensor.to(self.device)
result = self.model(img_tensor, mask_tensor)
result_cpu = result[0].cpu()
del img_tensor
del mask_tensor
if torch.cuda.is_available():
torch.cuda.empty_cache()
return result_cpu
def remove_object(self, images, masks, removal_strength, edge_smoothness):
try:
self.load_model()
results = []
for image, mask in zip(images, masks):
ori_image = tensor2pil(image)
w, h = ori_image.size
p_image = pad_image(ori_image)
mask_np = mask.cpu().numpy()
mask_pil = Image.fromarray((mask_np * 255).astype(np.uint8))
p_mask = pad_image(mask_pil, is_mask=True)
if p_mask.size != p_image.size:
try:
p_mask = p_mask.resize(p_image.size, Image.LANCZOS)
except AttributeError:
p_mask = p_mask.resize(p_image.size, Image.ANTIALIAS)
p_mask = ImageOps.invert(p_mask)
p_mask = p_mask.filter(ImageFilter.GaussianBlur(radius=edge_smoothness))
gray = p_mask.point(lambda x: 0 if x > removal_strength else 255)
img_tensor = torch.FloatTensor(np.array(p_image)).permute(2, 0, 1).unsqueeze(0) / 255.0
mask_tensor = torch.FloatTensor(np.array(gray)).unsqueeze(0).unsqueeze(0) / 255.0
result = self.process_with_model(img_tensor, mask_tensor)
result_img = self.to_pil(result.squeeze())
if result_img.width > w or result_img.height > h:
result_img = cropimage(result_img, w, h)
result_tensor = pil2comfy(result_img)
results.append(result_tensor)
del result
gc.collect()
return (torch.cat(results, dim=0),)
except Exception as e:
import traceback
print(traceback.format_exc())
raise RuntimeError(f"Error in object removal: {str(e)}")
finally:
if torch.cuda.is_available():
torch.cuda.empty_cache()
NODE_CLASS_MAPPINGS = {
"AILab_LamaRemover": AILab_LamaRemover,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"AILab_LamaRemover": "Lama Remover (RMBG)",
}
py/AILab_RMBG.py
+677
View File
@@ -0,0 +1,677 @@
# ComfyUI-RMBG
# This custom node for ComfyUI provides functionality for background removal using various models,
# including RMBG-2.0, INSPYRENET, BEN, BEN2 and BIREFNET-HR. It leverages deep learning techniques
# to process images and generate masks for background removal.
#
# Models 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)
# - BEN2: Apache-2.0 License (https://huggingface.co/PramaLLC/BEN2)
#
# 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/1038lab/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 transformers import AutoModelForImageSegmentation
import cv2
import types
device = "cuda" if torch.cuda.is_available() else "cpu"
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": "1038lab/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": "1038lab/BEN",
"files": {
"model.py": "model.py",
"BEN_Base.pth": "BEN_Base.pth"
},
"cache_dir": "BEN"
},
"BEN2": {
"type": "ben2",
"repo_id": "1038lab/BEN2",
"files": {
"BEN2_Base.pth": "BEN2_Base.pth",
"BEN2.py": "BEN2.py"
},
"cache_dir": "BEN2"
}
}
# 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):
cache_path = os.path.join(self.base_cache_dir, AVAILABLE_MODELS[model_name]["cache_dir"])
os.makedirs(cache_path, exist_ok=True)
return cache_path
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
)
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
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
self.model = None
self.current_model_version = None
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)
try:
# Primary path: Modern transformers compatibility mode (optimized for newer versions)
try:
from transformers import PreTrainedModel
import json
config_path = os.path.join(cache_dir, "config.json")
with open(config_path, 'r') as f:
config = json.load(f)
birefnet_path = os.path.join(cache_dir, "birefnet.py")
BiRefNetConfig_path = os.path.join(cache_dir, "BiRefNet_config.py")
# Load the BiRefNetConfig
config_spec = importlib.util.spec_from_file_location("BiRefNetConfig", BiRefNetConfig_path)
config_module = importlib.util.module_from_spec(config_spec)
sys.modules["BiRefNetConfig"] = config_module
config_spec.loader.exec_module(config_module)
# Fix and load birefnet module
with open(birefnet_path, 'r') as f:
birefnet_content = f.read()
birefnet_content = birefnet_content.replace(
"from .BiRefNet_config import BiRefNetConfig",
"from BiRefNetConfig import BiRefNetConfig"
)
module_name = f"custom_birefnet_model_{hash(birefnet_path)}"
module = types.ModuleType(module_name)
sys.modules[module_name] = module
exec(birefnet_content, module.__dict__)
for attr_name in dir(module):
attr = getattr(module, attr_name)
if isinstance(attr, type) and issubclass(attr, PreTrainedModel) and attr != PreTrainedModel:
BiRefNetConfig = getattr(config_module, "BiRefNetConfig")
model_config = BiRefNetConfig()
self.model = attr(model_config)
weights_path = os.path.join(cache_dir, "model.safetensors")
try:
try:
import safetensors.torch
self.model.load_state_dict(safetensors.torch.load_file(weights_path))
except ImportError:
from transformers.modeling_utils import load_state_dict
state_dict = load_state_dict(weights_path)
self.model.load_state_dict(state_dict)
except Exception as load_error:
pytorch_weights = os.path.join(cache_dir, "pytorch_model.bin")
if os.path.exists(pytorch_weights):
self.model.load_state_dict(torch.load(pytorch_weights, map_location="cpu"))
else:
raise RuntimeError(f"Failed to load weights: {str(load_error)}")
break
if self.model is None:
raise RuntimeError("Could not find suitable model class")
except Exception as modern_e:
print(f"[RMBG INFO] Using standard transformers loading (fallback mode)...")
try:
self.model = AutoModelForImageSegmentation.from_pretrained(
cache_dir,
trust_remote_code=True,
local_files_only=True
)
except Exception as standard_e:
handle_model_error(f"Failed to load model with both modern and standard methods. Modern error: {str(modern_e)}. Standard error: {str(standard_e)}")
except Exception as e:
handle_model_error(f"Error loading model: {str(e)}")
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])
])
if isinstance(images, torch.Tensor):
if len(images.shape) == 3:
images = [images]
else:
images = [img for img in images]
original_sizes = [tensor2pil(img).size for img in images]
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():
outputs = self.model(input_batch)
if isinstance(outputs, list) and len(outputs) > 0:
results = outputs[-1].sigmoid().cpu()
elif isinstance(outputs, dict) and 'logits' in outputs:
results = outputs['logits'].sigmoid().cpu()
elif isinstance(outputs, torch.Tensor):
results = outputs.sigmoid().cpu()
else:
try:
if hasattr(outputs, 'last_hidden_state'):
results = outputs.last_hidden_state.sigmoid().cpu()
else:
for k, v in outputs.items():
if isinstance(v, torch.Tensor):
results = v.sigmoid().cpu()
break
except:
handle_model_error("Unable to recognize model output format")
masks = []
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)
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 Exception as e:
handle_model_error(f"Failed to initialize 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
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)
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 BEN2Model(BaseModelLoader):
def __init__(self):
super().__init__()
def load_model(self, model_name):
if self.current_model_version != model_name:
self.clear_model()
try:
cache_dir = self.get_cache_dir(model_name)
model_path = os.path.join(cache_dir, "BEN2.py")
module_name = f"custom_ben2_model_{hash(model_path)}"
spec = importlib.util.spec_from_file_location(module_name, model_path)
ben2_module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = ben2_module
spec.loader.exec_module(ben2_module)
model_weights_path = os.path.join(cache_dir, "BEN2_Base.pth")
self.model = ben2_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
except Exception as e:
handle_model_error(f"Error loading BEN2 model: {str(e)}")
def process_image(self, images, model_name, params):
try:
self.load_model(model_name)
if isinstance(images, torch.Tensor):
if len(images.shape) == 3:
images = [images]
else:
images = [img for img in images]
batch_size = 3
all_masks = []
for i in range(0, len(images), batch_size):
batch_images = images[i:i + batch_size]
batch_pil_images = []
original_sizes = []
for img in batch_images:
orig_image = tensor2pil(img)
w, h = orig_image.size
original_sizes.append((w, h))
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")
batch_pil_images.append(processed_input)
with torch.no_grad():
try:
foregrounds = self.model.inference(batch_pil_images)
if not isinstance(foregrounds, list):
foregrounds = [foregrounds]
except Exception as e:
handle_model_error(f"Error in BEN2 inference: {str(e)}")
for foreground, (orig_w, orig_h) in zip(foregrounds, original_sizes):
foreground = foreground.resize((orig_w, orig_h), Image.LANCZOS)
mask = foreground.split()[-1]
all_masks.append(mask)
if len(all_masks) == 1:
return all_masks[0]
return all_masks
except Exception as e:
handle_model_error(f"Error in BEN2 processing: {str(e)}")
def refine_foreground(image_bchw, masks_b1hw):
b, c, h, w = image_bchw.shape
if b != masks_b1hw.shape[0]:
raise ValueError("images and masks must have the same batch size")
image_np = image_bchw.cpu().numpy()
mask_np = masks_b1hw.cpu().numpy()
refined_fg = []
for i in range(b):
mask = mask_np[i, 0]
thresh = 0.45
mask_binary = (mask > thresh).astype(np.float32)
edge_blur = cv2.GaussianBlur(mask_binary, (3, 3), 0)
transition_mask = np.logical_and(mask > 0.05, mask < 0.95)
alpha = 0.85
mask_refined = np.where(transition_mask,
alpha * mask + (1-alpha) * edge_blur,
mask_binary)
edge_region = np.logical_and(mask > 0.2, mask < 0.8)
mask_refined = np.where(edge_region,
mask_refined * 0.98,
mask_refined)
result = []
for c in range(image_np.shape[1]):
channel = image_np[i, c]
refined = channel * mask_refined
result.append(refined)
refined_fg.append(np.stack(result))
return torch.from_numpy(np.stack(refined_fg))
class RMBG:
def __init__(self):
self.models = {
"RMBG-2.0": RMBGModel(),
"INSPYRENET": InspyrenetModel(),
"BEN": BENModel(),
"BEN2": BEN2Model()
}
@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 output type: Alpha (transparent) or Color (custom background color).",
"background_color": "Pick background color (supports alpha, use color picker).",
"invert_output": "Enable to invert both the image and mask output (useful for certain effects).",
"refine_foreground": "Use Fast Foreground Colour Estimation to optimize transparent background"
}
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": 8, "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": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"refine_foreground": ("BOOLEAN", {"default": False, "tooltip": tooltips["refine_foreground"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "process_image"
CATEGORY = "🧪AILab/🧽RMBG"
def process_image(self, image, model, **params):
try:
processed_images = []
processed_masks = []
model_instance = self.models[model]
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")
model_type = AVAILABLE_MODELS[model]["type"]
def _process_pair(img, mask):
if isinstance(mask, list):
masks = [m.convert("L") for m in mask if isinstance(m, Image.Image)]
mask_local = masks[0] if masks else None
elif isinstance(mask, Image.Image):
mask_local = mask.convert("L")
else:
mask_local = mask
mask_tensor_local = pil2tensor(mask_local)
mask_tensor_local = mask_tensor_local * (1 + (1 - params["sensitivity"]))
mask_tensor_local = torch.clamp(mask_tensor_local, 0, 1)
mask_img_local = tensor2pil(mask_tensor_local)
if params["mask_blur"] > 0:
mask_img_local = mask_img_local.filter(ImageFilter.GaussianBlur(radius=params["mask_blur"]))
if params["mask_offset"] != 0:
if params["mask_offset"] > 0:
for _ in range(params["mask_offset"]):
mask_img_local = mask_img_local.filter(ImageFilter.MaxFilter(3))
else:
for _ in range(-params["mask_offset"]):
mask_img_local = mask_img_local.filter(ImageFilter.MinFilter(3))
if params["invert_output"]:
mask_img_local = Image.fromarray(255 - np.array(mask_img_local))
img_tensor_local = torch.from_numpy(np.array(tensor2pil(img))).permute(2, 0, 1).unsqueeze(0) / 255.0
mask_tensor_b1hw = torch.from_numpy(np.array(mask_img_local)).unsqueeze(0).unsqueeze(0) / 255.0
orig_image_local = tensor2pil(img)
if params.get("refine_foreground", False):
refined_fg_local = refine_foreground(img_tensor_local, mask_tensor_b1hw)
refined_fg_local = tensor2pil(refined_fg_local[0].permute(1, 2, 0))
r, g, b = refined_fg_local.split()
foreground_local = Image.merge('RGBA', (r, g, b, mask_img_local))
else:
orig_rgba_local = orig_image_local.convert("RGBA")
r, g, b, _ = orig_rgba_local.split()
foreground_local = Image.merge('RGBA', (r, g, b, mask_img_local))
if params["background"] == "Color":
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
background_color = params.get("background_color", "#222222")
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image_local.size, rgba)
composite_image = Image.alpha_composite(bg_image, foreground_local)
processed_images.append(pil2tensor(composite_image.convert("RGB")))
else:
processed_images.append(pil2tensor(foreground_local))
processed_masks.append(pil2tensor(mask_img_local))
if model_type in ("rmbg", "ben2"):
images_list = [img for img in image]
chunk_size = 4
for start in range(0, len(images_list), chunk_size):
batch_imgs = images_list[start:start + chunk_size]
masks = model_instance.process_image(batch_imgs, model, params)
if isinstance(masks, Image.Image):
masks = [masks]
for img_item, mask_item in zip(batch_imgs, masks):
_process_pair(img_item, mask_item)
else:
for img in image:
mask = model_instance.process_image(img, model, params)
_process_pair(img, mask)
mask_images = []
for mask_tensor in processed_masks:
mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
mask_images.append(mask_image)
mask_image_output = torch.cat(mask_images, dim=0)
return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0), mask_image_output)
except Exception as e:
handle_model_error(f"Error in image processing: {str(e)}")
empty_mask = torch.zeros((image.shape[0], image.shape[2], image.shape[3]))
empty_mask_image = empty_mask.reshape((-1, 1, empty_mask.shape[-2], empty_mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
return (image, empty_mask, empty_mask_image)
NODE_CLASS_MAPPINGS = {
"RMBG": RMBG
}
NODE_DISPLAY_NAME_MAPPINGS = {
"RMBG": "Remove Background (RMBG)"
}
py/AILab_SAM2Segment.py
+425
View File
@@ -0,0 +1,425 @@
import os
import sys
import copy
from pathlib import Path
import torch
import numpy as np
from PIL import Image, ImageFilter
from torch.hub import download_url_to_file
from safetensors.torch import load_file
import folder_paths
import comfy.model_management
from hydra import initialize_config_dir
from hydra.core.global_hydra import GlobalHydra
try:
from groundingdino.util.slconfig import SLConfig
from groundingdino.models import build_model
from groundingdino.util.utils import clean_state_dict
from groundingdino.util import box_ops
from groundingdino.datasets.transforms import Compose, RandomResize, ToTensor, Normalize
GROUNDINGDINO_AVAILABLE = True
except ImportError:
GROUNDINGDINO_AVAILABLE = False
print("Warning: GroundingDINO not available. Text prompts will use fallback method.")
# Disable PyTorch JIT compilation
# torch.jit._state._python_cu = None
# torch._C._jit_set_profiling_mode(False)
# torch._C._jit_set_profiling_executor(False)
# original_jit_script = torch.jit.script
# def patched_jit_script(obj):
# return obj
# torch.jit.script = patched_jit_script
current_dir = Path(__file__).resolve().parent
repo_root = current_dir.parent
sam2_path = repo_root / "models" / "sam2"
sys.path.insert(0, str(sam2_path))
import torch.serialization
original_torch_load = torch.load
from contextlib import contextmanager
@contextmanager
def _sam2_no_jit():
_orig = torch.jit.script
torch.jit.script = lambda x, *a, **k: x
try:
yield
finally:
torch.jit.script = _orig
def patched_torch_load(*args, **kwargs):
if 'weights_only' not in kwargs:
kwargs['weights_only'] = False
return original_torch_load(*args, **kwargs)
torch.load = patched_torch_load
from sam2_image_predictor import SAM2ImagePredictor
from AILab_ImageMaskTools import pil2tensor, tensor2pil
# SAM2 model definitions with FP32 and FP16 versions
SAM2_MODELS = {
"sam2.1_hiera_tiny": {
"fp32": {
"model_url": "https://huggingface.co/1038lab/sam2/resolve/main/sam2.1_hiera_tiny.safetensors",
"filename": "sam2.1_hiera_tiny.safetensors"
},
"fp16": {
"model_url": "https://huggingface.co/1038lab/sam2/resolve/main/sam2.1_hiera_tiny-fp16.safetensors",
"filename": "sam2.1_hiera_tiny-fp16.safetensors"
}
},
"sam2.1_hiera_small": {
"fp32": {
"model_url": "https://huggingface.co/1038lab/sam2/resolve/main/sam2.1_hiera_small.safetensors",
"filename": "sam2.1_hiera_small.safetensors"
},
"fp16": {
"model_url": "https://huggingface.co/1038lab/sam2/resolve/main/sam2.1_hiera_small-fp16.safetensors",
"filename": "sam2.1_hiera_small-fp16.safetensors"
}
},
"sam2.1_hiera_base_plus": {
"fp32": {
"model_url": "https://huggingface.co/1038lab/sam2/resolve/main/sam2.1_hiera_base_plus.safetensors",
"filename": "sam2.1_hiera_base_plus.safetensors"
},
"fp16": {
"model_url": "https://huggingface.co/1038lab/sam2/resolve/main/sam2.1_hiera_base_plus-fp16.safetensors",
"filename": "sam2.1_hiera_base_plus-fp16.safetensors"
}
},
"sam2.1_hiera_large": {
"fp32": {
"model_url": "https://huggingface.co/1038lab/sam2/resolve/main/sam2.1_hiera_large.safetensors",
"filename": "sam2.1_hiera_large.safetensors"
},
"fp16": {
"model_url": "https://huggingface.co/1038lab/sam2/resolve/main/sam2.1_hiera_large-fp16.safetensors",
"filename": "sam2.1_hiera_large-fp16.safetensors"
}
}
}
# GroundingDINO model definitions
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.safetensors",
"config_filename": "GroundingDINO_SwinT_OGC.cfg.py",
"model_filename": "groundingdino_swint_ogc.safetensors"
},
"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.safetensors",
"config_filename": "GroundingDINO_SwinB.cfg.py",
"model_filename": "groundingdino_swinb_cogcoor.safetensors"
}
}
def get_or_download_model_file(filename, url, dirname):
local_path = folder_paths.get_full_path(dirname, filename)
if local_path:
return local_path
folder = os.path.join(folder_paths.models_dir, dirname)
os.makedirs(folder, exist_ok=True)
local_path = os.path.join(folder, filename)
if not os.path.exists(local_path):
print(f"Downloading {filename} from {url} ...")
download_url_to_file(url, local_path)
return local_path
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 apply_background_color(image: Image.Image, mask_image: Image.Image,
background: str = "Alpha",
background_color: str = "#222222") -> Image.Image:
rgba_image = image.copy().convert('RGBA')
rgba_image.putalpha(mask_image.convert('L'))
if background == "Color":
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
return (r, g, b, 255)
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
return composite_image.convert('RGB')
return rgba_image
class SAM2Segment:
@classmethod
def INPUT_TYPES(cls):
tooltips = {
"prompt": "Enter text description of object to segment",
"sam2_model": "SAM2 model size: Tiny (fastest) to Large (best quality)",
"device": "Auto: smart detection, CPU: force CPU, GPU: force GPU",
"dino_model": "GroundingDINO model for text-to-box detection",
"threshold": "Detection threshold (higher = more strict)",
"mask_blur": "Blur mask edges (0 = disabled)",
"mask_offset": "Expand/shrink mask (positive = expand)",
"invert_output": "Invert the mask output",
"background": "Background type",
"background_color": "Background color (when not Alpha)",
}
return {
"required": {
"image": ("IMAGE",),
"prompt": ("STRING", {"default": "", "multiline": True, "placeholder": "Object to segment", "tooltip": tooltips["prompt"]}),
"sam2_model": (list(SAM2_MODELS.keys()), {"default": "sam2.1_hiera_tiny", "tooltip": tooltips["sam2_model"]}),
"dino_model": (list(DINO_MODELS.keys()), {"default": "GroundingDINO_SwinT_OGC (694MB)", "tooltip": tooltips["dino_model"]}),
"device": (["Auto", "CPU", "GPU"], {"default": "Auto", "tooltip": tooltips["device"]}),
},
"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": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "segment_v2"
CATEGORY = "🧪AILab/🧽RMBG"
def __init__(self):
self.dino_model_cache = {}
self.sam2_model_cache = {}
def load_sam2(self, model_name, device="Auto"):
cache_key = f"{model_name}_{device}"
if cache_key not in self.sam2_model_cache:
model_info = SAM2_MODELS[model_name]
device_obj = comfy.model_management.get_torch_device()
# Determine precision based on device preference
if device == "Auto":
precision = "fp16" if hasattr(device_obj, 'type') and device_obj.type == 'cuda' else "fp32"
elif device == "GPU":
precision = "fp16" if hasattr(device_obj, 'type') and device_obj.type == 'cuda' else "fp32"
else: # CPU
precision = "fp32"
print(f"Loading {model_name} in {precision.upper()} precision")
model_path = get_or_download_model_file(model_info[precision]["filename"], model_info[precision]["model_url"], "sam2")
# Clear any existing Hydra instance
if GlobalHydra().is_initialized():
GlobalHydra.instance().clear()
initialize_config_dir(config_dir=os.path.join(sam2_path, "configs"), job_name="sam2")
config_map = {
"sam2.1_hiera_tiny": "sam2.1/sam2.1_hiera_t.yaml",
"sam2.1_hiera_small": "sam2.1/sam2.1_hiera_s.yaml",
"sam2.1_hiera_base_plus": "sam2.1/sam2.1_hiera_b+.yaml",
"sam2.1_hiera_large": "sam2.1/sam2.1_hiera_l.yaml"
}
config_file = config_map[model_name]
sam_device = comfy.model_management.get_torch_device()
from sam2.build_sam import build_sam2
from hydra import compose
from omegaconf import OmegaConf
from hydra.utils import instantiate
cfg = compose(config_name=config_file)
OmegaConf.resolve(cfg)
sam_model = instantiate(cfg.model, _recursive_=True)
state_dict = load_file(model_path)
# Apply precision to model
dtype = {"fp16": torch.float16, "fp32": torch.float32}[precision]
sam_model.load_state_dict(state_dict, strict=False)
sam_model = sam_model.to(dtype).to(sam_device).eval()
# predictor = SAM2ImagePredictor(sam_model)
with _sam2_no_jit():
predictor = SAM2ImagePredictor(sam_model)
self.sam2_model_cache[cache_key] = predictor
return self.sam2_model_cache[cache_key]
def segment_v2(self, image, prompt, sam2_model, dino_model, device, threshold=0.35,
mask_blur=0, mask_offset=0, background="Alpha",
background_color="#222222", invert_output=False):
device_obj = comfy.model_management.get_torch_device()
# Process batch images
batch_size = image.shape[0] if len(image.shape) == 4 else 1
if len(image.shape) == 3:
image = image.unsqueeze(0)
result_images = []
result_masks = []
result_mask_images = []
for b in range(batch_size):
img_pil = tensor2pil(image[b])
img_np = np.array(img_pil.convert("RGB"))
# Load GroundingDINO config and weights
dino_info = DINO_MODELS[dino_model]
config_path = get_or_download_model_file(dino_info["config_filename"], dino_info["config_url"], "grounding-dino")
weights_path = get_or_download_model_file(dino_info["model_filename"], dino_info["model_url"], "grounding-dino")
# Load and cache GroundingDINO model
dino_key = (config_path, weights_path, device_obj)
if dino_key not in self.dino_model_cache:
args = SLConfig.fromfile(config_path)
model = build_model(args)
checkpoint = load_file(weights_path)
if isinstance(checkpoint, dict) and 'model' in checkpoint:
checkpoint = clean_state_dict(checkpoint['model'])
model.load_state_dict(checkpoint, strict=False)
model.eval()
model.to(device_obj)
self.dino_model_cache[dino_key] = model
dino = self.dino_model_cache[dino_key]
# Load SAM2 model
predictor = self.load_sam2(sam2_model, device)
# Preprocess image for DINO
transform = Compose([
RandomResize([800], max_size=1333),
ToTensor(),
Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
image_tensor, _ = transform(img_pil.convert("RGB"), None)
image_tensor = image_tensor.unsqueeze(0).to(device_obj)
# Prepare text prompt
text_prompt = prompt if prompt.endswith(".") else prompt + "."
# Forward pass
with torch.no_grad():
outputs = dino(image_tensor, captions=[text_prompt])
logits = outputs["pred_logits"].sigmoid()[0]
boxes = outputs["pred_boxes"][0]
# Filter boxes by threshold
filt_mask = logits.max(dim=1)[0] > threshold
boxes_filt = boxes[filt_mask]
# Handle case with no detected boxes
if boxes_filt.shape[0] == 0:
width, height = img_pil.size
empty_mask = torch.zeros((1, height, width), dtype=torch.float32, device="cpu")
empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3)
result_image = apply_background_color(img_pil, Image.fromarray((empty_mask[0].numpy() * 255).astype(np.uint8)), background, background_color)
result_images.append(pil2tensor(result_image))
result_masks.append(empty_mask)
result_mask_images.append(empty_mask_rgb)
continue
# Convert boxes to xyxy
H, W = img_pil.size[1], img_pil.size[0]
boxes_xyxy = box_ops.box_cxcywh_to_xyxy(boxes_filt)
boxes_xyxy = boxes_xyxy * torch.tensor([W, H, W, H], dtype=torch.float32, device=boxes_xyxy.device)
boxes_xyxy = boxes_xyxy.cpu().numpy()
# Set image and predict with autocast for precision handling
from contextlib import nullcontext
# Determine precision based on device preference
if device == "Auto":
precision = "fp16" if hasattr(device_obj, 'type') and device_obj.type == 'cuda' else "fp32"
elif device == "GPU":
precision = "fp16" if hasattr(device_obj, 'type') and device_obj.type == 'cuda' else "fp32"
else: # CPU
precision = "fp32"
autocast_condition = not comfy.model_management.is_device_mps(device_obj)
with torch.autocast(comfy.model_management.get_autocast_device(device_obj), dtype=torch.float16 if precision == "fp16" else torch.float32) if autocast_condition else nullcontext():
predictor.set_image(img_pil)
# Process each box individually and combine masks
all_masks = []
for box in boxes_xyxy:
with torch.no_grad():
masks, iou_predictions, low_res_masks = predictor.predict(
point_coords=None,
point_labels=None,
box=box,
multimask_output=False
)
all_masks.append(masks)
# Combine all masks into one
if len(all_masks) == 1:
mask = all_masks[0]
else:
combined_mask = np.zeros_like(all_masks[0])
for single_mask in all_masks:
combined_mask = np.maximum(combined_mask, single_mask)
mask = combined_mask
# Ensure mask is 2D
if mask.ndim > 2:
mask = mask.squeeze()
mask = (mask * 255).astype(np.uint8)
mask_pil = Image.fromarray(mask, mode="L")
# Process mask and apply background
mask_image = process_mask(mask_pil, invert_output, mask_blur, mask_offset)
result_image = apply_background_color(img_pil, mask_image, background, background_color)
if background == "Color":
result_image = result_image.convert("RGB")
else:
result_image = result_image.convert("RGBA")
# Convert to tensors
mask_tensor = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0)
mask_image_vis = mask_tensor.reshape((-1, 1, mask_image.height, mask_image.width)).movedim(1, -1).expand(-1, -1, -1, 3)
result_images.append(pil2tensor(result_image))
result_masks.append(mask_tensor)
result_mask_images.append(mask_image_vis)
# If no images were successfully processed, return empty results
if len(result_images) == 0:
width, height = tensor2pil(image[0]).size
empty_mask = torch.zeros((batch_size, 1, height, width), dtype=torch.float32, device="cpu")
empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3)
return (image, empty_mask, empty_mask_rgb)
# Combine all batch results
return (torch.cat(result_images, dim=0),
torch.cat(result_masks, dim=0),
torch.cat(result_mask_images, dim=0))
NODE_CLASS_MAPPINGS = {
"SAM2Segment": SAM2Segment,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SAM2Segment": "SAM2 Segmentation (RMBG)",
}
py/AILab_SAM3Segment.py
+332
View File
@@ -0,0 +1,332 @@
import os
import sys
from pathlib import Path
from contextlib import nullcontext
import numpy as np
import torch
from PIL import Image, ImageFilter
from torch.hub import download_url_to_file
import folder_paths
import comfy.model_management
from AILab_ImageMaskTools import pil2tensor, tensor2pil
CURRENT_DIR = Path(__file__).resolve().parent
REPO_ROOT = CURRENT_DIR.parent
SAM3_LOCAL_DIR = REPO_ROOT / "models" / "sam3"
if str(SAM3_LOCAL_DIR) not in sys.path:
sys.path.insert(0, str(SAM3_LOCAL_DIR))
MODELS_ROOT = REPO_ROOT / "models"
if str(MODELS_ROOT) not in sys.path:
sys.path.insert(0, str(MODELS_ROOT))
SAM3_BPE_PATH = SAM3_LOCAL_DIR / "assets" / "bpe_simple_vocab_16e6.txt.gz"
if not os.path.isfile(SAM3_BPE_PATH):
raise RuntimeError("SAM3 assets missing; ensure sam3/assets/bpe_simple_vocab_16e6.txt.gz exists.")
_DEFAULT_PT_ENTRY = {
"model_url": "https://huggingface.co/1038lab/sam3/resolve/main/sam3.pt",
"filename": "sam3.pt",
}
SAM3_MODELS = {
"sam3": _DEFAULT_PT_ENTRY.copy(),
}
def get_sam3_pt_models():
entry = SAM3_MODELS.get("sam3")
if entry and entry.get("filename", "").endswith(".pt"):
return {"sam3": entry}
for key, value in SAM3_MODELS.items():
if value.get("filename", "").endswith(".pt"):
return {"sam3": value}
if "sam3" in key and value:
candidate = value.copy()
candidate["model_url"] = _DEFAULT_PT_ENTRY["model_url"]
candidate["filename"] = _DEFAULT_PT_ENTRY["filename"]
return {"sam3": candidate}
return {"sam3": _DEFAULT_PT_ENTRY.copy()}
def process_mask(mask_image, invert_output=False, mask_blur=0, mask_offset=0):
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:
filt = 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(filt(size))
return mask_image
def apply_background_color(image, mask_image, background="Alpha", background_color="#222222"):
rgba_image = image.copy().convert("RGBA")
rgba_image.putalpha(mask_image.convert("L"))
if background == "Color":
hex_color = background_color.lstrip("#")
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
bg_image = Image.new("RGBA", image.size, (r, g, b, 255))
composite = Image.alpha_composite(bg_image, rgba_image)
return composite.convert("RGB")
return rgba_image
def get_or_download_model_file(filename, url):
local_path = None
if hasattr(folder_paths, "get_full_path"):
local_path = folder_paths.get_full_path("sam3", filename)
if local_path and os.path.isfile(local_path):
return local_path
base_models_dir = getattr(folder_paths, "models_dir", os.path.join(CURRENT_DIR, "models"))
models_dir = os.path.join(base_models_dir, "sam3")
os.makedirs(models_dir, exist_ok=True)
local_path = os.path.join(models_dir, filename)
if not os.path.exists(local_path):
print(f"Downloading {filename} from {url} ...")
download_url_to_file(url, local_path)
return local_path
def _resolve_device(user_choice):
auto_device = comfy.model_management.get_torch_device()
if user_choice == "CPU":
return torch.device("cpu")
if user_choice == "GPU":
if auto_device.type != "cuda":
raise RuntimeError("GPU unavailable")
return torch.device("cuda")
return auto_device
from sam3.model_builder import build_sam3_image_model # noqa: E402
from sam3.model.sam3_image_processor import Sam3Processor # noqa: E402
class SAM3Segment:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"prompt": ("STRING", {"default": "", "multiline": True, "placeholder": "Describe the concept"}),
"output_mode": (["Merged", "Separate"], {"default": "Merged"}),
"confidence_threshold": ("FLOAT", {"default": 0.5, "min": 0.05, "max": 0.95, "step": 0.01}),
},
"optional": {
"max_segments": ("INT", {"default": 0, "min": 0, "max": 128, "step": 1}),
"segment_pick": ("INT", {"default": 0, "min": 0, "max": 128, "step": 1}),
"mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1}),
"mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1}),
"device": (["Auto", "CPU", "GPU"], {"default": "Auto"}),
"invert_output": ("BOOLEAN", {"default": False}),
"unload_model": ("BOOLEAN", {"default": False}),
"background": (["Alpha", "Color"], {"default": "Alpha"}),
"background_color": ("COLORCODE", {"default": "#222222"}),
},
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "segment"
CATEGORY = "🧪AILab/🧽RMBG"
def __init__(self):
self.processor_cache = {}
def _load_processor(self, device_choice):
torch_device = _resolve_device(device_choice)
device_str = "cuda" if torch_device.type == "cuda" else "cpu"
cache_key = ("sam3", device_str)
if cache_key not in self.processor_cache:
model_info = SAM3_MODELS["sam3"]
ckpt_path = get_or_download_model_file(model_info["filename"], model_info["model_url"])
model = build_sam3_image_model(
bpe_path=SAM3_BPE_PATH,
device=device_str,
eval_mode=True,
checkpoint_path=ckpt_path,
load_from_HF=False,
enable_segmentation=True,
enable_inst_interactivity=False,
)
processor = Sam3Processor(model, device=device_str)
self.processor_cache[cache_key] = processor
return self.processor_cache[cache_key], torch_device
def _empty_result(self, img_pil, background, background_color):
w, h = img_pil.size
mask_image = Image.new("L", (w, h), 0)
result_image = apply_background_color(img_pil, mask_image, background, background_color)
result_image = result_image.convert("RGBA") if background == "Alpha" else result_image.convert("RGB")
empty_mask = torch.zeros((1, h, w), dtype=torch.float32)
mask_rgb = empty_mask.reshape((-1, 1, h, w)).movedim(1, -1).expand(-1, -1, -1, 3)
return result_image, empty_mask, mask_rgb
def _empty_batch(self, img_pil):
w, h = img_pil.size
empty_imgs = torch.zeros((0, h, w, 3), dtype=torch.float32)
empty_masks = torch.zeros((0, h, w), dtype=torch.float32)
empty_mask_images = torch.zeros((0, h, w, 3), dtype=torch.float32)
return empty_imgs, empty_masks, empty_mask_images
def _run_single_per_instance(self, processor, img_tensor, prompt, confidence, max_segments, segment_pick, mask_blur, mask_offset, invert, unload_model, background, background_color):
img_pil = tensor2pil(img_tensor)
text = prompt.strip() or "object"
state = processor.set_image(img_pil)
processor.reset_all_prompts(state)
processor.set_confidence_threshold(confidence, state)
state = processor.set_text_prompt(text, state)
masks = state.get("masks")
logits = state.get("masks_logits")
if masks is None or masks.numel() == 0:
return self._empty_batch(img_pil)
masks = masks.float()
if masks.ndim == 4:
masks = masks.squeeze(1)
scores = None
if logits is not None:
logits = logits.float()
if logits.ndim == 4:
logits = logits.squeeze(1)
scores = logits.mean(dim=(-2, -1))
if scores is None:
scores = torch.ones((masks.shape[0],), device=masks.device)
if max_segments > 0 and masks.shape[0] > max_segments:
topk = torch.topk(scores, k=max_segments)
masks = masks[topk.indices]
scores = scores[topk.indices]
sorted_idx = torch.argsort(scores, descending=True)
masks = masks[sorted_idx]
if segment_pick > 0:
idx = segment_pick - 1
if idx >= masks.shape[0]:
return self._empty_batch(img_pil)
masks = masks[idx : idx + 1]
mask_imgs, mask_tensors, mask_rgb_list = [], [], []
for single_mask in masks:
mask_np = (single_mask.clamp(0, 1).cpu().numpy() * 255).astype(np.uint8)
mask_image = Image.fromarray(mask_np, mode="L")
mask_image = process_mask(mask_image, invert, mask_blur, mask_offset)
composed = apply_background_color(img_pil, mask_image, background, background_color)
composed = composed.convert("RGBA") if background == "Alpha" else composed.convert("RGB")
mask_tensor = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0)
mask_rgb = mask_tensor.reshape((1, mask_image.height, mask_image.width, 1)).expand(-1, -1, -1, 3)
mask_imgs.append(pil2tensor(composed))
mask_tensors.append(mask_tensor)
mask_rgb_list.append(mask_rgb)
return (
torch.cat(mask_imgs, dim=0),
torch.cat(mask_tensors, dim=0),
torch.cat(mask_rgb_list, dim=0),
)
def _run_single_merged(self, processor, img_tensor, prompt, confidence, max_segments, segment_pick, mask_blur, mask_offset, invert, unload_model, background, background_color):
img_pil = tensor2pil(img_tensor)
imgs, masks, _ = self._run_single_per_instance(
processor,
img_tensor,
prompt,
confidence,
max_segments,
segment_pick,
mask_blur,
mask_offset,
invert,
unload_model,
background,
background_color,
)
if masks.shape[0] == 0:
return self._empty_result(img_pil, background, background_color)
merged = masks.amax(dim=0)
mask_np = (merged.clamp(0, 1).cpu().numpy() * 255).astype(np.uint8)
mask_image = Image.fromarray(mask_np, mode="L")
mask_image = process_mask(mask_image, invert, mask_blur, mask_offset)
result_image = apply_background_color(img_pil, mask_image, background, background_color)
result_image = result_image.convert("RGBA") if background == "Alpha" else result_image.convert("RGB")
mask_tensor = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0)
mask_rgb = mask_tensor.reshape((1, mask_image.height, mask_image.width, 1)).expand(-1, -1, -1, 3)
return result_image, mask_tensor, mask_rgb
def segment(self, image, prompt, device, confidence_threshold=0.5, max_segments=0, segment_pick=0, mask_blur=0, mask_offset=0, invert_output=False, unload_model=False, background="Alpha", background_color="#222222", output_mode="Merged"):
if image.ndim == 3:
image = image.unsqueeze(0)
processor, torch_device = self._load_processor(device)
autocast_device = comfy.model_management.get_autocast_device(torch_device)
autocast_enabled = torch_device.type == "cuda" and not comfy.model_management.is_device_mps(torch_device)
ctx = torch.autocast(autocast_device, dtype=torch.bfloat16) if autocast_enabled else nullcontext()
result_images, result_masks, result_mask_images = [], [], []
with ctx:
for tensor_img in image:
if output_mode == "Separate":
imgs_batch, masks_batch, mask_imgs_batch = self._run_single_per_instance(
processor,
tensor_img,
prompt,
confidence_threshold,
max_segments,
segment_pick,
mask_blur,
mask_offset,
invert_output,
unload_model,
background,
background_color,
)
result_images.append(imgs_batch)
result_masks.append(masks_batch)
result_mask_images.append(mask_imgs_batch)
else:
img_pil, mask_tensor, mask_rgb = self._run_single_merged(
processor,
tensor_img,
prompt,
confidence_threshold,
max_segments,
segment_pick,
mask_blur,
mask_offset,
invert_output,
unload_model,
background,
background_color,
)
result_images.append(pil2tensor(img_pil))
result_masks.append(mask_tensor)
result_mask_images.append(mask_rgb)
if unload_model:
device_str = "cuda" if torch_device.type == "cuda" else "cpu"
cache_key = (sam3_model, device_str)
if cache_key in self.processor_cache:
del self.processor_cache[cache_key]
if torch_device.type == "cuda":
torch.cuda.empty_cache()
# return torch.cat(result_images, dim=0), torch.cat(result_masks, dim=0), torch.cat(result_mask_images, dim=0)
# Handle empty results
final_images = torch.cat(result_images, dim=0)
final_masks = torch.cat(result_masks, dim=0)
final_mask_images = torch.cat(result_mask_images, dim=0)
# If no segments found in Separate mode, return at least one empty result
if final_images.shape[0] == 0:
# Use the first input image to get dimensions
img_pil = tensor2pil(image[0])
empty_img, empty_mask, empty_mask_img = self._empty_result(img_pil, background, background_color)
final_images = pil2tensor(empty_img)
final_masks = empty_mask
final_mask_images = empty_mask_img
return final_images, final_masks, final_mask_images
NODE_CLASS_MAPPINGS = {"SAM3Segment": SAM3Segment}
NODE_DISPLAY_NAME_MAPPINGS = {"SAM3Segment": "SAM3 Segmentation (RMBG)"}
py/AILab_SDMatte.py
+435
View File
@@ -0,0 +1,435 @@
# ComfyUI-RMBG
# This custom node for ComfyUI provides functionality for Object removal using SDMatte model.
#
# reference from https://github.com/vivoCameraResearch/SDMatte
# model: https://huggingface.co/1038lab/SDMatte
#
# 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
from pathlib import Path
if os.environ.get('SDMATTE_CPU_ONLY', '').lower() in ('1', 'true', 'yes'):
os.environ['CUDA_VISIBLE_DEVICES'] = ''
import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image, ImageFilter
from torch.hub import download_url_to_file
from torchvision import transforms
from torchvision.transforms import InterpolationMode
import folder_paths
try:
import comfy.model_management
COMFY_AVAILABLE = True
except Exception as e:
print(f"Warning: ComfyUI model management not available: {e}")
COMFY_AVAILABLE = False
class MockModelManagement:
@staticmethod
def get_torch_device():
return torch.device('cpu')
class MockComfy:
model_management = MockModelManagement()
comfy = MockComfy()
try:
from safetensors.torch import load_file
SAFETENSORS_AVAILABLE = True
except ImportError:
SAFETENSORS_AVAILABLE = False
print("Warning: safetensors not available. Will use torch.load for model loading.")
try:
import diffusers
import transformers
DIFFUSERS_AVAILABLE = True
except ImportError:
DIFFUSERS_AVAILABLE = False
print("Warning: diffusers/transformers not available. SDMatte functionality will be limited.")
current_dir = Path(__file__).resolve().parent
repo_root = current_dir.parent
sdmatte_path = repo_root / "models" / "SDMatte"
sys.path.insert(0, str(sdmatte_path))
SDMATTE_MODELS = {
"SDMatte": {
"model_url": "https://huggingface.co/1038lab/SDMatte/resolve/main/SDMatte.safetensors",
"filename": "SDMatte.safetensors",
"repo_id": "1038lab/SDMatte"
},
"SDMatte_plus": {
"model_url": "https://huggingface.co/1038lab/SDMatte/resolve/main/SDMatte_plus.safetensors",
"filename": "SDMatte_plus.safetensors",
"repo_id": "1038lab/SDMatte"
}
}
REQUIRED_COMPONENTS = ["scheduler", "text_encoder", "tokenizer", "unet", "vae"]
def get_or_download_model_file(filename, url, dirname):
local_path = folder_paths.get_full_path(dirname, filename)
if local_path:
return local_path
folder = os.path.join(folder_paths.models_dir, dirname)
os.makedirs(folder, exist_ok=True)
local_path = os.path.join(folder, filename)
if not os.path.exists(local_path):
print(f"Downloading {filename} from {url} ...")
try:
download_url_to_file(url, local_path)
except Exception as e:
raise RuntimeError(f"Failed to download {filename} from {url}: {e}")
return local_path
def ensure_model_components(model_name):
model_info = SDMATTE_MODELS[model_name]
repo_id = model_info["repo_id"]
components_dir = os.path.join(folder_paths.models_dir, "RMBG", "SDMatte")
missing_components = []
for component in REQUIRED_COMPONENTS:
component_path = os.path.join(components_dir, component)
if not os.path.exists(component_path) or not os.listdir(component_path):
missing_components.append(component)
if missing_components:
print(f"Downloading missing SDMatte components: {missing_components}")
base_url = f"https://huggingface.co/{repo_id}/resolve/main"
for component in missing_components:
component_dir = os.path.join(components_dir, component)
os.makedirs(component_dir, exist_ok=True)
if component == "scheduler":
files = ["scheduler_config.json"]
elif component == "text_encoder":
files = ["config.json"]
elif component == "tokenizer":
files = ["merges.txt", "special_tokens_map.json", "tokenizer_config.json", "vocab.json"]
elif component == "unet":
files = ["config.json"]
elif component == "vae":
files = ["config.json"]
for file in files:
file_url = f"{base_url}/{component}/{file}"
file_path = os.path.join(component_dir, file)
if not os.path.exists(file_path):
try:
print(f" Downloading {component}/{file}...")
download_url_to_file(file_url, file_path)
except Exception as e:
print(f" Warning: Failed to download {file}: {e}")
return components_dir
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 apply_background_color(image: Image.Image, mask_image: Image.Image,
background: str = "Alpha",
background_color: str = "#222222") -> Image.Image:
rgba_image = image.copy().convert('RGBA')
rgba_image.putalpha(mask_image.convert('L'))
if background == "Color":
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
return (r, g, b, 255)
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
return composite_image.convert('RGB')
return rgba_image
def pil2tensor(image):
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
def tensor2pil(image):
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
SDMatteCore = None
def _resize_norm_image_bchw(image_bchw: torch.Tensor, size_hw=(1024, 1024)) -> torch.Tensor:
if image_bchw.shape[1] == 4:
image_bchw = image_bchw[:, :3, :, :]
resize = transforms.Resize(size_hw, interpolation=InterpolationMode.BILINEAR, antialias=True)
norm = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
x = resize(image_bchw)
x = norm(x)
return x
def _resize_mask_b1hw(mask_b1hw: torch.Tensor, size_hw=(1024, 1024)) -> torch.Tensor:
resize = transforms.Resize(size_hw, interpolation=InterpolationMode.BILINEAR, antialias=True)
return resize(mask_b1hw)
class AILab_SDMatte:
@classmethod
def INPUT_TYPES(cls):
tooltips = {
"model": "SDMatte model variant: Standard or Plus version",
"image": "Input image for matting extraction",
"mask": "Mask: White=foreground, Black=background. If omitted and image has alpha, alpha will be used.",
"process_res": "Processing resolution: higher = better quality but slower",
"device": "Auto: smart detection, CPU: force CPU, GPU: force GPU",
"transparent_object": "Whether input image contains transparent objects",
"mask_refine": "Enable mask refinement using mask constraints",
"sensitivity": "Sensitivity for mask constraint (0.1-1.0): higher = more strict",
"mask_blur": "Blur mask edges (0 = disabled)",
"mask_offset": "Expand/shrink mask (positive = expand)",
"invert_output": "Invert the mask output",
"background": "Background type for output",
"background_color": "Background color (when not Alpha)",
}
return {
"required": {
"image": ("IMAGE",),
"model": (list(SDMATTE_MODELS.keys()), {"default": "SDMatte", "tooltip": tooltips["model"]}),
"device": (["Auto", "CPU", "GPU"], {"default": "Auto", "tooltip": tooltips["device"]}),
"process_res": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 8, "tooltip": tooltips["process_res"]}),
},
"optional": {
"mask": ("MASK", {"tooltip": tooltips["mask"]}),
"transparent_object": ("BOOLEAN", {"default": True, "tooltip": tooltips["transparent_object"]}),
"mask_refine": ("BOOLEAN", {"default": True, "tooltip": tooltips["mask_refine"]}),
"sensitivity": ("FLOAT", {"default": 0.9, "min": 0.1, "max": 1.0, "step": 0.1, "tooltip": tooltips["sensitivity"]}),
"mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}),
"mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "matting_inference"
CATEGORY = "🧪AILab/🧽RMBG"
def __init__(self):
self.model_cache = {}
def load_sdmatte_model(self, model_name, device="Auto"):
cache_key = f"{model_name}_{device}"
current_model_keys = [k for k in self.model_cache.keys() if k.startswith(model_name)]
if cache_key not in self.model_cache and len(self.model_cache) > 0:
for key in list(self.model_cache.keys()):
if key not in current_model_keys:
del self.model_cache[key]
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if cache_key not in self.model_cache:
if not DIFFUSERS_AVAILABLE:
raise ImportError("diffusers and transformers are required for SDMatte functionality")
global SDMatteCore
if SDMatteCore is None:
import sys
import os
current_dir = os.path.dirname(__file__)
if current_dir not in sys.path:
sys.path.insert(0, current_dir)
from SDMatte.modeling.SDMatte.meta_arch import SDMatte as SDMatteCore
model_info = SDMATTE_MODELS[model_name]
model_path = get_or_download_model_file(
model_info["filename"],
model_info["model_url"],
"RMBG/SDMatte"
)
pretrained_repo = ensure_model_components(model_name)
sdmatte_model = SDMatteCore(
pretrained_model_name_or_path=pretrained_repo,
load_weight=False,
use_aux_input=True,
aux_input="trimap",
use_encoder_hidden_states=True,
use_attention_mask=True,
add_noise=False,
)
self._load_model_weights(sdmatte_model, model_path)
device_obj = comfy.model_management.get_torch_device()
if device == "CPU":
device_obj = torch.device('cpu')
elif device == "GPU":
if not torch.cuda.is_available():
print("SDMatte: GPU requested but CUDA not available, falling back to CPU")
device_obj = torch.device('cpu')
else:
device_obj = comfy.model_management.get_torch_device()
sdmatte_model.eval()
sdmatte_model.to(device_obj)
if device_obj.type == 'cuda':
self._apply_memory_optimizations(sdmatte_model)
self.model_cache[cache_key] = sdmatte_model
return self.model_cache[cache_key]
def _load_model_weights(self, model, model_path):
if not SAFETENSORS_AVAILABLE:
raise ImportError("safetensors is required for SDMatte functionality")
try:
state_dict = load_file(model_path)
model.load_state_dict(state_dict, strict=False)
except Exception as e:
if os.path.exists(model_path):
print(f"[SDMatte] Model file appears corrupted, deleting: {model_path}")
os.remove(model_path)
raise RuntimeError(f"Failed to load model weights. File may be corrupted. Please try again to re-download. Error: {e}")
def _apply_memory_optimizations(self, model):
try:
torch.cuda.empty_cache()
except Exception:
pass
try:
unet = getattr(model, 'unet', None)
if unet is not None and hasattr(unet, 'set_attn_processor'):
from diffusers.models.attention_processor import SlicedAttnProcessor
unet.set_attn_processor(SlicedAttnProcessor(slice_size=1))
except Exception:
pass
def matting_inference(self, image, model, process_res, device="Auto",
mask=None, transparent_object=True, mask_refine=True,
sensitivity=0.8, mask_blur=0, mask_offset=0,
invert_output=False, background="Alpha", background_color="#222222"):
sdmatte_model = self.load_sdmatte_model(model, device)
device_obj = comfy.model_management.get_torch_device()
if device == "CPU":
device_obj = torch.device('cpu')
batch_size = image.shape[0]
result_masks = []
result_images = []
result_mask_images = []
for b in range(batch_size):
img_pil = tensor2pil(image[b])
B, H, W = 1, img_pil.height, img_pil.width
orig_h, orig_w = H, W
img_bchw = image[b:b+1].permute(0, 3, 1, 2).contiguous().to(device_obj)
img_in = _resize_norm_image_bchw(img_bchw, (int(process_res), int(process_res)))
if mask is not None:
mask_b1hw = mask[b:b+1].unsqueeze(1).contiguous().to(device_obj)
mask_for_refine = mask[b:b+1]
else:
if image.shape[-1] == 4:
alpha = image[b, :, :, 3]
mask_b1hw = alpha.unsqueeze(0).unsqueeze(0).contiguous().to(device_obj)
mask_for_refine = alpha.unsqueeze(0)
else:
raise ValueError("Mask required: provide a mask or use an image with alpha.")
tri = _resize_mask_b1hw(mask_b1hw, (int(process_res), int(process_res))) * 2 - 1
data = {"image": img_in,
"is_trans": torch.tensor([1 if transparent_object else 0], device=device_obj),
"caption": [""],
"trimap": tri,
"trimap_coords": torch.tensor([[0,0,1,1]], dtype=tri.dtype, device=device_obj)}
with torch.inference_mode():
if device_obj.type == 'cuda':
with torch.autocast(device_type='cuda', dtype=torch.float16):
pred_alpha = sdmatte_model(data)
else:
pred_alpha = sdmatte_model(data)
out = transforms.Resize((orig_h, orig_w), interpolation=InterpolationMode.BILINEAR, antialias=True)(pred_alpha)
out = out.squeeze(1).clamp(0, 1).detach().cpu()
if mask_refine:
out = self._refine_mask(out, mask_for_refine, sensitivity)
mask_pil = Image.fromarray((out[0].numpy() * 255).astype(np.uint8), mode="L")
mask_image = process_mask(mask_pil, invert_output, mask_blur, mask_offset)
result_image = apply_background_color(img_pil, mask_image, background, background_color)
if background == "Color":
result_image = result_image.convert("RGB")
else:
result_image = result_image.convert("RGBA")
mask_tensor = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0)
mask_image_vis = mask_tensor.reshape((-1, 1, mask_image.height, mask_image.width)).movedim(1, -1).expand(-1, -1, -1, 3)
result_masks.append(mask_tensor)
result_images.append(pil2tensor(result_image))
result_mask_images.append(mask_image_vis)
if device_obj.type == 'cuda':
torch.cuda.empty_cache()
import gc
gc.collect()
return (torch.cat(result_images, dim=0), torch.cat(result_masks, dim=0), torch.cat(result_mask_images, dim=0))
def _refine_mask(self, mask, trimap, constraint):
trimap_cpu = trimap.cpu()
foreground_regions = trimap_cpu > constraint
background_regions = trimap_cpu < (1.0 - constraint)
unknown_regions = ~(foreground_regions | background_regions)
refined_mask = mask.clone()
refined_mask[background_regions] = 0.0
refined_mask[foreground_regions] = torch.clamp(refined_mask[foreground_regions] * 1.2, 0, 1)
alpha_threshold = 0.3
low_confidence = (refined_mask < alpha_threshold) & unknown_regions
refined_mask[low_confidence] = 0.0
return refined_mask
NODE_CLASS_MAPPINGS = {
"AILab_SDMatte": AILab_SDMatte,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"AILab_SDMatte": "SDMatte Matting (RMBG)",
}
py/AILab_Segment.py
+385
View File
@@ -0,0 +1,385 @@
# 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.
# Models 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"
},
"sam_hq_vit_h (2.57GB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_hq_vit_h.pth",
"model_type": "vit_h"
},
"sam_hq_vit_l (1.25GB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_hq_vit_l.pth",
"model_type": "vit_l"
},
"sam_hq_vit_b (379MB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_hq_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: str = "Alpha",
background_color: str = "#222222") -> Image.Image:
rgba_image = image.copy().convert('RGBA')
rgba_image.putalpha(mask_image.convert('L'))
if background == "Color":
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
return (r, g, b, 255)
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', image.size, rgba)
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)",
"invert_output": "Invert the mask output",
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": "Choose background type"}),
"background_color": "Choose background color (Alpha = transparent)",
}
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.30, "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": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
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
self._sam_model_cache = {}
self._dino_model_cache = {}
def segment(self, image, prompt, sam_model, dino_model, threshold=0.35,
mask_blur=0, mask_offset=0, background="Alpha",
background_color="#222222", 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")
# Create empty RGB mask for visualization
empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3)
return (pil2tensor(image), empty_mask, empty_mask_rgb)
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")
# Create empty RGB mask for visualization
empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3)
return (pil2tensor(image), empty_mask, empty_mask_rgb)
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, background_color)
if background == "Color":
result_image = result_image.convert("RGB")
else:
result_image = result_image.convert("RGBA")
mask_tensor = image2mask(mask_image).unsqueeze(0)
print(f'Successfully created segment for: "{prompt}"')
# Create mask image for visualization (similar to other nodes)
mask_images = []
# Convert mask to RGB image format for visualization
mask_image_vis = mask_tensor.reshape((-1, 1, mask_image.height, mask_image.width)).movedim(1, -1).expand(-1, -1, -1, 3)
mask_images.append(mask_image_vis)
mask_image_output = torch.cat(mask_images, dim=0)
return (pil2tensor(result_image), mask_tensor, mask_image_output)
def load_sam(self, model_name):
if model_name in self._sam_model_cache:
return self._sam_model_cache[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]()
state_dict = torch.load(sam_checkpoint_path)
sam.load_state_dict(state_dict, strict=False)
sam_device = comfy.model_management.get_torch_device()
sam.to(device=sam_device)
sam.eval()
self._sam_model_cache[model_name] = sam
return sam
def load_groundingdino(self, model_name):
if model_name in self._dino_model_cache:
return self._dino_model_cache[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()
self._dino_model_cache[model_name] = dino
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": "Segmentation V1 (RMBG)"
}
py/AILab_SegmentV2.py
+285
View File
@@ -0,0 +1,285 @@
import os
import sys
import copy
import torch
import numpy as np
from PIL import Image, ImageFilter
from torch.hub import download_url_to_file
import folder_paths
from segment_anything import sam_model_registry, SamPredictor
from groundingdino.util.slconfig import SLConfig
from groundingdino.models import build_model
from groundingdino.util.utils import clean_state_dict
from groundingdino.util import box_ops
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
from AILab_ImageMaskTools import pil2tensor, tensor2pil
SAM_MODELS = {
"sam_vit_h (2.56GB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_vit_h.pth",
"model_type": "vit_h",
"filename": "sam_vit_h.pth"
},
"sam_vit_l (1.25GB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_vit_l.pth",
"model_type": "vit_l",
"filename": "sam_vit_l.pth"
},
"sam_vit_b (375MB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_vit_b.pth",
"model_type": "vit_b",
"filename": "sam_vit_b.pth"
},
"sam_hq_vit_h (2.57GB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_hq_vit_h.pth",
"model_type": "vit_h",
"filename": "sam_hq_vit_h.pth"
},
"sam_hq_vit_l (1.25GB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_hq_vit_l.pth",
"model_type": "vit_l",
"filename": "sam_hq_vit_l.pth"
},
"sam_hq_vit_b (379MB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_hq_vit_b.pth",
"model_type": "vit_b",
"filename": "sam_hq_vit_b.pth"
}
}
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",
"config_filename": "GroundingDINO_SwinT_OGC.cfg.py",
"model_filename": "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",
"config_filename": "GroundingDINO_SwinB.cfg.py",
"model_filename": "groundingdino_swinb_cogcoor.pth"
}
}
def get_or_download_model_file(filename, url, dirname):
local_path = folder_paths.get_full_path(dirname, filename)
if local_path:
return local_path
folder = os.path.join(folder_paths.models_dir, dirname)
os.makedirs(folder, exist_ok=True)
local_path = os.path.join(folder, filename)
if not os.path.exists(local_path):
print(f"Downloading {filename} from {url} ...")
download_url_to_file(url, local_path)
return local_path
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 apply_background_color(image: Image.Image, mask_image: Image.Image,
background: str = "Alpha",
background_color: str = "#222222") -> Image.Image:
rgba_image = image.copy().convert('RGBA')
rgba_image.putalpha(mask_image.convert('L'))
if background == "Color":
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
return (r, g, b, 255)
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
return composite_image.convert('RGB')
return rgba_image
def get_groundingdino_model(device):
processor = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-tiny")
model = AutoModelForZeroShotObjectDetection.from_pretrained("IDEA-Research/grounding-dino-tiny").to(device)
return processor, model
def get_boxes(processor, model, img_pil, prompt, threshold):
inputs = processor(images=img_pil, text=prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model(**inputs)
results = processor.post_process_grounded_object_detection(
outputs,
inputs.input_ids,
box_threshold=threshold,
text_threshold=threshold,
target_sizes=[img_pil.size[::-1]]
)
return results[0]["boxes"]
class SegmentV2:
@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)",
"invert_output": "Invert the mask output",
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": "Choose background type"}),
"background_color": "Choose background color (Alpha = transparent)",
}
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": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "segment_v2"
CATEGORY = "🧪AILab/🧽RMBG"
def __init__(self):
self.dino_model_cache = {}
self.sam_model_cache = {}
def segment_v2(self, image, prompt, sam_model, dino_model, threshold=0.30,
mask_blur=0, mask_offset=0, background="Alpha",
background_color="#222222", invert_output=False):
device = "cuda" if torch.cuda.is_available() else "cpu"
batch_size = image.shape[0] if len(image.shape) == 4 else 1
if len(image.shape) == 3:
image = image.unsqueeze(0)
result_images = []
result_masks = []
result_mask_images = []
for b in range(batch_size):
img_pil = tensor2pil(image[b])
img_np = np.array(img_pil.convert("RGB"))
dino_info = DINO_MODELS[dino_model]
config_path = get_or_download_model_file(dino_info["config_filename"], dino_info["config_url"], "grounding-dino")
weights_path = get_or_download_model_file(dino_info["model_filename"], dino_info["model_url"], "grounding-dino")
dino_key = (config_path, weights_path, device)
if dino_key not in self.dino_model_cache:
args = SLConfig.fromfile(config_path)
model = build_model(args)
checkpoint = torch.load(weights_path, map_location="cpu")
model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False)
model.eval()
model.to(device)
self.dino_model_cache[dino_key] = model
dino = self.dino_model_cache[dino_key]
sam_info = SAM_MODELS[sam_model]
sam_ckpt_path = get_or_download_model_file(sam_info["filename"], sam_info["model_url"], "SAM")
sam_key = (sam_info["model_type"], sam_ckpt_path, device)
if sam_key not in self.sam_model_cache:
try:
sam = sam_model_registry[sam_info["model_type"]]()
state_dict = torch.load(sam_ckpt_path, map_location="cpu")
sam.load_state_dict(state_dict, strict=False)
sam.to(device)
self.sam_model_cache[sam_key] = SamPredictor(sam)
except RuntimeError as e:
if "Unexpected key(s) in state_dict" in str(e):
print("Warning: SAM model loading issue detected, please try using SegmentV1 node instead")
print(f"Error details: {str(e)}")
width, height = img_pil.size
empty_mask = torch.zeros((1, height, width), dtype=torch.float32, device="cpu")
empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3)
result_image = apply_background_color(img_pil, Image.fromarray((empty_mask[0].numpy() * 255).astype(np.uint8)), background, background_color)
result_images.append(pil2tensor(result_image))
result_masks.append(empty_mask)
result_mask_images.append(empty_mask_rgb)
continue
else:
raise e
predictor = self.sam_model_cache[sam_key]
from groundingdino.datasets.transforms import Compose, RandomResize, ToTensor, Normalize
transform = Compose([
RandomResize([800], max_size=1333),
ToTensor(),
Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
image_tensor, _ = transform(img_pil.convert("RGB"), None)
image_tensor = image_tensor.unsqueeze(0).to(device)
text_prompt = prompt if prompt.endswith(".") else prompt + "."
with torch.no_grad():
outputs = dino(image_tensor, captions=[text_prompt])
logits = outputs["pred_logits"].sigmoid()[0]
boxes = outputs["pred_boxes"][0]
filt_mask = logits.max(dim=1)[0] > threshold
boxes_filt = boxes[filt_mask]
if boxes_filt.shape[0] == 0:
width, height = img_pil.size
empty_mask = torch.zeros((1, height, width), dtype=torch.float32, device="cpu")
empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3)
result_image = apply_background_color(img_pil, Image.fromarray((empty_mask[0].numpy() * 255).astype(np.uint8)), background, background_color)
result_images.append(pil2tensor(result_image))
result_masks.append(empty_mask)
result_mask_images.append(empty_mask_rgb)
continue
H, W = img_pil.size[1], img_pil.size[0]
boxes_xyxy = box_ops.box_cxcywh_to_xyxy(boxes_filt)
boxes_xyxy = boxes_xyxy * torch.tensor([W, H, W, H], dtype=torch.float32, device=boxes_xyxy.device)
boxes_xyxy = boxes_xyxy.cpu().numpy()
predictor.set_image(img_np)
boxes_tensor = torch.tensor(boxes_xyxy, dtype=torch.float32, device=predictor.device)
transformed_boxes = predictor.transform.apply_boxes_torch(boxes_tensor, img_np.shape[:2])
masks, _, _ = predictor.predict_torch(
point_coords=None,
point_labels=None,
boxes=transformed_boxes,
multimask_output=False
)
combined_mask = torch.max(masks, dim=0)[0]
mask = combined_mask.float().cpu().numpy()
mask = mask.squeeze(0)
mask = (mask * 255).astype(np.uint8)
mask_pil = Image.fromarray(mask, mode="L")
mask_image = process_mask(mask_pil, invert_output, mask_blur, mask_offset)
result_image = apply_background_color(img_pil, mask_image, background, background_color)
if background == "Color":
result_image = result_image.convert("RGB")
else:
result_image = result_image.convert("RGBA")
mask_tensor = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0)
mask_image_vis = mask_tensor.reshape((-1, 1, mask_image.height, mask_image.width)).movedim(1, -1).expand(-1, -1, -1, 3)
result_images.append(pil2tensor(result_image))
result_masks.append(mask_tensor)
result_mask_images.append(mask_image_vis)
if len(result_images) == 0:
width, height = tensor2pil(image[0]).size
empty_mask = torch.zeros((batch_size, 1, height, width), dtype=torch.float32, device="cpu")
empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3)
return (image, empty_mask, empty_mask_rgb)
return (torch.cat(result_images, dim=0),
torch.cat(result_masks, dim=0),
torch.cat(result_mask_images, dim=0))
NODE_CLASS_MAPPINGS = {
"SegmentV2": SegmentV2,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SegmentV2": "Segmentation V2 (RMBG)",
}
py/AILab_utils.py
+1
View File
@@ -1,3 +1,4 @@
# Shared utility helpers
import numpy as np
import torch
from PIL import Image