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

# Streamlit App Title
st.title("🖼️ Stable Diffusion Image Generator")
st.write("Generate images from text prompts using Stable Diffusion!")

# Sidebar for prompt input and button
st.sidebar.title("Image Generation Settings")
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")

# Cache the Stable Diffusion model to load it only once
@st.cache_resource
def load_model():
    model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
    model = model.to("cuda" if torch.cuda.is_available() else "cpu")
    return model

# Load the model
pipe = load_model()

# Generate the image when the button is clicked
if generate_button:
    if prompt.strip() == "":
        st.error("Prompt cannot be empty. Please enter a valid text prompt.")
    else:
        st.write(f"### Prompt: {prompt}")
        with st.spinner("Generating image... Please wait."):
            try:
                # Generate the image
                image = pipe(prompt).images[0]

                # Display the image
                st.image(image, caption="Generated Image", use_column_width=True)

                # Provide download option
                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"
                    )
            except Exception as e:
                st.error(f"An error occurred: {e}")