import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load the Stable Diffusion model (use CPU-compatible settings)
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)  # Use float32 for CPU

# Function to generate image from description
def generate_image(text_description):
    # Generate image using Stable Diffusion on CPU
    image = pipe(text_description, num_inference_steps=25, guidance_scale=7.5).images[0]  # Reduced steps for speed
    
    # 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", 20)
    except:
        font = ImageFont.load_default()
    text = f"Generated: {text_description}"
    draw.text((10, image.height - 100), text, font=font, fill=(255, 255, 255))
    
    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 a detailed image! (Note: Running on CPU may be slow)")
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(label="Description", placeholder="Type your image description here...", value="A serene lake surrounded by snow-capped mountains under a vibrant sunset sky with shades of orange and purple.")
            generate_btn = gr.Button("Generate Image")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
    
    # Connect the button to the function
    generate_btn.click(fn=generate_image, inputs=text_input, outputs=output_image)

# Launch the app
demo.launch()