import streamlit as st
import openai
import requests
from io import BytesIO

# Streamlit UI for input
st.title("Generate an Image with DALL-E 3")

# Input for OpenAI API key
api_key = st.text_input("Enter your OpenAI API Key:", type="password")

# Check if API key is provided
if api_key:
    # Set the API key directly in the OpenAI library
    openai.api_key = api_key

    # Prompt input field
    prompt = st.text_area("Enter your prompt:")

    # Variables to store the generated image
    image_url = None
    image_data = None

    # Button to generate image
    if st.button("Generate Image"):
        if prompt:
            st.write("Generating image...")
            try:
                # Call the OpenAI API to generate the image
                response = openai.Image.create(
                    prompt=prompt,
                    n=1,  # Number of images to generate
                    size="1024x1024"  # Available sizes: 256x256, 512x512, or 1024x1024
                )
                image_url = response['data'][0]['url']

                # Download the image
                image_response = requests.get(image_url)
                image_data = BytesIO(image_response.content)

                # Display the image
                st.image(image_url, caption="Generated Image", use_column_width=True)
                st.write("Image URL: " + image_url)
            except Exception as e:
                st.error(f"An error occurred: {e}")
        else:
            st.warning("Please enter a prompt.")

    # Button to download the image
    if image_data:
        st.download_button(
            label="Download Image",
            data=image_data,
            file_name="generated_image.png",
            mime="image/png"
        )
else:
    st.warning("Please enter your OpenAI API Key.")