98 lines
4.8 KiB
Markdown
98 lines
4.8 KiB
Markdown
# Documentation
|
|
- Class name: BiRefNet_Loader
|
|
- Category: BiRefNet🌟
|
|
- Output node: False
|
|
- Repo Ref: https://github.com/moon7star9/ComfyUI_BiRefNet_Universal
|
|
|
|
BiRefNet_Loader节点用于加载BiRefNet模型，这是一个专门用于图像分割和背景移除的深度学习模型。该节点支持多个预训练模型版本，可以根据不同的使用场景选择合适的模型。节点会自动处理模型的下载和加载，支持本地缓存，所下载的模型皆在models/BiRefNet文件夹下。
|
|
|
|
# Input types
|
|
## Required
|
|
- model_version
|
|
- 'model_version'参数允许用户选择要加载的BiRefNet模型版本。不同版本针对不同的使用场景和分辨率进行了优化，包括标准版本、高分辨率版本、轻量级版本等，名称与huggingface上的模型仓库相对应。选择建议：带HR字样的都是以2048*2048图像进行推理，因此边缘分辨率更高，通用抠图优先选择BiRefNet_HR比较高清，BiRefNet检测主体可能更好；有lite字样为轻量化模型，其中lite 2K模型更适合抠2k横屏图像；如果输入的图像尺寸较小，通用模型效果差的话，可以试试512x512模型；有matting字样的更加适合抠透明物体、纱、网、毛发等；有portrait字样的更适合抠人像；如果主体定位结果差，建议选择从DIS5K开始的几个模型，这些模型是以训练数据集命名的，偏二值化抠图，因此边缘没有很好的过渡，可以后续继续添加vitmatte等算法进一步处理。
|
|
- Comfy dtype: COMBO
|
|
- Python dtype: str
|
|
- device
|
|
- 'device'参数指定模型运行的硬件设备。可以设置为"auto"让系统自动选择最适合的设备，或手动指定特定设备（如"cuda"、"cpu"、"mps"等）。
|
|
- Comfy dtype: COMBO
|
|
- Python dtype: str
|
|
|
|
# Output types
|
|
- BIREFNET_MODEL
|
|
- 输出加载好的BiRefNet模型实例，包含模型本身、分辨率设置、设备信息和精度配置等。这个输出可以直接连接到BiRefNet_Remove_Background节点使用。
|
|
- Comfy dtype: BIREFNET_MODEL
|
|
- Python dtype: dict
|
|
|
|
# Usage tips
|
|
- Infra type: GPU/CPU
|
|
- 首次使用时会自动下载模型文件并缓存到本地
|
|
- 推荐使用GPU运行以获得更好的性能
|
|
- 当选择"auto"设备时，节点会自动检测并使用可用的最佳硬件
|
|
- 在CUDA设备上会自动使用半精度(FP16)以优化性能和内存使用
|
|
|
|
# Source code
|
|
```python
|
|
class BiRefNet_Loader:
|
|
def __init__(self):
|
|
self.loaded_model = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_version": (list(MODEL_VERSIONS.keys()), {"default": "BiRefNet_HR"}),
|
|
"device": (["auto", "cuda", "cpu", "mps", "xpu", "meta"], {"default": "auto"})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BIREFNET_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "BiRefNet🌟"
|
|
|
|
|
|
|
|
def load_model(self, model_version, device):
|
|
device = get_device_by_name(device)
|
|
model_name, resolution = MODEL_VERSIONS[model_version]
|
|
local_model_path = get_model_path(model_name)
|
|
|
|
# 首先尝试加载本地模型
|
|
try:
|
|
if os.path.exists(local_model_path):
|
|
print(f"\033[92mLoading local model from: {local_model_path}\033[0m")
|
|
model = AutoModelForImageSegmentation.from_pretrained(
|
|
local_model_path,
|
|
trust_remote_code=True
|
|
)
|
|
else:
|
|
print(f"\033[93mLocal model not found, downloading from HuggingFace: {model_name}\033[0m")
|
|
model = AutoModelForImageSegmentation.from_pretrained(
|
|
f"ZhengPeng7/{model_name}",
|
|
trust_remote_code=True,
|
|
cache_dir=local_model_path
|
|
)
|
|
except Exception as e:
|
|
print(f"\033[91mError loading local model: {str(e)}\033[0m")
|
|
print("\033[93mFallback to downloading from HuggingFace\033[0m")
|
|
try:
|
|
model = AutoModelForImageSegmentation.from_pretrained(
|
|
f"ZhengPeng7/{model_name}",
|
|
trust_remote_code=True,
|
|
cache_dir=local_model_path
|
|
)
|
|
except Exception as download_error:
|
|
raise RuntimeError(f"Failed to load model both locally and from HuggingFace: {str(download_error)}")
|
|
|
|
model.to(device)
|
|
model.eval()
|
|
if device == "cuda":
|
|
model.half()
|
|
|
|
return ({
|
|
"model": model,
|
|
"resolution": resolution,
|
|
"device": device,
|
|
"half_precision": (device == "cuda")
|
|
},)
|