File size: 1,701 Bytes
1e05993 7feaf20 f496fb3 6a72b4d 1e05993 6a67ddb 7feaf20 58ff319 e6b2e2d 1e05993 3b62fae 1a792e3 f496fb3 185c874 e6b2e2d 5701230 142f207 1e05993 eb098cf 928a2e1 142f207 eb098cf a3713b3 b97a340 3b62fae b97a340 928a2e1 |
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 39 40 41 42 43 |
from typing import Dict, Any
from PIL import Image
import requests
import flash-attn
import torch
import numpy as np
from transformers import AutoProcessor, LlavaForConditionalGeneration, BitsAndBytesConfig
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,
use_flash_attention_2=True
)
self.processor = AutoProcessor.from_pretrained(model_id)
def __call__(self, data: Dict[list, Any]):
parameters = data.pop("inputs", data)
givenprompt = data.pop("prompt", data)
outputs = []
print(parameters)
prompt = f"USER: <image>\n{givenprompt}?\nASSISTANT:"
for link in parameters:
try:
# Fetch image from URL
response = requests.get(link, stream=True)
response.raise_for_status() # Raise an exception for 4xx or 5xx status codes
raw_image = Image.open(response.raw)
# Process image and generate output
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)
readable = self.processor.decode(output[0][2:], skip_special_tokens=True)
outputs.append(readable)
except Exception as e:
# Handle any exceptions and log the error
outputs.append(f"Error processing image from {link}: {str(e)}")
return outputs
|