From 69931a15b4fa161c295d28f1f535c9ad02e0a817 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 13 Aug 2026 17:48:02 +0200 Subject: [PATCH] refactor: split matting module into package and add BRIA RMBG 2.0 backend --- README.md | 13 +- js/CanvasView.js | 14 +- pyproject.toml | 2 +- python/__init__.py | 2 +- python/matting.py | 780 --------------------------- python/matting/__init__.py | 85 +++ python/matting/api.py | 298 ++++++++++ python/matting/backends/__init__.py | 2 + python/matting/backends/birefnet.py | 237 ++++++++ python/matting/backends/rmbg.py | 235 ++++++++ python/matting/catalog.py | 124 +++++ python/matting/options.py | 115 ++++ python/matting/paths.py | 100 ++++ python/matting/service.py | 160 ++++++ requirements.txt | 2 + src/CanvasView.ts | 20 +- tests/test_entrypoint_integration.py | 191 +++++-- 17 files changed, 1541 insertions(+), 839 deletions(-) delete mode 100644 python/matting.py create mode 100644 python/matting/__init__.py create mode 100644 python/matting/api.py create mode 100644 python/matting/backends/__init__.py create mode 100644 python/matting/backends/birefnet.py create mode 100644 python/matting/backends/rmbg.py create mode 100644 python/matting/catalog.py create mode 100644 python/matting/options.py create mode 100644 python/matting/paths.py create mode 100644 python/matting/service.py diff --git a/README.md b/README.md index 3dd005a..0fe222a 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ https://github.com/user-attachments/assets/9c7ce1de-873b-4a3b-8579-0fc67642af3a - **Drag & Drop** image files to create new layers. - **Copy & Paste** images directly from your system clipboard. - Import the last generated image from your workflow with one click. -- **AI-Powered Matting:** Optional background removal for any layer using the `BiRefNet` model. +- **AI-Powered Matting:** Optional background removal for any layer using `BiRefNet` or local `BRIA RMBG 2.0`. - **Efficient Memory Management:** An automatic garbage collection system cleans up unused image data to keep the browser's storage footprint low. - **Inputs** @@ -224,9 +224,10 @@ LayerForge is designed to work with **any ComfyUI-compatible model**. The node o ## 🧠 Optional: Matting Model (for image cutout) The "Matting" feature allows you to automatically generate a cutout (alpha mask) for a selected layer. This is an -optional feature and requires a model. +optional feature and requires a model. LayerForge supports ComfyUI's native BiRefNet loader and can also run BRIA +RMBG 2.0 locally through Transformers. -> - **Model Name**: `BiRefNet` +> - **Model Names**: `BiRefNet`, `BRIA RMBG 2.0` > - **Download from**: > - [Hugging Face](https://huggingface.co/ZhengPeng7/BiRefNet/tree/main) (Recommended) @@ -234,14 +235,16 @@ optional feature and requires a model. > - **Installation Path**: Place a full BiRefNet checkpoint in `ComfyUI/models/background_removal/`. > - **Automatic download**: If no compatible checkpoint is found, LayerForge downloads the standard full BiRefNet checkpoint automatically on first use. > - **Compatibility**: LayerForge uses ComfyUI's native BiRefNet loader, so the checkpoint must use the full BiRefNet architecture. The `lite-*` variants are not compatible with this native loader. +> - **BRIA RMBG 2.0**: Select it in Matting settings to download the local Transformers model into `ComfyUI/models/background_removal/RMBG-2.0/`. The Hugging Face repository is gated, so access must be accepted there before downloading. +> - **BRIA license**: The self-hosted RMBG 2.0 weights are released for non-commercial use under the model's stated license. Check the [official model card](https://huggingface.co/briaai/RMBG-2.0) before using it commercially. Click the gear button next to **Matting** to configure its behavior. Settings are saved locally in the browser and include: -- **Model**: automatically select a compatible local checkpoint, choose an installed local checkpoint, or select an official BiRefNet variant for download on first use. +- **Model**: automatically select a compatible local checkpoint, choose an installed local checkpoint, or select an official BiRefNet or BRIA RMBG 2.0 model for download on first use. - **Processing mode**: remove the background, remove the detected foreground, or apply the generated mask to LayerForge's Draw Mask output-area mask. - **Mask threshold**: use `0` for a soft alpha mask or a higher value for a harder cutout. -The model selector separates checkpoints already installed locally from official full-size BiRefNet variants available for download. Downloaded variants are stored directly in `ComfyUI/models/background_removal/` with descriptive filenames such as `BiRefNet-general.safetensors` and `BiRefNet-portrait.safetensors`, then validated against ComfyUI's native loader before use. Lite variants are intentionally excluded because they use a different architecture. Hugging Face cache folders are ignored in the selector and are not deleted. +The model selector separates compatible local models from official models available for download. BiRefNet variants are stored directly in `ComfyUI/models/background_removal/` with descriptive filenames such as `BiRefNet-general.safetensors` and `BiRefNet-portrait.safetensors`. BRIA RMBG 2.0 is stored as a complete Transformers model under `ComfyUI/models/background_removal/RMBG-2.0/`. Every selected model is validated before use. Lite BiRefNet variants are intentionally excluded because they use a different architecture. Hugging Face cache folders are ignored in the selector and are not deleted. ### BiRefNet model guide diff --git a/js/CanvasView.js b/js/CanvasView.js index 32e8b46..6b2a07f 100644 --- a/js/CanvasView.js +++ b/js/CanvasView.js @@ -301,12 +301,14 @@ async function createCanvasWidget(node, widget, app) { modelDetails.hidden = false; return; } - modelDescription.textContent = selectedOption.description || 'Installed checkpoint validated by ComfyUI\'s native BiRefNet loader.'; + modelDescription.textContent = selectedOption.description || (selectedOption.backend === 'rmbg' + ? 'Local BRIA RMBG 2.0 model loaded through Transformers.' + : 'Installed checkpoint validated by ComfyUI\'s native BiRefNet loader.'); if (selectedOption.url) { modelLinks.appendChild(createModelLink('Model page', selectedOption.url)); } if (selectedOption.project_url) { - modelLinks.appendChild(createModelLink('BiRefNet project', selectedOption.project_url)); + modelLinks.appendChild(createModelLink(selectedOption.backend === 'rmbg' ? 'BRIA project' : 'BiRefNet project', selectedOption.project_url)); } modelDetails.hidden = false; }; @@ -346,7 +348,7 @@ async function createCanvasWidget(node, widget, app) { remoteCount > 0 ? `${remoteCount} official model(s) available for download` : '', ].filter(Boolean).join('; '); modelStatus.textContent = `${modelCounts}. ${modelStatusMessage}`; - body.append(createRow('Model', modelSelect, 'Only checkpoints accepted by the native ComfyUI BiRefNet loader are listed.'), modelDetails, createRow('Processing mode', modeSelect, 'The selected mode controls what the Matting button creates from the detected mask.'), createRow('Mask threshold', thresholdContainer, 'Set to 0 for a soft alpha mask; higher values create a harder cutout.'), modelStatus); + body.append(createRow('Model', modelSelect, 'Choose a local BiRefNet checkpoint or BRIA RMBG 2.0, or download an official model on first use.'), modelDetails, createRow('Processing mode', modeSelect, 'The selected mode controls what the Matting button creates from the detected mask.'), createRow('Mask threshold', thresholdContainer, 'Set to 0 for a soft alpha mask; higher values create a harder cutout.'), modelStatus); const actions = document.createElement('div'); actions.className = 'lf-matting-settings-actions'; const resetButton = document.createElement('button'); @@ -657,12 +659,12 @@ async function createCanvasWidget(node, widget, app) { showErrorNotification(modelStatus.message, 8000); return; case 'not_downloaded': - showWarningNotification("The selected BiRefNet model will be downloaded automatically (requires internet connection).", 5000); + showWarningNotification(modelStatus.message || "The selected background-removal model will be downloaded automatically.", 7000); // Ask user if they want to proceed with download - if (!confirm("The selected BiRefNet model needs to be downloaded. This is a one-time download and may be large. Do you want to proceed?")) { + if (!confirm(`${modelStatus.message || "The selected background-removal model needs to be downloaded."}\n\nThis is a one-time download and may be large. Do you want to proceed?`)) { return; } - showInfoNotification("Downloading the selected BiRefNet model... This may take a few minutes.", 10000); + showInfoNotification("Downloading the selected background-removal model... This may take a few minutes.", 10000); break; case 'selected_model_unavailable': showErrorNotification(modelStatus.message, 8000); diff --git a/pyproject.toml b/pyproject.toml index 3524a07..c8b06dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ description = "Photoshop-like layered canvas editor to your ComfyUI workflow. Th version = "1.5.13" license = { text = "MIT License" } requires-python = ">=3.10" -dependencies = ["torch", "torchvision", "huggingface_hub", "aiohttp", "numpy", "tqdm", "Pillow"] +dependencies = ["torch", "torchvision", "huggingface_hub", "aiohttp", "numpy", "tqdm", "Pillow", "transformers>=4.50.3", "kornia"] [project.urls] Repository = "https://github.com/Azornes/Comfyui-LayerForge" diff --git a/python/__init__.py b/python/__init__.py index 848b8d1..e636454 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -2,5 +2,5 @@ The runtime modules are intentionally imported explicitly by the entry point: ``node`` owns the ComfyUI node, ``routes`` owns HTTP/WebSocket registration, -``image_utils`` owns tensor conversions, and ``matting`` owns BiRefNet. +``image_utils`` owns tensor conversions, and ``matting`` owns background-removal backends. """ diff --git a/python/matting.py b/python/matting.py deleted file mode 100644 index 09a7986..0000000 --- a/python/matting.py +++ /dev/null @@ -1,780 +0,0 @@ -"""BiRefNet model discovery, download, inference, and HTTP endpoints.""" - -import hashlib -import os -import shutil -import threading -import traceback -from typing import ClassVar - -import folder_paths -import torch -import torch.nn.functional as F -from aiohttp import web -from server import PromptServer -from torchvision import transforms - -from .image_utils import convert_base64_to_tensor, convert_tensor_to_base64 -from .node import log - -_BIREFNET_REPOSITORY = "ZhengPeng7/BiRefNet" -_BIREFNET_FILENAME = "model.safetensors" -_BIREFNET_DEFAULT_LOCAL_FILENAME = "BiRefNet-general.safetensors" -_BIREFNET_PROJECT_URL = "https://github.com/ZhengPeng7/BiRefNet" -_BIREFNET_REMOTE_PREFIX = "remote:" -_BIREFNET_REQUIRED_KEYS = { - "bb.layers.1.blocks.0.attn.relative_position_index", - "bb.layers.2.blocks.17.attn.qkv.weight", -} -_BIREFNET_MODEL_CATALOG = ( - { - "id": "general", - "label": "BiRefNet — General", - "description": "The best starting point for everyday images and general background removal.", - "repo_id": "ZhengPeng7/BiRefNet", - "filename": "model.safetensors", - "local_filename": "BiRefNet-general.safetensors", - }, - { - "id": "high_resolution", - "label": "BiRefNet — High Resolution", - "description": "High-resolution segmentation for detailed edges and larger source images; it uses more memory.", - "repo_id": "ZhengPeng7/BiRefNet_HR", - "filename": "model.safetensors", - "local_filename": "BiRefNet-HR.safetensors", - }, - { - "id": "portrait", - "label": "BiRefNet — Portrait", - "description": "Portrait matting for people, hair, and portrait-focused cutouts.", - "repo_id": "ZhengPeng7/BiRefNet-portrait", - "filename": "model.safetensors", - "local_filename": "BiRefNet-portrait.safetensors", - }, - { - "id": "matting", - "label": "BiRefNet — Matting", - "description": "General matting with a focus on soft alpha edges such as hair and semi-transparent details.", - "repo_id": "ZhengPeng7/BiRefNet-matting", - "filename": "model.safetensors", - "local_filename": "BiRefNet-matting.safetensors", - }, - { - "id": "high_resolution_matting", - "label": "BiRefNet — High Resolution Matting", - "description": "High-resolution general matting for fine details; it requires more memory than the standard matting model.", - "repo_id": "ZhengPeng7/BiRefNet_HR-matting", - "filename": "model.safetensors", - "local_filename": "BiRefNet-HR-matting.safetensors", - }, - { - "id": "dynamic", - "label": "BiRefNet — Dynamic", - "description": "Dynamic-shape segmentation for inputs with varying aspect ratios and resolutions.", - "repo_id": "ZhengPeng7/BiRefNet_dynamic", - "filename": "model.safetensors", - "local_filename": "BiRefNet_dynamic.safetensors", - }, - { - "id": "dynamic_matting", - "label": "BiRefNet — Dynamic Matting", - "description": "Dynamic-shape matting for arbitrary input sizes, with a focus on soft alpha edges.", - "repo_id": "ZhengPeng7/BiRefNet_dynamic-matting", - "filename": "model.safetensors", - "local_filename": "BiRefNet_dynamic-matting.safetensors", - }, - { - "id": "hrsod", - "label": "BiRefNet — HRSOD", - "description": "High-resolution salient-object detection; useful when the main subject should stand out from its surroundings.", - "repo_id": "ZhengPeng7/BiRefNet-HRSOD", - "filename": "model.safetensors", - "local_filename": "BiRefNet-HRSOD.safetensors", - }, - { - "id": "dis5k", - "label": "BiRefNet — DIS5K", - "description": "Dichotomous image segmentation trained for clean foreground/background separation.", - "repo_id": "ZhengPeng7/BiRefNet-DIS5K", - "filename": "model.safetensors", - "local_filename": "BiRefNet-DIS5K.safetensors", - }, - { - "id": "cod", - "label": "BiRefNet — COD", - "description": "Camouflaged-object detection; use it for subjects that blend into their background.", - "repo_id": "ZhengPeng7/BiRefNet-COD", - "filename": "model.safetensors", - "local_filename": "BiRefNet-COD.safetensors", - }, -) - - -def _get_comfy_birefnet_loader(): - """Return ComfyUI's native BiRefNet loader when it is available.""" - try: - from comfy.bg_removal_model import load - - return load - except Exception as error: - log.debug(f"Native ComfyUI BiRefNet loader is unavailable: {error}") - return None - - -def _get_birefnet_base_paths(): - """Return native ComfyUI and legacy locations that may contain BiRefNet.""" - paths = [] - get_folder_paths = getattr(folder_paths, "get_folder_paths", None) - if callable(get_folder_paths): - try: - paths.extend(get_folder_paths("background_removal")) - except (KeyError, TypeError): - pass - - comfy_models_dir = getattr(folder_paths, "models_dir", None) - if comfy_models_dir: - paths.extend( - [ - os.path.join(comfy_models_dir, "background_removal"), - os.path.join(comfy_models_dir, "RMBG", "BiRefNet"), - os.path.join(comfy_models_dir, "BiRefNet"), - ] - ) - - unique_paths = [] - seen = set() - for path in paths: - normalized = os.path.normcase(os.path.normpath(path)) - if normalized not in seen: - seen.add(normalized) - unique_paths.append(path) - return unique_paths - - -def _is_native_birefnet_checkpoint(path): - """Check the checkpoint signature without loading all weights into memory.""" - if not os.path.isfile(path) or not path.lower().endswith(".safetensors"): - return False - - try: - from safetensors import safe_open - - with safe_open(path, framework="pt") as checkpoint: - return _BIREFNET_REQUIRED_KEYS.issubset(checkpoint.keys()) - except Exception as error: - log.debug(f"Unable to inspect BiRefNet checkpoint {path}: {error}") - return False - - -def _iter_birefnet_checkpoint_paths(): - """Yield candidate checkpoints from native and legacy model directories.""" - for base_path in _get_birefnet_base_paths(): - if not os.path.isdir(base_path): - continue - - for root, directories, files in os.walk(base_path): - directories[:] = [ - directory - for directory in directories - if directory not in { - ".git", - ".no_exist", - ".cache", - "__pycache__", - "refs", - "snapshots", - "blobs", - } - and not directory.startswith("models--") - ] - for filename in sorted(files): - if filename.lower().endswith(".safetensors"): - yield os.path.join(root, filename) - - -def _find_local_birefnet_model(model_path=None): - """Find a full BiRefNet checkpoint accepted by ComfyUI's native loader.""" - candidates = [] - seen = set() - for path in _iter_birefnet_checkpoint_paths(): - normalized = os.path.normcase(os.path.normpath(path)) - if normalized in seen: - continue - seen.add(normalized) - if _is_native_birefnet_checkpoint(path): - candidates.append(path) - - if model_path and model_path != "auto": - requested = os.path.normcase(os.path.normpath(os.path.abspath(model_path))) - for candidate in candidates: - if os.path.normcase(os.path.normpath(os.path.abspath(candidate))) == requested: - return candidate - return None - - if not candidates: - return None - - priority = { - "birefnet.safetensors": 0, - "model.safetensors": 1, - "birefnet-general.safetensors": 2, - "birefnet-hr.safetensors": 3, - } - return min( - candidates, - key=lambda path: (priority.get(os.path.basename(path).lower(), 10), path.lower()), - ) - - -def _get_birefnet_remote_model(model_path): - """Resolve a client-facing remote model identifier from the fixed catalog.""" - if not isinstance(model_path, str) or not model_path.startswith(_BIREFNET_REMOTE_PREFIX): - return None - - model_id = model_path[len(_BIREFNET_REMOTE_PREFIX) :] - return next((model for model in _BIREFNET_MODEL_CATALOG if model["id"] == model_id), None) - - -def _get_birefnet_remote_checkpoint_path(model): - """Return the managed local path for a catalog model.""" - base_paths = _get_birefnet_base_paths() - if not base_paths: - return None - return os.path.join(base_paths[0], model["local_filename"]) - - -def _migrate_birefnet_checkpoint( - checkpoint_path, - legacy_path, - *, - log_description, - success_suffix="to", -): - """Return an existing friendly checkpoint or migrate its legacy filename.""" - if checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path): - return checkpoint_path - - if ( - not checkpoint_path - or not legacy_path - or legacy_path == checkpoint_path - or not _is_native_birefnet_checkpoint(legacy_path) - ): - return None - - try: - os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True) - os.replace(legacy_path, checkpoint_path) - log.info(f"Renamed {log_description} {success_suffix} {checkpoint_path}") - return checkpoint_path - except OSError as error: - log.warning(f"Unable to rename {log_description} {legacy_path}: {error}") - return legacy_path - - -def _find_existing_birefnet_remote_checkpoint(model): - """Return an installed managed remote checkpoint, when available.""" - checkpoint_path = _get_birefnet_remote_checkpoint_path(model) - if checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path): - return checkpoint_path - return None - - -def _get_birefnet_model_options(): - """Return local checkpoints and official models available for download.""" - _find_existing_birefnet_default_checkpoint() - local_options = [] - seen = set() - base_paths = _get_birefnet_base_paths() - managed_paths = { - os.path.normcase(os.path.normpath(os.path.abspath(path))) - for base_path in base_paths - for model in _BIREFNET_MODEL_CATALOG - for path in [os.path.join(base_path, model["local_filename"])] - } - - for path in _iter_birefnet_checkpoint_paths(): - normalized = os.path.normcase(os.path.normpath(os.path.abspath(path))) - if normalized in seen or not _is_native_birefnet_checkpoint(path): - continue - if normalized in managed_paths: - continue - - seen.add(normalized) - label = os.path.basename(path) - for base_path in base_paths: - try: - relative_path = os.path.relpath(path, base_path) - except ValueError: - continue - if relative_path == os.pardir or relative_path.startswith(os.pardir + os.sep): - continue - label = relative_path.replace(os.sep, "/") - break - - local_options.append({ - "path": path, - "label": label, - "source": "local", - "downloaded": True, - }) - - remote_options = [] - for model in _BIREFNET_MODEL_CATALOG: - checkpoint_path = _find_existing_birefnet_remote_checkpoint(model) - remote_options.append({ - "path": f"{_BIREFNET_REMOTE_PREFIX}{model['id']}", - "label": model["label"], - "description": model["description"], - "url": f"https://huggingface.co/{model['repo_id']}", - "project_url": _BIREFNET_PROJECT_URL, - "source": "remote", - "downloaded": bool(checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path)), - }) - - local_options.sort(key=lambda option: option["label"].lower()) - return local_options + remote_options - - -def _get_birefnet_download_dir(): - paths = _get_birefnet_base_paths() - if not paths: - raise RuntimeError("ComfyUI did not expose a background_removal model directory") - - download_dir = paths[0] - os.makedirs(download_dir, exist_ok=True) - return download_dir - - -def _get_birefnet_default_checkpoint_path(): - """Return the friendly path used for the automatic General checkpoint.""" - base_paths = _get_birefnet_base_paths() - if not base_paths: - return None - return os.path.join(base_paths[0], _BIREFNET_DEFAULT_LOCAL_FILENAME) - - -def _find_existing_birefnet_default_checkpoint(): - """Find the automatic checkpoint and migrate its old generic filename.""" - checkpoint_path = _get_birefnet_default_checkpoint_path() - base_paths = _get_birefnet_base_paths() - if not base_paths: - return None - - legacy_path = os.path.join(base_paths[0], _BIREFNET_FILENAME) - return _migrate_birefnet_checkpoint( - checkpoint_path, - legacy_path, - log_description="the automatic BiRefNet checkpoint", - ) - - -def _download_birefnet_checkpoint(model=None): - """Download and validate a full BiRefNet checkpoint into ComfyUI's model path.""" - try: - from huggingface_hub import hf_hub_download - except ImportError as error: - raise RuntimeError( - "Automatic BiRefNet download requires the 'huggingface_hub' package. " - "Install the LayerForge requirements or place a compatible checkpoint in " - "ComfyUI/models/background_removal/." - ) from error - - if model is None: - repository = _BIREFNET_REPOSITORY - filename = _BIREFNET_FILENAME - download_dir = _get_birefnet_download_dir() - model_label = "BiRefNet — General" - target_path = os.path.join(download_dir, _BIREFNET_DEFAULT_LOCAL_FILENAME) - else: - repository = model["repo_id"] - filename = model["filename"] - download_dir = _get_birefnet_download_dir() - model_label = model["label"] - target_path = _get_birefnet_remote_checkpoint_path(model) - - log.info(f"Downloading {model_label} from Hugging Face into {download_dir}...") - try: - downloaded_path = hf_hub_download( - repo_id=repository, - filename=filename, - local_dir=download_dir, - local_dir_use_symlinks=False, - ) - except TypeError: - downloaded_path = hf_hub_download( - repo_id=repository, - filename=filename, - local_dir=download_dir, - ) - - if not _is_native_birefnet_checkpoint(downloaded_path): - raise RuntimeError( - f"Downloaded {model_label} is not a ComfyUI-compatible BiRefNet checkpoint: {downloaded_path}" - ) - - if target_path and os.path.normcase(os.path.abspath(downloaded_path)) != os.path.normcase( - os.path.abspath(target_path) - ): - os.makedirs(os.path.dirname(target_path), exist_ok=True) - try: - os.replace(downloaded_path, target_path) - except OSError: - shutil.copy2(downloaded_path, target_path) - downloaded_path = target_path - - if target_path and not _is_native_birefnet_checkpoint(downloaded_path): - raise RuntimeError( - f"Renamed {model_label} is not a ComfyUI-compatible BiRefNet checkpoint: {downloaded_path}" - ) - - log.info(f"{model_label} checkpoint is ready at {downloaded_path}") - return downloaded_path - - -def _ensure_birefnet_checkpoint(model_path=None): - remote_model = _get_birefnet_remote_model(model_path) - if remote_model: - checkpoint_path = _find_existing_birefnet_remote_checkpoint(remote_model) - if checkpoint_path: - return checkpoint_path - return _download_birefnet_checkpoint(remote_model) - - if not model_path or model_path == "auto": - _find_existing_birefnet_default_checkpoint() - - checkpoint_path = _find_local_birefnet_model(model_path) - if checkpoint_path: - return checkpoint_path - - if model_path and model_path != "auto": - raise RuntimeError("The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.") - - return _download_birefnet_checkpoint() - - -class BiRefNetMatting: - """Adapter around ComfyUI's native BiRefNet loader.""" - - _model_cache: ClassVar[dict] = {} - _model_cache_lock = threading.Lock() - - def __init__(self): - self.model = None - self.model_path = None - - def load_model(self, model_path=None): - loader = _get_comfy_birefnet_loader() - if loader is None: - raise RuntimeError( - "This ComfyUI version does not provide the native BiRefNet background-removal loader." - ) - - checkpoint_path = _ensure_birefnet_checkpoint(model_path) - with self._model_cache_lock: - if checkpoint_path not in self._model_cache: - log.info(f"Loading BiRefNet with ComfyUI's native loader from {checkpoint_path}") - model = loader(checkpoint_path) - if model is None: - raise RuntimeError( - f"ComfyUI did not recognize the BiRefNet checkpoint: {checkpoint_path}" - ) - self._model_cache[checkpoint_path] = model - else: - log.debug(f"Using cached native BiRefNet model from {checkpoint_path}") - - self.model = self._model_cache[checkpoint_path] - self.model_path = checkpoint_path - - return self.model - - def preprocess_image(self, image): - if not isinstance(image, torch.Tensor): - image = transforms.ToTensor()(image).unsqueeze(0) - - if image.dim() == 3: - if image.shape[0] in (1, 3, 4): - image = image.movedim(0, -1).unsqueeze(0) - else: - image = image.unsqueeze(0) - elif image.dim() == 4 and image.shape[1] in (1, 3, 4): - image = image.movedim(1, -1) - - if image.dim() != 4 or image.shape[-1] not in (1, 3, 4): - raise ValueError(f"Expected an image tensor in BCHW or BHWC format, got {tuple(image.shape)}") - - if image.shape[-1] == 1: - image = image.expand(-1, -1, -1, 3) - elif image.shape[-1] == 4: - image = image[..., :3] - - return image.to(dtype=torch.float32).contiguous() - - def execute(self, image, model_path, threshold=0.5, refinement=1, mode="remove_background"): - try: - PromptServer.instance.send_sync("matting_status", {"status": "processing"}) - del refinement - if mode not in {"remove_background", "remove_foreground", "mask_only"}: - raise ValueError(f"Unsupported matting mode: {mode}") - - threshold = float(threshold) - if not 0.0 <= threshold <= 1.0: - raise ValueError("Matting threshold must be between 0 and 1") - - image_tensor = self.preprocess_image(image) - original_size = (image_tensor.shape[1], image_tensor.shape[2]) - log.debug(f"Original size: {original_size}") - self.load_model(model_path) - log.debug(f"Processed image shape: {image_tensor.shape}, dtype: {image_tensor.dtype}") - - with torch.no_grad(): - result = self.model.encode_image(image_tensor) - if result.dim() == 3: - result = result.unsqueeze(1) - elif result.dim() == 2: - result = result.unsqueeze(0).unsqueeze(0) - else: - raise ValueError(f"Unexpected BiRefNet output shape: {tuple(result.shape)}") - - result = result.to(device=image_tensor.device, dtype=torch.float32) - if result.shape[-2:] != original_size: - result = F.interpolate( - result, - size=original_size, - mode="bilinear", - align_corners=False, - ) - result = result.clamp(0.0, 1.0) - log.debug(f"Native BiRefNet output shape: {result.shape}, dtype: {result.dtype}") - - if threshold > 0: - result = (result > threshold).to(dtype=torch.float32) - - alpha_mask = 1.0 - result if mode == "remove_foreground" else result - if mode == "mask_only": - masked_image = alpha_mask.expand(-1, 3, -1, -1) - else: - masked_image = image_tensor.movedim(-1, 1) * alpha_mask - - PromptServer.instance.send_sync("matting_status", {"status": "completed"}) - return masked_image, alpha_mask - except Exception: - PromptServer.instance.send_sync("matting_status", {"status": "error"}) - raise - - @classmethod - def IS_CHANGED(cls, image, model_path, threshold, refinement, mode="remove_background"): - digest = hashlib.md5() - digest.update(str(image).encode()) - digest.update(str(model_path).encode()) - digest.update(str(threshold).encode()) - digest.update(str(refinement).encode()) - digest.update(str(mode).encode()) - return digest.hexdigest() - - -_matting_lock = None - - -def _matting_status_response( - available, - reason, - message, - *, - status=200, - **details, -): - return web.json_response( - { - "available": available, - "reason": reason, - "message": message, - **details, - }, - status=status, - ) - - -async def check_matting_model(request): - """Report whether the selected native BiRefNet checkpoint is ready.""" - try: - if _get_comfy_birefnet_loader() is None: - return _matting_status_response( - False, - "unsupported_comfyui", - "This ComfyUI version does not provide the native BiRefNet background-removal loader.", - ) - - model_options = _get_birefnet_model_options() - requested_model = request.query.get("model_path") or "auto" - remote_model = _get_birefnet_remote_model(requested_model) - - if remote_model: - local_model_path = _find_existing_birefnet_remote_checkpoint(remote_model) - if local_model_path: - log.info(f"Selected BiRefNet model is ready at {local_model_path}") - return _matting_status_response( - True, - "ready", - "Selected model is ready to use", - model_path=local_model_path, - selected_model=remote_model["label"], - models=model_options, - ) - - log.info(f"Selected BiRefNet model is not downloaded: {remote_model['label']}") - return _matting_status_response( - False, - "not_downloaded", - f"{remote_model['label']} will be downloaded automatically on first use.", - model_path=requested_model, - selected_model=remote_model["label"], - models=model_options, - ) - - if requested_model != "auto": - local_model_path = _find_local_birefnet_model(requested_model) - if local_model_path: - log.info(f"Selected BiRefNet model is ready at {local_model_path}") - return _matting_status_response( - True, - "ready", - "Selected model is ready to use", - model_path=local_model_path, - selected_model=local_model_path, - models=model_options, - ) - - return _matting_status_response( - False, - "selected_model_unavailable", - "The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.", - model_path=requested_model, - models=model_options, - ) - - local_model_path = _find_local_birefnet_model() - if local_model_path: - log.info(f"BiRefNet model files detected at {local_model_path}") - return _matting_status_response( - True, - "ready", - "Model is ready to use", - model_path=local_model_path, - models=model_options, - ) - - searched_paths = _get_birefnet_base_paths() - log.info(f"BiRefNet model not found in any of: {searched_paths}") - return _matting_status_response( - False, - "not_downloaded", - "The BiRefNet checkpoint will be downloaded automatically on first use (requires internet connection).", - model_path=searched_paths[0] if searched_paths else None, - models=model_options, - ) - except Exception as error: - log.error(f"Error checking matting model: {error}") - return _matting_status_response( - False, - "error", - f"Error checking model status: {error}", - status=500, - ) - - -async def matting(request): - global _matting_lock - - if _matting_lock is not None: - log.warning("Matting already in progress, rejecting request") - return web.json_response({ - "error": "Another matting operation is in progress", - "details": "Please wait for the current operation to complete", - }, status=429) - - _matting_lock = True - try: - log.info("Received matting request") - data = await request.json() - matting_instance = BiRefNetMatting() - image_tensor, original_alpha = convert_base64_to_tensor(data["image"]) - log.debug(f"Input image shape: {image_tensor.shape}") - - mode = data.get("mode", "remove_background") - model_path = data.get("model_path") or "auto" - matted_image, alpha_mask = matting_instance.execute( - image_tensor, - model_path, - threshold=data.get("threshold", 0.5), - refinement=data.get("refinement", 1), - mode=mode, - ) - - if mode == "mask_only": - result_image = convert_tensor_to_base64(alpha_mask) - else: - result_image = convert_tensor_to_base64(matted_image, alpha_mask, original_alpha) - result_mask = convert_tensor_to_base64(alpha_mask) - # Draw Mask uses white pixels as the area to remove, while matting uses - # white pixels as the area to keep. - draw_mask = 1.0 - alpha_mask - - return web.json_response({ - "matted_image": result_image, - "alpha_mask": result_mask, - "draw_mask": convert_tensor_to_base64(draw_mask), - "mode": mode, - "model_path": matting_instance.model_path, - }) - except RuntimeError as error: - log.error(f"Runtime error during matting: {error}") - return web.json_response({ - "error": "Matting Model Error", - "details": str(error), - }, status=500) - except Exception as error: - log.exception(f"Error in matting endpoint: {error}") - error_text = str(error).lower() - if any(marker in error_text for marker in ( - "offline", "connection", "timed out", "huggingface", "localentrynotfound" - )): - return web.json_response({ - "error": "Network Connection Error", - "details": "Failed to download the BiRefNet model from Hugging Face. Please check your internet connection.", - }, status=400) - - return web.json_response({ - "error": "An unexpected error occurred", - "details": traceback.format_exc(), - }, status=500) - finally: - _matting_lock = None - log.debug("Matting lock released") - - -def register_matting_routes(): - """Register matting endpoints without import-time decorators.""" - PromptServer.instance.routes.get("/matting/check-model")(check_matting_model) - PromptServer.instance.routes.post("/matting")(matting) - - -__all__ = [ - "_BIREFNET_DEFAULT_LOCAL_FILENAME", - "_BIREFNET_FILENAME", - "_BIREFNET_MODEL_CATALOG", - "_BIREFNET_PROJECT_URL", - "_BIREFNET_REMOTE_PREFIX", - "BiRefNetMatting", - "_download_birefnet_checkpoint", - "_ensure_birefnet_checkpoint", - "_find_existing_birefnet_default_checkpoint", - "_find_existing_birefnet_remote_checkpoint", - "_find_local_birefnet_model", - "_get_birefnet_model_options", - "_get_comfy_birefnet_loader", - "_is_native_birefnet_checkpoint", - "check_matting_model", - "matting", - "register_matting_routes", -] diff --git a/python/matting/__init__.py b/python/matting/__init__.py new file mode 100644 index 0000000..9738f9e --- /dev/null +++ b/python/matting/__init__.py @@ -0,0 +1,85 @@ +"""Public compatibility facade for LayerForge background-removal support.""" + +import folder_paths +from aiohttp import web + +from .api import check_matting_model, matting, register_matting_routes +from .backends.birefnet import ( + _download_birefnet_checkpoint, + _ensure_birefnet_checkpoint, + _find_existing_birefnet_default_checkpoint, + _find_existing_birefnet_remote_checkpoint, + _find_local_birefnet_model, + _get_birefnet_remote_model, + _get_comfy_birefnet_loader, + _is_native_birefnet_checkpoint, + _iter_birefnet_checkpoint_paths, +) +from .backends.rmbg import ( + RMBG2Model, + _download_rmbg_model, + _ensure_rmbg_model, + _find_existing_rmbg_model, + _find_local_rmbg_model, + _get_rmbg_model_loader, + _get_rmbg_remote_model, + _is_rmbg_model_directory, +) +from .catalog import ( + _BIREFNET_DEFAULT_LOCAL_FILENAME, + _BIREFNET_FILENAME, + _BIREFNET_MODEL_CATALOG, + _BIREFNET_PROJECT_URL, + _BIREFNET_REMOTE_PREFIX, + _BIREFNET_REPOSITORY, + _BIREFNET_REQUIRED_KEYS, + _RMBG_MODEL_CATALOG, + _RMBG_REMOTE_PREFIX, +) +from .options import _get_birefnet_model_options +from .paths import ( + _get_birefnet_base_paths, + _get_birefnet_default_checkpoint_path, + _get_birefnet_download_dir, + _get_birefnet_remote_checkpoint_path, + _get_managed_birefnet_paths, + _get_rmbg_model_directory, + _get_rmbg_model_roots, +) +from .service import BiRefNetMatting + + +__all__ = [ + "BiRefNetMatting", + "RMBG2Model", + "check_matting_model", + "matting", + "register_matting_routes", + "_BIREFNET_DEFAULT_LOCAL_FILENAME", + "_BIREFNET_FILENAME", + "_BIREFNET_MODEL_CATALOG", + "_BIREFNET_PROJECT_URL", + "_BIREFNET_REMOTE_PREFIX", + "_BIREFNET_REPOSITORY", + "_BIREFNET_REQUIRED_KEYS", + "_RMBG_MODEL_CATALOG", + "_RMBG_REMOTE_PREFIX", + "_download_birefnet_checkpoint", + "_download_rmbg_model", + "_ensure_birefnet_checkpoint", + "_ensure_rmbg_model", + "_find_existing_birefnet_default_checkpoint", + "_find_existing_birefnet_remote_checkpoint", + "_find_existing_rmbg_model", + "_find_local_birefnet_model", + "_find_local_rmbg_model", + "_get_birefnet_model_options", + "_get_birefnet_remote_checkpoint_path", + "_get_birefnet_remote_model", + "_get_comfy_birefnet_loader", + "_get_rmbg_model_directory", + "_get_rmbg_model_loader", + "_get_rmbg_remote_model", + "_is_native_birefnet_checkpoint", + "_is_rmbg_model_directory", +] diff --git a/python/matting/api.py b/python/matting/api.py new file mode 100644 index 0000000..6486ed2 --- /dev/null +++ b/python/matting/api.py @@ -0,0 +1,298 @@ +"""HTTP endpoints for model status and background-removal execution.""" + +import traceback + +from aiohttp import web +from server import PromptServer + +from ..image_utils import convert_base64_to_tensor, convert_tensor_to_base64 +from ..node import log +from .backends.birefnet import ( + _find_existing_birefnet_remote_checkpoint, + _find_local_birefnet_model, + _get_birefnet_remote_model, + _get_comfy_birefnet_loader, +) +from .backends.rmbg import ( + _find_existing_rmbg_model, + _find_local_rmbg_model, + _get_rmbg_model_loader, + _get_rmbg_remote_model, +) +from .options import _get_birefnet_model_options +from .paths import _get_birefnet_base_paths +from .service import BiRefNetMatting + +_matting_lock = None + + +def _matting_status_response( + available, + reason, + message, + *, + status=200, + **details, +): + return web.json_response( + { + "available": available, + "reason": reason, + "message": message, + **details, + }, + status=status, + ) + + +async def check_matting_model(request): + """Report whether the selected background-removal model is ready.""" + try: + model_options = _get_birefnet_model_options() + requested_model = request.query.get("model_path") or "auto" + remote_birefnet_model = _get_birefnet_remote_model(requested_model) + remote_rmbg_model = _get_rmbg_remote_model(requested_model) + + local_rmbg_path = None + if not remote_rmbg_model and requested_model != "auto": + local_rmbg_path = _find_local_rmbg_model(requested_model) + + if remote_rmbg_model or local_rmbg_path: + if _get_rmbg_model_loader() is None: + return _matting_status_response( + False, + "unsupported_rmbg", + "BRIA RMBG 2.0 requires the 'transformers' package in the active ComfyUI environment.", + model_path=requested_model, + models=model_options, + ) + + if remote_rmbg_model: + local_model_path = _find_existing_rmbg_model(remote_rmbg_model) + if local_model_path: + log.info(f"Selected BRIA RMBG model is ready at {local_model_path}") + return _matting_status_response( + True, + "ready", + "Selected model is ready to use", + model_path=local_model_path, + selected_model=remote_rmbg_model["label"], + models=model_options, + ) + + log.info(f"Selected BRIA RMBG model is not downloaded: {remote_rmbg_model['label']}") + return _matting_status_response( + False, + "not_downloaded", + f"{remote_rmbg_model['label']} will be downloaded automatically on first use.", + model_path=requested_model, + selected_model=remote_rmbg_model["label"], + models=model_options, + ) + + return _matting_status_response( + True, + "ready", + "Selected model is ready to use", + model_path=local_rmbg_path, + selected_model=local_rmbg_path, + models=model_options, + ) + + if _get_comfy_birefnet_loader() is None: + return _matting_status_response( + False, + "unsupported_comfyui", + "This ComfyUI version does not provide the native BiRefNet background-removal loader.", + models=model_options, + ) + + if remote_birefnet_model: + local_model_path = _find_existing_birefnet_remote_checkpoint(remote_birefnet_model) + if local_model_path: + log.info(f"Selected BiRefNet model is ready at {local_model_path}") + return _matting_status_response( + True, + "ready", + "Selected model is ready to use", + model_path=local_model_path, + selected_model=remote_birefnet_model["label"], + models=model_options, + ) + + log.info(f"Selected BiRefNet model is not downloaded: {remote_birefnet_model['label']}") + return _matting_status_response( + False, + "not_downloaded", + f"{remote_birefnet_model['label']} will be downloaded automatically on first use.", + model_path=requested_model, + selected_model=remote_birefnet_model["label"], + models=model_options, + ) + + if requested_model != "auto": + local_model_path = _find_local_birefnet_model(requested_model) + if local_model_path: + log.info(f"Selected BiRefNet model is ready at {local_model_path}") + return _matting_status_response( + True, + "ready", + "Selected model is ready to use", + model_path=local_model_path, + selected_model=local_model_path, + models=model_options, + ) + + return _matting_status_response( + False, + "selected_model_unavailable", + "The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.", + model_path=requested_model, + models=model_options, + ) + + local_model_path = _find_local_birefnet_model() + if local_model_path: + log.info(f"BiRefNet model files detected at {local_model_path}") + return _matting_status_response( + True, + "ready", + "Model is ready to use", + model_path=local_model_path, + models=model_options, + ) + + searched_paths = _get_birefnet_base_paths() + log.info(f"BiRefNet model not found in any of: {searched_paths}") + return _matting_status_response( + False, + "not_downloaded", + "The BiRefNet checkpoint will be downloaded automatically on first use (requires internet connection).", + model_path=searched_paths[0] if searched_paths else None, + models=model_options, + ) + except Exception as error: + log.error(f"Error checking matting model: {error}") + return _matting_status_response( + False, + "error", + f"Error checking model status: {error}", + status=500, + ) + + +def _is_model_download_error(error): + error_text = str(error).lower() + return any( + marker in error_text + for marker in ( + "offline", + "connection", + "timed out", + "huggingface", + "localentrynotfound", + "gated", + "401", + "access approval", + "credentials", + "permission", + ) + ) + + +async def matting(request): + global _matting_lock + + if _matting_lock is not None: + log.warning("Matting already in progress, rejecting request") + return web.json_response( + { + "error": "Another matting operation is in progress", + "details": "Please wait for the current operation to complete", + }, + status=429, + ) + + _matting_lock = True + try: + log.info("Received matting request") + data = await request.json() + matting_instance = BiRefNetMatting() + image_tensor, original_alpha = convert_base64_to_tensor(data["image"]) + log.debug(f"Input image shape: {image_tensor.shape}") + + mode = data.get("mode", "remove_background") + model_path = data.get("model_path") or "auto" + matted_image, alpha_mask = matting_instance.execute( + image_tensor, + model_path, + threshold=data.get("threshold", 0.5), + refinement=data.get("refinement", 1), + mode=mode, + ) + + if mode == "mask_only": + result_image = convert_tensor_to_base64(alpha_mask) + else: + result_image = convert_tensor_to_base64(matted_image, alpha_mask, original_alpha) + result_mask = convert_tensor_to_base64(alpha_mask) + # Draw Mask uses white pixels as the area to remove, while matting uses + # white pixels as the area to keep. + draw_mask = 1.0 - alpha_mask + + return web.json_response( + { + "matted_image": result_image, + "alpha_mask": result_mask, + "draw_mask": convert_tensor_to_base64(draw_mask), + "mode": mode, + "model_path": matting_instance.model_path, + } + ) + except RuntimeError as error: + log.error(f"Runtime error during matting: {error}") + if _is_model_download_error(error): + return web.json_response( + { + "error": "Model Download Error", + "details": "Failed to download the selected background-removal model from Hugging Face. Check internet access, gated-model approval, and Hugging Face credentials.", + }, + status=400, + ) + return web.json_response( + { + "error": "Matting Model Error", + "details": str(error), + }, + status=500, + ) + except Exception as error: + log.exception(f"Error in matting endpoint: {error}") + if _is_model_download_error(error): + return web.json_response( + { + "error": "Network Connection Error", + "details": "Failed to download the selected background-removal model from Hugging Face. Check internet access, gated-model approval, and Hugging Face credentials.", + }, + status=400, + ) + + return web.json_response( + { + "error": "An unexpected error occurred", + "details": traceback.format_exc(), + }, + status=500, + ) + finally: + _matting_lock = None + log.debug("Matting lock released") + + +def register_matting_routes(): + """Register matting endpoints without import-time decorators.""" + PromptServer.instance.routes.get("/matting/check-model")(check_matting_model) + PromptServer.instance.routes.post("/matting")(matting) + + +__all__ = ["check_matting_model", "matting", "register_matting_routes"] diff --git a/python/matting/backends/__init__.py b/python/matting/backends/__init__.py new file mode 100644 index 0000000..95bf1ff --- /dev/null +++ b/python/matting/backends/__init__.py @@ -0,0 +1,2 @@ +"""Model-specific background-removal backends.""" +"""Background-removal backend implementations.""" diff --git a/python/matting/backends/birefnet.py b/python/matting/backends/birefnet.py new file mode 100644 index 0000000..cdc768d --- /dev/null +++ b/python/matting/backends/birefnet.py @@ -0,0 +1,237 @@ +"""BiRefNet backend using ComfyUI's native background-removal loader.""" + +import os +import shutil + +from ...node import log +from ..catalog import ( + _BIREFNET_DEFAULT_LOCAL_FILENAME, + _BIREFNET_FILENAME, + _BIREFNET_MODEL_CATALOG, + _BIREFNET_REPOSITORY, + _BIREFNET_REQUIRED_KEYS, + _BIREFNET_REMOTE_PREFIX, +) +from ..paths import ( + _get_birefnet_base_paths, + _get_birefnet_default_checkpoint_path, + _get_birefnet_download_dir, + _get_birefnet_remote_checkpoint_path, + _get_rmbg_model_roots, +) + + +def _get_comfy_birefnet_loader(): + """Return ComfyUI's native BiRefNet loader when it is available.""" + try: + from comfy.bg_removal_model import load + + return load + except Exception as error: + log.debug(f"Native ComfyUI BiRefNet loader is unavailable: {error}") + return None + + +def _is_native_birefnet_checkpoint(path): + """Check the checkpoint signature without loading all weights into memory.""" + if not os.path.isfile(path) or not path.lower().endswith(".safetensors"): + return False + + try: + from safetensors import safe_open + + with safe_open(path, framework="pt") as checkpoint: + return _BIREFNET_REQUIRED_KEYS.issubset(checkpoint.keys()) + except Exception as error: + log.debug(f"Unable to inspect BiRefNet checkpoint {path}: {error}") + return False + + +def _iter_birefnet_checkpoint_paths(): + """Yield candidate checkpoints from current ComfyUI model directories.""" + for base_path in _get_birefnet_base_paths(): + if not os.path.isdir(base_path): + continue + + for root, directories, files in os.walk(base_path): + directories[:] = [ + directory + for directory in directories + if directory not in { + ".git", + ".no_exist", + ".cache", + "__pycache__", + "refs", + "snapshots", + "blobs", + } + and not directory.startswith("models--") + ] + for filename in sorted(files): + if filename.lower().endswith(".safetensors"): + yield os.path.join(root, filename) + + +def _find_local_birefnet_model(model_path=None): + """Find a full BiRefNet checkpoint accepted by ComfyUI's native loader.""" + candidates = [] + seen = set() + managed_rmbg_roots = _get_rmbg_model_roots() + for path in _iter_birefnet_checkpoint_paths(): + normalized = os.path.normcase(os.path.normpath(path)) + if normalized in seen: + continue + seen.add(normalized) + normalized_absolute = os.path.normcase(os.path.normpath(os.path.abspath(path))) + if any( + normalized_absolute == root or normalized_absolute.startswith(root + os.sep) + for root in managed_rmbg_roots + ): + continue + if _is_native_birefnet_checkpoint(path): + candidates.append(path) + + if model_path and model_path != "auto": + requested = os.path.normcase(os.path.normpath(os.path.abspath(model_path))) + for candidate in candidates: + if os.path.normcase(os.path.normpath(os.path.abspath(candidate))) == requested: + return candidate + return None + + if not candidates: + return None + + priority = { + "birefnet.safetensors": 0, + "model.safetensors": 1, + "birefnet-general.safetensors": 2, + "birefnet-hr.safetensors": 3, + } + return min( + candidates, + key=lambda path: (priority.get(os.path.basename(path).lower(), 10), path.lower()), + ) + + +def _get_birefnet_remote_model(model_path): + """Resolve a client-facing remote model identifier from the fixed catalog.""" + if not isinstance(model_path, str) or not model_path.startswith(_BIREFNET_REMOTE_PREFIX): + return None + + model_id = model_path[len(_BIREFNET_REMOTE_PREFIX) :] + return next((model for model in _BIREFNET_MODEL_CATALOG if model["id"] == model_id), None) + + +def _find_existing_birefnet_remote_checkpoint(model): + """Return an installed managed remote checkpoint, when available.""" + checkpoint_path = _get_birefnet_remote_checkpoint_path(model) + if checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path): + return checkpoint_path + return None + + +def _find_existing_birefnet_default_checkpoint(): + """Return the automatic checkpoint when its current friendly filename exists.""" + checkpoint_path = _get_birefnet_default_checkpoint_path() + if checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path): + return checkpoint_path + return None + + +def _download_birefnet_checkpoint(model=None): + """Download and validate a full BiRefNet checkpoint into ComfyUI's model path.""" + try: + from huggingface_hub import hf_hub_download + except ImportError as error: + raise RuntimeError( + "Automatic BiRefNet download requires the 'huggingface_hub' package. " + "Install the LayerForge requirements or place a compatible checkpoint in " + "ComfyUI/models/background_removal/." + ) from error + + if model is None: + repository = _BIREFNET_REPOSITORY + filename = _BIREFNET_FILENAME + download_dir = _get_birefnet_download_dir() + model_label = "BiRefNet — General" + target_path = os.path.join(download_dir, _BIREFNET_DEFAULT_LOCAL_FILENAME) + else: + repository = model["repo_id"] + filename = model["filename"] + download_dir = _get_birefnet_download_dir() + model_label = model["label"] + target_path = _get_birefnet_remote_checkpoint_path(model) + + log.info(f"Downloading {model_label} from Hugging Face into {download_dir}...") + try: + downloaded_path = hf_hub_download( + repo_id=repository, + filename=filename, + local_dir=download_dir, + local_dir_use_symlinks=False, + ) + except TypeError: + downloaded_path = hf_hub_download( + repo_id=repository, + filename=filename, + local_dir=download_dir, + ) + + if not _is_native_birefnet_checkpoint(downloaded_path): + raise RuntimeError( + f"Downloaded {model_label} is not a ComfyUI-compatible BiRefNet checkpoint: {downloaded_path}" + ) + + if target_path and os.path.normcase(os.path.abspath(downloaded_path)) != os.path.normcase( + os.path.abspath(target_path) + ): + os.makedirs(os.path.dirname(target_path), exist_ok=True) + try: + os.replace(downloaded_path, target_path) + except OSError: + shutil.copy2(downloaded_path, target_path) + downloaded_path = target_path + + if target_path and not _is_native_birefnet_checkpoint(downloaded_path): + raise RuntimeError( + f"Renamed {model_label} is not a ComfyUI-compatible BiRefNet checkpoint: {downloaded_path}" + ) + + log.info(f"{model_label} checkpoint is ready at {downloaded_path}") + return downloaded_path + + +def _ensure_birefnet_checkpoint(model_path=None): + """Resolve or download the selected BiRefNet checkpoint.""" + remote_model = _get_birefnet_remote_model(model_path) + if remote_model: + checkpoint_path = _find_existing_birefnet_remote_checkpoint(remote_model) + if checkpoint_path: + return checkpoint_path + return _download_birefnet_checkpoint(remote_model) + + if not model_path or model_path == "auto": + _find_existing_birefnet_default_checkpoint() + + checkpoint_path = _find_local_birefnet_model(model_path) + if checkpoint_path: + return checkpoint_path + + if model_path and model_path != "auto": + raise RuntimeError("The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.") + + return _download_birefnet_checkpoint() + + +__all__ = [ + "_download_birefnet_checkpoint", + "_ensure_birefnet_checkpoint", + "_find_existing_birefnet_default_checkpoint", + "_find_existing_birefnet_remote_checkpoint", + "_find_local_birefnet_model", + "_get_comfy_birefnet_loader", + "_get_birefnet_remote_model", + "_is_native_birefnet_checkpoint", + "_iter_birefnet_checkpoint_paths", +] diff --git a/python/matting/backends/rmbg.py b/python/matting/backends/rmbg.py new file mode 100644 index 0000000..0f91470 --- /dev/null +++ b/python/matting/backends/rmbg.py @@ -0,0 +1,235 @@ +"""BRIA RMBG 2.0 backend using the local Transformers model.""" + +import os +import threading +from typing import ClassVar + +import torch +import torch.nn.functional as F + +from ...node import log +from ..catalog import _RMBG_MODEL_CATALOG, _RMBG_REMOTE_PREFIX +from ..paths import _get_rmbg_model_directory + + +def _get_rmbg_model_loader(): + """Return the Transformers loader used by BRIA RMBG 2.0, when installed.""" + try: + from transformers import AutoModelForImageSegmentation + + return AutoModelForImageSegmentation + except Exception as error: + log.debug(f"BRIA RMBG 2.0 Transformers loader is unavailable: {error}") + return None + + +def _is_rmbg_model_directory(path): + """Check whether a directory contains the files needed by RMBG 2.0.""" + if not path or not os.path.isdir(path): + return False + + required_files = ( + "config.json", + "preprocessor_config.json", + "birefnet.py", + "BiRefNet_config.py", + ) + has_weights = any( + os.path.isfile(os.path.join(path, filename)) + for filename in ("model.safetensors", "pytorch_model.bin") + ) + return has_weights and all(os.path.isfile(os.path.join(path, filename)) for filename in required_files) + + +def _get_rmbg_remote_model(model_path): + """Resolve a client-facing BRIA RMBG model identifier from the catalog.""" + if not isinstance(model_path, str) or not model_path.startswith(_RMBG_REMOTE_PREFIX): + return None + + model_id = model_path[len(_RMBG_REMOTE_PREFIX) :] + return next((model for model in _RMBG_MODEL_CATALOG if model["id"] == model_id), None) + + +def _find_existing_rmbg_model(model): + """Return an installed BRIA RMBG model directory, when available.""" + model_directory = _get_rmbg_model_directory(model) + if _is_rmbg_model_directory(model_directory): + return model_directory + return None + + +def _find_local_rmbg_model(model_path=None): + """Find a local BRIA RMBG model directory, optionally matching a selection.""" + if model_path and model_path != "auto": + requested = os.path.abspath(model_path) + return requested if _is_rmbg_model_directory(requested) else None + + for model in _RMBG_MODEL_CATALOG: + model_directory = _find_existing_rmbg_model(model) + if model_directory: + return model_directory + return None + + +def _download_rmbg_model(model): + """Download the BRIA RMBG 2.0 repository into ComfyUI's model path.""" + try: + from huggingface_hub import snapshot_download + except ImportError as error: + raise RuntimeError( + "BRIA RMBG 2.0 download requires the 'huggingface_hub' package. " + "Install the LayerForge requirements or place the model files in " + "ComfyUI/models/background_removal/RMBG-2.0/." + ) from error + + model_directory = _get_rmbg_model_directory(model) + if not model_directory: + raise RuntimeError("ComfyUI did not expose a background_removal model directory") + + os.makedirs(model_directory, exist_ok=True) + log.info(f"Downloading {model['label']} from Hugging Face into {model_directory}...") + try: + snapshot_download( + repo_id=model["repo_id"], + local_dir=model_directory, + allow_patterns=["*.json", "*.py", "*.safetensors"], + ) + except Exception as error: + raise RuntimeError( + "Unable to download BRIA RMBG 2.0. Accept the model's gated access on " + "Hugging Face and ensure the configured Hugging Face credentials can access it." + ) from error + + if not _is_rmbg_model_directory(model_directory): + raise RuntimeError( + f"Downloaded {model['label']} is incomplete or is not a valid Transformers model: " + f"{model_directory}" + ) + + log.info(f"{model['label']} model is ready at {model_directory}") + return model_directory + + +def _ensure_rmbg_model(model_path): + """Resolve or download the selected BRIA RMBG 2.0 model directory.""" + remote_model = _get_rmbg_remote_model(model_path) + if remote_model: + model_directory = _find_existing_rmbg_model(remote_model) + if model_directory: + return model_directory + return _download_rmbg_model(remote_model) + + model_directory = _find_local_rmbg_model(model_path) + if model_directory: + return model_directory + + raise RuntimeError( + "The selected BRIA RMBG 2.0 model is not available or is incomplete. " + "Place its files in ComfyUI/models/background_removal/RMBG-2.0/ or download it from the model selector." + ) + + +class RMBG2Model: + """Adapter for BRIA RMBG 2.0 loaded through Transformers.""" + + _model_cache: ClassVar[dict] = {} + _model_cache_lock = threading.Lock() + _image_size = 1024 + + def __init__(self, model, device): + self.model = model + self.device = device + + @classmethod + def load(cls, model_directory): + loader = _get_rmbg_model_loader() + if loader is None: + raise RuntimeError( + "BRIA RMBG 2.0 requires the 'transformers' package. " + "Install the LayerForge requirements in the ComfyUI environment." + ) + + cache_key = os.path.normcase(os.path.normpath(os.path.abspath(model_directory))) + with cls._model_cache_lock: + if cache_key not in cls._model_cache: + device = _get_rmbg_device() + log.info(f"Loading BRIA RMBG 2.0 from {model_directory} on {device}") + model = loader.from_pretrained( + model_directory, + trust_remote_code=True, + local_files_only=True, + ) + model.eval().to(device) + cls._model_cache[cache_key] = cls(model, device) + else: + log.debug(f"Using cached BRIA RMBG 2.0 model from {model_directory}") + + return cls._model_cache[cache_key] + + def encode_image(self, image): + if image.dim() != 4 or image.shape[-1] not in (1, 3, 4): + raise ValueError(f"Expected a BHWC image tensor, got {tuple(image.shape)}") + + original_size = (image.shape[1], image.shape[2]) + input_image = image.movedim(-1, 1) + if input_image.shape[1] == 1: + input_image = input_image.expand(-1, 3, -1, -1) + elif input_image.shape[1] == 4: + input_image = input_image[:, :3] + + input_image = input_image.to(device=self.device, dtype=torch.float32) + input_image = F.interpolate( + input_image, + size=(self._image_size, self._image_size), + mode="bilinear", + align_corners=False, + ) + mean = input_image.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = input_image.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + + with torch.no_grad(): + prediction = self.model((input_image - mean) / std) + if isinstance(prediction, (tuple, list)): + prediction = prediction[-1] + elif hasattr(prediction, "logits"): + prediction = prediction.logits + + if not isinstance(prediction, torch.Tensor): + raise ValueError("BRIA RMBG 2.0 returned an unsupported prediction type") + if prediction.dim() == 3: + prediction = prediction.unsqueeze(1) + if prediction.dim() != 4: + raise ValueError(f"Unexpected BRIA RMBG 2.0 output shape: {tuple(prediction.shape)}") + + alpha_mask = prediction.sigmoid() + alpha_mask = F.interpolate( + alpha_mask, + size=original_size, + mode="bilinear", + align_corners=False, + ) + + return alpha_mask[:, 0].clamp(0.0, 1.0).to(device=image.device, dtype=torch.float32) + + +def _get_rmbg_device(): + """Use ComfyUI's active torch device when available.""" + try: + import comfy.model_management as model_management + + return model_management.get_torch_device() + except Exception as error: + log.debug(f"Unable to resolve ComfyUI's model device for RMBG 2.0: {error}") + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +__all__ = [ + "RMBG2Model", + "_download_rmbg_model", + "_ensure_rmbg_model", + "_find_existing_rmbg_model", + "_find_local_rmbg_model", + "_get_rmbg_model_loader", + "_get_rmbg_remote_model", + "_is_rmbg_model_directory", +] diff --git a/python/matting/catalog.py b/python/matting/catalog.py new file mode 100644 index 0000000..71e519b --- /dev/null +++ b/python/matting/catalog.py @@ -0,0 +1,124 @@ +"""Catalogs and identifiers for LayerForge background-removal models.""" + +_BIREFNET_REPOSITORY = "ZhengPeng7/BiRefNet" +_BIREFNET_FILENAME = "model.safetensors" +_BIREFNET_DEFAULT_LOCAL_FILENAME = "BiRefNet-general.safetensors" +_BIREFNET_PROJECT_URL = "https://github.com/ZhengPeng7/BiRefNet" +_BIREFNET_REMOTE_PREFIX = "remote:" +_BIREFNET_REQUIRED_KEYS = { + "bb.layers.1.blocks.0.attn.relative_position_index", + "bb.layers.2.blocks.17.attn.qkv.weight", +} + +_RMBG_REMOTE_PREFIX = "remote:" +_RMBG_MODEL_CATALOG = ( + { + "id": "rmbg_2_0", + "label": "BRIA RMBG 2.0", + "description": ( + "Local BRIA background removal with strong general-purpose segmentation. " + "Hugging Face access approval is required before the gated weights can be downloaded." + ), + "repo_id": "briaai/RMBG-2.0", + "local_directory": "RMBG-2.0", + "url": "https://huggingface.co/briaai/RMBG-2.0", + "project_url": "https://github.com/Bria-AI/RMBG-2.0", + "backend": "rmbg", + }, +) + +_BIREFNET_MODEL_CATALOG = ( + { + "id": "general", + "label": "BiRefNet — General", + "description": "The best starting point for everyday images and general background removal.", + "repo_id": "ZhengPeng7/BiRefNet", + "filename": "model.safetensors", + "local_filename": "BiRefNet-general.safetensors", + }, + { + "id": "high_resolution", + "label": "BiRefNet — High Resolution", + "description": "High-resolution segmentation for detailed edges and larger source images; it uses more memory.", + "repo_id": "ZhengPeng7/BiRefNet_HR", + "filename": "model.safetensors", + "local_filename": "BiRefNet-HR.safetensors", + }, + { + "id": "portrait", + "label": "BiRefNet — Portrait", + "description": "Portrait matting for people, hair, and portrait-focused cutouts.", + "repo_id": "ZhengPeng7/BiRefNet-portrait", + "filename": "model.safetensors", + "local_filename": "BiRefNet-portrait.safetensors", + }, + { + "id": "matting", + "label": "BiRefNet — Matting", + "description": "General matting with a focus on soft alpha edges such as hair and semi-transparent details.", + "repo_id": "ZhengPeng7/BiRefNet-matting", + "filename": "model.safetensors", + "local_filename": "BiRefNet-matting.safetensors", + }, + { + "id": "high_resolution_matting", + "label": "BiRefNet — High Resolution Matting", + "description": "High-resolution general matting for fine details; it requires more memory than the standard matting model.", + "repo_id": "ZhengPeng7/BiRefNet_HR-matting", + "filename": "model.safetensors", + "local_filename": "BiRefNet-HR-matting.safetensors", + }, + { + "id": "dynamic", + "label": "BiRefNet — Dynamic", + "description": "Dynamic-shape segmentation for inputs with varying aspect ratios and resolutions.", + "repo_id": "ZhengPeng7/BiRefNet_dynamic", + "filename": "model.safetensors", + "local_filename": "BiRefNet_dynamic.safetensors", + }, + { + "id": "dynamic_matting", + "label": "BiRefNet — Dynamic Matting", + "description": "Dynamic-shape matting for arbitrary input sizes, with a focus on soft alpha edges.", + "repo_id": "ZhengPeng7/BiRefNet_dynamic-matting", + "filename": "model.safetensors", + "local_filename": "BiRefNet_dynamic-matting.safetensors", + }, + { + "id": "hrsod", + "label": "BiRefNet — HRSOD", + "description": "High-resolution salient-object detection; useful when the main subject should stand out from its surroundings.", + "repo_id": "ZhengPeng7/BiRefNet-HRSOD", + "filename": "model.safetensors", + "local_filename": "BiRefNet-HRSOD.safetensors", + }, + { + "id": "dis5k", + "label": "BiRefNet — DIS5K", + "description": "Dichotomous image segmentation trained for clean foreground/background separation.", + "repo_id": "ZhengPeng7/BiRefNet-DIS5K", + "filename": "model.safetensors", + "local_filename": "BiRefNet-DIS5K.safetensors", + }, + { + "id": "cod", + "label": "BiRefNet — COD", + "description": "Camouflaged-object detection; use it for subjects that blend into the background.", + "repo_id": "ZhengPeng7/BiRefNet-COD", + "filename": "model.safetensors", + "local_filename": "BiRefNet-COD.safetensors", + }, +) + + +__all__ = [ + "_BIREFNET_DEFAULT_LOCAL_FILENAME", + "_BIREFNET_FILENAME", + "_BIREFNET_MODEL_CATALOG", + "_BIREFNET_PROJECT_URL", + "_BIREFNET_REMOTE_PREFIX", + "_BIREFNET_REPOSITORY", + "_BIREFNET_REQUIRED_KEYS", + "_RMBG_MODEL_CATALOG", + "_RMBG_REMOTE_PREFIX", +] diff --git a/python/matting/options.py b/python/matting/options.py new file mode 100644 index 0000000..01f34bf --- /dev/null +++ b/python/matting/options.py @@ -0,0 +1,115 @@ +"""Model-selector options assembled from local files and supported catalogs.""" + +import os + +from .backends.birefnet import ( + _find_existing_birefnet_default_checkpoint, + _find_existing_birefnet_remote_checkpoint, + _is_native_birefnet_checkpoint, + _iter_birefnet_checkpoint_paths, +) +from .backends.rmbg import _find_existing_rmbg_model +from .catalog import _BIREFNET_MODEL_CATALOG, _BIREFNET_PROJECT_URL, _BIREFNET_REMOTE_PREFIX +from .catalog import _RMBG_MODEL_CATALOG, _RMBG_REMOTE_PREFIX +from .paths import _get_birefnet_base_paths, _get_managed_birefnet_paths, _get_rmbg_model_roots + + +def _is_path_under_root(path, root): + normalized_path = os.path.normcase(os.path.normpath(os.path.abspath(path))) + return normalized_path == root or normalized_path.startswith(root + os.sep) + + +def _get_birefnet_model_options(): + """Return local models and official BiRefNet/RMBG models for the selector.""" + _find_existing_birefnet_default_checkpoint() + + local_options = [] + seen = set() + base_paths = _get_birefnet_base_paths() + managed_paths = _get_managed_birefnet_paths(base_paths) + managed_rmbg_roots = _get_rmbg_model_roots(base_paths) + + for path in _iter_birefnet_checkpoint_paths(): + normalized = os.path.normcase(os.path.normpath(os.path.abspath(path))) + if normalized in seen or not _is_native_birefnet_checkpoint(path): + continue + if normalized in managed_paths or any( + _is_path_under_root(path, root) for root in managed_rmbg_roots + ): + continue + + seen.add(normalized) + label = os.path.basename(path) + for base_path in base_paths: + try: + relative_path = os.path.relpath(path, base_path) + except ValueError: + continue + if relative_path == os.pardir or relative_path.startswith(os.pardir + os.sep): + continue + label = relative_path.replace(os.sep, "/") + break + + local_options.append( + { + "path": path, + "label": label, + "source": "local", + "backend": "birefnet", + "downloaded": True, + } + ) + + for model in _RMBG_MODEL_CATALOG: + model_directory = _find_existing_rmbg_model(model) + if not model_directory: + continue + local_options.append( + { + "path": model_directory, + "label": model["label"], + "description": model["description"], + "url": model["url"], + "project_url": model["project_url"], + "source": "local", + "backend": model["backend"], + "downloaded": True, + } + ) + + remote_options = [] + for model in _BIREFNET_MODEL_CATALOG: + checkpoint_path = _find_existing_birefnet_remote_checkpoint(model) + remote_options.append( + { + "path": f"{_BIREFNET_REMOTE_PREFIX}{model['id']}", + "label": model["label"], + "description": model["description"], + "url": f"https://huggingface.co/{model['repo_id']}", + "project_url": _BIREFNET_PROJECT_URL, + "source": "remote", + "backend": "birefnet", + "downloaded": bool(checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path)), + } + ) + + for model in _RMBG_MODEL_CATALOG: + model_directory = _find_existing_rmbg_model(model) + remote_options.append( + { + "path": f"{_RMBG_REMOTE_PREFIX}{model['id']}", + "label": model["label"], + "description": model["description"], + "url": model["url"], + "project_url": model["project_url"], + "source": "remote", + "backend": model["backend"], + "downloaded": bool(model_directory), + } + ) + + local_options.sort(key=lambda option: option["label"].lower()) + return local_options + remote_options + + +__all__ = ["_get_birefnet_model_options"] diff --git a/python/matting/paths.py b/python/matting/paths.py new file mode 100644 index 0000000..cd6fef9 --- /dev/null +++ b/python/matting/paths.py @@ -0,0 +1,100 @@ +"""Model-directory and checkpoint-path helpers shared by matting backends.""" + +import os + +import folder_paths + +from .catalog import _BIREFNET_DEFAULT_LOCAL_FILENAME, _BIREFNET_MODEL_CATALOG +from .catalog import _RMBG_MODEL_CATALOG + + +def _get_birefnet_base_paths(): + """Return the current ComfyUI background-removal model locations.""" + paths = [] + get_folder_paths = getattr(folder_paths, "get_folder_paths", None) + if callable(get_folder_paths): + try: + paths.extend(get_folder_paths("background_removal")) + except (KeyError, TypeError): + pass + + comfy_models_dir = getattr(folder_paths, "models_dir", None) + if comfy_models_dir: + paths.append(os.path.join(comfy_models_dir, "background_removal")) + + unique_paths = [] + seen = set() + for path in paths: + normalized = os.path.normcase(os.path.normpath(path)) + if normalized not in seen: + seen.add(normalized) + unique_paths.append(path) + return unique_paths + + +def _get_birefnet_remote_checkpoint_path(model): + """Return the managed local path for a catalog BiRefNet model.""" + base_paths = _get_birefnet_base_paths() + if not base_paths: + return None + return os.path.join(base_paths[0], model["local_filename"]) + + +def _get_rmbg_model_directory(model): + """Return the managed directory containing a local BRIA RMBG model.""" + base_paths = _get_birefnet_base_paths() + if not base_paths: + return None + return os.path.join(base_paths[0], model["local_directory"]) + + +def _get_rmbg_model_roots(base_paths=None): + """Return normalized roots reserved for managed BRIA RMBG model files.""" + base_paths = _get_birefnet_base_paths() if base_paths is None else base_paths + return { + os.path.normcase(os.path.normpath(os.path.abspath(path))) + for base_path in base_paths + for model in _RMBG_MODEL_CATALOG + for path in [os.path.join(base_path, model["local_directory"])] + } + + +def _get_birefnet_download_dir(): + """Return and create the shared ComfyUI background-removal directory.""" + paths = _get_birefnet_base_paths() + if not paths: + raise RuntimeError("ComfyUI did not expose a background_removal model directory") + + download_dir = paths[0] + os.makedirs(download_dir, exist_ok=True) + return download_dir + + +def _get_birefnet_default_checkpoint_path(): + """Return the friendly path used for the automatic General checkpoint.""" + base_paths = _get_birefnet_base_paths() + if not base_paths: + return None + return os.path.join(base_paths[0], _BIREFNET_DEFAULT_LOCAL_FILENAME) + + +def _get_managed_birefnet_paths(base_paths=None): + """Return normalized paths reserved for catalog BiRefNet checkpoints.""" + base_paths = _get_birefnet_base_paths() if base_paths is None else base_paths + return { + os.path.normcase(os.path.normpath(os.path.abspath(path))) + for base_path in base_paths + for model in _BIREFNET_MODEL_CATALOG + for path in [os.path.join(base_path, model["local_filename"])] + } + + +__all__ = [ + "_get_birefnet_base_paths", + "_get_birefnet_default_checkpoint_path", + "_get_birefnet_download_dir", + "_get_birefnet_remote_checkpoint_path", + "_get_managed_birefnet_paths", + "_get_rmbg_model_directory", + "_get_rmbg_model_roots", +] diff --git a/python/matting/service.py b/python/matting/service.py new file mode 100644 index 0000000..c9cf049 --- /dev/null +++ b/python/matting/service.py @@ -0,0 +1,160 @@ +"""Shared model selection, preprocessing, inference, and output handling.""" + +import hashlib +import threading +from typing import ClassVar + +import torch +import torch.nn.functional as F +from server import PromptServer +from torchvision import transforms + +from ..node import log +from .backends.birefnet import _ensure_birefnet_checkpoint, _get_comfy_birefnet_loader +from .backends.rmbg import ( + RMBG2Model, + _ensure_rmbg_model, + _find_local_rmbg_model, + _get_rmbg_model_loader, + _get_rmbg_remote_model, +) + + +class BiRefNetMatting: + """Background-removal adapter for native BiRefNet and BRIA RMBG 2.0.""" + + _model_cache: ClassVar[dict] = {} + _model_cache_lock = threading.Lock() + + def __init__(self): + self.model = None + self.model_path = None + + def load_model(self, model_path=None): + """Load the selected backend and reuse its process-wide model cache.""" + rmbg_model = _get_rmbg_remote_model(model_path) + local_rmbg_path = None + if not rmbg_model and model_path and model_path != "auto": + local_rmbg_path = _find_local_rmbg_model(model_path) + + if rmbg_model or local_rmbg_path: + if _get_rmbg_model_loader() is None: + raise RuntimeError( + "BRIA RMBG 2.0 requires the 'transformers' package. " + "Install the LayerForge requirements in the active ComfyUI environment." + ) + model_directory = _ensure_rmbg_model(model_path) + self.model = RMBG2Model.load(model_directory) + self.model_path = model_directory + return self.model + + loader = _get_comfy_birefnet_loader() + if loader is None: + raise RuntimeError( + "This ComfyUI version does not provide the native BiRefNet background-removal loader." + ) + + checkpoint_path = _ensure_birefnet_checkpoint(model_path) + with self._model_cache_lock: + if checkpoint_path not in self._model_cache: + log.info(f"Loading BiRefNet with ComfyUI's native loader from {checkpoint_path}") + model = loader(checkpoint_path) + if model is None: + raise RuntimeError( + f"ComfyUI did not recognize the BiRefNet checkpoint: {checkpoint_path}" + ) + self._model_cache[checkpoint_path] = model + else: + log.debug(f"Using cached native BiRefNet model from {checkpoint_path}") + + self.model = self._model_cache[checkpoint_path] + self.model_path = checkpoint_path + + return self.model + + def preprocess_image(self, image): + if not isinstance(image, torch.Tensor): + image = transforms.ToTensor()(image).unsqueeze(0) + + if image.dim() == 3: + if image.shape[0] in (1, 3, 4): + image = image.movedim(0, -1).unsqueeze(0) + else: + image = image.unsqueeze(0) + elif image.dim() == 4 and image.shape[1] in (1, 3, 4): + image = image.movedim(1, -1) + + if image.dim() != 4 or image.shape[-1] not in (1, 3, 4): + raise ValueError(f"Expected an image tensor in BCHW or BHWC format, got {tuple(image.shape)}") + + if image.shape[-1] == 1: + image = image.expand(-1, -1, -1, 3) + elif image.shape[-1] == 4: + image = image[..., :3] + + return image.to(dtype=torch.float32).contiguous() + + def execute(self, image, model_path, threshold=0.5, refinement=1, mode="remove_background"): + try: + PromptServer.instance.send_sync("matting_status", {"status": "processing"}) + del refinement + if mode not in {"remove_background", "remove_foreground", "mask_only"}: + raise ValueError(f"Unsupported matting mode: {mode}") + + threshold = float(threshold) + if not 0.0 <= threshold <= 1.0: + raise ValueError("Matting threshold must be between 0 and 1") + + image_tensor = self.preprocess_image(image) + original_size = (image_tensor.shape[1], image_tensor.shape[2]) + log.debug(f"Original size: {original_size}") + self.load_model(model_path) + log.debug(f"Processed image shape: {image_tensor.shape}, dtype: {image_tensor.dtype}") + + with torch.no_grad(): + result = self.model.encode_image(image_tensor) + if result.dim() == 3: + result = result.unsqueeze(1) + elif result.dim() == 2: + result = result.unsqueeze(0).unsqueeze(0) + else: + raise ValueError(f"Unexpected background-removal output shape: {tuple(result.shape)}") + + result = result.to(device=image_tensor.device, dtype=torch.float32) + if result.shape[-2:] != original_size: + result = F.interpolate( + result, + size=original_size, + mode="bilinear", + align_corners=False, + ) + result = result.clamp(0.0, 1.0) + log.debug(f"Background-removal output shape: {result.shape}, dtype: {result.dtype}") + + if threshold > 0: + result = (result > threshold).to(dtype=torch.float32) + + alpha_mask = 1.0 - result if mode == "remove_foreground" else result + if mode == "mask_only": + masked_image = alpha_mask.expand(-1, 3, -1, -1) + else: + masked_image = image_tensor.movedim(-1, 1) * alpha_mask + + PromptServer.instance.send_sync("matting_status", {"status": "completed"}) + return masked_image, alpha_mask + except Exception: + PromptServer.instance.send_sync("matting_status", {"status": "error"}) + raise + + @classmethod + def IS_CHANGED(cls, image, model_path, threshold, refinement, mode="remove_background"): + digest = hashlib.md5() + digest.update(str(image).encode()) + digest.update(str(model_path).encode()) + digest.update(str(threshold).encode()) + digest.update(str(refinement).encode()) + digest.update(str(mode).encode()) + return digest.hexdigest() + + +__all__ = ["BiRefNetMatting"] diff --git a/requirements.txt b/requirements.txt index aa25c8b..494123d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,5 @@ aiohttp numpy tqdm Pillow +transformers>=4.50.3 +kornia diff --git a/src/CanvasView.ts b/src/CanvasView.ts index dfb9d52..3db8395 100644 --- a/src/CanvasView.ts +++ b/src/CanvasView.ts @@ -43,6 +43,7 @@ interface MattingModelOption { url?: string; project_url?: string; source?: 'local' | 'remote'; + backend?: 'birefnet' | 'rmbg'; downloaded?: boolean; } @@ -396,12 +397,19 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): return; } - modelDescription.textContent = selectedOption.description || 'Installed checkpoint validated by ComfyUI\'s native BiRefNet loader.'; + modelDescription.textContent = selectedOption.description || ( + selectedOption.backend === 'rmbg' + ? 'Local BRIA RMBG 2.0 model loaded through Transformers.' + : 'Installed checkpoint validated by ComfyUI\'s native BiRefNet loader.' + ); if (selectedOption.url) { modelLinks.appendChild(createModelLink('Model page', selectedOption.url)); } if (selectedOption.project_url) { - modelLinks.appendChild(createModelLink('BiRefNet project', selectedOption.project_url)); + modelLinks.appendChild(createModelLink( + selectedOption.backend === 'rmbg' ? 'BRIA project' : 'BiRefNet project', + selectedOption.project_url, + )); } modelDetails.hidden = false; }; @@ -451,7 +459,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): modelStatus.textContent = `${modelCounts}. ${modelStatusMessage}`; body.append( - createRow('Model', modelSelect, 'Only checkpoints accepted by the native ComfyUI BiRefNet loader are listed.'), + createRow('Model', modelSelect, 'Choose a local BiRefNet checkpoint or BRIA RMBG 2.0, or download an official model on first use.'), modelDetails, createRow('Processing mode', modeSelect, 'The selected mode controls what the Matting button creates from the detected mask.'), createRow('Mask threshold', thresholdContainer, 'Set to 0 for a soft alpha mask; higher values create a harder cutout.'), @@ -820,13 +828,13 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): return; case 'not_downloaded': - showWarningNotification("The selected BiRefNet model will be downloaded automatically (requires internet connection).", 5000); + showWarningNotification(modelStatus.message || "The selected background-removal model will be downloaded automatically.", 7000); // Ask user if they want to proceed with download - if (!confirm("The selected BiRefNet model needs to be downloaded. This is a one-time download and may be large. Do you want to proceed?")) { + if (!confirm(`${modelStatus.message || "The selected background-removal model needs to be downloaded."}\n\nThis is a one-time download and may be large. Do you want to proceed?`)) { return; } - showInfoNotification("Downloading the selected BiRefNet model... This may take a few minutes.", 10000); + showInfoNotification("Downloading the selected background-removal model... This may take a few minutes.", 10000); break; case 'selected_model_unavailable': diff --git a/tests/test_entrypoint_integration.py b/tests/test_entrypoint_integration.py index ffd51ce..39fed8a 100644 --- a/tests/test_entrypoint_integration.py +++ b/tests/test_entrypoint_integration.py @@ -113,6 +113,10 @@ def _import_layerforge(monkeypatch, tmp_path): node=sys.modules[f"{package_name}.python.node"], image_utils=sys.modules[f"{package_name}.python.image_utils"], matting=sys.modules[f"{package_name}.python.matting"], + matting_api=sys.modules[f"{package_name}.python.matting.api"], + matting_birefnet=sys.modules[f"{package_name}.python.matting.backends.birefnet"], + matting_rmbg=sys.modules[f"{package_name}.python.matting.backends.rmbg"], + matting_service=sys.modules[f"{package_name}.python.matting.service"], routes=routes, package_name=package_name, ) @@ -334,6 +338,7 @@ def test_matting_adapter_uses_native_loader_and_bhwc_input(layerforge_runtime, m import torch node_module = layerforge_runtime.matting + service_module = layerforge_runtime.matting_service calls = {} class NativeBiRefNet: @@ -343,12 +348,12 @@ def test_matting_adapter_uses_native_loader_and_bhwc_input(layerforge_runtime, m return torch.full((image.shape[0], image.shape[1], image.shape[2]), 0.75) monkeypatch.setattr( - node_module, + service_module, "_get_comfy_birefnet_loader", lambda: lambda path: NativeBiRefNet(), ) monkeypatch.setattr( - node_module, + service_module, "_ensure_birefnet_checkpoint", lambda model_path=None: "native-birefnet.safetensors", ) @@ -369,24 +374,26 @@ def test_matting_adapter_uses_native_loader_and_bhwc_input(layerforge_runtime, m def test_matting_model_status_preserves_unsupported_and_error_responses(layerforge_runtime, monkeypatch): node_module = layerforge_runtime.matting + api_module = layerforge_runtime.matting_api responses = _record_matting_responses(monkeypatch, node_module) - monkeypatch.setattr(node_module, "_get_comfy_birefnet_loader", lambda: None) + monkeypatch.setattr(api_module, "_get_comfy_birefnet_loader", lambda: None) unsupported = _run_matting_model_check(node_module) - assert unsupported.payload == { - "available": False, - "reason": "unsupported_comfyui", - "message": "This ComfyUI version does not provide the native BiRefNet background-removal loader.", - } + assert unsupported.payload["available"] is False + assert unsupported.payload["reason"] == "unsupported_comfyui" + assert unsupported.payload["message"] == ( + "This ComfyUI version does not provide the native BiRefNet background-removal loader." + ) + assert unsupported.payload["models"] assert unsupported.status == 200 - monkeypatch.setattr(node_module, "_get_comfy_birefnet_loader", lambda: object()) + monkeypatch.setattr(api_module, "_get_comfy_birefnet_loader", lambda: object()) def fail_to_read_options(): raise RuntimeError("catalog unavailable") - monkeypatch.setattr(node_module, "_get_birefnet_model_options", fail_to_read_options) + monkeypatch.setattr(api_module, "_get_birefnet_model_options", fail_to_read_options) failed = _run_matting_model_check(node_module) assert failed.payload == { @@ -400,15 +407,16 @@ def test_matting_model_status_preserves_unsupported_and_error_responses(layerfor def test_matting_model_status_preserves_remote_model_responses(layerforge_runtime, monkeypatch): node_module = layerforge_runtime.matting + api_module = layerforge_runtime.matting_api responses = _record_matting_responses(monkeypatch, node_module) model_options = [{"path": "remote:portrait", "label": "Portrait"}] remote_model = {"label": "Portrait"} - monkeypatch.setattr(node_module, "_get_comfy_birefnet_loader", lambda: object()) - monkeypatch.setattr(node_module, "_get_birefnet_model_options", lambda: model_options) - monkeypatch.setattr(node_module, "_get_birefnet_remote_model", lambda path: remote_model) + monkeypatch.setattr(api_module, "_get_comfy_birefnet_loader", lambda: object()) + monkeypatch.setattr(api_module, "_get_birefnet_model_options", lambda: model_options) + monkeypatch.setattr(api_module, "_get_birefnet_remote_model", lambda path: remote_model) monkeypatch.setattr( - node_module, + api_module, "_find_existing_birefnet_remote_checkpoint", lambda model: "/models/portrait.safetensors", ) @@ -425,7 +433,7 @@ def test_matting_model_status_preserves_remote_model_responses(layerforge_runtim } assert ready.status == 200 - monkeypatch.setattr(node_module, "_find_existing_birefnet_remote_checkpoint", lambda model: None) + monkeypatch.setattr(api_module, "_find_existing_birefnet_remote_checkpoint", lambda model: None) missing = _run_matting_model_check(node_module, "remote:portrait") assert missing.payload == { @@ -442,14 +450,15 @@ def test_matting_model_status_preserves_remote_model_responses(layerforge_runtim def test_matting_model_status_preserves_local_and_automatic_responses(layerforge_runtime, monkeypatch): node_module = layerforge_runtime.matting + api_module = layerforge_runtime.matting_api responses = _record_matting_responses(monkeypatch, node_module) model_options = [{"path": "/models/custom.safetensors", "label": "custom.safetensors"}] - monkeypatch.setattr(node_module, "_get_comfy_birefnet_loader", lambda: object()) - monkeypatch.setattr(node_module, "_get_birefnet_model_options", lambda: model_options) - monkeypatch.setattr(node_module, "_get_birefnet_remote_model", lambda path: None) + monkeypatch.setattr(api_module, "_get_comfy_birefnet_loader", lambda: object()) + monkeypatch.setattr(api_module, "_get_birefnet_model_options", lambda: model_options) + monkeypatch.setattr(api_module, "_get_birefnet_remote_model", lambda path: None) monkeypatch.setattr( - node_module, + api_module, "_find_local_birefnet_model", lambda model_path=None: "/models/custom.safetensors", ) @@ -466,7 +475,7 @@ def test_matting_model_status_preserves_local_and_automatic_responses(layerforge } assert selected.status == 200 - monkeypatch.setattr(node_module, "_find_local_birefnet_model", lambda model_path=None: None) + monkeypatch.setattr(api_module, "_find_local_birefnet_model", lambda model_path=None: None) unavailable = _run_matting_model_check(node_module, "/models/missing.safetensors") assert unavailable.payload == { @@ -479,7 +488,7 @@ def test_matting_model_status_preserves_local_and_automatic_responses(layerforge assert unavailable.status == 200 monkeypatch.setattr( - node_module, + api_module, "_find_local_birefnet_model", lambda model_path=None: "/models/automatic.safetensors", ) @@ -494,9 +503,9 @@ def test_matting_model_status_preserves_local_and_automatic_responses(layerforge } assert automatic.status == 200 - monkeypatch.setattr(node_module, "_find_local_birefnet_model", lambda model_path=None: None) + monkeypatch.setattr(api_module, "_find_local_birefnet_model", lambda model_path=None: None) monkeypatch.setattr( - node_module, + api_module, "_get_birefnet_base_paths", lambda: ["/models/background_removal"], ) @@ -517,18 +526,19 @@ def test_matting_adapter_supports_inverted_and_mask_only_modes(layerforge_runtim import torch node_module = layerforge_runtime.matting + service_module = layerforge_runtime.matting_service class NativeBiRefNet: def encode_image(self, image): return torch.full((image.shape[0], image.shape[1], image.shape[2]), 0.75) monkeypatch.setattr( - node_module, + service_module, "_get_comfy_birefnet_loader", lambda: lambda path: NativeBiRefNet(), ) monkeypatch.setattr( - node_module, + service_module, "_ensure_birefnet_checkpoint", lambda model_path=None: "native-birefnet.safetensors", ) @@ -560,8 +570,9 @@ def test_matting_adapter_supports_inverted_and_mask_only_modes(layerforge_runtim def test_matting_model_options_include_downloadable_official_variants(layerforge_runtime, monkeypatch): node_module = layerforge_runtime.matting + birefnet_module = layerforge_runtime.matting_birefnet - monkeypatch.setattr(node_module, "_iter_birefnet_checkpoint_paths", lambda: iter(())) + monkeypatch.setattr(birefnet_module, "_iter_birefnet_checkpoint_paths", lambda: iter(())) options = node_module._get_birefnet_model_options() remote_options = [option for option in options if option["source"] == "remote"] @@ -571,28 +582,35 @@ def test_matting_model_options_include_downloadable_official_variants(layerforge assert all(option["downloaded"] is False for option in remote_options) assert all(option["description"] for option in remote_options) assert all(option["url"].startswith("https://huggingface.co/") for option in remote_options) - assert all(option["project_url"] == "https://github.com/ZhengPeng7/BiRefNet" for option in remote_options) + birefnet_options = [option for option in remote_options if option["backend"] == "birefnet"] + assert all(option["project_url"] == "https://github.com/ZhengPeng7/BiRefNet" for option in birefnet_options) assert any(option["path"] == "remote:portrait" for option in remote_options) portrait = next(option for option in node_module._BIREFNET_MODEL_CATALOG if option["id"] == "portrait") assert portrait["local_filename"] == "BiRefNet-portrait.safetensors" + rmbg_option = next(option for option in remote_options if option["path"] == "remote:rmbg_2_0") + assert rmbg_option["backend"] == "rmbg" + assert rmbg_option["label"] == "BRIA RMBG 2.0" + assert rmbg_option["project_url"] == "https://github.com/Bria-AI/RMBG-2.0" + def test_selected_remote_matting_model_is_sent_to_downloader(layerforge_runtime, monkeypatch): node_module = layerforge_runtime.matting + birefnet_module = layerforge_runtime.matting_birefnet selected_model = next( model for model in node_module._BIREFNET_MODEL_CATALOG if model["id"] == "portrait" ) downloaded = {} - monkeypatch.setattr(node_module, "_is_native_birefnet_checkpoint", lambda path: False) + monkeypatch.setattr(birefnet_module, "_is_native_birefnet_checkpoint", lambda path: False) def fake_download(model=None): downloaded["model"] = model return "downloaded-portrait.safetensors" - monkeypatch.setattr(node_module, "_download_birefnet_checkpoint", fake_download) + monkeypatch.setattr(birefnet_module, "_download_birefnet_checkpoint", fake_download) - result = node_module._ensure_birefnet_checkpoint("remote:portrait") + result = birefnet_module._ensure_birefnet_checkpoint("remote:portrait") assert result == "downloaded-portrait.safetensors" assert downloaded["model"] is selected_model @@ -611,6 +629,7 @@ def test_remote_matting_checkpoint_path_uses_background_removal_root(layerforge_ def test_remote_matting_download_uses_background_removal_root(layerforge_runtime, monkeypatch, tmp_path): node_module = layerforge_runtime.matting + birefnet_module = layerforge_runtime.matting_birefnet model = next(model for model in node_module._BIREFNET_MODEL_CATALOG if model["id"] == "portrait") download = {} @@ -623,9 +642,9 @@ def test_remote_matting_download_uses_background_removal_root(layerforge_runtime huggingface_hub = ModuleType("huggingface_hub") huggingface_hub.hf_hub_download = fake_hf_hub_download monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub) - monkeypatch.setattr(node_module, "_is_native_birefnet_checkpoint", lambda path: True) + monkeypatch.setattr(birefnet_module, "_is_native_birefnet_checkpoint", lambda path: True) - result = node_module._download_birefnet_checkpoint(model) + result = birefnet_module._download_birefnet_checkpoint(model) expected_dir = Path(tmp_path) / "models" / "background_removal" assert Path(download["local_dir"]) == expected_dir @@ -634,25 +653,117 @@ def test_remote_matting_download_uses_background_removal_root(layerforge_runtime assert Path(result).exists() -def test_legacy_automatic_checkpoint_gets_friendly_filename(layerforge_runtime, monkeypatch, tmp_path): +def test_rmbg_model_path_uses_background_removal_subdirectory(layerforge_runtime): node_module = layerforge_runtime.matting + model = node_module._RMBG_MODEL_CATALOG[0] + + expected = Path(layerforge_runtime.matting.folder_paths.models_dir) / "background_removal" / "RMBG-2.0" + + assert Path(node_module._get_rmbg_model_directory(model)) == expected + + +def test_rmbg_checkpoint_is_not_selected_as_native_birefnet(layerforge_runtime, monkeypatch, tmp_path): + node_module = layerforge_runtime.matting + birefnet_module = layerforge_runtime.matting_birefnet + model_dir = Path(tmp_path) / "models" / "background_removal" / "RMBG-2.0" + model_dir.mkdir(parents=True) + checkpoint_path = model_dir / "model.safetensors" + checkpoint_path.write_bytes(b"RMBG checkpoint") + + monkeypatch.setattr(birefnet_module, "_get_birefnet_base_paths", lambda: [str(model_dir.parent)]) + monkeypatch.setattr(birefnet_module, "_iter_birefnet_checkpoint_paths", lambda: iter((str(checkpoint_path),))) + monkeypatch.setattr(birefnet_module, "_is_native_birefnet_checkpoint", lambda path: True) + + assert birefnet_module._find_local_birefnet_model() is None + + +def test_rmbg_download_uses_background_removal_subdirectory(layerforge_runtime, monkeypatch, tmp_path): + node_module = layerforge_runtime.matting + model = node_module._RMBG_MODEL_CATALOG[0] + download = {} + + def fake_snapshot_download(**kwargs): + download.update(kwargs) + model_dir = Path(kwargs["local_dir"]) + model_dir.mkdir(parents=True, exist_ok=True) + for filename in ( + "config.json", + "preprocessor_config.json", + "birefnet.py", + "BiRefNet_config.py", + "model.safetensors", + ): + (model_dir / filename).write_bytes(b"model") + return str(model_dir) + + huggingface_hub = ModuleType("huggingface_hub") + huggingface_hub.snapshot_download = fake_snapshot_download + monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub) + + result = node_module._download_rmbg_model(model) + expected_dir = Path(tmp_path) / "models" / "background_removal" / "RMBG-2.0" + + assert Path(download["local_dir"]) == expected_dir + assert download["repo_id"] == "briaai/RMBG-2.0" + assert "*.safetensors" in download["allow_patterns"] + assert Path(result) == expected_dir + assert node_module._is_rmbg_model_directory(result) + + +def test_rmbg_adapter_returns_batched_hw_mask(layerforge_runtime): + import torch + + class FakeRMBGModel: + def __call__(self, image): + return (None, torch.zeros((image.shape[0], 1, 1024, 1024), device=image.device)) + + adapter = layerforge_runtime.matting.RMBG2Model( + FakeRMBGModel(), + torch.device("cpu"), + ) + image = torch.rand((1, 2, 3, 4), dtype=torch.float32) + + mask = adapter.encode_image(image) + + assert tuple(mask.shape) == (1, 2, 3) + assert torch.allclose(mask, torch.full((1, 2, 3), 0.5)) + + +def test_rmbg_model_status_does_not_require_native_birefnet_loader(layerforge_runtime, monkeypatch): + node_module = layerforge_runtime.matting + api_module = layerforge_runtime.matting_api + _record_matting_responses(monkeypatch, node_module) + + monkeypatch.setattr(api_module, "_get_rmbg_model_loader", lambda: object()) + monkeypatch.setattr(api_module, "_find_existing_rmbg_model", lambda model: None) + + status = _run_matting_model_check(node_module, "remote:rmbg_2_0") + + assert status.payload["available"] is False + assert status.payload["reason"] == "not_downloaded" + assert status.payload["selected_model"] == "BRIA RMBG 2.0" + assert status.payload["model_path"] == "remote:rmbg_2_0" + + +def test_legacy_automatic_checkpoint_is_not_migrated(layerforge_runtime, monkeypatch, tmp_path): + birefnet_module = layerforge_runtime.matting_birefnet model_dir = Path(tmp_path) / "models" / "background_removal" model_dir.mkdir(parents=True) legacy_path = model_dir / "model.safetensors" + friendly_path = model_dir / "BiRefNet-general.safetensors" legacy_path.write_bytes(b"legacy checkpoint") monkeypatch.setattr( - node_module, + birefnet_module, "_is_native_birefnet_checkpoint", - lambda path: Path(path).resolve() == legacy_path.resolve(), + lambda path: Path(path).resolve() == friendly_path.resolve() and Path(path).is_file(), ) - result = node_module._find_existing_birefnet_default_checkpoint() - friendly_path = model_dir / "BiRefNet-general.safetensors" + result = birefnet_module._find_existing_birefnet_default_checkpoint() - assert Path(result) == friendly_path - assert friendly_path.exists() - assert not legacy_path.exists() + assert result is None + assert legacy_path.exists() + assert not friendly_path.exists() def test_empty_execution_returns_comfyui_compatible_fallback_tensors(layerforge_runtime):