Upload handler.py with huggingface_hub

handler.py

@@ -9,10 +9,36 @@ import shutil
 from typing import Dict, Any, List
 import json
 
+# Try to import cairosvg for SVG to PNG conversion
+try:
+    import cairosvg
+    CAIROSVG_AVAILABLE = True
+except ImportError:
+    CAIROSVG_AVAILABLE = False
+
 # Add current directory to path for imports
 current_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.insert(0, current_dir)
 
+
+def svg_to_pil_image(svg_string: str, width: int = 224, height: int = 224) -> Image.Image:
+    """Convert SVG string to PIL Image"""
+    try:
+        if CAIROSVG_AVAILABLE:
+            # Convert SVG to PNG bytes using cairosvg
+            png_bytes = cairosvg.svg2png(bytestring=svg_string.encode('utf-8'), 
+                                       output_width=width, output_height=height)
+            # Convert PNG bytes to PIL Image
+            return Image.open(io.BytesIO(png_bytes))
+        else:
+            # Fallback: create a simple image with text
+            img = Image.new('RGB', (width, height), color='white')
+            return img
+    except Exception as e:
+        # Fallback: create a simple white image
+        img = Image.new('RGB', (width, height), color='white')
+        return img
+
 try:
     import pydiffvg
     from diffusers import StableDiffusionPipeline
@@ -71,9 +97,9 @@ class EndpointHandler:
         except Exception as e:
             print(f"Warning: Could not initialize pydiffvg: {e}")
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any]) -> Image.Image:
         """
-        Process the input data and return the edited SVG.
+        Process the input data and return the edited SVG as PIL Image.
         
         Args:
             data: Dictionary containing:
@@ -81,13 +107,15 @@ class EndpointHandler:
                 - parameters: Optional parameters including input_svg, edit_instruction, etc.
         
         Returns:
-            List containing the edited SVG as base64 encoded string
+            PIL Image of the edited SVG
         """
         try:
             # Extract inputs
             prompt = data.get("inputs", "")
             if not prompt:
-                return [{"error": "No prompt provided"}]
+                # Return a white image with error text
+                img = Image.new('RGB', (256, 256), color='white')
+                return img
             
             # If dependencies aren't available, return a mock response
             if not DEPENDENCIES_AVAILABLE:
@@ -97,13 +125,7 @@ class EndpointHandler:
                         Mock DiffSketchEdit for: {prompt}
                     </text>
                 </svg>'''
-                return [{
-                    "svg": mock_svg,
-                    "svg_base64": base64.b64encode(mock_svg.encode()).decode(),
-                    "prompt": prompt,
-                    "status": "mock_response",
-                    "message": "This is a mock response. Full model not available."
-                }]
+                return svg_to_pil_image(mock_svg, 256, 256)
             
             # Extract parameters
             parameters = data.get("parameters", {})
@@ -154,24 +176,12 @@ class EndpointHandler:
                     </text>
                 </svg>'''
             
-            return [{
-                "svg": edited_svg,
-                "svg_base64": base64.b64encode(edited_svg.encode()).decode(),
-                "prompt": prompt,
-                "edit_instruction": edit_instruction,
-                "parameters": {
-                    "num_paths": num_paths,
-                    "num_iter": num_iter,
-                    "guidance_scale": guidance_scale,
-                    "edit_strength": edit_strength,
-                    "canvas_size": canvas_size
-                },
-                "status": "simplified_response",
-                "message": "Simplified SVG edit generated. Full DiffSketchEdit pipeline requires additional setup."
-            }]
+            return svg_to_pil_image(edited_svg, canvas_size, canvas_size)
                     
         except Exception as e:
-            return [{"error": f"Error during SVG editing: {str(e)}"}]
+            # Return a white image on error
+            img = Image.new('RGB', (256, 256), color='white')
+            return img
 
 
 # For testing