import streamlit as st
from diffusers import AutoPipelineForText2Image
import torch

# Load model and move it to CUDA
pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
pipe.to("cpu")

# Cache the image generation function
@st.cache(suppress_st_warning=True)
def generate_image(prompt):
    # Generate image using the model
    image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
    return image

# Streamlit app
st.title("Text to Image Generation App")

# User input prompt
prompt = st.text_area("Enter a prompt for image generation:")

if st.button("Generate Image"):
    if prompt:
        # Generate and display the image
        st.image(generate_image(prompt), caption="Generated Image", use_column_width=True)
    else:
        st.warning("Please enter a prompt.")