Fix serialization: Return PIL.Image directly instead of list

handler.py

@@ -1,8 +1,6 @@
 from typing import Dict, List, Any, Union
 import torch
 from PIL import Image, ImageDraw
-import io
-import base64
 import random
 import math
 
@@ -12,7 +10,7 @@ class EndpointHandler:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         print(f"DiffSketchEdit handler initialized on {self.device}")
     
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any]) -> Image.Image:
         """
         Process the request and return generated image
         
@@ -22,7 +20,7 @@ class EndpointHandler:
                 - parameters: Optional parameters (num_paths, num_iter, etc.)
         
         Returns:
-            List containing dictionary with base64 encoded image
+            PIL Image object
         """
         try:
             # Extract inputs
@@ -43,12 +41,7 @@ class EndpointHandler:
             # Generate sequential sketch edits
             image = self._generate_sequential_sketch(inputs, num_paths)
             
-            # Convert to base64
-            buffered = io.BytesIO()
-            image.save(buffered, format="PNG")
-            img_base64 = base64.b64encode(buffered.getvalue()).decode()
-            
-            return [{"generated_image": img_base64}]
+            return image
             
         except Exception as e:
             print(f"Error in DiffSketchEdit handler: {e}")
@@ -57,11 +50,7 @@ class EndpointHandler:
             draw = ImageDraw.Draw(error_img)
             draw.text((10, 100), f"Error: {str(e)[:30]}", fill='white')
             
-            buffered = io.BytesIO()
-            error_img.save(buffered, format="PNG")
-            img_base64 = base64.b64encode(buffered.getvalue()).decode()
-            
-            return [{"generated_image": img_base64}]
+            return error_img
     
     def _generate_sequential_sketch(self, prompts: List[str], num_paths: int) -> Image.Image:
         """Generate a sketch that shows sequential editing based on prompts"""