| from typing import Any, Dict |
| from transformers import Blip2Processor, Blip2ForConditionalGeneration |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| |
| self.processor = Blip2Processor.from_pretrained(path) |
| self.model = Blip2ForConditionalGeneration.from_pretrained(path) |
| self.device = "cuda" |
|
|
| self.model.to(self.device) |
|
|
| def __call__(self, data: Dict[str, Any]) -> Dict[str, str]: |
| |
| image = data.pop("image", data) |
| text = data.pop("text", data) |
|
|
| inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device) |
| generated_ids = self.model.generate(**inputs) |
| generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() |
|
|
| return [{"answer": generated_text}] |
