support BirefNet_HR and float16
This commit is contained in:
刘雪峰
@@ -24,6 +24,7 @@ Support the use of new and old versions of BiRefNet models
|
||||
### The available newest models are:
|
||||
|
||||
- General: A pre-trained model for general use cases.
|
||||
- General-HR: A pre-trained model for general use cases which shows great performance on higher resolution images (2048x2048).
|
||||
- General-Lite: A light pre-trained model for general use cases.
|
||||
- General-Lite-2K: A light pre-trained model for general use cases in high resolution (2560x1440).
|
||||
- Portrait: A pre-trained model for human portraits.
|
||||
@@ -37,6 +38,7 @@ Model files go here (when use AutoDownloadBiRefNetModel automatically downloaded
|
||||
|
||||
If necessary, they can be downloaded from:
|
||||
- [General](https://huggingface.co/ZhengPeng7/BiRefNet/resolve/main/model.safetensors) ➔ `model.safetensors` must be renamed `General.safetensors`
|
||||
- [General-HR](https://huggingface.co/ZhengPeng7/BiRefNet_HR/resolve/main/model.safetensors) ➔ `model.safetensors` must be renamed `General-HR.safetensors`
|
||||
- [General-Lite](https://huggingface.co/ZhengPeng7/BiRefNet_T/resolve/main/model.safetensors) ➔ `model.safetensors` must be renamed `General-Lite.safetensors`
|
||||
- [General-Lite-2K](https://huggingface.co/ZhengPeng7/BiRefNet_lite-2K/resolve/main/model.safetensors) ➔ `model.safetensors` must be renamed `General-Lite-2K.safetensors`
|
||||
- [Portrait](https://huggingface.co/ZhengPeng7/BiRefNet-portrait/resolve/main/model.safetensors) ➔ `model.safetensors` must be renamed `Portrait.safetensors`
|
||||
|
||||
+4
-2
@@ -24,8 +24,9 @@
|
||||
### 最新的模型:
|
||||
|
||||
- General: 用于一般用例的预训练模型。
|
||||
- General-HR: 用于一般用例的预训练模型，在更高分辨率的图像上表现出色(训练分辨率2048x2048)。
|
||||
- General-Lite: 用于一般用例的轻量级预训练模型。
|
||||
- General-Lite-2K: 用于一般用例的轻量级预训练模型，适用于高分辨率图像。 (最佳分辨率2560x1440).
|
||||
- General-Lite-2K: 用于一般用例的轻量级预训练模型，适用于高分辨率图像(最佳分辨率2560x1440)。
|
||||
- Portrait: 人物肖像预训练模型。
|
||||
- Matting: 一种使用无trimap matting的预训练模型。
|
||||
- DIS: 一种用于二分图像分割（DIS）的预训练模型。
|
||||
@@ -33,10 +34,11 @@
|
||||
- COD: 一种用于隐蔽目标检测（COD）的预训练模型。
|
||||
- DIS-TR_TEs: 具有大量数据集的预训练模型。
|
||||
|
||||
模型文件放在`${comfyui_rootpath}/models/BiRefNet`（当使用AutoDownloadBiRefNetModel时，则会自动下载模型）.
|
||||
模型文件放在`${comfyui_rootpath}/models/BiRefNet`（当使用AutoDownloadBiRefNetModel时，则会自动下载模型）。
|
||||
|
||||
也可以手动下载模型:
|
||||
- [General](https://huggingface.co/ZhengPeng7/BiRefNet/resolve/main/model.safetensors) ➔ `model.safetensors` 重命名为 `General.safetensors`
|
||||
- [General-HR](https://huggingface.co/ZhengPeng7/BiRefNet_HR/resolve/main/model.safetensors) ➔ `model.safetensors` 重命名为 `General-HR.safetensors`
|
||||
- [General-Lite](https://huggingface.co/ZhengPeng7/BiRefNet_T/resolve/main/model.safetensors) ➔ `model.safetensors` 重命名为 `General-Lite.safetensors`
|
||||
- [General-Lite-2K](https://huggingface.co/ZhengPeng7/BiRefNet_lite-2K/resolve/main/model.safetensors) ➔ `model.safetensors` 重命名为 `General-Lite-2K.safetensors`
|
||||
- [Portrait](https://huggingface.co/ZhengPeng7/BiRefNet-portrait/resolve/main/model.safetensors) ➔ `model.safetensors` 重命名为 `Portrait.safetensors`
|
||||
|
||||
+1
-4
@@ -61,7 +61,7 @@ class Config:
|
||||
'DIS5K': -40,
|
||||
'COD': -20,
|
||||
'HRSOD': -20,
|
||||
'General': -40,
|
||||
'General': -20,
|
||||
'General-2K': -20,
|
||||
'Matting': -20,
|
||||
}[self.task]
|
||||
@@ -154,9 +154,6 @@ class Config:
|
||||
self.lambdas_cls = {
|
||||
'ce': 5.0
|
||||
}
|
||||
# Adv
|
||||
self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training
|
||||
self.lambda_adv_d = 3. * (self.lambda_adv_g > 0)
|
||||
|
||||
# PATH settings - inactive
|
||||
# https://drive.google.com/drive/folders/1cmce_emsS8A5ha5XT2c_CZiJzlLM81ms
|
||||
|
||||
@@ -68,9 +68,9 @@ def window_reverse(windows, window_size, H, W):
|
||||
Returns:
|
||||
x: (B, H, W, C)
|
||||
"""
|
||||
B = int(windows.shape[0] / (H * W / window_size / window_size))
|
||||
x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
|
||||
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
|
||||
C = int(windows.shape[-1])
|
||||
x = windows.view(-1, H // window_size, W // window_size, window_size, window_size, C)
|
||||
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, H, W, C)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
+21
-9
@@ -9,7 +9,7 @@ import folder_paths
|
||||
from birefnet.models.birefnet import BiRefNet
|
||||
from birefnet_old.models.birefnet import BiRefNet as OldBiRefNet
|
||||
from birefnet.utils import check_state_dict
|
||||
from .util import tensor_to_pil, apply_mask_to_image, normalize_mask, refine_foreground, filter_mask, add_mask_as_alpha
|
||||
from .util import refine_foreground, filter_mask, add_mask_as_alpha
|
||||
deviceType = model_management.get_torch_device().type
|
||||
|
||||
models_dir_key = "birefnet"
|
||||
@@ -18,6 +18,7 @@ models_path_default = folder_paths.get_folder_paths(models_dir_key)[0]
|
||||
|
||||
usage_to_weights_file = {
|
||||
'General': 'BiRefNet',
|
||||
'General-HR': 'BiRefNet_HR',
|
||||
'General-Lite': 'BiRefNet_T',
|
||||
'General-Lite-2K': 'BiRefNet_lite-2K',
|
||||
'Portrait': 'BiRefNet-portrait',
|
||||
@@ -28,7 +29,7 @@ usage_to_weights_file = {
|
||||
'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs'
|
||||
}
|
||||
|
||||
modelNameList = ['General', 'General-Lite', 'General-Lite-2K', 'Portrait', 'Matting', 'DIS', 'HRSOD', 'COD', 'DIS-TR_TEs']
|
||||
modelNameList = ['General', 'General-HR', 'General-Lite', 'General-Lite-2K', 'Portrait', 'Matting', 'DIS', 'HRSOD', 'COD', 'DIS-TR_TEs']
|
||||
|
||||
|
||||
def get_model_path(model_name):
|
||||
@@ -90,6 +91,11 @@ class ImagePreprocessor:
|
||||
VERSION = ["old", "v1"]
|
||||
old_models_name = ["BiRefNet-DIS_ep580.pth", "BiRefNet-ep480.pth"]
|
||||
|
||||
torch_dtype={
|
||||
"float16": torch.float16,
|
||||
"float32": torch.float32,
|
||||
"bfloat16": torch.bfloat16,
|
||||
}
|
||||
|
||||
class AutoDownloadBiRefNetModel:
|
||||
|
||||
@@ -99,6 +105,9 @@ class AutoDownloadBiRefNetModel:
|
||||
"required": {
|
||||
"model_name": (modelNameList,),
|
||||
"device": (["AUTO", "CPU"],)
|
||||
},
|
||||
"optional": {
|
||||
"dtype": (["float32", "float16"], {"default": "float32"})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -108,7 +117,7 @@ class AutoDownloadBiRefNetModel:
|
||||
CATEGORY = "image/BiRefNet"
|
||||
DESCRIPTION = "Auto download BiRefNet model from huggingface to models/BiRefNet/{model_name}.safetensors"
|
||||
|
||||
def load_model(self, model_name, device):
|
||||
def load_model(self, model_name, device, dtype="float32"):
|
||||
bb_index = 3 if model_name == "General-Lite" or model_name == "General-Lite-2K" else 6
|
||||
biRefNet_model = BiRefNet(bb_pretrained=False, bb_index=bb_index)
|
||||
model_file_name = f'{model_name}.safetensors'
|
||||
@@ -122,7 +131,7 @@ class AutoDownloadBiRefNetModel:
|
||||
device_type = "cpu"
|
||||
state_dict = safetensors.torch.load_file(model_full_path, device=device_type)
|
||||
biRefNet_model.load_state_dict(state_dict)
|
||||
biRefNet_model.to(device_type)
|
||||
biRefNet_model.to(device_type, dtype=torch_dtype[dtype])
|
||||
biRefNet_model.eval()
|
||||
return [(biRefNet_model, VERSION[1])]
|
||||
|
||||
@@ -137,7 +146,8 @@ class LoadRembgByBiRefNetModel:
|
||||
"device": (["AUTO", "CPU"], )
|
||||
},
|
||||
"optional": {
|
||||
"use_weight": ("BOOLEAN", {"default": False})
|
||||
"use_weight": ("BOOLEAN", {"default": False}),
|
||||
"dtype": (["float32", "float16"], {"default": "float32"})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -147,7 +157,7 @@ class LoadRembgByBiRefNetModel:
|
||||
CATEGORY = "rembg/BiRefNet"
|
||||
DESCRIPTION = "Load BiRefNet model from folder models/BiRefNet or the path of birefnet configured in the extra YAML file"
|
||||
|
||||
def load_model(self, model, device, use_weight=False):
|
||||
def load_model(self, model, device, use_weight=False, dtype="float32"):
|
||||
if model in old_models_name:
|
||||
version = VERSION[0]
|
||||
biRefNet_model = OldBiRefNet(bb_pretrained=use_weight)
|
||||
@@ -168,7 +178,7 @@ class LoadRembgByBiRefNetModel:
|
||||
state_dict = check_state_dict(state_dict)
|
||||
|
||||
biRefNet_model.load_state_dict(state_dict)
|
||||
biRefNet_model.to(device_type)
|
||||
biRefNet_model.to(device_type, dtype=torch_dtype[dtype])
|
||||
biRefNet_model.eval()
|
||||
return [(biRefNet_model, version)]
|
||||
|
||||
@@ -211,7 +221,9 @@ class GetMaskByBiRefNet:
|
||||
|
||||
def get_mask(self, model, images, width=1024, height=1024, upscale_method='bilinear', mask_threshold=0.000):
|
||||
model, version = model
|
||||
model_device_type = next(model.parameters()).device.type
|
||||
one_torch = next(model.parameters())
|
||||
model_device_type = one_torch.device.type
|
||||
model_dtype = one_torch.dtype
|
||||
b, h, w, c = images.shape
|
||||
image_bchw = images.permute(0, 3, 1, 2)
|
||||
|
||||
@@ -226,7 +238,7 @@ class GetMaskByBiRefNet:
|
||||
_mask_bchw = []
|
||||
for each_image in im_tensor:
|
||||
with torch.no_grad():
|
||||
each_mask = model(each_image.unsqueeze(0).to(model_device_type))[-1].sigmoid().cpu()
|
||||
each_mask = model(each_image.unsqueeze(0).to(model_device_type, dtype=model_dtype))[-1].sigmoid().cpu().float()
|
||||
_mask_bchw.append(each_mask)
|
||||
del each_mask
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "comfyui_birefnet_ll"
|
||||
description = "Sync with version of BiRefNet. NODES:AutoDownloadBiRefNetModel, LoadRembgByBiRefNetModel, RembgByBiRefNet, RembgByBiRefNetAdvanced, GetMaskByBiRefNet, BlurFusionForegroundEstimation."
|
||||
version = "1.1.0"
|
||||
version = "1.1.1"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["numpy", "opencv-python", "timm"]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user