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("cuda")  # Move the model to GPU for faster processing

# Create a Streamlit app
st.title("Image Generation App")

# Add a text input for the prompt
prompt = st.text_input("Prompt", "a photo of an astronaut riding a horse on mars")

# Generate the image based on the prompt
if st.button("Generate Image"):
    # 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)