app.py

import streamlit as st
from diffusers import StableDiffusionPipeline
import torch

# UI loads first
st.title("Stable Diffusion Image Generator")
st.write("Enter a description, and AI will generate an image for you!")

# Load model in the background
@st.cache_resource
def load_model():
    st.write("Loading the model... Please wait.")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
    return pipe

pipe = load_model()

# User input
prompt = st.text_input("Enter your text prompt:")

# Generate button
if st.button("Generate Image"):
    if prompt:
        with st.spinner("Generating image... Please wait."):
            image = pipe(prompt).images[0]  # Generate image
            st.image(image, caption="Generated Image", use_column_width=True)
    else:
        st.error("Please enter a prompt to generate an image.")