handler_fallback.py

from PIL import Image, ImageDraw
import io
import random
import math

class EndpointHandler:
    def __init__(self):
        """Initialize the DiffSketcher handler with fallback PIL drawing"""
        pass
    
    def __call__(self, data):
        """
        Generate a sketch-style image using PIL drawing (fallback method)
        
        Args:
            data (dict): Input data containing:
                - inputs (str): Text prompt
                - parameters (dict): Generation parameters
        
        Returns:
            PIL.Image.Image: Generated sketch image
        """
        try:
            # Extract inputs
            prompt = data.get("inputs", "")
            parameters = data.get("parameters", {})
            
            # Extract parameters
            width = parameters.get("width", 224)
            height = parameters.get("height", 224)
            guidance_scale = parameters.get("guidance_scale", 7.5)
            seed = parameters.get("seed", 42)
            
            # Set random seed for reproducibility
            random.seed(seed)
            
            # Create white background
            image = Image.new('RGB', (width, height), 'white')
            draw = ImageDraw.Draw(image)
            
            # Generate sketch based on prompt keywords
            self._draw_sketch_from_prompt(draw, prompt, width, height)
            
            return image
            
        except Exception as e:
            # Return error image
            error_image = Image.new('RGB', (224, 224), 'white')
            error_draw = ImageDraw.Draw(error_image)
            error_draw.text((10, 100), f"Error: {str(e)[:30]}", fill='red')
            return error_image
    
    def _draw_sketch_from_prompt(self, draw, prompt, width, height):
        """Draw a simple sketch based on prompt keywords"""
        prompt_lower = prompt.lower()
        
        # Define colors for sketching
        colors = ['black', 'gray', 'darkgray']
        
        if any(word in prompt_lower for word in ['car', 'vehicle', 'automobile']):
            self._draw_car(draw, width, height, colors)
        elif any(word in prompt_lower for word in ['house', 'building', 'home']):
            self._draw_house(draw, width, height, colors)
        elif any(word in prompt_lower for word in ['flower', 'plant', 'bloom']):
            self._draw_flower(draw, width, height, colors)
        elif any(word in prompt_lower for word in ['tree', 'forest']):
            self._draw_tree(draw, width, height, colors)
        elif any(word in prompt_lower for word in ['mountain', 'landscape']):
            self._draw_mountain(draw, width, height, colors)
        else:
            self._draw_abstract(draw, width, height, colors)
    
    def _draw_car(self, draw, width, height, colors):
        """Draw a simple car sketch"""
        cx, cy = width // 2, height // 2
        
        # Car body
        draw.rectangle([cx-60, cy-20, cx+60, cy+20], outline=colors[0], width=2)
        
        # Car roof
        draw.rectangle([cx-40, cy-40, cx+40, cy-20], outline=colors[0], width=2)
        
        # Wheels
        draw.ellipse([cx-50, cy+10, cx-30, cy+30], outline=colors[0], width=2)
        draw.ellipse([cx+30, cy+10, cx+50, cy+30], outline=colors[0], width=2)
        
        # Windows
        draw.rectangle([cx-35, cy-35, cx+35, cy-25], outline=colors[1], width=1)
    
    def _draw_house(self, draw, width, height, colors):
        """Draw a simple house sketch"""
        cx, cy = width // 2, height // 2
        
        # House base
        draw.rectangle([cx-50, cy-10, cx+50, cy+40], outline=colors[0], width=2)
        
        # Roof
        draw.polygon([cx-60, cy-10, cx, cy-50, cx+60, cy-10], outline=colors[0], width=2)
        
        # Door
        draw.rectangle([cx-15, cy+10, cx+15, cy+40], outline=colors[1], width=2)
        
        # Windows
        draw.rectangle([cx-40, cy-5, cx-25, cy+10], outline=colors[1], width=1)
        draw.rectangle([cx+25, cy-5, cx+40, cy+10], outline=colors[1], width=1)
    
    def _draw_flower(self, draw, width, height, colors):
        """Draw a simple flower sketch"""
        cx, cy = width // 2, height // 2
        
        # Stem
        draw.line([cx, cy+20, cx, cy+60], fill=colors[0], width=3)
        
        # Petals
        for i in range(6):
            angle = i * 60 * math.pi / 180
            x = cx + 25 * math.cos(angle)
            y = cy + 25 * math.sin(angle)
            draw.ellipse([x-8, y-8, x+8, y+8], outline=colors[0], width=2)
        
        # Center
        draw.ellipse([cx-8, cy-8, cx+8, cy+8], fill=colors[1], outline=colors[0], width=2)
        
        # Leaves
        draw.ellipse([cx-10, cy+30, cx+10, cy+50], outline=colors[0], width=2)
    
    def _draw_tree(self, draw, width, height, colors):
        """Draw a simple tree sketch"""
        cx, cy = width // 2, height // 2
        
        # Trunk
        draw.rectangle([cx-8, cy+10, cx+8, cy+60], outline=colors[0], width=2)
        
        # Tree crown (circle)
        draw.ellipse([cx-40, cy-40, cx+40, cy+20], outline=colors[0], width=2)
        
        # Branches
        for i in range(5):
            angle = (i * 72 - 90) * math.pi / 180
            x = cx + 30 * math.cos(angle)
            y = cy + 30 * math.sin(angle)
            draw.line([cx, cy, x, y], fill=colors[1], width=1)
    
    def _draw_mountain(self, draw, width, height, colors):
        """Draw a simple mountain landscape"""
        cx, cy = width // 2, height // 2
        
        # Mountains
        draw.polygon([20, cy+30, 80, cy-40, 140, cy+30], outline=colors[0], width=2)
        draw.polygon([100, cy+30, 160, cy-20, 200, cy+30], outline=colors[0], width=2)
        
        # Ground line
        draw.line([0, cy+30, width, cy+30], fill=colors[1], width=1)
        
        # Sun
        draw.ellipse([width-60, 20, width-20, 60], outline=colors[1], width=2)
    
    def _draw_abstract(self, draw, width, height, colors):
        """Draw abstract shapes for unknown prompts"""
        cx, cy = width // 2, height // 2
        
        # Random geometric shapes
        for i in range(5):
            x = random.randint(20, width-20)
            y = random.randint(20, height-20)
            size = random.randint(10, 30)
            
            if i % 3 == 0:
                draw.ellipse([x-size, y-size, x+size, y+size], outline=colors[i%len(colors)], width=2)
            elif i % 3 == 1:
                draw.rectangle([x-size, y-size, x+size, y+size], outline=colors[i%len(colors)], width=2)
            else:
                points = []
                for j in range(3):
                    px = x + size * math.cos(j * 120 * math.pi / 180)
                    py = y + size * math.sin(j * 120 * math.pi / 180)
                    points.extend([px, py])
                draw.polygon(points, outline=colors[i%len(colors)], width=2)