|
from typing import Dict, List, Any |
|
from transformers import AutoProcessor, Blip2ForConditionalGeneration |
|
import base64 |
|
import torch |
|
|
|
class EndpointHandler: |
|
def __init__(self, path=""): |
|
|
|
self.processor = AutoProcessor.from_pretrained(path) |
|
self.model = Blip2ForConditionalGeneration.from_pretrained(path, device_map="auto", load_in_8bit=True).to("cuda") |
|
|
|
def __call__(self, inputs: Dict[str, Any]) -> Dict[str, str]: |
|
""" |
|
Args: |
|
inputs: |
|
Dict of image and text inputs. |
|
""" |
|
|
|
inputs = data.pop("inputs", data) |
|
image = base64.b64decode(inputs["image"]) |
|
inputs = processor(images=image, text=inputs["text"], return_tensors="pt").to("cuda", torch.float16) |
|
generated_ids = model.generate( |
|
**inputs, |
|
do_sample=decoding_method == "Nucleus sampling", |
|
temperature=1.0, |
|
length_penalty=1.0, |
|
repetition_penalty=1.5, |
|
max_length=30, |
|
min_length=1, |
|
num_beams=5, |
|
top_p=0.9, |
|
) |
|
result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() |
|
if output and output[-1] not in string.punctuation: |
|
output += "." |
|
|
|
|
|
return [{"generated_text": output}] |