import gradio as gr

from app_integration import generate_caption


def gradio_generate_caption(image, model_variant, use_cache):
    result = generate_caption(
        image=image,
        model_variant=model_variant,
        use_cache=use_cache,
    )
    return (
        result["caption"],
        result["total_time"],
        result["preprocess_time"],
        result["model_load_time"],
        result["inference_time"],
        result["postprocess_time"],
        result["approx_gpu_time"],
        str(result["cache_hit"]),
        result["model"],
    )


demo = gr.Interface(
    fn=gradio_generate_caption,
    inputs=[
        gr.Image(type="pil", label="Upload image"),
        gr.Dropdown(choices=["base", "large"], value="base", label="Model variant"),
        gr.Checkbox(value=True, label="Use cache"),
    ],
    outputs=[
        gr.Textbox(label="Generated caption"),
        gr.Number(label="Total time"),
        gr.Number(label="Preprocessing time"),
        gr.Number(label="Model load time"),
        gr.Number(label="Inference time"),
        gr.Number(label="Postprocessing time"),
        gr.Number(label="Approx. GPU time"),
        gr.Textbox(label="Cache hit"),
        gr.Textbox(label="Model used"),
    ],
    title="ZeroGPU Image Captioning",
    description="Upload an image, choose a model variant, and view captioning performance metrics.",
)

demo.launch()