From e4441dc832f8d257533b63dc5e04227db3c66782 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 13 Aug 2026 17:57:47 +0200 Subject: [PATCH] feat: improve BRIA RMBG 2.0 transformers compatibility and error handling --- pyproject.toml | 2 +- python/matting/__init__.py | 4 ++ python/matting/api.py | 3 +- python/matting/backends/rmbg.py | 97 +++++++++++++++++++++++----- python/matting/service.py | 6 +- requirements.txt | 2 +- src/CanvasView.ts | 4 ++ tests/test_entrypoint_integration.py | 15 +++++ 8 files changed, 110 insertions(+), 23 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c8b06dd..d0251e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ description = "Photoshop-like layered canvas editor to your ComfyUI workflow. Th version = "1.5.13" license = { text = "MIT License" } requires-python = ">=3.10" -dependencies = ["torch", "torchvision", "huggingface_hub", "aiohttp", "numpy", "tqdm", "Pillow", "transformers>=4.50.3", "kornia"] +dependencies = ["torch", "torchvision", "huggingface_hub", "aiohttp", "numpy", "tqdm", "Pillow", "transformers", "kornia"] [project.urls] Repository = "https://github.com/Azornes/Comfyui-LayerForge" diff --git a/python/matting/__init__.py b/python/matting/__init__.py index 9738f9e..70a4adf 100644 --- a/python/matting/__init__.py +++ b/python/matting/__init__.py @@ -22,6 +22,8 @@ from .backends.rmbg import ( _find_existing_rmbg_model, _find_local_rmbg_model, _get_rmbg_model_loader, + _get_rmbg_model_status_message, + _get_rmbg_transformers_status, _get_rmbg_remote_model, _is_rmbg_model_directory, ) @@ -79,6 +81,8 @@ __all__ = [ "_get_comfy_birefnet_loader", "_get_rmbg_model_directory", "_get_rmbg_model_loader", + "_get_rmbg_model_status_message", + "_get_rmbg_transformers_status", "_get_rmbg_remote_model", "_is_native_birefnet_checkpoint", "_is_rmbg_model_directory", diff --git a/python/matting/api.py b/python/matting/api.py index 6486ed2..4ad84ca 100644 --- a/python/matting/api.py +++ b/python/matting/api.py @@ -17,6 +17,7 @@ from .backends.rmbg import ( _find_existing_rmbg_model, _find_local_rmbg_model, _get_rmbg_model_loader, + _get_rmbg_model_status_message, _get_rmbg_remote_model, ) from .options import _get_birefnet_model_options @@ -62,7 +63,7 @@ async def check_matting_model(request): return _matting_status_response( False, "unsupported_rmbg", - "BRIA RMBG 2.0 requires the 'transformers' package in the active ComfyUI environment.", + _get_rmbg_model_status_message(), model_path=requested_model, models=model_options, ) diff --git a/python/matting/backends/rmbg.py b/python/matting/backends/rmbg.py index 0f91470..40a00b8 100644 --- a/python/matting/backends/rmbg.py +++ b/python/matting/backends/rmbg.py @@ -12,15 +12,72 @@ from ..catalog import _RMBG_MODEL_CATALOG, _RMBG_REMOTE_PREFIX from ..paths import _get_rmbg_model_directory +def _get_rmbg_transformers_status(): + """Inspect the installed Transformers API without changing the environment.""" + try: + import transformers + except ModuleNotFoundError as error: + if error.name == "transformers": + return { + "loader": None, + "message": ( + "BRIA RMBG 2.0 requires the 'transformers' package in the active ComfyUI environment. " + "Install or update it manually if you want to use this model." + ), + } + return { + "loader": None, + "message": ( + "The installed Transformers package could not be loaded for BRIA RMBG 2.0: " + f"{error}. Check the active ComfyUI environment and update Transformers manually if needed." + ), + } + except Exception as error: + return { + "loader": None, + "message": ( + "The installed Transformers package could not be loaded for BRIA RMBG 2.0: " + f"{error}. Update Transformers manually if needed." + ), + } + + version = getattr(transformers, "__version__", "unknown") + try: + loader = getattr(transformers, "AutoModelForImageSegmentation", None) + from_pretrained = getattr(loader, "from_pretrained", None) + except Exception as error: + return { + "loader": None, + "message": ( + f"Installed Transformers version {version} is not supported for BRIA RMBG 2.0: {error}. " + "Update Transformers manually if you want to use this model." + ), + } + + if not callable(loader) or not callable(from_pretrained): + return { + "loader": None, + "message": ( + f"Installed Transformers version {version} is not supported for BRIA RMBG 2.0. " + "It does not provide the required AutoModelForImageSegmentation API. " + "Update Transformers manually if you want to use this model." + ), + } + + return {"loader": loader, "message": None, "version": version} + + def _get_rmbg_model_loader(): """Return the Transformers loader used by BRIA RMBG 2.0, when installed.""" - try: - from transformers import AutoModelForImageSegmentation + status = _get_rmbg_transformers_status() + if status["loader"] is None: + log.debug(status["message"]) + return status["loader"] - return AutoModelForImageSegmentation - except Exception as error: - log.debug(f"BRIA RMBG 2.0 Transformers loader is unavailable: {error}") - return None + +def _get_rmbg_model_status_message(): + """Return an actionable message for a missing or incompatible Transformers install.""" + return _get_rmbg_transformers_status()["message"] def _is_rmbg_model_directory(path): @@ -142,23 +199,29 @@ class RMBG2Model: @classmethod def load(cls, model_directory): - loader = _get_rmbg_model_loader() + status = _get_rmbg_transformers_status() + loader = status["loader"] if loader is None: - raise RuntimeError( - "BRIA RMBG 2.0 requires the 'transformers' package. " - "Install the LayerForge requirements in the ComfyUI environment." - ) + raise RuntimeError(status["message"]) cache_key = os.path.normcase(os.path.normpath(os.path.abspath(model_directory))) with cls._model_cache_lock: if cache_key not in cls._model_cache: device = _get_rmbg_device() log.info(f"Loading BRIA RMBG 2.0 from {model_directory} on {device}") - model = loader.from_pretrained( - model_directory, - trust_remote_code=True, - local_files_only=True, - ) + try: + model = loader.from_pretrained( + model_directory, + trust_remote_code=True, + local_files_only=True, + ) + except Exception as error: + version = status.get("version", "unknown") + raise RuntimeError( + f"BRIA RMBG 2.0 could not be loaded with installed Transformers version {version}. " + "This version may not support the model's remote code; update Transformers manually " + f"if needed. Original error: {error}" + ) from error model.eval().to(device) cls._model_cache[cache_key] = cls(model, device) else: @@ -230,6 +293,8 @@ __all__ = [ "_find_existing_rmbg_model", "_find_local_rmbg_model", "_get_rmbg_model_loader", + "_get_rmbg_model_status_message", + "_get_rmbg_transformers_status", "_get_rmbg_remote_model", "_is_rmbg_model_directory", ] diff --git a/python/matting/service.py b/python/matting/service.py index c9cf049..b345440 100644 --- a/python/matting/service.py +++ b/python/matting/service.py @@ -16,6 +16,7 @@ from .backends.rmbg import ( _ensure_rmbg_model, _find_local_rmbg_model, _get_rmbg_model_loader, + _get_rmbg_model_status_message, _get_rmbg_remote_model, ) @@ -39,10 +40,7 @@ class BiRefNetMatting: if rmbg_model or local_rmbg_path: if _get_rmbg_model_loader() is None: - raise RuntimeError( - "BRIA RMBG 2.0 requires the 'transformers' package. " - "Install the LayerForge requirements in the active ComfyUI environment." - ) + raise RuntimeError(_get_rmbg_model_status_message()) model_directory = _ensure_rmbg_model(model_path) self.model = RMBG2Model.load(model_directory) self.model_path = model_directory diff --git a/requirements.txt b/requirements.txt index 494123d..c96d8ea 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,5 +5,5 @@ aiohttp numpy tqdm Pillow -transformers>=4.50.3 +transformers kornia diff --git a/src/CanvasView.ts b/src/CanvasView.ts index 3db8395..58f2a35 100644 --- a/src/CanvasView.ts +++ b/src/CanvasView.ts @@ -826,6 +826,10 @@ async function createCanvasWidget(node: ComfyNode, widget: any, app: ComfyApp): case 'unsupported_comfyui': showErrorNotification(modelStatus.message, 8000); return; + + case 'unsupported_rmbg': + showErrorNotification(modelStatus.message, 8000); + return; case 'not_downloaded': showWarningNotification(modelStatus.message || "The selected background-removal model will be downloaded automatically.", 7000); diff --git a/tests/test_entrypoint_integration.py b/tests/test_entrypoint_integration.py index 39fed8a..c6e671f 100644 --- a/tests/test_entrypoint_integration.py +++ b/tests/test_entrypoint_integration.py @@ -745,6 +745,21 @@ def test_rmbg_model_status_does_not_require_native_birefnet_loader(layerforge_ru assert status.payload["model_path"] == "remote:rmbg_2_0" +def test_rmbg_reports_an_unsupported_transformers_api(layerforge_runtime, monkeypatch): + rmbg_module = layerforge_runtime.matting_rmbg + transformers = ModuleType("transformers") + transformers.__version__ = "0.0-test" + monkeypatch.setitem(sys.modules, "transformers", transformers) + + status = rmbg_module._get_rmbg_transformers_status() + + assert status["loader"] is None + assert "0.0-test" in status["message"] + assert "not supported" in status["message"] + assert "AutoModelForImageSegmentation" in status["message"] + assert rmbg_module._get_rmbg_model_loader() is None + + def test_legacy_automatic_checkpoint_is_not_migrated(layerforge_runtime, monkeypatch, tmp_path): birefnet_module = layerforge_runtime.matting_birefnet model_dir = Path(tmp_path) / "models" / "background_removal"