feat: support onnx model and rename

This commit is contained in:
duhaifeng
2024-12-10 14:59:58 +08:00
parent 3bd65e2f6a
commit 1f651d34fd
17 changed files with 1026 additions and 24 deletions
+43 -9
View File
@@ -1,4 +1,4 @@
<h1 align="center">ComfyUI-BiRefNet-lite</h1>
<h1 align="center">ComfyUI-BiRefNet-Super</h1>
<p align="center">
<br><font size=5>English</font> | <a href="README_CN.md">中文</a>
@@ -18,14 +18,17 @@ Supports putalpha, naive, and alpha_matting cropping methods.
![slot](./assets/feature2.png)<br>
## News
- Nov 19, 2023: Add `mask_precision_threshold` parameter to control the accuracy threshold of the mask, default to 0.1<br>
- Nov 19, 2024: Add `mask_precision_threshold` parameter to control the accuracy threshold of the mask, default to 0.1<br>
- Dec 10, 2024:
- Add `BiRefNet_onnx` node to Support `onnx` model <br>
- Modify the name of repository from `ComfyUI-BiRefNet-lite` to `ComfyUI-BiRefNet-Super`<br>
## 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`.<br/>
Example workflows are placed in `ComfyUI-BiRefNet-Super/workflow`.<br/>
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).<br/>
@@ -46,7 +51,7 @@ ___
![](./assets/model_path.png)
Model download links:<br/>
#### Model download links:<br/>
BiRefNet: https://huggingface.co/ZhengPeng7/BiRefNet/tree/main<br/>
BiRefNet_lite-2K: https://huggingface.co/ZhengPeng7/BiRefNet_lite-2K/tree/main<br/>
BiRefNet-portrait: https://huggingface.co/ZhengPeng7/BiRefNet-portrait/tree/main<br/>
@@ -60,10 +65,39 @@ Usage of workflow.json<br/>
___
Usage of video_workflow.json<br/>
[Workflow Address](./workflow/video_workflow.json)
[Workflow Address](./workflow/video_workflow.json)
### BiRefNet_onnx
Example workflows are placed in [workflow](./workflow/workflow-onnx.json).<br/>
Place the model file in `ComfyUI/models/birefnet`.<br/>
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:<br/>
birefnet-genernal: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-epoch_244.onnx<br/>
birefnet-genernal-lite: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx<br/>
birefnet-portrait: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-portrait-epoch_150.onnx <br/>
birefnet_massive: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-massive-TR_DIS5K_TR_TEs-epoch_420.onnx<br/>
birefnet-hrsod: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-HRSOD_DHU-epoch_115.onnx <br/>
birefnet-dis: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-DIS-epoch_590.onnx <br/>
birefnet-cod: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-COD-epoch_125.onnx <br/>
## 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!
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!
+39 -7
View File
@@ -1,4 +1,4 @@
<h1 align="center">ComfyUI-BiRefNet-lite</h1>
<h1 align="center">ComfyUI-BiRefNet-Super</h1>
<p align="center">
<br> <a href="README.md">English</a> | <font size=5>中文</font>
@@ -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`中<br/>
### BiRefNet_Super | BiRefNet_Lite
示例工作流放置在`ComfyUI-BiRefNet-Super/workflow`中<br/>
加载模型有两种方式可选:一是自动下载并加载远程模型,二是加载本地模型(此时需将`load_local_model`设为true,并将`local_model_path`设置为指向models/birefnet目录下的本地模型路径,例如BiRefNet文件夹)<br/>
@@ -49,7 +51,7 @@ ___
![](./assets/model_path.png)
模型下载地址:<br/>
#### 模型下载地址:<br/>
BiRefNet: https://huggingface.co/ZhengPeng7/BiRefNet/tree/main<br/>
BiRefNet_lite-2K: https://huggingface.co/ZhengPeng7/BiRefNet_lite-2K/tree/main<br/>
BiRefNet-portrait: https://huggingface.co/ZhengPeng7/BiRefNet-portrait/tree/main<br/>
@@ -65,8 +67,38 @@ ___
工作流video_workflow.json的使用<br/>
[workflow地址](./workflow/video_workflow.json)
### BiRefNet_onnx
示例工作流放置在 [workflow](./workflow/workflow-onnx.json)中<br/>
模型下载完放置在 `ComfyUI/models/birefnet` 目录下.<br/>
模型名称如下:
- 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)
#### 模型下载地址:<br/>
birefnet-genernal: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-epoch_244.onnx<br/>
birefnet-genernal-lite: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx<br/>
birefnet-portrait: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-portrait-epoch_150.onnx <br/>
birefnet_massive: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-massive-TR_DIS5K_TR_TEs-epoch_420.onnx<br/>
birefnet-hrsod: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-HRSOD_DHU-epoch_115.onnx <br/>
birefnet-dis: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-DIS-epoch_590.onnx <br/>
birefnet-cod: https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-COD-epoch_125.onnx <br/>
## 感谢
感谢BiRefNet仓库的所有作者 [ZhengPeng7/BiRefNet](https://github.com/zhengpeng7/birefnet)
部分代码参考了 [MoonHugo/ComfyUI-BiRefNet-Hugo](https://github.com/MoonHugo/ComfyUI-BiRefNet-Hugo) 感谢!
部分代码参考了 [MoonHugo/ComfyUI-BiRefNet-Hugo](https://github.com/MoonHugo/ComfyUI-BiRefNet-Hugo) 感谢!
部分代码参考了 [danielgatis/rembg](https://github.com/danielgatis/rembg). 感谢!
BIN
View File
Binary file not shown.

After

Width:  |  Height:  |  Size: 2.1 MiB

BIN
View File
Binary file not shown.

After

Width:  |  Height:  |  Size: 153 KiB

+120 -4
View File
@@ -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"
}
+4 -4
View File
@@ -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 = ""
+52
View File
@@ -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"
+52
View File
@@ -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"
+93
View File
@@ -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"
+52
View File
@@ -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"
+52
View File
@@ -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"
+52
View File
@@ -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"
+53
View File
@@ -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"
+46
View File
@@ -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)
+43
View File
@@ -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())
+94
View File
@@ -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
+231
View File
@@ -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
}