BiRefNet_pipe.py

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()