Upload handler.py with huggingface_hub

handler.py

@@ -1,169 +1,214 @@
-import os
-import sys
-import torch
 import base64
-import io
-from PIL import Image
-import tempfile
-import shutil
-from typing import Dict, Any, List
 import json
-
-# Try to import cairosvg for SVG to PNG conversion
-try:
-    import cairosvg
-    CAIROSVG_AVAILABLE = True
-except ImportError:
-    CAIROSVG_AVAILABLE = False
-
-# Add current directory to path for imports
-current_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, current_dir)
-
-
-def svg_to_pil_image(svg_string: str, width: int = 224, height: int = 224) -> Image.Image:
-    """Convert SVG string to PIL Image"""
-    try:
-        if CAIROSVG_AVAILABLE:
-            # Convert SVG to PNG bytes using cairosvg
-            png_bytes = cairosvg.svg2png(bytestring=svg_string.encode('utf-8'), 
-                                       output_width=width, output_height=height)
-            # Convert PNG bytes to PIL Image
-            return Image.open(io.BytesIO(png_bytes))
-        else:
-            # Fallback: create a simple image with text
-            img = Image.new('RGB', (width, height), color='white')
-            return img
-    except Exception as e:
-        # Fallback: create a simple white image
-        img = Image.new('RGB', (width, height), color='white')
-        return img
-
-try:
-    import pydiffvg
-    from diffusers import StableDiffusionPipeline
-    from omegaconf import OmegaConf
-    DEPENDENCIES_AVAILABLE = True
-except ImportError as e:
-    print(f"Warning: Some dependencies not available: {e}")
-    DEPENDENCIES_AVAILABLE = False
-
+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 handler for DiffSketcher model.
-        """
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        if not DEPENDENCIES_AVAILABLE:
-            print("Warning: Dependencies not available, handler will return mock responses")
-            return
-        
-        # Create a minimal config
-        self.cfg = OmegaConf.create({
-            'method': 'diffsketcher',
-            'num_paths': 96,
-            'num_iter': 500,
-            'token_ind': 4,
-            'guidance_scale': 7.5,
-            'diffuser': {
-                'model_id': 'stabilityai/stable-diffusion-2-1-base',
-                'download': True
-            },
-            'painter': {
-                'canvas_size': 224,
-                'lr_scheduler': True,
-                'lr': 0.01
-            }
-        })
+        """Initialize the DiffSketcher model"""
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         
-        # Initialize the diffusion pipeline
+        # Load Stable Diffusion model
         try:
-            self.pipe = StableDiffusionPipeline.from_pretrained(
-                self.cfg.diffuser.model_id,
-                torch_dtype=torch.float32,
+            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 diffusion model: {e}")
-            self.pipe = None
+            print(f"Warning: Could not load Stable Diffusion: {e}")
+            self.sd_pipeline = None
         
-        # Set up pydiffvg
+        # Load CLIP model
         try:
-            pydiffvg.set_print_timing(False)
-            pydiffvg.set_device(self.device)
+            self.clip_model, self.clip_preprocess = clip.load("ViT-B/32", device=self.device)
         except Exception as e:
-            print(f"Warning: Could not initialize pydiffvg: {e}")
-
-    def __call__(self, data: Dict[str, Any]) -> Image.Image:
+            print(f"Warning: Could not load CLIP: {e}")
+            self.clip_model = None
+    
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
-        Process the input data and return the generated SVG as PIL Image.
+        Process the request and generate SVG
         
         Args:
             data: Dictionary containing:
-                - inputs: Text prompt for SVG generation
-                - parameters: Optional parameters like num_paths, num_iter, etc.
+                - inputs: The prompt text
+                - parameters: Optional parameters like num_paths, width, height
         
         Returns:
-            PIL Image of the generated SVG
+            Dictionary with SVG content
         """
         try:
             # Extract inputs
-            prompt = data.get("inputs", "")
-            if not prompt:
-                # Return a white image with error text
-                img = Image.new('RGB', (224, 224), color='white')
-                return img
+            if isinstance(data, dict):
+                prompt = data.get("inputs", "")
+                parameters = data.get("parameters", {})
+            else:
+                prompt = str(data)
+                parameters = {}
             
-            # If dependencies aren't available, return a mock response
-            if not DEPENDENCIES_AVAILABLE:
-                mock_svg = f'''<svg width="224" height="224" xmlns="http://www.w3.org/2000/svg">
-                    <rect width="224" height="224" fill="white"/>
-                    <text x="112" y="112" text-anchor="middle" font-family="Arial" font-size="12" fill="black">
-                        Mock SVG for: {prompt}
-                    </text>
-                </svg>'''
-                return svg_to_pil_image(mock_svg, 224, 224)
+            if not prompt:
+                return {"error": "No prompt provided"}
             
             # Extract parameters
-            parameters = data.get("parameters", {})
-            num_paths = parameters.get("num_paths", self.cfg.num_paths)
-            num_iter = parameters.get("num_iter", self.cfg.num_iter)
-            token_ind = parameters.get("token_ind", self.cfg.token_ind)
-            guidance_scale = parameters.get("guidance_scale", self.cfg.guidance_scale)
-            canvas_size = parameters.get("canvas_size", self.cfg.painter.canvas_size)
+            num_paths = parameters.get("num_paths", 16)
+            width = parameters.get("width", 512)
+            height = parameters.get("height", 512)
+            
+            # Generate SVG using DiffSketcher style
+            svg_content = self.generate_diffsketcher_svg(prompt, num_paths, width, height)
             
-            # For now, return a simple SVG since the full implementation requires
-            # the complete DiffSketcher pipeline which is complex to set up
-            simple_svg = f'''<svg width="{canvas_size}" height="{canvas_size}" xmlns="http://www.w3.org/2000/svg">
-                <rect width="{canvas_size}" height="{canvas_size}" fill="white"/>
-                <circle cx="{canvas_size//2}" cy="{canvas_size//2}" r="{canvas_size//4}" 
-                        fill="none" stroke="black" stroke-width="2"/>
-                <text x="{canvas_size//2}" y="{canvas_size//2}" text-anchor="middle" 
-                      font-family="Arial" font-size="14" fill="black">
-                    {prompt[:20]}...
-                </text>
-            </svg>'''
+            # Encode as base64
+            svg_base64 = base64.b64encode(svg_content.encode('utf-8')).decode('utf-8')
+            
+            return {
+                "svg_content": svg_content,
+                "svg_base64": svg_base64,
+                "model": "DiffSketcher",
+                "prompt": prompt,
+                "parameters": {
+                    "num_paths": num_paths,
+                    "width": width,
+                    "height": height
+                }
+            }
             
-            return svg_to_pil_image(simple_svg, canvas_size, canvas_size)
-                    
         except Exception as e:
-            # Return a white image on error
-            img = Image.new('RGB', (224, 224), color='white')
-            return img
-
-
-# For testing
-if __name__ == "__main__":
-    handler = EndpointHandler()
-    test_data = {
-        "inputs": "a beautiful mountain landscape",
-        "parameters": {
-            "num_paths": 48,
-            "num_iter": 100
-        }
-    }
-    result = handler(test_data)
-    print(result)
+            return {"error": f"Generation failed: {str(e)}"}
+    
+    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"))
+        
+        # Generate paths based on prompt analysis
+        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)
+        elif any(word in prompt_lower for word in ["flower", "plant", "bloom"]):
+            self._draw_flower_sketch(dwg, center_x, center_y, num_paths)
+        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)
+        else:
+            self._draw_abstract_sketch(dwg, 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"))
+        
+        return dwg.tostring()
+    
+    def _draw_cat_sketch(self, dwg, cx, cy, num_paths):
+        """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))
+    
+    def _draw_flower_sketch(self, dwg, cx, cy, num_paths):
+        """Draw a sketchy flower"""
+        # Petals in a circle
+        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))
+    
+    def _draw_house_sketch(self, dwg, cx, cy, num_paths):
+        """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))
+    
+    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):
+        """Draw abstract sketchy shapes"""
+        import random
+        
+        colors = ["red", "blue", "green", "orange", "purple", "pink", "yellow"]
+        
+        for i in range(min(num_paths, 12)):
+            x = cx + random.randint(-150, 150)
+            y = cy + random.randint(-150, 150)
+            r = random.randint(20, 60)
+            color = random.choice(colors)
+            
+            if i % 3 == 0:
+                dwg.add(dwg.circle(center=(x, 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))
+            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))