app.py

import streamlit as st
from diffusers import DiffusionPipeline
import torch
from PIL import Image

# Title and description of the app
st.title("🖼️ Stable Diffusion Image Generator")
st.write("Generate images from text using the Stable Diffusion v1.5 model!")

# Sidebar for user inputs
st.sidebar.title("Input Options")
prompt = st.sidebar.text_input("Enter your prompt", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k")
generate_button = st.sidebar.button("Generate Image")

# Load the pipeline when the app starts
@st.cache_resource
def load_pipeline():
    pipe = DiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5", 
        torch_dtype=torch.float16
    )
    device = "cuda" if torch.cuda.is_available() else "cpu"
    return pipe.to(device)

pipe = load_pipeline()

# Generate image when button is clicked
if generate_button:
    st.write(f"### Prompt: {prompt}")
    with st.spinner("Generating image... Please wait."):
        # Generate the image
        image = pipe(prompt).images[0]
        
        # Display the generated image
        st.image(image, caption="Generated Image", use_column_width=True)
        
        # Option to download the image
        img_path = "generated_image.png"
        image.save(img_path)
        with open(img_path, "rb") as img_file:
            st.download_button(
                label="Download Image",
                data=img_file,
                file_name="generated_image.png",
                mime="image/png"
            )