import spaces
import gradio as gr
import torch
from PIL import Image
from utils import WaifuScorer

SCORER = None

@spaces.GPU(duration=30)
def score_image(image: Image.Image) -> float:
    global SCORER
    if SCORER is None:
        SCORER = WaifuScorer(
            device='cuda' if torch.cuda.is_available() else 'cpu',
            verbose=True,
        )
    return SCORER([image])[0]


demo = gr.Interface(
    fn=score_image,
    inputs=gr.Image(type='pil', label='Image', height=512),
    outputs=gr.Number(label='Score', precision=2),
    title='Waifu Scorer V3',
    description='''Score ranges from 0 to 10, higher is better\n
[Github](https://github.com/Eugeoter/waifu-scorer) | [Model](https://huggingface.co/Eugeoter/waifu-scorer-v3) | [Inspiration](https://github.com/christophschuhmann/improved-aesthetic-predictor)''',
    examples=[
        'examples/example_0.jpg',
        'examples/example_1.jpg',
        'examples/example_2.jpg',
        'examples/example_3.jpg',
        'examples/example_4.jpg',
    ],
    cache_examples=True,
)
demo.queue().launch()