diff --git a/README.md b/README.md index acb920f..cad4615 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -

ComfyUI-BiRefNet-lite

+

ComfyUI-BiRefNet-Super


English | 中文 @@ -18,14 +18,17 @@ Supports putalpha, naive, and alpha_matting cropping methods. ![slot](./assets/feature2.png)
## News -- Nov 19, 2023: Add `mask_precision_threshold` parameter to control the accuracy threshold of the mask, default to 0.1
+- Nov 19, 2024: Add `mask_precision_threshold` parameter to control the accuracy threshold of the mask, default to 0.1
+- Dec 10, 2024: + - Add `BiRefNet_onnx` node to Support `onnx` model
+ - Modify the name of repository from `ComfyUI-BiRefNet-lite` to `ComfyUI-BiRefNet-Super`
## Installation #### Method 1: 1. Navigate to the node directory, `ComfyUI/custom_nodes/` -2. `git clone https://github.com/rubi-du/ComfyUI-BiRefNet-lite.git` -3. `cd ComfyUI-BiRefNet-lite` +2. `git clone https://github.com/rubi-du/ComfyUI-BiRefNet-Super.git` +3. `cd ComfyUI-BiRefNet-Super` 4. `pip install -r requirements.txt` 5. Restart ComfyUI @@ -33,11 +36,13 @@ Supports putalpha, naive, and alpha_matting cropping methods. Directly download the node source code package, unzip it into the `custom_nodes` directory, and then restart ComfyUI. #### Method 3: -Install via ComfyUI-Manager by searching for "ComfyUI-BiRefNet-lite". +Install via ComfyUI-Manager by searching for "ComfyUI-BiRefNet-Super". ## Usage + +### BiRefNet_Super | BiRefNet_Lite -Example workflows are placed in `ComfyUI-BiRefNet-lite/workflow`.
+Example workflows are placed in `ComfyUI-BiRefNet-Super/workflow`.
There are two options for loading models: one is to automatically download and load a remote model, and the other is to load a local model (in which case you need to set `load_local_model` to true and set `local_model_path` to the local model path under the models/birefnet directory, for example, the BiRefNet folder).
@@ -46,7 +51,7 @@ ___ ![](./assets/model_path.png) -Model download links:
+#### Model download links:
BiRefNet: https://huggingface.co/ZhengPeng7/BiRefNet/tree/main
BiRefNet_lite-2K: https://huggingface.co/ZhengPeng7/BiRefNet_lite-2K/tree/main
BiRefNet-portrait: https://huggingface.co/ZhengPeng7/BiRefNet-portrait/tree/main
@@ -60,10 +65,39 @@ Usage of workflow.json
___ Usage of video_workflow.json
-[Workflow Address](./workflow/video_workflow.json) +[Workflow Address](./workflow/video_workflow.json) + +### BiRefNet_onnx +Example workflows are placed in [workflow](./workflow/workflow-onnx.json).
+ +Place the model file in `ComfyUI/models/birefnet`.
+ +Model file name should be one of the following: +- birefnet-general.onnx +- birefnet-general-lite.onnx +- birefnet-portrait.onnx +- birefnet_massive.onnx +- birefnet-hrsod.onnx +- birefnet-dis.onnx +- birefnet-cod.onnx + + +![plot](./assets/demo2.png) + +#### Model download links:
+birefnet-genernal: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-epoch_244.onnx
+birefnet-genernal-lite: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx
+birefnet-portrait: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-portrait-epoch_150.onnx
+birefnet_massive: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-massive-TR_DIS5K_TR_TEs-epoch_420.onnx
+birefnet-hrsod: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-HRSOD_DHU-epoch_115.onnx
+birefnet-dis: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-DIS-epoch_590.onnx
+birefnet-cod: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-COD-epoch_125.onnx
+ ## Acknowledgments Thanks to all the authors of the BiRefNet repository [ZhengPeng7/BiRefNet](https://github.com/zhengpeng7/birefnet). -Some code was referenced from [MoonHugo/ComfyUI-BiRefNet-Hugo](https://github.com/MoonHugo/ComfyUI-BiRefNet-Hugo). Thank you! \ No newline at end of file +Some code was referenced from [MoonHugo/ComfyUI-BiRefNet-Hugo](https://github.com/MoonHugo/ComfyUI-BiRefNet-Hugo). Thank you! + +Some code was referenced from [danielgatis/rembg](https://github.com/danielgatis/rembg). Thank you! \ No newline at end of file diff --git a/README_CN.md b/README_CN.md index 2317fca..41b9a9d 100644 --- a/README_CN.md +++ b/README_CN.md @@ -1,4 +1,4 @@ -

ComfyUI-BiRefNet-lite

+

ComfyUI-BiRefNet-Super


English | 中文 @@ -26,8 +26,8 @@ #### 方法1: 1. 进入节点目录, `ComfyUI/custom_nodes/` -2. `git clone https://github.com/rubi-du/ComfyUI-BiRefNet-lite.git` -3. `cd ComfyUI-BiRefNet-lite` +2. `git clone https://github.com/rubi-du/ComfyUI-BiRefNet-Super.git` +3. `cd ComfyUI-BiRefNet-Super` 4. `pip install -r requirements.txt` 5. 重启ComfyUI @@ -35,11 +35,13 @@ 直接下载节点源码包，然后解压到custom_nodes目录下，最后重启ComfyUI #### 方法3： -通过ComfyUI-Manager安装，搜索“ComfyUI-BiRefNet-lite”进行安装 +通过ComfyUI-Manager安装，搜索“ComfyUI-BiRefNet-Super”进行安装 ## 使用 -示例工作流放置在`ComfyUI-BiRefNet-lite/workflow`中
+### BiRefNet_Super | BiRefNet_Lite + +示例工作流放置在`ComfyUI-BiRefNet-Super/workflow`中
加载模型有两种方式可选：一是自动下载并加载远程模型，二是加载本地模型（此时需将`load_local_model`设为true，并将`local_model_path`设置为指向models/birefnet目录下的本地模型路径，例如BiRefNet文件夹）
@@ -49,7 +51,7 @@ ___ ![](./assets/model_path.png) -模型下载地址：
+#### 模型下载地址：
BiRefNet: https://huggingface.co/ZhengPeng7/BiRefNet/tree/main
BiRefNet_lite-2K: https://huggingface.co/ZhengPeng7/BiRefNet_lite-2K/tree/main
BiRefNet-portrait: https://huggingface.co/ZhengPeng7/BiRefNet-portrait/tree/main
@@ -65,8 +67,38 @@ ___ 工作流video_workflow.json的使用
[workflow地址](./workflow/video_workflow.json) + + +### BiRefNet_onnx +示例工作流放置在 [workflow](./workflow/workflow-onnx.json)中
+ +模型下载完放置在 `ComfyUI/models/birefnet` 目录下.
+ +模型名称如下: +- birefnet-general.onnx +- birefnet-general-lite.onnx +- birefnet-portrait.onnx +- birefnet_massive.onnx +- birefnet-hrsod.onnx +- birefnet-dis.onnx +- birefnet-cod.onnx + + +![plot](./assets/demo2.png) + +#### 模型下载地址：
+birefnet-genernal: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-epoch_244.onnx
+birefnet-genernal-lite: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx
+birefnet-portrait: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-portrait-epoch_150.onnx
+birefnet_massive: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-massive-TR_DIS5K_TR_TEs-epoch_420.onnx
+birefnet-hrsod: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-HRSOD_DHU-epoch_115.onnx
+birefnet-dis: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-DIS-epoch_590.onnx
+birefnet-cod: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-COD-epoch_125.onnx
+ ## 感谢 感谢BiRefNet仓库的所有作者 [ZhengPeng7/BiRefNet](https://github.com/zhengpeng7/birefnet) -部分代码参考了 [MoonHugo/ComfyUI-BiRefNet-Hugo](https://github.com/MoonHugo/ComfyUI-BiRefNet-Hugo) 感谢！ \ No newline at end of file +部分代码参考了 [MoonHugo/ComfyUI-BiRefNet-Hugo](https://github.com/MoonHugo/ComfyUI-BiRefNet-Hugo) 感谢！ + +部分代码参考了 [danielgatis/rembg](https://github.com/danielgatis/rembg). 感谢！ \ No newline at end of file diff --git a/assets/demo2.png b/assets/demo2.png new file mode 100644 index 0000000..06fb36a Binary files /dev/null and b/assets/demo2.png differ diff --git a/assets/onnx.png b/assets/onnx.png new file mode 100644 index 0000000..f5906b5 Binary files /dev/null and b/assets/onnx.png differ diff --git a/nodes.py b/nodes.py index 2d732c9..6cc7c79 100644 --- a/nodes.py +++ b/nodes.py @@ -10,6 +10,7 @@ from scipy.ndimage import binary_erosion from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml from pymatting.util.util import stack_images +from .rmbg.session_factory import new_session from folder_paths import models_dir, add_model_folder_path, get_folder_paths torch.set_float32_matmul_precision(["high", "highest"][0]) @@ -154,7 +155,112 @@ def get_device_by_name(device): print("\033[93mUse Device(使用设备):", device, "\033[0m") return device +def get_device_provider(device): + providers = [] + cuda_provider_options = { + "gpu_mem_limit": 9 * 1024 * 1024 * 1024, # 9GB. + "arena_extend_strategy": "kSameAsRequested", + "cudnn_conv_use_max_workspace": "0" + } + if device == 'cuda' and torch.cuda.is_available(): + providers.append(("CUDAExecutionProvider", cuda_provider_options)) + + if len(providers) == 0: + providers.append(("CPUExecutionProvider", {})) + + return providers + + _birefnet_model = None +_birefnet_model_name = "" +_birefnet_onnx_model = None +_birefnet_model_name_onnx = "" + +class BiRefNet_onnx: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "device": (["cpu", "cuda"],{"default": "cpu"}), + "model_name": ([ + "birefnet-general", + "birefnet-general-lite", + "birefnet-portrait", + "birefnet-cod", + "birefnet-dis", + "birefnet-hrsod", + "birefnet-massive" + ], {"default": "birefnet-general"}), + "cutout_func": (["putalpha", "naive", "alpha_matting"],{"default": "putalpha"}), + # "cached": ("BOOLEAN", {"default": True}), + "alpha_matting_foreground_threshold": ("INT", {"default": 240}), + "alpha_matting_background_threshold": ("INT", {"default": 10}), + "alpha_matting_erode_size": ("INT", {"default": 10}), + }, + } + RETURN_TYPES = ("IMAGE", "MASK",) + RETURN_NAMES = ("image", "mask",) + FUNCTION = "background_remove" + CATEGORY = "🔥BiRefNet" + + def background_remove(self, + image, + device, + model_name, + cutout_func, + # cached, + alpha_matting_foreground_threshold, + alpha_matting_background_threshold, + alpha_matting_erode_size, + ): + providers = get_device_provider(device) + processed_images = [] + processed_masks = [] + global _birefnet_onnx_model, _birefnet_model_name_onnx + + ## TODO: cache model + if _birefnet_onnx_model is not None and _birefnet_model_name_onnx == model_name: + session = _birefnet_onnx_model + print("Loading from cache...") + else: + session = new_session(model_name, providers=providers) + if _birefnet_onnx_model is not None: + _birefnet_onnx_model = None + _birefnet_model_name_onnx = '' + clear_memory() + # if cached: + # _birefnet_onnx_model = session + # _birefnet_model_name_onnx = model_name + + for image in image: + orig_image = tensor2pil(image) + masks = session.predict(orig_image) + pil_im = masks[0] + if cutout_func == 'putalpha': + new_im = putalpha_cutout(orig_image, pil_im) + elif cutout_func == 'naive': + new_im = naive_cutout(orig_image, pil_im) + elif cutout_func == 'alpha_matting': + new_im = alpha_matting_cutout( + orig_image, + pil_im, + foreground_threshold=alpha_matting_foreground_threshold, + background_threshold=alpha_matting_background_threshold, + erode_structure_size=alpha_matting_erode_size + ) + + new_im_tensor = pil2tensor(new_im) + pil_im_tensor = pil2tensor(pil_im) + processed_images.append(new_im_tensor) + processed_masks.append(pil_im_tensor) + + clear_memory() + + new_ims = torch.cat(processed_images, dim=0) + new_masks = torch.cat(processed_masks, dim=0) + + return new_ims, new_masks class BiRefNet_Lite: def __init__(self): @@ -205,15 +311,18 @@ class BiRefNet_Lite: device = get_device_by_name(device) birefnet = None global _birefnet_model - if cached and _birefnet_model is not None: + global _birefnet_model_name + + local_model_path = kwargs.get("local_model_path", model_path) + if cached and _birefnet_model is not None and _birefnet_model_name == local_model_path: birefnet = _birefnet_model else: if _birefnet_model is not None: _birefnet_model = None + _birefnet_model_name = None clear_memory() if load_local_model: - local_model_path = kwargs.get("local_model_path", model_path) local_model_path = os.path.join(get_folder_paths('birefnet')[-1], local_model_path) # 判断是否开启双卡支持 # 获取不同显卡的内存可用量 @@ -244,6 +353,9 @@ class BiRefNet_Lite: print('No need to delete:', e) birefnet = AutoModelForImageSegmentation.from_pretrained(local_model_path,trust_remote_code=True, **spare_params) + if cached: + _birefnet_model = birefnet + _birefnet_model_name = local_model_path else: birefnet = AutoModelForImageSegmentation.from_pretrained( "ZhengPeng7/BiRefNet", trust_remote_code=True @@ -297,10 +409,14 @@ class BiRefNet_Lite: NODE_CLASS_MAPPINGS = { - "BiRefNet_Lite": BiRefNet_Lite + "BiRefNet_Super": BiRefNet_Lite, + "BiRefNet_Lite": BiRefNet_Lite, + "BiRefNet_onnx": BiRefNet_onnx } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { - "BiRefNet_Lite": "🔥BiRefNet_Lite" + "BiRefNet_Lite": "🔥BiRefNet_Lite", + "BiRefNet_Super": "🔥BiRefNet_Super", + "BiRefNet_onnx": "🔥BiRefNet_onnx" } diff --git a/pyproject.toml b/pyproject.toml index 72dbc15..32c33f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,16 +1,16 @@ [project] -name = "comfyui-birefnet-lite" +name = "comfyui-BiRefNet-Super" description = "This repository packages the latest BiRefNet model as a ComfyUI node for use, supporting chunked loading on both CPU and GPU, as well as model caching features." -version = "1.0.5" +version = "1.0.6" license = {file = "LICENSE"} [project.urls] -Repository = "https://github.com/rubi-du/ComfyUI-BiRefNet-lite" +Repository = "https://github.com/rubi-du/ComfyUI-BiRefNet-Super" # Used by Comfy Registry https://comfyregistry.org dependencies = ["numpy", "timm", "pymatting", "scikit-image","scipy"] [tool.comfy] PublisherId = "duhaifeng" -DisplayName = "ComfyUI-BiRefNet-lite" +DisplayName = "ComfyUI-BiRefNet-Super" Icon = "" diff --git a/rmbg/birefnet_cod.py b/rmbg/birefnet_cod.py new file mode 100644 index 0000000..4fa13dc --- /dev/null +++ b/rmbg/birefnet_cod.py @@ -0,0 +1,52 @@ +import os + +import pooch + +from .birefnet_general import BiRefNetSessionGeneral + + +class BiRefNetSessionCOD(BiRefNetSessionGeneral): + """ + This class represents a BiRefNet-COD session, which is a subclass of BiRefNetSessionGeneral. + """ + + @classmethod + def download_models(cls, *args, **kwargs): + """ + Downloads the BiRefNet-COD model file from a specific URL and saves it. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The path to the downloaded model file. + """ + fname = f"{cls.name(*args, **kwargs)}.onnx" + pooch.retrieve( + "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-COD-epoch_125.onnx", + ( + None + if cls.checksum_disabled(*args, **kwargs) + else "md5:f6d0d21ca89d287f17e7afe9f5fd3b45" + ), + fname=fname, + path=cls.u2net_home(*args, **kwargs), + progressbar=True, + ) + + return os.path.join(cls.u2net_home(*args, **kwargs), fname) + + @classmethod + def name(cls, *args, **kwargs): + """ + Returns the name of the BiRefNet-COD session. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The name of the session. + """ + return "birefnet-cod" diff --git a/rmbg/birefnet_dis.py b/rmbg/birefnet_dis.py new file mode 100644 index 0000000..eecf147 --- /dev/null +++ b/rmbg/birefnet_dis.py @@ -0,0 +1,52 @@ +import os + +import pooch + +from .birefnet_general import BiRefNetSessionGeneral + + +class BiRefNetSessionDIS(BiRefNetSessionGeneral): + """ + This class represents a BiRefNet-DIS session, which is a subclass of BiRefNetSessionGeneral. + """ + + @classmethod + def download_models(cls, *args, **kwargs): + """ + Downloads the BiRefNet-DIS model file from a specific URL and saves it. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The path to the downloaded model file. + """ + fname = f"{cls.name(*args, **kwargs)}.onnx" + pooch.retrieve( + "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-DIS-epoch_590.onnx", + ( + None + if cls.checksum_disabled(*args, **kwargs) + else "md5:2d4d44102b446f33a4ebb2e56c051f2b" + ), + fname=fname, + path=cls.u2net_home(*args, **kwargs), + progressbar=True, + ) + + return os.path.join(cls.u2net_home(*args, **kwargs), fname) + + @classmethod + def name(cls, *args, **kwargs): + """ + Returns the name of the BiRefNet-DIS session. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The name of the session. + """ + return "birefnet-dis" diff --git a/rmbg/birefnet_general.py b/rmbg/birefnet_general.py new file mode 100644 index 0000000..1f5be70 --- /dev/null +++ b/rmbg/birefnet_general.py @@ -0,0 +1,93 @@ +""" BiRefNet General Session """ +import os +from typing import List + +import numpy as np +import pooch +from PIL import Image +from PIL.Image import Image as PILImage + +from .sessions.base import BaseSession + + +class BiRefNetSessionGeneral(BaseSession): + """ + This class represents a BiRefNet-General session, which is a subclass of BaseSession. + """ + + def sigmoid(self, mat): + """ Sigmoidal function applied on an array 'mat' """ + return 1 / (1 + np.exp(-mat)) + + def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]: + """ + Predicts the output masks for the input image using the inner session. + + Parameters: + img (PILImage): The input image. + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + List[PILImage]: The list of output masks. + """ + ort_outs = self.inner_session.run( + None, + self.normalize( + img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (1024, 1024) + ), + ) + + pred = self.sigmoid(ort_outs[0][:, 0, :, :]) + + ma = np.max(pred) + mi = np.min(pred) + + pred = (pred - mi) / (ma - mi) + pred = np.squeeze(pred) + + mask = Image.fromarray((pred * 255).astype("uint8"), mode="L") + mask = mask.resize(img.size, Image.Resampling.LANCZOS) + + return [mask] + + @classmethod + def download_models(cls, *args, **kwargs): + """ + Downloads the BiRefNet-General model file from a specific URL and saves it. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The path to the downloaded model file. + """ + fname = f"{cls.name(*args, **kwargs)}.onnx" + pooch.retrieve( + "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-epoch_244.onnx", + ( + None + if cls.checksum_disabled(*args, **kwargs) + else "md5:7a35a0141cbbc80de11d9c9a28f52697" + ), + fname=fname, + path=cls.u2net_home(*args, **kwargs), + progressbar=True, + ) + + return os.path.join(cls.u2net_home(*args, **kwargs), fname) + + @classmethod + def name(cls, *args, **kwargs): + """ + Returns the name of the BiRefNet-General session. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The name of the session. + """ + return "birefnet-general" diff --git a/rmbg/birefnet_general_lite.py b/rmbg/birefnet_general_lite.py new file mode 100644 index 0000000..416f821 --- /dev/null +++ b/rmbg/birefnet_general_lite.py @@ -0,0 +1,52 @@ +import os + +import pooch + +from .birefnet_general import BiRefNetSessionGeneral + + +class BiRefNetSessionGeneralLite(BiRefNetSessionGeneral): + """ + This class represents a BiRefNet-General-Lite session, which is a subclass of BiRefNetSessionGeneral. + """ + + @classmethod + def download_models(cls, *args, **kwargs): + """ + Downloads the BiRefNet-General-Lite model file from a specific URL and saves it. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The path to the downloaded model file. + """ + fname = f"{cls.name(*args, **kwargs)}.onnx" + pooch.retrieve( + "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx", + ( + None + if cls.checksum_disabled(*args, **kwargs) + else "md5:4fab47adc4ff364be1713e97b7e66334" + ), + fname=fname, + path=cls.u2net_home(*args, **kwargs), + progressbar=True, + ) + + return os.path.join(cls.u2net_home(*args, **kwargs), fname) + + @classmethod + def name(cls, *args, **kwargs): + """ + Returns the name of the BiRefNet-General-Lite session. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The name of the session. + """ + return "birefnet-general-lite" diff --git a/rmbg/birefnet_hrsod.py b/rmbg/birefnet_hrsod.py new file mode 100644 index 0000000..6ab72fc --- /dev/null +++ b/rmbg/birefnet_hrsod.py @@ -0,0 +1,52 @@ +import os + +import pooch + +from .birefnet_general import BiRefNetSessionGeneral + + +class BiRefNetSessionHRSOD(BiRefNetSessionGeneral): + """ + This class represents a BiRefNet-HRSOD session, which is a subclass of BiRefNetSessionGeneral. + """ + + @classmethod + def download_models(cls, *args, **kwargs): + """ + Downloads the BiRefNet-HRSOD model file from a specific URL and saves it. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The path to the downloaded model file. + """ + fname = f"{cls.name(*args, **kwargs)}.onnx" + pooch.retrieve( + "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-HRSOD_DHU-epoch_115.onnx", + ( + None + if cls.checksum_disabled(*args, **kwargs) + else "md5:c017ade5de8a50ff0fd74d790d268dda" + ), + fname=fname, + path=cls.u2net_home(*args, **kwargs), + progressbar=True, + ) + + return os.path.join(cls.u2net_home(*args, **kwargs), fname) + + @classmethod + def name(cls, *args, **kwargs): + """ + Returns the name of the BiRefNet-HRSOD session. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The name of the session. + """ + return "birefnet-hrsod" diff --git a/rmbg/birefnet_massive.py b/rmbg/birefnet_massive.py new file mode 100644 index 0000000..cdd3fa1 --- /dev/null +++ b/rmbg/birefnet_massive.py @@ -0,0 +1,52 @@ +import os + +import pooch + +from .birefnet_general import BiRefNetSessionGeneral + + +class BiRefNetSessionMassive(BiRefNetSessionGeneral): + """ + This class represents a BiRefNet-Massive session, which is a subclass of BiRefNetSessionGeneral. + """ + + @classmethod + def download_models(cls, *args, **kwargs): + """ + Downloads the BiRefNet-Massive model file from a specific URL and saves it. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The path to the downloaded model file. + """ + fname = f"{cls.name(*args, **kwargs)}.onnx" + pooch.retrieve( + "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-massive-TR_DIS5K_TR_TEs-epoch_420.onnx", + ( + None + if cls.checksum_disabled(*args, **kwargs) + else "md5:33e726a2136a3d59eb0fdf613e31e3e9" + ), + fname=fname, + path=cls.u2net_home(*args, **kwargs), + progressbar=True, + ) + + return os.path.join(cls.u2net_home(*args, **kwargs), fname) + + @classmethod + def name(cls, *args, **kwargs): + """ + Returns the name of the BiRefNet-Massive session. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The name of the session. + """ + return "birefnet-massive" diff --git a/rmbg/birefnet_portrait.py b/rmbg/birefnet_portrait.py new file mode 100644 index 0000000..12a5994 --- /dev/null +++ b/rmbg/birefnet_portrait.py @@ -0,0 +1,53 @@ +""" birefnet_portrait """ +import os + +import pooch + +from .birefnet_general import BiRefNetSessionGeneral + + +class BiRefNetSessionPortrait(BiRefNetSessionGeneral): + """ + This class represents a BiRefNet-Portrait session, which is a subclass of BiRefNetSessionGeneral. + """ + + @classmethod + def download_models(cls, *args, **kwargs): + """ + Downloads the BiRefNet-Portrait model file from a specific URL and saves it. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The path to the downloaded model file. + """ + fname = f"{cls.name(*args, **kwargs)}.onnx" + pooch.retrieve( + "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-portrait-epoch_150.onnx", + ( + None + if cls.checksum_disabled(*args, **kwargs) + else "md5:c3a64a6abf20250d090cd055f12a3b67" + ), + fname=fname, + path=cls.u2net_home(*args, **kwargs), + progressbar=True, + ) + + return os.path.join(cls.u2net_home(*args, **kwargs), fname) + + @classmethod + def name(cls, *args, **kwargs): + """ + Returns the name of the BiRefNet-Portrait session. + + Parameters: + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + str: The name of the session. + """ + return "birefnet-portrait" diff --git a/rmbg/session_factory.py b/rmbg/session_factory.py new file mode 100644 index 0000000..8700d16 --- /dev/null +++ b/rmbg/session_factory.py @@ -0,0 +1,46 @@ +import os +from typing import Type + +import onnxruntime as ort + +from .sessions import sessions_class +from .sessions.base import BaseSession + + +def new_session( + model_name:str, providers=None, *args, **kwargs +) -> BaseSession: + """ + Create a new session object based on the specified model name. + + This function searches for the session class based on the model name in the 'sessions_class' list. + It then creates an instance of the session class with the provided arguments. + The 'sess_opts' object is created using the 'ort.SessionOptions()' constructor. + If the 'OMP_NUM_THREADS' environment variable is set, the 'inter_op_num_threads' option of 'sess_opts' is set to its value. + + Parameters: + model_name (str): The name of the model. + providers: The providers for the session. + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + BaseSession: The created session object. + """ + session_class: Type[BaseSession] + + for sc in sessions_class: + if sc.name() == model_name: + session_class = sc + break + + if not session_class: + raise ValueError("Model not found") + + sess_opts = ort.SessionOptions() + + if "OMP_NUM_THREADS" in os.environ: + sess_opts.inter_op_num_threads = int(os.environ["OMP_NUM_THREADS"]) + sess_opts.intra_op_num_threads = int(os.environ["OMP_NUM_THREADS"]) + + return session_class(model_name, sess_opts, providers, *args, **kwargs) diff --git a/rmbg/sessions/__init__.py b/rmbg/sessions/__init__.py new file mode 100644 index 0000000..11f8846 --- /dev/null +++ b/rmbg/sessions/__init__.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from typing import List + +from .base import BaseSession + +sessions_class: List[type[BaseSession]] = [] +sessions_names: List[str] = [] + +from ..birefnet_general import BiRefNetSessionGeneral + +sessions_class.append(BiRefNetSessionGeneral) +sessions_names.append(BiRefNetSessionGeneral.name()) + +from ..birefnet_general_lite import BiRefNetSessionGeneralLite + +sessions_class.append(BiRefNetSessionGeneralLite) +sessions_names.append(BiRefNetSessionGeneralLite.name()) + +from ..birefnet_portrait import BiRefNetSessionPortrait + +sessions_class.append(BiRefNetSessionPortrait) +sessions_names.append(BiRefNetSessionPortrait.name()) + +from ..birefnet_dis import BiRefNetSessionDIS + +sessions_class.append(BiRefNetSessionDIS) +sessions_names.append(BiRefNetSessionDIS.name()) + +from ..birefnet_hrsod import BiRefNetSessionHRSOD + +sessions_class.append(BiRefNetSessionHRSOD) +sessions_names.append(BiRefNetSessionHRSOD.name()) + +from ..birefnet_cod import BiRefNetSessionCOD + +sessions_class.append(BiRefNetSessionCOD) +sessions_names.append(BiRefNetSessionCOD.name()) + +from ..birefnet_massive import BiRefNetSessionMassive + +sessions_class.append(BiRefNetSessionMassive) +sessions_names.append(BiRefNetSessionMassive.name()) diff --git a/rmbg/sessions/base.py b/rmbg/sessions/base.py new file mode 100644 index 0000000..3799066 --- /dev/null +++ b/rmbg/sessions/base.py @@ -0,0 +1,94 @@ +""" Base Session Class """ +import os +from typing import Dict, List, Tuple + +import numpy as np +import onnxruntime as ort +from PIL import Image +from PIL.Image import Image as PILImage +from folder_paths import models_dir + +class BaseSession: + """This is a base class for managing a session with a machine learning model.""" + + def __init__( + self, + model_name: str, + sess_opts: ort.SessionOptions, + providers=None, + *args, + **kwargs + ): + """Initialize an instance of the BaseSession class.""" + self.model_name = model_name + + self.providers = [] + + if providers: + if type(providers) == list: + self.providers = providers + else: + self.providers.append(providers) + + self.inner_session = ort.InferenceSession( + str(self.__class__.download_models(*args, **kwargs)), + providers=self.providers, + sess_options=sess_opts, + ) + + def normalize( + self, + img: PILImage, + mean: Tuple[float, float, float], + std: Tuple[float, float, float], + size: Tuple[int, int], + *args, + **kwargs + ) -> Dict[str, np.ndarray]: + """ Normalizes and resizes image to be used in prediction """ + im = img.convert("RGB").resize(size, Image.Resampling.LANCZOS) + + im_ary = np.array(im) + im_ary = im_ary / np.max(im_ary) + + tmpImg = np.zeros((im_ary.shape[0], im_ary.shape[1], 3)) + tmpImg[:, :, 0] = (im_ary[:, :, 0] - mean[0]) / std[0] + tmpImg[:, :, 1] = (im_ary[:, :, 1] - mean[1]) / std[1] + tmpImg[:, :, 2] = (im_ary[:, :, 2] - mean[2]) / std[2] + + tmpImg = tmpImg.transpose((2, 0, 1)) + + return { + self.inner_session.get_inputs()[0] + .name: np.expand_dims(tmpImg, 0) + .astype(np.float32) + } + + def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]: + """Predict using the loaded model""" + raise NotImplementedError + + @classmethod + def checksum_disabled(cls, *args, **kwargs): + """Checks whether checksums are disabled""" + return os.getenv("MODEL_CHECKSUM_DISABLED", None) is not None + + @classmethod + def u2net_home(cls, *args, **kwargs): + """Returns U2Net home directory""" + return os.path.join(models_dir, 'birefnet') + # return os.path.expanduser( + # os.getenv( + # "U2NET_HOME", os.path.join(os.getenv("XDG_DATA_HOME", "~"), ".u2net") + # ) + # ) + + @classmethod + def download_models(cls, *args, **kwargs): + """Download models""" + raise NotImplementedError + + @classmethod + def name(cls, *args, **kwargs): + """Model Name""" + raise NotImplementedError diff --git a/workflow/workflow-onnx.json b/workflow/workflow-onnx.json new file mode 100644 index 0000000..731a732 --- /dev/null +++ b/workflow/workflow-onnx.json @@ -0,0 +1,231 @@ +{ + "last_node_id": 12, + "last_link_id": 17, + "nodes": [ + { + "id": 12, + "type": "PreviewImage", + "pos": { + "0": 1554, + "1": 515 + }, + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 17 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 6, + "type": "LoadImage", + "pos": { + "0": 526, + "1": 30 + }, + "size": { + "0": 513.1171875, + "1": 501.84765625 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 14 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "测试截图.png", + "image" + ] + }, + { + "id": 8, + "type": "PreviewImage", + "pos": { + "0": 1533, + "1": -18 + }, + "size": { + "0": 405.41015625, + "1": 454.2578125 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 15 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 10, + "type": "BiRefNet_onnx", + "pos": { + "0": 1118, + "1": 116 + }, + "size": { + "0": 315, + "1": 222 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 14 + } + ], + "outputs": [ + { + "name": "image", + "type": "IMAGE", + "links": [ + 15 + ], + "slot_index": 0 + }, + { + "name": "mask", + "type": "MASK", + "links": [ + 16 + ], + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "BiRefNet_onnx" + }, + "widgets_values": [ + "cpu", + "birefnet-general", + "putalpha", + 240, + 10, + 10 + ] + }, + { + "id": 11, + "type": "MaskToImage", + "pos": { + "0": 1294, + "1": 511 + }, + "size": { + "0": 264.5999755859375, + "1": 26 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 16 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 17 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "MaskToImage" + }, + "widgets_values": [] + } + ], + "links": [ + [ + 14, + 6, + 0, + 10, + 0, + "IMAGE" + ], + [ + 15, + 10, + 0, + 8, + 0, + "IMAGE" + ], + [ + 16, + 10, + 1, + 11, + 0, + "MASK" + ], + [ + 17, + 11, + 0, + 12, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.8390545288824016, + "offset": [ + -524.8041373279959, + 67.63914579313717 + ] + } + }, + "version": 0.4 +} \ No newline at end of file