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.float32  # Use float32 for CPU support
    )
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipeline.to(device)
    return pipeline

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

    # 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  # ✅ Fixed indentation

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

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

if __name__ == "__main__":
    main()