import numpy as np from transformers import Blip2Processor, Blip2ForConditionalGeneration, BlipForQuestionAnswering, BitsAndBytesConfig from transformers import AutoProcessor, AutoModelForCausalLM from typing import Dict, List, Any from PIL import Image from transformers import pipeline import requests import torch from io import BytesIO import base64 class EndpointHandler(): def __init__(self, path=""): self.device = "cuda:0" if torch.cuda.is_available() else "cpu" print("device:",self.device) # self.model_base = "Salesforce/blip2-opt-2.7b" # self.model_name = "sooh-j/blip2-vizwizqa" self.model_name = "Salesforce/blip2-opt-2.7b" self.processor = AutoProcessor.from_pretrained(self.model_name) self.model = Blip2ForConditionalGeneration.from_pretrained(self.model_name, device_map="auto", ).to(self.device) def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: """ data args: inputs (:obj: `str` | `PIL.Image` | `np.array`) kwargs Return: A :obj:`list` | `dict`: will be serialized and returned """ # await hf.visualQuestionAnswering({ # model: 'dandelin/vilt-b32-finetuned-vqa', # inputs: { # question: 'How many cats are lying down?', # image: await (await fetch('https://placekitten.com/300/300')).blob() # } # }) inputs = data.get("inputs") imageBase64 = inputs.get("image") question = inputs.get("question") if ('http:' in imageBase64) or ('https:' in imageBase64): image = Image.open(requests.get(imageBase64, stream=True).raw) else: image = Image.open(BytesIO(base64.b64decode(imageBase64.split(",")[0].encode()))) prompt = f"Question: {question}, Answer:" processed = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device) with torch.no_grad(): out = self.model.generate(**processed, max_new_tokens=20, temperature = 0.5, do_sample=True, top_k=50, top_p=0.9, repetition_penalty=1.2 ).to(self.device) result = {} text_output = self.processor.decode(out[0], skip_special_tokens=True) result["text_output"] = text_output score = 0 return [{"answer":text_output,"score":score}]