Add matting download progress, interrupt handling, and inverted mask mode

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",

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",

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__ = [

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:

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",
+]

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"]

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()