|
from typing import Dict, List, Any |
|
from transformers import AutoProcessor, AutoModelForVision2Seq |
|
import io |
|
import base64 |
|
from PIL import Image |
|
|
|
class EndpointHandler(): |
|
def __init__(self, path=""): |
|
|
|
self.device = "cuda:0" |
|
self.model = AutoModelForVision2Seq.from_pretrained("HuggingFaceM4/idefics2-8b").to(self.device) |
|
self.processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b") |
|
|
|
def __call__(self, data: dict[str, Any]) -> str: |
|
""" |
|
example: |
|
{"inputs": |
|
messages: [{ |
|
"role": "user", |
|
"content": [ |
|
{"type": "text", "text": "What’s the difference between these two images?"}, |
|
{"type": "image"}, |
|
{"type": "image"}, |
|
], |
|
}] |
|
images: [] |
|
} |
|
""" |
|
text = self.processor.apply_chat_template(data["inputs"]["messages"], add_generation_prompt=False) |
|
images = [self.decode_image_base64(img) for img in data["inputs"]["images"]] |
|
inputs = self.processor(images=images, text=text, return_tensors="pt") |
|
inputs = {k: v.to(self.device) for k,v in inputs.items()} |
|
generated_ids = self.model.generate(**inputs, max_new_tokens=500) |
|
generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True) |
|
return generated_text |
|
|
|
|
|
|
|
def decode_image_base64(self, encoded_image): |
|
""" |
|
Decodes a base64-encoded image back into a PIL image. |
|
""" |
|
|
|
img_data = base64.b64decode(encoded_image.encode("utf-8")) |
|
|
|
|
|
img_io = io.BytesIO(img_data) |
|
|
|
|
|
image = Image.open(img_io) |
|
|
|
return image |