Update app.py

app.py

@@ -1,54 +1,34 @@
 import gradio as gr
-from PIL import Image, ImageDraw, ImageFont
-import random
-import textwrap
+from diffusers import StableDiffusionPipeline
+import torch
 
-# Enhanced function for image generation with better visual representation
+# Load the Stable Diffusion model (use a lightweight version for Spaces)
+model_id = "runwayml/stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe = pipe.to("cuda")  # Use GPU if available
+
+# Function to generate image from description
 def generate_image(text_description):
-    # Create a blank image with a gradient background to simulate a sunset
-    img = Image.new('RGB', (512, 512), color=(0, 0, 0))  # Start with black
-    draw = ImageDraw.Draw(img)
-    
-    # Create a simple gradient for sunset (orange to purple)
-    for y in range(512):
-        r = int(255 * (1 - y / 512) * 0.8)  # Fade red/orange
-        g = int(100 * (1 - y / 512))        # Minimal green
-        b = int(150 + 105 * (y / 512))      # Increase blue/purple
-        draw.line([(0, y), (512, y)], fill=(r, g, b))
+    # Generate image using Stable Diffusion
+    image = pipe(text_description, num_inference_steps=50, guidance_scale=7.5).images[0]
     
-    # Load a default font (fallback to system font if arial.ttf is unavailable)
+    # Optional: Add text to the image
+    from PIL import Image, ImageDraw, ImageFont
+    img_with_text = image.copy()
+    draw = ImageDraw.Draw(img_with_text)
     try:
-        font = ImageFont.truetype("arial.ttf", 25)
+        font = ImageFont.truetype("arial.ttf", 20)
     except:
         font = ImageFont.load_default()
+    text = f"Generated: {text_description}"
+    draw.text((10, image.height - 100), text, font=font, fill=(255, 255, 255))
     
-    # Wrap text to fit within the image
-    wrapper = textwrap.TextWrapper(width=20)  # Adjust width for better wrapping
-    wrapped_text = wrapper.fill(text=text_description)
-    lines = wrapped_text.split('\n')
-    
-    # Draw each line of text at the bottom
-    y_text = 400
-    for line in lines:
-        bbox = draw.textbbox((0, 0), line, font=font)
-        line_width = bbox[2] - bbox[0]
-        line_height = bbox[3] - bbox[1]
-        draw.text((20, y_text), line, font=font, fill=(255, 255, 255))
-        y_text += line_height + 5
-    
-    # Add graphical elements based on description
-    if "lake" in text_description.lower():
-        draw.rectangle([150, 300, 350, 400], fill=(0, 191, 255))  # Blue lake
-    if "mountains" in text_description.lower():
-        draw.polygon([(200, 250), (250, 200), (300, 250)], fill=(255, 255, 255))  # White mountain
-        draw.polygon([(250, 200), (300, 150), (350, 200)], fill=(255, 255, 255))  # Second mountain
-    
-    return img
+    return img_with_text
 
 # Gradio interface
 with gr.Blocks(title="Text-to-Image Generator") as demo:
     gr.Markdown("# Text-to-Image Generator")
-    gr.Markdown("Enter a description below and generate an image!")
+    gr.Markdown("Enter a description below and generate a detailed image!")
     
     with gr.Row():
         with gr.Column():