refactor: split matting module into package and add BRIA RMBG 2.0 backend
This commit is contained in:
Dariusz L
+1
-1
@@ -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.
|
||||
"""
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -0,0 +1,2 @@
|
||||
"""Model-specific background-removal backends."""
|
||||
"""Background-removal backend implementations."""
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user