from PIL import Image
import torch
import gradio as gr

model2 = torch.hub.load(
    "AK391/animegan2-pytorch:main",
    "generator",
    pretrained=True,
    device="cpu",
    progress=False
)

model1 = torch.hub.load(
    "AK391/animegan2-pytorch:main",
    "generator",
    pretrained="face_paint_512_v1",
    device="cpu"
)

face2paint = torch.hub.load(
    'AK391/animegan2-pytorch:main',
    'face2paint',
    size=512,
    device="cpu",
    side_by_side=False
)

def inference(img, ver):
    if ver == 'versão 2 (🔺 robustez, 🔻 estilização)':
        out = face2paint(model2, img)
    else:
        out = face2paint(model1, img)
    return out

title = "AnimeGANv2"
description = "Demonstração do AnimeGanv2 para retratos de rostos. Para usá-lo, simplesmente faça o upload da sua imagem."
article = "<p style='text-align: center'><a href='https://github.com/bryandlee/animegan2-pytorch' target='_blank'>Repositório do Github Pytorch</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=akhaliq_animegan' alt='insígnia de visitantes'></center></p>"

gr.Interface(
    fn=inference,
    inputs=[
        gr.inputs.Image(type="pil"),
        gr.inputs.Radio(
            ['versão 1 (🔺 estilização, 🔻 robustez)', 'versão 2 (🔺 robustez, 🔻 estilização)'],
            type="value",
            default='versão 2 (🔺 robustez, 🔻 estilização)',
            label='versão'
        )
    ],
    outputs=gr.outputs.Image(type="pil"),
    title=title,
    description=description,
    article=article,
    allow_flagging=False,
    allow_screenshot=False
).launch()