import model
import torch
from PIL import Image
from huggingface_hub import hf_hub_download


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

file = "./image.png" # input image

model = model.BEN_Base().to(device).eval() #init pipeline

# 从 Hugging Face 下载模型
model_path = hf_hub_download(
    repo_id="PramaLLC/BEN",
    filename="BEN_Base.pth",
    cache_dir="./models"  # 缓存目录，避免重复下载
)
model.loadcheckpoints(model_path)

image = Image.open(file)
mask, foreground = model.inference(image)

mask.save("./mask.png")
foreground.save("./foreground.png")
foreground.save("./foreground.png")