Add files via upload

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)

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)",
+}

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)"
+}

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)"
+}

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": "<REFERRING_EXPRESSION_SEGMENTATION>",
-        "mode": "polygon",
+        "type": "polygon",
         "allows_text": True,
     },
-    "Phrase Grounding (text boxes)": {
+    "Phrase Grounding (boxes)": {
         "token": "<CAPTION_TO_PHRASE_GROUNDING>",
-        "mode": "bbox",
+        "type": "bbox",
         "allows_text": True,
     },
     "Region Proposals (boxes only)": {
         "token": "<REGION_PROPOSAL>",
-        "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)",
 }

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
 

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

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)"}

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)"
+}

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)",
+}

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