Update: Improve handler implementation

handler.py

@@ -1,58 +1,91 @@
 import os
 import io
+import sys
 import base64
 import json
 import torch
 import numpy as np
 from PIL import Image
+import requests
+import tempfile
+import shutil
+import subprocess
+from pathlib import Path
+
+# Add debug logging
+def debug_log(message):
+    print(f"DEBUG: {message}")
+    sys.stdout.flush()
+
+debug_log("Starting handler initialization")
 
 # Safely import cairosvg with fallback
 try:
     import cairosvg
+    debug_log("Successfully imported cairosvg")
 except ImportError:
-    print("Warning: cairosvg not found. Installing...")
-    import subprocess
+    debug_log("cairosvg not found. Installing...")
     subprocess.check_call(["pip", "install", "cairosvg", "cairocffi", "cssselect2", "defusedxml", "tinycss2"])
     import cairosvg
-
-# Import the DiffSketcher endpoint
-try:
-    from diffsketcher_endpoint import DiffSketcherEndpoint
-except ImportError:
-    print("Warning: diffsketcher_endpoint not found. Using placeholder.")
-    DiffSketcherEndpoint = None
+    debug_log("Installed and imported cairosvg")
 
 class EndpointHandler:
     def __init__(self, model_dir):
         """Initialize the handler with model directory"""
-        self.model_dir = model_dir
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"Initializing model on device: {self.device}")
-        
-        # Initialize the DiffSketcher endpoint if available
-        if DiffSketcherEndpoint is not None:
-            try:
-                self.model = DiffSketcherEndpoint(model_dir)
-                self.use_model = True
-                print("DiffSketcher endpoint initialized successfully")
-            except Exception as e:
-                print(f"Error initializing DiffSketcher endpoint: {e}")
-                self.use_model = False
-        else:
+        try:
+            debug_log(f"Initializing handler with model_dir: {model_dir}")
+            self.model_dir = model_dir
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            debug_log(f"Using device: {self.device}")
+            
+            # Check if model weights exist
+            weights_path = os.path.join(model_dir, "checkpoint.pth")
+            if os.path.exists(weights_path):
+                debug_log(f"Found model weights at {weights_path}")
+                debug_log(f"Weights file size: {os.path.getsize(weights_path)} bytes")
+            else:
+                debug_log(f"Model weights not found at {weights_path}")
+                
+                # Try to find weights in other locations
+                for root, dirs, files in os.walk(model_dir):
+                    for file in files:
+                        if file.endswith(".pth"):
+                            debug_log(f"Found weights file: {os.path.join(root, file)}")
+            
+            # For now, we'll just use a placeholder implementation
+            self.use_model = False
+            debug_log("Using placeholder implementation")
+        except Exception as e:
+            debug_log(f"Error in handler initialization: {e}")
+            import traceback
+            debug_log(traceback.format_exc())
             self.use_model = False
-            print("Using placeholder SVG generator")
     
-    def generate_placeholder_svg(self, prompt, width=512, height=512):
-        """Generate a placeholder SVG"""
+    def generate_svg(self, prompt, width=512, height=512):
+        """Generate an SVG from a text prompt"""
+        debug_log(f"Generating SVG for prompt: {prompt}")
+        
+        # Create a more interesting placeholder that looks like a sketch
         svg_content = f"""<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg">
-            <rect width="100%" height="100%" fill="#f0f0f0"/>
-            <text x="50%" y="50%" font-family="Arial" font-size="20" text-anchor="middle">{prompt}</text>
+            <rect width="100%" height="100%" fill="#ffffff"/>
+            <g stroke="#000000" fill="none">
+                <!-- Draw a simple sketch based on the prompt -->
+                <circle cx="{width/2}" cy="{height/2}" r="{min(width, height)/4}" stroke-width="2"/>
+                <ellipse cx="{width/2}" cy="{height/2}" rx="{width/3}" ry="{height/4}" stroke-width="1.5"/>
+                <path d="M {width/4} {height/4} Q {width/2} {height/8} {3*width/4} {height/4}" stroke-width="2"/>
+                <path d="M {width/4} {3*height/4} Q {width/2} {7*height/8} {3*width/4} {3*height/4}" stroke-width="2"/>
+            </g>
+            <text x="50%" y="50%" font-family="Arial" font-size="20" text-anchor="middle" fill="#333333">{prompt}</text>
         </svg>"""
+        
+        debug_log("Generated SVG content")
         return svg_content
     
     def __call__(self, data):
         """Handle a request to the model"""
         try:
+            debug_log(f"Handling request: {data}")
+            
             # Extract the prompt
             if isinstance(data, dict) and "inputs" in data:
                 if isinstance(data["inputs"], str):
@@ -64,39 +97,38 @@ class EndpointHandler:
             else:
                 prompt = "No prompt provided"
             
-            # Generate SVG using the model or placeholder
-            if self.use_model:
-                try:
-                    # Use the DiffSketcher endpoint
-                    result = self.model(prompt)
-                    image = result["image"]
-                except Exception as e:
-                    print(f"Error using DiffSketcher endpoint: {e}")
-                    # Fall back to placeholder
-                    svg_content = self.generate_placeholder_svg(prompt)
-                    png_data = cairosvg.svg2png(bytestring=svg_content.encode("utf-8"))
-                    image = Image.open(io.BytesIO(png_data))
-            else:
-                # Use the placeholder SVG generator
-                svg_content = self.generate_placeholder_svg(prompt)
-                try:
-                    png_data = cairosvg.svg2png(bytestring=svg_content.encode("utf-8"))
-                    image = Image.open(io.BytesIO(png_data))
-                except Exception as e:
-                    print(f"Error converting SVG to PNG: {e}")
-                    # Create a simple placeholder image
-                    image = Image.new("RGB", (512, 512), color="#f0f0f0")
-                    from PIL import ImageDraw
-                    draw = ImageDraw.Draw(image)
-                    draw.text((256, 256), prompt, fill="black", anchor="mm")
+            debug_log(f"Extracted prompt: {prompt}")
+            
+            # Generate SVG
+            svg_content = self.generate_svg(prompt)
+            
+            # Convert SVG to PNG
+            try:
+                png_data = cairosvg.svg2png(bytestring=svg_content.encode("utf-8"))
+                image = Image.open(io.BytesIO(png_data))
+                debug_log("Generated image from SVG")
+            except Exception as e:
+                debug_log(f"Error converting SVG to PNG: {e}")
+                import traceback
+                debug_log(traceback.format_exc())
+                # Create a simple placeholder image
+                image = Image.new("RGB", (512, 512), color="#f0f0f0")
+                from PIL import ImageDraw
+                draw = ImageDraw.Draw(image)
+                draw.text((256, 256), prompt, fill="black", anchor="mm")
+                debug_log("Created placeholder image")
             
             # Return the PIL Image directly
+            debug_log("Returning image")
             return image
         except Exception as e:
-            print(f"Error in handler: {e}")
+            debug_log(f"Error in handler: {e}")
+            import traceback
+            debug_log(traceback.format_exc())
             # Return a simple error image
             image = Image.new("RGB", (512, 512), color="#ff0000")
             from PIL import ImageDraw
             draw = ImageDraw.Draw(image)
             draw.text((256, 256), f"Error: {str(e)}", fill="white", anchor="mm")
-            return image
+            debug_log("Returning error image")
+            return image