Update handler.py

handler.py

@@ -50,60 +50,75 @@ class EndpointHandler:
             use_auth_token=hf_token
         ).eval().cuda()
 
+    def load_image(self, image_input):
+        """
+        Load image from URL, base64 or file input.
+        """
+        if image_input.startswith("http"):
+            # Load image from URL
+            try:
+                resp = requests.get(image_input, verify=False)
+                image = Image.open(BytesIO(resp.content)).convert("RGB")
+                return image
+            except Exception as e:
+                raise ValueError(f"Failed to fetch image from URL: {e}")
+        
+        elif image_input.startswith("data:image"):
+            # Load base64 encoded image
+            try:
+                image = Image.open(BytesIO(base64.b64decode(image_input.split(",")[1]))).convert("RGB")
+                return image
+            except Exception as e:
+                raise ValueError(f"Invalid base64 image format: {e}")
+        
+        else:
+            # Load image from file
+            try:
+                image = Image.open(image_input).convert("RGB")
+                return image
+            except Exception as e:
+                raise ValueError(f"Failed to open image from file path: {e}")
+    
     def predict(self, request):
         """
         Expected input:
         {
-            "inputs": {
-                "image": "<image URL or base64>",
-                "msgs": [{"role": "user", "content": "<image>\nWhat is this?"}],
-                "stream": false
-            }
+            "image": "<image URL, file path, or base64>",
+            "question": "What is this?",
+            "stream": false
         }
         """
-        # Extract data from the "inputs" key in the request
-        data = request.get("inputs", {})  # ✅ Extract from "inputs" key
-        image_input = data.get("image")
-        msgs = data.get("msgs", [])
-        stream = data.get("stream", False)
+        image_input = request.get("image")
+        question = request.get("question", "What is in the image?")
+        stream = request.get("stream", False)
 
-        if not image_input or not msgs:
-            return {"error": "Missing 'image' or 'msgs'."}
+        if not image_input:
+            return {"error": "Missing image."}
 
         try:
-            # Load image from URL or base64
-            if image_input.startswith("http"):
-                resp = requests.get(image_input, verify=False)
-                image = Image.open(BytesIO(resp.content)).convert("RGB")
-            else:
-                image = Image.open(BytesIO(base64.b64decode(image_input))).convert("RGB")
-        except Exception as e:
-            return {"error": f"Invalid image format or URL: {e}"}
+            # Load image using the new load_image function
+            image = self.load_image(image_input)
 
-        try:
-            # Run inference with or without streaming
-            if stream:
-                generated_text = ""
-                for chunk in self.model.chat(
-                    image=image,
-                    msgs=msgs,
-                    tokenizer=self.tokenizer,
-                    sampling=True,
-                    stream=True
-                ):
-                    generated_text += chunk
-                return {"output": generated_text}
-            else:
-                output = self.model.chat(
-                    image=image,
-                    msgs=msgs,
-                    tokenizer=self.tokenizer,
-                    sampling=True,
-                    stream=False
-                )
-                return {"output": output}
-        except Exception as e:
-            return {"error": f"Inference failed: {e}"}
+            # Prepare message with <image> placeholder
+            msgs = [{"role": "user", "content": f"<image>\n{question}"}]
+
+            try:
+                if stream:
+                    generated_text = ""
+                    for chunk in self.model.chat(
+                        image=image, msgs=msgs, tokenizer=self.tokenizer,
+                        sampling=True, stream=True
+                    ):
+                        generated_text += chunk
+                    return {"output": generated_text}
+                else:
+                    output = self.model.chat(image=image, msgs=msgs, tokenizer=self.tokenizer)
+                    return {"output": output}
+            except Exception as e:
+                return {"error": f"Inference failed: {e}"}
+
+        except ValueError as e:
+            return {"error": f"Image processing error: {e}"}
 
     def __call__(self, data):
         """