import gradio as gr
import openai
import matplotlib.pyplot as plt
import io
import numpy as np
import base64
from PIL import Image

# Convert image to base64 for display in gradio
def get_image_data(plt):
    buf = io.BytesIO()
    plt.savefig(buf, format='PNG')
    buf.seek(0)
    img = Image.open(buf)
    return img

# Execute Python code and generate images
def execute_code(code):
    exec(code)
    return get_image_data(plt)

def gpt_inference(prompt, model, openai_key):
    openai.api_key = openai_key
    response = openai.Completion.create(
      engine=model, 
      prompt=prompt, 
      max_tokens=100
    )
    code = response.choices[0].text.strip()
    img = execute_code(code)
    return img

iface = gr.Interface(
    fn=gpt_inference, 
    inputs=["text", gr.inputs.Dropdown(choices=["gpt3.5-turbo", "gpt4"], label="Model"), "text"], 
    outputs=gr.outputs.Image(type="pil"),
    input_labels=["Prompt", "Model", "OpenAI Key"]
)
iface.launch()