From 8cc87fb8238fe94cbddf3929084570a5a654cf8b Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 13 Aug 2026 16:58:09 +0200 Subject: [PATCH] refactor: store BiRefNet checkpoints directly under background_removal --- README.md | 2 +- python/matting.py | 55 +++++++--------------- tests/test_entrypoint_integration.py | 68 +++++++++++++++------------- 3 files changed, 53 insertions(+), 72 deletions(-) diff --git a/README.md b/README.md index 3051d40..3dd005a 100644 --- a/README.md +++ b/README.md @@ -241,7 +241,7 @@ Click the gear button next to **Matting** to configure its behavior. Settings ar - **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. -The model selector separates checkpoints already installed locally from official full-size BiRefNet variants available for download. Downloaded variants are stored in separate subdirectories under `ComfyUI/models/background_removal/` with descriptive filenames such as `BiRefNet-general.safetensors` and `BiRefNet-portrait.safetensors`, then validated against ComfyUI's native loader before use. Lite variants are intentionally excluded because they use a different architecture. Hugging Face cache folders are ignored in the selector and are not deleted. +The model selector separates checkpoints already installed locally from official full-size BiRefNet variants available for download. Downloaded variants are stored directly in `ComfyUI/models/background_removal/` with descriptive filenames such as `BiRefNet-general.safetensors` and `BiRefNet-portrait.safetensors`, then validated against ComfyUI's native loader before use. Lite variants are intentionally excluded because they use a different architecture. Hugging Face cache folders are ignored in the selector and are not deleted. ### BiRefNet model guide diff --git a/python/matting.py b/python/matting.py index b0c8c85..09a7986 100644 --- a/python/matting.py +++ b/python/matting.py @@ -22,7 +22,6 @@ _BIREFNET_FILENAME = "model.safetensors" _BIREFNET_DEFAULT_LOCAL_FILENAME = "BiRefNet-general.safetensors" _BIREFNET_PROJECT_URL = "https://github.com/ZhengPeng7/BiRefNet" _BIREFNET_REMOTE_PREFIX = "remote:" -_BIREFNET_REMOTE_DIRECTORY = "layerforge_birefnet" _BIREFNET_REQUIRED_KEYS = { "bb.layers.1.blocks.0.attn.relative_position_index", "bb.layers.2.blocks.17.attn.qkv.weight", @@ -241,25 +240,7 @@ def _get_birefnet_remote_checkpoint_path(model): base_paths = _get_birefnet_base_paths() if not base_paths: return None - return os.path.join( - base_paths[0], - _BIREFNET_REMOTE_DIRECTORY, - model["id"], - model["local_filename"], - ) - - -def _get_birefnet_remote_legacy_checkpoint_path(model): - """Return the pre-friendly-name path used by earlier LayerForge builds.""" - base_paths = _get_birefnet_base_paths() - if not base_paths: - return None - return os.path.join( - base_paths[0], - _BIREFNET_REMOTE_DIRECTORY, - model["id"], - model["filename"], - ) + return os.path.join(base_paths[0], model["local_filename"]) def _migrate_birefnet_checkpoint( @@ -292,13 +273,11 @@ def _migrate_birefnet_checkpoint( def _find_existing_birefnet_remote_checkpoint(model): - """Find a managed remote checkpoint and migrate its old generic filename.""" - return _migrate_birefnet_checkpoint( - _get_birefnet_remote_checkpoint_path(model), - _get_birefnet_remote_legacy_checkpoint_path(model), - log_description="BiRefNet checkpoint", - success_suffix="to the friendly filename", - ) + """Return an installed managed remote checkpoint, when available.""" + checkpoint_path = _get_birefnet_remote_checkpoint_path(model) + if checkpoint_path and _is_native_birefnet_checkpoint(checkpoint_path): + return checkpoint_path + return None def _get_birefnet_model_options(): @@ -306,21 +285,24 @@ def _get_birefnet_model_options(): _find_existing_birefnet_default_checkpoint() local_options = [] seen = set() - managed_roots = [ - os.path.normcase(os.path.normpath(os.path.join(base_path, _BIREFNET_REMOTE_DIRECTORY))) - for base_path in _get_birefnet_base_paths() - ] + base_paths = _get_birefnet_base_paths() + managed_paths = { + os.path.normcase(os.path.normpath(os.path.abspath(path))) + for base_path in base_paths + for model in _BIREFNET_MODEL_CATALOG + for path in [os.path.join(base_path, model["local_filename"])] + } for path in _iter_birefnet_checkpoint_paths(): normalized = os.path.normcase(os.path.normpath(os.path.abspath(path))) if normalized in seen or not _is_native_birefnet_checkpoint(path): continue - if any(normalized == root or normalized.startswith(root + os.sep) for root in managed_roots): + if normalized in managed_paths: continue seen.add(normalized) label = os.path.basename(path) - for base_path in _get_birefnet_base_paths(): + for base_path in base_paths: try: relative_path = os.path.relpath(path, base_path) except ValueError: @@ -407,12 +389,7 @@ def _download_birefnet_checkpoint(model=None): else: repository = model["repo_id"] filename = model["filename"] - download_dir = os.path.join( - _get_birefnet_download_dir(), - _BIREFNET_REMOTE_DIRECTORY, - model["id"], - ) - os.makedirs(download_dir, exist_ok=True) + download_dir = _get_birefnet_download_dir() model_label = model["label"] target_path = _get_birefnet_remote_checkpoint_path(model) diff --git a/tests/test_entrypoint_integration.py b/tests/test_entrypoint_integration.py index b3f2816..ffd51ce 100644 --- a/tests/test_entrypoint_integration.py +++ b/tests/test_entrypoint_integration.py @@ -598,6 +598,42 @@ def test_selected_remote_matting_model_is_sent_to_downloader(layerforge_runtime, assert downloaded["model"] is selected_model +def test_remote_matting_checkpoint_path_uses_background_removal_root(layerforge_runtime): + node_module = layerforge_runtime.matting + model = next(model for model in node_module._BIREFNET_MODEL_CATALOG if model["id"] == "portrait") + + expected = Path(layerforge_runtime.matting.folder_paths.models_dir) / "background_removal" / model[ + "local_filename" + ] + + assert Path(node_module._get_birefnet_remote_checkpoint_path(model)) == expected + + +def test_remote_matting_download_uses_background_removal_root(layerforge_runtime, monkeypatch, tmp_path): + node_module = layerforge_runtime.matting + model = next(model for model in node_module._BIREFNET_MODEL_CATALOG if model["id"] == "portrait") + download = {} + + def fake_hf_hub_download(**kwargs): + download.update(kwargs) + downloaded_path = Path(kwargs["local_dir"]) / kwargs["filename"] + downloaded_path.write_bytes(b"checkpoint") + return str(downloaded_path) + + huggingface_hub = ModuleType("huggingface_hub") + huggingface_hub.hf_hub_download = fake_hf_hub_download + monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub) + monkeypatch.setattr(node_module, "_is_native_birefnet_checkpoint", lambda path: True) + + result = node_module._download_birefnet_checkpoint(model) + expected_dir = Path(tmp_path) / "models" / "background_removal" + + assert Path(download["local_dir"]) == expected_dir + assert "layerforge_birefnet" not in Path(download["local_dir"]).parts + assert Path(result) == expected_dir / model["local_filename"] + assert Path(result).exists() + + def test_legacy_automatic_checkpoint_gets_friendly_filename(layerforge_runtime, monkeypatch, tmp_path): node_module = layerforge_runtime.matting model_dir = Path(tmp_path) / "models" / "background_removal" @@ -619,38 +655,6 @@ def test_legacy_automatic_checkpoint_gets_friendly_filename(layerforge_runtime, assert not legacy_path.exists() -def test_legacy_remote_checkpoint_gets_friendly_filename(layerforge_runtime, monkeypatch, tmp_path): - node_module = layerforge_runtime.matting - model = next(model for model in node_module._BIREFNET_MODEL_CATALOG if model["id"] == "portrait") - managed_dir = Path(tmp_path) / "managed" / model["id"] - managed_dir.mkdir(parents=True) - legacy_path = managed_dir / model["filename"] - friendly_path = managed_dir / model["local_filename"] - legacy_path.write_bytes(b"legacy checkpoint") - - monkeypatch.setattr( - node_module, - "_get_birefnet_remote_checkpoint_path", - lambda selected_model: str(friendly_path) if selected_model is model else None, - ) - monkeypatch.setattr( - node_module, - "_get_birefnet_remote_legacy_checkpoint_path", - lambda selected_model: str(legacy_path) if selected_model is model else None, - ) - monkeypatch.setattr( - node_module, - "_is_native_birefnet_checkpoint", - lambda path: Path(path).resolve() == legacy_path.resolve(), - ) - - result = node_module._find_existing_birefnet_remote_checkpoint(model) - - assert Path(result) == friendly_path - assert friendly_path.exists() - assert not legacy_path.exists() - - def test_empty_execution_returns_comfyui_compatible_fallback_tensors(layerforge_runtime): node_class = layerforge_runtime.node.LayerForgeNode node_class._canvas_data_storage.clear()