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

Add files via upload

Browse Source
This commit is contained in:
AI Lab
2026-09-30 12:40:42 -07:00
committed by GitHub
parent 93ab85f059
commit d5ee9818ac
16 changed files with 2841 additions and 2639 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/AILab_BiRefNet.py
+8 -2
View File
@@ -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 = {
py/AILab_BodySegment.py
+7 -28
View File
@@ -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
py/AILab_ClothSegment.py
+6 -28
View File
@@ -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
py/AILab_FaceSegment.py
+274 -294
View File
@@ -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)"
}
py/AILab_FashionSegment.py
+6 -23
View File
@@ -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()
py/AILab_Florence2.py
+352 -352
View File
@@ -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": "<REFERRING_EXPRESSION_SEGMENTATION>",
"mode": "polygon",
"allows_text": True,
},
"Phrase Grounding (text boxes)": {
"token": "<CAPTION_TO_PHRASE_GROUNDING>",
"mode": "bbox",
"allows_text": True,
},
"Region Proposals (boxes only)": {
"token": "<REGION_PROPOSAL>",
"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": "<REFERRING_EXPRESSION_SEGMENTATION>",
"mode": "polygon",
"allows_text": True,
},
"Phrase Grounding (text boxes)": {
"token": "<CAPTION_TO_PHRASE_GROUNDING>",
"mode": "bbox",
"allows_text": True,
},
"Region Proposals (boxes only)": {
"token": "<REGION_PROPOSAL>",
"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)",
}
py/AILab_ImageCompareView.py
+201
View File
@@ -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) 🔀🖼️",
}
py/AILab_LamaRemover.py
+17 -21
View File
@@ -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,
py/AILab_RMBG.py
+705 -706
View File
File diff suppressed because it is too large Load Diff
py/AILab_SAM2Segment.py
+41 -67
View File
@@ -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,
py/AILab_SAM3Segment.py
+100 -121
View File
@@ -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)"}
py/AILab_SDMatte.py
+404 -435
View File
@@ -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)",
}
py/AILab_Segment.py
+56 -96
View File
@@ -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(
py/AILab_SegmentV2.py
+26 -58
View File
@@ -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,
py/AILab_YoloV8.py
+289 -279
View File
@@ -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)",
}
py/AILab_utils.py
+349 -129
View File
@@ -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}"