"""Gradio app for Hugging Face Spaces: upload an image, get a generated caption."""
from __future__ import annotations

import spaces
import gradio as gr
from huggingface_hub import hf_hub_download

from src.inference.predict import Predictor

MODEL_REPO_ID = "AdhamAshraf/image-caption-generator"

print("Downloading model artifacts from Hugging Face Hub...")
checkpoint_path = hf_hub_download(
    repo_id=MODEL_REPO_ID,
    filename="resnet_attention_lstm/best_model.pt",
)
vocab_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename="vocab.json")

predictor = Predictor(checkpoint_path=checkpoint_path, vocab_path=vocab_path, device="cpu")
print("Model loaded.")


@spaces.GPU
def generate_caption(image):
    if image is None:
        return "Please upload an image."
    return predictor.predict(image)


demo = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=gr.Textbox(label="Generated Caption"),
    title="🖼️ Image Caption Generator",
    description="ResNet50 (frozen) + LSTM baseline, trained on Flickr8k.",
)

if __name__ == "__main__":
    demo.launch()