|
import base64 |
|
from io import BytesIO |
|
from typing import Dict, List, Any |
|
from transformers import Pix2StructForConditionalGeneration, AutoProcessor |
|
from PIL import Image |
|
import torch |
|
|
|
class EndpointHandler(): |
|
|
|
def __init__(self): |
|
|
|
model_name = "google/pix2struct-infographics-vqa-large" |
|
|
|
|
|
self.model = Pix2StructForConditionalGeneration.from_pretrained(model_name) |
|
self.processor = AutoProcessor.from_pretrained(model_name) |
|
self.text_prompt = None |
|
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
self.model.to(self.device) |
|
|
|
|
|
def __call__(self, data: Any) -> List[List[Dict[str, float]]]: |
|
""" |
|
Args: |
|
data (:obj:): |
|
includes the input data and the parameters for the inference. |
|
Return: |
|
a dictionary with the output of the model. The only key is `output` and the |
|
value is a list of str. |
|
""" |
|
inputs = data.pop("inputs", data) |
|
parameters = data.pop("parameters", {}) |
|
|
|
if isinstance(inputs["image"], list): |
|
img = [Image.open(BytesIO(base64.b64decode(img))) for img in inputs['image']] |
|
else: |
|
img = Image.open(BytesIO(base64.b64decode(inputs['image']))) |
|
|
|
question = inputs['question'] |
|
|
|
|
|
|
|
with torch.inference_mode(): |
|
model_inputs = self.processor(images=img, text=question, return_tensors="pt").to(self.device) |
|
|
|
raw_output = self.model.generate(**model_inputs, **parameters) |
|
|
|
decoded_output = self.processor.batch_decode(raw_output, skip_special_tokens=True) |
|
|
|
|
|
|
|
return { |
|
"output": decoded_output |
|
} |
|
|