Update handler.py

handler.py

@@ -5,17 +5,18 @@ import requests
 
 class EndpointHandler():
     def __init__(self, path=""):
-        self.pipeline = pipeline("image-to-text",model=path)
+        print("init")
+        # self.pipeline = pipeline("image-to-text", model=path)
 
+    def __call__(self, data: Dict[str, Any]):
+        totalarr = []
+        inputs = data.pop("inputs", [])
+        prompt_base = data.pop("prompt", "")
         
-    def __call__(self, data: Dict[str, Union[List[str], str]]):
-            totalarr = []
-            input = data.pop("inputs", data)
-            prompt = data.pop("prompt", "")
-            for link in input:
-                image = Image.open(requests.get(link, stream=True).raw)
-                prompt = "USER: <image>\n" + prompt + ".Answer in one word\nASSISTANT:"
-                outputs = self.pipeline(image, prompt=prompt, generate_kwargs={"max_new_tokens": 200})
-                totalarr.append(outputs)
-            return totalarr
-
+        for link in inputs:
+            image = Image.open(requests.get(link, 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