# app.py

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

# Set up the Streamlit app
st.set_page_config(page_title="Huggingface Diffusers Showcase", page_icon=":guardsman:", layout="wide")

# Add a sidebar
st.sidebar.title("Navigation")
st.sidebar.markdown("## Choose a Model")
model_options = ["stable-diffusion", "ddim", "ddpm"]
selected_model = st.sidebar.selectbox("Select a model", model_options)

# Load the selected model
if selected_model == "stable-diffusion":
    pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-512-v2")
elif selected_model == "ddim":
    pipe = DiffusionPipeline.from_pretrained("CompVis/ddim-512-fp16")
else:
    pipe = DiffusionPipeline.from_pretrained("CompVis/ddpm-512-fp16")

# Add a text input for the user to enter a prompt
prompt = st.text_input("Enter a text prompt:", "A sunset over a mountain")

# Add a button to generate the image
if st.button("Generate Image"):
    # Generate the image
    image = pipe(prompt, num_inference_steps=50).images[0]

    # Display the generated image
    col1, col2 = st.columns(2)
    with col1:
        st.subheader("Prompt")
        st.markdown(prompt)
    with col2:
        st.subheader("Generated Image")
        st.image(image, caption="Generated using the selected diffusion model.", use_column_width=True)

# Add a section for the app info
st.sidebar.markdown("---")
st.sidebar.markdown("## App Info")
st.sidebar.markdown("This app showcases the Huggingface Diffusers library and its features.")
st.sidebar.markdown("You can select a model from the sidebar and enter a text prompt to generate an image.")
st.sidebar.markdown("Visit the [Huggingface Diffusers documentation](https://huggingface.co/docs/diffusers) for more information.")