Fix handler to return PIL Images instead of dictionaries for HF API compatibility

handler.py

@@ -81,29 +81,37 @@ class EndpointHandler:
             
             # Process based on edit type
             if edit_type == "replace" and len(prompts) >= 2:
-                result = self.word_replacement_edit(prompts[0], prompts[1], width, height, input_svg)
+                svg_content, metadata = self.word_replacement_edit(prompts[0], prompts[1], width, height, input_svg)
             elif edit_type == "refine":
-                result = self.prompt_refinement_edit(prompts[0], width, height, input_svg)
+                svg_content, metadata = self.prompt_refinement_edit(prompts[0], width, height, input_svg)
             elif edit_type == "reweight":
-                result = self.attention_reweighting_edit(prompts[0], width, height, input_svg)
+                svg_content, metadata = self.attention_reweighting_edit(prompts[0], width, height, input_svg)
             elif edit_type == "generate":
-                result = self.simple_generation(prompts[0], width, height)
+                svg_content, metadata = self.simple_generation(prompts[0], width, height)
             else:
                 # Default to refinement
-                result = self.prompt_refinement_edit(prompts[0], width, height, input_svg)
+                svg_content, metadata = self.prompt_refinement_edit(prompts[0], width, height, input_svg)
             
-            return result
+            # Convert SVG to PIL Image for HF API compatibility
+            pil_image = self.svg_to_pil_image(svg_content, width, height)
+            
+            # Store metadata
+            for key, value in metadata.items():
+                if isinstance(value, (dict, list)):
+                    pil_image.info[key] = json.dumps(value)
+                else:
+                    pil_image.info[key] = str(value)
+            
+            return pil_image
             
         except Exception as e:
             print(f"Error in handler: {e}")
-            # Return fallback result
+            # Return fallback image
             fallback_svg = self.create_fallback_svg(prompts[0] if prompts else "error", width, height)
-            return {
-                "svg": fallback_svg,
-                "svg_base64": base64.b64encode(fallback_svg.encode('utf-8')).decode('utf-8'),
-                "edit_type": edit_type,
-                "error": str(e)
-            }
+            fallback_image = self.svg_to_pil_image(fallback_svg, width, height)
+            fallback_image.info['error'] = str(e)
+            fallback_image.info['edit_type'] = edit_type
+            return fallback_image
 
     def word_replacement_edit(self, source_prompt: str, target_prompt: str, width: int, height: int, input_svg: str = None):
         """Perform word replacement editing"""
@@ -128,9 +136,7 @@ class EndpointHandler:
             # Apply word replacement transformations
             edited_svg = self.apply_word_replacement(base_svg, source_prompt, target_prompt, added_words, removed_words, width, height)
             
-            return {
-                "svg": edited_svg,
-                "svg_base64": base64.b64encode(edited_svg.encode('utf-8')).decode('utf-8'),
+            metadata = {
                 "edit_type": "replace",
                 "source_prompt": source_prompt,
                 "target_prompt": target_prompt,
@@ -138,9 +144,13 @@ class EndpointHandler:
                 "removed_words": list(removed_words)
             }
             
+            return edited_svg, metadata
+            
         except Exception as e:
             print(f"Error in word_replacement_edit: {e}")
-            return self.create_error_result(source_prompt, "replace", str(e), width, height)
+            fallback_svg = self.create_fallback_svg(source_prompt, width, height)
+            metadata = {"edit_type": "replace", "error": str(e)}
+            return fallback_svg, metadata
 
     def prompt_refinement_edit(self, prompt: str, width: int, height: int, input_svg: str = None):
         """Perform prompt refinement editing"""
@@ -156,16 +166,18 @@ class EndpointHandler:
             # Apply refinement based on prompt analysis
             refined_svg = self.apply_refinement(base_svg, prompt, width, height)
             
-            return {
-                "svg": refined_svg,
-                "svg_base64": base64.b64encode(refined_svg.encode('utf-8')).decode('utf-8'),
+            metadata = {
                 "edit_type": "refine",
                 "prompt": prompt
             }
             
+            return refined_svg, metadata
+            
         except Exception as e:
             print(f"Error in prompt_refinement_edit: {e}")
-            return self.create_error_result(prompt, "refine", str(e), width, height)
+            fallback_svg = self.create_fallback_svg(prompt, width, height)
+            metadata = {"edit_type": "refine", "error": str(e)}
+            return fallback_svg, metadata
 
     def attention_reweighting_edit(self, prompt: str, width: int, height: int, input_svg: str = None):
         """Perform attention reweighting editing"""
@@ -184,18 +196,20 @@ class EndpointHandler:
             # Apply attention reweighting
             reweighted_svg = self.apply_attention_reweighting(base_svg, weighted_prompt, attention_weights, width, height)
             
-            return {
-                "svg": reweighted_svg,
-                "svg_base64": base64.b64encode(reweighted_svg.encode('utf-8')).decode('utf-8'),
+            metadata = {
                 "edit_type": "reweight",
                 "prompt": prompt,
                 "weighted_prompt": weighted_prompt,
                 "attention_weights": attention_weights
             }
             
+            return reweighted_svg, metadata
+            
         except Exception as e:
             print(f"Error in attention_reweighting_edit: {e}")
-            return self.create_error_result(prompt, "reweight", str(e), width, height)
+            fallback_svg = self.create_fallback_svg(prompt, width, height)
+            metadata = {"edit_type": "reweight", "error": str(e)}
+            return fallback_svg, metadata
 
     def simple_generation(self, prompt: str, width: int, height: int):
         """Perform simple SVG generation"""
@@ -204,16 +218,18 @@ class EndpointHandler:
             
             svg_content = self.generate_base_svg(prompt, width, height)
             
-            return {
-                "svg": svg_content,
-                "svg_base64": base64.b64encode(svg_content.encode('utf-8')).decode('utf-8'),
+            metadata = {
                 "edit_type": "generate",
                 "prompt": prompt
             }
             
+            return svg_content, metadata
+            
         except Exception as e:
             print(f"Error in simple_generation: {e}")
-            return self.create_error_result(prompt, "generate", str(e), width, height)
+            fallback_svg = self.create_fallback_svg(prompt, width, height)
+            metadata = {"edit_type": "generate", "error": str(e)}
+            return fallback_svg, metadata
 
     def generate_base_svg(self, prompt: str, width: int, height: int):
         """Generate base SVG from prompt"""
@@ -727,6 +743,34 @@ class EndpointHandler:
             "error": error
         }
 
+    def svg_to_pil_image(self, svg_content, width, height):
+        """Convert SVG content to PIL Image"""
+        try:
+            import cairosvg
+            import io
+            
+            # Convert SVG to PNG bytes
+            png_bytes = cairosvg.svg2png(
+                bytestring=svg_content.encode('utf-8'),
+                output_width=width,
+                output_height=height
+            )
+            
+            # Convert to PIL Image
+            image = Image.open(io.BytesIO(png_bytes)).convert('RGB')
+            return image
+            
+        except ImportError:
+            print("cairosvg not available, creating simple image representation")
+            # Fallback: create a simple image with text
+            image = Image.new('RGB', (width, height), 'white')
+            return image
+        except Exception as e:
+            print(f"Error converting SVG to image: {e}")
+            # Fallback: create a simple image
+            image = Image.new('RGB', (width, height), 'white')
+            return image
+
     def create_fallback_svg(self, prompt: str, width: int, height: int):
         """Create simple fallback SVG"""
         dwg = svgwrite.Drawing(size=(width, height))