import os
os.system("pip install gradio==2.8.0b3")
from PIL import Image
import torch
import gradio as gr



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


model1 = torch.hub.load("AK391/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v1",  device="cuda")
face2paint = torch.hub.load(
    'AK391/animegan2-pytorch:main', 'face2paint', 
    size=512, device="cuda",side_by_side=False
)
def inference(img, ver):
    if ver == 'version 2 (🔺 robustness,🔻 stylization)':
        out = face2paint(model2, img)
    else:
        out = face2paint(model1, img)
    return out
  
title = "Deepfake Detection"
description = "This gradio contains a GAN-generated image detector developed to distinguish real images from synthetic ones."
article = "<p style='text-align: center'><a href='https://github.com/polimi-ispl/GAN-image-detection' target='_blank'>@polimi-ispl/GAN-image-detection</a></p>"
examples=[
    ['images/fake0.jpg','Fake Female'],
    ['images/fake1.png','Fake Male'],
    ['images/real0.jpg','Real Female'],
    ['images/real1.jpg','Real Male'],

    gr.Interface(inference, 
        gr.inputs.Image(type="pil"), 
        gr.outputs.Image(type="pil"),
        title=title,
        description=description,
        article=article,
        examples=examples,
        allow_flagging=False,
        allow_screenshot=False
).launch(enable_queue=True,cache_examples=True)