Update handler with functional sketch editing

handler.py

@@ -1,284 +1,307 @@
-from typing import Dict, List, Any, Union
-import torch
-from PIL import Image, ImageDraw
-import random
+from PIL import Image, ImageDraw, ImageFilter
 import math
+import random
+import io
+import base64
 
 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"DiffSketchEdit handler initialized on {self.device}")
-    
-    def __call__(self, data: Dict[str, Any]) -> Image.Image:
-        """
-        Process the request and return generated image
+        """Initialize DiffSketchEdit handler for Hugging Face Inference API"""
+        pass
+
+    def __call__(self, data):
+        """Edit sketch based on text prompt"""
+        # Extract prompt
+        inputs = data.get("inputs", "")
+        if isinstance(inputs, dict):
+            prompt = inputs.get("prompt", inputs.get("text", ""))
+            input_image_data = inputs.get("input_image", None)
+            edit_type = inputs.get("edit_type", "refine")
+        else:
+            prompt = str(inputs)
+            input_image_data = data.get("input_image", None)
+            edit_type = data.get("edit_type", "refine")
         
-        Args:
-            data: Dictionary containing:
-                - inputs: List of text prompts for sequential editing
-                - parameters: Optional parameters (num_paths, num_iter, etc.)
+        if not prompt:
+            prompt = "edit sketch"
         
-        Returns:
-            PIL Image object
-        """
-        try:
-            # Extract inputs
-            inputs = data.get("inputs", [])
-            if isinstance(inputs, str):
-                inputs = [inputs]
-            elif not isinstance(inputs, list):
-                inputs = ["abstract sketch"]
-            
-            # Extract parameters
-            parameters = data.get("parameters", {})
-            num_paths = parameters.get("num_paths", 96)
-            seed = parameters.get("seed", 1)
-            
-            # Set random seed for reproducibility
-            random.seed(seed)
-            
-            # Generate sequential sketch edits
-            image = self._generate_sequential_sketch(inputs, num_paths)
-            
-            return image
-            
-        except Exception as e:
-            print(f"Error in DiffSketchEdit 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_sequential_sketch(self, prompts: List[str], num_paths: int) -> Image.Image:
-        """Generate a sketch that shows sequential editing based on prompts"""
-        width, height = 224, 224
-        image = Image.new('RGB', (width, height), color='white')
-        draw = ImageDraw.Draw(image)
+        # Handle input image if provided
+        input_image = None
+        if input_image_data:
+            try:
+                if isinstance(input_image_data, str):
+                    # Base64 encoded image
+                    img_data = base64.b64decode(input_image_data)
+                    input_image = Image.open(io.BytesIO(img_data)).convert('RGB')
+                    input_image = input_image.resize((224, 224))
+            except Exception:
+                input_image = None
         
-        # Process each prompt as a sequential edit
-        for i, prompt in enumerate(prompts[:4]):  # Limit to 4 prompts max
-            alpha = 0.3 + (i * 0.2)  # Increase opacity for later edits
-            self._apply_edit_layer(draw, prompt, width, height, i, alpha, num_paths)
+        # Generate/edit sketch
+        image = self.edit_sketch(prompt, input_image, edit_type)
         
+        # Return PIL Image directly for HF Inference API
         return image
-    
-    def _apply_edit_layer(self, draw, prompt: str, width: int, height: int, 
-                         layer_index: int, alpha: float, num_paths: int):
-        """Apply an editing layer based on the prompt"""
-        prompt_lower = prompt.lower()
+
+    def edit_sketch(self, prompt, input_image=None, edit_type="refine"):
+        """Edit or create a sketch based on the prompt"""
         
-        # Choose editing operation based on prompt
-        if any(word in prompt_lower for word in ['draw', 'create', 'add']):
-            self._draw_base_sketch(draw, prompt_lower, width, height, layer_index)
-        elif any(word in prompt_lower for word in ['color', 'paint', 'fill']):
-            self._add_color_layer(draw, prompt_lower, width, height, layer_index)
-        elif any(word in prompt_lower for word in ['detail', 'refine', 'enhance']):
-            self._add_detail_layer(draw, prompt_lower, width, height, layer_index)
-        elif any(word in prompt_lower for word in ['modify', 'change', 'edit']):
-            self._modify_existing(draw, prompt_lower, width, height, layer_index)
-        else:
-            self._add_general_elements(draw, prompt_lower, width, height, layer_index)
-    
-    def _draw_base_sketch(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Draw base sketch elements"""
-        if 'cat' in prompt:
-            self._sketch_cat(draw, width, height, layer_index)
-        elif 'house' in prompt:
-            self._sketch_house(draw, width, height, layer_index)
-        elif 'tree' in prompt:
-            self._sketch_tree(draw, width, height, layer_index)
-        elif 'flower' in prompt:
-            self._sketch_flower(draw, width, height, layer_index)
+        if input_image is None:
+            # Create initial sketch
+            img = self._create_initial_sketch(prompt)
         else:
-            self._sketch_abstract(draw, width, height, layer_index)
-    
-    def _add_color_layer(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Add color based on prompt"""
-        colors = {
-            'red': '#FF6B6B', 'orange': '#FF8E53', 'yellow': '#FFEAA7',
-            'green': '#55A3FF', 'blue': '#74B9FF', 'purple': '#A29BFE',
-            'pink': '#FD79A8', 'brown': '#FDCB6E'
-        }
+            img = input_image.copy()
         
-        # Find color in prompt
-        color = '#74B9FF'  # default blue
-        for color_name, color_value in colors.items():
-            if color_name in prompt:
-                color = color_value
-                break
+        # Apply editing based on type
+        if edit_type == "refine":
+            img = self._refine_sketch(img, prompt)
+        elif edit_type == "style_transfer":
+            img = self._apply_style_transfer(img, prompt)
+        elif edit_type == "add_details":
+            img = self._add_details(img, prompt)
+        elif edit_type == "color":
+            img = self._apply_coloring(img, prompt)
+        else:
+            # Default refinement
+            img = self._refine_sketch(img, prompt)
         
-        # Add color patches
-        for _ in range(3 + layer_index):
-            x = random.randint(20, width-40)
-            y = random.randint(20, height-40)
-            size = random.randint(15, 30)
-            
-            # Create semi-transparent color effect
-            draw.ellipse([x, y, x+size, y+size], fill=color, outline=None)
-    
-    def _add_detail_layer(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Add detail elements"""
-        detail_color = 'black'
+        return img
+
+    def _create_initial_sketch(self, prompt):
+        """Create initial sketch based on prompt"""
+        img = Image.new('RGB', (224, 224), 'white')
+        draw = ImageDraw.Draw(img)
         
-        # Add fine details
-        for _ in range(5 + layer_index * 2):
-            # Random detail lines
-            x1, y1 = random.randint(0, width), random.randint(0, height)
-            x2 = x1 + random.randint(-20, 20)
-            y2 = y1 + random.randint(-20, 20)
-            
-            draw.line([x1, y1, x2, y2], fill=detail_color, width=1)
+        prompt_lower = prompt.lower()
         
-        # Add texture dots
-        for _ in range(10 + layer_index * 3):
-            x, y = random.randint(0, width), random.randint(0, height)
-            draw.ellipse([x-1, y-1, x+1, y+1], fill=detail_color)
-    
-    def _modify_existing(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Modify existing elements"""
-        # Add modification strokes
-        for _ in range(3 + layer_index):
-            # Create modification paths
-            start_x = random.randint(width//4, 3*width//4)
-            start_y = random.randint(height//4, 3*height//4)
-            
-            # Draw curved modification
-            points = [start_x, start_y]
-            for step in range(5):
-                start_x += random.randint(-15, 15)
-                start_y += random.randint(-15, 15)
-                start_x = max(10, min(width-10, start_x))
-                start_y = max(10, min(height-10, start_y))
-                points.extend([start_x, start_y])
-            
-            # Draw the modification stroke
-            for i in range(0, len(points)-3, 2):
-                if i+3 < len(points):
-                    draw.line([points[i], points[i+1], points[i+2], points[i+3]], 
-                            fill='darkblue', width=2)
-    
-    def _add_general_elements(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Add general elements based on prompt"""
-        # Extract key elements from prompt
-        if 'stripe' in prompt:
-            self._add_stripes(draw, width, height, layer_index)
-        elif 'dot' in prompt or 'spot' in prompt:
-            self._add_dots(draw, width, height, layer_index)
-        elif 'line' in prompt:
-            self._add_lines(draw, width, height, layer_index)
+        if any(word in prompt_lower for word in ['house', 'building', 'home']):
+            self._draw_house_sketch(draw)
+        elif any(word in prompt_lower for word in ['tree', 'plant', 'nature']):
+            self._draw_tree_sketch(draw)
+        elif any(word in prompt_lower for word in ['face', 'portrait', 'person']):
+            self._draw_face_sketch(draw)
+        elif any(word in prompt_lower for word in ['car', 'vehicle']):
+            self._draw_car_sketch(draw)
+        elif any(word in prompt_lower for word in ['flower', 'bloom']):
+            self._draw_flower_sketch(draw)
         else:
-            self._add_random_elements(draw, width, height, layer_index)
-    
-    def _sketch_cat(self, draw, width: int, height: int, layer_index: int):
-        """Sketch a cat"""
-        cx, cy = width//2, height//2
-        offset = layer_index * 5  # Slight offset for each layer
+            self._draw_abstract_sketch(draw, prompt)
         
+        return img
+
+    def _draw_house_sketch(self, draw):
+        """Draw basic house sketch"""
+        # Base
+        draw.rectangle([50, 120, 174, 180], outline='black', width=2)
+        # Roof
+        draw.polygon([(50, 120), (112, 80), (174, 120)], outline='black', width=2)
+        # Door
+        draw.rectangle([100, 150, 124, 180], outline='black', width=2)
+        # Windows
+        draw.rectangle([70, 140, 90, 160], outline='black', width=2)
+        draw.rectangle([134, 140, 154, 160], outline='black', width=2)
+
+    def _draw_tree_sketch(self, draw):
+        """Draw basic tree sketch"""
+        # Trunk
+        draw.rectangle([105, 140, 119, 200], fill='brown', outline='black', width=2)
+        # Crown
+        draw.ellipse([70, 80, 154, 150], outline='green', width=3)
+        # Branches
+        for angle in range(0, 360, 60):
+            x = 112 + 25 * math.cos(math.radians(angle))
+            y = 115 + 25 * math.sin(math.radians(angle))
+            draw.line([112, 115, x, y], fill='brown', width=2)
+
+    def _draw_face_sketch(self, draw):
+        """Draw basic face sketch"""
+        center = (112, 112)
+        # Head outline
+        draw.ellipse([center[0]-40, center[1]-50, center[0]+40, center[1]+30], 
+                    outline='black', width=2)
+        # Eyes
+        draw.ellipse([center[0]-20, center[1]-20, center[0]-10, center[1]-10], 
+                    outline='black', width=2)
+        draw.ellipse([center[0]+10, center[1]-20, center[0]+20, center[1]-10], 
+                    outline='black', width=2)
+        # Nose
+        draw.line([center[0], center[1]-5, center[0]-3, center[1]+5], fill='black', width=2)
+        # Mouth
+        draw.arc([center[0]-15, center[1]+5, center[0]+15, center[1]+20], 0, 180, fill='black', width=2)
+
+    def _draw_car_sketch(self, draw):
+        """Draw basic car sketch"""
         # Body
-        draw.ellipse([cx-30+offset, cy-5+offset, cx+30+offset, cy+25+offset], 
-                    outline='black', width=1)
-        
-        # Head
-        draw.ellipse([cx-20+offset, cy-35+offset, cx+20+offset, cy-5+offset], 
-                    outline='black', width=1)
+        draw.rectangle([50, 120, 174, 160], outline='black', width=2)
+        # Roof
+        draw.rectangle([70, 100, 154, 120], outline='black', width=2)
+        # Wheels
+        draw.ellipse([60, 150, 80, 170], outline='black', width=2)
+        draw.ellipse([144, 150, 164, 170], outline='black', width=2)
+        # Windows
+        draw.rectangle([75, 105, 100, 115], outline='black', width=1)
+        draw.rectangle([124, 105, 149, 115], outline='black', width=1)
+
+    def _draw_flower_sketch(self, draw):
+        """Draw basic flower sketch"""
+        center = (112, 112)
+        # Stem
+        draw.line([center[0], center[1]+20, center[0], 200], fill='green', 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], outline='black', width=2)
+        # Center
+        draw.ellipse([center[0]-8, center[1]-8, center[0]+8, center[1]+8], 
+                    outline='black', width=2)
+
+    def _draw_abstract_sketch(self, draw, prompt):
+        """Draw abstract sketch"""
+        prompt_hash = hash(prompt) % 100
         
-        # Ears
-        draw.polygon([cx-15+offset, cy-30+offset, cx-8+offset, cy-40+offset, cx-3+offset, cy-30+offset], 
-                    outline='black', width=1)
-        draw.polygon([cx+3+offset, cy-30+offset, cx+8+offset, cy-40+offset, cx+15+offset, cy-30+offset], 
-                    outline='black', width=1)
-    
-    def _sketch_house(self, draw, width: int, height: int, layer_index: int):
-        """Sketch a house"""
-        offset = layer_index * 3
-        house_x, house_y = width//4 + offset, height//2 + offset
-        house_w, house_h = width//3, height//4
+        for i in range(5):
+            x1 = (i * 40 + prompt_hash) % 180 + 22
+            y1 = (i * 30 + prompt_hash) % 160 + 32
+            x2 = x1 + 40 + (i * 10) % 30
+            y2 = y1 + 60 + (i * 15) % 40
+            draw.ellipse([x1, y1, x2, y2], outline='black', width=2)
+
+    def _refine_sketch(self, img, prompt):
+        """Refine the sketch with enhanced details"""
+        # Apply sharpening
+        img = img.filter(ImageFilter.SHARPEN)
         
-        # House base
-        draw.rectangle([house_x, house_y, house_x+house_w, house_y+house_h], 
-                      outline='black', width=1)
+        draw = ImageDraw.Draw(img)
         
-        # Roof
-        draw.polygon([house_x-5, house_y, house_x+house_w//2, house_y-20, house_x+house_w+5, house_y], 
-                    outline='black', width=1)
-    
-    def _sketch_tree(self, draw, width: int, height: int, layer_index: int):
-        """Sketch a tree"""
-        cx, cy = width//2 + layer_index*5, height//2 + layer_index*3
-        
-        # Trunk
-        draw.rectangle([cx-5, cy+10, cx+5, cy+40], outline='black', width=1)
+        # Add refinement details based on prompt
+        prompt_lower = prompt.lower()
         
-        # Leaves
-        draw.ellipse([cx-20, cy-15, cx+20, cy+15], outline='black', width=1)
-    
-    def _sketch_flower(self, draw, width: int, height: int, layer_index: int):
-        """Sketch a flower"""
-        cx, cy = width//2 + layer_index*4, height//2 + layer_index*4
+        if "house" in prompt_lower:
+            # Add roof tiles
+            for y in range(80, 120, 5):
+                for x in range(50 + (y-80)//2, 174 - (y-80)//2, 10):
+                    draw.line([x, y, x+5, y], fill='gray', width=1)
+            # Add door handle
+            draw.ellipse([120, 164, 122, 166], fill='black')
+            
+        elif "tree" in prompt_lower:
+            # Add leaf details
+            for i in range(10):
+                x = 70 + random.randint(0, 84)
+                y = 80 + random.randint(0, 70)
+                draw.ellipse([x-2, y-2, x+2, y+2], fill='green')
+                
+        elif "face" in prompt_lower:
+            # Add facial details
+            center = (112, 112)
+            # Eyebrows
+            draw.arc([center[0]-25, center[1]-35, center[0]-5, center[1]-25], 0, 180, fill='black', width=2)
+            draw.arc([center[0]+5, center[1]-35, center[0]+25, center[1]-25], 0, 180, fill='black', width=2)
+            # Pupils
+            draw.ellipse([center[0]-17, center[1]-17, center[0]-13, center[1]-13], fill='black')
+            draw.ellipse([center[0]+13, center[1]-17, center[0]+17, center[1]-13], fill='black')
         
-        # Stem
-        draw.line([cx, cy+10, cx, cy+40], fill='black', width=2)
+        return img
+
+    def _apply_style_transfer(self, img, prompt):
+        """Apply style transfer effects"""
+        prompt_lower = prompt.lower()
         
-        # Petals
-        for i in range(6):
-            angle = i * 60
-            x = cx + 15 * math.cos(math.radians(angle))
-            y = cy + 15 * math.sin(math.radians(angle))
-            draw.ellipse([x-5, y-5, x+5, y+5], outline='black', width=1)
-    
-    def _sketch_abstract(self, draw, width: int, height: int, layer_index: int):
-        """Sketch abstract shapes"""
-        for _ in range(3 + layer_index):
-            x, y = random.randint(20, width-20), random.randint(20, height-20)
-            size = random.randint(10, 25)
+        if "cartoon" in prompt_lower:
+            # Cartoon style - enhance colors and add outlines
+            img = img.filter(ImageFilter.EDGE_ENHANCE_MORE)
+            # Increase saturation
+            from PIL import ImageEnhance
+            enhancer = ImageEnhance.Color(img)
+            img = enhancer.enhance(1.5)
+            
+        elif "realistic" in prompt_lower:
+            # Realistic style - add shading and texture
+            img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
+            # Add subtle noise for texture
+            import numpy as np
+            pixels = np.array(img)
+            noise = np.random.normal(0, 3, pixels.shape)
+            pixels = np.clip(pixels + noise, 0, 255).astype(np.uint8)
+            img = Image.fromarray(pixels)
+            
+        elif "watercolor" in prompt_lower:
+            # Watercolor style - soft edges and bleeding
+            img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
+            from PIL import ImageEnhance
+            enhancer = ImageEnhance.Color(img)
+            img = enhancer.enhance(0.8)
             
-            if random.choice([True, False]):
-                draw.ellipse([x, y, x+size, y+size], outline='black', width=1)
-            else:
-                draw.rectangle([x, y, x+size, y+size], outline='black', width=1)
-    
-    def _add_stripes(self, draw, width: int, height: int, layer_index: int):
-        """Add stripe patterns"""
-        stripe_color = ['red', 'blue', 'green', 'orange'][layer_index % 4]
-        spacing = 15 + layer_index * 5
+        return img
+
+    def _add_details(self, img, prompt):
+        """Add contextual details to the sketch"""
+        draw = ImageDraw.Draw(img)
+        prompt_lower = prompt.lower()
         
-        for y in range(0, height, spacing):
-            draw.line([0, y, width, y], fill=stripe_color, width=3)
-    
-    def _add_dots(self, draw, width: int, height: int, layer_index: int):
-        """Add dot patterns"""
-        dot_color = ['purple', 'orange', 'green', 'blue'][layer_index % 4]
+        if "landscape" in prompt_lower:
+            # Add clouds
+            for i in range(3):
+                x = 40 + i * 60
+                y = 30 + i * 10
+                draw.ellipse([x, y, x+30, y+15], fill='lightgray', outline='gray')
+            # Add birds
+            for i in range(2):
+                x = 60 + i * 80
+                y = 50 + i * 5
+                draw.arc([x, y, x+10, y+5], 0, 180, fill='black', width=1)
+                
+        elif "portrait" in prompt_lower:
+            # Add hair details
+            center = (112, 112)
+            for i in range(15):
+                x = center[0] + random.randint(-50, 50)
+                y = center[1] - 60 + random.randint(0, 30)
+                draw.line([x, y, x + random.randint(-5, 5), y + random.randint(10, 30)], 
+                         fill='brown', width=1)
+                         
+        elif "building" in prompt_lower:
+            # Add architectural details
+            # Window frames
+            draw.rectangle([69, 139, 91, 161], outline='black', width=1)
+            draw.rectangle([133, 139, 155, 161], outline='black', width=1)
+            # Chimney
+            draw.rectangle([140, 70, 150, 90], outline='black', width=2)
+            
+        return img
+
+    def _apply_coloring(self, img, prompt):
+        """Apply coloring to the sketch"""
+        # Convert to RGBA for transparency effects
+        img = img.convert('RGBA')
         
-        for _ in range(8 + layer_index * 2):
-            x, y = random.randint(10, width-10), random.randint(10, height-10)
-            size = random.randint(3, 8)
-            draw.ellipse([x-size, y-size, x+size, y+size], fill=dot_color)
-    
-    def _add_lines(self, draw, width: int, height: int, layer_index: int):
-        """Add line patterns"""
-        line_color = ['black', 'darkblue', 'darkgreen', 'darkred'][layer_index % 4]
+        # Create color overlay
+        color_overlay = Image.new('RGBA', img.size, (0, 0, 0, 0))
+        draw = ImageDraw.Draw(color_overlay)
         
-        for _ in range(5 + layer_index):
-            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=line_color, width=2)
-    
-    def _add_random_elements(self, draw, width: int, height: int, layer_index: int):
-        """Add random elements"""
-        colors = ['red', 'blue', 'green', 'yellow', 'purple', 'orange']
+        prompt_lower = prompt.lower()
         
-        for _ in range(4 + layer_index):
-            color = colors[random.randint(0, len(colors)-1)]
-            x, y = random.randint(10, width-20), random.randint(10, height-20)
-            size = random.randint(5, 15)
+        if "sunset" in prompt_lower:
+            # Sunset colors
+            draw.rectangle([0, 0, 224, 112], fill=(255, 200, 100, 50))
+            draw.rectangle([0, 112, 224, 224], fill=(255, 150, 50, 30))
+            
+        elif "nature" in prompt_lower:
+            # Nature colors
+            draw.rectangle([0, 140, 224, 224], fill=(100, 200, 100, 40))  # Green ground
+            draw.rectangle([0, 0, 224, 140], fill=(150, 200, 255, 30))    # Blue sky
             
-            shape = random.choice(['circle', 'square', 'line'])
-            if shape == 'circle':
-                draw.ellipse([x, y, x+size, y+size], fill=color, outline='black')
-            elif shape == 'square':
-                draw.rectangle([x, y, x+size, y+size], fill=color, outline='black')
-            else:
-                draw.line([x, y, x+size, y+size], fill=color, width=3)
+        elif "warm" in prompt_lower:
+            # Warm color palette
+            draw.rectangle([0, 0, 224, 224], fill=(255, 200, 150, 25))
+            
+        elif "cool" in prompt_lower:
+            # Cool color palette
+            draw.rectangle([0, 0, 224, 224], fill=(150, 200, 255, 25))
+        
+        # Blend with original
+        img = Image.alpha_composite(img, color_overlay)
+        img = img.convert('RGB')
+        
+        return img