import torch
import os

import gradio as gr
from huggingface_hub import hf_hub_download

from DummyModel import DummyModel


file_path = hf_hub_download("lfolle/DeepNAPSIModel", "dummy_model.pth",
                            use_auth_token=os.environ['DeepNAPSIModel'])
model = DummyModel().load_state_dict(torch.load(file_path))

def predict(data):
    return model(data)

iface = gr.Interface(
    predict,
    title="DeepNAPSI Application",
    inputs=gr.Image(),
    outputs=gr.Number(label="DeepNAPSI prediction"),
    description=""
)
iface.launch()