diff --git a/py/AILab_BiRefNet.py b/py/AILab_BiRefNet.py index 4e2329f..152800a 100644 --- a/py/AILab_BiRefNet.py +++ b/py/AILab_BiRefNet.py @@ -21,9 +21,8 @@ from huggingface_hub import hf_hub_download import sys import importlib.util from safetensors.torch import load_file -import cv2 -from AILab_utils import get_device +from AILab_utils import get_device, clean_vram, unload_model device = get_device() # Add model path @@ -214,6 +213,7 @@ def handle_model_error(message): raise RuntimeError(message) def refine_foreground(image_bchw, masks_b1hw): + import cv2 b, c, h, w = image_bchw.shape if b != masks_b1hw.shape[0]: raise ValueError("images and masks must have the same batch size") @@ -406,6 +406,7 @@ class BiRefNetRMBG: "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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}), "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), } @@ -497,6 +498,11 @@ class BiRefNetRMBG: 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)}") + finally: + if params.get("unload_model", False): + unload_model(self.model) + else: + clean_vram() # Node Mapping NODE_CLASS_MAPPINGS = { diff --git a/py/AILab_BodySegment.py b/py/AILab_BodySegment.py index f10eed1..fb441bd 100644 --- a/py/AILab_BodySegment.py +++ b/py/AILab_BodySegment.py @@ -16,36 +16,11 @@ 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'))) - -from AILab_utils import get_device +from AILab_utils import get_device, clean_vram, unload_model, pil2tensor, tensor2pil, RGB2RGBA, mask2image, image2mask device = get_device() folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) @@ -83,6 +58,7 @@ class BodySegment: "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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}), "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), }, @@ -124,7 +100,7 @@ class BodySegment: 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): + def segment_body(self, images, mask_blur=0, mask_offset=0, invert_output=False, unload_model=False, background="Alpha", background_color="#222222", **class_selections): try: # Check and download model if needed cache_status, message = self.check_model_cache() @@ -136,6 +112,7 @@ class BodySegment: # Load model if needed if self.model is None: + import onnxruntime self.model = onnxruntime.InferenceSession( os.path.join(self.cache_dir, self.model_file) ) @@ -241,7 +218,9 @@ class BodySegment: self.clear_model() raise RuntimeError(f"Error in Body Segmentation processing: {str(e)}") finally: - self.clear_model() + if unload_model: + self.clear_model() + clean_vram() NODE_CLASS_MAPPINGS = { "BodySegment": BodySegment diff --git a/py/AILab_ClothSegment.py b/py/AILab_ClothSegment.py index 96813d6..9a85cf8 100644 --- a/py/AILab_ClothSegment.py +++ b/py/AILab_ClothSegment.py @@ -15,36 +15,11 @@ 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'))) - -from AILab_utils import get_device +from AILab_utils import get_device, clean_vram, unload_model, pil2tensor, tensor2pil, RGB2RGBA, mask2image, image2mask device = get_device() folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) @@ -83,6 +58,7 @@ class ClothesSegment: "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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}), "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), }, @@ -144,7 +120,7 @@ class ClothesSegment: 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): + def segment_clothes(self, images, process_res=1024, mask_blur=0, mask_offset=0, invert_output=False, unload_model=False, background="Alpha", background_color="#222222", **class_selections): try: # Check and download model if needed cache_status, message = self.check_model_cache() @@ -156,6 +132,7 @@ class ClothesSegment: # Load model if needed if self.processor is None: + from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation self.processor = SegformerImageProcessor.from_pretrained(self.cache_dir) self.model = AutoModelForSemanticSegmentation.from_pretrained(self.cache_dir) self.model.eval() @@ -277,8 +254,9 @@ class ClothesSegment: self.clear_model() raise RuntimeError(f"Error in Clothes Segformer processing: {str(e)}") finally: - if not getattr(self.model, "training", False): + if unload_model: self.clear_model() + clean_vram() NODE_CLASS_MAPPINGS = { "ClothesSegment": ClothesSegment diff --git a/py/AILab_FaceSegment.py b/py/AILab_FaceSegment.py index 2e9c512..563639a 100644 --- a/py/AILab_FaceSegment.py +++ b/py/AILab_FaceSegment.py @@ -1,294 +1,274 @@ -# 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'))) - -from AILab_utils import get_device -device = get_device() - -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)" -} +# 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 +import folder_paths +from huggingface_hub import hf_hub_download +import shutil +from torchvision import transforms + +from torchvision import transforms +from AILab_utils import get_device, clean_vram, unload_model, pil2tensor, tensor2pil, RGB2RGBA, mask2image, image2mask +device = get_device() + +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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), + "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, invert_output=False, unload_model=False, background="Alpha", background_color="#222222", **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: + from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation + 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 unload_model: + self.clear_model() + clean_vram() + +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 index 2278dc6..e072d69 100644 --- a/py/AILab_FashionSegment.py +++ b/py/AILab_FashionSegment.py @@ -14,31 +14,11 @@ 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) - -from AILab_utils import get_device +from AILab_utils import get_device, clean_vram, unload_model, pil2tensor, tensor2pil, RGB2RGBA, mask2image device = get_device() folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) @@ -176,6 +156,7 @@ class FashionSegmentClothing: "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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}), "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), }, @@ -238,7 +219,7 @@ class FashionSegmentClothing: 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): + invert_output=False, unload_model=False, background="Alpha", background_color="#222222", **class_selections): if accessories_options is None: accessories_options = [] try: @@ -252,6 +233,7 @@ class FashionSegmentClothing: # Load model if needed if self.processor is None: + from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation self.processor = SegformerImageProcessor.from_pretrained(self.cache_dir) self.model = AutoModelForSemanticSegmentation.from_pretrained(self.cache_dir) self.model.eval() @@ -362,8 +344,9 @@ class FashionSegmentClothing: self.clear_model() raise RuntimeError(f"Error in fashion segmentation: {str(e)}") finally: - if self.model is not None and not self.model.training: + if unload_model: self.clear_model() + clean_vram() def __del__(self): self.clear_model() diff --git a/py/AILab_Florence2.py b/py/AILab_Florence2.py index 68af3eb..5160245 100644 --- a/py/AILab_Florence2.py +++ b/py/AILab_Florence2.py @@ -1,352 +1,352 @@ -import io -import os -import random -from typing import Any, Dict, List, Tuple -from unittest.mock import patch - -import matplotlib - -matplotlib.use("Agg") -import matplotlib.patches as patches -import matplotlib.pyplot as plt -import numpy as np -import torch -import torchvision.transforms.functional as TF -from PIL import Image, ImageColor, ImageDraw - -import comfy.model_management as mm -from comfy.utils import ProgressBar -import folder_paths -import transformers -from transformers import AutoModelForCausalLM, AutoProcessor -from transformers.dynamic_module_utils import get_imports - -MODEL_DIR = os.path.join(folder_paths.models_dir, "LLM") -os.makedirs(MODEL_DIR, exist_ok=True) -folder_paths.add_model_folder_path("LLM", MODEL_DIR) - -TASK_CONFIGS = { - "Polygon Mask (text prompt)": { - "token": "", - "mode": "polygon", - "allows_text": True, - }, - "Phrase Grounding (text boxes)": { - "token": "", - "mode": "bbox", - "allows_text": True, - }, - "Region Proposals (boxes only)": { - "token": "", - "mode": "bbox", - "allows_text": False, - }, -} - -TASK_CHOICES = tuple(TASK_CONFIGS.keys()) - -GENERATION_CONFIG = { - "max_new_tokens": 512, - "num_beams": 3, - "do_sample": True, -} - -MODEL_CHOICES = ( - "microsoft/Florence-2-base", - "microsoft/Florence-2-base-ft", - "microsoft/Florence-2-large", - "microsoft/Florence-2-large-ft", - "thwri/CogFlorence-2.1-Large", - "thwri/CogFlorence-2.2-Large", -) - -COLOR_BANK = ["blue", "orange", "green", "purple", "pink", "cyan"] - - -def _fixed_get_imports(filename): - try: - if not str(filename).endswith("modeling_florence2.py"): - return get_imports(filename) - imports = get_imports(filename) - if "flash_attn" in imports: - imports.remove("flash_attn") - return imports - except Exception: - return get_imports(filename) - - -class AILab_Florence2: - CATEGORY = "🧪AILab/🧽RMBG" - RETURN_TYPES = ("IMAGE", "MASK", "JSON") - RETURN_NAMES = ("IMAGE", "MASK", "DATA") - FUNCTION = "analyze" - MODEL_CACHE: Dict[Tuple[str, str, str], Dict[str, Any]] = {} - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "model_name": ( - MODEL_CHOICES, - { - "default": MODEL_CHOICES[0], - "tooltip": "Base = stable, +ft = fine-tuned captions, CogFlorence = sharper phrase alignment.", - }, - ), - "task": ( - TASK_CHOICES, - { - "default": TASK_CHOICES[0], - "tooltip": "Polygon masks use prompts; phrase grounding/region proposals return boxes.", - }, - ), - "precision": ( - ("fp16", "bf16", "fp32"), - {"default": "fp16", "tooltip": "Lower precision saves VRAM; fp32 is safest if you hit NaNs."}, - ), - "attention": ( - ("flash_attention_2", "sdpa", "eager"), - {"default": "sdpa", "tooltip": "flash_attn2 needs PyTorch 2.1+; use eager if kernels fail."}, - ), - "fill_mask": ( - "BOOLEAN", - {"default": True, "tooltip": "When true, bbox tasks also output filled mask tensors."}, - ), - }, - "optional": { - "output_mask_select": ( - "STRING", - { - "default": "", - "tooltip": "Comma-separated indices or labels (e.g. 0,2,person) to limit masks.", - }, - ), - "keep_model_loaded": ( - "BOOLEAN", - {"default": False, "tooltip": "Keep weights on the current device after execution."}, - ), - "text_prompt": ( - "STRING", - { - "default": "", - "multiline": True, - "placeholder": "Prompt: e.g. a person wearing red coat", - "tooltip": "Used for polygon masks or phrase grounding; ignored for region proposals.", - }, - ), - }, - } - - def _ensure_weights(self, model_name: str) -> str: - target = os.path.join(MODEL_DIR, model_name.split("/")[-1]) - if os.path.exists(target): - return target - from huggingface_hub import snapshot_download - - snapshot_download(model_name, local_dir=target, local_dir_use_symlinks=False) - return target - - def _get_model(self, model_name: str, precision: str, attention: str) -> Dict[str, Any]: - key = (model_name, precision, attention) - if key in self.MODEL_CACHE: - return self.MODEL_CACHE[key] - - dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[precision] - weights = self._ensure_weights(model_name) - offload = mm.unet_offload_device() - - if transformers.__version__ < "4.51.0": - with patch("transformers.dynamic_module_utils.get_imports", _fixed_get_imports): - model = AutoModelForCausalLM.from_pretrained( - weights, - attn_implementation=attention, - torch_dtype=dtype, - trust_remote_code=True, - ).to(offload) - else: - from models.modeling_florence2 import Florence2ForConditionalGeneration - - model = Florence2ForConditionalGeneration.from_pretrained( - weights, - attn_implementation=attention, - torch_dtype=dtype, - ).to(offload) - - processor = AutoProcessor.from_pretrained(weights, trust_remote_code=True) - bundle = {"model": model, "processor": processor, "dtype": dtype} - self.MODEL_CACHE[key] = bundle - return bundle - - @staticmethod - def _prepare_prompt(task: str, text_prompt: str) -> str: - config = TASK_CONFIGS[task] - text = text_prompt.strip() - if text and not config["allows_text"]: - text = "" - base = config["token"] - return f"{base} {text}" if text else base - - def _draw_regions( - self, - image_pil: Image.Image, - predictions: Dict[str, Any], - fill_mask: bool, - select_filter: List[str], - ) -> Tuple[torch.Tensor, torch.Tensor]: - width, height = image_pil.size - fig, ax = plt.subplots(figsize=(width / 100, height / 100), dpi=100) - fig.subplots_adjust(left=0, right=1, top=1, bottom=0) - ax.imshow(image_pil) - - mask_layer = Image.new("RGB", image_pil.size, (0, 0, 0)) if fill_mask else None - mask_draw = ImageDraw.Draw(mask_layer) if mask_layer else None - - for index, (bbox, label) in enumerate(zip(predictions["bboxes"], predictions["labels"])): - x0, y0, x1, y1 = bbox - if y1 < y0: - y0, y1 = y1, y0 - if x1 < x0: - x0, x1 = x1, x0 - filter_hit = not select_filter or str(index) in select_filter or label in select_filter - if fill_mask and filter_hit: - mask_draw.rectangle([x0, y0, x1, y1], fill=(255, 255, 255)) - rect = patches.Rectangle((x0, y0), x1 - x0, y1 - y0, linewidth=1, edgecolor="red", facecolor="none") - ax.add_patch(rect) - ax.text( - x0, - max(0, y0 - 12), - f"{index}.{label}", - color="white", - fontsize=12, - bbox=dict(facecolor=random.choice(COLOR_BANK), alpha=0.5), - ) - - ax.axis("off") - buf = io.BytesIO() - plt.savefig(buf, format="png", pad_inches=0) - buf.seek(0) - annotated = Image.open(buf) - plt.close(fig) - - annotated_tensor = TF.to_tensor(annotated)[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() - if mask_layer is not None: - mask_tensor = TF.to_tensor(mask_layer).unsqueeze(0).permute(0, 2, 3, 1).cpu().float() - mask_tensor = mask_tensor.mean(dim=0, keepdim=True)[:, :, :, 0] - else: - mask_tensor = torch.zeros((1, annotated_tensor.shape[1], annotated_tensor.shape[2]), dtype=torch.float32) - - return annotated_tensor, mask_tensor - - def _segment_polygons( - self, - image_pil: Image.Image, - predictions: Dict[str, Any], - fill_mask: bool, - ) -> Tuple[torch.Tensor, torch.Tensor]: - width, height = image_pil.size - mask_image = Image.new("RGB", (width, height), "black") - mask_draw = ImageDraw.Draw(mask_image) - - for polygons, label in zip(predictions["polygons"], predictions["labels"]): - color = random.choice(COLOR_BANK) - for polygon in polygons: - polygon = np.array(polygon).reshape(-1, 2) - polygon = np.clip(polygon, [0, 0], [width - 1, height - 1]) - if len(polygon) < 3: - continue - pts = polygon.reshape(-1).tolist() - if fill_mask: - overlay = Image.new("RGBA", image_pil.size, (255, 255, 255, 0)) - draw = ImageDraw.Draw(overlay) - rgba = ImageColor.getrgb(color) + (180,) - draw.polygon(pts, outline=color, fill=rgba, width=3) - image_pil = Image.alpha_composite(image_pil.convert("RGBA"), overlay).convert("RGB") - else: - draw = ImageDraw.Draw(image_pil) - draw.polygon(pts, outline=color, width=3) - mask_draw.polygon(pts, outline="white", fill="white") - - image_tensor = TF.to_tensor(image_pil)[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() - mask_tensor = TF.to_tensor(mask_image).unsqueeze(0).permute(0, 2, 3, 1).cpu().float() - mask_tensor = mask_tensor.mean(dim=0, keepdim=True)[:, :, :, 0] - return image_tensor, mask_tensor - - def analyze( - self, - image: torch.Tensor, - model_name: str, - task: str, - precision: str, - attention: str, - fill_mask: bool, - output_mask_select: str = "", - keep_model_loaded: bool = False, - text_prompt: str = "", - ): - bundle = self._get_model(model_name, precision, attention) - model = bundle["model"] - processor = bundle["processor"] - dtype = bundle["dtype"] - - device = mm.get_torch_device() - offload = mm.unet_offload_device() - model.to(device) - images = image.permute(0, 3, 1, 2) - - task_config = TASK_CONFIGS[task] - prompt = self._prepare_prompt(task, text_prompt) - out_images, out_masks, out_data = [], [], [] - pbar = ProgressBar(len(images)) - select_filter = [s.strip() for s in output_mask_select.split(",") if s.strip()] - - for tensor in images: - image_pil = TF.to_pil_image(tensor) - inputs = processor( - text=prompt, - images=image_pil, - return_tensors="pt", - do_rescale=False, - ).to(dtype).to(device) - - generated_ids = model.generate( - input_ids=inputs["input_ids"], - pixel_values=inputs["pixel_values"], - max_new_tokens=GENERATION_CONFIG["max_new_tokens"], - do_sample=GENERATION_CONFIG["do_sample"], - num_beams=GENERATION_CONFIG["num_beams"], - use_cache=False, - ) - raw = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] - parsed = processor.post_process_generation(raw, task=task_config["token"], image_size=image_pil.size) - predictions = parsed[task_config["token"]] - - if task_config["mode"] == "bbox": - annotated, mask = self._draw_regions(image_pil, predictions, fill_mask, select_filter) - out_images.append(annotated) - out_masks.append(mask) - else: - image_tensor, mask_tensor = self._segment_polygons(image_pil, predictions, fill_mask) - out_images.append(image_tensor) - out_masks.append(mask_tensor) - - out_data.append(predictions) - pbar.update(1) - - image_out = torch.cat(out_images, dim=0) - mask_out = torch.cat(out_masks, dim=0) - - if not keep_model_loaded: - model.to(offload) - mm.soft_empty_cache() - - return image_out, mask_out, out_data - - -NODE_CLASS_MAPPINGS = { - "AILab_Florence2": AILab_Florence2, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "AILab_Florence2": "Florence2 (RMBG)", -} +import io +import os +import random +from typing import Any, Dict, List, Tuple +from unittest.mock import patch + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.patches as patches +import matplotlib.pyplot as plt +import numpy as np +import torch +import torchvision.transforms.functional as TF +from PIL import Image, ImageColor, ImageDraw + +import comfy.model_management as mm +from comfy.utils import ProgressBar +import folder_paths +import transformers +from transformers import AutoModelForCausalLM, AutoProcessor +from transformers.dynamic_module_utils import get_imports + +MODEL_DIR = os.path.join(folder_paths.models_dir, "LLM") +os.makedirs(MODEL_DIR, exist_ok=True) +folder_paths.add_model_folder_path("LLM", MODEL_DIR) + +TASK_CONFIGS = { + "Polygon Mask (text prompt)": { + "token": "", + "mode": "polygon", + "allows_text": True, + }, + "Phrase Grounding (text boxes)": { + "token": "", + "mode": "bbox", + "allows_text": True, + }, + "Region Proposals (boxes only)": { + "token": "", + "mode": "bbox", + "allows_text": False, + }, +} + +TASK_CHOICES = tuple(TASK_CONFIGS.keys()) + +GENERATION_CONFIG = { + "max_new_tokens": 512, + "num_beams": 3, + "do_sample": True, +} + +MODEL_CHOICES = ( + "microsoft/Florence-2-base", + "microsoft/Florence-2-base-ft", + "microsoft/Florence-2-large", + "microsoft/Florence-2-large-ft", + "thwri/CogFlorence-2.1-Large", + "thwri/CogFlorence-2.2-Large", +) + +COLOR_BANK = ["blue", "orange", "green", "purple", "pink", "cyan"] + + +def _fixed_get_imports(filename): + try: + if not str(filename).endswith("modeling_florence2.py"): + return get_imports(filename) + imports = get_imports(filename) + if "flash_attn" in imports: + imports.remove("flash_attn") + return imports + except Exception: + return get_imports(filename) + + +class AILab_Florence2: + CATEGORY = "🧪AILab/🧽RMBG" + RETURN_TYPES = ("IMAGE", "MASK", "JSON") + RETURN_NAMES = ("IMAGE", "MASK", "DATA") + FUNCTION = "analyze" + MODEL_CACHE: Dict[Tuple[str, str, str], Dict[str, Any]] = {} + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "model_name": ( + MODEL_CHOICES, + { + "default": MODEL_CHOICES[0], + "tooltip": "Base = stable, +ft = fine-tuned captions, CogFlorence = sharper phrase alignment.", + }, + ), + "task": ( + TASK_CHOICES, + { + "default": TASK_CHOICES[0], + "tooltip": "Polygon masks use prompts; phrase grounding/region proposals return boxes.", + }, + ), + "precision": ( + ("fp16", "bf16", "fp32"), + {"default": "fp16", "tooltip": "Lower precision saves VRAM; fp32 is safest if you hit NaNs."}, + ), + "attention": ( + ("flash_attention_2", "sdpa", "eager"), + {"default": "sdpa", "tooltip": "flash_attn2 needs PyTorch 2.1+; use eager if kernels fail."}, + ), + "fill_mask": ( + "BOOLEAN", + {"default": True, "tooltip": "When true, bbox tasks also output filled mask tensors."}, + ), + }, + "optional": { + "output_mask_select": ( + "STRING", + { + "default": "", + "tooltip": "Comma-separated indices or labels (e.g. 0,2,person) to limit masks.", + }, + ), + "keep_model_loaded": ( + "BOOLEAN", + {"default": False, "tooltip": "Keep weights on the current device after execution."}, + ), + "text_prompt": ( + "STRING", + { + "default": "", + "multiline": True, + "placeholder": "Prompt: e.g. a person wearing red coat", + "tooltip": "Used for polygon masks or phrase grounding; ignored for region proposals.", + }, + ), + }, + } + + def _ensure_weights(self, model_name: str) -> str: + target = os.path.join(MODEL_DIR, model_name.split("/")[-1]) + if os.path.exists(target): + return target + from huggingface_hub import snapshot_download + + snapshot_download(model_name, local_dir=target, local_dir_use_symlinks=False) + return target + + def _get_model(self, model_name: str, precision: str, attention: str) -> Dict[str, Any]: + key = (model_name, precision, attention) + if key in self.MODEL_CACHE: + return self.MODEL_CACHE[key] + + dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[precision] + weights = self._ensure_weights(model_name) + offload = mm.unet_offload_device() + + if transformers.__version__ < "4.51.0": + with patch("transformers.dynamic_module_utils.get_imports", _fixed_get_imports): + model = AutoModelForCausalLM.from_pretrained( + weights, + attn_implementation=attention, + torch_dtype=dtype, + trust_remote_code=True, + ).to(offload) + else: + from models.modeling_florence2 import Florence2ForConditionalGeneration + + model = Florence2ForConditionalGeneration.from_pretrained( + weights, + attn_implementation=attention, + torch_dtype=dtype, + ).to(offload) + + processor = AutoProcessor.from_pretrained(weights, trust_remote_code=True) + bundle = {"model": model, "processor": processor, "dtype": dtype} + self.MODEL_CACHE[key] = bundle + return bundle + + @staticmethod + def _prepare_prompt(task: str, text_prompt: str) -> str: + config = TASK_CONFIGS[task] + text = text_prompt.strip() + if text and not config["allows_text"]: + text = "" + base = config["token"] + return f"{base} {text}" if text else base + + def _draw_regions( + self, + image_pil: Image.Image, + predictions: Dict[str, Any], + fill_mask: bool, + select_filter: List[str], + ) -> Tuple[torch.Tensor, torch.Tensor]: + width, height = image_pil.size + fig, ax = plt.subplots(figsize=(width / 100, height / 100), dpi=100) + fig.subplots_adjust(left=0, right=1, top=1, bottom=0) + ax.imshow(image_pil) + + mask_layer = Image.new("RGB", image_pil.size, (0, 0, 0)) if fill_mask else None + mask_draw = ImageDraw.Draw(mask_layer) if mask_layer else None + + for index, (bbox, label) in enumerate(zip(predictions["bboxes"], predictions["labels"])): + x0, y0, x1, y1 = bbox + if y1 < y0: + y0, y1 = y1, y0 + if x1 < x0: + x0, x1 = x1, x0 + filter_hit = not select_filter or str(index) in select_filter or label in select_filter + if fill_mask and filter_hit: + mask_draw.rectangle([x0, y0, x1, y1], fill=(255, 255, 255)) + rect = patches.Rectangle((x0, y0), x1 - x0, y1 - y0, linewidth=1, edgecolor="red", facecolor="none") + ax.add_patch(rect) + ax.text( + x0, + max(0, y0 - 12), + f"{index}.{label}", + color="white", + fontsize=12, + bbox=dict(facecolor=random.choice(COLOR_BANK), alpha=0.5), + ) + + ax.axis("off") + buf = io.BytesIO() + plt.savefig(buf, format="png", pad_inches=0) + buf.seek(0) + annotated = Image.open(buf) + plt.close(fig) + + annotated_tensor = TF.to_tensor(annotated)[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() + if mask_layer is not None: + mask_tensor = TF.to_tensor(mask_layer).unsqueeze(0).permute(0, 2, 3, 1).cpu().float() + mask_tensor = mask_tensor.mean(dim=0, keepdim=True)[:, :, :, 0] + else: + mask_tensor = torch.zeros((1, annotated_tensor.shape[1], annotated_tensor.shape[2]), dtype=torch.float32) + + return annotated_tensor, mask_tensor + + def _segment_polygons( + self, + image_pil: Image.Image, + predictions: Dict[str, Any], + fill_mask: bool, + ) -> Tuple[torch.Tensor, torch.Tensor]: + width, height = image_pil.size + mask_image = Image.new("RGB", (width, height), "black") + mask_draw = ImageDraw.Draw(mask_image) + + for polygons, label in zip(predictions["polygons"], predictions["labels"]): + color = random.choice(COLOR_BANK) + for polygon in polygons: + polygon = np.array(polygon).reshape(-1, 2) + polygon = np.clip(polygon, [0, 0], [width - 1, height - 1]) + if len(polygon) < 3: + continue + pts = polygon.reshape(-1).tolist() + if fill_mask: + overlay = Image.new("RGBA", image_pil.size, (255, 255, 255, 0)) + draw = ImageDraw.Draw(overlay) + rgba = ImageColor.getrgb(color) + (180,) + draw.polygon(pts, outline=color, fill=rgba, width=3) + image_pil = Image.alpha_composite(image_pil.convert("RGBA"), overlay).convert("RGB") + else: + draw = ImageDraw.Draw(image_pil) + draw.polygon(pts, outline=color, width=3) + mask_draw.polygon(pts, outline="white", fill="white") + + image_tensor = TF.to_tensor(image_pil)[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() + mask_tensor = TF.to_tensor(mask_image).unsqueeze(0).permute(0, 2, 3, 1).cpu().float() + mask_tensor = mask_tensor.mean(dim=0, keepdim=True)[:, :, :, 0] + return image_tensor, mask_tensor + + def analyze( + self, + image: torch.Tensor, + model_name: str, + task: str, + precision: str, + attention: str, + fill_mask: bool, + output_mask_select: str = "", + keep_model_loaded: bool = False, + text_prompt: str = "", + ): + bundle = self._get_model(model_name, precision, attention) + model = bundle["model"] + processor = bundle["processor"] + dtype = bundle["dtype"] + + device = mm.get_torch_device() + offload = mm.unet_offload_device() + model.to(device) + images = image.permute(0, 3, 1, 2) + + task_config = TASK_CONFIGS[task] + prompt = self._prepare_prompt(task, text_prompt) + out_images, out_masks, out_data = [], [], [] + pbar = ProgressBar(len(images)) + select_filter = [s.strip() for s in output_mask_select.split(",") if s.strip()] + + for tensor in images: + image_pil = TF.to_pil_image(tensor) + inputs = processor( + text=prompt, + images=image_pil, + return_tensors="pt", + do_rescale=False, + ).to(dtype).to(device) + + generated_ids = model.generate( + input_ids=inputs["input_ids"], + pixel_values=inputs["pixel_values"], + max_new_tokens=GENERATION_CONFIG["max_new_tokens"], + do_sample=GENERATION_CONFIG["do_sample"], + num_beams=GENERATION_CONFIG["num_beams"], + use_cache=False, + ) + raw = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] + parsed = processor.post_process_generation(raw, task=task_config["token"], image_size=image_pil.size) + predictions = parsed[task_config["token"]] + + if task_config["mode"] == "bbox": + annotated, mask = self._draw_regions(image_pil, predictions, fill_mask, select_filter) + out_images.append(annotated) + out_masks.append(mask) + else: + image_tensor, mask_tensor = self._segment_polygons(image_pil, predictions, fill_mask) + out_images.append(image_tensor) + out_masks.append(mask_tensor) + + out_data.append(predictions) + pbar.update(1) + + image_out = torch.cat(out_images, dim=0) + mask_out = torch.cat(out_masks, dim=0) + + if not keep_model_loaded: + model.to(offload) + mm.soft_empty_cache() + + return image_out, mask_out, out_data + + +NODE_CLASS_MAPPINGS = { + "AILab_Florence2": AILab_Florence2, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AILab_Florence2": "Florence2 (RMBG)", +} diff --git a/py/AILab_ImageCompareView.py b/py/AILab_ImageCompareView.py new file mode 100644 index 0000000..7ee1a28 --- /dev/null +++ b/py/AILab_ImageCompareView.py @@ -0,0 +1,201 @@ +# ComfyUI-RMBG +# Interactive Image Compare Node (ComfyUI Standard) +# License: GPL-3.0 + +import os +import random + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +import folder_paths + + +class AILab_ImageCompareView: + """Real-time Interactive Image Compare node. + + Renders two images side-by-side or with various comparison modes directly + on the node canvas. All UI controls (mode dropdown, match_size checkbox) + are native ComfyUI widgets — the frontend JS handles canvas rendering and + mouse interaction only. + + Outputs ``IMAGE`` (the composed result for the current mode) and + ``DIFF_MASK`` (the absolute-difference mask between the two inputs). + """ + + DESCRIPTION = ( + "Real-time Interactive Image Compare (RMBG) - Compare two images on-node. " + "Supports wipe (L/R & T/B), overlay opacity blend, pixel difference, " + "side-by-side, and highlight-diff (red on grey). Uses native ComfyUI " + "widgets for mode selection; the frontend handles interactive canvas " + "rendering." + ) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image1": ("IMAGE", {"tooltip": "First image to compare (e.g. before / original)."}), + "image2": ("IMAGE", {"tooltip": "Second image to compare (e.g. after / processed)."}), + "mode": ( + [ + "left_right", + "up_down", + "overlay", + "difference", + "side_by_side", + "highlight_diff", + ], + {"default": "left_right", "tooltip": "Comparison mode."}, + ), + "match_size": ( + "BOOLEAN", + { + "default": True, + "tooltip": "If True, resize image2 to match image1 dimensions when they differ.", + }, + ), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("IMAGE", "DIFF_MASK") + FUNCTION = "compare_images" + OUTPUT_NODE = True + CATEGORY = "🧪AILab/🖼️IMAGE" + + def __init__(self): + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" + self.prefix_append = "_cmp_" + "".join( + random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5) + ) + self.compress_level = 4 + + # ── Core entry point ─────────────────────────────────────── + def compare_images(self, image1, image2, mode="left_right", + match_size=True, prompt=None, extra_pnginfo=None): + # ── Both images provided (required inputs) ────────────── + # Ensure RGB (strip alpha if present) + t1 = image1[..., :3] if image1.shape[-1] == 4 else image1 + t2 = image2[..., :3] if image2.shape[-1] == 4 else image2 + + # Align batch size + b1, b2 = t1.shape[0], t2.shape[0] + if b1 != b2: + if b1 == 1: + t1 = t1.repeat(b2, 1, 1, 1) + elif b2 == 1: + t2 = t2.repeat(b1, 1, 1, 1) + else: + min_b = min(b1, b2) + t1 = t1[:min_b] + t2 = t2[:min_b] + + # Match spatial size if requested (resize image2 → image1) + if match_size and t1.shape[1:3] != t2.shape[1:3]: + nchw2 = t2.permute(0, 3, 1, 2) + resized2 = F.interpolate( + nchw2, size=(t1.shape[1], t1.shape[2]), + mode="bicubic", align_corners=False, + ) + t2 = torch.clamp(resized2.permute(0, 2, 3, 1), 0.0, 1.0) + + # Safety: if sizes still differ (match_size=False), crop both to the + # common minimum so _compose and diff_mask never crash on shape mismatch. + if t1.shape[1:3] != t2.shape[1:3]: + min_h = min(t1.shape[1], t2.shape[1]) + min_w = min(t1.shape[2], t2.shape[2]) + t1 = t1[:, :min_h, :min_w, :] + t2 = t2[:, :min_h, :min_w, :] + + # Save both images to temp directory for frontend interaction + results = self._save_images([(1, t1), (2, t2)]) + + # Difference mask (used by all modes) + diff_mask = torch.clamp(torch.abs(t1 - t2).mean(dim=-1), 0.0, 1.0) + + # Compose output based on mode + out_img = self._compose(mode, t1, t2) + + return {"ui": {"images": results}, "result": (out_img, diff_mask)} + + # ── Mode composition ────────────────────────────────────── + @staticmethod + def _compose(mode, t1, t2): + """Return the composed output tensor for *mode*.""" + if mode == "left_right": + mid = t1.shape[2] // 2 + out = t2.clone() + out[:, :, :mid, :] = t1[:, :, :mid, :] + elif mode == "up_down": + mid = t1.shape[1] // 2 + out = t2.clone() + out[:, :mid, :, :] = t1[:, :mid, :, :] + elif mode == "overlay": + out = torch.clamp(0.5 * t1 + 0.5 * t2, 0.0, 1.0) + elif mode == "difference": + out = torch.clamp(torch.abs(t1 - t2), 0.0, 1.0) + elif mode == "side_by_side": + out = torch.cat([t1, t2], dim=2) + elif mode == "highlight_diff": + # grey base (luminance of image1), red where pixels differ + grey = t1.mean(dim=-1, keepdim=True).repeat(1, 1, 1, 3) + diff = torch.abs(t1 - t2).mean(dim=-1, keepdim=True) + red = torch.zeros_like(t1) + red[..., 0] = 1.0 # pure red channel + threshold = 0.15 + mask = (diff > threshold).expand_as(t1).float() + out = grey * (1 - mask) + red * mask + else: + out = t1 + return out + + # ── Image saving helper ──────────────────────────────────── + def _save_images(self, slot_tensors): + """Save tensors to temp dir, tagged with *slot* (1 or 2). + + The frontend uses ``slot`` to distinguish image1 from image2. + """ + results = [] + os.makedirs(self.output_dir, exist_ok=True) + first_tensor = slot_tensors[0][1] + prefix = "ailab_cmp" + self.prefix_append + w = first_tensor.shape[2] + h = first_tensor.shape[1] + + full_output_folder, filename, counter, subfolder, _ = ( + folder_paths.get_save_image_path(prefix, self.output_dir, w, h) + ) + + for slot, tensor in slot_tensors: + arr = np.clip(255.0 * tensor[0].cpu().numpy(), 0, 255).astype(np.uint8) + img = Image.fromarray(arr) + file = f"{filename}_{counter:05}_{slot}.png" + img.save( + os.path.join(full_output_folder, file), + compress_level=self.compress_level, + ) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type, + "slot": slot, + }) + counter += 1 + + return results + + +NODE_CLASS_MAPPINGS = { + "AILab_ImageCompareView": AILab_ImageCompareView, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AILab_ImageCompareView": "Image Compare (RMBG) 🔀🖼️", +} diff --git a/py/AILab_LamaRemover.py b/py/AILab_LamaRemover.py index 24c9d1f..f066c6e 100644 --- a/py/AILab_LamaRemover.py +++ b/py/AILab_LamaRemover.py @@ -20,12 +20,7 @@ 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) +from AILab_utils import clean_vram, unload_model, tensor2pil, pil2tensor def pil2comfy(image): img_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0) @@ -68,6 +63,9 @@ class AILab_LamaRemover: "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"]}), }, + "optional": { + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), + }, } CATEGORY = "🧪AILab/🧽RMBG" @@ -100,9 +98,8 @@ class AILab_LamaRemover: self.model.to(self.device) def download_model(self): - print("Downloading Big-Lama model...") os.makedirs(self.cache_dir, exist_ok=True) - + print("Downloading Big-Lama model...") try: downloaded_path = hf_hub_download( repo_id="1038lab/Lama", @@ -110,22 +107,20 @@ class AILab_LamaRemover: 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") + target_path = os.path.join(self.cache_dir, "big-lama.pt") + shutil.move(downloaded_path, target_path) + print("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) - + img_tensor = img_tensor.to(self.device) + mask_tensor = mask_tensor.to(self.device) + + with torch.no_grad(): result = self.model(img_tensor, mask_tensor) - result_cpu = result[0].cpu() - + result_cpu = result.cpu() del img_tensor del mask_tensor if torch.cuda.is_available(): @@ -133,7 +128,7 @@ class AILab_LamaRemover: return result_cpu - def remove_object(self, images, masks, removal_strength, edge_smoothness): + def remove_object(self, images, masks, removal_strength, edge_smoothness, unload_model=False): try: self.load_model() results = [] @@ -179,8 +174,9 @@ class AILab_LamaRemover: print(traceback.format_exc()) raise RuntimeError(f"Error in object removal: {str(e)}") finally: - if torch.cuda.is_available(): - torch.cuda.empty_cache() + if unload_model: + self.model = None + clean_vram() NODE_CLASS_MAPPINGS = { "AILab_LamaRemover": AILab_LamaRemover, diff --git a/py/AILab_RMBG.py b/py/AILab_RMBG.py index 2ffe2d8..f18698f 100644 --- a/py/AILab_RMBG.py +++ b/py/AILab_RMBG.py @@ -1,707 +1,706 @@ -# 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 - -from AILab_utils import get_device -device = get_device() - -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) - -def normalize_process_res(process_res, step=128, minimum=256, maximum=2048): - process_res = max(minimum, min(maximum, int(process_res))) - normalized = max(minimum, (process_res // step) * step) - return min(maximum, normalized) - -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) - - normalized_process_res = normalize_process_res(params["process_res"]) - - def extract_results(outputs): - if isinstance(outputs, list) and len(outputs) > 0: - return outputs[-1].sigmoid().cpu() - if isinstance(outputs, dict) and 'logits' in outputs: - return outputs['logits'].sigmoid().cpu() - if isinstance(outputs, torch.Tensor): - return outputs.sigmoid().cpu() - - try: - if hasattr(outputs, 'last_hidden_state'): - return outputs.last_hidden_state.sigmoid().cpu() - for _, value in outputs.items(): - if isinstance(value, torch.Tensor): - return value.sigmoid().cpu() - except Exception: - pass - - handle_model_error("Unable to recognize model output format") - - # Prepare batch processing - transform_image = transforms.Compose([ - transforms.Resize((normalized_process_res, normalized_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(): - try: - results = extract_results(self.model(input_batch)) - - masks = [] - - for result, (orig_w, orig_h) in 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 batch_error: - if len(images) == 1 or "Sizes of tensors must match" not in str(batch_error): - raise - - print("[RMBG INFO] Batch inference failed due to tensor size mismatch; retrying images one by one.") - masks = [] - for img, (orig_w, orig_h) in zip(images, original_sizes): - single_input = transform_image(tensor2pil(img)).unsqueeze(0).to(device) - single_results = extract_results(self.model(single_input)) - result = single_results[0].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": 128, "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)" +# 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 +import types + +from AILab_utils import get_device, unload_model, tensor2pil, pil2tensor +device = get_device() + +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 handle_model_error(message): + print(f"[RMBG ERROR] {message}") + raise RuntimeError(message) + +def normalize_process_res(process_res, step=128, minimum=256, maximum=2048): + process_res = max(minimum, min(maximum, int(process_res))) + normalized = max(minimum, (process_res // step) * step) + return min(maximum, normalized) + +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) + module.__file__ = birefnet_path + 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: + from transformers import AutoModelForImageSegmentation + 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) + + normalized_process_res = normalize_process_res(params["process_res"]) + + def extract_results(outputs): + if isinstance(outputs, list) and len(outputs) > 0: + return outputs[-1].sigmoid().cpu() + if isinstance(outputs, dict) and 'logits' in outputs: + return outputs['logits'].sigmoid().cpu() + if isinstance(outputs, torch.Tensor): + return outputs.sigmoid().cpu() + + try: + if hasattr(outputs, 'last_hidden_state'): + return outputs.last_hidden_state.sigmoid().cpu() + for _, value in outputs.items(): + if isinstance(value, torch.Tensor): + return value.sigmoid().cpu() + except Exception: + pass + + handle_model_error("Unable to recognize model output format") + + # Prepare batch processing + transform_image = transforms.Compose([ + transforms.Resize((normalized_process_res, normalized_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(): + try: + results = extract_results(self.model(input_batch)) + + masks = [] + for result, (orig_w, orig_h) in 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 batch_error: + if len(images) == 1 or "Sizes of tensors must match" not in str(batch_error): + raise + + print("[RMBG INFO] Batch inference failed due to tensor size mismatch; retrying images one by one.") + masks = [] + for img, (orig_w, orig_h) in zip(images, original_sizes): + single_input = transform_image(tensor2pil(img)).unsqueeze(0).to(device) + single_results = extract_results(self.model(single_input)) + result = single_results[0].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): + import cv2 + 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).", + "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", + "unload_model": "Unload model from VRAM after execution to free up memory.", + "background": "Choose output type: Alpha (transparent) or Color (custom background color).", + "background_color": "Pick background color (supports alpha, use color picker)." + } + + 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": 128, "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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": tooltips["unload_model"]}), + "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) + finally: + if params.get("unload_model", False): + unload_model(model_instance, "model") + +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 index 4e9a29d..318b878 100644 --- a/py/AILab_SAM2Segment.py +++ b/py/AILab_SAM2Segment.py @@ -13,17 +13,10 @@ import comfy.model_management from hydra import initialize_config_dir from hydra.core.global_hydra import GlobalHydra +from AILab_utils import patch_transformers_bert, clean_vram, unload_model, get_or_download_model_file, process_mask, apply_background_color -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.") +# GroundingDINO imports are deferred to first use (segment_v2) to avoid ~12.5s import at startup. +GROUNDINGDINO_AVAILABLE = None # Tri-state: None = not yet checked, True/False = checked current_dir = Path(__file__).resolve().parent repo_root = current_dir.parent @@ -105,50 +98,6 @@ DINO_MODELS = { } } -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): @@ -177,6 +126,7 @@ class SAM2Segment: "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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}), "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), } @@ -249,8 +199,25 @@ class SAM2Segment: 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): + mask_blur=0, mask_offset=0, invert_output=False, unload_model=False, + background="Alpha", background_color="#222222"): + # Lazy import groundingdino (~12.5s) -- only when the node is actually executed + global GROUNDINGDINO_AVAILABLE + if GROUNDINGDINO_AVAILABLE is None: + try: + from groundingdino.util.slconfig import SLConfig # noqa: F401 + GROUNDINGDINO_AVAILABLE = True + except ImportError: + GROUNDINGDINO_AVAILABLE = False + print("Warning: GroundingDINO not available. Text prompts will use fallback method.") + if not GROUNDINGDINO_AVAILABLE: + raise RuntimeError("GroundingDINO is required for SAM2Segment but is not installed.") + 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 + device_obj = comfy.model_management.get_torch_device() # Process batch images @@ -274,6 +241,7 @@ class SAM2Segment: # Load and cache GroundingDINO model dino_key = (config_path, weights_path, device_obj) if dino_key not in self.dino_model_cache: + patch_transformers_bert() args = SLConfig.fromfile(config_path) model = build_model(args) checkpoint = load_file(weights_path) @@ -385,17 +353,23 @@ class SAM2Segment: 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)) + try: + # 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)) + finally: + if unload_model: + self.dino_model_cache.clear() + self.sam2_model_cache.clear() + clean_vram() NODE_CLASS_MAPPINGS = { "SAM2Segment": SAM2Segment, diff --git a/py/AILab_SAM3Segment.py b/py/AILab_SAM3Segment.py index 6f42573..96fa9cb 100644 --- a/py/AILab_SAM3Segment.py +++ b/py/AILab_SAM3Segment.py @@ -12,86 +12,42 @@ import folder_paths import comfy.model_management from AILab_ImageMaskTools import pil2tensor, tensor2pil +from AILab_utils import clean_vram, process_mask, apply_background_color, get_or_download_model_file +# ── paths ────────────────────────────────────────────────────────────── 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: + +if "sam3" not in sys.modules: + sys.path.insert(0, str(SAM3_LOCAL_DIR)) 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", +# ── safetensors import ───────────────────────────────────────────────── +try: + from safetensors.torch import load_file as safetensors_load_file +except ImportError: + safetensors_load_file = None + +# ── model registry ───────────────────────────────────────────────────── +SAM3_SAFETENSOR_MODELS = { + "sam3.1_multiplex_fp16": { + "model_url": "https://huggingface.co/1038lab/sam3/resolve/main/sam3.1_multiplex_fp16.safetensors", + "filename": "sam3.1_multiplex_fp16.safetensors", + "dtype": "float16", + }, + "sam3": { + "model_url": "https://huggingface.co/1038lab/sam3/resolve/main/sam3.safetensors", + "filename": "sam3.safetensors", + "dtype": "float32", + }, } -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() @@ -104,6 +60,43 @@ def _resolve_device(user_choice): return auto_device +# ── safetensors-aware checkpoint loader ──────────────────────────────── +def _load_safetensors_checkpoint(model, checkpoint_path): + """Load safetensors checkpoint with smart key-mapping.""" + if safetensors_load_file is None: + raise ImportError( + "safetensors is required for SAM3 Multiplex. " + "Install with: pip install safetensors" + ) + state_dict = safetensors_load_file(checkpoint_path, device="cpu") + + # Sniff key structure: if keys have 'detector.' prefix, strip it + has_detector_prefix = any(k.startswith("detector.") for k in state_dict) + if has_detector_prefix: + sam3_image_ckpt = { + k.replace("detector.", ""): v + for k, v in state_dict.items() + if "detector" in k + } + if model.inst_interactive_predictor is not None: + sam3_image_ckpt.update({ + k.replace("tracker.", "inst_interactive_predictor.model."): v + for k, v in state_dict.items() + if "tracker" in k + }) + else: + # Flat keys — load as-is + sam3_image_ckpt = state_dict + + missing_keys, unexpected_keys = model.load_state_dict(sam3_image_ckpt, strict=False) + if missing_keys: + print(f"[SAM3 Multiplex] loaded {checkpoint_path}") + print(f" missing keys ({len(missing_keys)}): {missing_keys[:10]}{'...' if len(missing_keys) > 10 else ''}") + if unexpected_keys: + print(f" unexpected keys ({len(unexpected_keys)}): {unexpected_keys[:10]}{'...' if len(unexpected_keys) > 10 else ''}") + + +# ── model builder import ─────────────────────────────────────────────── from sam3.model_builder import build_sam3_image_model # noqa: E402 from sam3.model.sam3_image_processor import Sam3Processor # noqa: E402 @@ -111,9 +104,11 @@ from sam3.model.sam3_image_processor import Sam3Processor # noqa: E402 class SAM3Segment: @classmethod def INPUT_TYPES(cls): + model_names = list(SAM3_SAFETENSOR_MODELS.keys()) return { "required": { "image": ("IMAGE",), + "model_name": (model_names, {"default": model_names[0]}), "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}), @@ -139,24 +134,36 @@ class SAM3Segment: def __init__(self): self.processor_cache = {} - def _load_processor(self, device_choice): + def _load_processor(self, model_name, device_choice): torch_device = _resolve_device(device_choice) device_str = "cuda" if torch_device.type == "cuda" else "cpu" - cache_key = ("sam3", device_str) + cache_key = (model_name, 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_info = SAM3_SAFETENSOR_MODELS[model_name] + ckpt_path = get_or_download_model_file(model_info["filename"], model_info["model_url"], "sam3") + + # Build the model architecture (without loading weights) model = build_sam3_image_model( bpe_path=SAM3_BPE_PATH, - device=device_str, + device="cpu", # load on CPU first, move later eval_mode=True, - checkpoint_path=ckpt_path, + checkpoint_path=None, # we load manually below load_from_HF=False, enable_segmentation=True, enable_inst_interactivity=False, ) + + # Load safetensors checkpoint + _load_safetensors_checkpoint(model, ckpt_path) + + # Move to target device (keep FP32 — autocast handles precision at inference) + if device_str == "cuda": + model = model.cuda() + model.eval() + processor = Sam3Processor(model, device=device_str) self.processor_cache[cache_key] = processor + print(f"[SAM3 Multiplex] Loaded {model_name} on {device_str} ({model_info['dtype']})") return self.processor_cache[cache_key], torch_device def _empty_result(self, img_pil, background, background_color): @@ -230,18 +237,8 @@ class SAM3Segment: 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, + 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) @@ -255,48 +252,35 @@ class SAM3Segment: 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"): + def segment(self, image, model_name, 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) + processor, torch_device = self._load_processor(model_name, 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() + + # Use bfloat16 for original sam3, float16 for multiplex fp16 + model_info = SAM3_SAFETENSOR_MODELS[model_name] + autocast_dtype = torch.float16 if model_info["dtype"] == "float16" else torch.bfloat16 + ctx = torch.autocast(autocast_device, dtype=autocast_dtype) 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, + 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, + 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) @@ -304,30 +288,25 @@ class SAM3Segment: if unload_model: device_str = "cuda" if torch_device.type == "cuda" else "cpu" - cache_key = ("sam3", device_str) + cache_key = (model_name, device_str) if cache_key in self.processor_cache: del self.processor_cache[cache_key] - if torch_device.type == "cuda": - torch.cuda.empty_cache() + clean_vram() - # 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 index f4ebf94..a32e107 100644 --- a/py/AILab_SDMatte.py +++ b/py/AILab_SDMatte.py @@ -1,435 +1,404 @@ -# 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)", -} +# 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.") + +# Lazy: diffusers/transformers availability is checked on first use to avoid ~3.7s import at startup +DIFFUSERS_AVAILABLE = None # Tri-state: None = not yet checked + +from AILab_utils import get_device, clean_vram, unload_model, get_or_download_model_file, process_mask, apply_background_color, pil2tensor, tensor2pil + +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 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 + +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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), + "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: + global DIFFUSERS_AVAILABLE + if DIFFUSERS_AVAILABLE is None: + try: + import diffusers # noqa: F401 + import transformers # noqa: F401 + DIFFUSERS_AVAILABLE = True + except ImportError: + DIFFUSERS_AVAILABLE = False + 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, unload_model=False, background="Alpha", background_color="#222222"): + if device == "CPU": + device_obj = torch.device('cpu') + elif device == "GPU": + device_obj = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + else: + device_obj = get_device() + + sdmatte_model = self.load_sdmatte_model(model, device) + + batch_size = image.shape[0] + + result_masks = [] + result_images = [] + result_mask_images = [] + + try: + 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: + if mask.ndim == 2: + mask_single = mask.unsqueeze(0).unsqueeze(0) + elif mask.ndim == 3: + idx = min(b, mask.shape[0] - 1) + mask_single = mask[idx:idx+1].unsqueeze(1) + elif mask.ndim == 4: + idx = min(b, mask.shape[0] - 1) + mask_single = mask[idx:idx+1] + else: + mask_single = mask.view(1, 1, H, W) + mask_b1hw = mask_single.contiguous().to(device_obj) + mask_for_refine = mask_b1hw.squeeze(1).cpu() + elif 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).cpu() + else: + print("[SDMatte INFO] No input mask provided; using full foreground mask as fallback. For best results, connect a MASK from RMBG/BiRefNet/SAM node.") + mask_b1hw = torch.ones((1, 1, H, W), dtype=torch.float32, device=device_obj) + mask_for_refine = torch.ones((1, H, W), dtype=torch.float32) + + 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) + + return (torch.cat(result_images, dim=0), torch.cat(result_masks, dim=0), torch.cat(result_mask_images, dim=0)) + finally: + if unload_model: + self.model_cache.clear() + clean_vram() + + 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 index 189dad5..3c220ac 100644 --- a/py/AILab_Segment.py +++ b/py/AILab_Segment.py @@ -28,6 +28,7 @@ from torch.hub import download_url_to_file import folder_paths import comfy.model_management from segment_anything import sam_model_registry, SamPredictor +from AILab_utils import get_device, patch_transformers_bert, clean_vram, unload_model, process_mask, apply_background_color, image2mask, pil2tensor, tensor2pil SAM_MODELS = { "sam_vit_h (2.56GB)": { @@ -95,55 +96,6 @@ def split_image_mask(image): 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): @@ -168,6 +120,7 @@ class Segment: "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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}), "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), } @@ -192,55 +145,61 @@ class Segment: 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_blur=0, mask_offset=0, invert_output=False, unload_model=False, + background="Alpha", background_color="#222222"): + try: + print(f'Processing create segment for: "{prompt}"...') - mask_tensor = image2mask(mask_image).unsqueeze(0) + 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) - 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) + 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) + finally: + if unload_model: + self._sam_model_cache.clear() + self._dino_model_cache.clear() + clean_vram() def load_sam(self, model_name): if model_name in self._sam_model_cache: @@ -275,6 +234,7 @@ class Segment: "grounding-dino" ) ) + patch_transformers_bert() dino = self.build_model(dino_model_args) checkpoint = torch.load( self.get_local_filepath( diff --git a/py/AILab_SegmentV2.py b/py/AILab_SegmentV2.py index 1b8bbcf..a67a57c 100644 --- a/py/AILab_SegmentV2.py +++ b/py/AILab_SegmentV2.py @@ -8,14 +8,9 @@ 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 -from AILab_utils import get_device +from AILab_utils import get_device, patch_transformers_bert, clean_vram, unload_model, get_or_download_model_file, process_mask, apply_background_color SAM_MODELS = { "sam_vit_h (2.56GB)": { @@ -65,49 +60,8 @@ DINO_MODELS = { } } -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): + from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection processor = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-tiny") model = AutoModelForZeroShotObjectDetection.from_pretrained("IDEA-Research/grounding-dino-tiny").to(device) return processor, model @@ -149,6 +103,7 @@ class SegmentV2: "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"]}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}), "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), } @@ -164,8 +119,14 @@ class SegmentV2: 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): + mask_blur=0, mask_offset=0, invert_output=False, unload_model=False, + background="Alpha", background_color="#222222"): + # Lazy import groundingdino (~12.5s) -- only when the node is actually executed + 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 + device = get_device() batch_size = image.shape[0] if len(image.shape) == 4 else 1 @@ -182,6 +143,7 @@ class SegmentV2: 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: + patch_transformers_bert() args = SLConfig.fromfile(config_path) model = build_model(args) checkpoint = torch.load(weights_path, map_location="cpu") @@ -268,14 +230,20 @@ class SegmentV2: 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)) + try: + 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)) + finally: + if unload_model: + self.dino_model_cache.clear() + self.sam_model_cache.clear() + clean_vram() NODE_CLASS_MAPPINGS = { "SegmentV2": SegmentV2, diff --git a/py/AILab_YoloV8.py b/py/AILab_YoloV8.py index fe96094..85a3d28 100644 --- a/py/AILab_YoloV8.py +++ b/py/AILab_YoloV8.py @@ -1,279 +1,289 @@ -import os -from typing import Dict, List, Optional - -import numpy as np -import torch -from PIL import Image, ImageDraw - -import folder_paths -from AILab_ImageMaskTools import pil2tensor, tensor2pil - -ULTRALYTICS_DIR = os.path.join(folder_paths.models_dir, "ultralytics") -YOLO_LEGACY_DIR = os.path.join(folder_paths.models_dir, "yolo") -os.makedirs(ULTRALYTICS_DIR, exist_ok=True) -os.makedirs(YOLO_LEGACY_DIR, exist_ok=True) -folder_paths.add_model_folder_path("ultralytics", ULTRALYTICS_DIR, is_default=True) -folder_paths.add_model_folder_path("ultralytics", YOLO_LEGACY_DIR) - -DEVICE_CHOICES = ("auto", "cuda", "cpu", "mps") -MASK_COUNT_CHOICES = ("all", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10") -MASK_INDEX_CHOICES = ("none", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10") - - -class AILab_YoloV8Adv: - CATEGORY = "🧪AILab/🧽RMBG" - RETURN_TYPES = ("IMAGE", "MASK", "MASK") - RETURN_NAMES = ("ANNOTATED_IMAGE", "MASK", "MASK_LIST") - FUNCTION = "yolo_detect" - - _MODEL_CACHE: Dict[str, "YOLO"] = {} - - @classmethod - def _list_models(cls) -> List[str]: - files = folder_paths.get_filename_list("ultralytics") - return sorted(f for f in files if f.lower().endswith(".pt")) - - @classmethod - def INPUT_TYPES(cls): - models = cls._list_models() - if not models: - models = [f"Put .pt models into {ULTRALYTICS_DIR}"] - default_model = models[0] - - return { - "required": { - "images": ("IMAGE",), - "yolo_model": (tuple(models), {"default": default_model, "tooltip": f"YOLOv8 weights stored under {ULTRALYTICS_DIR} (subfolders allowed)."}), - "mask_count": (MASK_COUNT_CHOICES, {"default": "all", "tooltip": "Merge this many detections. 'all' merges everything (or just the selected index when specified)."}), - }, - "optional": { - "select_mask_index": (MASK_INDEX_CHOICES, {"default": "none", "tooltip": "1-based index of the first mask to keep. Use 'none' to start from the first detection."}), - "conf": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Confidence threshold forwarded to Ultralytics."}), - "iou": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "IOU used during NMS."}), - "classes": ("STRING", {"default": "", "placeholder": "e.g. 0,2,5-7", "tooltip": "Comma list or ranges of class IDs; empty keeps every class."}), - "device": (DEVICE_CHOICES, {"default": "auto", "tooltip": "Force a device or auto-detect CUDA → MPS → CPU."}), - "max_det": ("INT", {"default": 300, "min": 1, "max": 1000, "step": 1, "tooltip": "Maximum detections per image."}), - "retina_masks": ("BOOLEAN", {"default": True, "tooltip": "Use high-resolution masks (Ultralytics retina_masks flag)."}), - "agnostic_nms": ("BOOLEAN", {"default": False, "tooltip": "Enable class-agnostic NMS."}), - }, - } - - def _resolve_device(self, requested: str) -> str: - if requested != "auto": - return requested - if torch.cuda.is_available(): - return "cuda" - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return "mps" - return "cpu" - - def _parse_classes(self, value: str) -> Optional[List[int]]: - if not value or not value.strip(): - return None - classes: List[int] = [] - try: - for chunk in value.split(","): - chunk = chunk.strip() - if not chunk: - continue - if "-" in chunk: - start, end = [int(x) for x in chunk.split("-", 1)] - if start > end: - start, end = end, start - classes.extend(range(start, end + 1)) - else: - classes.append(int(chunk)) - return sorted(set(classes)) - except ValueError: - print(f"[AILab_YoloV8] Invalid classes string: {value}. Ignoring filter.") - return None - - def _resolve_model_path(self, name: str) -> str: - return folder_paths.get_full_path_or_raise("ultralytics", name) - - def _get_model(self, model_path: str): - model = self._MODEL_CACHE.get(model_path) - if model is None: - from ultralytics import YOLO - model = YOLO(model_path) - self._MODEL_CACHE[model_path] = model - return model - - def _result_to_tensor(self, result) -> torch.Tensor: - plotted = result.plot() - rgb = plotted[..., ::-1] - return pil2tensor(Image.fromarray(rgb)) - - def _mask_from_tensor(self, mask_tensor: torch.Tensor, size: Image.Image.size): - mask_np = mask_tensor.detach().cpu().numpy() - mask_img = Image.fromarray((mask_np * 255).astype(np.uint8)) - if mask_img.size != size: - mask_img = mask_img.resize(size, Image.Resampling.NEAREST) - return torch.from_numpy(np.array(mask_img).astype(np.float32) / 255.0) - - def _collect_masks(self, result, size) -> List[torch.Tensor]: - width, height = size - masks: List[torch.Tensor] = [] - - if getattr(result, "masks", None) is not None and result.masks.data is not None: - for mask_tensor in result.masks.data: - masks.append(self._mask_from_tensor(mask_tensor, size)) - - elif getattr(result, "boxes", None) is not None and len(result.boxes.xyxy) > 0: - for box in result.boxes: - x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()] - mask_img = Image.new("L", size, 0) - draw = ImageDraw.Draw(mask_img) - draw.rectangle([x1, y1, x2, y2], fill=255) - masks.append(torch.from_numpy(np.array(mask_img).astype(np.float32) / 255.0)) - - if not masks: - masks.append(torch.zeros((height, width), dtype=torch.float32)) - - return masks - - def _merge_masks(self, masks: List[torch.Tensor]) -> torch.Tensor: - if not masks: - raise ValueError("Cannot merge empty mask list.") - merged = torch.zeros_like(masks[0]) - - for mask in masks: - merged = torch.maximum(merged, mask) - - return merged - - def yolo_detect( - self, - images, - yolo_model, - mask_count="all", - conf=0.25, - iou=0.45, - classes="", - device="auto", - max_det=300, - retina_masks=True, - agnostic_nms=False, - select_mask_index: str = "none", - ): - model_path = self._resolve_model_path(yolo_model) - model = self._get_model(model_path) - device_target = self._resolve_device(device) - class_filter = self._parse_classes(classes) - - merged_masks: List[torch.Tensor] = [] - annotated_images: List[torch.Tensor] = [] - mask_list: List[torch.Tensor] = [] - - count_limit = 0 if mask_count == "all" else max(0, int(mask_count)) - chosen_index: Optional[int] = None - if select_mask_index != "none": - chosen_index = int(select_mask_index) - 1 - - for idx in range(images.shape[0]): - image_pil = tensor2pil(images[idx]) - - results = model( - image_pil, - conf=conf, - iou=iou, - classes=class_filter, - device=device_target, - max_det=max_det, - retina_masks=retina_masks, - agnostic_nms=agnostic_nms, - ) - - if not results: - continue - - result = results[0] - annotated_images.append(self._result_to_tensor(result)) - - frame_masks = self._collect_masks(result, image_pil.size) - - selected_masks: List[torch.Tensor] - if chosen_index is None: - if count_limit <= 0 or count_limit >= len(frame_masks): - selected_masks = frame_masks - else: - selected_masks = frame_masks[:count_limit] - else: - if chosen_index >= len(frame_masks): - selected_masks = [] - else: - span = count_limit if count_limit > 0 else 1 - selected_masks = frame_masks[chosen_index : chosen_index + span] - - if selected_masks: - merged_masks.append(self._merge_masks(selected_masks)) - mask_list.extend(selected_masks) - else: - fallback = torch.zeros_like(frame_masks[0]) - merged_masks.append(fallback) - mask_list.append(fallback) - - if not merged_masks: - width, height = tensor2pil(images[0]).size - merged_masks = [torch.zeros((height, width), dtype=torch.float32)] - - if not mask_list: - width, height = merged_masks[0].shape[1], merged_masks[0].shape[0] - mask_list = [torch.zeros((height, width), dtype=torch.float32)] - - if not annotated_images: - annotated_images = [images] - - merged_tensor = torch.stack(merged_masks, dim=0) - annotated_tensor = torch.cat(annotated_images, dim=0) - mask_tensor = torch.stack(mask_list, dim=0) - - return annotated_tensor, merged_tensor, mask_tensor - - -class AILab_YoloV8(AILab_YoloV8Adv): - FUNCTION = "yolo_detect_simple" - - @classmethod - def INPUT_TYPES(cls): - models = cls._list_models() - if not models: - models = [f"Put .pt models into {ULTRALYTICS_DIR}"] - default_model = models[0] - - return { - "required": { - "images": ("IMAGE",), - "yolo_model": (tuple(models), {"default": default_model, "tooltip": f"YOLOv8 weights stored under {ULTRALYTICS_DIR}. Advanced controls available on YOLOv8 Adv."}), - "mask_count": (MASK_COUNT_CHOICES, {"default": "all", "tooltip": "Merge this many detections. 'all' merges everything (or just the selected index when specified)."}), - }, - "optional": { - "select_mask_index": (MASK_INDEX_CHOICES, {"default": "none", "tooltip": "1-based index of the first mask to keep. Use 'none' to start from the first detection."}), - }, - } - - def yolo_detect_simple(self, images, yolo_model, mask_count="all", select_mask_index="none"): - return super().yolo_detect( - images=images, - yolo_model=yolo_model, - mask_count=mask_count, - conf=0.25, - iou=0.45, - classes="", - device="auto", - max_det=300, - retina_masks=True, - agnostic_nms=False, - select_mask_index=select_mask_index, - ) - - -NODE_CLASS_MAPPINGS = { - "AILab_YoloV8": AILab_YoloV8, - "AILab_YoloV8Adv": AILab_YoloV8Adv, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "AILab_YoloV8": "YOLOv8 (RMBG)", - "AILab_YoloV8Adv": "YOLOv8 Adv (RMBG)", -} +import os +from typing import Dict, List, Optional + +import numpy as np +import torch +from PIL import Image, ImageDraw + +import folder_paths +from AILab_ImageMaskTools import pil2tensor, tensor2pil +from AILab_utils import clean_vram + +ULTRALYTICS_DIR = os.path.join(folder_paths.models_dir, "ultralytics") +YOLO_LEGACY_DIR = os.path.join(folder_paths.models_dir, "yolo") +os.makedirs(ULTRALYTICS_DIR, exist_ok=True) +os.makedirs(YOLO_LEGACY_DIR, exist_ok=True) +folder_paths.add_model_folder_path("ultralytics", ULTRALYTICS_DIR, is_default=True) +folder_paths.add_model_folder_path("ultralytics", YOLO_LEGACY_DIR) + +DEVICE_CHOICES = ("auto", "cuda", "cpu", "mps") +MASK_COUNT_CHOICES = ("all", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10") +MASK_INDEX_CHOICES = ("none", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10") + + +class AILab_YoloV8Adv: + CATEGORY = "🧪AILab/🧽RMBG" + RETURN_TYPES = ("IMAGE", "MASK", "MASK") + RETURN_NAMES = ("ANNOTATED_IMAGE", "MASK", "MASK_LIST") + FUNCTION = "yolo_detect" + + _MODEL_CACHE: Dict[str, "YOLO"] = {} + + @classmethod + def _list_models(cls) -> List[str]: + files = folder_paths.get_filename_list("ultralytics") + return sorted(f for f in files if f.lower().endswith(".pt")) + + @classmethod + def INPUT_TYPES(cls): + models = cls._list_models() + if not models: + models = [f"Put .pt models into {ULTRALYTICS_DIR}"] + default_model = models[0] + + return { + "required": { + "images": ("IMAGE",), + "yolo_model": (tuple(models), {"default": default_model, "tooltip": f"YOLOv8 weights stored under {ULTRALYTICS_DIR} (subfolders allowed)."}), + "mask_count": (MASK_COUNT_CHOICES, {"default": "all", "tooltip": "Merge this many detections. 'all' merges everything (or just the selected index when specified)."}), + }, + "optional": { + "select_mask_index": (MASK_INDEX_CHOICES, {"default": "none", "tooltip": "1-based index of the first mask to keep. Use 'none' to start from the first detection."}), + "conf": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Confidence threshold forwarded to Ultralytics."}), + "iou": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "IOU used during NMS."}), + "classes": ("STRING", {"default": "", "placeholder": "e.g. 0,2,5-7", "tooltip": "Comma list or ranges of class IDs; empty keeps every class."}), + "device": (DEVICE_CHOICES, {"default": "auto", "tooltip": "Force a device or auto-detect CUDA → MPS → CPU."}), + "max_det": ("INT", {"default": 300, "min": 1, "max": 1000, "step": 1, "tooltip": "Maximum detections per image."}), + "retina_masks": ("BOOLEAN", {"default": True, "tooltip": "Use high-resolution masks (Ultralytics retina_masks flag)."}), + "agnostic_nms": ("BOOLEAN", {"default": False, "tooltip": "Enable class-agnostic NMS."}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), + }, + } + + def _resolve_device(self, requested: str) -> str: + if requested != "auto": + return requested + if torch.cuda.is_available(): + return "cuda" + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return "mps" + return "cpu" + + def _parse_classes(self, value: str) -> Optional[List[int]]: + if not value or not value.strip(): + return None + classes: List[int] = [] + try: + for chunk in value.split(","): + chunk = chunk.strip() + if not chunk: + continue + if "-" in chunk: + start, end = [int(x) for x in chunk.split("-", 1)] + if start > end: + start, end = end, start + classes.extend(range(start, end + 1)) + else: + classes.append(int(chunk)) + return sorted(set(classes)) + except ValueError: + print(f"[AILab_YoloV8] Invalid classes string: {value}. Ignoring filter.") + return None + + def _resolve_model_path(self, name: str) -> str: + return folder_paths.get_full_path_or_raise("ultralytics", name) + + def _get_model(self, model_path: str): + model = self._MODEL_CACHE.get(model_path) + if model is None: + from ultralytics import YOLO + model = YOLO(model_path) + self._MODEL_CACHE[model_path] = model + return model + + def _result_to_tensor(self, result) -> torch.Tensor: + plotted = result.plot() + rgb = plotted[..., ::-1] + return pil2tensor(Image.fromarray(rgb)) + + def _mask_from_tensor(self, mask_tensor: torch.Tensor, size: Image.Image.size): + mask_np = mask_tensor.detach().cpu().numpy() + mask_img = Image.fromarray((mask_np * 255).astype(np.uint8)) + if mask_img.size != size: + mask_img = mask_img.resize(size, Image.Resampling.NEAREST) + return torch.from_numpy(np.array(mask_img).astype(np.float32) / 255.0) + + def _collect_masks(self, result, size) -> List[torch.Tensor]: + width, height = size + masks: List[torch.Tensor] = [] + + if getattr(result, "masks", None) is not None and result.masks.data is not None: + for mask_tensor in result.masks.data: + masks.append(self._mask_from_tensor(mask_tensor, size)) + + elif getattr(result, "boxes", None) is not None and len(result.boxes.xyxy) > 0: + for box in result.boxes: + x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()] + mask_img = Image.new("L", size, 0) + draw = ImageDraw.Draw(mask_img) + draw.rectangle([x1, y1, x2, y2], fill=255) + masks.append(torch.from_numpy(np.array(mask_img).astype(np.float32) / 255.0)) + + if not masks: + masks.append(torch.zeros((height, width), dtype=torch.float32)) + + return masks + + def _merge_masks(self, masks: List[torch.Tensor]) -> torch.Tensor: + if not masks: + raise ValueError("Cannot merge empty mask list.") + merged = torch.zeros_like(masks[0]) + + for mask in masks: + merged = torch.maximum(merged, mask) + + return merged + + def yolo_detect( + self, + images, + yolo_model, + mask_count="all", + conf=0.25, + iou=0.45, + classes="", + device="auto", + max_det=300, + retina_masks=True, + agnostic_nms=False, + select_mask_index: str = "none", + unload_model: bool = False, + ): + try: + model_path = self._resolve_model_path(yolo_model) + model = self._get_model(model_path) + device_target = self._resolve_device(device) + class_filter = self._parse_classes(classes) + + merged_masks: List[torch.Tensor] = [] + annotated_images: List[torch.Tensor] = [] + mask_list: List[torch.Tensor] = [] + + count_limit = 0 if mask_count == "all" else max(0, int(mask_count)) + chosen_index: Optional[int] = None + if select_mask_index != "none": + chosen_index = int(select_mask_index) - 1 + + for idx in range(images.shape[0]): + image_pil = tensor2pil(images[idx]) + + results = model( + image_pil, + conf=conf, + iou=iou, + classes=class_filter, + device=device_target, + max_det=max_det, + retina_masks=retina_masks, + agnostic_nms=agnostic_nms, + ) + + if not results: + continue + + result = results[0] + annotated_images.append(self._result_to_tensor(result)) + + frame_masks = self._collect_masks(result, image_pil.size) + + selected_masks: List[torch.Tensor] + if chosen_index is None: + if count_limit <= 0 or count_limit >= len(frame_masks): + selected_masks = frame_masks + else: + selected_masks = frame_masks[:count_limit] + else: + if chosen_index >= len(frame_masks): + selected_masks = [] + else: + span = count_limit if count_limit > 0 else 1 + selected_masks = frame_masks[chosen_index : chosen_index + span] + + if selected_masks: + merged_masks.append(self._merge_masks(selected_masks)) + mask_list.extend(selected_masks) + else: + fallback = torch.zeros_like(frame_masks[0]) + merged_masks.append(fallback) + mask_list.append(fallback) + + if not merged_masks: + width, height = tensor2pil(images[0]).size + merged_masks = [torch.zeros((height, width), dtype=torch.float32)] + + if not mask_list: + width, height = merged_masks[0].shape[1], merged_masks[0].shape[0] + mask_list = [torch.zeros((height, width), dtype=torch.float32)] + + if not annotated_images: + annotated_images = [images] + + merged_tensor = torch.stack(merged_masks, dim=0) + annotated_tensor = torch.cat(annotated_images, dim=0) + mask_tensor = torch.stack(mask_list, dim=0) + + return annotated_tensor, merged_tensor, mask_tensor + finally: + if unload_model: + self._MODEL_CACHE.clear() + clean_vram() + + +class AILab_YoloV8(AILab_YoloV8Adv): + FUNCTION = "yolo_detect_simple" + + @classmethod + def INPUT_TYPES(cls): + models = cls._list_models() + if not models: + models = [f"Put .pt models into {ULTRALYTICS_DIR}"] + default_model = models[0] + + return { + "required": { + "images": ("IMAGE",), + "yolo_model": (tuple(models), {"default": default_model, "tooltip": f"YOLOv8 weights stored under {ULTRALYTICS_DIR}. Advanced controls available on YOLOv8 Adv."}), + "mask_count": (MASK_COUNT_CHOICES, {"default": "all", "tooltip": "Merge this many detections. 'all' merges everything (or just the selected index when specified)."}), + }, + "optional": { + "select_mask_index": (MASK_INDEX_CHOICES, {"default": "none", "tooltip": "1-based index of the first mask to keep. Use 'none' to start from the first detection."}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM after execution"}), + }, + } + + def yolo_detect_simple(self, images, yolo_model, mask_count="all", select_mask_index="none", unload_model=False): + return super().yolo_detect( + images=images, + yolo_model=yolo_model, + mask_count=mask_count, + conf=0.25, + iou=0.45, + classes="", + device="auto", + max_det=300, + retina_masks=True, + agnostic_nms=False, + select_mask_index=select_mask_index, + unload_model=unload_model, + ) + + +NODE_CLASS_MAPPINGS = { + "AILab_YoloV8": AILab_YoloV8, + "AILab_YoloV8Adv": AILab_YoloV8Adv, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AILab_YoloV8": "YOLOv8 (RMBG)", + "AILab_YoloV8Adv": "YOLOv8 Adv (RMBG)", +} diff --git a/py/AILab_utils.py b/py/AILab_utils.py index 4b66c81..9edc01d 100644 --- a/py/AILab_utils.py +++ b/py/AILab_utils.py @@ -1,129 +1,349 @@ -# Shared utility helpers -import numpy as np -import torch -from PIL import Image -from comfy.utils import common_upscale -import comfy.model_management as mm - - -def get_device(): - """Return the torch device ComfyUI selected (CUDA, MPS, or CPU). - - Defers to ComfyUI's device management instead of assuming CUDA-or-CPU, so nodes - follow whatever backend ComfyUI chose. Fixes models silently running on CPU on - Apple Silicon (MPS). See issues #200 and #135. - """ - return mm.get_torch_device() - - -def tensor2pil(image: torch.Tensor) -> Image.Image: - return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) - - -def pil2tensor(image: Image.Image) -> torch.Tensor: - return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) - - -def pil2mask(image: Image.Image) -> torch.Tensor: - return torch.from_numpy(np.array(image.convert("L")).astype(np.float32) / 255.0).unsqueeze(0) - - -def resize_image(img: Image.Image, width: int, height: int) -> Image.Image: - return img.resize((width, height), resample=Image.LANCZOS) - - -def blend_overlay(img_1: Image.Image, img_2: Image.Image) -> Image.Image: - arr1 = np.array(img_1).astype(float) / 255.0 - arr2 = np.array(img_2).astype(float) / 255.0 - mask = arr2 < 0.5 - result = np.zeros_like(arr1) - result[mask] = 2 * arr1[mask] * arr2[mask] - result[~mask] = 1 - 2 * (1 - arr1[~mask]) * (1 - arr2[~mask]) - return Image.fromarray(np.clip(result * 255, 0, 255).astype(np.uint8)) - - -def fill_mask(width: int, height: int, mask: Image.Image, box=(0, 0), color=0) -> Image.Image: - bg = Image.new("L", (width, height), color) - bg.paste(mask, box, mask) - return bg - - -def empty_image(width: int, height: int, batch_size: int = 1) -> torch.Tensor: - return torch.zeros([batch_size, height, width, 3]) - - -def upscale_mask(mask: torch.Tensor, width: int, height: int) -> torch.Tensor: - if mask.ndim == 3: - mask = mask.unsqueeze(1) - mask = common_upscale(mask, width, height, "bicubic", "disabled") - mask = mask.squeeze(1) - return mask - - -def extract_alpha_mask(image: torch.Tensor) -> torch.Tensor: - alpha = image[..., 3] - if alpha.max() > 1.0: - alpha = alpha / 255.0 - if len(alpha.shape) == 4: - alpha = alpha[:, :, :, 0] - return alpha.unsqueeze(1) if alpha.ndim == 3 else alpha - - -def ensure_mask_shape(mask: torch.Tensor | None) -> torch.Tensor | None: - if mask is None: - return None - if mask.ndim == 2: - return mask.unsqueeze(0) - if mask.ndim == 4 and mask.shape[1] == 1: - return mask.squeeze(1) - return mask - - -COLOR_PRESETS = { - "black": "#000000", - "white": "#FFFFFF", - "red": "#FF0000", - "green": "#00FF00", - "blue": "#0000FF", - "yellow": "#FFFF00", - "cyan": "#00FFFF", - "magenta": "#FF00FF", - "gray": "#808080", - "silver": "#C0C0C0", - "maroon": "#800000", - "olive": "#808000", - "purple": "#800080", - "teal": "#008080", - "navy": "#000080", - "orange": "#FFA500", - "pink": "#FFC0CB", - "brown": "#A52A2A", - "violet": "#EE82EE", - "indigo": "#4B0082", - "light_gray": "#D3D3D3", - "dark_gray": "#A9A9A9", - "light_blue": "#ADD8E6", - "dark_blue": "#00008B", - "light_green": "#90EE90", - "dark_green": "#006400", -} - - -def color_format(color: str) -> str: - if not color: - return "" - - color = color.strip().upper() - if not color.startswith("#"): - color = f"#{color}" - - color = color[1:] - if len(color) == 3: - r, g, b = color[0], color[1], color[2] - return f"#{r}{r}{g}{g}{b}{b}" - if len(color) < 6: - raise ValueError(f"Invalid color format: {color}") - if len(color) > 6: - color = color[:6] - - return f"#{color}" +# Shared utility helpers +import os +import gc +from typing import Union, Optional +import numpy as np +import torch +from PIL import Image, ImageFilter +from torch.hub import download_url_to_file +from comfy.utils import common_upscale +import comfy.model_management as mm +import folder_paths + + +def get_device(): + """Return the torch device ComfyUI selected (CUDA, MPS, or CPU).""" + return mm.get_torch_device() + + +def clean_vram(): + """Safely clear cached VRAM across CUDA, MPS, and CPU backends.""" + gc.collect() + mm.soft_empty_cache() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if hasattr(torch.cuda, "ipc_collect"): + torch.cuda.ipc_collect() + elif hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): + try: + torch.mps.empty_cache() + except Exception: + pass + + +def unload_model(loader_or_model, attr_name="model"): + """Unload a model from memory and release VRAM. + + Accepts either: + 1. A loader object that owns a model attribute (e.g. loader, attr_name='model') + 2. A raw torch.nn.Module instance directly. + """ + if hasattr(loader_or_model, attr_name): + model = getattr(loader_or_model, attr_name) + if model is not None: + if hasattr(model, "to"): + try: + model.to("cpu") + except Exception: + pass + setattr(loader_or_model, attr_name, None) + if hasattr(loader_or_model, "current_model_version"): + loader_or_model.current_model_version = None + elif isinstance(loader_or_model, torch.nn.Module): + try: + loader_or_model.to("cpu") + except Exception: + pass + del loader_or_model + + clean_vram() + + +def patch_transformers_bert(): + """Ensure BertModel and BertModelWarper compatibility for GroundingDINO on transformers >= 5.0.""" + try: + def _get_head_mask(self, head_mask=None, num_hidden_layers=12, is_attention_chunked=False): + if head_mask is not None: + if hasattr(head_mask, "dim"): + if head_mask.dim() == 1: + head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) + head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1) + elif head_mask.dim() == 2: + head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) + if is_attention_chunked and hasattr(head_mask, "unsqueeze"): + head_mask = head_mask.unsqueeze(-1) + return head_mask + return [None] * num_hidden_layers + + def _adapted_get_extended_attention_mask(self, attention_mask, input_shape=None, *args, **kwargs): + device = kwargs.pop("device", None) + dtype = kwargs.pop("dtype", None) + if len(args) >= 1: + arg = args[0] + if isinstance(arg, torch.device) or (isinstance(arg, str) and str(arg).startswith(("cuda", "cpu", "mps"))): + device = arg + elif isinstance(arg, torch.dtype): + dtype = arg + if len(args) >= 2: + arg2 = args[1] + if isinstance(arg2, torch.dtype): + dtype = arg2 + elif isinstance(arg2, (torch.device, str)): + device = arg2 + + if dtype is None: + dtype = getattr(self, "dtype", torch.float32) + + if attention_mask.dim() == 3: + extended_attention_mask = attention_mask[:, None, :, :] + elif attention_mask.dim() == 2: + extended_attention_mask = attention_mask[:, None, None, :] + else: + raise ValueError(f"Wrong shape for attention_mask (shape {attention_mask.shape})") + + if device is not None: + extended_attention_mask = extended_attention_mask.to(device=device) + if dtype is not None: + extended_attention_mask = extended_attention_mask.to(dtype=dtype) + + extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min + return extended_attention_mask + + # 1. Patch BertModelWarper directly in GroundingDINO + try: + import groundingdino.models.GroundingDINO.bertwarper as bw + if not getattr(bw, "_rmbg_patched", False): + orig_init = bw.BertModelWarper.__init__ + def new_init(self, bert_model): + if not hasattr(bert_model, "get_head_mask"): + bert_model.get_head_mask = lambda *a, **k: _get_head_mask(bert_model, *a, **k) + bert_model.get_extended_attention_mask = lambda *a, **k: _adapted_get_extended_attention_mask(bert_model, *a, **k) + orig_init(self, bert_model) + if not hasattr(self, "get_head_mask") or self.get_head_mask is None: + self.get_head_mask = lambda *a, **k: _get_head_mask(self, *a, **k) + self.get_extended_attention_mask = lambda *a, **k: _adapted_get_extended_attention_mask(self, *a, **k) + bw.BertModelWarper.__init__ = new_init + bw._rmbg_patched = True + except Exception: + pass + + # 2. Patch BertModel in transformers + try: + import transformers + if hasattr(transformers, "BertModel"): + if not hasattr(transformers.BertModel, "get_head_mask"): + transformers.BertModel.get_head_mask = _get_head_mask + transformers.BertModel.get_extended_attention_mask = _adapted_get_extended_attention_mask + if hasattr(transformers, "PreTrainedModel"): + if not hasattr(transformers.PreTrainedModel, "get_head_mask"): + transformers.PreTrainedModel.get_head_mask = _get_head_mask + transformers.PreTrainedModel.get_extended_attention_mask = _adapted_get_extended_attention_mask + from transformers.models.bert.modeling_bert import BertModel + if not hasattr(BertModel, "get_head_mask"): + BertModel.get_head_mask = _get_head_mask + BertModel.get_extended_attention_mask = _adapted_get_extended_attention_mask + except Exception: + pass + + except Exception: + pass + + +def tensor2pil(image: torch.Tensor) -> Image.Image: + return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + + +def pil2tensor(image: Image.Image) -> torch.Tensor: + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + + +def pil2mask(image: Image.Image) -> torch.Tensor: + return torch.from_numpy(np.array(image.convert("L")).astype(np.float32) / 255.0).unsqueeze(0) + + +def mask2image(mask: torch.Tensor) -> Image.Image: + if len(mask.shape) == 2: + mask = mask.unsqueeze(0) + return tensor2pil(mask) + + +def image2mask(image: Image.Image) -> torch.Tensor: + return torch.from_numpy(np.array(image.convert("L")).astype(np.float32) / 255.0) + + +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 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_or_download_model_file(filename: str, url: str, dirname: str = "rmbg") -> str: + 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 resize_image(img: Image.Image, width: int, height: int) -> Image.Image: + return img.resize((width, height), resample=Image.LANCZOS) + + +def blend_overlay(img_1: Image.Image, img_2: Image.Image) -> Image.Image: + arr1 = np.array(img_1).astype(float) / 255.0 + arr2 = np.array(img_2).astype(float) / 255.0 + mask = arr2 < 0.5 + result = np.zeros_like(arr1) + result[mask] = 2 * arr1[mask] * arr2[mask] + result[~mask] = 1 - 2 * (1 - arr1[~mask]) * (1 - arr2[~mask]) + return Image.fromarray(np.clip(result * 255, 0, 255).astype(np.uint8)) + + +def fill_mask(width: int, height: int, mask: Image.Image, box=(0, 0), color=0) -> Image.Image: + bg = Image.new("L", (width, height), color) + bg.paste(mask, box, mask) + return bg + + +def empty_image(width: int, height: int, batch_size: int = 1) -> torch.Tensor: + return torch.zeros([batch_size, height, width, 3]) + + +def upscale_mask(mask: torch.Tensor, width: int, height: int) -> torch.Tensor: + if mask.ndim == 3: + mask = mask.unsqueeze(1) + mask = common_upscale(mask, width, height, "bicubic", "disabled") + mask = mask.squeeze(1) + return mask + + +def extract_alpha_mask(image: torch.Tensor) -> torch.Tensor: + alpha = image[..., 3] + if alpha.max() > 1.0: + alpha = alpha / 255.0 + if len(alpha.shape) == 4: + alpha = alpha[:, :, :, 0] + return alpha.unsqueeze(1) if alpha.ndim == 3 else alpha + + +def ensure_mask_shape(mask: torch.Tensor | None) -> torch.Tensor | None: + if mask is None: + return None + if mask.ndim == 2: + return mask.unsqueeze(0) + if mask.ndim == 4 and mask.shape[1] == 1: + return mask.squeeze(1) + return mask + + +COLOR_PRESETS = { + "black": "#000000", + "white": "#FFFFFF", + "red": "#FF0000", + "green": "#00FF00", + "blue": "#0000FF", + "yellow": "#FFFF00", + "cyan": "#00FFFF", + "magenta": "#FF00FF", + "gray": "#808080", + "silver": "#C0C0C0", + "maroon": "#800000", + "olive": "#808000", + "purple": "#800080", + "teal": "#008080", + "navy": "#000080", + "orange": "#FFA500", + "pink": "#FFC0CB", + "brown": "#A52A2A", + "violet": "#EE82EE", + "indigo": "#4B0082", + "light_gray": "#D3D3D3", + "dark_gray": "#A9A9A9", + "light_blue": "#ADD8E6", + "dark_blue": "#00008B", + "light_green": "#90EE90", + "dark_green": "#006400", +} + +ASPECT_RATIOS = { + # Square + "1:1 (Square)": (1, 1), + + # Landscape + "5:4 (Landscape)": (5, 4), + "4:3 (Standard)": (4, 3), + "3:2 (Photo)": (3, 2), + "16:10 (Landscape)": (16, 10), + "16:9 (Widescreen)": (16, 9), + "1.85:1 (Film)": (185, 100), + "2:1 (Univisium)": (2, 1), + "21:9 (Ultrawide)": (21, 9), + "2.35:1 (Anamorphic)": (235, 100), + "2.39:1 (Panavision)": (239, 100), + "32:9 (Super Ultrawide)": (32, 9), + "3:1 (Panorama)": (3, 1), + + # Portrait (mirrors of the landscape ratios) + "4:5 (Portrait)": (4, 5), + "3:4 (Portrait Standard)": (3, 4), + "2:3 (Portrait Photo)": (2, 3), + "10:16 (Portrait)": (10, 16), + "9:16 (Portrait Widescreen)": (9, 16), + "9:21 (Portrait Ultrawide)": (9, 21), + "1:3 (Portrait Panorama)": (1, 3), +} + +def color_format(color: str) -> str: + if not color: + return "" + + color = color.strip().upper() + if not color.startswith("#"): + color = f"#{color}" + + color = color[1:] + if len(color) == 3: + r, g, b = color[0], color[1], color[2] + return f"#{r}{r}{g}{g}{b}{b}" + if len(color) < 6: + raise ValueError(f"Invalid color format: {color}") + if len(color) > 6: + color = color[:6] + + return f"#{color}"