From a830da759368a93ba3bbe1800bef1da83ef55cba Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 13 Aug 2026 18:20:52 +0200 Subject: [PATCH] Store Matting settings server-side and support optional Hugging Face token --- .gitignore | 3 + README.md | 4 +- js/CanvasView.js | 114 +++++++++++++++++------ js/css/canvas_view.css | 30 ++++++ js/utils/MattingUtils.js | 14 +++ python/matting/__init__.py | 24 ++++- python/matting/api.py | 45 ++++++++- python/matting/backends/birefnet.py | 24 ++--- python/matting/backends/rmbg.py | 15 ++- python/matting/settings.py | 128 ++++++++++++++++++++++++++ src/CanvasView.ts | 131 +++++++++++++++++++++------ src/css/canvas_view.css | 30 ++++++ src/utils/MattingUtils.ts | 42 +++++++++ tests/test_entrypoint_integration.py | 33 +++++++ 14 files changed, 559 insertions(+), 78 deletions(-) create mode 100644 python/matting/settings.py diff --git a/.gitignore b/.gitignore index 68966e2..7a165c1 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,9 @@ tmp/ *.swp *.tmp +# local LayerForge settings may contain a Hugging Face token +layerforge_settings.json + *.txt Zbierz .gitignore diff --git a/README.md b/README.md index 0fe222a..f8f9cd1 100644 --- a/README.md +++ b/README.md @@ -236,13 +236,15 @@ RMBG 2.0 locally through Transformers. > - **Automatic download**: If no compatible checkpoint is found, LayerForge downloads the standard full BiRefNet checkpoint automatically on first use. > - **Compatibility**: LayerForge uses ComfyUI's native BiRefNet loader, so the checkpoint must use the full BiRefNet architecture. The `lite-*` variants are not compatible with this native loader. > - **BRIA RMBG 2.0**: Select it in Matting settings to download the local Transformers model into `ComfyUI/models/background_removal/RMBG-2.0/`. The Hugging Face repository is gated, so access must be accepted there before downloading. +> - **Hugging Face token**: Paste an optional read token in Matting settings for gated downloads. It is stored server-side in `layerforge_settings.json`, never in the workflow or browser `localStorage`. > - **BRIA license**: The self-hosted RMBG 2.0 weights are released for non-commercial use under the model's stated license. Check the [official model card](https://huggingface.co/briaai/RMBG-2.0) before using it commercially. -Click the gear button next to **Matting** to configure its behavior. Settings are saved locally in the browser and include: +Click the gear button next to **Matting** to configure its behavior. Settings are saved server-side in `layerforge_settings.json` and include: - **Model**: automatically select a compatible local checkpoint, choose an installed local checkpoint, or select an official BiRefNet or BRIA RMBG 2.0 model for download on first use. - **Processing mode**: remove the background, remove the detected foreground, or apply the generated mask to LayerForge's Draw Mask output-area mask. - **Mask threshold**: use `0` for a soft alpha mask or a higher value for a harder cutout. +- **Hugging Face token**: an optional read token for gated repositories such as BRIA RMBG 2.0; leave the field blank to keep the saved token. The model selector separates compatible local models from official models available for download. BiRefNet variants are stored directly in `ComfyUI/models/background_removal/` with descriptive filenames such as `BiRefNet-general.safetensors` and `BiRefNet-portrait.safetensors`. BRIA RMBG 2.0 is stored as a complete Transformers model under `ComfyUI/models/background_removal/RMBG-2.0/`. Every selected model is validated before use. Lite BiRefNet variants are intentionally excluded because they use a different architecture. Hugging Face cache folders are ignored in the selector and are not deleted. diff --git a/js/CanvasView.js b/js/CanvasView.js index dc48113..3c30812 100644 --- a/js/CanvasView.js +++ b/js/CanvasView.js @@ -16,40 +16,60 @@ import { exportCanvasImage } from "./utils/CanvasExportUtils.js"; import { getFlattenedCanvasBlob } from "./utils/CanvasBlobUtils.js"; import { loadPreviewImage } from "./utils/PreviewUtils.js"; import { getImageAddMode } from "./utils/CanvasInputUtils.js"; -import { fetchMattingModelStatus } from "./utils/MattingUtils.js"; +import { fetchMattingModelStatus, fetchMattingSettings, saveMattingSettings as saveMattingSettingsToServer, } from "./utils/MattingUtils.js"; import { setupSAMDetectorHook } from "./SAMDetectorIntegration.js"; const log = createModuleLogger('Canvas_view'); -const MATTING_SETTINGS_STORAGE_KEY = 'layerforge.matting.settings'; const DEFAULT_MATTING_SETTINGS = { modelPath: '', mode: 'remove_background', threshold: 0.5, + hfTokenConfigured: false, }; const isMattingMode = (value) => { return value === 'remove_background' || value === 'remove_foreground' || value === 'mask_only'; }; -const loadMattingSettings = () => { - try { - const stored = JSON.parse(localStorage.getItem(MATTING_SETTINGS_STORAGE_KEY) || '{}'); - const threshold = Number(stored.threshold); - return { - modelPath: typeof stored.modelPath === 'string' ? stored.modelPath : DEFAULT_MATTING_SETTINGS.modelPath, - mode: isMattingMode(stored.mode) ? stored.mode : DEFAULT_MATTING_SETTINGS.mode, - threshold: Number.isFinite(threshold) ? Math.min(1, Math.max(0, threshold)) : DEFAULT_MATTING_SETTINGS.threshold, - }; - } - catch (error) { - log.warn('Unable to load Matting settings:', error); - return { ...DEFAULT_MATTING_SETTINGS }; - } +const normalizeMattingSettings = (settings) => { + const threshold = Number(settings.threshold); + return { + modelPath: typeof settings.modelPath === 'string' ? settings.modelPath : DEFAULT_MATTING_SETTINGS.modelPath, + mode: isMattingMode(settings.mode) ? settings.mode : DEFAULT_MATTING_SETTINGS.mode, + threshold: Number.isFinite(threshold) ? Math.min(1, Math.max(0, threshold)) : DEFAULT_MATTING_SETTINGS.threshold, + hfTokenConfigured: settings.hfTokenConfigured === true, + }; }; -const saveMattingSettings = (settings) => { +const fromServerMattingSettings = (settings) => normalizeMattingSettings({ + modelPath: settings.model_path, + mode: settings.mode, + threshold: settings.threshold, + hfTokenConfigured: settings.hf_token_configured, +}); +const loadMattingSettings = async () => { try { - localStorage.setItem(MATTING_SETTINGS_STORAGE_KEY, JSON.stringify(settings)); + const response = await fetchMattingSettings(); + if (response.ok && response.data.settings) { + return fromServerMattingSettings(response.data.settings); + } } catch (error) { - log.warn('Unable to save Matting settings:', error); + log.warn('Unable to load Matting settings from ComfyUI:', error); } + return { ...DEFAULT_MATTING_SETTINGS }; +}; +const persistMattingSettings = async (settings, token, clearToken) => { + const payload = { + model_path: settings.modelPath, + mode: settings.mode, + threshold: settings.threshold, + }; + if (token.trim()) + payload.hf_token = token.trim(); + if (clearToken) + payload.clear_hf_token = true; + const response = await saveMattingSettingsToServer(payload); + if (!response.ok || !response.data.settings) { + throw new Error(response.data.error || 'Unable to save Matting settings on the ComfyUI server.'); + } + return fromServerMattingSettings(response.data.settings); }; const getMattingModeLabel = (mode) => { switch (mode) { @@ -187,7 +207,7 @@ async function createCanvasWidget(node, widget, app) { const openMattingSettings = async () => { if (mattingSettingsBackdrop) return; - const settings = loadMattingSettings(); + const settings = await loadMattingSettings(); let modelOptions = []; let modelStatusMessage = 'Model options are loaded from ComfyUI background-removal storage.'; try { @@ -339,6 +359,29 @@ async function createCanvasWidget(node, widget, app) { thresholdValue.textContent = value.toFixed(2); }; thresholdContainer.append(thresholdInput, thresholdValue); + const tokenContainer = document.createElement('div'); + tokenContainer.className = 'lf-matting-settings-token'; + const tokenInput = document.createElement('input'); + tokenInput.type = 'password'; + tokenInput.className = 'lf-matting-settings-input'; + tokenInput.autocomplete = 'off'; + tokenInput.placeholder = settings.hfTokenConfigured + ? 'Token saved — leave blank to keep it' + : 'Paste a Hugging Face read token'; + tokenInput.spellcheck = false; + const clearTokenLabel = document.createElement('label'); + clearTokenLabel.className = 'lf-matting-settings-token-clear'; + const clearTokenInput = document.createElement('input'); + clearTokenInput.type = 'checkbox'; + clearTokenInput.disabled = !settings.hfTokenConfigured; + const clearTokenText = document.createElement('span'); + clearTokenText.textContent = 'Clear saved token'; + clearTokenLabel.append(clearTokenInput, clearTokenText); + tokenInput.oninput = () => { + if (tokenInput.value.trim()) + clearTokenInput.checked = false; + }; + tokenContainer.append(tokenInput, clearTokenLabel); const modelStatus = document.createElement('p'); modelStatus.className = 'lf-matting-settings-status'; const localCount = localModelOptions.length; @@ -348,7 +391,7 @@ async function createCanvasWidget(node, widget, app) { remoteCount > 0 ? `${remoteCount} official model(s) available for download` : '', ].filter(Boolean).join('; '); modelStatus.textContent = `${modelCounts}. ${modelStatusMessage}`; - body.append(createRow('Model', modelSelect, 'Choose a local BiRefNet checkpoint or BRIA RMBG 2.0, or download an official model on first use.'), modelDetails, createRow('Processing mode', modeSelect, 'The selected mode controls what the Matting button creates from the detected mask.'), createRow('Mask threshold', thresholdContainer, 'Set to 0 for a soft alpha mask; higher values create a harder cutout.'), modelStatus); + body.append(createRow('Model', modelSelect, 'Choose a local BiRefNet checkpoint or BRIA RMBG 2.0, or download an official model on first use.'), modelDetails, createRow('Processing mode', modeSelect, 'The selected mode controls what the Matting button creates from the detected mask.'), createRow('Mask threshold', thresholdContainer, 'Set to 0 for a soft alpha mask; higher values create a harder cutout.'), createRow('Hugging Face token', tokenContainer, 'Optional read token for gated models such as BRIA RMBG 2.0. It is stored only in the ComfyUI custom node settings file.'), modelStatus); const actions = document.createElement('div'); actions.className = 'lf-matting-settings-actions'; const resetButton = document.createElement('button'); @@ -366,14 +409,25 @@ async function createCanvasWidget(node, widget, app) { saveButton.type = 'button'; saveButton.className = 'lf-matting-settings-primary'; saveButton.textContent = 'Save settings'; - saveButton.onclick = () => { - saveMattingSettings({ - modelPath: modelSelect.value === 'auto' ? '' : modelSelect.value, - mode: modeSelect.value, - threshold: Number(thresholdInput.value), - }); - closeMattingSettings(); - showInfoNotification('Matting settings saved.', 2000); + saveButton.onclick = async () => { + saveButton.disabled = true; + try { + await persistMattingSettings({ + modelPath: modelSelect.value === 'auto' ? '' : modelSelect.value, + mode: modeSelect.value, + threshold: Number(thresholdInput.value), + hfTokenConfigured: settings.hfTokenConfigured, + }, tokenInput.value, clearTokenInput.checked); + closeMattingSettings(); + showInfoNotification('Matting settings saved.', 2000); + } + catch (error) { + log.error('Unable to save Matting settings:', error); + showErrorNotification(error instanceof Error ? error.message : 'Unable to save Matting settings.', 8000); + } + finally { + saveButton.disabled = false; + } }; actions.append(resetButton, saveButton); dialog.append(header, body, actions); @@ -646,7 +700,7 @@ async function createCanvasWidget(node, widget, app) { const button = e.target.closest('.lf-matting-button'); if (button.classList.contains('lf-loading')) return; - const mattingSettings = loadMattingSettings(); + const mattingSettings = await loadMattingSettings(); try { // First check if model is available const { data: modelStatus } = await fetchMattingModelStatus(mattingSettings.modelPath); diff --git a/js/css/canvas_view.css b/js/css/canvas_view.css index d64b455..cd6ffc7 100644 --- a/js/css/canvas_view.css +++ b/js/css/canvas_view.css @@ -726,7 +726,37 @@ border-radius: 4px; } +.lf-matting-settings-input { + width: 100%; + min-height: 34px; + box-sizing: border-box; + padding: 6px 9px; + color: #f4f4f4; + background: #252525; + border: 1px solid #5a5a5a; + border-radius: 4px; +} + +.lf-matting-settings-token { + display: flex; + flex-direction: column; + gap: 8px; +} + +.lf-matting-settings-token-clear { + display: flex; + align-items: center; + gap: 7px; + color: #bdbdbd; + font-size: 11px; +} + +.lf-matting-settings-token-clear input { + margin: 0; +} + .lf-matting-settings-select:focus, +.lf-matting-settings-input:focus, .lf-matting-settings-threshold input:focus { outline: 2px solid #4b8fe8; outline-offset: 1px; diff --git a/js/utils/MattingUtils.js b/js/utils/MattingUtils.js index d2deb0b..18ec186 100644 --- a/js/utils/MattingUtils.js +++ b/js/utils/MattingUtils.js @@ -6,3 +6,17 @@ export async function fetchMattingModelStatus(modelPath) { const data = await response.json(); return { ok: response.ok, data }; } +export async function fetchMattingSettings() { + const response = await fetch('/matting/settings'); + const data = await response.json(); + return { ok: response.ok, data }; +} +export async function saveMattingSettings(settings) { + const response = await fetch('/matting/settings', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(settings), + }); + const data = await response.json(); + return { ok: response.ok, data }; +} diff --git a/python/matting/__init__.py b/python/matting/__init__.py index 70a4adf..e180816 100644 --- a/python/matting/__init__.py +++ b/python/matting/__init__.py @@ -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", diff --git a/python/matting/api.py b/python/matting/api.py index 4ad84ca..2433e5b 100644 --- a/python/matting/api.py +++ b/python/matting/api.py @@ -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", +] diff --git a/python/matting/backends/birefnet.py b/python/matting/backends/birefnet.py index cdc768d..d68d875 100644 --- a/python/matting/backends/birefnet.py +++ b/python/matting/backends/birefnet.py @@ -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( diff --git a/python/matting/backends/rmbg.py b/python/matting/backends/rmbg.py index 40a00b8..0971a29 100644 --- a/python/matting/backends/rmbg.py +++ b/python/matting/backends/rmbg.py @@ -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 " diff --git a/python/matting/settings.py b/python/matting/settings.py new file mode 100644 index 0000000..41921a2 --- /dev/null +++ b/python/matting/settings.py @@ -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", +] diff --git a/src/CanvasView.ts b/src/CanvasView.ts index 58f2a35..b8c6822 100644 --- a/src/CanvasView.ts +++ b/src/CanvasView.ts @@ -22,7 +22,13 @@ import { exportCanvasImage, type CanvasExportAction } from "./utils/CanvasExport import { getFlattenedCanvasBlob, type CanvasBlobVariant } from "./utils/CanvasBlobUtils.js"; import { loadPreviewImage } from "./utils/PreviewUtils.js"; import { getImageAddMode } from "./utils/CanvasInputUtils.js"; -import { fetchMattingModelStatus } from "./utils/MattingUtils.js"; +import { + fetchMattingModelStatus, + fetchMattingSettings, + saveMattingSettings as saveMattingSettingsToServer, + type MattingServerSettings, + type MattingSettingsUpdate, +} from "./utils/MattingUtils.js"; import { registerImageInClipspace, startSAMDetectorMonitoring, setupSAMDetectorHook } from "./SAMDetectorIntegration.js"; import type { ComfyNode, Layer, AddMode } from './types'; @@ -34,6 +40,7 @@ interface MattingSettings { modelPath: string; mode: MattingMode; threshold: number; + hfTokenConfigured: boolean; } interface MattingModelOption { @@ -47,39 +54,66 @@ interface MattingModelOption { downloaded?: boolean; } -const MATTING_SETTINGS_STORAGE_KEY = 'layerforge.matting.settings'; const DEFAULT_MATTING_SETTINGS: MattingSettings = { modelPath: '', mode: 'remove_background', threshold: 0.5, + hfTokenConfigured: false, }; const isMattingMode = (value: unknown): value is MattingMode => { return value === 'remove_background' || value === 'remove_foreground' || value === 'mask_only'; }; -const loadMattingSettings = (): MattingSettings => { - try { - const stored = JSON.parse(localStorage.getItem(MATTING_SETTINGS_STORAGE_KEY) || '{}') as Partial; - const threshold = Number(stored.threshold); - - return { - modelPath: typeof stored.modelPath === 'string' ? stored.modelPath : DEFAULT_MATTING_SETTINGS.modelPath, - mode: isMattingMode(stored.mode) ? stored.mode : DEFAULT_MATTING_SETTINGS.mode, - threshold: Number.isFinite(threshold) ? Math.min(1, Math.max(0, threshold)) : DEFAULT_MATTING_SETTINGS.threshold, - }; - } catch (error) { - log.warn('Unable to load Matting settings:', error); - return { ...DEFAULT_MATTING_SETTINGS }; - } +const normalizeMattingSettings = (settings: Partial): MattingSettings => { + const threshold = Number(settings.threshold); + return { + modelPath: typeof settings.modelPath === 'string' ? settings.modelPath : DEFAULT_MATTING_SETTINGS.modelPath, + mode: isMattingMode(settings.mode) ? settings.mode : DEFAULT_MATTING_SETTINGS.mode, + threshold: Number.isFinite(threshold) ? Math.min(1, Math.max(0, threshold)) : DEFAULT_MATTING_SETTINGS.threshold, + hfTokenConfigured: settings.hfTokenConfigured === true, + }; }; -const saveMattingSettings = (settings: MattingSettings): void => { +const fromServerMattingSettings = (settings: MattingServerSettings): MattingSettings => normalizeMattingSettings({ + modelPath: settings.model_path, + mode: settings.mode as MattingMode, + threshold: settings.threshold, + hfTokenConfigured: settings.hf_token_configured, +}); + +const loadMattingSettings = async (): Promise => { try { - localStorage.setItem(MATTING_SETTINGS_STORAGE_KEY, JSON.stringify(settings)); + const response = await fetchMattingSettings(); + if (response.ok && response.data.settings) { + return fromServerMattingSettings(response.data.settings); + } } catch (error) { - log.warn('Unable to save Matting settings:', error); + log.warn('Unable to load Matting settings from ComfyUI:', error); } + + return { ...DEFAULT_MATTING_SETTINGS }; +}; + +const persistMattingSettings = async ( + settings: MattingSettings, + token: string, + clearToken: boolean, +): Promise => { + const payload: MattingSettingsUpdate = { + model_path: settings.modelPath, + mode: settings.mode, + threshold: settings.threshold, + }; + if (token.trim()) payload.hf_token = token.trim(); + if (clearToken) payload.clear_hf_token = true; + + const response = await saveMattingSettingsToServer(payload); + if (!response.ok || !response.data.settings) { + throw new Error(response.data.error || 'Unable to save Matting settings on the ComfyUI server.'); + } + + return fromServerMattingSettings(response.data.settings); }; const getMattingModeLabel = (mode: MattingMode): string => { @@ -260,7 +294,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): const openMattingSettings = async (): Promise => { if (mattingSettingsBackdrop) return; - const settings = loadMattingSettings(); + const settings = await loadMattingSettings(); let modelOptions: MattingModelOption[] = []; let modelStatusMessage = 'Model options are loaded from ComfyUI background-removal storage.'; @@ -448,6 +482,31 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): thresholdContainer.append(thresholdInput, thresholdValue); + const tokenContainer = document.createElement('div'); + tokenContainer.className = 'lf-matting-settings-token'; + + const tokenInput = document.createElement('input'); + tokenInput.type = 'password'; + tokenInput.className = 'lf-matting-settings-input'; + tokenInput.autocomplete = 'off'; + tokenInput.placeholder = settings.hfTokenConfigured + ? 'Token saved — leave blank to keep it' + : 'Paste a Hugging Face read token'; + tokenInput.spellcheck = false; + + const clearTokenLabel = document.createElement('label'); + clearTokenLabel.className = 'lf-matting-settings-token-clear'; + const clearTokenInput = document.createElement('input'); + clearTokenInput.type = 'checkbox'; + clearTokenInput.disabled = !settings.hfTokenConfigured; + const clearTokenText = document.createElement('span'); + clearTokenText.textContent = 'Clear saved token'; + clearTokenLabel.append(clearTokenInput, clearTokenText); + tokenInput.oninput = () => { + if (tokenInput.value.trim()) clearTokenInput.checked = false; + }; + tokenContainer.append(tokenInput, clearTokenLabel); + const modelStatus = document.createElement('p'); modelStatus.className = 'lf-matting-settings-status'; const localCount = localModelOptions.length; @@ -463,6 +522,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): modelDetails, createRow('Processing mode', modeSelect, 'The selected mode controls what the Matting button creates from the detected mask.'), createRow('Mask threshold', thresholdContainer, 'Set to 0 for a soft alpha mask; higher values create a harder cutout.'), + createRow('Hugging Face token', tokenContainer, 'Optional read token for gated models such as BRIA RMBG 2.0. It is stored only in the ComfyUI custom node settings file.'), modelStatus, ); @@ -485,14 +545,27 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): saveButton.type = 'button'; saveButton.className = 'lf-matting-settings-primary'; saveButton.textContent = 'Save settings'; - saveButton.onclick = () => { - saveMattingSettings({ - modelPath: modelSelect.value === 'auto' ? '' : modelSelect.value, - mode: modeSelect.value as MattingMode, - threshold: Number(thresholdInput.value), - }); - closeMattingSettings(); - showInfoNotification('Matting settings saved.', 2000); + saveButton.onclick = async () => { + saveButton.disabled = true; + try { + await persistMattingSettings( + { + modelPath: modelSelect.value === 'auto' ? '' : modelSelect.value, + mode: modeSelect.value as MattingMode, + threshold: Number(thresholdInput.value), + hfTokenConfigured: settings.hfTokenConfigured, + }, + tokenInput.value, + clearTokenInput.checked, + ); + closeMattingSettings(); + showInfoNotification('Matting settings saved.', 2000); + } catch (error) { + log.error('Unable to save Matting settings:', error); + showErrorNotification(error instanceof Error ? error.message : 'Unable to save Matting settings.', 8000); + } finally { + saveButton.disabled = false; + } }; actions.append(resetButton, saveButton); @@ -811,7 +884,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): const button = (e.target as HTMLElement).closest('.lf-matting-button') as HTMLButtonElement; if (button.classList.contains('lf-loading')) return; - const mattingSettings = loadMattingSettings(); + const mattingSettings = await loadMattingSettings(); try { // First check if model is available diff --git a/src/css/canvas_view.css b/src/css/canvas_view.css index d64b455..cd6ffc7 100644 --- a/src/css/canvas_view.css +++ b/src/css/canvas_view.css @@ -726,7 +726,37 @@ border-radius: 4px; } +.lf-matting-settings-input { + width: 100%; + min-height: 34px; + box-sizing: border-box; + padding: 6px 9px; + color: #f4f4f4; + background: #252525; + border: 1px solid #5a5a5a; + border-radius: 4px; +} + +.lf-matting-settings-token { + display: flex; + flex-direction: column; + gap: 8px; +} + +.lf-matting-settings-token-clear { + display: flex; + align-items: center; + gap: 7px; + color: #bdbdbd; + font-size: 11px; +} + +.lf-matting-settings-token-clear input { + margin: 0; +} + .lf-matting-settings-select:focus, +.lf-matting-settings-input:focus, .lf-matting-settings-threshold input:focus { outline: 2px solid #4b8fe8; outline-offset: 1px; diff --git a/src/utils/MattingUtils.ts b/src/utils/MattingUtils.ts index 52c173a..04cb0b0 100644 --- a/src/utils/MattingUtils.ts +++ b/src/utils/MattingUtils.ts @@ -10,6 +10,30 @@ export interface MattingModelStatusResult { data: MattingModelStatus; } +export interface MattingServerSettings { + model_path: string; + mode: string; + threshold: number; + hf_token_configured: boolean; + configured: boolean; +} + +export interface MattingSettingsResponse { + ok: boolean; + data: { + settings?: MattingServerSettings; + error?: string; + }; +} + +export interface MattingSettingsUpdate { + model_path?: string; + mode?: string; + threshold?: number; + hf_token?: string; + clear_hf_token?: boolean; +} + export async function fetchMattingModelStatus( modelPath?: string, ): Promise> { @@ -21,3 +45,21 @@ export async function fetchMattingModelStatus( return { ok: response.ok, data }; } + +export async function fetchMattingSettings(): Promise { + const response = await fetch('/matting/settings'); + const data = await response.json() as MattingSettingsResponse['data']; + return { ok: response.ok, data }; +} + +export async function saveMattingSettings( + settings: MattingSettingsUpdate, +): Promise { + const response = await fetch('/matting/settings', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(settings), + }); + const data = await response.json() as MattingSettingsResponse['data']; + return { ok: response.ok, data }; +} diff --git a/tests/test_entrypoint_integration.py b/tests/test_entrypoint_integration.py index c6e671f..a7309fc 100644 --- a/tests/test_entrypoint_integration.py +++ b/tests/test_entrypoint_integration.py @@ -116,6 +116,7 @@ def _import_layerforge(monkeypatch, tmp_path): matting_api=sys.modules[f"{package_name}.python.matting.api"], matting_birefnet=sys.modules[f"{package_name}.python.matting.backends.birefnet"], matting_rmbg=sys.modules[f"{package_name}.python.matting.backends.rmbg"], + matting_settings=sys.modules[f"{package_name}.python.matting.settings"], matting_service=sys.modules[f"{package_name}.python.matting.service"], routes=routes, package_name=package_name, @@ -169,6 +170,8 @@ def test_entrypoint_registers_backend_route_contract(layerforge_runtime): ("GET", "/layerforge/get-latest-images/{since}"), ("GET", "/ycnode/get_latest_image"), ("POST", "/ycnode/load_image_from_path"), + ("GET", "/matting/settings"), + ("POST", "/matting/settings"), ("GET", "/matting/check-model"), ("POST", "/matting"), } @@ -176,6 +179,33 @@ def test_entrypoint_registers_backend_route_contract(layerforge_runtime): assert expected <= registered +def test_matting_settings_are_persisted_without_exposing_the_token(layerforge_runtime, monkeypatch, tmp_path): + settings_module = layerforge_runtime.matting_settings + settings_file = tmp_path / "layerforge_settings.json" + monkeypatch.setattr(settings_module, "SETTINGS_FILE", settings_file) + + saved = settings_module.save_settings( + { + "model_path": "remote:rmbg_2_0", + "mode": "mask_only", + "threshold": 0.75, + "hf_token": "hf-test-token", + } + ) + + assert saved["model_path"] == "remote:rmbg_2_0" + assert saved["mode"] == "mask_only" + assert saved["threshold"] == 0.75 + assert settings_module.get_huggingface_token() == "hf-test-token" + public_settings = settings_module.get_public_settings() + assert public_settings["hf_token_configured"] is True + assert "hf_token" not in public_settings + assert settings_file.exists() + + settings_module.save_settings({"clear_hf_token": True}) + assert settings_module.get_huggingface_token() == "" + + def test_tensor_input_normalization_preserves_comfyui_shapes(layerforge_runtime): import torch @@ -679,6 +709,7 @@ def test_rmbg_checkpoint_is_not_selected_as_native_birefnet(layerforge_runtime, def test_rmbg_download_uses_background_removal_subdirectory(layerforge_runtime, monkeypatch, tmp_path): node_module = layerforge_runtime.matting + rmbg_module = layerforge_runtime.matting_rmbg model = node_module._RMBG_MODEL_CATALOG[0] download = {} @@ -699,6 +730,7 @@ def test_rmbg_download_uses_background_removal_subdirectory(layerforge_runtime, huggingface_hub = ModuleType("huggingface_hub") huggingface_hub.snapshot_download = fake_snapshot_download monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub) + monkeypatch.setattr(rmbg_module, "get_huggingface_token", lambda: "hf-test-token") result = node_module._download_rmbg_model(model) expected_dir = Path(tmp_path) / "models" / "background_removal" / "RMBG-2.0" @@ -706,6 +738,7 @@ def test_rmbg_download_uses_background_removal_subdirectory(layerforge_runtime, assert Path(download["local_dir"]) == expected_dir assert download["repo_id"] == "briaai/RMBG-2.0" assert "*.safetensors" in download["allow_patterns"] + assert download["token"] == "hf-test-token" assert Path(result) == expected_dir assert node_module._is_rmbg_model_directory(result)