handler_fresh_1748635848.py

from PIL import Image, ImageDraw
import math
import random

class EndpointHandler:
    def __init__(self, path=""):
        """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)
        
        if not prompt:
            prompt = "a simple sketch"
        
        # 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)
        
        prompt_lower = prompt.lower()
        
        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)
        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, prompt)
        
        return img

    def _draw_landscape(self, draw):
        """Draw mountain landscape"""
        # 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([center[0]-30, center[1], center[0]+30, center[1]+50], 
                    fill=(200, 200, 200), outline=(0, 0, 0), width=2)
        # Head
        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([(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([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([(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
        
        # 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
            
            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))
            
            # 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)