import streamlit as st
from diffusers import DiffusionPipeline
import torch

# Load the model
@st.cache(allow_output_mutation=True)
def load_model():
    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
    pipe.to("cuda")
    return pipe

# Streamlit app
def main():
    st.title("Image Generation with Diffusion Models")
    prompt = st.text_input("Enter a prompt:", "An astronaut riding a green horse")
    if st.button("Generate"):
        pipe = load_model()
        images = pipe(prompt=prompt).images[0]
        st.image(images, caption=prompt, use_column_width=True)

if __name__ == "__main__":
    main()