Fix handler to support 'inputs' key

handler.py

@@ -13,22 +13,21 @@ class EndpointHandler:
         self.model.to(self.device)
 
     def __call__(self, data):
-        # Expecting: { "image_url": "https://example.com/image.jpg" }
-        if isinstance(data, dict) and "image_url" in data:
-            image_url = data["image_url"]
-            try:
-                response = requests.get(image_url)
-                response.raise_for_status()
-                image = Image.open(io.BytesIO(response.content)).convert("RGB")
-            except Exception as e:
-                return {"error": f"Failed to load image from URL: {str(e)}"}
-        else:
+        # Unwrap Hugging Face "inputs" key if present
+        if isinstance(data, dict) and "inputs" in data:
+            data = data["inputs"]
+
+        image_url = data.get("image_url")
+        if not image_url:
             return {"error": "Please provide an 'image_url' in the JSON payload."}
 
-        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+        try:
+            response = requests.get(image_url)
+            response.raise_for_status()
+            image = Image.open(io.BytesIO(response.content)).convert("RGB")
+        except Exception as e:
+            return {"error": f"Failed to load image: {str(e)}"}
 
+        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
         with torch.no_grad():
             out = self.model.generate(**inputs)
-
-        caption = self.processor.decode(out[0], skip_special_tokens=True)
-        return {"caption": caption}