import gradio as gr
from modules.predict import WaifuScorer


def ui():
    scorer = WaifuScorer()

    def predict(img):
        return scorer(img)

    interface = gr.Interface(title='Waifu Scorer', description='A model that scores an anime illustration (0 ~ 10).',
                             fn=predict, inputs=gr.Image(sources='upload', type='pil', height=512), outputs=gr.Number(precision=3),
                             allow_flagging='never', examples=[['./examples/1.png'], ['./examples/2.png']]
                             )

    return interface


if __name__ == '__main__':
    ui().launch()