Store Matting settings server-side and support optional Hugging Face token
This commit is contained in:
@@ -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
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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,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 "
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
Reference in New Issue
Block a user