import streamlit as st
from diffusers import DiffusionPipeline
import torch

# Load the pre-trained model
@st.cache_resource
def load_model():
    # Load the Stable Diffusion model from Hugging Face
    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
    pipe.to("cuda")  # Make sure to move the model to GPU if available
    return pipe

# Initialize the model
pipe = load_model()

# Streamlit Interface
st.title("Text to Image Generator using Stable Diffusion 2.1")
st.write("Enter a description below, and the model will generate an image based on it.")

# Text input field for user to enter prompt
user_input = st.text_area("Enter the text prompt", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k")

if st.button("Generate Image"):
    if user_input:
        with st.spinner("Generating image..."):
            # Generate image based on user input
            generated_image = pipe(user_input).images