BiRefNet / BiRefNet_pipe.py
ZhengPeng7's picture
Move all BiRefNet github codes to the first level directory.
79a6660
raw
history blame
242 Bytes
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()