import gradio as gr
import functools
from inference import CaptionGenerator


@functools.lru_cache(maxsize=1)
def load_model():
    return CaptionGenerator(
        model_path="best_model.pth",
        vocab_path="vocab.pkl",
        use_vit=True
    )


generator = load_model()



def caption_image(img):
    return generator.generate(img)


demo = gr.Interface(
    fn=caption_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(label="Generated Caption"),
    title="Image Captioning with Transformers",
    description="Upload an image and get an AI-generated caption."
)

demo.launch()