Vvladmandiccleanup rembg
7588ec06创建于 3月20日历史提交
model = None


def remove(image, refine: bool = True):
    global model # pylint: disable=global-statement
    from modules import shared, devices

    if model is None:
        from huggingface_hub import hf_hub_download
        from modules.rembg.ben2_model import BEN_Base
        model = BEN_Base()
        model_file = hf_hub_download(
            repo_id='PramaLLC/BEN2',
            filename='BEN2_Base.pth',
            cache_dir=shared.opts.hfcache_dir)
        model.loadcheckpoints(model_file)
        model = model.to(device=devices.device, dtype=devices.dtype).eval()

    model = model.to(device=devices.device)
    foreground = model.inference(image, refine_foreground=refine)
    model = model.to(device=devices.cpu)
    if foreground is None:
        return image
    return foreground