import gradio as gr
import openai
import urllib.request
from PIL import Image
import os
import nltk
#nltk.download('punkt')

def generate_image(api_key, prompt, resolution):
    if not api_key:
        print("Error: API Key is required.")
        return
    openai.api_key = api_key
    response = openai.Image.create(
        prompt=prompt,
        n=1,
        size=resolution
    )

    image_url = response['data'][0]['url']

    # Open the URL image, resize it to the chosen resolution and return it
    with urllib.request.urlopen(image_url) as url:
        with open('temp.jpg', 'wb') as f:
            f.write(url.read())
        img = Image.open('temp.jpg')

    return img

iface = gr.Interface(
    fn=generate_image, 
    inputs=[
        gr.inputs.Textbox(lines=1, label="API Key", type="password"), 
        gr.inputs.Textbox(lines=1, label="Prompt"),
        gr.inputs.Radio(choices=["256x256", "512x512", "1024x1024"], label="Resolution")
    ], 
    outputs=gr.outputs.Image(type="pil"), 
    title="DALL-E Image Generator",
    description="Enter your API key, a prompt, and choose a resolution to generate an image from DALL-E."
)
    

iface.launch()