Heartsync/TRELLIS2

1from typing import *
2from transformers import AutoModelForImageSegmentation
3import torch
4from torchvision import transforms
5from PIL import Image
6
7
8class BiRefNet:
9    def __init__(self, model_name: str = "ZhengPeng7/BiRefNet"):
10        self.model = AutoModelForImageSegmentation.from_pretrained(
11            model_name, trust_remote_code=True
12        )
13        self.model.eval()
14        self.transform_image = transforms.Compose(
15            [
16                transforms.Resize((1024, 1024)),
17                transforms.ToTensor(),
18                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
19            ]
20        )
21    
22    def to(self, device: str):
23        self.model.to(device)
24
25    def cuda(self):
26        self.model.cuda()
27
28    def cpu(self):
29        self.model.cpu()
30        
31    def __call__(self, image: Image.Image) -> Image.Image:
32        image_size = image.size
33        input_images = self.transform_image(image).unsqueeze(0).to("cuda")
34        # Prediction
35        with torch.no_grad():
36            preds = self.model(input_images)[-1].sigmoid().cpu()
37        pred = preds[0].squeeze()
38        pred_pil = transforms.ToPILImage()(pred)
39        mask = pred_pil.resize(image_size)
40        image.putalpha(mask)
41        return image
42