Needs to be inputs

handler.py

@@ -13,7 +13,8 @@ class EndpointHandler():
         self.model = Pix2StructVisionModel.from_pretrained("google/pix2struct-textcaps-base")
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
-        image = Image.open(requests.get(data["url"], stream=True).raw)
+        url = data.pop("inputs", data)
+        image = Image.open(requests.get(url, stream=True).raw)
         inputs = self.processor(images=image, return_tensors="pt")
 
         with torch.no_grad():
@@ -23,8 +24,6 @@ class EndpointHandler():
         embedding = torch.mean(last_hidden_state, dim=1).flatten().tolist()
         return {"embedding": embedding}
 
-"""
 handler = EndpointHandler()
-output = handler({"url": "https://figma-staging-api.s3.us-west-2.amazonaws.com/images/a8c6a0cc-c022-4f3a-9fc5-ac8582c964dd"})
+output = handler({"inputs": "https://figma-staging-api.s3.us-west-2.amazonaws.com/images/a8c6a0cc-c022-4f3a-9fc5-ac8582c964dd"})
 print(output)
-"""