import os

import gradio as gr
import spaces

import core

engine = core.load()
# @spaces.GPU runs in a separate process, so a global flag would not survive; a file does
WARM = "/tmp/vev-warm"


def duration(*args):
    return 15 if os.path.exists(WARM) else 60   # the first request compiles the attention kernels


@spaces.GPU(duration=duration)
def ask(image, context, *qs):
    rows = list(zip(qs[0::2], qs[1::2]))
    try:
        out, ms = core.ask(engine, image, context, rows)
    except ValueError as e:
        raise gr.Error(str(e))
    open(WARM, "w").close()
    labels = [gr.Label(value=o, visible=o is not None) for o in out]
    return *labels, f"{len([o for o in out if o])} question(s), one request, {ms:.0f} ms on the GPU."


with gr.Blocks(title="Vev") as demo:
    gr.Markdown(
        "# Vev\n"
        "Ask yes/no or multiple-choice questions about an image and get a probability for every answer. "
        "Nothing is generated, so the answer is always one of your options.\n\n"
        "Model: [CountingSheep/vev-4b](https://huggingface.co/CountingSheep/vev-4b) · "
        "Code and API: [Xiaooolong/vev](https://github.com/Xiaooolong/vev)\n\n"
        "The first request after the Space starts takes about half a minute while GPU kernels compile; later ones take a few seconds.")
    with gr.Row():
        with gr.Column():
            image = gr.Image(type="pil", label="Image")
            context = gr.Textbox(label="Text (optional)", placeholder="Anything that should go into the state with the image")
            qs = []
            for i in range(core.N_QUESTIONS):
                with gr.Row():
                    qs.append(gr.Textbox(label=f"Question {i + 1}", scale=3))
                    qs.append(gr.Textbox(label="Options", placeholder="a; b; c  (empty = yes/no)", scale=2))
            button = gr.Button("Ask", variant="primary")
        with gr.Column():
            outputs = [gr.Label(label=f"Question {i + 1}", visible=False) for i in range(core.N_QUESTIONS)]
            note = gr.Markdown()
    gr.Examples([core.EXAMPLE], inputs=[image, context, *qs])
    button.click(ask, inputs=[image, context, *qs], outputs=[*outputs, note])

demo.launch()