# coding: utf-8
import os
#os.system("pip install gradio==2.8.0b3")
import torch
import gradio as gr
from PIL import ImageFont, ImageDraw

device = "cuda" if torch.cuda.is_available() else "cpu"
model1 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", device=device, pretrained="face_paint_512_v1")
model2 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", device=device, pretrained="face_paint_512_v2")
face2paint = torch.hub.load("bryandlee/animegan2-pytorch:main", "face2paint", device=device)


def inference(img, ver):
    if ver == "Version 2 (Réaliste)":
        out = face2paint(model2, img)
    else:
        out = face2paint(model1, img)
    return out


title = "Temp Late"
gr.Interface(
    inference,
    [
        gr.inputs.Image(type="pil", source="upload"),
        gr.inputs.Radio(
            ["Version 1 (Stylisé)", "Version 2 (Réaliste)"],
            type="value",
            default="Version 2 (Réaliste)",
            label="Version",
        ),
    ],
    gr.outputs.Image(type="pil"),
    allow_flagging="auto",
    allow_screenshot=False,
    flagging_dir="flagged",
).launch(enable_queue=True)