| from typing import Dict, List, Any | |
| from PIL import Image | |
| import requests | |
| import torch | |
| from transformers import AutoProcessor, LlavaForConditionalGeneration | |
| class EndpointHandler(): | |
| def __init__(self, path=""): | |
| model_id = path | |
| self.model = LlavaForConditionalGeneration.from_pretrained( | |
| model_id, | |
| torch_dtype=torch.float16, | |
| low_cpu_mem_usage=True, | |
| ).to(0) | |
| self.processor = AutoProcessor.from_pretrained(model_id) | |
| def __call__(self, data: Dict[str, Any]): | |
| parameters = data.pop("inputs",data) | |
| inputs = data.pop("inputs", data) | |
| if parameters is not None: | |
| url = "http://images.cocodataset.org/val2017/000000039769.jpg" | |
| prompt = "USER: <image>\nWhat are these?\nASSISTANT:" | |
| raw_image = Image.open(requests.get(url, stream=True).raw) | |
| inputs = self.processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16) | |
| output = self.model.generate(**inputs, max_new_tokens=200, do_sample=False) | |
| return output | |