Upload handler.py with huggingface_hub

handler.py

@@ -1,48 +1,15 @@
 import base64
 import json
-import torch
-import svgwrite
 import math
 from typing import Dict, Any
-from diffusers import StableDiffusionPipeline
-import clip
 
 class EndpointHandler:
     def __init__(self, path=""):
         """Initialize the DiffSketcher model"""
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        
-        # Load Stable Diffusion model
-        try:
-            self.sd_pipeline = StableDiffusionPipeline.from_pretrained(
-                "runwayml/stable-diffusion-v1-5",
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                safety_checker=None,
-                requires_safety_checker=False
-            ).to(self.device)
-        except Exception as e:
-            print(f"Warning: Could not load Stable Diffusion: {e}")
-            self.sd_pipeline = None
-        
-        # Load CLIP model
-        try:
-            self.clip_model, self.clip_preprocess = clip.load("ViT-B/32", device=self.device)
-        except Exception as e:
-            print(f"Warning: Could not load CLIP: {e}")
-            self.clip_model = None
+        print("DiffSketcher handler initialized")
     
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Process the request and generate SVG
-        
-        Args:
-            data: Dictionary containing:
-                - inputs: The prompt text
-                - parameters: Optional parameters like num_paths, width, height
-        
-        Returns:
-            Dictionary with SVG content
-        """
+        """Generate SVG using DiffSketcher style"""
         try:
             # Extract inputs
             if isinstance(data, dict):
@@ -60,7 +27,7 @@ class EndpointHandler:
             width = parameters.get("width", 512)
             height = parameters.get("height", 512)
             
-            # Generate SVG using DiffSketcher style
+            # Generate SVG content
             svg_content = self.generate_diffsketcher_svg(prompt, num_paths, width, height)
             
             # Encode as base64
@@ -83,120 +50,76 @@ class EndpointHandler:
     
     def generate_diffsketcher_svg(self, prompt, num_paths, width, height):
         """Generate SVG in DiffSketcher style (painterly, sketchy)"""
-        dwg = svgwrite.Drawing(size=(f"{width}px", f"{height}px"))
-        
-        # White background
-        dwg.add(dwg.rect(insert=(0, 0), size=("100%", "100%"), fill="white"))
+        svg_parts = [
+            f'<svg baseProfile="full" height="{height}px" version="1.1" width="{width}px" xmlns="http://www.w3.org/2000/svg">',
+            f'<rect fill="white" height="100%" width="100%" x="0" y="0" />',
+        ]
         
-        # Generate paths based on prompt analysis
+        # Generate content based on prompt
         center_x, center_y = width // 2, height // 2
-        
-        # Analyze prompt for content type
         prompt_lower = prompt.lower()
         
         if any(word in prompt_lower for word in ["cat", "animal", "pet"]):
-            self._draw_cat_sketch(dwg, center_x, center_y, num_paths)
+            svg_parts.extend(self._draw_cat_sketch(center_x, center_y))
         elif any(word in prompt_lower for word in ["flower", "plant", "bloom"]):
-            self._draw_flower_sketch(dwg, center_x, center_y, num_paths)
+            svg_parts.extend(self._draw_flower_sketch(center_x, center_y))
         elif any(word in prompt_lower for word in ["house", "building", "home"]):
-            self._draw_house_sketch(dwg, center_x, center_y, num_paths)
-        elif any(word in prompt_lower for word in ["mountain", "landscape", "nature"]):
-            self._draw_landscape_sketch(dwg, center_x, center_y, num_paths)
+            svg_parts.extend(self._draw_house_sketch(center_x, center_y))
         else:
-            self._draw_abstract_sketch(dwg, center_x, center_y, num_paths)
+            svg_parts.extend(self._draw_abstract_sketch(center_x, center_y, num_paths))
         
-        # Add prompt as text
-        dwg.add(dwg.text(f"DiffSketcher: {prompt}", insert=(10, height-10), 
-                        fill="gray", font_size="12px"))
+        # Add prompt text
+        svg_parts.append(f'<text fill="gray" font-size="12px" x="10" y="{height-10}">DiffSketcher: {prompt}</text>')
+        svg_parts.append('</svg>')
         
-        return dwg.tostring()
+        return ''.join(svg_parts)
     
-    def _draw_cat_sketch(self, dwg, cx, cy, num_paths):
+    def _draw_cat_sketch(self, cx, cy):
         """Draw a sketchy cat"""
-        # Head
-        dwg.add(dwg.circle(center=(cx, cy-20), r=60, fill="none", stroke="black", stroke_width=3))
-        
-        # Ears
-        dwg.add(dwg.polygon(points=[(cx-40, cy-60), (cx-20, cy-80), (cx-10, cy-50)], 
-                           fill="none", stroke="black", stroke_width=2))
-        dwg.add(dwg.polygon(points=[(cx+40, cy-60), (cx+20, cy-80), (cx+10, cy-50)], 
-                           fill="none", stroke="black", stroke_width=2))
-        
-        # Eyes
-        dwg.add(dwg.circle(center=(cx-20, cy-10), r=8, fill="black"))
-        dwg.add(dwg.circle(center=(cx+20, cy-10), r=8, fill="black"))
-        
-        # Nose
-        dwg.add(dwg.polygon(points=[(cx-5, cy+10), (cx+5, cy+10), (cx, cy+20)], fill="pink"))
-        
-        # Whiskers
-        dwg.add(dwg.line(start=(cx-50, cy), end=(cx-70, cy-5), stroke="black", stroke_width=1))
-        dwg.add(dwg.line(start=(cx+50, cy), end=(cx+70, cy-5), stroke="black", stroke_width=1))
-        
-        # Body
-        dwg.add(dwg.ellipse(center=(cx, cy+80), r=(40, 60), fill="none", stroke="black", stroke_width=3))
+        return [
+            f'<circle cx="{cx}" cy="{cy-20}" r="60" fill="none" stroke="black" stroke-width="3" />',
+            f'<polygon points="{cx-40},{cy-60} {cx-20},{cy-80} {cx-10},{cy-50}" fill="none" stroke="black" stroke-width="2" />',
+            f'<polygon points="{cx+40},{cy-60} {cx+20},{cy-80} {cx+10},{cy-50}" fill="none" stroke="black" stroke-width="2" />',
+            f'<circle cx="{cx-20}" cy="{cy-10}" r="8" fill="black" />',
+            f'<circle cx="{cx+20}" cy="{cy-10}" r="8" fill="black" />',
+            f'<polygon points="{cx-5},{cy+10} {cx+5},{cy+10} {cx},{cy+20}" fill="pink" />',
+            f'<line x1="{cx-50}" y1="{cy}" x2="{cx-70}" y2="{cy-5}" stroke="black" stroke-width="1" />',
+            f'<line x1="{cx+50}" y1="{cy}" x2="{cx+70}" y2="{cy-5}" stroke="black" stroke-width="1" />',
+            f'<ellipse cx="{cx}" cy="{cy+80}" rx="40" ry="60" fill="none" stroke="black" stroke-width="3" />',
+        ]
     
-    def _draw_flower_sketch(self, dwg, cx, cy, num_paths):
+    def _draw_flower_sketch(self, cx, cy):
         """Draw a sketchy flower"""
-        # Petals in a circle
+        petals = []
         for i in range(8):
             angle = i * 45
             petal_x = cx + 50 * math.cos(math.radians(angle))
             petal_y = cy + 50 * math.sin(math.radians(angle))
-            
-            dwg.add(dwg.ellipse(
-                center=(petal_x, petal_y), 
-                r=(20, 35),
-                fill="pink", 
-                stroke="red", 
-                stroke_width=2,
-                transform=f"rotate({angle} {petal_x} {petal_y})"
-            ))
-        
-        # Center
-        dwg.add(dwg.circle(center=(cx, cy), r=15, fill="yellow", stroke="orange", stroke_width=2))
-        
-        # Stem
-        dwg.add(dwg.line(start=(cx, cy+15), end=(cx, cy+120), stroke="green", stroke_width=4))
-        
-        # Leaves
-        dwg.add(dwg.ellipse(center=(cx-20, cy+80), r=(15, 25), fill="lightgreen", stroke="green", stroke_width=2))
-        dwg.add(dwg.ellipse(center=(cx+20, cy+90), r=(15, 25), fill="lightgreen", stroke="green", stroke_width=2))
+            petals.append(f'<ellipse cx="{petal_x}" cy="{petal_y}" rx="20" ry="35" fill="pink" stroke="red" stroke-width="2" transform="rotate({angle} {petal_x} {petal_y})" />')
+        
+        return petals + [
+            f'<circle cx="{cx}" cy="{cy}" r="15" fill="yellow" stroke="orange" stroke-width="2" />',
+            f'<line x1="{cx}" y1="{cy+15}" x2="{cx}" y2="{cy+120}" stroke="green" stroke-width="4" />',
+            f'<ellipse cx="{cx-20}" cy="{cy+80}" rx="15" ry="25" fill="lightgreen" stroke="green" stroke-width="2" />',
+            f'<ellipse cx="{cx+20}" cy="{cy+90}" rx="15" ry="25" fill="lightgreen" stroke="green" stroke-width="2" />',
+        ]
     
-    def _draw_house_sketch(self, dwg, cx, cy, num_paths):
+    def _draw_house_sketch(self, cx, cy):
         """Draw a sketchy house"""
-        # House base
-        dwg.add(dwg.rect(insert=(cx-50, cy), size=(100, 60), fill="lightblue", stroke="blue", stroke_width=3))
-        
-        # Roof
-        dwg.add(dwg.polygon(points=[(cx-60, cy), (cx, cy-50), (cx+60, cy)], fill="red", stroke="darkred", stroke_width=2))
-        
-        # Door
-        dwg.add(dwg.rect(insert=(cx-15, cy+20), size=(30, 40), fill="brown"))
-        
-        # Windows
-        dwg.add(dwg.rect(insert=(cx-40, cy+15), size=(20, 20), fill="lightblue", stroke="blue", stroke_width=2))
-        dwg.add(dwg.rect(insert=(cx+20, cy+15), size=(20, 20), fill="lightblue", stroke="blue", stroke_width=2))
+        return [
+            f'<rect x="{cx-50}" y="{cy}" width="100" height="60" fill="lightblue" stroke="blue" stroke-width="3" />',
+            f'<polygon points="{cx-60},{cy} {cx},{cy-50} {cx+60},{cy}" fill="red" stroke="darkred" stroke-width="2" />',
+            f'<rect x="{cx-15}" y="{cy+20}" width="30" height="40" fill="brown" />',
+            f'<rect x="{cx-40}" y="{cy+15}" width="20" height="20" fill="lightblue" stroke="blue" stroke-width="2" />',
+            f'<rect x="{cx+20}" y="{cy+15}" width="20" height="20" fill="lightblue" stroke="blue" stroke-width="2" />',
+        ]
     
-    def _draw_landscape_sketch(self, dwg, cx, cy, num_paths):
-        """Draw a sketchy landscape"""
-        # Mountains
-        dwg.add(dwg.polygon(points=[(0, cy), (cx-100, cy-100), (cx-50, cy-80), (cx, cy-120), 
-                                   (cx+50, cy-90), (cx+100, cy-110), (cx+200, cy)], 
-                           fill="lightgray", stroke="gray", stroke_width=2))
-        
-        # Trees
-        for i, tree_x in enumerate([cx-80, cx-20, cx+40, cx+100]):
-            tree_y = cy + 20
-            # Trunk
-            dwg.add(dwg.rect(insert=(tree_x-5, tree_y), size=(10, 40), fill="brown"))
-            # Foliage
-            dwg.add(dwg.circle(center=(tree_x, tree_y-10), r=25, fill="green", stroke="darkgreen", stroke_width=2))
-    
-    def _draw_abstract_sketch(self, dwg, cx, cy, num_paths):
+    def _draw_abstract_sketch(self, cx, cy, num_paths):
         """Draw abstract sketchy shapes"""
         import random
+        random.seed(42)  # For consistent results
         
+        shapes = []
         colors = ["red", "blue", "green", "orange", "purple", "pink", "yellow"]
         
         for i in range(min(num_paths, 12)):
@@ -206,9 +129,11 @@ class EndpointHandler:
             color = random.choice(colors)
             
             if i % 3 == 0:
-                dwg.add(dwg.circle(center=(x, y), r=r, fill="none", stroke=color, stroke_width=3))
+                shapes.append(f'<circle cx="{x}" cy="{y}" r="{r}" fill="none" stroke="{color}" stroke-width="3" />')
             elif i % 3 == 1:
-                dwg.add(dwg.rect(insert=(x-r//2, y-r//2), size=(r, r), fill="none", stroke=color, stroke_width=2))
+                shapes.append(f'<rect x="{x-r//2}" y="{y-r//2}" width="{r}" height="{r}" fill="none" stroke="{color}" stroke-width="2" />')
             else:
-                points = [(x, y-r), (x+r, y+r), (x-r, y+r)]
-                dwg.add(dwg.polygon(points=points, fill="none", stroke=color, stroke_width=2))
+                points = f"{x},{y-r} {x+r},{y+r} {x-r},{y+r}"
+                shapes.append(f'<polygon points="{points}" fill="none" stroke="{color}" stroke-width="2" />')
+        
+        return shapes