Update handler with functional sketch generation

handler.py

@@ -1,227 +1,171 @@
-from typing import Dict, List, Any
-import torch
 from PIL import Image, ImageDraw
-import random
 import math
+import random
 
 class EndpointHandler:
     def __init__(self, path=""):
-        """Initialize the handler with minimal dependencies"""
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"DiffSketcher handler initialized on {self.device}")
-    
-    def __call__(self, data: Dict[str, Any]) -> Image.Image:
-        """
-        Process the request and return generated image
+        """Initialize DiffSketcher handler for Hugging Face Inference API"""
+        pass
+
+    def __call__(self, data):
+        """Generate sketch from text prompt"""
+        # Extract prompt
+        inputs = data.get("inputs", "")
+        if isinstance(inputs, dict):
+            prompt = inputs.get("prompt", inputs.get("text", ""))
+        else:
+            prompt = str(inputs)
         
-        Args:
-            data: Dictionary containing:
-                - inputs: Text prompt for sketch generation
-                - parameters: Optional parameters (num_paths, num_iter, etc.)
+        if not prompt:
+            prompt = "a simple sketch"
         
-        Returns:
-            PIL Image object
-        """
-        try:
-            # Extract inputs
-            inputs = data.get("inputs", "")
-            if isinstance(inputs, list):
-                inputs = inputs[0] if inputs else ""
-            
-            # Extract parameters
-            parameters = data.get("parameters", {})
-            num_paths = parameters.get("num_paths", 64)
-            seed = parameters.get("seed", 42)
-            
-            # Set random seed for reproducibility
-            random.seed(seed)
-            
-            # Generate SVG-style sketch
-            image = self._generate_sketch(inputs, num_paths)
-            
-            return image
-            
-        except Exception as e:
-            print(f"Error in DiffSketcher handler: {e}")
-            # Return error image
-            error_img = Image.new('RGB', (224, 224), color='lightcoral')
-            draw = ImageDraw.Draw(error_img)
-            draw.text((10, 100), f"Error: {str(e)[:30]}", fill='white')
-            
-            return error_img
-    
-    def _generate_sketch(self, prompt: str, num_paths: int) -> Image.Image:
-        """Generate a sketch-style image based on the prompt"""
-        # Create canvas
-        width, height = 224, 224
-        image = Image.new('RGB', (width, height), color='white')
-        draw = ImageDraw.Draw(image)
+        # Generate sketch
+        image = self.generate_sketch(prompt)
+        
+        # Return PIL Image directly for HF Inference API
+        return image
+
+    def generate_sketch(self, prompt):
+        """Generate a sketch based on the prompt"""
+        # Create 224x224 white canvas
+        img = Image.new('RGB', (224, 224), 'white')
+        draw = ImageDraw.Draw(img)
         
-        # Simple prompt-based sketch generation
         prompt_lower = prompt.lower()
         
-        # Generate sketch elements based on prompt keywords
-        if any(word in prompt_lower for word in ['mountain', 'landscape', 'hill']):
-            self._draw_mountains(draw, width, height, num_paths)
-        elif any(word in prompt_lower for word in ['cat', 'animal', 'pet']):
-            self._draw_cat(draw, width, height, num_paths)
-        elif any(word in prompt_lower for word in ['flower', 'plant', 'garden']):
-            self._draw_flower(draw, width, height, num_paths)
+        if any(word in prompt_lower for word in ['mountain', 'landscape', 'nature']):
+            self._draw_landscape(draw)
+        elif any(word in prompt_lower for word in ['car', 'vehicle', 'automobile']):
+            self._draw_car(draw)
         elif any(word in prompt_lower for word in ['house', 'building', 'home']):
-            self._draw_house(draw, width, height, num_paths)
-        elif any(word in prompt_lower for word in ['tree', 'forest', 'wood']):
-            self._draw_tree(draw, width, height, num_paths)
+            self._draw_house(draw)
+        elif any(word in prompt_lower for word in ['flower', 'rose', 'bloom']):
+            self._draw_flower(draw)
+        elif any(word in prompt_lower for word in ['face', 'portrait', 'person']):
+            self._draw_face(draw)
+        elif any(word in prompt_lower for word in ['cat', 'animal', 'pet']):
+            self._draw_cat(draw)
         else:
-            self._draw_abstract(draw, width, height, num_paths)
+            self._draw_abstract(draw, prompt)
         
-        return image
-    
-    def _draw_mountains(self, draw, width, height, num_paths):
+        return img
+
+    def _draw_landscape(self, draw):
         """Draw mountain landscape"""
-        # Background mountains
-        for i in range(3):
-            y_offset = height // 3 + i * 20
-            points = []
-            for x in range(0, width + 20, 20):
-                y = y_offset + random.randint(-30, 10)
-                points.extend([x, y])
-            if len(points) >= 6:
-                draw.polygon(points + [width, height, 0, height], 
-                           fill=f'rgb({200-i*30},{220-i*20},{240-i*10})', 
-                           outline='gray')
-        
-        # Add some trees
-        for _ in range(num_paths // 20):
-            x = random.randint(10, width-10)
-            y = random.randint(height//2, height-20)
-            draw.ellipse([x-5, y-15, x+5, y], fill='darkgreen')
-            draw.line([x, y, x, y+15], fill='brown', width=2)
-    
-    def _draw_cat(self, draw, width, height, num_paths):
-        """Draw a simple cat"""
-        cx, cy = width//2, height//2
-        
+        # Sky
+        draw.rectangle([0, 0, 224, 120], fill=(200, 220, 255))
+        
+        # Mountains
+        points = [(0, 120)]
+        for x in range(0, 224, 30):
+            y = 120 + 40 * math.sin(x * 0.02) + 20 * math.sin(x * 0.05)
+            points.append((x, int(y)))
+        points.extend([(224, 224), (0, 224)])
+        draw.polygon(points, fill=(100, 150, 100), outline=(50, 100, 50))
+        
+        # Trees
+        for i in range(4):
+            x = 40 + i * 45
+            y = 140 + i * 5
+            # Trunk
+            draw.rectangle([x-3, y, x+3, y+30], fill=(101, 67, 33))
+            # Leaves
+            draw.ellipse([x-15, y-20, x+15, y+5], fill=(34, 139, 34), outline=(0, 100, 0))
+
+    def _draw_car(self, draw):
+        """Draw car sketch"""
+        # Main body
+        draw.rectangle([50, 120, 174, 160], fill=(150, 150, 150), outline=(0, 0, 0), width=2)
+        # Roof
+        draw.rectangle([70, 100, 154, 120], fill=(120, 120, 120), outline=(0, 0, 0), width=2)
+        # Wheels
+        draw.ellipse([60, 150, 80, 170], fill=(50, 50, 50), outline=(0, 0, 0), width=2)
+        draw.ellipse([144, 150, 164, 170], fill=(50, 50, 50), outline=(0, 0, 0), width=2)
+        # Windows
+        draw.rectangle([75, 105, 100, 115], fill=(150, 200, 255), outline=(0, 0, 0))
+        draw.rectangle([124, 105, 149, 115], fill=(150, 200, 255), outline=(0, 0, 0))
+
+    def _draw_house(self, draw):
+        """Draw house sketch"""
+        # Base
+        draw.rectangle([60, 130, 164, 190], fill=(200, 180, 140), outline=(0, 0, 0), width=2)
+        # Roof
+        draw.polygon([(60, 130), (112, 80), (164, 130)], fill=(180, 100, 100), outline=(0, 0, 0), width=2)
+        # Door
+        draw.rectangle([100, 160, 124, 190], fill=(139, 69, 19), outline=(0, 0, 0), width=2)
+        # Windows
+        draw.rectangle([70, 145, 90, 165], fill=(150, 200, 255), outline=(0, 0, 0), width=2)
+        draw.rectangle([134, 145, 154, 165], fill=(150, 200, 255), outline=(0, 0, 0), width=2)
+
+    def _draw_flower(self, draw):
+        """Draw flower sketch"""
+        center = (112, 112)
+        # Stem
+        draw.line([center[0], center[1]+20, center[0], 200], fill=(34, 139, 34), width=4)
+        # Petals
+        for angle in range(0, 360, 45):
+            x = center[0] + 25 * math.cos(math.radians(angle))
+            y = center[1] + 25 * math.sin(math.radians(angle))
+            draw.ellipse([x-8, y-15, x+8, y+5], fill=(255, 100, 150), outline=(200, 50, 100))
+        # Center
+        draw.ellipse([center[0]-8, center[1]-8, center[0]+8, center[1]+8], 
+                    fill=(255, 255, 0), outline=(200, 200, 0))
+
+    def _draw_face(self, draw):
+        """Draw face sketch"""
+        center = (112, 112)
+        # Head
+        draw.ellipse([center[0]-40, center[1]-50, center[0]+40, center[1]+30], 
+                    fill=(255, 220, 177), outline=(0, 0, 0), width=2)
+        # Eyes
+        draw.ellipse([center[0]-20, center[1]-20, center[0]-10, center[1]-10], 
+                    fill=(255, 255, 255), outline=(0, 0, 0))
+        draw.ellipse([center[0]+10, center[1]-20, center[0]+20, center[1]-10], 
+                    fill=(255, 255, 255), outline=(0, 0, 0))
+        # Pupils
+        draw.ellipse([center[0]-17, center[1]-17, center[0]-13, center[1]-13], fill=(0, 0, 0))
+        draw.ellipse([center[0]+13, center[1]-17, center[0]+17, center[1]-13], fill=(0, 0, 0))
+        # Nose
+        draw.line([center[0], center[1]-5, center[0]-3, center[1]+5], fill=(0, 0, 0), width=2)
+        # Mouth
+        draw.arc([center[0]-15, center[1]+5, center[0]+15, center[1]+20], 0, 180, fill=(0, 0, 0), width=2)
+
+    def _draw_cat(self, draw):
+        """Draw cat sketch"""
+        center = (112, 112)
         # Body
-        draw.ellipse([cx-40, cy-10, cx+40, cy+30], outline='black', width=2)
-        
+        draw.ellipse([center[0]-30, center[1], center[0]+30, center[1]+50], 
+                    fill=(200, 200, 200), outline=(0, 0, 0), width=2)
         # Head
-        draw.ellipse([cx-25, cy-40, cx+25, cy-10], outline='black', width=2)
-        
+        draw.ellipse([center[0]-20, center[1]-30, center[0]+20, center[1]+5], 
+                    fill=(200, 200, 200), outline=(0, 0, 0), width=2)
         # Ears
-        draw.polygon([cx-20, cy-35, cx-10, cy-50, cx-5, cy-35], outline='black', width=2)
-        draw.polygon([cx+5, cy-35, cx+10, cy-50, cx+20, cy-35], outline='black', width=2)
-        
+        draw.polygon([(center[0]-15, center[1]-25), (center[0]-10, center[1]-40), (center[0]-5, center[1]-30)], 
+                    fill=(200, 200, 200), outline=(0, 0, 0), width=2)
+        draw.polygon([(center[0]+5, center[1]-30), (center[0]+10, center[1]-40), (center[0]+15, center[1]-25)], 
+                    fill=(200, 200, 200), outline=(0, 0, 0), width=2)
         # Eyes
-        draw.ellipse([cx-15, cy-30, cx-10, cy-25], fill='black')
-        draw.ellipse([cx+10, cy-30, cx+15, cy-25], fill='black')
-        
+        draw.ellipse([center[0]-10, center[1]-20, center[0]-5, center[1]-15], fill=(0, 0, 0))
+        draw.ellipse([center[0]+5, center[1]-20, center[0]+10, center[1]-15], fill=(0, 0, 0))
         # Nose
-        draw.polygon([cx-2, cy-20, cx+2, cy-20, cx, cy-15], fill='pink')
-        
-        # Whiskers
-        for i in range(3):
-            y_pos = cy-18 + i*3
-            draw.line([cx-25, y_pos, cx-35, y_pos], fill='black', width=1)
-            draw.line([cx+25, y_pos, cx+35, y_pos], fill='black', width=1)
-        
-        # Tail
-        draw.arc([cx+30, cy-5, cx+60, cy+25], 0, 180, fill='black', width=3)
-    
-    def _draw_flower(self, draw, width, height, num_paths):
-        """Draw a simple flower"""
-        cx, cy = width//2, height//2
-        
-        # Stem
-        draw.line([cx, cy+20, cx, height-20], fill='green', width=4)
+        draw.polygon([(center[0]-1, center[1]-10), (center[0]+1, center[1]-10), (center[0], center[1]-7)], fill=(255, 100, 100))
+
+    def _draw_abstract(self, draw, prompt):
+        """Draw abstract sketch"""
+        prompt_hash = hash(prompt) % 100
         
-        # Petals
-        petal_colors = ['red', 'pink', 'yellow', 'orange', 'purple']
-        for i in range(6):
-            angle = i * 60
-            x1 = cx + 20 * math.cos(math.radians(angle))
-            y1 = cy + 20 * math.sin(math.radians(angle))
-            x2 = cx + 35 * math.cos(math.radians(angle))
-            y2 = cy + 35 * math.sin(math.radians(angle))
+        # Create flowing lines
+        for i in range(8):
+            points = []
+            start_x = (i * 30 + prompt_hash) % 200 + 12
+            start_y = (i * 25 + prompt_hash) % 180 + 22
             
-            color = random.choice(petal_colors)
-            draw.ellipse([x2-8, y2-8, x2+8, y2+8], fill=color, outline='darkred')
-        
-        # Center
-        draw.ellipse([cx-8, cy-8, cx+8, cy+8], fill='yellow', outline='orange')
-        
-        # Leaves
-        draw.ellipse([cx-15, cy+30, cx-5, cy+50], fill='green', outline='darkgreen')
-        draw.ellipse([cx+5, cy+35, cx+15, cy+55], fill='green', outline='darkgreen')
-    
-    def _draw_house(self, draw, width, height, num_paths):
-        """Draw a simple house"""
-        # House base
-        house_x, house_y = width//4, height//2
-        house_w, house_h = width//2, height//3
-        
-        draw.rectangle([house_x, house_y, house_x+house_w, house_y+house_h], 
-                      outline='black', width=2)
-        
-        # Roof
-        draw.polygon([house_x-10, house_y, width//2, house_y-40, house_x+house_w+10, house_y], 
-                    fill='red', outline='darkred', width=2)
-        
-        # Door
-        door_w, door_h = 20, 40
-        door_x = house_x + house_w//2 - door_w//2
-        door_y = house_y + house_h - door_h
-        draw.rectangle([door_x, door_y, door_x+door_w, door_y+door_h], 
-                      fill='brown', outline='black', width=2)
-        
-        # Windows
-        win_size = 15
-        draw.rectangle([house_x+10, house_y+15, house_x+10+win_size, house_y+15+win_size], 
-                      fill='lightblue', outline='black', width=2)
-        draw.rectangle([house_x+house_w-25, house_y+15, house_x+house_w-10, house_y+30], 
-                      fill='lightblue', outline='black', width=2)
-        
-        # Chimney
-        draw.rectangle([house_x+house_w-20, house_y-35, house_x+house_w-10, house_y-10], 
-                      fill='gray', outline='black', width=2)
-    
-    def _draw_tree(self, draw, width, height, num_paths):
-        """Draw a simple tree"""
-        cx, cy = width//2, height//2
-        
-        # Trunk
-        trunk_w, trunk_h = 15, 60
-        draw.rectangle([cx-trunk_w//2, cy+20, cx+trunk_w//2, cy+20+trunk_h], 
-                      fill='brown', outline='black', width=2)
-        
-        # Leaves (multiple circles for fuller look)
-        leaf_positions = [
-            (cx, cy-20), (cx-25, cy-10), (cx+25, cy-10),
-            (cx-15, cy+5), (cx+15, cy+5), (cx, cy+15)
-        ]
-        
-        for lx, ly in leaf_positions:
-            size = random.randint(20, 35)
-            draw.ellipse([lx-size//2, ly-size//2, lx+size//2, ly+size//2], 
-                        fill='green', outline='darkgreen', width=1)
-    
-    def _draw_abstract(self, draw, width, height, num_paths):
-        """Draw abstract shapes"""
-        colors = ['red', 'blue', 'green', 'yellow', 'purple', 'orange', 'pink']
-        
-        # Draw random shapes
-        for _ in range(min(num_paths//10, 8)):
-            color = random.choice(colors)
-            shape_type = random.choice(['circle', 'rectangle', 'line'])
+            for j in range(4):
+                x = start_x + j * 25 + 15 * math.sin((i + j + prompt_hash) * 0.5)
+                y = start_y + j * 20 + 15 * math.cos((i + j + prompt_hash) * 0.3)
+                points.append((int(x) % 224, int(y) % 224))
             
-            if shape_type == 'circle':
-                x, y = random.randint(20, width-20), random.randint(20, height-20)
-                r = random.randint(10, 30)
-                draw.ellipse([x-r, y-r, x+r, y+r], fill=color, outline='black')
-            elif shape_type == 'rectangle':
-                x1, y1 = random.randint(10, width//2), random.randint(10, height//2)
-                x2, y2 = random.randint(width//2, width-10), random.randint(height//2, height-10)
-                draw.rectangle([x1, y1, x2, y2], fill=color, outline='black')
-            else:  # line
-                x1, y1 = random.randint(0, width), random.randint(0, height)
-                x2, y2 = random.randint(0, width), random.randint(0, height)
-                draw.line([x1, y1, x2, y2], fill=color, width=random.randint(2, 5))
+            # Draw connected lines
+            for k in range(len(points)-1):
+                color = (50 + (i * 25) % 150, 50 + (i * 35) % 150, 50 + (i * 45) % 150)
+                draw.line([points[k], points[k+1]], fill=color, width=2)