safer parsing

handler.py

@@ -31,14 +31,14 @@ class EndpointHandler:
             A :obj:`list` | `dict`: will be serialized and returned
         """
         # get inputs
-        inputs = data.pop("inputs", None)
+        inputs = data.get("inputs")
 
         if not inputs:
             return f"Inputs not in payload got {data}"
         
         # get additional date field0
-        prompt = inputs.pop("prompt", None)
-        image_url = inputs.pop("image", None)
+        prompt = inputs.get("prompt")
+        image_url = inputs.get("image")
 
         if image_url is None:
             return "You need to upload an image URL for LLaVA to work."
@@ -72,6 +72,9 @@ class EndpointHandler:
 
                 output = self.model.generate(**inputs, max_new_tokens=100)
 
+                if not output:
+                    return 'Model failed to generate'
+
                 clean = self.processor.decode(output[0], skip_special_tokens=True)
 
                 return clean