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():
    # You can load any Stable Diffusion model here
    model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v-1-4-original")
    model.to("cpu")  # Ensure it uses CPU for inference
    return model

# Initialize Streamlit components
st.title("Text-to-Image Generator")

st.write("Enter a description and generate an image!")

# User input for text prompt
prompt = st.text_input("Enter your text prompt here:")

if prompt:
    st.write("Generating image... Please wait.")

    # Load model once (from the Hugging Face model hub)
    model = load_model()

    # Generate the image from the text prompt
    with torch.no_grad():
        image = model(prompt).images[0]

    # Show the generated image in the app
    st.image(image, caption="Generated Image", use_column_width=True)