#
# SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
# SPDX-License-Identifier: Apache-2.0
#

import os
import openai
import base64
from io import BytesIO
from PIL import Image
from config import MODEL, SIZE, DESCRIPTION, EXAMPLES
import gradio as gr

def playground(model, size, prompt):
    client = openai.OpenAI(
        base_url=os.getenv("OPENAI_API_BASE_URL"),
        api_key=os.getenv("OPENAI_API_KEY")
    )
    
    response = client.images.generate(
        model=model,
        prompt=prompt,
        size=size,
        n=1,
        response_format="b64_json"
    )
    
    return Image.open(
        BytesIO(
            base64.b64decode(
                response.data[0].b64_json
            )
        )
    )

with gr.Blocks() as app:
    with gr.Sidebar():
        gr.HTML(DESCRIPTION)
    gr.Interface(
        fn=playground,
        inputs=[
            gr.Dropdown(MODEL, label="MODEL"),
            gr.Dropdown(SIZE, label="IMAGE SIZE"),
            gr.Textbox(
                label="INSTRUCTIONS",
                placeholder="Insert your prompt here..."
            )
        ],
        outputs=gr.Image(
            label="GENERATED IMAGE", type="pil"
        ),
        examples=EXAMPLES,
        show_api=False,
        concurrency_limit=3
    )

app.launch(
    server_name="0.0.0.0",
    pwa=True
)