import streamlit as st
from transformers import pipeline
from diffusers import StableDiffusionPipeline
import torch

# Load the pre-trained Stable Diffusion model
@st.cache_resource
def load_model():
    model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v-1-4-original", torch_dtype=torch.float16)
    model.to("cuda")
    return model

# Initialize model
model = load_model()

# Streamlit Interface
st.title("Text to Image Generator using Stable Diffusion")
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", "A beautiful landscape with mountains and lakes.")

if st.button("Generate Image"):
    if user_input:
        with st.spinner("Generating image..."):
            # Generate image based on user input
            generated_image = model(user_input).images[0]
            st.image(generated_image)
    else:
        st.error("Please enter a text prompt!")