246 lines
8.5 KiB
Python
246 lines
8.5 KiB
Python
"""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_REMOTE_PREFIX,
|
|
_BIREFNET_REPOSITORY,
|
|
_BIREFNET_REQUIRED_KEYS,
|
|
)
|
|
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,
|
|
)
|
|
from ..progress import call_huggingface_download
|
|
from ..settings import get_huggingface_token
|
|
|
|
|
|
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, node_id=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}...")
|
|
download_kwargs = {
|
|
"repo_id": repository,
|
|
"filename": filename,
|
|
"local_dir": download_dir,
|
|
"local_dir_use_symlinks": False,
|
|
}
|
|
token = get_huggingface_token()
|
|
if token:
|
|
download_kwargs["token"] = token
|
|
downloaded_path = call_huggingface_download(
|
|
hf_hub_download,
|
|
download_kwargs,
|
|
model_label,
|
|
node_id=node_id,
|
|
)
|
|
|
|
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, node_id=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
|
|
if node_id is None:
|
|
return _download_birefnet_checkpoint(remote_model)
|
|
return _download_birefnet_checkpoint(remote_model, node_id=node_id)
|
|
|
|
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.")
|
|
|
|
if node_id is None:
|
|
return _download_birefnet_checkpoint()
|
|
return _download_birefnet_checkpoint(node_id=node_id)
|
|
|
|
|
|
__all__ = [
|
|
"_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",
|
|
]
|