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

# Function to load the model
@st.cache(allow_output_mutation=True)
def load_model():
    # Replace 'your_token_here' with your actual Hugging Face token
    YOUR_TOKEN = "your_token_here"
    pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
    return pipe

# Load the model
pipe = load_model()

# Streamlit interface
st.title("Text-to-Image Generation with Stable Diffusion")

# Text input for the prompt
prompt = st.text_input("Enter your prompt:")

# Generate button
if st.button("Generate Image"):
    with st.spinner('Generating image...'):
        # Generate image
        image = pipe(prompt).images[0]

        # Convert to PIL Image to display
        img = Image.fromarray(image)
        st.image(img, caption="Generated Image")