Update app.py

app.py

@@ -2,46 +2,101 @@ import gradio as gr
 from PIL import Image, ImageDraw, ImageFont
 import random
 
-def draw_rat(draw):
+def draw_cat(draw):
     # Body
     draw.ellipse((170, 240, 340, 380), outline="white", width=6)
     # Head
     draw.ellipse((120, 260, 210, 330), outline="white", width=6)
     # Ears
-    draw.ellipse((130, 240, 160, 270), outline="white", width=6)
-    draw.ellipse((180, 240, 210, 270), outline="white", width=6)
+    draw.polygon([(130, 260), (150, 230), (170, 260)], outline="white", fill=None)
+    draw.polygon([(180, 260), (200, 230), (220, 260)], outline="white", fill=None)
     # Tail
     draw.line((340, 320, 450, 260), fill="white", width=6)
     # Eyes
     draw.ellipse((155, 285, 165, 295), fill="white")
     draw.ellipse((185, 285, 195, 295), fill="white")
 
+def draw_rat(draw):
+    draw.ellipse((170, 240, 340, 380), outline="white", width=6)
+    draw.ellipse((120, 260, 210, 330), outline="white", width=6)
+    draw.ellipse((130, 240, 160, 270), outline="white", width=6)
+    draw.ellipse((180, 240, 210, 270), outline="white", width=6)
+    draw.line((340, 320, 460, 260), fill="white", width=6)
+    draw.ellipse((155, 285, 165, 295), fill="white")
+    draw.ellipse((185, 285, 195, 295), fill="white")
+
+def draw_dog(draw):
+    draw.ellipse((170, 240, 340, 380), outline="white", width=6)
+    draw.ellipse((120, 260, 210, 330), outline="white", width=6)
+    draw.ellipse((110, 250, 140, 290), outline="white", width=6)
+    draw.ellipse((200, 250, 230, 290), outline="white", width=6)
+    draw.line((340, 320, 460, 260), fill="white", width=6)
+    draw.ellipse((155, 285, 165, 295), fill="white")
+    draw.ellipse((185, 285, 195, 295), fill="white")
+
+def draw_lion(draw):
+    draw.ellipse((170, 240, 340, 380), outline="white", width=6)
+    draw.ellipse((120, 260, 210, 330), outline="white", width=6)
+    # Mane
+    for i in range(0, 360, 20):
+        x1 = 165 + 90 * random.random()
+        y1 = 260 + 90 * random.random()
+        draw.arc((x1, y1, x1 + 120, y1 + 120), start=i, end=i+10, fill="white", width=6)
+    draw.ellipse((155, 285, 165, 295), fill="white")
+    draw.ellipse((185, 285, 195, 295), fill="white")
+
+def draw_bird(draw):
+    draw.ellipse((170, 240, 340, 380), outline="white", width=6)
+    draw.polygon([(120, 300), (170, 250), (170, 350)], outline="white", fill=None)
+    draw.polygon([(340, 300), (290, 250), (290, 350)], outline="white", fill=None)
+    draw.line((250, 380, 250, 450), fill="white", width=6)
+    draw.ellipse((155, 285, 165, 295), fill="white")
+    draw.ellipse((185, 285, 195, 295), fill="white")
+
+def draw_fish(draw):
+    draw.ellipse((170, 250, 340, 350), outline="white", width=6)
+    draw.polygon([(340, 300), (410, 270), (410, 330)], outline="white", fill=None)
+    draw.line((200, 300, 140, 270), fill="white", width=6)
+    draw.ellipse((240, 285, 250, 295), fill="white")
+
+def draw_turtle(draw):
+    draw.ellipse((170, 250, 340, 380), outline="white", width=6)
+    draw.ellipse((120, 270, 190, 330), outline="white", width=6)
+    draw.ellipse((340, 270, 410, 330), outline="white", width=6)
+    draw.ellipse((200, 380, 240, 430), outline="white", width=6)
+    draw.ellipse((270, 380, 310, 430), outline="white", width=6)
+    draw.ellipse((155, 285, 165, 295), fill="white")
+
 def generate_image(prompt):
     img = Image.new("RGB", (512, 512), (20, 20, 30))
     draw = ImageDraw.Draw(img)
 
-    # Background
+    # Gradient background
     for i in range(512):
-        color = (20 + i // 20, 20 + i // 25, 30 + i // 15)
+        color = (30 + i // 18, 20 + i // 30, 50 + i // 20)
         draw.line([(0, i), (512, i)], fill=color)
 
-    # Basic prompt detection
     prompt_lower = prompt.lower()
+
     if "rat" in prompt_lower:
         draw_rat(draw)
     elif "cat" in prompt_lower:
-        draw.ellipse((170, 240, 340, 380), outline="white", width=6)
-        draw.ellipse((120, 260, 210, 330), outline="white", width=6)
-        draw.ellipse((130, 240, 160, 270), outline="white", width=6)
-        draw.ellipse((180, 240, 210, 270), outline="white", width=6)
-        draw.line((340, 320, 450, 260), fill="white", width=6)
-        draw.ellipse((155, 285, 165, 295), fill="white")
-        draw.ellipse((185, 285, 195, 295), fill="white")
+        draw_cat(draw)
+    elif "dog" in prompt_lower:
+        draw_dog(draw)
+    elif "lion" in prompt_lower:
+        draw_lion(draw)
+    elif "bird" in prompt_lower:
+        draw_bird(draw)
+    elif "fish" in prompt_lower:
+        draw_fish(draw)
+    elif "turtle" in prompt_lower:
+        draw_turtle(draw)
     else:
-        draw.text((50, 250), "I can draw only cats & rats", fill="white")
+        draw.text((30, 250), "I can draw: cat, rat, dog, lion, bird, fish, turtle", fill="white")
 
     draw.rectangle([0, 470, 512, 512], fill=(0, 0, 0))
-    draw.text((10, 472), f"Prompt: {prompt[:30]}", fill="white")
+    draw.text((10, 472), f"Prompt: {prompt[:40]}", fill="white")
 
     return img