Add: diffsketcher handler.py with original implementation

handler.py

@@ -1,16 +1,10 @@
 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
+import traceback
 
 # Add debug logging
 def debug_log(message):
@@ -25,10 +19,24 @@ try:
     debug_log("Successfully imported cairosvg")
 except ImportError:
     debug_log("cairosvg not found. Installing...")
+    import subprocess
     subprocess.check_call(["pip", "install", "cairosvg", "cairocffi", "cssselect2", "defusedxml", "tinycss2"])
     import cairosvg
     debug_log("Installed and imported cairosvg")
 
+# Add the model directory to the path
+sys.path.append('/code/diffsketcher')
+
+# Try to import the model
+try:
+    from models.clip_model import ClipModel
+    from models.diffusion_model import DiffusionModel
+    from models.sketch_model import SketchModel
+    debug_log("Successfully imported DiffSketcher models")
+except ImportError as e:
+    debug_log(f"Error importing DiffSketcher models: {e}")
+    debug_log(traceback.format_exc())
+
 class EndpointHandler:
     def __init__(self, model_dir):
         """Initialize the handler with model directory"""
@@ -38,26 +46,29 @@ class EndpointHandler:
             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}")
+            # Initialize the model
+            try:
+                self.clip_model = ClipModel(device=self.device)
+                self.diffusion_model = DiffusionModel(device=self.device)
+                self.sketch_model = SketchModel(device=self.device)
                 
-                # 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")
+                # Load checkpoint if available
+                weights_path = os.path.join(model_dir, "checkpoint.pth")
+                if os.path.exists(weights_path):
+                    debug_log(f"Loading checkpoint from {weights_path}")
+                    checkpoint = torch.load(weights_path, map_location=self.device)
+                    self.sketch_model.load_state_dict(checkpoint['sketch_model'])
+                    debug_log("Successfully loaded checkpoint")
+                    self.use_model = True
+                else:
+                    debug_log(f"Checkpoint not found at {weights_path}, using model without pre-trained weights")
+                    self.use_model = True
+            except Exception as e:
+                debug_log(f"Error initializing model: {e}")
+                debug_log(traceback.format_exc())
+                self.use_model = False
         except Exception as e:
             debug_log(f"Error in handler initialization: {e}")
-            import traceback
             debug_log(traceback.format_exc())
             self.use_model = False
     
@@ -65,6 +76,28 @@ class EndpointHandler:
         """Generate an SVG from a text prompt"""
         debug_log(f"Generating SVG for prompt: {prompt}")
         
+        if self.use_model:
+            try:
+                debug_log("Using initialized model")
+                
+                # Generate SVG using DiffSketcher
+                text_features = self.clip_model.encode_text(prompt)
+                latent = self.diffusion_model.generate(text_features)
+                svg_data = self.sketch_model.generate(latent, num_paths=20, width=width, height=height)
+                debug_log("Generated SVG using DiffSketcher")
+                return svg_data
+            except Exception as e:
+                debug_log(f"Error generating SVG with model: {e}")
+                debug_log(traceback.format_exc())
+                return self._generate_placeholder_svg(prompt, width, height)
+        else:
+            debug_log("Model not initialized, using placeholder")
+            return self._generate_placeholder_svg(prompt, width, height)
+    
+    def _generate_placeholder_svg(self, prompt, width=512, height=512):
+        """Generate a placeholder SVG"""
+        debug_log(f"Generating placeholder 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="#ffffff"/>
@@ -78,7 +111,7 @@ class EndpointHandler:
             <text x="50%" y="50%" font-family="Arial" font-size="20" text-anchor="middle" fill="#333333">{prompt}</text>
         </svg>"""
         
-        debug_log("Generated SVG content")
+        debug_log("Generated placeholder SVG")
         return svg_content
     
     def __call__(self, data):
@@ -109,7 +142,6 @@ class EndpointHandler:
                 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")
@@ -123,7 +155,6 @@ class EndpointHandler:
             return image
         except Exception as 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")
@@ -131,4 +162,4 @@ class EndpointHandler:
             draw = ImageDraw.Draw(image)
             draw.text((256, 256), f"Error: {str(e)}", fill="white", anchor="mm")
             debug_log("Returning error image")
-            return image
+            return image