commit ed0cbc7ef08ea1332194bfe2efa15d35c56e35de Author: d3cker Date: Sat Jan 3 11:21:56 2026 +0000 initial commit diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..d4a2c44 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,21 @@ +# http://editorconfig.org + +root = true + +[*] +indent_style = space +indent_size = 4 +trim_trailing_whitespace = true +insert_final_newline = true +charset = utf-8 +end_of_line = lf + +[*.bat] +indent_style = tab +end_of_line = crlf + +[LICENSE] +insert_final_newline = false + +[Makefile] +indent_style = tab diff --git a/.github/ISSUE_TEMPLATE.md b/.github/ISSUE_TEMPLATE.md new file mode 100644 index 0000000..255a874 --- /dev/null +++ b/.github/ISSUE_TEMPLATE.md @@ -0,0 +1,17 @@ +* comfyui_remove_background version: +* ComfyUI version: +* Python version: +* Operating System: +* Device: + +### Description + +Describe what you were trying to get done. +Tell us what happened, what went wrong, and what you expected to happen. + +### What I Did + +``` +Paste the command(s) you ran and the output. +If there was a crash, please include the traceback here. +``` diff --git a/.github/workflows/build-pipeline.yml b/.github/workflows/build-pipeline.yml new file mode 100644 index 0000000..43a5a19 --- /dev/null +++ b/.github/workflows/build-pipeline.yml @@ -0,0 +1,33 @@ +name: CI build + +on: + pull_request: + branches: + - master + - main +jobs: + build: + runs-on: ${{ matrix.os }} + env: + PYTHONIOENCODING: "utf8" + strategy: + matrix: + os: [ubuntu-latest] + python-version: ["3.12"] + + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install .[dev] + - name: Run Linting + run: | + ruff check . + - name: Run Tests + run: | + pytest tests/ diff --git a/.github/workflows/publish_node.yml b/.github/workflows/publish_node.yml new file mode 100644 index 0000000..0118542 --- /dev/null +++ b/.github/workflows/publish_node.yml @@ -0,0 +1,21 @@ +name: 📦 Publish to Comfy registry +on: + workflow_dispatch: + push: + tags: + - '*' + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: ♻️ Check out code + uses: actions/checkout@v4 + - name: 📦 Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/.github/workflows/validate.yml b/.github/workflows/validate.yml new file mode 100644 index 0000000..8d3b5d4 --- /dev/null +++ b/.github/workflows/validate.yml @@ -0,0 +1,13 @@ +name: Validate backwards compatibility + +on: + pull_request: + branches: + - master + - main + +jobs: + validate: + runs-on: ubuntu-latest + steps: + - uses: comfy-org/node-diff@main diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3141aae --- /dev/null +++ b/.gitignore @@ -0,0 +1,116 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# OSX useful to ignore +*.DS_Store +.AppleDouble +.LSOverride + +# Thumbnails +._* + +# Files that might appear in the root of a volume +.DocumentRevisions-V100 +.fseventsd +.Spotlight-V100 +.TemporaryItems +.Trashes +.VolumeIcon.icns +.com.apple.timemachine.donotpresent + +# Directories potentially created on remote AFP share +.AppleDB +.AppleDesktop +Network Trash Folder +Temporary Items +.apdisk + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +venv/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log + +# Sphinx documentation +docs/_build/ + +# IntelliJ Idea +.idea +*.iml +*.ipr +*.iws + +# PyBuilder +target/ + +# Cookiecutter +output/ +python_boilerplate/ +cookiecutter-pypackage-env/ + +# vscode settings +.history/ +*.code-workspace + +# Frontend extension +node_modules/ +.env +.env.local +.env.development.local +.env.test.local +.env.production.local +npm-debug.log* +yarn-debug.log* +yarn-error.log* +node.zip +.vscode/ +.claude/ + +.ruff_cache diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..2539924 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,10 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.4.9 + hooks: + # Run the linter. + - id: ruff + args: [ --fix ] + # Run the formatter. + - id: ruff-format diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..0cd0487 --- /dev/null +++ b/LICENSE @@ -0,0 +1,32 @@ + + +BSD License + +Copyright (c) 2026, ComfyUI Background Removal +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, this + list of conditions and the following disclaimer in the documentation and/or + other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from this + software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. +IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, +INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE +OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED +OF THE POSSIBILITY OF SUCH DAMAGE. + diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..6514960 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,9 @@ +include LICENSE +include README.md + +recursive-exclude * __pycache__ +recursive-exclude * *.py[co] + +recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif + + diff --git a/README.md b/README.md new file mode 100644 index 0000000..536c77d --- /dev/null +++ b/README.md @@ -0,0 +1,67 @@ +# ComfyUI Background Removal + +Custom ComfyUI node for background removal in images. Built with ONNX models and Pillow (PIL). + +> [!NOTE] +> This projected was created with a [cookiecutter](https://github.com/Comfy-Org/cookiecutter-comfy-extension) template. It helps you start writing custom nodes without worrying about the Python setup. + +## Quickstart + +1. Install [ComfyUI](https://docs.comfy.org/get_started). +1. Install [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) +1. Look up this extension in ComfyUI-Manager. If you are installing manually, clone this repository under `ComfyUI/custom_nodes`. +1. Restart ComfyUI. + +# Features + +- A list of features + +## Develop + +To install the dev dependencies and pre-commit (will run the ruff hook), do: + +```bash +cd comfyui_remove_background +pip install -e .[dev] +pre-commit install +``` + +The `-e` flag above will result in a "live" install, in the sense that any changes you make to your node extension will automatically be picked up the next time you run ComfyUI. + +## Publish to Github + +Install Github Desktop or follow these [instructions](https://docs.github.com/en/authentication/connecting-to-github-with-ssh/generating-a-new-ssh-key-and-adding-it-to-the-ssh-agent) for ssh. + +1. Create a Github repository that matches the directory name. +2. Push the files to Git +``` +git add . +git commit -m "project scaffolding" +git push +``` + +## Writing custom nodes + +An example custom node is located in [node.py](src/comfyui_remove_background/nodes.py). To learn more, read the [docs](https://docs.comfy.org/essentials/custom_node_overview). + + +## Tests + +This repo contains unit tests written in Pytest in the `tests/` directory. It is recommended to unit test your custom node. + +- [build-pipeline.yml](.github/workflows/build-pipeline.yml) will run pytest and linter on any open PRs +- [validate.yml](.github/workflows/validate.yml) will run [node-diff](https://github.com/Comfy-Org/node-diff) to check for breaking changes + +## Publishing to Registry + +If you wish to share this custom node with others in the community, you can publish it to the registry. We've already auto-populated some fields in `pyproject.toml` under `tool.comfy`, but please double-check that they are correct. + +You need to make an account on https://registry.comfy.org and create an API key token. + +- [ ] Go to the [registry](https://registry.comfy.org). Login and create a publisher id (everything after the `@` sign on your registry profile). +- [ ] Add the publisher id into the pyproject.toml file. +- [ ] Create an api key on the Registry for publishing from Github. [Instructions](https://docs.comfy.org/registry/publishing#create-an-api-key-for-publishing). +- [ ] Add it to your Github Repository Secrets as `REGISTRY_ACCESS_TOKEN`. + +A Github action will run on every git push. You can also run the Github action manually. Full instructions [here](https://docs.comfy.org/registry/publishing). Join our [discord](https://discord.com/invite/comfyorg) if you have any questions! + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..43211fc --- /dev/null +++ b/__init__.py @@ -0,0 +1,16 @@ +"""Top-level package for comfyui_remove_background.""" + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + +] + +__author__ = """ComfyUI Background Removal""" +__email__ = "l33tdev.marcinkowski@gmail.com" +__version__ = "0.0.1" + +from .src.comfyui_remove_background.nodes import NODE_CLASS_MAPPINGS +from .src.comfyui_remove_background.nodes import NODE_DISPLAY_NAME_MAPPINGS + + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..6898584 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,93 @@ +[build-system] +requires = ["setuptools>=70.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "comfyui_remove_background" +version = "0.0.1" +description = "Custom ComfyUI node for background removal in images. Built with ONNX models and Pillow (PIL)." +authors = [ + {name = "ComfyUI Background Removal", email = "l33tdev.marcinkowski@gmail.com"} +] +readme = "README.md" +license = {text = "BSD license"} +requires-python = ">=3.10" +classifiers = [] +dependencies = [ + +] + + + +[project.optional-dependencies] +dev = [ + "bump-my-version", + "coverage", # testing + "mypy", # linting + "pre-commit", # runs linting on commit + "pytest", # testing + "ruff", # linting +] + +[project.urls] +Repository = "https://github.com/d3cker/comfyui_remove_background" +BugTracker = "https://github.com/d3cker/comfyui_remove_background/issues" +Documentation = "https://github.com/d3cker/comfyui_remove_background/wiki" + + +[tool.comfy] +PublisherId = "d3cker" +DisplayName = "ComfyUI Background Removal" +Icon = "" +Tags = [] +Repository = "https://github.com/d3cker/comfyui_remove_background" + +includes = [] + +[tool.setuptools] +package-dir = {"" = "src"} + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-data] +"*" = ["*.*"] + +[tool.pytest.ini_options] +minversion = "8.0" +testpaths = [ + "tests", +] + +[tool.mypy] +files = "." + +# Use strict defaults +strict = true +warn_unreachable = true +warn_no_return = true + +[[tool.mypy.overrides]] +# Don't require test functions to include types +module = "tests.*" +allow_untyped_defs = true +disable_error_code = "attr-defined" + +[tool.ruff] +# extend-exclude = ["static", "ci/templates"] +line-length = 140 +src = ["src", "tests"] +target-version = "py39" + +# Add rules to ban exec/eval +[tool.ruff.lint] +select = [ + "S102", # exec-builtin + "S307", # eval-used + "W293", + "F", # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names. + # See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f +] + +[tool.ruff.lint.flake8-quotes] +inline-quotes = "double" diff --git a/src/comfyui_remove_background/__init__.py b/src/comfyui_remove_background/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/comfyui_remove_background/nodes.py b/src/comfyui_remove_background/nodes.py new file mode 100644 index 0000000..21ae17e --- /dev/null +++ b/src/comfyui_remove_background/nodes.py @@ -0,0 +1,191 @@ +# ComfyUI custom node: Remove Background (rembg) -> outputs RGB IMAGE + alpha MASK +# deps: +# pip install rembg onnxruntime pillow +# (for GPU: onnxruntime-gpu, and for Windows preload_dlls may be helpful) + +import onnxruntime as ort +import numpy as np +import torch +from PIL import Image + +try: + from rembg import remove, new_session +except Exception as e: + remove = None + new_session = None + _rembg_import_error = e + + +def _device_choices(): + choices = ["auto", "cpu"] + try: + if torch.cuda.is_available(): + for i in range(torch.cuda.device_count()): + name = torch.cuda.get_device_name(i) + choices.append(f"cuda:{i} ({name})") + except Exception: + pass + return choices + + +def _tensor_to_pil(img: torch.Tensor) -> Image.Image: + img = img.detach().cpu().clamp(0, 1) + arr = (img.numpy() * 255.0).round().astype(np.uint8) + + if arr.ndim != 3 or arr.shape[2] not in (3, 4): + raise ValueError(f"Expected image tensor [H,W,3/4], got {arr.shape}") + + if arr.shape[2] == 3: + return Image.fromarray(arr, mode="RGB").convert("RGBA") + return Image.fromarray(arr, mode="RGBA") + + +def _pil_to_image_and_mask(pil_rgba: Image.Image): + pil_rgba = pil_rgba.convert("RGBA") + rgba = np.array(pil_rgba, dtype=np.uint8) + + rgb = rgba[:, :, :3].astype(np.float32) / 255.0 + alpha = rgba[:, :, 3].astype(np.float32) / 255.0 + + return torch.from_numpy(rgb), torch.from_numpy(alpha) + + +def _parse_cuda_device_id(device_str: str) -> int: + # "cuda:0 (RTX ...)" -> 0 + s = device_str.strip() + s = s.split()[0] # "cuda:0" + _, idx = s.split(":") # "0" + return int(idx) + + +class RemoveBackground: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "model": ( + [ + "u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg", "silueta", "isnet-general-use", "isnet-anime", "sam", + "birefnet-general", "birefnet-general-lite", "birefnet-portrait", "birefnet-dis", "birefnet-hrsod", "birefnet-cod", "birefnet-massive" + ], + {"default": "birefnet-portrait"}, + ), + }, + "optional": { + # NOWE: + "device": (_device_choices(), {"default": "auto"}), + "allow_cpu_fallback": ("BOOLEAN", {"default": False}), + "preload_dlls": ("BOOLEAN", {"default": False}), # may help with Windows + + "alpha_matting": ("BOOLEAN", {"default": False}), + "alpha_matting_foreground_threshold": ("INT", {"default": 240, "min": 0, "max": 255}), + "alpha_matting_background_threshold": ("INT", {"default": 10, "min": 0, "max": 255}), + "alpha_matting_erode_size": ("INT", {"default": 10, "min": 0, "max": 64}), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "alpha") + FUNCTION = "run" + CATEGORY = "d3cker/Image" + + def __init__(self): + self._session = None + self._session_key = None # (model, device, allow_cpu_fallback) + + def _get_session(self, model: str, device: str, allow_cpu_fallback: bool, preload_dlls: bool): + if remove is None or new_session is None: + raise RuntimeError( + f"rembg import failed: {_rembg_import_error}\n" + "Install dependencies in ComfyUI environment:\n" + " pip install rembg onnxruntime pillow" + ) + + key = (model, device, bool(allow_cpu_fallback)) + if self._session is not None and self._session_key == key: + return self._session + + # preload DLLs (onnxruntime-gpu >= 1.21) + if preload_dlls and hasattr(ort, "preload_dlls"): + try: + ort.preload_dlls() + ort.print_debug_info() + except Exception as e: + print(f"[rembg] preload_dlls failed: {e}") + + avail = ort.get_available_providers() + + if device == "cpu": + providers = ["CPUExecutionProvider"] + elif device.startswith("cuda:"): + if "CUDAExecutionProvider" not in avail: + raise RuntimeError(f"CUDAExecutionProvider not available. available={avail}") + device_id = _parse_cuda_device_id(device) + providers = [("CUDAExecutionProvider", {"device_id": device_id})] + if allow_cpu_fallback: + providers.append("CPUExecutionProvider") + else: + # auto: prefer CUDA if available, then DML, then CPU + providers = [] + if "CUDAExecutionProvider" in avail: + providers.append(("CUDAExecutionProvider", {"device_id": 0})) + elif "DmlExecutionProvider" in avail: + providers.append("DmlExecutionProvider") + providers.append("CPUExecutionProvider") + + print(f"[rembg] available providers: {avail}") + print(f"[rembg] selected device: {device} | allow_cpu_fallback={allow_cpu_fallback}") + print(f"[rembg] providers arg: {providers}") + + self._session = new_session(model, providers=providers) + self._session_key = key + + # rembg BaseSession ma inner_session = ort.InferenceSession(...) + inner = getattr(self._session, "inner_session", None) + if inner is not None and hasattr(inner, "get_providers"): + print(f"[rembg] ONNX session providers in use: {inner.get_providers()}") + + return self._session + + def run( + self, + image, + model="birefnet-portrait", + device="auto", + allow_cpu_fallback=True, + preload_dlls=False, + alpha_matting=False, + alpha_matting_foreground_threshold=240, + alpha_matting_background_threshold=10, + alpha_matting_erode_size=10, + ): + session = self._get_session(model, device, allow_cpu_fallback, preload_dlls) + + if image.dim() != 4: + raise ValueError(f"Expected batched IMAGE tensor [B,H,W,C], got {tuple(image.shape)}") + + out_images = [] + out_masks = [] + + for i in range(image.shape[0]): + pil_in = _tensor_to_pil(image[i]) + + pil_out = remove( + pil_in, + session=session, + alpha_matting=bool(alpha_matting), + alpha_matting_foreground_threshold=int(alpha_matting_foreground_threshold), + alpha_matting_background_threshold=int(alpha_matting_background_threshold), + alpha_matting_erode_size=int(alpha_matting_erode_size), + ) + + img_t, mask_t = _pil_to_image_and_mask(pil_out) + out_images.append(img_t) + out_masks.append(mask_t) + + return (torch.stack(out_images, dim=0), torch.stack(out_masks, dim=0)) + + +NODE_CLASS_MAPPINGS = {"RemoveBackground": RemoveBackground} +NODE_DISPLAY_NAME_MAPPINGS = {"RemoveBackground": "Remove Background (Image + Alpha)"} diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..abd422b --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Unit test package for comfyui_remove_background.""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..2baa0ae --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,6 @@ +import os +import sys + +# Add the src directory to Python path +# This allows the tests to import the project +project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 0000000..e08e215 --- /dev/null +++ b/tests/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +testpaths = . +python_files = test_*.py +norecursedirs = .. diff --git a/tests/test_comfyui_remove_background.py b/tests/test_comfyui_remove_background.py new file mode 100644 index 0000000..e69de29