From 11d448cdc8bf3b32b3330555b51e999fa88ca622 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 13 Aug 2026 20:47:10 +0200 Subject: [PATCH] Add matting download progress, interrupt handling, and inverted mask mode --- README.md | 3 +- js/CanvasView.js | 182 ++++++++++++++++++-- js/css/canvas_view.css | 78 ++++++++- python/matting/__init__.py | 2 + python/matting/api.py | 27 ++- python/matting/backends/birefnet.py | 24 ++- python/matting/backends/rmbg.py | 16 +- python/matting/progress.py | 239 +++++++++++++++++++++++++++ python/matting/service.py | 59 +++++-- python/matting/settings.py | 2 +- src/CanvasView.ts | 202 ++++++++++++++++++++-- src/css/canvas_view.css | 78 ++++++++- tests/test_entrypoint_integration.py | 103 +++++++++++- 13 files changed, 942 insertions(+), 73 deletions(-) create mode 100644 python/matting/progress.py diff --git a/README.md b/README.md index f8f9cd1..1a72a97 100644 --- a/README.md +++ b/README.md @@ -234,6 +234,7 @@ RMBG 2.0 locally through Transformers. - [Google Drive](https://drive.google.com/drive/folders/1BCLInCLH89fmTpYoP8Sgs_Eqww28f_wq?usp=sharing) > - **Installation Path**: Place a full BiRefNet checkpoint in `ComfyUI/models/background_removal/`. > - **Automatic download**: If no compatible checkpoint is found, LayerForge downloads the standard full BiRefNet checkpoint automatically on first use. +> - **Download progress**: A thin green progress bar inside the Matting button shows the real Hugging Face download progress without changing the button's layout. > - **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`. @@ -242,7 +243,7 @@ RMBG 2.0 locally through Transformers. 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. +- **Processing mode**: remove the background, remove the detected foreground, apply the generated background mask to LayerForge's Draw Mask output-area mask, or apply its inverted foreground 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. diff --git a/js/CanvasView.js b/js/CanvasView.js index 3c30812..67d88e5 100644 --- a/js/CanvasView.js +++ b/js/CanvasView.js @@ -1,6 +1,8 @@ // @ts-ignore import { app } from "../../scripts/app.js"; // @ts-ignore +import { api } from "../../scripts/api.js"; +// @ts-ignore import { ChangeTracker } from "../../scripts/changeTracker.js"; // @ts-ignore import { $el } from "../../scripts/ui.js"; @@ -26,7 +28,10 @@ const DEFAULT_MATTING_SETTINGS = { hfTokenConfigured: false, }; const isMattingMode = (value) => { - return value === 'remove_background' || value === 'remove_foreground' || value === 'mask_only'; + return value === 'remove_background' + || value === 'remove_foreground' + || value === 'mask_only' + || value === 'mask_only_inverted'; }; const normalizeMattingSettings = (settings) => { const threshold = Number(settings.threshold); @@ -77,6 +82,8 @@ const getMattingModeLabel = (mode) => { return 'Remove detected foreground / keep background'; case 'mask_only': return 'Apply generated mask to Draw Mask'; + case 'mask_only_inverted': + return 'Apply inverted mask to Draw Mask'; default: return 'Remove background / keep foreground'; } @@ -194,9 +201,38 @@ async function createCanvasWidget(node, widget, app) { const hideTooltip = () => { helpTooltip.style.display = 'none'; }; + const showMattingTooltip = (target, content) => { + helpTooltip.classList.add('lf-matting-tooltip'); + showTooltip(target, content); + }; + const hideMattingTooltip = () => { + hideTooltip(); + helpTooltip.classList.remove('lf-matting-tooltip'); + }; + const createMattingTooltipBadge = (labelText, tooltipText) => { + const badge = document.createElement('span'); + badge.className = 'lf-matting-tooltip-badge'; + badge.textContent = '?'; + badge.tabIndex = 0; + badge.dataset.tooltip = tooltipText; + badge.setAttribute('aria-label', `More information about ${labelText}`); + const show = () => { + const content = badge.dataset.tooltip; + if (content) + showMattingTooltip(badge, content); + }; + const hide = () => hideMattingTooltip(); + badge.addEventListener('mouseenter', show); + badge.addEventListener('focus', show); + badge.addEventListener('mouseleave', hide); + badge.addEventListener('blur', hide); + badge.addEventListener('click', (event) => event.preventDefault()); + return badge; + }; let mattingSettingsBackdrop = null; let mattingSettingsEscapeHandler = null; const closeMattingSettings = () => { + hideMattingTooltip(); if (mattingSettingsEscapeHandler) { document.removeEventListener('keydown', mattingSettingsEscapeHandler); mattingSettingsEscapeHandler = null; @@ -248,20 +284,17 @@ async function createCanvasWidget(node, widget, app) { header.append(title, closeButton); const body = document.createElement('div'); body.className = 'lf-matting-settings-body'; - const createRow = (labelText, control, description) => { + const createRow = (labelText, control, tooltipText) => { const row = document.createElement('label'); row.className = 'lf-matting-settings-row'; const label = document.createElement('span'); label.className = 'lf-matting-settings-label'; - label.textContent = labelText; + label.appendChild(document.createTextNode(labelText)); + if (tooltipText) { + label.appendChild(createMattingTooltipBadge(labelText, tooltipText)); + } row.appendChild(label); row.appendChild(control); - if (description) { - const hint = document.createElement('small'); - hint.className = 'lf-matting-settings-hint'; - hint.textContent = description; - row.appendChild(hint); - } return row; }; const modelSelect = document.createElement('select'); @@ -312,7 +345,7 @@ async function createCanvasWidget(node, widget, app) { modelLinks.replaceChildren(); const selectedOption = modelOptions.find((option) => option.path === modelSelect.value); if (modelSelect.value === 'auto') { - modelDescription.textContent = 'LayerForge selects the best compatible installed checkpoint. If none is available, the standard General model is downloaded.'; + modelDescription.textContent = 'Automatic mode picks the best compatible installed model and downloads the default model if needed.'; modelDetails.hidden = false; return; } @@ -337,7 +370,7 @@ async function createCanvasWidget(node, widget, app) { updateModelDetails(); const modeSelect = document.createElement('select'); modeSelect.className = 'lf-matting-settings-select'; - ['remove_background', 'remove_foreground', 'mask_only'].forEach((mode) => { + ['remove_background', 'remove_foreground', 'mask_only', 'mask_only_inverted'].forEach((mode) => { modeSelect.appendChild(new Option(getMattingModeLabel(mode), mode)); }); modeSelect.value = settings.mode; @@ -404,6 +437,7 @@ async function createCanvasWidget(node, widget, app) { thresholdInput.value = String(DEFAULT_MATTING_SETTINGS.threshold); thresholdValue.value = DEFAULT_MATTING_SETTINGS.threshold.toFixed(2); thresholdValue.textContent = DEFAULT_MATTING_SETTINGS.threshold.toFixed(2); + updateModelDetails(); }; const saveButton = document.createElement('button'); saveButton.type = 'button'; @@ -738,6 +772,11 @@ async function createCanvasWidget(node, widget, app) { const spinner = $el("div.lf-matting-spinner"); button.appendChild(spinner); button.classList.add('lf-loading'); + startMattingSpinner(); + if (modelStatus.reason === 'not_downloaded') { + setMattingDownloadProgress(0); + startMattingProgressPolling(); + } if (modelStatus.available) { showInfoNotification(`Starting ${getMattingModeLabel(mattingSettings.mode).toLowerCase()}...`, 2000); } @@ -755,6 +794,7 @@ async function createCanvasWidget(node, widget, app) { model_path: mattingSettings.modelPath || "auto", mode: mattingSettings.mode, threshold: mattingSettings.threshold, + node_id: String(node.id), }) }); const result = await response.json(); @@ -766,6 +806,10 @@ async function createCanvasWidget(node, widget, app) { showErrorNotification("Failed to download the matting model. Please check your internet connection and try again.", 8000); return; } + else if (result.error === "Matting Interrupted") { + showWarningNotification(result.details || "Matting was interrupted by ComfyUI.", 6000); + return; + } else if (result.error === "Matting Model Error") { showErrorNotification(result.details || "Model loading error. Please check the console for details.", 8000); return; @@ -778,7 +822,7 @@ async function createCanvasWidget(node, widget, app) { } throw new Error(errorMsg); } - if (mattingSettings.mode === 'mask_only') { + if (mattingSettings.mode === 'mask_only' || mattingSettings.mode === 'mask_only_inverted') { if (typeof result.draw_mask !== 'string') { throw new Error('Matting response did not contain a Draw Mask image.'); } @@ -786,7 +830,9 @@ async function createCanvasWidget(node, widget, app) { drawMaskImage.src = result.draw_mask; await drawMaskImage.decode(); canvas.maskTool.setMaskForLayer(drawMaskImage, selectedLayer); - showSuccessNotification('Generated mask applied to Draw Mask.'); + showSuccessNotification(mattingSettings.mode === 'mask_only_inverted' + ? 'Inverted mask applied to Draw Mask.' + : 'Generated mask applied to Draw Mask.'); return; } const mattedImage = new Image(); @@ -812,6 +858,9 @@ async function createCanvasWidget(node, widget, app) { } } finally { + stopMattingProgressPolling(); + stopMattingSpinner(); + setMattingDownloadProgress(null); button.classList.remove('lf-loading'); const spinner = button.querySelector('.lf-matting-spinner'); if (spinner && button.contains(spinner)) { @@ -1061,6 +1110,108 @@ async function createCanvasWidget(node, widget, app) { ]), $el("div.lf-painter-separator") ]); + const mattingButton = controlPanel.querySelector('.lf-matting-button'); + const mattingProgressTrack = document.createElement('span'); + mattingProgressTrack.className = 'lf-matting-download-progress'; + mattingProgressTrack.setAttribute('aria-hidden', 'true'); + const mattingProgressFill = document.createElement('span'); + mattingProgressFill.className = 'lf-matting-download-progress-fill'; + mattingProgressTrack.appendChild(mattingProgressFill); + mattingButton?.appendChild(mattingProgressTrack); + let mattingSpinnerAnimationFrame = null; + let mattingProgressPollTimer = null; + let mattingProgressPolling = false; + const stopMattingSpinner = () => { + if (mattingSpinnerAnimationFrame !== null) { + window.cancelAnimationFrame(mattingSpinnerAnimationFrame); + mattingSpinnerAnimationFrame = null; + } + const spinner = mattingButton?.querySelector('.lf-matting-spinner'); + if (spinner) { + spinner.style.transform = ''; + } + }; + const startMattingSpinner = () => { + stopMattingSpinner(); + const spinner = mattingButton?.querySelector('.lf-matting-spinner'); + if (!spinner) + return; + const startedAt = performance.now(); + const animate = (timestamp) => { + if (!mattingButton?.classList.contains('lf-loading') || !spinner.isConnected) { + stopMattingSpinner(); + return; + } + const rotation = ((timestamp - startedAt) * 0.45) % 360; + spinner.style.transform = `translateY(-50%) rotate(${rotation}deg)`; + mattingSpinnerAnimationFrame = window.requestAnimationFrame(animate); + }; + mattingSpinnerAnimationFrame = window.requestAnimationFrame(animate); + }; + const setMattingDownloadProgress = (progress) => { + if (!mattingButton) + return; + if (progress === null) { + mattingButton.classList.remove('lf-downloading'); + mattingProgressFill.style.width = '0%'; + return; + } + const normalizedProgress = Math.min(100, Math.max(0, Number(progress) || 0)); + mattingButton.classList.add('lf-downloading'); + mattingProgressFill.style.width = `${normalizedProgress}%`; + }; + const handleMattingStatus = (event) => { + const eventPayload = event?.detail && typeof event.detail === 'object' ? event.detail : event; + const payload = eventPayload?.data && typeof eventPayload.data === 'object' + ? eventPayload.data + : eventPayload; + if (!payload || typeof payload.status !== 'string') + return; + if (payload.node_id !== undefined && String(payload.node_id) !== String(node.id)) { + return; + } + if (payload.status === 'downloading') { + setMattingDownloadProgress(Number(payload.progress) || 0); + } + else if (payload.status === 'completed' || payload.status === 'error') { + setMattingDownloadProgress(null); + } + }; + const stopMattingProgressPolling = () => { + mattingProgressPolling = false; + if (mattingProgressPollTimer !== null) { + window.clearTimeout(mattingProgressPollTimer); + mattingProgressPollTimer = null; + } + }; + const pollMattingProgress = async () => { + if (!mattingProgressPolling) + return; + try { + const response = await fetch(`/matting/progress?node_id=${encodeURIComponent(String(node.id))}`, { cache: 'no-store' }); + if (response.ok) { + handleMattingStatus({ detail: await response.json() }); + } + } + catch { + // WebSocket events remain the primary path; polling is a best-effort fallback. + } + finally { + if (mattingProgressPolling) { + mattingProgressPollTimer = window.setTimeout(() => { + void pollMattingProgress(); + }, 250); + } + } + }; + const startMattingProgressPolling = () => { + stopMattingProgressPolling(); + mattingProgressPolling = true; + void pollMattingProgress(); + }; + if (mattingButton) { + api.addEventListener('matting_status', handleMattingStatus); + } // Function to create mask icon const createMaskIcon = () => { const iconContainer = document.createElement('div'); @@ -1483,6 +1634,11 @@ async function createCanvasWidget(node, widget, app) { panel: controlPanel, destroy: () => { closeMattingSettings(); + stopMattingProgressPolling(); + stopMattingSpinner(); + if (mattingButton) { + api.removeEventListener('matting_status', handleMattingStatus); + } mainContainer.removeEventListener('copy', stopEditableClipboardLeak); mainContainer.removeEventListener('cut', stopEditableClipboardLeak); mainContainer.removeEventListener('paste', stopEditableClipboardLeak); diff --git a/js/css/canvas_view.css b/js/css/canvas_view.css index cd6ffc7..377553f 100644 --- a/js/css/canvas_view.css +++ b/js/css/canvas_view.css @@ -601,6 +601,7 @@ .lf-painter-button.lf-matting-button { position: relative; transition: all 0.3s ease; + overflow: hidden; } .lf-painter-button.lf-matting-button.lf-loading { @@ -619,17 +620,37 @@ border: 2px solid rgba(255, 255, 255, 0.3); border-radius: 50%; border-top-color: #fff; - animation: matting-spin 1s linear infinite; + will-change: transform; } .lf-painter-button.lf-matting-button.lf-loading .lf-matting-spinner { display: block; } -@keyframes matting-spin { - to { - transform: translateY(-50%) rotate(360deg); - } +.lf-painter-button.lf-matting-button .lf-matting-download-progress { + display: none; + position: absolute; + left: 0; + right: 0; + bottom: 0; + height: 3px; + overflow: hidden; + background: rgba(255, 255, 255, 0.18); + pointer-events: none; + z-index: 1; +} + +.lf-painter-button.lf-matting-button.lf-downloading .lf-matting-download-progress { + display: block; +} + +.lf-painter-button.lf-matting-button .lf-matting-download-progress-fill { + display: block; + width: 0; + height: 100%; + background: #39d353; + box-shadow: 0 0 5px rgba(57, 211, 83, 0.85); + transition: width 0.12s linear; } .lf-matting-settings-button { @@ -710,11 +731,57 @@ } .lf-matting-settings-label { + display: inline-flex; + align-items: center; + gap: 6px; color: #fff; font-size: 13px; font-weight: 600; } +.lf-matting-tooltip-badge { + display: inline-flex; + align-items: center; + justify-content: center; + width: 18px; + height: 18px; + flex: 0 0 auto; + margin-left: 2px; + color: rgba(255, 255, 255, 0.7); + font-size: 11px; + font-weight: 700; + line-height: 1; + cursor: help; + user-select: none; + background: rgba(255, 255, 255, 0.03); + border: 1px solid rgba(255, 255, 255, 0.12); + border-radius: 50%; +} + +.lf-matting-tooltip-badge:hover, +.lf-matting-tooltip-badge:focus-visible { + color: #d9ecff; + outline: none; + background: rgba(78, 161, 255, 0.18); + border-color: rgba(78, 161, 255, 0.65); +} + +.lf-painter-tooltip.lf-matting-tooltip { + min-width: 220px; + max-width: min(320px, calc(100vw - 24px)); + padding: 8px 10px; + color: #f2f2f2; + font-size: 12px; + font-weight: 400; + line-height: 1.5; + white-space: pre-line; + overflow-wrap: anywhere; + background: rgba(34, 34, 38, 0.98); + border: 1px solid #555; + border-radius: 8px; + z-index: 1000001; +} + .lf-matting-settings-select { width: 100%; min-height: 34px; @@ -762,7 +829,6 @@ outline-offset: 1px; } -.lf-matting-settings-hint, .lf-matting-settings-status { margin: 0; color: #bdbdbd; diff --git a/python/matting/__init__.py b/python/matting/__init__.py index e180816..9922394 100644 --- a/python/matting/__init__.py +++ b/python/matting/__init__.py @@ -5,6 +5,7 @@ from aiohttp import web from .api import ( check_matting_model, + get_matting_progress, get_matting_settings, matting, register_matting_routes, @@ -69,6 +70,7 @@ __all__ = [ "BiRefNetMatting", "RMBG2Model", "check_matting_model", + "get_matting_progress", "get_matting_settings", "matting", "register_matting_routes", diff --git a/python/matting/api.py b/python/matting/api.py index 2433e5b..d7f4a9c 100644 --- a/python/matting/api.py +++ b/python/matting/api.py @@ -1,5 +1,6 @@ """HTTP endpoints for model status and background-removal execution.""" +import asyncio import traceback from aiohttp import web @@ -22,7 +23,8 @@ from .backends.rmbg import ( ) from .options import _get_birefnet_model_options from .paths import _get_birefnet_base_paths -from .service import BiRefNetMatting +from .progress import get_matting_status +from .service import BiRefNetMatting, MattingInterruptedError from .settings import get_public_settings, save_settings _matting_lock = None @@ -41,6 +43,11 @@ async def get_matting_settings(request): ) +async def get_matting_progress(request): + """Return the latest download status for the requested LayerForge node.""" + return web.json_response(get_matting_status(request.query.get("node_id"))) + + async def save_matting_settings(request): """Persist Matting preferences and an optional Hugging Face token.""" try: @@ -259,15 +266,18 @@ async def matting(request): mode = data.get("mode", "remove_background") model_path = data.get("model_path") or "auto" - matted_image, alpha_mask = matting_instance.execute( + node_id = data.get("node_id") + matted_image, alpha_mask = await asyncio.to_thread( + matting_instance.execute, image_tensor, model_path, threshold=data.get("threshold", 0.5), refinement=data.get("refinement", 1), mode=mode, + node_id=node_id, ) - if mode == "mask_only": + if mode in {"mask_only", "mask_only_inverted"}: result_image = convert_tensor_to_base64(alpha_mask) else: result_image = convert_tensor_to_base64(matted_image, alpha_mask, original_alpha) @@ -285,6 +295,15 @@ async def matting(request): "model_path": matting_instance.model_path, } ) + except MattingInterruptedError as error: + log.info("Matting was interrupted by ComfyUI") + return web.json_response( + { + "error": "Matting Interrupted", + "details": str(error), + }, + status=409, + ) except RuntimeError as error: log.error(f"Runtime error during matting: {error}") if _is_model_download_error(error): @@ -330,12 +349,14 @@ def register_matting_routes(): 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.get("/matting/progress")(get_matting_progress) PromptServer.instance.routes.post("/matting")(matting) __all__ = [ "check_matting_model", "get_matting_settings", + "get_matting_progress", "matting", "register_matting_routes", "save_matting_settings", diff --git a/python/matting/backends/birefnet.py b/python/matting/backends/birefnet.py index d68d875..192aec6 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 ..progress import call_huggingface_download from ..settings import get_huggingface_token @@ -140,7 +141,7 @@ def _find_existing_birefnet_default_checkpoint(): return None -def _download_birefnet_checkpoint(model=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 @@ -174,11 +175,12 @@ def _download_birefnet_checkpoint(model=None): token = get_huggingface_token() if token: download_kwargs["token"] = token - try: - downloaded_path = hf_hub_download(**download_kwargs) - except TypeError: - download_kwargs.pop("local_dir_use_symlinks", None) - downloaded_path = hf_hub_download(**download_kwargs) + 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( @@ -204,14 +206,16 @@ def _download_birefnet_checkpoint(model=None): return downloaded_path -def _ensure_birefnet_checkpoint(model_path=None): +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 - return _download_birefnet_checkpoint(remote_model) + 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() @@ -223,7 +227,9 @@ def _ensure_birefnet_checkpoint(model_path=None): if model_path and model_path != "auto": raise RuntimeError("The selected BiRefNet checkpoint is not available or is not compatible with ComfyUI.") - return _download_birefnet_checkpoint() + if node_id is None: + return _download_birefnet_checkpoint() + return _download_birefnet_checkpoint(node_id=node_id) __all__ = [ diff --git a/python/matting/backends/rmbg.py b/python/matting/backends/rmbg.py index 0971a29..1ee8e3c 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 ..progress import call_huggingface_download from ..settings import get_huggingface_token @@ -129,7 +130,7 @@ def _find_local_rmbg_model(model_path=None): return None -def _download_rmbg_model(model): +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 @@ -155,7 +156,12 @@ def _download_rmbg_model(model): if token: download_kwargs["token"] = token try: - snapshot_download(**download_kwargs) + 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 " @@ -172,14 +178,16 @@ def _download_rmbg_model(model): return model_directory -def _ensure_rmbg_model(model_path): +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 - return _download_rmbg_model(remote_model) + 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: diff --git a/python/matting/progress.py b/python/matting/progress.py new file mode 100644 index 0000000..d4d344c --- /dev/null +++ b/python/matting/progress.py @@ -0,0 +1,239 @@ +"""Progress reporting for background-removal model downloads.""" + +import threading +import time + + +_MIN_PROGRESS_INTERVAL = 0.1 +_MIN_PROGRESS_DELTA = 0.5 +_PROGRESS_STATE_TTL = 120.0 +_progress_state_lock = threading.Lock() +_progress_state = {} + + +def _progress_state_key(node_id): + return str(node_id) if node_id is not None else "_global" + + +def send_matting_status(status, *, node_id=None, **details): + """Send a best-effort Matting status event to the ComfyUI frontend.""" + payload = {"status": status} + if node_id is not None: + payload["node_id"] = str(node_id) + payload.update({key: value for key, value in details.items() if value is not None}) + + state_key = _progress_state_key(node_id) + with _progress_state_lock: + _progress_state[state_key] = { + **payload, + "updated_at": time.monotonic(), + } + + try: + from server import PromptServer + + PromptServer.instance.send_sync("matting_status", payload) + except Exception: + # Progress reporting must never interrupt a model download or inference. + pass + + +def get_matting_status(node_id=None): + """Return the latest in-process status for a Matting operation.""" + state_key = _progress_state_key(node_id) + now = time.monotonic() + with _progress_state_lock: + state = _progress_state.get(state_key) + if state is None: + return {"status": "idle"} + if now - state["updated_at"] > _PROGRESS_STATE_TTL: + _progress_state.pop(state_key, None) + return {"status": "idle"} + return {key: value for key, value in state.items() if key != "updated_at"} + + +class _FallbackTqdm: + """Small fallback used only when an older environment lacks tqdm imports.""" + + def __init__(self, *args, **kwargs): + del args + self.total = kwargs.get("total") + self.n = kwargs.get("initial", 0) + self.unit = kwargs.get("unit") or "it" + + def update(self, amount=1): + self.n += amount + + def close(self): + return None + + def refresh(self): + return None + + def set_description(self, *args, **kwargs): + del args, kwargs + + def __enter__(self): + return self + + def __exit__(self, *args): + del args + self.close() + + +def _get_tqdm_base(): + try: + from huggingface_hub.utils import tqdm as huggingface_tqdm + + return huggingface_tqdm + except Exception: + try: + from tqdm.auto import tqdm + + return tqdm + except Exception: + return _FallbackTqdm + + +class MattingProgressAdapter: + """Track downloaded bytes independently of tqdm's terminal state.""" + + def __init__(self, model_label, node_id=None, total_size=0, initial=0): + self.model_label = model_label + self.node_id = node_id + self.total_size = self._coerce_bytes(total_size) + self.downloaded = self._coerce_bytes(initial) + self.last_reported = -1.0 + self.last_report_time = 0.0 + self.last_update_time = time.monotonic() + self.last_downloaded = self.downloaded + self._publish(force=True) + + @staticmethod + def _coerce_bytes(value): + try: + return max(0, int(float(value or 0))) + except (TypeError, ValueError): + return 0 + + def set_total(self, total_size): + total = self._coerce_bytes(total_size) + if total > self.total_size: + self.total_size = total + + def _publish(self, downloaded=None, speed=0, force=False): + if downloaded is not None: + self.downloaded = max(self.downloaded, self._coerce_bytes(downloaded)) + if self.total_size > 0: + self.downloaded = min(self.downloaded, self.total_size) + + progress = ( + min(100.0, round(self.downloaded / self.total_size * 100, 1)) + if self.total_size > 0 + else 0.0 + ) + + now = time.monotonic() + if not force and progress < 100.0: + if now - self.last_report_time < _MIN_PROGRESS_INTERVAL: + if abs(progress - self.last_reported) < _MIN_PROGRESS_DELTA: + return + + self.last_reported = progress + self.last_report_time = now + payload = { + "model": self.model_label, + "progress": progress, + "downloaded_bytes": self.downloaded, + "total_bytes": self.total_size, + "speed": max(0, self._coerce_bytes(speed)), + } + send_matting_status("downloading", node_id=self.node_id, **payload) + + def update(self, byte_delta=1): + """Accept a byte delta from an HTTP or Xet transfer.""" + delta = self._coerce_bytes(byte_delta) + now = time.monotonic() + elapsed = now - self.last_update_time + downloaded = self.downloaded + delta + speed = (downloaded - self.last_downloaded) / elapsed if elapsed > 0 else 0 + self.last_update_time = now + self.last_downloaded = downloaded + self._publish(downloaded, speed=speed) + + def close(self): + self._publish(force=True) + + +def create_huggingface_tqdm_class(model_label, node_id=None): + """Create a Hugging Face-compatible tqdm class that reports byte progress.""" + base_tqdm = _get_tqdm_base() + + class MattingProgressTqdm(base_tqdm): + def __init__(self, *args, **kwargs): + self._matting_unit = kwargs.get("unit") + self._matting_progress = None + super().__init__(*args, **kwargs) + if self._is_byte_unit(): + self._matting_progress = MattingProgressAdapter( + model_label, + node_id, + total_size=getattr(self, "total", kwargs.get("total", 0)), + initial=kwargs.get("initial", 0), + ) + + def update(self, amount=1): + result = super().update(amount) + if self._matting_progress is not None: + self._matting_progress.set_total(getattr(self, "total", 0)) + self._matting_progress.update(amount) + return result + + def close(self): + if self._matting_progress is not None: + self._matting_progress.set_total(getattr(self, "total", 0)) + self._matting_progress.close() + return super().close() + + def _is_byte_unit(self): + unit = str( + self._matting_unit + or getattr(self, "unit", None) + or getattr(self, "_unit", None) + or "" + ).lower() + return unit in {"b", "bytes"} + + return MattingProgressTqdm + + +def call_huggingface_download(download_function, download_kwargs, model_label, node_id=None): + """Call a Hub downloader with progress, falling back for older Hub APIs.""" + call_kwargs = dict(download_kwargs) + call_kwargs["tqdm_class"] = create_huggingface_tqdm_class(model_label, node_id) + + while True: + try: + return download_function(**call_kwargs) + except TypeError as error: + error_text = str(error) + unsupported_keyword = next( + ( + keyword + for keyword in ("tqdm_class", "local_dir_use_symlinks") + if keyword in call_kwargs and keyword in error_text + ), + None, + ) + if unsupported_keyword is None: + raise + call_kwargs.pop(unsupported_keyword) + + +__all__ = [ + "call_huggingface_download", + "create_huggingface_tqdm_class", + "get_matting_status", + "MattingProgressAdapter", + "send_matting_status", +] diff --git a/python/matting/service.py b/python/matting/service.py index b345440..47f3126 100644 --- a/python/matting/service.py +++ b/python/matting/service.py @@ -6,7 +6,6 @@ from typing import ClassVar import torch import torch.nn.functional as F -from server import PromptServer from torchvision import transforms from ..node import log @@ -19,6 +18,19 @@ from .backends.rmbg import ( _get_rmbg_model_status_message, _get_rmbg_remote_model, ) +from .progress import send_matting_status + + +class MattingInterruptedError(RuntimeError): + """Raised when ComfyUI interrupts the native model during inference.""" + + +def _is_comfyui_processing_interrupt(error): + try: + from comfy.model_management import InterruptProcessingException + except (ImportError, AttributeError): + return False + return isinstance(error, InterruptProcessingException) class BiRefNetMatting: @@ -31,7 +43,7 @@ class BiRefNetMatting: self.model = None self.model_path = None - def load_model(self, model_path=None): + def load_model(self, model_path=None, node_id=None): """Load the selected backend and reuse its process-wide model cache.""" rmbg_model = _get_rmbg_remote_model(model_path) local_rmbg_path = None @@ -41,7 +53,10 @@ class BiRefNetMatting: if rmbg_model or local_rmbg_path: if _get_rmbg_model_loader() is None: raise RuntimeError(_get_rmbg_model_status_message()) - model_directory = _ensure_rmbg_model(model_path) + if node_id is None: + model_directory = _ensure_rmbg_model(model_path) + else: + model_directory = _ensure_rmbg_model(model_path, node_id=node_id) self.model = RMBG2Model.load(model_directory) self.model_path = model_directory return self.model @@ -52,7 +67,10 @@ class BiRefNetMatting: "This ComfyUI version does not provide the native BiRefNet background-removal loader." ) - checkpoint_path = _ensure_birefnet_checkpoint(model_path) + if node_id is None: + checkpoint_path = _ensure_birefnet_checkpoint(model_path) + else: + checkpoint_path = _ensure_birefnet_checkpoint(model_path, node_id=node_id) with self._model_cache_lock: if checkpoint_path not in self._model_cache: log.info(f"Loading BiRefNet with ComfyUI's native loader from {checkpoint_path}") @@ -92,11 +110,19 @@ class BiRefNetMatting: return image.to(dtype=torch.float32).contiguous() - def execute(self, image, model_path, threshold=0.5, refinement=1, mode="remove_background"): + def execute( + self, + image, + model_path, + threshold=0.5, + refinement=1, + mode="remove_background", + node_id=None, + ): try: - PromptServer.instance.send_sync("matting_status", {"status": "processing"}) + send_matting_status("processing", node_id=node_id) del refinement - if mode not in {"remove_background", "remove_foreground", "mask_only"}: + if mode not in {"remove_background", "remove_foreground", "mask_only", "mask_only_inverted"}: raise ValueError(f"Unsupported matting mode: {mode}") threshold = float(threshold) @@ -106,7 +132,7 @@ class BiRefNetMatting: image_tensor = self.preprocess_image(image) original_size = (image_tensor.shape[1], image_tensor.shape[2]) log.debug(f"Original size: {original_size}") - self.load_model(model_path) + self.load_model(model_path, node_id=node_id) log.debug(f"Processed image shape: {image_tensor.shape}, dtype: {image_tensor.dtype}") with torch.no_grad(): @@ -132,16 +158,23 @@ class BiRefNetMatting: if threshold > 0: result = (result > threshold).to(dtype=torch.float32) - alpha_mask = 1.0 - result if mode == "remove_foreground" else result - if mode == "mask_only": + alpha_mask = 1.0 - result if mode in {"remove_foreground", "mask_only_inverted"} else result + if mode in {"mask_only", "mask_only_inverted"}: masked_image = alpha_mask.expand(-1, 3, -1, -1) else: masked_image = image_tensor.movedim(-1, 1) * alpha_mask - PromptServer.instance.send_sync("matting_status", {"status": "completed"}) + send_matting_status("completed", node_id=node_id) return masked_image, alpha_mask except Exception: - PromptServer.instance.send_sync("matting_status", {"status": "error"}) + send_matting_status("error", node_id=node_id) + raise + except BaseException as error: + if _is_comfyui_processing_interrupt(error): + send_matting_status("error", node_id=node_id, reason="interrupted") + raise MattingInterruptedError( + "ComfyUI interrupted the Matting operation." + ) from error raise @classmethod @@ -155,4 +188,4 @@ class BiRefNetMatting: return digest.hexdigest() -__all__ = ["BiRefNetMatting"] +__all__ = ["BiRefNetMatting", "MattingInterruptedError"] diff --git a/python/matting/settings.py b/python/matting/settings.py index 41921a2..66777e1 100644 --- a/python/matting/settings.py +++ b/python/matting/settings.py @@ -17,7 +17,7 @@ DEFAULT_SETTINGS = { "threshold": 0.5, "hf_token": "", } -_VALID_MODES = {"remove_background", "remove_foreground", "mask_only"} +_VALID_MODES = {"remove_background", "remove_foreground", "mask_only", "mask_only_inverted"} _SETTINGS_LOCK = threading.RLock() diff --git a/src/CanvasView.ts b/src/CanvasView.ts index b8c6822..ae099c2 100644 --- a/src/CanvasView.ts +++ b/src/CanvasView.ts @@ -34,7 +34,7 @@ import type { ComfyNode, Layer, AddMode } from './types'; const log = createModuleLogger('Canvas_view'); -type MattingMode = 'remove_background' | 'remove_foreground' | 'mask_only'; +type MattingMode = 'remove_background' | 'remove_foreground' | 'mask_only' | 'mask_only_inverted'; interface MattingSettings { modelPath: string; @@ -62,7 +62,10 @@ const DEFAULT_MATTING_SETTINGS: MattingSettings = { }; const isMattingMode = (value: unknown): value is MattingMode => { - return value === 'remove_background' || value === 'remove_foreground' || value === 'mask_only'; + return value === 'remove_background' + || value === 'remove_foreground' + || value === 'mask_only' + || value === 'mask_only_inverted'; }; const normalizeMattingSettings = (settings: Partial): MattingSettings => { @@ -122,6 +125,8 @@ const getMattingModeLabel = (mode: MattingMode): string => { return 'Remove detected foreground / keep background'; case 'mask_only': return 'Apply generated mask to Draw Mask'; + case 'mask_only_inverted': + return 'Apply inverted mask to Draw Mask'; default: return 'Remove background / keep foreground'; } @@ -278,10 +283,43 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): helpTooltip.style.display = 'none'; }; + const showMattingTooltip = (target: HTMLElement, content: string): void => { + helpTooltip.classList.add('lf-matting-tooltip'); + showTooltip(target, content); + }; + + const hideMattingTooltip = (): void => { + hideTooltip(); + helpTooltip.classList.remove('lf-matting-tooltip'); + }; + + const createMattingTooltipBadge = (labelText: string, tooltipText: string): HTMLSpanElement => { + const badge = document.createElement('span'); + badge.className = 'lf-matting-tooltip-badge'; + badge.textContent = '?'; + badge.tabIndex = 0; + badge.dataset.tooltip = tooltipText; + badge.setAttribute('aria-label', `More information about ${labelText}`); + + const show = (): void => { + const content = badge.dataset.tooltip; + if (content) showMattingTooltip(badge, content); + }; + const hide = (): void => hideMattingTooltip(); + + badge.addEventListener('mouseenter', show); + badge.addEventListener('focus', show); + badge.addEventListener('mouseleave', hide); + badge.addEventListener('blur', hide); + badge.addEventListener('click', (event) => event.preventDefault()); + return badge; + }; + let mattingSettingsBackdrop: HTMLDivElement | null = null; let mattingSettingsEscapeHandler: ((event: KeyboardEvent) => void) | null = null; const closeMattingSettings = () => { + hideMattingTooltip(); if (mattingSettingsEscapeHandler) { document.removeEventListener('keydown', mattingSettingsEscapeHandler); mattingSettingsEscapeHandler = null; @@ -344,23 +382,19 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): const body = document.createElement('div'); body.className = 'lf-matting-settings-body'; - const createRow = (labelText: string, control: HTMLElement, description?: string): HTMLLabelElement => { + const createRow = (labelText: string, control: HTMLElement, tooltipText?: string): HTMLLabelElement => { const row = document.createElement('label'); row.className = 'lf-matting-settings-row'; const label = document.createElement('span'); label.className = 'lf-matting-settings-label'; - label.textContent = labelText; + label.appendChild(document.createTextNode(labelText)); + if (tooltipText) { + label.appendChild(createMattingTooltipBadge(labelText, tooltipText)); + } row.appendChild(label); row.appendChild(control); - if (description) { - const hint = document.createElement('small'); - hint.className = 'lf-matting-settings-hint'; - hint.textContent = description; - row.appendChild(hint); - } - return row; }; @@ -420,7 +454,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): const selectedOption = modelOptions.find((option) => option.path === modelSelect.value); if (modelSelect.value === 'auto') { - modelDescription.textContent = 'LayerForge selects the best compatible installed checkpoint. If none is available, the standard General model is downloaded.'; + modelDescription.textContent = 'Automatic mode picks the best compatible installed model and downloads the default model if needed.'; modelDetails.hidden = false; return; } @@ -454,7 +488,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): const modeSelect = document.createElement('select'); modeSelect.className = 'lf-matting-settings-select'; - (['remove_background', 'remove_foreground', 'mask_only'] as MattingMode[]).forEach((mode) => { + (['remove_background', 'remove_foreground', 'mask_only', 'mask_only_inverted'] as MattingMode[]).forEach((mode) => { modeSelect.appendChild(new Option(getMattingModeLabel(mode), mode)); }); modeSelect.value = settings.mode; @@ -539,6 +573,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): thresholdInput.value = String(DEFAULT_MATTING_SETTINGS.threshold); thresholdValue.value = DEFAULT_MATTING_SETTINGS.threshold.toFixed(2); thresholdValue.textContent = DEFAULT_MATTING_SETTINGS.threshold.toFixed(2); + updateModelDetails(); }; const saveButton = document.createElement('button'); @@ -932,6 +967,11 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): const spinner = $el("div.lf-matting-spinner") as HTMLDivElement; button.appendChild(spinner); button.classList.add('lf-loading'); + startMattingSpinner(); + if (modelStatus.reason === 'not_downloaded') { + setMattingDownloadProgress(0); + startMattingProgressPolling(); + } if (modelStatus.available) { showInfoNotification(`Starting ${getMattingModeLabel(mattingSettings.mode).toLowerCase()}...`, 2000); @@ -952,6 +992,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): model_path: mattingSettings.modelPath || "auto", mode: mattingSettings.mode, threshold: mattingSettings.threshold, + node_id: String(node.id), }) }); @@ -964,6 +1005,9 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): if (result.error === "Network Connection Error") { showErrorNotification("Failed to download the matting model. Please check your internet connection and try again.", 8000); return; + } else if (result.error === "Matting Interrupted") { + showWarningNotification(result.details || "Matting was interrupted by ComfyUI.", 6000); + return; } else if (result.error === "Matting Model Error") { showErrorNotification(result.details || "Model loading error. Please check the console for details.", 8000); return; @@ -976,7 +1020,7 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): throw new Error(errorMsg); } - if (mattingSettings.mode === 'mask_only') { + if (mattingSettings.mode === 'mask_only' || mattingSettings.mode === 'mask_only_inverted') { if (typeof result.draw_mask !== 'string') { throw new Error('Matting response did not contain a Draw Mask image.'); } @@ -985,7 +1029,11 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): drawMaskImage.src = result.draw_mask; await drawMaskImage.decode(); canvas.maskTool.setMaskForLayer(drawMaskImage, selectedLayer); - showSuccessNotification('Generated mask applied to Draw Mask.'); + showSuccessNotification( + mattingSettings.mode === 'mask_only_inverted' + ? 'Inverted mask applied to Draw Mask.' + : 'Generated mask applied to Draw Mask.', + ); return; } @@ -1015,6 +1063,9 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): showErrorNotification(`Matting Failed: ${errorMessage}`); } } finally { + stopMattingProgressPolling(); + stopMattingSpinner(); + setMattingDownloadProgress(null); button.classList.remove('lf-loading'); const spinner = button.querySelector('.lf-matting-spinner'); if (spinner && button.contains(spinner)) { @@ -1263,6 +1314,122 @@ $el("label.lf-clipboard-switch.lf-mask-switch", { $el("div.lf-painter-separator") ]); + const mattingButton = controlPanel.querySelector('.lf-matting-button') as HTMLButtonElement | null; + const mattingProgressTrack = document.createElement('span'); + mattingProgressTrack.className = 'lf-matting-download-progress'; + mattingProgressTrack.setAttribute('aria-hidden', 'true'); + const mattingProgressFill = document.createElement('span'); + mattingProgressFill.className = 'lf-matting-download-progress-fill'; + mattingProgressTrack.appendChild(mattingProgressFill); + mattingButton?.appendChild(mattingProgressTrack); + + let mattingSpinnerAnimationFrame: number | null = null; + let mattingProgressPollTimer: number | null = null; + let mattingProgressPolling = false; + + const stopMattingSpinner = (): void => { + if (mattingSpinnerAnimationFrame !== null) { + window.cancelAnimationFrame(mattingSpinnerAnimationFrame); + mattingSpinnerAnimationFrame = null; + } + + const spinner = mattingButton?.querySelector('.lf-matting-spinner') as HTMLDivElement | null; + if (spinner) { + spinner.style.transform = ''; + } + }; + + const startMattingSpinner = (): void => { + stopMattingSpinner(); + const spinner = mattingButton?.querySelector('.lf-matting-spinner') as HTMLDivElement | null; + if (!spinner) return; + + const startedAt = performance.now(); + const animate = (timestamp: number): void => { + if (!mattingButton?.classList.contains('lf-loading') || !spinner.isConnected) { + stopMattingSpinner(); + return; + } + + const rotation = ((timestamp - startedAt) * 0.45) % 360; + spinner.style.transform = `translateY(-50%) rotate(${rotation}deg)`; + mattingSpinnerAnimationFrame = window.requestAnimationFrame(animate); + }; + + mattingSpinnerAnimationFrame = window.requestAnimationFrame(animate); + }; + + const setMattingDownloadProgress = (progress: number | null): void => { + if (!mattingButton) return; + + if (progress === null) { + mattingButton.classList.remove('lf-downloading'); + mattingProgressFill.style.width = '0%'; + return; + } + + const normalizedProgress = Math.min(100, Math.max(0, Number(progress) || 0)); + mattingButton.classList.add('lf-downloading'); + mattingProgressFill.style.width = `${normalizedProgress}%`; + }; + + const handleMattingStatus = (event: any): void => { + const eventPayload = event?.detail && typeof event.detail === 'object' ? event.detail : event; + const payload = eventPayload?.data && typeof eventPayload.data === 'object' + ? eventPayload.data + : eventPayload; + if (!payload || typeof payload.status !== 'string') return; + + if (payload.node_id !== undefined && String(payload.node_id) !== String(node.id)) { + return; + } + + if (payload.status === 'downloading') { + setMattingDownloadProgress(Number(payload.progress) || 0); + } else if (payload.status === 'completed' || payload.status === 'error') { + setMattingDownloadProgress(null); + } + }; + + const stopMattingProgressPolling = (): void => { + mattingProgressPolling = false; + if (mattingProgressPollTimer !== null) { + window.clearTimeout(mattingProgressPollTimer); + mattingProgressPollTimer = null; + } + }; + + const pollMattingProgress = async (): Promise => { + if (!mattingProgressPolling) return; + + try { + const response = await fetch( + `/matting/progress?node_id=${encodeURIComponent(String(node.id))}`, + { cache: 'no-store' }, + ); + if (response.ok) { + handleMattingStatus({ detail: await response.json() }); + } + } catch { + // WebSocket events remain the primary path; polling is a best-effort fallback. + } finally { + if (mattingProgressPolling) { + mattingProgressPollTimer = window.setTimeout(() => { + void pollMattingProgress(); + }, 250); + } + } + }; + + const startMattingProgressPolling = (): void => { + stopMattingProgressPolling(); + mattingProgressPolling = true; + void pollMattingProgress(); + }; + + if (mattingButton) { + api.addEventListener('matting_status', handleMattingStatus); + } // Function to create mask icon const createMaskIcon = (): HTMLElement => { @@ -1761,6 +1928,11 @@ $el("label.lf-clipboard-switch.lf-mask-switch", { panel: controlPanel, destroy: () => { closeMattingSettings(); + stopMattingProgressPolling(); + stopMattingSpinner(); + if (mattingButton) { + api.removeEventListener('matting_status', handleMattingStatus); + } mainContainer.removeEventListener('copy', stopEditableClipboardLeak); mainContainer.removeEventListener('cut', stopEditableClipboardLeak); mainContainer.removeEventListener('paste', stopEditableClipboardLeak); diff --git a/src/css/canvas_view.css b/src/css/canvas_view.css index cd6ffc7..377553f 100644 --- a/src/css/canvas_view.css +++ b/src/css/canvas_view.css @@ -601,6 +601,7 @@ .lf-painter-button.lf-matting-button { position: relative; transition: all 0.3s ease; + overflow: hidden; } .lf-painter-button.lf-matting-button.lf-loading { @@ -619,17 +620,37 @@ border: 2px solid rgba(255, 255, 255, 0.3); border-radius: 50%; border-top-color: #fff; - animation: matting-spin 1s linear infinite; + will-change: transform; } .lf-painter-button.lf-matting-button.lf-loading .lf-matting-spinner { display: block; } -@keyframes matting-spin { - to { - transform: translateY(-50%) rotate(360deg); - } +.lf-painter-button.lf-matting-button .lf-matting-download-progress { + display: none; + position: absolute; + left: 0; + right: 0; + bottom: 0; + height: 3px; + overflow: hidden; + background: rgba(255, 255, 255, 0.18); + pointer-events: none; + z-index: 1; +} + +.lf-painter-button.lf-matting-button.lf-downloading .lf-matting-download-progress { + display: block; +} + +.lf-painter-button.lf-matting-button .lf-matting-download-progress-fill { + display: block; + width: 0; + height: 100%; + background: #39d353; + box-shadow: 0 0 5px rgba(57, 211, 83, 0.85); + transition: width 0.12s linear; } .lf-matting-settings-button { @@ -710,11 +731,57 @@ } .lf-matting-settings-label { + display: inline-flex; + align-items: center; + gap: 6px; color: #fff; font-size: 13px; font-weight: 600; } +.lf-matting-tooltip-badge { + display: inline-flex; + align-items: center; + justify-content: center; + width: 18px; + height: 18px; + flex: 0 0 auto; + margin-left: 2px; + color: rgba(255, 255, 255, 0.7); + font-size: 11px; + font-weight: 700; + line-height: 1; + cursor: help; + user-select: none; + background: rgba(255, 255, 255, 0.03); + border: 1px solid rgba(255, 255, 255, 0.12); + border-radius: 50%; +} + +.lf-matting-tooltip-badge:hover, +.lf-matting-tooltip-badge:focus-visible { + color: #d9ecff; + outline: none; + background: rgba(78, 161, 255, 0.18); + border-color: rgba(78, 161, 255, 0.65); +} + +.lf-painter-tooltip.lf-matting-tooltip { + min-width: 220px; + max-width: min(320px, calc(100vw - 24px)); + padding: 8px 10px; + color: #f2f2f2; + font-size: 12px; + font-weight: 400; + line-height: 1.5; + white-space: pre-line; + overflow-wrap: anywhere; + background: rgba(34, 34, 38, 0.98); + border: 1px solid #555; + border-radius: 8px; + z-index: 1000001; +} + .lf-matting-settings-select { width: 100%; min-height: 34px; @@ -762,7 +829,6 @@ outline-offset: 1px; } -.lf-matting-settings-hint, .lf-matting-settings-status { margin: 0; color: #bdbdbd; diff --git a/tests/test_entrypoint_integration.py b/tests/test_entrypoint_integration.py index a7309fc..1e041a8 100644 --- a/tests/test_entrypoint_integration.py +++ b/tests/test_entrypoint_integration.py @@ -118,6 +118,7 @@ def _import_layerforge(monkeypatch, tmp_path): 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"], + matting_progress=sys.modules[f"{package_name}.python.matting.progress"], routes=routes, package_name=package_name, ) @@ -173,6 +174,7 @@ def test_entrypoint_registers_backend_route_contract(layerforge_runtime): ("GET", "/matting/settings"), ("POST", "/matting/settings"), ("GET", "/matting/check-model"), + ("GET", "/matting/progress"), ("POST", "/matting"), } @@ -187,14 +189,14 @@ def test_matting_settings_are_persisted_without_exposing_the_token(layerforge_ru saved = settings_module.save_settings( { "model_path": "remote:rmbg_2_0", - "mode": "mask_only", + "mode": "mask_only_inverted", "threshold": 0.75, "hf_token": "hf-test-token", } ) assert saved["model_path"] == "remote:rmbg_2_0" - assert saved["mode"] == "mask_only" + assert saved["mode"] == "mask_only_inverted" assert saved["threshold"] == 0.75 assert settings_module.get_huggingface_token() == "hf-test-token" public_settings = settings_module.get_public_settings() @@ -402,6 +404,48 @@ def test_matting_adapter_uses_native_loader_and_bhwc_input(layerforge_runtime, m assert tuple(alpha_mask.shape) == (1, 1, 2, 4) +def test_matting_adapter_translates_comfyui_interrupt(layerforge_runtime, monkeypatch): + import torch + + node_module = layerforge_runtime.matting + service_module = layerforge_runtime.matting_service + + class FakeInterruptProcessingException(BaseException): + pass + + comfy_module = ModuleType("comfy") + model_management = ModuleType("comfy.model_management") + model_management.InterruptProcessingException = FakeInterruptProcessingException + comfy_module.model_management = model_management + monkeypatch.setitem(sys.modules, "comfy", comfy_module) + monkeypatch.setitem(sys.modules, "comfy.model_management", model_management) + + class InterruptingBiRefNet: + def encode_image(self, image): + del image + raise FakeInterruptProcessingException() + + monkeypatch.setattr( + service_module, + "_get_comfy_birefnet_loader", + lambda: lambda path: InterruptingBiRefNet(), + ) + monkeypatch.setattr( + service_module, + "_ensure_birefnet_checkpoint", + lambda model_path=None: "interrupting-birefnet.safetensors", + ) + node_module.BiRefNetMatting._model_cache.clear() + + with pytest.raises(service_module.MattingInterruptedError, match="interrupted"): + node_module.BiRefNetMatting().execute( + torch.ones((1, 3, 2, 4), dtype=torch.float32), + model_path=None, + threshold=0, + refinement=1, + ) + + def test_matting_model_status_preserves_unsupported_and_error_responses(layerforge_runtime, monkeypatch): node_module = layerforge_runtime.matting api_module = layerforge_runtime.matting_api @@ -591,11 +635,20 @@ def test_matting_adapter_supports_inverted_and_mask_only_modes(layerforge_runtim refinement=1, mode="mask_only", ) + inverted_mask_preview, inverted_preview_mask = matting.execute( + image, + model_path=None, + threshold=0.5, + refinement=1, + mode="mask_only_inverted", + ) assert torch.allclose(removed_foreground, torch.zeros_like(removed_foreground)) assert torch.allclose(inverted_mask, torch.zeros_like(inverted_mask)) assert torch.allclose(mask_preview, torch.ones_like(mask_preview)) assert torch.allclose(preview_mask, torch.ones_like(preview_mask)) + assert torch.allclose(inverted_mask_preview, torch.zeros_like(inverted_mask_preview)) + assert torch.allclose(inverted_preview_mask, torch.zeros_like(inverted_preview_mask)) def test_matting_model_options_include_downloadable_official_variants(layerforge_runtime, monkeypatch): @@ -683,6 +736,52 @@ def test_remote_matting_download_uses_background_removal_root(layerforge_runtime assert Path(result).exists() +def test_model_download_progress_reports_bytes_to_the_frontend(layerforge_runtime, monkeypatch): + progress_module = layerforge_runtime.matting_progress + events = [] + monkeypatch.setattr( + sys.modules["server"].PromptServer.instance, + "send_sync", + lambda *args: events.append(args), + ) + + progress_class = progress_module.create_huggingface_tqdm_class("Test model", node_id="42") + progress = progress_class(total=100, initial=0, unit="B", disable=True) + progress.update(25) + progress.update(75) + progress.close() + + download_events = [event for event in events if event[1]["status"] == "downloading"] + assert download_events + assert download_events[-1][1]["node_id"] == "42" + assert download_events[-1][1]["progress"] == 100.0 + assert download_events[-1][1]["downloaded_bytes"] == 100 + assert download_events[-1][1]["total_bytes"] == 100 + progress_module.send_matting_status("completed", node_id="42") + assert progress_module.get_matting_status("42")["status"] == "completed" + + +def test_matting_progress_status_is_available_without_websocket_delivery(layerforge_runtime): + progress_module = layerforge_runtime.matting_progress + + progress_module.send_matting_status( + "downloading", + node_id="poll-node", + progress=37.5, + downloaded_bytes=375, + total_bytes=1000, + ) + + assert progress_module.get_matting_status("poll-node") == { + "status": "downloading", + "node_id": "poll-node", + "progress": 37.5, + "downloaded_bytes": 375, + "total_bytes": 1000, + } + progress_module.send_matting_status("completed", node_id="poll-node") + + def test_rmbg_model_path_uses_background_removal_subdirectory(layerforge_runtime): node_module = layerforge_runtime.matting model = node_module._RMBG_MODEL_CATALOG[0]