Update handler.py

handler.py

@@ -1,29 +1,28 @@
-from typing import Dict, List, Any
+from typing import Dict, Any
 from PIL import Image    
 import requests
 import torch
+import numpy as np
 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, 
+            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)
+        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)
+            # Convert Tensor to NumPy array or list before returning
+            output = output.cpu().numpy().tolist()
         return output
-
-