import streamlit as st
import torch
from diffusers import StableDiffusionPipeline

@st.cache_resource
def load_model():
    model_id = "CompVis/stable-diffusion-v1-4"  # Correct model ID

    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float32  # Ensure CPU compatibility
    )

    pipe.to("cpu")  # Force CPU execution
    return pipe

st.title("Text-to-Image Generator")
prompt = st.text_input("Enter a prompt:")

if prompt:
    st.write("Generating image... Please wait.")
    model = load_model()
    with torch.no_grad():
        image = model(prompt).images[0]
    st.image(image, caption="Generated Image", use_column_width=True)