import streamlit as st
import openai

st.title("DALL-E2 API Image Generation Demo")
openai.api_key = st.text_input("Enter Your OpenAI Key", "")

st.write("Enter a prompt to generate an image")

prompt = st.text_area("Prompt", "An old man wearing a t-shirt with the word 'SMILE' printed.")

num_images = st.slider("Number of images to generate", min_value=1, max_value=10, value=1)

image_size = st.selectbox(
    "Select an image size",
    ["256x256", "512x512", "1024x1024"]
)

def generate_images(prompt, num_images, image_size):
    response = openai.Image.create(
        prompt=prompt,
        n=num_images,
        size=image_size,
    )
    return response["data"]

if st.button("Generate Images"):
    with st.spinner("Generating images..."):
        image_data = generate_images(prompt, num_images, image_size)
        for idx, image in enumerate(image_data):
            st.image(image['url'], caption=f"Image {idx+1}", width=400)