Update handler.py

handler.py

@@ -169,24 +169,24 @@ class EndpointHandler:
                 - metadata: Dictionary with actual values used for generation
         """
         # Get inputs from request data
-        prompt = data.pop("inputs", None)
+        prompt = data.get("inputs", None)
         if not prompt:
             raise ValueError("No prompt provided in the 'inputs' field")
 
         # Get and validate resolution
-        width = data.pop("width", self.DEFAULT_WIDTH)
-        height = data.pop("height", self.DEFAULT_HEIGHT)
+        width = data.get("width", self.DEFAULT_WIDTH)
+        height = data.get("height", self.DEFAULT_HEIGHT)
         width, height = self._validate_and_adjust_resolution(width, height)
 
         # Get and validate frames and FPS
-        num_frames = data.pop("num_frames", self.DEFAULT_NUM_FRAMES)
-        fps = data.pop("fps", self.DEFAULT_FPS)
+        num_frames = data.get("num_frames", self.DEFAULT_NUM_FRAMES)
+        fps = data.get("fps", self.DEFAULT_FPS)
         num_frames, fps = self._validate_and_adjust_frames(num_frames, fps)
 
         # Get other parameters with defaults
-        guidance_scale = data.pop("guidance_scale", 7.5)
-        num_inference_steps = data.pop("num_inference_steps", self.DEFAULT_NUM_STEPS)
-        seed = data.pop("seed", -1)
+        guidance_scale = data.get("guidance_scale", 7.5)
+        num_inference_steps = data.get("num_inference_steps", self.DEFAULT_NUM_STEPS)
+        seed = data.get("seed", -1)
         seed = None if seed == -1 else int(seed)
 
         logger.info(f"Generating video with prompt: '{prompt}'")
@@ -208,7 +208,8 @@ class EndpointHandler:
                 # Check if image is provided for image-to-video generation
                 image_data = data.get("image")
                 if image_data:
-                    # Decode base64 image
+                    if image_data.startswith('data:'):
+                        image_data = image_data.split(',', 1)[1]
                     image_bytes = base64.b64decode(image_data)
                     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
                     logger.info("Using image-to-video generation mode")