import tempfile
import gradio as gr
from main import load_artifacts_and_helpers, device

# Load model and helpers once when the Space starts
final_pipeline, predict_fn, vocab, idx_to_answer, model, encode_fn = \
    load_artifacts_and_helpers(prefix="vqa_custom", map_location=device)

def vqa_interface(image, question):
    # image is a PIL image from Gradio
    with tempfile.NamedTemporaryFile(suffix=".jpg") as f:
        image.save(f.name)
        answer = final_pipeline(
            f.name,
            question,
            open_vqa_fn=None,     # BLIP disabled in Space for speed
            translate_fn=None
        )
    return answer

demo = gr.Interface(
    fn=vqa_interface,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Textbox(lines=1, label="Question", value="What is in the image?")
    ],
    outputs=gr.Textbox(label="Answer"),
    title="VQA-RAD Demo",
    description="Custom ResNet18 + LSTM VQA model (top-50 answers)."
)

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