Add: diffsketcher handler.py with original implementation

handler.py

@@ -36,83 +36,58 @@ try:
 except ImportError as e:
     debug_log(f"Error importing DiffSketcher models: {e}")
     debug_log(traceback.format_exc())
+    raise ImportError(f"Failed to import DiffSketcher models: {e}")
 
 class EndpointHandler:
     def __init__(self, model_dir):
         """Initialize the handler with model directory"""
-        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}")
-            
-            # Initialize the model
+        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}")
+        
+        # Initialize the model
+        self.clip_model = ClipModel(device=self.device)
+        self.diffusion_model = DiffusionModel(device=self.device)
+        self.sketch_model = SketchModel(device=self.device)
+        
+        # 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")
+        else:
+            debug_log(f"Checkpoint not found at {weights_path}, using model without pre-trained weights")
+            # Download the checkpoint if not available
             try:
-                self.clip_model = ClipModel(device=self.device)
-                self.diffusion_model = DiffusionModel(device=self.device)
-                self.sketch_model = SketchModel(device=self.device)
-                
-                # 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
+                debug_log("Attempting to download checkpoint...")
+                import urllib.request
+                os.makedirs(os.path.dirname(weights_path), exist_ok=True)
+                urllib.request.urlretrieve(
+                    "https://github.com/ximinng/DiffSketcher/releases/download/v0.1-weights/diffvg_checkpoint.pth",
+                    weights_path
+                )
+                debug_log(f"Downloaded checkpoint to {weights_path}")
+                checkpoint = torch.load(weights_path, map_location=self.device)
+                self.sketch_model.load_state_dict(checkpoint['sketch_model'])
+                debug_log("Successfully loaded downloaded checkpoint")
             except Exception as e:
-                debug_log(f"Error initializing model: {e}")
+                debug_log(f"Error downloading checkpoint: {e}")
                 debug_log(traceback.format_exc())
-                self.use_model = False
-        except Exception as e:
-            debug_log(f"Error in handler initialization: {e}")
-            debug_log(traceback.format_exc())
-            self.use_model = False
+                debug_log("Continuing with uninitialized weights")
     
     def generate_svg(self, prompt, width=512, height=512):
         """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"/>
-            <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 placeholder SVG")
-        return svg_content
+        # 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
     
     def __call__(self, data):
         """Handle a request to the model"""
@@ -136,19 +111,9 @@ class EndpointHandler:
             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}")
-                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")
+            png_data = cairosvg.svg2png(bytestring=svg_content.encode("utf-8"))
+            image = Image.open(io.BytesIO(png_data))
+            debug_log("Generated image from SVG")
             
             # Return the PIL Image directly
             debug_log("Returning image")
@@ -156,10 +121,4 @@ class EndpointHandler:
         except Exception as e:
             debug_log(f"Error in handler: {e}")
             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")
-            debug_log("Returning error image")
-            return image
+            raise Exception(f"Error generating image: {str(e)}")