app.py

import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load model (smaller for faster loading)
model = StableDiffusionPipeline.from_pretrained(
    "OFA-Sys/small-stable-diffusion-v0",  # Lightweight model
    torch_dtype=torch.float32
)

def generate_image(prompt):
    return model(prompt).images[0]

# Create Gradio interface
demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Describe your image"),
    outputs=gr.Image(label="Generated Image"),
    title="🎨 AI Image Generator",
    examples=[["a cat astronaut"], ["colorful jellyfish in space"]]
)

demo.launch()