File size: 1,368 Bytes
57be487 aa93800 57be487 7d619b2 57be487 |
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 32 33 34 35 36 37 38 |
from typing import Dict, List, Any
from transformers import AutoProcessor, Blip2ForConditionalGeneration
import base64
import torch
class EndpointHandler:
def __init__(self, path=""):
# load model and processor from 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.
"""
# process input
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}] |