From 84312251db96178469c0192360ff3a05070d33d5 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Sun, 5 Oct 2025 01:06:04 -0700 Subject: [PATCH] Add files via upload --- AILab_BiRefNet.py | 2 +- AILab_BodySegment.py | 2 +- AILab_ClothSegment.py | 2 +- AILab_FaceSegment.py | 2 +- AILab_FashionSegment.py | 2 +- AILab_InpaintTools.py | 4 +- AILab_RMBG.py | 4 +- AILab_SAM2Segment.py | 2 +- AILab_SDMatte.py | 867 ++++++++++++++++++++-------------------- AILab_Segment.py | 2 +- AILab_SegmentV2.py | 2 +- 11 files changed, 445 insertions(+), 446 deletions(-) diff --git a/AILab_BiRefNet.py b/AILab_BiRefNet.py index 414541b..94e2a7e 100644 --- a/AILab_BiRefNet.py +++ b/AILab_BiRefNet.py @@ -389,7 +389,7 @@ class BiRefNetRMBG: "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": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), } } diff --git a/AILab_BodySegment.py b/AILab_BodySegment.py index 9b2ac84..8e0082c 100644 --- a/AILab_BodySegment.py +++ b/AILab_BodySegment.py @@ -83,7 +83,7 @@ class BodySegment: "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": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), }, } diff --git a/AILab_ClothSegment.py b/AILab_ClothSegment.py index 742431b..c22fe99 100644 --- a/AILab_ClothSegment.py +++ b/AILab_ClothSegment.py @@ -83,7 +83,7 @@ class ClothesSegment: "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": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), }, } diff --git a/AILab_FaceSegment.py b/AILab_FaceSegment.py index 3d53a7b..bebd1d2 100644 --- a/AILab_FaceSegment.py +++ b/AILab_FaceSegment.py @@ -83,7 +83,7 @@ class FaceSegment: "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": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), }, } diff --git a/AILab_FashionSegment.py b/AILab_FashionSegment.py index 6d3520f..0fcb829 100644 --- a/AILab_FashionSegment.py +++ b/AILab_FashionSegment.py @@ -176,7 +176,7 @@ class FashionSegmentClothing: "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": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), }, } diff --git a/AILab_InpaintTools.py b/AILab_InpaintTools.py index 996d303..f2764bb 100644 --- a/AILab_InpaintTools.py +++ b/AILab_InpaintTools.py @@ -1,4 +1,4 @@ -# ComfyUI-RMBG v2.6.0 +# ComfyUI-RMBG # # AILab Inpaint Tools # A collection of specialized nodes for inpainting tasks in ComfyUI. @@ -100,7 +100,7 @@ class AILab_ReferenceLatentMask: RETURN_TYPES = ("CONDITIONING", "LATENT", "MASK") RETURN_NAMES = ("CONDITIONING", "LATENT", "MASK") FUNCTION = "prepare_inpaint_conditioning" - CATEGORY = "🧪AILab/🧽RMBG/🎭Inpaint" + CATEGORY = "🧪AILab/🎭Inpaint" def add_latent_to_conditioning(self, conditioning, latent=None): if latent is not None: diff --git a/AILab_RMBG.py b/AILab_RMBG.py index af906bd..06c5a63 100644 --- a/AILab_RMBG.py +++ b/AILab_RMBG.py @@ -1,4 +1,4 @@ -# ComfyUI-RMBG V2.9.1 +# 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. @@ -547,7 +547,7 @@ class RMBG: "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": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), } } diff --git a/AILab_SAM2Segment.py b/AILab_SAM2Segment.py index 58844e6..f9198e9 100644 --- a/AILab_SAM2Segment.py +++ b/AILab_SAM2Segment.py @@ -184,7 +184,7 @@ class SAM2Segment: "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": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), } } diff --git a/AILab_SDMatte.py b/AILab_SDMatte.py index d0f7603..70e7389 100644 --- a/AILab_SDMatte.py +++ b/AILab_SDMatte.py @@ -1,434 +1,433 @@ -# 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 - -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 = os.path.dirname(__file__) -if str(current_dir) not in sys.path: - sys.path.insert(0, current_dir) - -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": ("COLOR", {"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 + +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 = os.path.dirname(__file__) +if str(current_dir) not in sys.path: + sys.path.insert(0, current_dir) + +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)", +} \ No newline at end of file diff --git a/AILab_Segment.py b/AILab_Segment.py index 2317908..189dad5 100644 --- a/AILab_Segment.py +++ b/AILab_Segment.py @@ -169,7 +169,7 @@ class Segment: "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": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), } } diff --git a/AILab_SegmentV2.py b/AILab_SegmentV2.py index 865f841..1d4ea58 100644 --- a/AILab_SegmentV2.py +++ b/AILab_SegmentV2.py @@ -149,7 +149,7 @@ class SegmentV2: "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": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": tooltips["background_color"]}), } }