Update handler.py

handler.py

@@ -3,7 +3,8 @@ from PIL import Image
 import requests
 import torch
 import numpy as np
-from transformers import AutoProcessor, LlavaForConditionalGeneration
+from transformers import AutoProcessor, LlavaForConditionalGeneration, BitsAndBytesConfig
+
 
 class EndpointHandler():
     def __init__(self, path=""):
@@ -12,16 +13,16 @@ class EndpointHandler():
             model_id, 
             torch_dtype=torch.float16, 
             low_cpu_mem_usage=True, 
+            load_in_4bit=True
         ).to(0)
         self.processor = AutoProcessor.from_pretrained(model_id)
 
     def __call__(self, data: Dict[str, Any]):
         parameters = 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)
-            readable = self.processor.decode(output[0][2:], skip_special_tokens=True)
-        return readable
+        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)
+        print(self.processor.decode(output[0][2:], skip_special_tokens=True))
+        return output