diff --git a/py/AILab_BiRefNet.py b/py/AILab_BiRefNet.py index 152800a..2ae599e 100644 --- a/py/AILab_BiRefNet.py +++ b/py/AILab_BiRefNet.py @@ -321,6 +321,10 @@ class BiRefNetModel: with open(model_path, 'r', encoding='utf-8') as f: model_content = f.read() model_content = model_content.replace("from .BiRefNet_config", "from BiRefNet_config") + model_content = model_content.replace( + "attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)", + "attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.to(attn.dtype).unsqueeze(1).unsqueeze(0)" + ) with open(model_path, 'w', encoding='utf-8') as f: f.write(model_content) @@ -367,7 +371,7 @@ class BiRefNetModel: with torch.no_grad(): preds = self.model(input_tensor) - pred = preds[-1].sigmoid().cpu() + pred = preds[-1].float().sigmoid().cpu() pred = pred[0].squeeze() pred_pil = transforms.ToPILImage()(pred) diff --git a/py/AILab_Canvas.py b/py/AILab_Canvas.py new file mode 100644 index 0000000..78c8c0a --- /dev/null +++ b/py/AILab_Canvas.py @@ -0,0 +1,361 @@ +import json +import os +import torch +import torchvision.transforms.functional as TF +from typing import Dict, Any, Tuple, List +from PIL import Image +import numpy as np +import folder_paths +from AILab_utils import clean_vram, get_device + +def hex_to_rgba(hex_str: str) -> Tuple[float, float, float, float]: + if not isinstance(hex_str, str): + return (0.0, 0.0, 0.0, 1.0) + cleaned = hex_str.strip().lstrip("#") + if len(cleaned) == 3: + cleaned = "".join([c * 2 for c in cleaned]) + "FF" + elif len(cleaned) == 6: + cleaned = cleaned + "FF" + elif len(cleaned) != 8: + return (0.0, 0.0, 0.0, 1.0) + try: + r = int(cleaned[0:2], 16) / 255.0 + g = int(cleaned[2:4], 16) / 255.0 + b = int(cleaned[4:6], 16) / 255.0 + a = int(cleaned[6:8], 16) / 255.0 + return (r, g, b, a) + except Exception: + return (0.0, 0.0, 0.0, 1.0) + +class AILab_Canvas: + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + return { + "required": { + "canvas_width": ( + "INT", + { + "default": 640, + "min": 64, + "max": 8192, + "step": 8, + "tooltip": "Canvas width in pixels when not fitting background image size.", + }, + ), + "canvas_height": ( + "INT", + { + "default": 360, + "min": 64, + "max": 8192, + "step": 8, + "tooltip": "Canvas height in pixels when not fitting background image size.", + }, + ), + "background": ( + ["Alpha", "Color"], + { + "default": "Alpha", + "tooltip": "Choose background type: Alpha (transparent) or Color (custom background color).", + }, + ), + "background_color": ( + "COLORCODE", + { + "default": "#222222", + "tooltip": "Pick background color (use color picker).", + }, + ), + "fit_background_size": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When True, canvas dimensions match the background image. When False, canvas dimensions use canvas_width and canvas_height.", + }, + ), + "canvas_state": ( + "STRING", + { + "default": "{}", + "multiline": True, + "tooltip": "Serialized JSON layout containing layer coordinates, scales, rotations, and order.", + }, + ), + }, + "optional": { + "Background": ( + "IMAGE", + {"tooltip": "Optional background image input."}, + ), + "layer_1": ( + "IMAGE", + {"tooltip": "Layer 1 image input."}, + ), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("IMAGE", "MASK") + FUNCTION = "composite" + OUTPUT_NODE = True + CATEGORY = "🧪AILab/🖼️IMAGE" + + def _save_images(self, slot_tensors: List[Tuple[str, torch.Tensor]]) -> List[Dict[str, Any]]: + results = [] + output_dir = folder_paths.get_temp_directory() + os.makedirs(output_dir, exist_ok=True) + for slot_name, tensor in slot_tensors: + if tensor is None or not isinstance(tensor, torch.Tensor): + continue + t = tensor[0].detach().cpu() + h, w = t.shape[0], t.shape[1] + full_out_dir, base_filename, counter, subfolder, _ = folder_paths.get_save_image_path( + f"ailab_canvas_{slot_name.lower()}", output_dir, w, h + ) + arr = np.clip(255.0 * t.numpy(), 0, 255).astype(np.uint8) + if arr.shape[-1] == 1: + arr = np.squeeze(arr, axis=-1) + img = Image.fromarray(arr, mode="L") + elif arr.shape[-1] == 3: + img = Image.fromarray(arr, mode="RGB") + elif arr.shape[-1] == 4: + img = Image.fromarray(arr, mode="RGBA") + else: + img = Image.fromarray(arr[..., :3], mode="RGB") + + file = f"{base_filename}_{counter:05}_{slot_name}.png" + img.save(os.path.join(full_out_dir, file), compress_level=1) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": "temp", + "slot": slot_name, + "width": w, + "height": h, + }) + return results + + def composite( + self, + canvas_width: int, + canvas_height: int, + background: str = "Alpha", + background_color: str = "#222222", + fit_background_size: bool = True, + canvas_state: str = "{}", + **kwargs: Any, + ): + try: + bg_tensor = kwargs.get("Background", kwargs.get("background_image", None)) + if not canvas_state or canvas_state == "{}": + canvas_state = kwargs.get("canvas_state", "{}") + + layer_dict = {} + for i in range(1, 6): + key = f"layer_{i}" + if key in kwargs and isinstance(kwargs[key], torch.Tensor): + layer_dict[key] = kwargs[key] + + for k, v in kwargs.items(): + k_clean = k.lower().replace(" ", "_") + if k_clean.startswith("layer_") and isinstance(v, torch.Tensor): + layer_dict[k_clean] = v + + device = bg_tensor.device if bg_tensor is not None else ( + next(iter(layer_dict.values())).device if layer_dict else get_device() + ) + + if bg_tensor is not None and fit_background_size: + target_h = int(bg_tensor.shape[1]) + target_w = int(bg_tensor.shape[2]) + else: + target_h = max(64, int(canvas_height)) + target_w = max(64, int(canvas_width)) + target_h = (target_h // 8) * 8 + target_w = (target_w // 8) * 8 + + batch_size = 1 + if bg_tensor is not None: + batch_size = max(batch_size, bg_tensor.shape[0]) + for l_tensor in layer_dict.values(): + batch_size = max(batch_size, l_tensor.shape[0]) + + canvas = torch.zeros((batch_size, 4, target_h, target_w), dtype=torch.float32, device=device) + fg_mask = torch.zeros((batch_size, 1, target_h, target_w), dtype=torch.float32, device=device) + + if bg_tensor is not None: + bg = bg_tensor.to(device=device, dtype=torch.float32) + if bg.shape[0] < batch_size: + bg = bg.repeat(batch_size, 1, 1, 1) + + if bg.shape[1] != target_h or bg.shape[2] != target_w: + bg_perm = bg.permute(0, 3, 1, 2) + bg_resized = TF.resize(bg_perm, [target_h, target_w], interpolation=TF.InterpolationMode.BILINEAR) + bg = bg_resized.permute(0, 2, 3, 1) + + if bg.shape[-1] >= 4: + canvas[:, :3, :, :] = bg[..., :3].permute(0, 3, 1, 2) + canvas[:, 3:4, :, :] = bg[..., 3:4].permute(0, 3, 1, 2) + else: + canvas[:, :3, :, :] = bg[..., :3].permute(0, 3, 1, 2) + canvas[:, 3:4, :, :] = 1.0 + else: + if background == "Alpha": + canvas[:, :, :, :] = 0.0 + else: + r, g, b, a = hex_to_rgba(background_color) + canvas[:, 0, :, :] = r + canvas[:, 1, :, :] = g + canvas[:, 2, :, :] = b + canvas[:, 3, :, :] = a + + raw_layers = [] + for k in sorted(layer_dict.keys(), key=lambda x: int(x.split('_')[1]) if '_' in x and x.split('_')[1].isdigit() else 99): + v = layer_dict[k] + if v is not None and isinstance(v, torch.Tensor): + raw_layers.append((k, v)) + + parsed_state = {} + if canvas_state and isinstance(canvas_state, str) and canvas_state.strip(): + try: + parsed_state = json.loads(canvas_state) + except Exception: + parsed_state = {} + + state_layer_map = {} + for item in parsed_state.get("layers", []): + if isinstance(item, dict) and "id" in item: + state_layer_map[str(item["id"]).lower()] = item + + state_width = float(parsed_state.get("canvas_width", target_w) or target_w) + state_height = float(parsed_state.get("canvas_height", target_h) or target_h) + scale_ratio_x = target_w / max(1.0, state_width) + scale_ratio_y = target_h / max(1.0, state_height) + + ordered_layers = [] + for layer_name, tensor in raw_layers: + layer_cfg = state_layer_map.get(layer_name.lower(), {}) + z_index = int(layer_cfg.get("zIndex", layer_cfg.get("z_index", 0))) + ordered_layers.append((z_index, layer_name, tensor, layer_cfg)) + + ordered_layers.sort(key=lambda x: x[0]) + + for _, _, fg_tensor, cfg in ordered_layers: + fg = fg_tensor.to(device=device, dtype=torch.float32) + if fg.shape[0] < batch_size: + fg = fg.repeat(batch_size, 1, 1, 1) + + orig_h, orig_w = int(fg.shape[1]), int(fg.shape[2]) + + if fg.shape[-1] >= 4: + fg_rgb = fg[..., :3].permute(0, 3, 1, 2) + fg_a = fg[..., 3:4].permute(0, 3, 1, 2) + else: + fg_rgb = fg[..., :3].permute(0, 3, 1, 2) + fg_a = torch.ones((batch_size, 1, orig_h, orig_w), dtype=torch.float32, device=device) + + opacity = float(cfg.get("opacity", 1.0)) + opacity = max(0.0, min(1.0, opacity)) + fg_a = fg_a * opacity + + flip_x = bool(cfg.get("flipX", cfg.get("flip_x", False))) + flip_y = bool(cfg.get("flipY", cfg.get("flip_y", False))) + if flip_x: + fg_rgb = torch.flip(fg_rgb, dims=[3]) + fg_a = torch.flip(fg_a, dims=[3]) + if flip_y: + fg_rgb = torch.flip(fg_rgb, dims=[2]) + fg_a = torch.flip(fg_a, dims=[2]) + + max_w = target_w * 0.85 + max_h = target_h * 0.85 + if orig_w > max_w or orig_h > max_h: + default_scale = min(max_w / max(1, orig_w), max_h / max(1, orig_h)) + else: + default_scale = 1.0 + + scale = float(cfg.get("scaleX", cfg.get("scale_x", cfg.get("scaleY", default_scale)))) + scale_x = scale + scale_y = scale + + new_w = max(1, int(round(orig_w * scale_x * scale_ratio_x))) + new_h = max(1, int(round(orig_h * scale_y * scale_ratio_y))) + + fg_rgba = torch.cat([fg_rgb, fg_a], dim=1) + if new_w != orig_w or new_h != orig_h: + fg_rgba = TF.resize(fg_rgba, [new_h, new_w], interpolation=TF.InterpolationMode.BILINEAR) + + angle = float(cfg.get("angle", 0.0)) + if abs(angle) > 1e-3: + fg_rgba = TF.rotate(fg_rgba, angle=-angle, expand=True, interpolation=TF.InterpolationMode.BILINEAR) + + rot_h, rot_w = fg_rgba.shape[2], fg_rgba.shape[3] + + if "x" in cfg and "y" in cfg: + raw_cx = float(cfg["x"]) * scale_ratio_x + raw_cy = float(cfg["y"]) * scale_ratio_y + else: + raw_cx = target_w / 2.0 + raw_cy = target_h / 2.0 + + x0 = int(round(raw_cx - rot_w / 2.0)) + y0 = int(round(raw_cy - rot_h / 2.0)) + x1 = x0 + rot_w + y1 = y0 + rot_h + + dst_x0, dst_y0 = max(0, x0), max(0, y0) + dst_x1, dst_y1 = min(target_w, x1), min(target_h, y1) + src_x0 = max(0, -x0) + src_y0 = max(0, -y0) + src_x1 = src_x0 + (dst_x1 - dst_x0) + src_y1 = src_y0 + (dst_y1 - dst_y0) + + if dst_x1 > dst_x0 and dst_y1 > dst_y0: + fg_patch = fg_rgba[:, :, src_y0:src_y1, src_x0:src_x1] + bg_patch = canvas[:, :, dst_y0:dst_y1, dst_x0:dst_x1] + + p_fg_rgb = fg_patch[:, :3] + p_fg_a = fg_patch[:, 3:4] + p_bg_rgb = bg_patch[:, :3] + p_bg_a = bg_patch[:, 3:4] + + out_a = p_fg_a + p_bg_a * (1.0 - p_fg_a) + out_rgb = torch.where( + out_a > 1e-6, + (p_fg_rgb * p_fg_a + p_bg_rgb * p_bg_a * (1.0 - p_fg_a)) / out_a, + torch.zeros_like(p_fg_rgb), + ) + + canvas[:, :3, dst_y0:dst_y1, dst_x0:dst_x1] = out_rgb + canvas[:, 3:4, dst_y0:dst_y1, dst_x0:dst_x1] = out_a + + curr_fg_a = fg_mask[:, :, dst_y0:dst_y1, dst_x0:dst_x1] + fg_mask[:, :, dst_y0:dst_y1, dst_x0:dst_x1] = p_fg_a + curr_fg_a * (1.0 - p_fg_a) + + if background == "Alpha" and (bg_tensor is None or bg_tensor.shape[-1] >= 4): + out_image = canvas.permute(0, 2, 3, 1).clamp(0.0, 1.0) + else: + out_image = canvas[:, :3, :, :].permute(0, 2, 3, 1).clamp(0.0, 1.0) + out_mask = fg_mask.squeeze(1).clamp(0.0, 1.0) + + slot_tensors = [("Background", bg_tensor)] + for i in range(1, 6): + slot_tensors.append((f"layer_{i}", layer_dict.get(f"layer_{i}"))) + + ui_images = self._save_images(slot_tensors) + + return { + "ui": {"images": ui_images}, + "result": (out_image.cpu(), out_mask.cpu()), + } + + finally: + clean_vram() + +NODE_CLASS_MAPPINGS = { + "AILab_Canvas": AILab_Canvas, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AILab_Canvas": "Canvas (RMBG)", +} \ No newline at end of file diff --git a/py/AILab_FeyNoBg.py b/py/AILab_FeyNoBg.py new file mode 100644 index 0000000..c7b94d0 --- /dev/null +++ b/py/AILab_FeyNoBg.py @@ -0,0 +1,225 @@ +import os +import torch +import numpy as np +from PIL import Image, ImageFilter +import cv2 +import json +import safetensors.torch +import folder_paths +import comfy.model_management as mm +from huggingface_hub import hf_hub_download + +from AILab_utils import get_device, clean_vram, unload_model as release_model, tensor2pil, pil2tensor + +device = get_device() +folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) + +MODEL_REPO = "1038lab/FeyNobg" +MODEL_CACHE_NAME = "FeyNoBg" +REQUIRED_FILES = ["config.json", "FeyNobg.safetensors", "preprocessor_config.json"] + + +def _refine_foreground_edges(image_bchw, masks_b1hw): + b, c, h, w = image_bchw.shape + if b != masks_b1hw.shape[0]: + raise ValueError("images and masks must have the same batch size") + + image_np = image_bchw.cpu().numpy() + mask_np = masks_b1hw.cpu().numpy() + + refined_fg = [] + for i in range(b): + mask = mask_np[i, 0] + thresh = 0.45 + mask_binary = (mask > thresh).astype(np.float32) + + edge_blur = cv2.GaussianBlur(mask_binary, (3, 3), 0) + transition_mask = np.logical_and(mask > 0.05, mask < 0.95) + + alpha = 0.85 + mask_refined = np.where(transition_mask, + alpha * mask + (1-alpha) * edge_blur, + mask_binary) + + edge_region = np.logical_and(mask > 0.2, mask < 0.8) + mask_refined = np.where(edge_region, + mask_refined * 0.98, + mask_refined) + + result = [] + for ch in range(image_np.shape[1]): + channel = image_np[i, ch] + refined = channel * mask_refined + result.append(refined) + + refined_fg.append(np.stack(result)) + + return torch.from_numpy(np.stack(refined_fg)) + + +class AILab_FeyNoBg: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", {"tooltip": "Input image batch to remove background from."}), + }, + "optional": { + "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": "Gaussian blur radius for softening mask edges."}), + "mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": "Expand (positive) or shrink (negative) the mask boundary in pixels."}), + "invert_output": ("BOOLEAN", {"default": False, "tooltip": "Invert the output mask."}), + "refine_foreground": ("BOOLEAN", {"default": True, "tooltip": "Sub-pixel foreground boundary refinement for fine hairs and semi-transparent edges."}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model weights from VRAM immediately after execution."}), + "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": "Output format: Alpha for transparent PNG, or Color to composite against background."}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": "Hex background color used when background mode is set to Color."}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") + FUNCTION = "process" + CATEGORY = "🧪AILab/🧽RMBG" + + def __init__(self): + self.model = None + self.processor = None + self.target_device = device + + def _ensure_model(self): + if self.model is not None and self.processor is not None: + return self.model, self.processor + + try: + from nobg.birefnet.modeling_birefnet import BiRefNet, BiRefNetConfig + from nobg import BiRefNetImageProcessor + except ImportError as e: + raise RuntimeError("The 'nobg' library is required for FeyNoBg. Please run: pip install nobg") from e + + local_cache_dir = os.path.join(folder_paths.models_dir, "RMBG", MODEL_CACHE_NAME) + os.makedirs(local_cache_dir, exist_ok=True) + + for filename in REQUIRED_FILES: + filepath = os.path.join(local_cache_dir, filename) + if not os.path.exists(filepath): + hf_hub_download( + repo_id=MODEL_REPO, + filename=filename, + local_dir=local_cache_dir, + local_dir_use_symlinks=False + ) + + config_path = os.path.join(local_cache_dir, "config.json") + weights_path = os.path.join(local_cache_dir, "FeyNobg.safetensors") + + with open(config_path, "r", encoding="utf-8") as f: + config_data = json.load(f) + + config = BiRefNetConfig(**config_data) + self.model = BiRefNet(config) + state_dict = safetensors.torch.load_file(weights_path) + self.model.load_state_dict(state_dict) + self.model.eval() + + self.processor = BiRefNetImageProcessor.from_pretrained(local_cache_dir) + + if torch.cuda.is_available() and str(self.target_device).startswith("cuda"): + self.model = self.model.half().to(self.target_device) + else: + self.model = self.model.to(self.target_device) + + return self.model, self.processor + + def process(self, image, mask_blur=0, mask_offset=0, invert_output=False, + refine_foreground=True, unload_model=False, + background="Alpha", background_color="#222222", **kwargs): + do_unload = kwargs.get("unload_model", unload_model) + refine = kwargs.get("refine_foreground", refine_foreground) + + model, processor = self._ensure_model() + + processed_images = [] + processed_masks = [] + + try: + for single_img in image: + pil_img = tensor2pil(single_img).convert("RGB") + orig_w, orig_h = pil_img.size + + inputs = processor(pil_img, return_tensors="pt") + pixel_values = inputs["pixel_values"].to(self.target_device) + if next(model.parameters()).dtype == torch.float16: + pixel_values = pixel_values.half() + + with torch.no_grad(): + outputs = model(pixel_values=pixel_values) + alpha = processor.post_process_alpha_matting(outputs, target_sizes=[(orig_h, orig_w)])[0] + + if isinstance(alpha, torch.Tensor): + alpha_np = alpha.detach().cpu().float().numpy() + elif isinstance(alpha, Image.Image): + alpha_np = np.array(alpha).astype(np.float32) / 255.0 + else: + alpha_np = np.array(alpha).astype(np.float32) + if alpha_np.max() > 1.0: + alpha_np /= 255.0 + + mask_img = Image.fromarray((np.clip(alpha_np, 0.0, 1.0) * 255.0).astype(np.uint8), mode="L") + + if mask_offset != 0: + offset_k = abs(mask_offset) * 2 + 1 + filter_type = ImageFilter.MaxFilter if mask_offset > 0 else ImageFilter.MinFilter + mask_img = mask_img.filter(filter_type(offset_k)) + + if mask_blur > 0: + mask_img = mask_img.filter(ImageFilter.GaussianBlur(radius=mask_blur)) + + if invert_output: + mask_np_proc = 255.0 - np.array(mask_img, dtype=np.float32) + mask_img = Image.fromarray(np.clip(mask_np_proc, 0, 255).astype(np.uint8), mode="L") + + mask_tensor = pil2tensor(mask_img) + processed_masks.append(mask_tensor) + + if background == "Alpha": + if refine: + refined_single = _refine_foreground_edges(single_img.unsqueeze(0).permute(0, 3, 1, 2), mask_tensor.unsqueeze(1)) + refined_tensor = refined_single.squeeze(0).permute(1, 2, 0) + refined_pil = tensor2pil(refined_tensor) + else: + refined_pil = pil_img + + cutout = Image.new("RGBA", (orig_w, orig_h), (0, 0, 0, 0)) + cutout.paste(refined_pil, (0, 0), mask_img) + processed_images.append(pil2tensor(cutout)) + else: + bg_hex = background_color.lstrip("#") + if len(bg_hex) == 6: + bg_rgb = tuple(int(bg_hex[i:i+2], 16) for i in (0, 2, 4)) + else: + bg_rgb = (34, 34, 34) + + bg_canvas = Image.new("RGB", (orig_w, orig_h), bg_rgb) + bg_canvas.paste(pil_img, (0, 0), mask_img) + processed_images.append(pil2tensor(bg_canvas)) + + out_image = torch.cat(processed_images, dim=0) + out_mask = torch.cat(processed_masks, dim=0) + out_mask_image = out_mask.unsqueeze(-1).expand(-1, -1, -1, 3) + + return (out_image, out_mask, out_mask_image) + + finally: + if do_unload: + release_model(self, "model") + self.processor = None + else: + clean_vram() + + +NODE_CLASS_MAPPINGS = { + "AILab_FeyNoBg": AILab_FeyNoBg +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AILab_FeyNoBg": "FeyNoBg (RMBG)" +} diff --git a/py/AILab_FinegrainSegment.py b/py/AILab_FinegrainSegment.py new file mode 100644 index 0000000..e7fb967 --- /dev/null +++ b/py/AILab_FinegrainSegment.py @@ -0,0 +1,225 @@ +import os +import json +from pathlib import Path +from typing import Any, List, Tuple, Optional +import torch +import numpy as np +from PIL import Image, ImageFilter +from huggingface_hub import hf_hub_download +import folder_paths + +from AILab_utils import get_device, clean_vram, unload_model as release_model, tensor2pil, pil2tensor + +device = get_device() +folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) + +MODEL_SUBDIR = os.path.join(folder_paths.models_dir, "RMBG", "finegrain-box-segmenter") +REPO_ID = "1038lab/finegrain-box-segmenter" +WEIGHTS_FILENAME = "finegrain-box-segmenter.safetensors" + + +def parse_bounding_boxes(bboxes_input: Any, box_str: str, img_w: int, img_h: int) -> List[Tuple[int, int, int, int]]: + boxes = [] + + if bboxes_input is not None: + raw_list = [] + if isinstance(bboxes_input, str): + cleaned = bboxes_input.strip() + if cleaned: + try: + parsed = json.loads(cleaned) + if isinstance(parsed, list): + raw_list = parsed + elif isinstance(parsed, dict): + raw_list = parsed.get("bboxes", parsed.get("boxes", [])) + except Exception: + pass + elif isinstance(bboxes_input, dict): + raw_list = bboxes_input.get("bboxes", bboxes_input.get("boxes", [])) + elif isinstance(bboxes_input, (list, tuple)): + if len(bboxes_input) == 4 and all(isinstance(x, (int, float)) for x in bboxes_input): + raw_list = [bboxes_input] + else: + raw_list = bboxes_input + + for item in raw_list: + if isinstance(item, (list, tuple)) and len(item) >= 4: + x1, y1, x2, y2 = float(item[0]), float(item[1]), float(item[2]), float(item[3]) + if max(x1, y1, x2, y2) <= 1.0: + x1, y1, x2, y2 = x1 * img_w, y1 * img_h, x2 * img_w, y2 * img_h + x1_c = max(0, min(img_w, int(round(min(x1, x2))))) + y1_c = max(0, min(img_h, int(round(min(y1, y2))))) + x2_c = max(0, min(img_w, int(round(max(x1, x2))))) + y2_c = max(0, min(img_h, int(round(max(y1, y2))))) + if x2_c > x1_c and y2_c > y1_c: + boxes.append((x1_c, y1_c, x2_c, y2_c)) + + if not boxes and box_str: + cleaned = box_str.strip().strip("[]()").replace(" ", "") + parts = [p for p in cleaned.split(",") if p] + if len(parts) >= 4: + try: + x1, y1, x2, y2 = float(parts[0]), float(parts[1]), float(parts[2]), float(parts[3]) + if max(x1, y1, x2, y2) <= 1.0: + x1, y1, x2, y2 = x1 * img_w, y1 * img_h, x2 * img_w, y2 * img_h + x1_c = max(0, min(img_w, int(round(min(x1, x2))))) + y1_c = max(0, min(img_h, int(round(min(y1, y2))))) + x2_c = max(0, min(img_w, int(round(max(x1, x2))))) + y2_c = max(0, min(img_h, int(round(max(y1, y2))))) + if x2_c > x1_c and y2_c > y1_c: + boxes.append((x1_c, y1_c, x2_c, y2_c)) + except ValueError: + pass + + return boxes + + +class AILab_FinegrainBoxSegment: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", {"tooltip": "Input image batch to segment with high definition 1024px quality."}), + }, + "optional": { + "bboxes": ("BBOX", {"tooltip": "Bounding box(es) from Florence-2 or YOLO detection nodes."}), + "box_coordinates": ("STRING", {"default": "", "tooltip": "Manual bounding box as 'x1, y1, x2, y2' in pixel or normalized [0..1] coordinates."}), + "margin": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 0.5, "step": 0.01, "tooltip": "Safety padding margin around bounding boxes."}), + "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": "Blur mask edges (0 to disable)."}), + "mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": "Expand (positive) or shrink (negative) mask edges in pixels."}), + "invert_output": ("BOOLEAN", {"default": False, "tooltip": "Invert the output mask."}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model from VRAM immediately after execution."}), + "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": "Background type: Alpha (transparent) or Color (custom solid background)."}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": "Hex color code for background when Color mode is selected."}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") + FUNCTION = "segment" + CATEGORY = "🧪AILab/🧽RMBG" + + def __init__(self): + self.segmenter = None + self.current_weights_path = None + self.target_device = device + + def _ensure_model(self, margin=0.05): + if self.segmenter is not None: + self.segmenter.margin = margin + return self.segmenter + + try: + from refiners.solutions import BoxSegmenter + except ImportError as e: + raise RuntimeError("The 'refiners' library is required for Finegrain Box Segmenter. Please run: pip install refiners") from e + + os.makedirs(MODEL_SUBDIR, exist_ok=True) + local_weights = os.path.join(MODEL_SUBDIR, WEIGHTS_FILENAME) + + if not os.path.exists(local_weights): + print(f"[Finegrain] Downloading {WEIGHTS_FILENAME} from {REPO_ID} to {MODEL_SUBDIR}...") + local_weights = hf_hub_download( + repo_id=REPO_ID, + filename=WEIGHTS_FILENAME, + local_dir=MODEL_SUBDIR, + local_dir_use_symlinks=False + ) + + self.segmenter = BoxSegmenter( + margin=margin, + weights=local_weights, + device=str(self.target_device) + ) + self.current_weights_path = local_weights + return self.segmenter + + def segment(self, image, bboxes=None, box_coordinates="", margin=0.05, + mask_blur=0, mask_offset=0, invert_output=False, unload_model=False, + background="Alpha", background_color="#222222", **kwargs): + do_unload = kwargs.get("unload_model", unload_model) + segmenter = self._ensure_model(margin=margin) + + processed_images = [] + processed_masks = [] + + try: + for single_img in image: + pil_img = tensor2pil(single_img).convert("RGB") + orig_w, orig_h = pil_img.size + + parsed_boxes = parse_bounding_boxes(bboxes, box_coordinates, orig_w, orig_h) + + if not parsed_boxes: + # Segment whole image / default subject + mask_pil = segmenter(pil_img, box_prompt=None) + combined_mask_np = np.array(mask_pil, dtype=np.float32) + else: + combined_mask_np = np.zeros((orig_h, orig_w), dtype=np.float32) + for box in parsed_boxes: + mask_box = segmenter(pil_img, box_prompt=box) + mask_np = np.array(mask_box, dtype=np.float32) + combined_mask_np = np.maximum(combined_mask_np, mask_np) + + if combined_mask_np.max() > 1.0: + combined_mask_np = combined_mask_np / 255.0 + + mask_img = Image.fromarray((np.clip(combined_mask_np, 0.0, 1.0) * 255.0).astype(np.uint8), mode="L") + + if mask_offset != 0: + offset_k = abs(mask_offset) * 2 + 1 + filter_type = ImageFilter.MaxFilter if mask_offset > 0 else ImageFilter.MinFilter + mask_img = mask_img.filter(filter_type(offset_k)) + + if mask_blur > 0: + mask_img = mask_img.filter(ImageFilter.GaussianBlur(radius=mask_blur)) + + if invert_output: + mask_inv = 255.0 - np.array(mask_img, dtype=np.float32) + mask_img = Image.fromarray(np.clip(mask_inv, 0, 255).astype(np.uint8), mode="L") + + mask_tensor = pil2tensor(mask_img) + processed_masks.append(mask_tensor) + + if background == "Alpha": + cutout = Image.new("RGBA", (orig_w, orig_h), (0, 0, 0, 0)) + cutout.paste(pil_img, (0, 0), mask_img) + processed_images.append(pil2tensor(cutout)) + else: + bg_hex = background_color.lstrip("#") + if len(bg_hex) == 6: + bg_rgb = tuple(int(bg_hex[i:i+2], 16) for i in (0, 2, 4)) + else: + bg_rgb = (34, 34, 34) + + bg_canvas = Image.new("RGB", (orig_w, orig_h), bg_rgb) + bg_canvas.paste(pil_img, (0, 0), mask_img) + processed_images.append(pil2tensor(bg_canvas)) + + out_image = torch.cat(processed_images, dim=0) + out_mask = torch.cat(processed_masks, dim=0) + out_mask_image = out_mask.unsqueeze(-1).expand(-1, -1, -1, 3) + + return (out_image, out_mask, out_mask_image) + + finally: + if do_unload: + if self.segmenter is not None: + if hasattr(self.segmenter, "model") and hasattr(self.segmenter.model, "to"): + try: + self.segmenter.model.to("cpu") + except Exception: + pass + self.segmenter = None + clean_vram() + else: + clean_vram() + + +NODE_CLASS_MAPPINGS = { + "AILab_FinegrainBoxSegment": AILab_FinegrainBoxSegment +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AILab_FinegrainBoxSegment": "Finegrain Box Segmenter (RMBG)" +} diff --git a/py/AILab_Florence2.py b/py/AILab_Florence2.py index 5160245..9ccce22 100644 --- a/py/AILab_Florence2.py +++ b/py/AILab_Florence2.py @@ -1,55 +1,49 @@ +import hashlib 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 +from PIL import Image, ImageColor, ImageDraw, ImageFont import comfy.model_management as mm -from comfy.utils import ProgressBar +from comfy.utils import ProgressBar, load_torch_file +import comfy.ops import folder_paths -import transformers -from transformers import AutoModelForCausalLM, AutoProcessor -from transformers.dynamic_module_utils import get_imports + +from models.florence2.config import Florence2Config +from models.florence2.model import Florence2 +from models.florence2.processing import Processor 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)": { +MODE_CONFIGS = { + "Polygon Mask (prompt)": { "token": "", - "mode": "polygon", + "type": "polygon", "allows_text": True, }, - "Phrase Grounding (text boxes)": { + "Phrase Grounding (boxes)": { "token": "", - "mode": "bbox", + "type": "bbox", "allows_text": True, }, "Region Proposals (boxes only)": { "token": "", - "mode": "bbox", + "type": "bbox", "allows_text": False, }, } -TASK_CHOICES = tuple(TASK_CONFIGS.keys()) +MODE_CHOICES = tuple(MODE_CONFIGS.keys()) +TASK_CONFIGS = MODE_CONFIGS +TASK_CHOICES = MODE_CHOICES -GENERATION_CONFIG = { - "max_new_tokens": 512, - "num_beams": 3, - "do_sample": True, -} MODEL_CHOICES = ( "microsoft/Florence-2-base", @@ -63,22 +57,13 @@ MODEL_CHOICES = ( 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") + RETURN_TYPES = ("IMAGE", "MASK", "JSON", "BBOX") + RETURN_NAMES = ("IMAGE", "MASK", "DATA", "BBOXES") FUNCTION = "analyze" MODEL_CACHE: Dict[Tuple[str, str, str], Dict[str, Any]] = {} @@ -87,17 +72,26 @@ class AILab_Florence2: return { "required": { "image": ("IMAGE",), - "model_name": ( + "prompt": ( + "STRING", + { + "default": "", + "multiline": True, + "placeholder": "Prompt: e.g. face, a person wearing red coat", + "tooltip": "Used for polygon masks or phrase grounding; ignored for region proposals.", + }, + ), + "model": ( MODEL_CHOICES, { "default": MODEL_CHOICES[0], "tooltip": "Base = stable, +ft = fine-tuned captions, CogFlorence = sharper phrase alignment.", }, ), - "task": ( - TASK_CHOICES, + "mode": ( + MODE_CHOICES, { - "default": TASK_CHOICES[0], + "default": MODE_CHOICES[0], "tooltip": "Polygon masks use prompts; phrase grounding/region proposals return boxes.", }, ), @@ -122,18 +116,13 @@ class AILab_Florence2: "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."}, + "max_new_tokens": ( + "INT", + {"default": 1024, "min": 256, "max": 4096, "step": 32}, ), - "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.", - }, + "unload_model": ( + "BOOLEAN", + {"default": False, "tooltip": "Unload model from VRAM after execution to free memory."}, ), }, } @@ -148,7 +137,7 @@ class AILab_Florence2: return target def _get_model(self, model_name: str, precision: str, attention: str) -> Dict[str, Any]: - key = (model_name, precision, attention) + key = (model_name, precision) if key in self.MODEL_CACHE: return self.MODEL_CACHE[key] @@ -156,31 +145,37 @@ class AILab_Florence2: 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 + config_path = os.path.join(weights, "config.json") + checkpoint_path = os.path.join(weights, "model.safetensors") + if not os.path.exists(checkpoint_path): + checkpoint_path = os.path.join(weights, "pytorch_model.bin") + if not os.path.exists(checkpoint_path): + raise FileNotFoundError(f"No model weights found at {weights}") - model = Florence2ForConditionalGeneration.from_pretrained( - weights, - attn_implementation=attention, - torch_dtype=dtype, - ).to(offload) + config = Florence2Config.from_json(config_path) + model = Florence2(config, dtype=dtype, device="cpu", operations=comfy.ops.manual_cast) - processor = AutoProcessor.from_pretrained(weights, trust_remote_code=True) + sd = load_torch_file(checkpoint_path) + for key_name in [ + "language_model.model.encoder.embed_tokens.weight", + "language_model.model.decoder.embed_tokens.weight", + ]: + if key_name in sd and "language_model.model.shared.weight" in sd: + sd.pop(key_name, None) + + model.load_state_dict(sd, strict=False) + del sd + model.language_model.tie_weights() + model = model.eval().to(offload) + + processor = Processor(model_path=weights) 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] + def _prepare_prompt(mode: str, text_prompt: str) -> str: + config = MODE_CONFIGS.get(mode, MODE_CONFIGS[MODE_CHOICES[0]]) text = text_prompt.strip() if text and not config["allows_text"]: text = "" @@ -194,14 +189,17 @@ class AILab_Florence2: 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) + annotated = image_pil.copy() + draw = ImageDraw.Draw(annotated) 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 + try: + font = ImageFont.load_default(size=14) + except Exception: + font = ImageFont.load_default() + for index, (bbox, label) in enumerate(zip(predictions["bboxes"], predictions["labels"])): x0, y0, x1, y1 = bbox if y1 < y0: @@ -211,23 +209,14 @@ class AILab_Florence2: 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) + draw.rectangle([x0, y0, x1, y1], outline="red", width=2) + label_text = f"{index}.{label}" + color = random.choice(COLOR_BANK) + text_pos = (x0, max(0, y0 - 16)) + text_bbox = draw.textbbox(text_pos, label_text, font=font) + draw.rectangle(text_bbox, fill=color) + draw.text(text_pos, label_text, fill="white", font=font) annotated_tensor = TF.to_tensor(annotated)[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() if mask_layer is not None: @@ -275,27 +264,42 @@ class AILab_Florence2: def analyze( self, image: torch.Tensor, - model_name: str, - task: str, - precision: str, - attention: str, - fill_mask: bool, + prompt: str = "", + model: str = MODEL_CHOICES[0], + mode: str = MODE_CHOICES[0], + precision: str = "fp16", + attention: str = "sdpa", + fill_mask: bool = True, output_mask_select: str = "", - keep_model_loaded: bool = False, - text_prompt: str = "", + max_new_tokens: int = 1024, + unload_model: bool = False, + model_name: str = None, + task: str = None, + text_prompt: str = None, + keep_model_loaded: bool = None, + **kwargs, ): - bundle = self._get_model(model_name, precision, attention) - model = bundle["model"] + if model_name is not None and (model is None or model == MODEL_CHOICES[0]): + model = model_name + if task is not None and not mode: + mode = task + if text_prompt is not None and not prompt: + prompt = text_prompt + if keep_model_loaded is not None: + unload_model = not keep_model_loaded + + bundle = self._get_model(model, precision, attention) + florence_model = bundle["model"] processor = bundle["processor"] dtype = bundle["dtype"] device = mm.get_torch_device() offload = mm.unet_offload_device() - model.to(device) + florence_model.to(device) images = image.permute(0, 3, 1, 2) - task_config = TASK_CONFIGS[task] - prompt = self._prepare_prompt(task, text_prompt) + mode_config = MODE_CONFIGS.get(mode, MODE_CONFIGS[MODE_CHOICES[0]]) + prepared_prompt = self._prepare_prompt(mode, 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()] @@ -303,25 +307,22 @@ class AILab_Florence2: 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) + text=prepared_prompt, + images=tensor.unsqueeze(0), + ) - 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, + generated_ids = florence_model.generate( + input_ids=inputs["input_ids"].to(device), + pixel_values=inputs["pixel_values"].to(dtype=dtype, device=device), + max_new_tokens=max_new_tokens, + do_sample=False, + num_beams=3, ) 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"]] + parsed = processor.post_process_generation(raw, task=mode_config["token"], image_size=image_pil.size) + predictions = parsed[mode_config["token"]] - if task_config["mode"] == "bbox": + if mode_config["type"] == "bbox": annotated, mask = self._draw_regions(image_pil, predictions, fill_mask, select_filter) out_images.append(annotated) out_masks.append(mask) @@ -336,17 +337,21 @@ class AILab_Florence2: image_out = torch.cat(out_images, dim=0) mask_out = torch.cat(out_masks, dim=0) - if not keep_model_loaded: - model.to(offload) + if unload_model: + florence_model.to(offload) mm.soft_empty_cache() - return image_out, mask_out, out_data + all_bboxes = [] + for entry in out_data: + if isinstance(entry, dict) and "bboxes" in entry: + all_bboxes.extend(entry["bboxes"]) + + return image_out, mask_out, out_data, all_bboxes NODE_CLASS_MAPPINGS = { "AILab_Florence2": AILab_Florence2, } - NODE_DISPLAY_NAME_MAPPINGS = { "AILab_Florence2": "Florence2 (RMBG)", } diff --git a/py/AILab_ImageMaskTools.py b/py/AILab_ImageMaskTools.py index 58b6e60..def0660 100644 --- a/py/AILab_ImageMaskTools.py +++ b/py/AILab_ImageMaskTools.py @@ -2297,9 +2297,6 @@ class AILab_ImageResize: elif custom_width > 0 and custom_height > 0: target_width = custom_width target_height = custom_height - elif scale_by != 1.0: - target_width = int(orig_width * scale_by) - target_height = int(orig_height * scale_by) final_width, final_height = target_width, target_height diff --git a/py/AILab_RMBG.py b/py/AILab_RMBG.py index f18698f..5b6b764 100644 --- a/py/AILab_RMBG.py +++ b/py/AILab_RMBG.py @@ -29,7 +29,7 @@ import sys import importlib.util import types -from AILab_utils import get_device, unload_model, tensor2pil, pil2tensor +from AILab_utils import get_device, clean_vram, unload_model, tensor2pil, pil2tensor device = get_device() folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) @@ -182,6 +182,10 @@ class RMBGModel(BaseModelLoader): "from .BiRefNet_config import BiRefNetConfig", "from BiRefNetConfig import BiRefNetConfig" ) + birefnet_content = birefnet_content.replace( + "attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)", + "attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.to(attn.dtype).unsqueeze(1).unsqueeze(0)" + ) module_name = f"custom_birefnet_model_{hash(birefnet_path)}" module = types.ModuleType(module_name) @@ -235,6 +239,9 @@ class RMBGModel(BaseModelLoader): for param in self.model.parameters(): param.requires_grad = False + if hasattr(device, "type") and device.type == "cuda": + self.model.half() + torch.set_float32_matmul_precision('high') self.model.to(device) self.current_model_version = model_name @@ -247,18 +254,18 @@ class RMBGModel(BaseModelLoader): def extract_results(outputs): if isinstance(outputs, list) and len(outputs) > 0: - return outputs[-1].sigmoid().cpu() + return outputs[-1].float().sigmoid().cpu() if isinstance(outputs, dict) and 'logits' in outputs: - return outputs['logits'].sigmoid().cpu() + return outputs['logits'].float().sigmoid().cpu() if isinstance(outputs, torch.Tensor): - return outputs.sigmoid().cpu() + return outputs.float().sigmoid().cpu() try: if hasattr(outputs, 'last_hidden_state'): - return outputs.last_hidden_state.sigmoid().cpu() + return outputs.last_hidden_state.float().sigmoid().cpu() for _, value in outputs.items(): if isinstance(value, torch.Tensor): - return value.sigmoid().cpu() + return value.float().sigmoid().cpu() except Exception: pass @@ -279,8 +286,10 @@ class RMBGModel(BaseModelLoader): original_sizes = [tensor2pil(img).size for img in images] - input_tensors = [transform_image(tensor2pil(img)).unsqueeze(0) for img in images] + input_tensors = [transform_image(tensor2pil(img).convert("RGB")).unsqueeze(0) for img in images] input_batch = torch.cat(input_tensors, dim=0).to(device) + if hasattr(device, "type") and device.type == "cuda": + input_batch = input_batch.half() with torch.no_grad(): try: @@ -306,7 +315,9 @@ class RMBGModel(BaseModelLoader): print("[RMBG INFO] Batch inference failed due to tensor size mismatch; retrying images one by one.") masks = [] for img, (orig_w, orig_h) in zip(images, original_sizes): - single_input = transform_image(tensor2pil(img)).unsqueeze(0).to(device) + single_input = transform_image(tensor2pil(img).convert("RGB")).unsqueeze(0).to(device) + if hasattr(device, "type") and device.type == "cuda": + single_input = single_input.half() single_results = extract_results(self.model(single_input)) result = single_results[0].squeeze() result = result * (1 + (1 - params["sensitivity"])) @@ -636,7 +647,7 @@ class RMBG: if params.get("refine_foreground", False): refined_fg_local = refine_foreground(img_tensor_local, mask_tensor_b1hw) refined_fg_local = tensor2pil(refined_fg_local[0].permute(1, 2, 0)) - r, g, b = refined_fg_local.split() + r, g, b = refined_fg_local.convert("RGB").split() foreground_local = Image.merge('RGBA', (r, g, b, mask_img_local)) else: orig_rgba_local = orig_image_local.convert("RGBA") @@ -674,6 +685,8 @@ class RMBG: masks = [masks] for img_item, mask_item in zip(batch_imgs, masks): _process_pair(img_item, mask_item) + if torch.cuda.is_available(): + torch.cuda.empty_cache() else: for img in image: mask = model_instance.process_image(img, model, params) @@ -696,6 +709,8 @@ class RMBG: finally: if params.get("unload_model", False): unload_model(model_instance, "model") + else: + clean_vram() NODE_CLASS_MAPPINGS = { "RMBG": RMBG diff --git a/py/AILab_SAM3Segment.py b/py/AILab_SAM3Segment.py index 96fa9cb..8e63d43 100644 --- a/py/AILab_SAM3Segment.py +++ b/py/AILab_SAM3Segment.py @@ -126,8 +126,8 @@ class SAM3Segment: }, } - RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") - RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", "BBOX") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE", "BBOXES") FUNCTION = "segment" CATEGORY = "🧪AILab/🧽RMBG" @@ -173,14 +173,14 @@ class SAM3Segment: result_image = result_image.convert("RGBA") if background == "Alpha" else result_image.convert("RGB") empty_mask = torch.zeros((1, h, w), dtype=torch.float32) mask_rgb = empty_mask.reshape((-1, 1, h, w)).movedim(1, -1).expand(-1, -1, -1, 3) - return result_image, empty_mask, mask_rgb + return result_image, empty_mask, mask_rgb, [] def _empty_batch(self, img_pil): w, h = img_pil.size empty_imgs = torch.zeros((0, h, w, 3), dtype=torch.float32) empty_masks = torch.zeros((0, h, w), dtype=torch.float32) empty_mask_images = torch.zeros((0, h, w, 3), dtype=torch.float32) - return empty_imgs, empty_masks, empty_mask_images + return empty_imgs, empty_masks, empty_mask_images, [] def _run_single_per_instance(self, processor, img_tensor, prompt, confidence, max_segments, segment_pick, mask_blur, mask_offset, invert, unload_model, background, background_color): img_pil = tensor2pil(img_tensor) @@ -191,6 +191,7 @@ class SAM3Segment: state = processor.set_text_prompt(text, state) masks = state.get("masks") logits = state.get("masks_logits") + boxes = state.get("boxes") if masks is None or masks.numel() == 0: return self._empty_batch(img_pil) masks = masks.float() @@ -208,13 +209,24 @@ class SAM3Segment: topk = torch.topk(scores, k=max_segments) masks = masks[topk.indices] scores = scores[topk.indices] + if boxes is not None and boxes.shape[0] > 0: + boxes = boxes[topk.indices] sorted_idx = torch.argsort(scores, descending=True) masks = masks[sorted_idx] + if boxes is not None and boxes.shape[0] > 0: + boxes = boxes[sorted_idx] if segment_pick > 0: idx = segment_pick - 1 if idx >= masks.shape[0]: return self._empty_batch(img_pil) masks = masks[idx : idx + 1] + if boxes is not None and boxes.shape[0] > 0: + boxes = boxes[idx : idx + 1] + + bboxes_list = [] + if boxes is not None and boxes.shape[0] > 0: + for b in boxes: + bboxes_list.append([round(float(v), 2) for v in b.tolist()[:4]]) mask_imgs, mask_tensors, mask_rgb_list = [], [], [] for single_mask in masks: @@ -232,11 +244,12 @@ class SAM3Segment: torch.cat(mask_imgs, dim=0), torch.cat(mask_tensors, dim=0), torch.cat(mask_rgb_list, dim=0), + bboxes_list, ) 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( + imgs, masks, _, bboxes_list = self._run_single_per_instance( processor, img_tensor, prompt, confidence, max_segments, segment_pick, mask_blur, mask_offset, invert, unload_model, background, background_color, ) @@ -250,7 +263,7 @@ class SAM3Segment: result_image = result_image.convert("RGBA") if background == "Alpha" else result_image.convert("RGB") mask_tensor = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0) mask_rgb = mask_tensor.reshape((1, mask_image.height, mask_image.width, 1)).expand(-1, -1, -1, 3) - return result_image, mask_tensor, mask_rgb + return result_image, mask_tensor, mask_rgb, bboxes_list 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: @@ -265,48 +278,52 @@ class SAM3Segment: 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, - ) - result_images.append(imgs_batch) - result_masks.append(masks_batch) - result_mask_images.append(mask_imgs_batch) - else: - img_pil, mask_tensor, mask_rgb = self._run_single_merged( - processor, tensor_img, prompt, confidence_threshold, - max_segments, segment_pick, mask_blur, mask_offset, - invert_output, unload_model, background, background_color, - ) - result_images.append(pil2tensor(img_pil)) - result_masks.append(mask_tensor) - result_mask_images.append(mask_rgb) + result_bboxes = [] + try: + with ctx: + for tensor_img in image: + if output_mode == "Separate": + imgs_batch, masks_batch, mask_imgs_batch, bboxes_batch = self._run_single_per_instance( + processor, tensor_img, prompt, confidence_threshold, + max_segments, segment_pick, mask_blur, mask_offset, + invert_output, unload_model, background, background_color, + ) + result_images.append(imgs_batch) + result_masks.append(masks_batch) + result_mask_images.append(mask_imgs_batch) + result_bboxes.extend(bboxes_batch) + else: + img_pil, mask_tensor, mask_rgb, bboxes_batch = self._run_single_merged( + processor, tensor_img, prompt, confidence_threshold, + max_segments, segment_pick, mask_blur, mask_offset, + invert_output, unload_model, background, background_color, + ) + result_images.append(pil2tensor(img_pil)) + result_masks.append(mask_tensor) + result_mask_images.append(mask_rgb) + result_bboxes.extend(bboxes_batch) - if unload_model: - device_str = "cuda" if torch_device.type == "cuda" else "cpu" - cache_key = (model_name, device_str) - if cache_key in self.processor_cache: - del self.processor_cache[cache_key] + 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: + 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, result_bboxes + finally: + if unload_model: + device_str = "cuda" if torch_device.type == "cuda" else "cpu" + cache_key = (model_name, device_str) + if cache_key in self.processor_cache: + del self.processor_cache[cache_key] clean_vram() - 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: - img_pil = tensor2pil(image[0]) - empty_img, empty_mask, empty_mask_img = self._empty_result(img_pil, background, background_color) - final_images = pil2tensor(empty_img) - final_masks = empty_mask - final_mask_images = empty_mask_img - - return final_images, final_masks, final_mask_images - NODE_CLASS_MAPPINGS = {"SAM3Segment": SAM3Segment} NODE_DISPLAY_NAME_MAPPINGS = {"SAM3Segment": "SAM3 Segmentation (RMBG)"} diff --git a/py/AILab_VRMBG.py b/py/AILab_VRMBG.py new file mode 100644 index 0000000..55ee2a1 --- /dev/null +++ b/py/AILab_VRMBG.py @@ -0,0 +1,211 @@ +# ComfyUI-RMBG +# VRMBG Video Background Removal Node +# +# Model License Notice: +# - VRMBG-3.0: Provided by BRIA AI (https://huggingface.co/briaai/VRMBG-3.0) +# Please respect the original BRIA AI model license and usage terms. +# +# This integration script follows GPL-3.0 License. +# Source: https://github.com/1038lab/ComfyUI-RMBG + +import os +from typing import Tuple, Optional +import torch +import torch.nn.functional as F +import numpy as np +from PIL import Image, ImageFilter +from torchvision import transforms +import folder_paths +import comfy.model_management as mm + +from AILab_utils import get_device, clean_vram, unload_model as release_model, tensor2pil, pil2tensor + +device = get_device() +folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) + +MODEL_SUBDIR = os.path.join(folder_paths.models_dir, "RMBG", "VRMBG-3.0") +REPO_URL = "https://huggingface.co/briaai/VRMBG-3.0" + +INFER_SIZE = 1024 +REQUIRED_FILES = ["config.json", "model.safetensors", "model.py", "vrmbg3_config.py"] + + +class AILab_VRMBG: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE", {"tooltip": "Batch of consecutive video frames [B, H, W, C] for temporal background removal."}), + }, + "optional": { + "reset_state": ("BOOLEAN", {"default": True, "tooltip": "Reset temporal prior state at the beginning of this batch. Keep True for standalone clips, or False when feeding consecutive chunks."}), + "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": "Gaussian blur radius for softening mask edges."}), + "mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": "Expand (positive) or shrink (negative) the mask boundary in pixels."}), + "invert_output": ("BOOLEAN", {"default": False, "tooltip": "Invert the output mask."}), + "unload_model": ("BOOLEAN", {"default": False, "tooltip": "Unload model weights from VRAM immediately after execution."}), + "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": "Output format: Alpha for transparent cutout, or Color to composite against solid background."}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": "Hex background color used when background mode is set to Color."}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") + FUNCTION = "process" + CATEGORY = "🧪AILab/🧽RMBG" + + def __init__(self): + self.model = None + self.prev_rgb_t = None + self.prev_alpha = None + self.target_device = device + self.normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + + def _find_model_dir(self) -> str: + for base in folder_paths.get_folder_paths("rmbg"): + cand = os.path.join(base, "VRMBG-3.0") + if os.path.isdir(cand) and all(os.path.exists(os.path.join(cand, f)) for f in REQUIRED_FILES): + return cand + return MODEL_SUBDIR + + def _ensure_model(self): + if self.model is not None: + return self.model + + model_dir = self._find_model_dir() + missing_files = [f for f in REQUIRED_FILES if not os.path.exists(os.path.join(model_dir, f))] + + if missing_files: + raise FileNotFoundError( + f"[VRMBG-3.0] Model files missing: {missing_files}\n" + f"Automatic download has been disabled. Please download the model manually and place all required files in:\n" + f" {model_dir}\n\n" + f"Required files:\n" + f" - " + "\n - ".join(REQUIRED_FILES) + "\n\n" + f"Download link:\n" + f" - {REPO_URL}" + ) + + from transformers import AutoModelForImageSegmentation, PreTrainedModel + if not hasattr(PreTrainedModel, "all_tied_weights_keys"): + PreTrainedModel.all_tied_weights_keys = {} + + try: + self.model = AutoModelForImageSegmentation.from_pretrained( + model_dir, + trust_remote_code=True, + local_files_only=True + ).eval() + except Exception as e: + raise RuntimeError(f"[VRMBG-3.0] Failed to load local model from '{model_dir}': {e}") + + if not hasattr(self.model, "all_tied_weights_keys"): + self.model.all_tied_weights_keys = {} + + if torch.cuda.is_available() and str(self.target_device).startswith("cuda"): + self.model = self.model.half().to(self.target_device) + else: + self.model = self.model.to(self.target_device) + + return self.model + + def process(self, images, reset_state=True, mask_blur=0, mask_offset=0, + invert_output=False, unload_model=False, + background="Alpha", background_color="#222222", **kwargs): + do_unload = kwargs.get("unload_model", unload_model) + model = self._ensure_model() + + model_dtype = next(model.parameters()).dtype + + if reset_state or self.prev_rgb_t is None or self.prev_alpha is None: + self.prev_rgb_t = torch.zeros(3, INFER_SIZE, INFER_SIZE, device=self.target_device, dtype=model_dtype) + self.prev_alpha = torch.zeros(1, INFER_SIZE, INFER_SIZE, device=self.target_device, dtype=model_dtype) + + processed_images = [] + processed_masks = [] + + try: + for single_img in images: + orig_h, orig_w = single_img.shape[0], single_img.shape[1] + pil_img = tensor2pil(single_img).convert("RGB") + + # Permute [H, W, C] -> [C, H, W] + img_chw = single_img.permute(2, 0, 1).to(self.target_device, dtype=torch.float32) + # Resize to 1024x1024 + img_resized = F.interpolate(img_chw.unsqueeze(0), size=(INFER_SIZE, INFER_SIZE), mode="bilinear", align_corners=False).squeeze(0) + current_t = self.normalize(img_resized).to(dtype=model_dtype) + + # Pair with temporal prior: [current_RGB, prev_RGB * prev_alpha] + temporal_prior = (self.prev_rgb_t * self.prev_alpha).to(dtype=model_dtype) + paired = torch.cat([current_t, temporal_prior], dim=0).unsqueeze(0) + paired = paired.contiguous(memory_format=torch.channels_last) + + with torch.no_grad(): + raw_out = model(paired) + if isinstance(raw_out, (list, tuple)): + out_tensor = raw_out[-1] + else: + out_tensor = raw_out + pred_alpha = out_tensor.sigmoid().squeeze(0) # [1, 1024, 1024] + + # Update autoregressive temporal states for next frame + self.prev_rgb_t = current_t + self.prev_alpha = pred_alpha + + # Resize alpha back to native resolution + alpha_resized = F.interpolate(pred_alpha.unsqueeze(0).float(), size=(orig_h, orig_w), mode="bilinear", align_corners=False).squeeze(0).squeeze(0) + alpha_np = alpha_resized.detach().cpu().numpy() + + mask_img = Image.fromarray((np.clip(alpha_np, 0.0, 1.0) * 255.0).astype(np.uint8), mode="L") + + if mask_offset != 0: + offset_k = abs(mask_offset) * 2 + 1 + filter_type = ImageFilter.MaxFilter if mask_offset > 0 else ImageFilter.MinFilter + mask_img = mask_img.filter(filter_type(offset_k)) + + if mask_blur > 0: + mask_img = mask_img.filter(ImageFilter.GaussianBlur(radius=mask_blur)) + + if invert_output: + mask_np_inv = 255.0 - np.array(mask_img, dtype=np.float32) + mask_img = Image.fromarray(np.clip(mask_np_inv, 0, 255).astype(np.uint8), mode="L") + + mask_tensor = pil2tensor(mask_img) + processed_masks.append(mask_tensor) + + if background == "Alpha": + cutout = Image.new("RGBA", (orig_w, orig_h), (0, 0, 0, 0)) + cutout.paste(pil_img, (0, 0), mask_img) + processed_images.append(pil2tensor(cutout)) + else: + bg_hex = background_color.lstrip("#") + if len(bg_hex) == 6: + bg_rgb = tuple(int(bg_hex[i:i+2], 16) for i in (0, 2, 4)) + else: + bg_rgb = (34, 34, 34) + + bg_canvas = Image.new("RGB", (orig_w, orig_h), bg_rgb) + bg_canvas.paste(pil_img, (0, 0), mask_img) + processed_images.append(pil2tensor(bg_canvas)) + + out_image = torch.cat(processed_images, dim=0) + out_mask = torch.cat(processed_masks, dim=0) + out_mask_image = out_mask.unsqueeze(-1).expand(-1, -1, -1, 3) + + return (out_image, out_mask, out_mask_image) + + finally: + if do_unload: + release_model(self, "model") + self.prev_rgb_t = None + self.prev_alpha = None + else: + clean_vram() + + +NODE_CLASS_MAPPINGS = { + "AILab_VRMBG": AILab_VRMBG +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AILab_VRMBG": "VRMBG (Video Background Removal)" +} diff --git a/py/AILab_Yolo.py b/py/AILab_Yolo.py new file mode 100644 index 0000000..390bace --- /dev/null +++ b/py/AILab_Yolo.py @@ -0,0 +1,352 @@ +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 apply_background_color, clean_vram, get_or_download_model_file, process_mask + +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") +ANNOTATION_MODES = ("Mask & Box", "Box Only", "Mask Only", "None") +BACKGROUND_CHOICES = ("Alpha", "Color") + +YOLO_MODELS = { + "yolo26n-seg (Nano / 6.7MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo26n-seg.pt", + "filename": "yolo26n-seg.pt", + }, + "yolo26s-seg (Small / 23.5MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo26s-seg.pt", + "filename": "yolo26s-seg.pt", + }, + "yolo26m-seg (Medium / 54.8MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo26m-seg.pt", + "filename": "yolo26m-seg.pt", + }, + "yolo26l-seg (Large / 63.7MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo26l-seg.pt", + "filename": "yolo26l-seg.pt", + }, + "yolo26x-seg (Extra / 142.1MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo26x-seg.pt", + "filename": "yolo26x-seg.pt", + }, + "yolo11n-seg (Nano / 6.2MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo11n-seg.pt", + "filename": "yolo11n-seg.pt", + }, + "yolo11s-seg (Small / 19.8MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo11s-seg.pt", + "filename": "yolo11s-seg.pt", + }, + "yolo11m-seg (Medium / 44.8MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo11m-seg.pt", + "filename": "yolo11m-seg.pt", + }, + "yolo11l-seg (Large / 56.7MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo11l-seg.pt", + "filename": "yolo11l-seg.pt", + }, + "yolo11x-seg (Extra / 125.8MB)": { + "url": "https://huggingface.co/1038lab/YOLO/resolve/main/yolo11x-seg.pt", + "filename": "yolo11x-seg.pt", + }, +} + + +class AILab_Yolo: + CATEGORY = "🧪AILab/🧽RMBG" + RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "MASK") + RETURN_NAMES = ("IMAGE", "ANNOTATED_IMAGE", "MASK", "MASK_LIST") + FUNCTION = "yolo_detect" + + _MODEL_CACHE: Dict[str, "YOLO"] = {} + + @classmethod + def _list_models(cls) -> List[str]: + preset_models = list(YOLO_MODELS.keys()) + preset_filenames = {cfg["filename"].lower() for cfg in YOLO_MODELS.values()} + local_files = folder_paths.get_filename_list("ultralytics") + custom_models = [ + f + for f in local_files + if f.lower().endswith(".pt") + and os.path.basename(f).lower() not in preset_filenames + and f.lower() not in preset_filenames + ] + return preset_models + sorted(custom_models) + + @classmethod + def INPUT_TYPES(cls): + models = cls._list_models() + default_model = "yolo26m-seg (Medium / 54.8MB)" if "yolo26m-seg (Medium / 54.8MB)" in models else models[0] + + return { + "required": { + "images": ("IMAGE",), + "model": (tuple(models), {"default": default_model, "tooltip": f"YOLO model weights stored under {ULTRALYTICS_DIR} (presets download automatically)."}), + "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."}), + "annotation_mode": (ANNOTATION_MODES, {"default": "Mask & Box", "tooltip": "Display mode for ANNOTATED_IMAGE output."}), + "background": (BACKGROUND_CHOICES, {"default": "Alpha", "tooltip": "Background type for cutout IMAGE output."}), + "background_color": ("COLORCODE", {"default": "#222222", "tooltip": "Background color when background is set to Color."}), + "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": "Apply Gaussian blur to mask edges."}), + "mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": "Expand/Shrink mask boundary (positive = expand, negative = shrink)."}), + "invert_output": ("BOOLEAN", {"default": False, "tooltip": "Invert the mask output."}), + "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_Yolo] Invalid classes string: {value}. Ignoring filter.") + return None + + def _resolve_model_path(self, name: str) -> str: + if name in YOLO_MODELS: + filename = YOLO_MODELS[name]["filename"] + url = YOLO_MODELS[name]["url"] + return get_or_download_model_file(filename, url, "ultralytics") + for cfg in YOLO_MODELS.values(): + if name == cfg["filename"]: + return get_or_download_model_file(cfg["filename"], cfg["url"], "ultralytics") + local_path = folder_paths.get_full_path("ultralytics", name) + if local_path and os.path.exists(local_path): + return local_path + 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_annotated(self, result, image_pil: Image.Image, annotation_mode: str) -> torch.Tensor: + if annotation_mode == "None": + return pil2tensor(image_pil) + show_boxes = annotation_mode in ("Mask & Box", "Box Only") + show_masks = annotation_mode in ("Mask & Box", "Mask Only") + plotted = result.plot(boxes=show_boxes, masks=show_masks) + 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, + model=None, + mask_count="all", + select_mask_index: str = "none", + conf=0.25, + iou=0.45, + classes="", + annotation_mode="Mask & Box", + background="Alpha", + background_color="#222222", + mask_blur=0, + mask_offset=0, + invert_output=False, + device="auto", + max_det=300, + retina_masks=True, + agnostic_nms=False, + unload_model: bool = False, + yolo_model=None, + ): + try: + model_selected = model or yolo_model + model_path = self._resolve_model_path(model_selected) + model_inst = self._get_model(model_path) + device_target = self._resolve_device(device) + class_filter = self._parse_classes(classes) + + cutout_images: List[torch.Tensor] = [] + annotated_images: List[torch.Tensor] = [] + merged_masks: 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]) + width, height = image_pil.size + + results = model_inst( + image_pil, + conf=conf, + iou=iou, + classes=class_filter, + device=device_target, + max_det=max_det, + retina_masks=retina_masks, + agnostic_nms=agnostic_nms, + ) + + result = results[0] if results else None + if result is not None: + annotated_images.append(self._result_to_annotated(result, image_pil, annotation_mode)) + frame_masks = self._collect_masks(result, (width, height)) + else: + annotated_images.append(pil2tensor(image_pil)) + frame_masks = [torch.zeros((height, width), dtype=torch.float32)] + + 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: + raw_merged = self._merge_masks(selected_masks) + mask_list.extend(selected_masks) + else: + raw_merged = torch.zeros((height, width), dtype=torch.float32) + mask_list.append(raw_merged) + + mask_np = (raw_merged.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8) + mask_pil = Image.fromarray(mask_np, mode="L") + processed_mask_pil = process_mask( + mask_pil, + invert_output=invert_output, + mask_blur=mask_blur, + mask_offset=mask_offset, + ) + + cutout_pil = apply_background_color( + image_pil, + processed_mask_pil, + background=background, + background_color=background_color, + ) + if background == "Color": + cutout_pil = cutout_pil.convert("RGB") + else: + cutout_pil = cutout_pil.convert("RGBA") + cutout_images.append(pil2tensor(cutout_pil)) + + processed_mask_tensor = torch.from_numpy(np.array(processed_mask_pil).astype(np.float32) / 255.0) + merged_masks.append(processed_mask_tensor) + + if not cutout_images: + cutout_images = [images] + if not annotated_images: + annotated_images = [images] + 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)] + + cutout_tensor = torch.cat(cutout_images, dim=0) + annotated_tensor = torch.cat(annotated_images, dim=0) + merged_tensor = torch.stack(merged_masks, dim=0) + mask_tensor = torch.stack(mask_list, dim=0) + + return cutout_tensor, annotated_tensor, merged_tensor, mask_tensor + finally: + if unload_model: + self._MODEL_CACHE.clear() + clean_vram() + + +NODE_CLASS_MAPPINGS = { + "AILab_Yolo": AILab_Yolo, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AILab_Yolo": "YOLO Segmentation (RMBG)", +} diff --git a/py/AILab_utils.py b/py/AILab_utils.py index 9edc01d..6bf3247 100644 --- a/py/AILab_utils.py +++ b/py/AILab_utils.py @@ -323,6 +323,7 @@ ASPECT_RATIOS = { "4:5 (Portrait)": (4, 5), "3:4 (Portrait Standard)": (3, 4), "2:3 (Portrait Photo)": (2, 3), + "1:2 (Portrait Univisium)": (1, 2), "10:16 (Portrait)": (10, 16), "9:16 (Portrait Widescreen)": (9, 16), "9:21 (Portrait Ultrawide)": (9, 21), @@ -347,3 +348,17 @@ def color_format(color: str) -> str: color = color[:6] return f"#{color}" + + +def is_transformers_supported(min_version=(5, 0, 0)) -> bool: + try: + import transformers + parts = [] + for p in transformers.__version__.split(".")[:3]: + digits = "".join(c for c in p if c.isdigit()) + parts.append(int(digits) if digits else 0) + while len(parts) < 3: + parts.append(0) + return tuple(parts[:3]) >= min_version + except Exception: + return False