import os
os.system("pip install torch==1.11.0+cpu torchvision==0.12.0+cpu -f https://download.pytorch.org/whl/cpu/torch_stable.html Pillow")
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 version2(img):
    out = face2paint(model2, img)
    return out
    
def version1(img):
    out = face2paint(model1, img)
    return out
  
block = gr.Blocks()

with block:
    gr.Markdown("Gradio Demo for AnimeGanv2 Face Portrait. To use it, simply upload your image. Try out the different versions below. Please use a cropped portrait picture for best results.")
    
    with gr.Tabs():
        with gr.TabItem("version 2 (🔺 robustness,🔻 stylization)"):
            with gr.Row():
                facepaint2 = gr.inputs.Image(type="pil",shape=(512,512))
                faceout2 = gr.outputs.Image(type="pil")
            face_run = gr.Button("Run")
            face_run.click(version2, inputs=facepaint2, outputs=faceout2)

    with gr.Tabs():
        with gr.TabItem("version 1 (🔺 stylization, 🔻 robustness)"):
            with gr.Row():
                facepaint1 = gr.inputs.Image(type="pil")
                faceout1 = gr.outputs.Image(type="pil")
            face_run = gr.Button("Run")
            face_run.click(version1, inputs=facepaint1, outputs=faceout1)
    
block.launch()