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

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")  
        },)

4.8 KiB Raw Permalink Blame History

Documentation

Input types

Required

Output types

Usage tips

Source code

4.8 KiB

Raw Permalink Blame History