File size: 7,481 Bytes
5c522de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
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) |