Update handler.py

handler.py

@@ -23,23 +23,32 @@ class EndpointHandler:
     def __call__(self, data: dict):
         """Custom call function for Hugging Face Inference Endpoints."""
         try:
-            inputs = data.pop("inputs", data)
-            encoded_image = data.pop("image", None)
-            encoded_mask_image = data.pop("mask_image", None)
-            
-            num_inference_steps = data.pop("num_inference_steps", 25)
-            guidance_scale = data.pop("guidance_scale", 7.5)
-            negative_prompt = data.pop("negative_prompt", None)
-            height = data.pop("height", None)
-            width = data.pop("width", None)
-            
-            # Process images
-            if encoded_image and encoded_mask_image:
-                image = self.decode_base64_image(encoded_image)
-                mask_image = self.decode_base64_image(encoded_mask_image)
-            else:
-                raise ValueError("Both image and mask_image are required")
-            
+            # Extract inputs from JSON payload
+            inputs = data.get("inputs", "")
+            encoded_image = data.get("image", None)
+            encoded_mask_image = data.get("mask_image", None)
+
+            # Extract optional parameters with default values
+            num_inference_steps = data.get("num_inference_steps", 25)
+            guidance_scale = data.get("guidance_scale", 7.5)
+            negative_prompt = data.get("negative_prompt", None)
+            height = data.get("height", None)
+            width = data.get("width", None)
+
+            # Ensure both images are provided
+            if not encoded_image or not encoded_mask_image:
+                raise ValueError("Both 'image' and 'mask_image' are required in base64 format.")
+
+            # Decode base64 images
+            image = self.decode_base64_image(encoded_image)
+            mask_image = self.decode_base64_image(encoded_mask_image)
+
+            print("\n--- Running Inference ---")
+            print(f"Prompt: {inputs}")
+            print(f"Steps: {num_inference_steps}, Guidance Scale: {guidance_scale}")
+            print(f"Negative Prompt: {negative_prompt}")
+            print(f"Image Size: {image.size}, Mask Size: {mask_image.size}")
+
             # Run inference
             output_image = self.pipeline(
                 prompt=inputs,
@@ -52,17 +61,22 @@ class EndpointHandler:
                 height=height,
                 width=width
             ).images[0]
-            
+
+            # Return base64-encoded image
             return json.dumps({"output": self.encode_base64_image(output_image)})
+
         except Exception as e:
             return json.dumps({"error": str(e)})
     
     def decode_base64_image(self, image_string):
-        """Decode base64 encoded image."""
-        base64_image = base64.b64decode(image_string)
-        buffer = io.BytesIO(base64_image)
-        return Image.open(buffer).convert("RGB")
-    
+        """Decode base64-encoded image to a PIL Image."""
+        try:
+            base64_image = base64.b64decode(image_string)
+            buffer = io.BytesIO(base64_image)
+            return Image.open(buffer).convert("RGB")
+        except Exception as e:
+            raise ValueError(f"Failed to decode base64 image: {e}")
+
     def encode_base64_image(self, image):
         """Encode PIL image to base64."""
         buffered = io.BytesIO()