Spaces:
Runtime error
Runtime error
File size: 1,131 Bytes
67d6f5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
from dotenv import load_dotenv
from transformers import BlipForConditionalGeneration, BlipProcessor
import torch
import litserve as ls
import os
load_dotenv()
hf_token = os.getenv("HUGGINGFACE")
class RedionesBlipModel():
def __init__(self):
self.model_name = "Salesforce/blip-image-captioning-base"
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.token = hf_token
def setup(self, device):
device = self.device
self.model = BlipForConditionalGeneration.from_pretrained(self.model_name,
use_auth_token=self.token,
)
self.tokenizer = BlipProcessor.from_pretrained(self.model_name, use_auth_token=self.token)
self.model.to(device)
self.model.eval()
def predict(self, image):
input_text = self.tokenizer(image, return_tensors="pt")
outputs = self.model.generate(input_ids = input_text["input_ids"].to(self.device), max_new_tokens=50)
return outputs
|