From d6fe942a1e63fbd84ec34e605fdbe62c4cbea12d Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Thu, 1 Jan 2026 19:18:04 -0800 Subject: [PATCH] Add files via upload --- py/AILab_BiRefNet.py | 484 ++++++++++++++++++++++++++ py/AILab_BodySegment.py | 251 ++++++++++++++ py/AILab_ClothSegment.py | 288 ++++++++++++++++ py/AILab_FaceSegment.py | 293 ++++++++++++++++ py/AILab_FashionSegment.py | 378 +++++++++++++++++++++ py/AILab_InpaintTools.py | 146 ++++++++ py/AILab_LamaRemover.py | 191 +++++++++++ py/AILab_RMBG.py | 677 +++++++++++++++++++++++++++++++++++++ py/AILab_SAM2Segment.py | 425 +++++++++++++++++++++++ py/AILab_SAM3Segment.py | 332 ++++++++++++++++++ py/AILab_SDMatte.py | 435 ++++++++++++++++++++++++ py/AILab_Segment.py | 385 +++++++++++++++++++++ py/AILab_SegmentV2.py | 285 ++++++++++++++++ py/AILab_utils.py | 1 + 14 files changed, 4571 insertions(+) create mode 100644 py/AILab_BiRefNet.py create mode 100644 py/AILab_BodySegment.py create mode 100644 py/AILab_ClothSegment.py create mode 100644 py/AILab_FaceSegment.py create mode 100644 py/AILab_FashionSegment.py create mode 100644 py/AILab_InpaintTools.py create mode 100644 py/AILab_LamaRemover.py create mode 100644 py/AILab_RMBG.py create mode 100644 py/AILab_SAM2Segment.py create mode 100644 py/AILab_SAM3Segment.py create mode 100644 py/AILab_SDMatte.py create mode 100644 py/AILab_Segment.py create mode 100644 py/AILab_SegmentV2.py diff --git a/py/AILab_BiRefNet.py b/py/AILab_BiRefNet.py new file mode 100644 index 0000000..94e2a7e --- /dev/null +++ b/py/AILab_BiRefNet.py @@ -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)" +} \ No newline at end of file diff --git a/py/AILab_BodySegment.py b/py/AILab_BodySegment.py new file mode 100644 index 0000000..8e0082c --- /dev/null +++ b/py/AILab_BodySegment.py @@ -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)" +} \ No newline at end of file diff --git a/py/AILab_ClothSegment.py b/py/AILab_ClothSegment.py new file mode 100644 index 0000000..62bd5bd --- /dev/null +++ b/py/AILab_ClothSegment.py @@ -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)" +} \ No newline at end of file diff --git a/py/AILab_FaceSegment.py b/py/AILab_FaceSegment.py new file mode 100644 index 0000000..50c9e75 --- /dev/null +++ b/py/AILab_FaceSegment.py @@ -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)" +} diff --git a/py/AILab_FashionSegment.py b/py/AILab_FashionSegment.py new file mode 100644 index 0000000..0fcb829 --- /dev/null +++ b/py/AILab_FashionSegment.py @@ -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)" +} \ No newline at end of file diff --git a/py/AILab_InpaintTools.py b/py/AILab_InpaintTools.py new file mode 100644 index 0000000..f2764bb --- /dev/null +++ b/py/AILab_InpaintTools.py @@ -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) 🖼️🎭", +} \ No newline at end of file diff --git a/py/AILab_LamaRemover.py b/py/AILab_LamaRemover.py new file mode 100644 index 0000000..24c9d1f --- /dev/null +++ b/py/AILab_LamaRemover.py @@ -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)", +} \ No newline at end of file diff --git a/py/AILab_RMBG.py b/py/AILab_RMBG.py new file mode 100644 index 0000000..06c5a63 --- /dev/null +++ b/py/AILab_RMBG.py @@ -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)" +} \ No newline at end of file diff --git a/py/AILab_SAM2Segment.py b/py/AILab_SAM2Segment.py new file mode 100644 index 0000000..a5e55c9 --- /dev/null +++ b/py/AILab_SAM2Segment.py @@ -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)", +} diff --git a/py/AILab_SAM3Segment.py b/py/AILab_SAM3Segment.py new file mode 100644 index 0000000..bcb6ae7 --- /dev/null +++ b/py/AILab_SAM3Segment.py @@ -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)"} diff --git a/py/AILab_SDMatte.py b/py/AILab_SDMatte.py new file mode 100644 index 0000000..f4ebf94 --- /dev/null +++ b/py/AILab_SDMatte.py @@ -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)", +} diff --git a/py/AILab_Segment.py b/py/AILab_Segment.py new file mode 100644 index 0000000..189dad5 --- /dev/null +++ b/py/AILab_Segment.py @@ -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)" +} \ No newline at end of file diff --git a/py/AILab_SegmentV2.py b/py/AILab_SegmentV2.py new file mode 100644 index 0000000..1d4ea58 --- /dev/null +++ b/py/AILab_SegmentV2.py @@ -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)", +} diff --git a/py/AILab_utils.py b/py/AILab_utils.py index d2ad40a..acc2e47 100644 --- a/py/AILab_utils.py +++ b/py/AILab_utils.py @@ -1,3 +1,4 @@ +# Shared utility helpers import numpy as np import torch from PIL import Image