# Import necessary libraries
import streamlit as st
import torch
from diffusers import StableDiffusionXLPipeline

# Load the stable diffusion model
pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
)
pipe = pipe.to("cpu")  # Move the model to CPU for processing

# Create a Streamlit app
st.title("Stable Diffusion XL Image Generation")

# Add a text input for the prompt
prompt = st.text_area("Enter your prompt here", "Type your prompt here...")

# Add an info text to provide guidance
st.markdown("This app uses Stable Diffusion XL to generate an image based on the given prompt.")

# Generate the image based on the prompt
if st.button("Generate"):
    with st.spinner("Generating..."):  # Display a spinner while generating the image
        try:
            # Use the model to generate the image
            image = pipe(prompt).images[0]

            # Display the generated image
            st.image(image, caption="Generated Image", use_column_width=True)
        except Exception as e:
            st.error(f"An error occurred: {e}")