Store Matting settings server-side and support optional Hugging Face token

This commit is contained in:
Dariusz L
2026-08-13 18:20:52 +02:00
parent 91a16edde4
commit a830da7593
14 changed files with 559 additions and 78 deletions
+23 -1
View File
@@ -3,7 +3,13 @@
import folder_paths
from aiohttp import web
from .api import check_matting_model, matting, register_matting_routes
from .api import (
check_matting_model,
get_matting_settings,
matting,
register_matting_routes,
save_matting_settings,
)
from .backends.birefnet import (
_download_birefnet_checkpoint,
_ensure_birefnet_checkpoint,
@@ -49,14 +55,30 @@ from .paths import (
_get_rmbg_model_roots,
)
from .service import BiRefNetMatting
from .settings import (
DEFAULT_SETTINGS,
SETTINGS_FILE,
get_huggingface_token,
get_public_settings,
load_settings,
save_settings,
)
__all__ = [
"BiRefNetMatting",
"RMBG2Model",
"check_matting_model",
"get_matting_settings",
"matting",
"register_matting_routes",
"save_matting_settings",
"DEFAULT_SETTINGS",
"SETTINGS_FILE",
"get_huggingface_token",
"get_public_settings",
"load_settings",
"save_settings",
"_BIREFNET_DEFAULT_LOCAL_FILENAME",
"_BIREFNET_FILENAME",
"_BIREFNET_MODEL_CATALOG",
+44 -1
View File
@@ -23,10 +23,45 @@ from .backends.rmbg import (
from .options import _get_birefnet_model_options
from .paths import _get_birefnet_base_paths
from .service import BiRefNetMatting
from .settings import get_public_settings, save_settings
_matting_lock = None
async def get_matting_settings(request):
"""Return persisted Matting preferences without exposing the Hugging Face token."""
del request
try:
return web.json_response({"settings": get_public_settings()})
except Exception as error:
log.error(f"Error reading Matting settings: {error}")
return web.json_response(
{"error": "Unable to read Matting settings"},
status=500,
)
async def save_matting_settings(request):
"""Persist Matting preferences and an optional Hugging Face token."""
try:
payload = await request.json()
except Exception:
return web.json_response({"error": "Expected valid JSON"}, status=400)
if not isinstance(payload, dict):
return web.json_response({"error": "Expected JSON object"}, status=400)
try:
save_settings(payload)
return web.json_response({"success": True, "settings": get_public_settings()})
except Exception as error:
log.error(f"Error saving Matting settings: {error}")
return web.json_response(
{"error": "Unable to save Matting settings"},
status=500,
)
def _matting_status_response(
available,
reason,
@@ -292,8 +327,16 @@ async def matting(request):
def register_matting_routes():
"""Register matting endpoints without import-time decorators."""
PromptServer.instance.routes.get("/matting/settings")(get_matting_settings)
PromptServer.instance.routes.post("/matting/settings")(save_matting_settings)
PromptServer.instance.routes.get("/matting/check-model")(check_matting_model)
PromptServer.instance.routes.post("/matting")(matting)
__all__ = ["check_matting_model", "matting", "register_matting_routes"]
__all__ = [
"check_matting_model",
"get_matting_settings",
"matting",
"register_matting_routes",
"save_matting_settings",
]
+13 -11
View File
@@ -19,6 +19,7 @@ from ..paths import (
_get_birefnet_remote_checkpoint_path,
_get_rmbg_model_roots,
)
from ..settings import get_huggingface_token
def _get_comfy_birefnet_loader():
@@ -164,19 +165,20 @@ def _download_birefnet_checkpoint(model=None):
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
try:
downloaded_path = hf_hub_download(
repo_id=repository,
filename=filename,
local_dir=download_dir,
local_dir_use_symlinks=False,
)
downloaded_path = hf_hub_download(**download_kwargs)
except TypeError:
downloaded_path = hf_hub_download(
repo_id=repository,
filename=filename,
local_dir=download_dir,
)
download_kwargs.pop("local_dir_use_symlinks", None)
downloaded_path = hf_hub_download(**download_kwargs)
if not _is_native_birefnet_checkpoint(downloaded_path):
raise RuntimeError(
+10 -5
View File
@@ -10,6 +10,7 @@ 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 ..settings import get_huggingface_token
def _get_rmbg_transformers_status():
@@ -145,12 +146,16 @@ def _download_rmbg_model(model):
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:
snapshot_download(
repo_id=model["repo_id"],
local_dir=model_directory,
allow_patterns=["*.json", "*.py", "*.safetensors"],
)
snapshot_download(**download_kwargs)
except Exception as error:
raise RuntimeError(
"Unable to download BRIA RMBG 2.0. Accept the model's gated access on "
+128
View File
@@ -0,0 +1,128 @@
"""Persistent Matting settings stored next to the LayerForge custom node."""
from __future__ import annotations
import json
import os
import tempfile
import threading
from pathlib import Path
from typing import Any, Mapping
SETTINGS_FILE = Path(__file__).resolve().parents[2] / "layerforge_settings.json"
DEFAULT_SETTINGS = {
"model_path": "",
"mode": "remove_background",
"threshold": 0.5,
"hf_token": "",
}
_VALID_MODES = {"remove_background", "remove_foreground", "mask_only"}
_SETTINGS_LOCK = threading.RLock()
def _read_settings_file() -> dict[str, Any]:
try:
with SETTINGS_FILE.open("r", encoding="utf-8") as settings_file:
data = json.load(settings_file)
except (OSError, ValueError, TypeError):
return {}
return data if isinstance(data, dict) else {}
def _normalize_settings(settings: Mapping[str, Any] | None) -> dict[str, Any]:
data = dict(DEFAULT_SETTINGS)
raw = dict(settings or {})
model_path = raw.get("model_path", raw.get("modelPath", DEFAULT_SETTINGS["model_path"]))
data["model_path"] = str(model_path or "").strip()
mode = str(raw.get("mode") or DEFAULT_SETTINGS["mode"]).strip()
data["mode"] = mode if mode in _VALID_MODES else DEFAULT_SETTINGS["mode"]
try:
threshold = float(raw.get("threshold", DEFAULT_SETTINGS["threshold"]))
except (TypeError, ValueError):
threshold = DEFAULT_SETTINGS["threshold"]
data["threshold"] = min(1.0, max(0.0, threshold))
data["hf_token"] = str(raw.get("hf_token") or "").strip()
return data
def load_settings() -> dict[str, Any]:
"""Load and normalize all settings, including the private Hugging Face token."""
with _SETTINGS_LOCK:
return _normalize_settings(_read_settings_file())
def get_public_settings() -> dict[str, Any]:
"""Return settings safe for the frontend without exposing the token."""
settings = load_settings()
return {
"model_path": settings["model_path"],
"mode": settings["mode"],
"threshold": settings["threshold"],
"hf_token_configured": bool(settings["hf_token"]),
"configured": SETTINGS_FILE.is_file(),
}
def get_huggingface_token() -> str:
"""Return the configured token for authenticated Hugging Face downloads."""
return str(load_settings().get("hf_token") or "")
def _write_settings_file(settings: Mapping[str, Any]) -> None:
SETTINGS_FILE.parent.mkdir(parents=True, exist_ok=True)
file_descriptor, temporary_name = tempfile.mkstemp(
prefix=f".{SETTINGS_FILE.stem}.",
suffix=".tmp",
dir=SETTINGS_FILE.parent,
text=True,
)
try:
with os.fdopen(file_descriptor, "w", encoding="utf-8") as settings_file:
json.dump(dict(settings), settings_file, indent=2)
settings_file.write("\n")
settings_file.flush()
os.fsync(settings_file.fileno())
os.replace(temporary_name, SETTINGS_FILE)
finally:
try:
os.unlink(temporary_name)
except FileNotFoundError:
pass
def save_settings(payload: Mapping[str, Any]) -> dict[str, Any]:
"""Validate and persist Matting settings without logging secret values."""
if not isinstance(payload, Mapping):
raise TypeError("Matting settings must be a JSON object")
with _SETTINGS_LOCK:
current = load_settings()
for key in ("model_path", "mode", "threshold"):
if key in payload:
current[key] = payload[key]
if "modelPath" in payload and "model_path" not in payload:
current["model_path"] = payload["modelPath"]
if payload.get("clear_hf_token"):
current["hf_token"] = ""
elif str(payload.get("hf_token") or "").strip():
current["hf_token"] = str(payload["hf_token"]).strip()
normalized = _normalize_settings(current)
_write_settings_file(normalized)
return normalized
__all__ = [
"DEFAULT_SETTINGS",
"SETTINGS_FILE",
"get_huggingface_token",
"get_public_settings",
"load_settings",
"save_settings",
]