Update handler.py

handler.py

@@ -1,24 +1,45 @@
 from typing import Dict, List, Any
-from transformers import pipeline, BitsAndBytesConfig
 from PIL import Image    
 import requests
 import torch
+from transformers import AutoProcessor, LlavaForConditionalGeneration
 
 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=path, model_kwargs={"quantization_config": quantization_config})
-
-    def __call__(self, data: Dict[Any, Any]):
-        totalarr = []
-        inputs = data.pop("inputs", '')
-        prompt_base = data.pop("prompt", "")
-        for x in inputs:
-            image = Image.open(requests.get(x, stream=True).raw)
-            prompt = f"USER: <image>\n{prompt_base}.Answer in one word\nASSISTANT:"
-            outputs = self.pipeline(image, prompt=prompt, generate_kwargs={"max_new_tokens": 200})
-            totalarr.append(outputs)
-        return totalarr
+        model_id = ""
+        model = LlavaForConditionalGeneration.from_pretrained(
+        model_id, 
+        torch_dtype=torch.float16, 
+        low_cpu_mem_usage=True, 
+        ).to(0)
+        processor = AutoProcessor.from_pretrained(model_id)
+
+    def __call__(self, data: Dict[str, Any]) -> List[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"
+            image = Image.open(requests.get(url, stream=True).raw)
+            prompt = "USER: <image>\nWhat are these?\nASSISTANT:"
+            output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
+        return output
+
+    
+prompt = "USER: <image>\nWhat are these?\nASSISTANT:"
+image_file = "http://images.cocodataset.org/val2017/000000039769.jpg"
+
+model = LlavaForConditionalGeneration.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float16, 
+    low_cpu_mem_usage=True, 
+).to(0)
+
+processor = AutoProcessor.from_pretrained(model_id)
+
+
+raw_image = Image.open(requests.get(image_file, stream=True).raw)
+inputs = processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)
+
+output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
+print(processor.decode(output[0][2:], skip_special_tokens=True))
+