File size: 251 Bytes
d73d557
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
import torch
from transformers import Pipeline


class BiRefNetPipe(Pipeline):
    def __init__(self, **kwargs):
        Pipeline.__init__(self, **kwargs)
        self.model.to(['cpu', 0][torch.cuda.is_available()])
        self.model.eval()