# 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):
    title_font = ImageFont.truetype("LEMONMILK-Regular.otf", 20)
    title_text = "@_temp.late_"
    padding = 10
    x, y = 15, 15  # 180, 412
    w, h = title_font.getsize(title_text)
    if ver == "Version 2 (🔺 Réaliste,🔻 Stylé)":
        out = face2paint(model2, img)
    else:
        out = face2paint(model1, img)
    image_editable = ImageDraw.Draw(out)
    image_editable.rectangle((x, y, x + w + padding, y + h + padding), fill="white")
    image_editable.text((x + padding / 2, y + padding / 2), title_text, (0, 0, 0), font=title_font)
    return out


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