File size: 971 Bytes
a19a224 |
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 |
from typing import Dict, List, Any
import torch
from transformers import BitsAndBytesConfig, pipeline
device = "cuda" if torch.cuda.is_available() else "cpu"
class EndpointHandler():
def __init__(self, path=""):
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16
)
# self.pipeline = pipeline("image-to-text", model="llava-hf/llava-1.5-7b-hf",
# model_kwargs={"quantization_config": quantization_config})
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
print(data)
inputs = data.pop("inputs", data)
image = data.pop("image", None)
prompt = data.pop("prompt", None)
# outputs = self.pipeline(image, prompt=prompt, generate_kwargs={
# "max_new_tokens": 200})
# return {"caption": outputs[0]["generated_text"]}
return {"image": image, "prompt": prompt}
|