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)