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

# Optional: Log in to Hugging Face (if using private models)
# login(token="your_huggingface_token")

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

# Load the diffusion pipeline (make sure you have access to the model)
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float32)
pipe.to("cpu")  # Ensure it's using CPU

# 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...'):
        try:
            # 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!")
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")