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

# Streamlit app title
st.title("Stable Diffusion Image Generator")

# Load the diffusion pipeline
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float32)
pipe.to("cpu")  # Use CPU for inference

# Get user input (prompt)
prompt = st.text_input("What do you want to see?", "A beautiful landscape")

# Button to generate image
if st.button('Generate Image'):
    with st.spinner('Generating...'):
        # Generate image
        image = pipe(prompt).images[0]
        
        # Display image in Streamlit
        st.image(image, caption="Generated Image", use_column_width=True)

    st.success("Image generated successfully!")