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

# Load the Stable Diffusion pipeline
@st.cache(allow_output_mutation=True)
def load_pipeline():
    pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
    pipeline = pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
    return pipeline

def main():
    st.title("Stable Diffusion Image Generator")
    st.write("Generate images from text prompts using Stable Diffusion 2.1")

    # Initialize the pipeline
    pipeline = load_pipeline()

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

    # Generate button
    if st.button("Generate"):
        if not prompt:
            st.warning("Please enter a prompt first.")
            return

        st.write("Generating your image...")
        with torch.no_grad():
            image = pipeline(prompt).images[0]

        st.write("Generated Image:")
        st.image(image, use_column_width=True)

if __name__ == "__main__":
    main()