314 lines
11 KiB
Python
314 lines
11 KiB
Python
"""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
|
|
from ..progress import call_huggingface_download
|
|
from ..settings import get_huggingface_token
|
|
|
|
|
|
def _get_rmbg_transformers_status():
|
|
"""Inspect the installed Transformers API without changing the environment."""
|
|
try:
|
|
import transformers
|
|
except ModuleNotFoundError as error:
|
|
if error.name == "transformers":
|
|
return {
|
|
"loader": None,
|
|
"message": (
|
|
"BRIA RMBG 2.0 requires the 'transformers' package in the active ComfyUI environment. "
|
|
"Install or update it manually if you want to use this model."
|
|
),
|
|
}
|
|
return {
|
|
"loader": None,
|
|
"message": (
|
|
"The installed Transformers package could not be loaded for BRIA RMBG 2.0: "
|
|
f"{error}. Check the active ComfyUI environment and update Transformers manually if needed."
|
|
),
|
|
}
|
|
except Exception as error:
|
|
return {
|
|
"loader": None,
|
|
"message": (
|
|
"The installed Transformers package could not be loaded for BRIA RMBG 2.0: "
|
|
f"{error}. Update Transformers manually if needed."
|
|
),
|
|
}
|
|
|
|
version = getattr(transformers, "__version__", "unknown")
|
|
try:
|
|
loader = getattr(transformers, "AutoModelForImageSegmentation", None)
|
|
from_pretrained = getattr(loader, "from_pretrained", None)
|
|
except Exception as error:
|
|
return {
|
|
"loader": None,
|
|
"message": (
|
|
f"Installed Transformers version {version} is not supported for BRIA RMBG 2.0: {error}. "
|
|
"Update Transformers manually if you want to use this model."
|
|
),
|
|
}
|
|
|
|
if not callable(loader) or not callable(from_pretrained):
|
|
return {
|
|
"loader": None,
|
|
"message": (
|
|
f"Installed Transformers version {version} is not supported for BRIA RMBG 2.0. "
|
|
"It does not provide the required AutoModelForImageSegmentation API. "
|
|
"Update Transformers manually if you want to use this model."
|
|
),
|
|
}
|
|
|
|
return {"loader": loader, "message": None, "version": version}
|
|
|
|
|
|
def _get_rmbg_model_loader():
|
|
"""Return the Transformers loader used by BRIA RMBG 2.0, when installed."""
|
|
status = _get_rmbg_transformers_status()
|
|
if status["loader"] is None:
|
|
log.debug(status["message"])
|
|
return status["loader"]
|
|
|
|
|
|
def _get_rmbg_model_status_message():
|
|
"""Return an actionable message for a missing or incompatible Transformers install."""
|
|
return _get_rmbg_transformers_status()["message"]
|
|
|
|
|
|
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, node_id=None):
|
|
"""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}...")
|
|
download_kwargs = {
|
|
"repo_id": model["repo_id"],
|
|
"local_dir": model_directory,
|
|
"allow_patterns": ["*.json", "*.py", "*.safetensors"],
|
|
}
|
|
token = get_huggingface_token()
|
|
if token:
|
|
download_kwargs["token"] = token
|
|
try:
|
|
call_huggingface_download(
|
|
snapshot_download,
|
|
download_kwargs,
|
|
model["label"],
|
|
node_id=node_id,
|
|
)
|
|
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, node_id=None):
|
|
"""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
|
|
if node_id is None:
|
|
return _download_rmbg_model(remote_model)
|
|
return _download_rmbg_model(remote_model, node_id=node_id)
|
|
|
|
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):
|
|
status = _get_rmbg_transformers_status()
|
|
loader = status["loader"]
|
|
if loader is None:
|
|
raise RuntimeError(status["message"])
|
|
|
|
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}")
|
|
try:
|
|
model = loader.from_pretrained(
|
|
model_directory,
|
|
trust_remote_code=True,
|
|
local_files_only=True,
|
|
)
|
|
except Exception as error:
|
|
version = status.get("version", "unknown")
|
|
raise RuntimeError(
|
|
f"BRIA RMBG 2.0 could not be loaded with installed Transformers version {version}. "
|
|
"This version may not support the model's remote code; update Transformers manually "
|
|
f"if needed. Original error: {error}"
|
|
) from error
|
|
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_model_status_message",
|
|
"_get_rmbg_remote_model",
|
|
"_get_rmbg_transformers_status",
|
|
"_is_rmbg_model_directory",
|
|
]
|