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

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

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

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

face2paint = torch.hub.load(
    'AK391/animegan2-pytorch:main', 
    'face2paint', 
    size = 512, 
    device = DEVICE,
    #trust_repo = True,
    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,
#             inputs=gr.inputs.Image(type="pil"),
#             outputs=gr.outputs.Image(type="pil"),
#             title=title,
# description=description,
#             article=article,
#             examples=examples,
#             allow_flagging= 'auto',
#             allow_screenshot=False
#).launch(enable_queue=True,cache_examples=False)


#interface = 
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)

"""
        fn = inference,
        inputs=[
            gr.inputs.Image(label="Input Image", type="pil"),
            "text"
        ],
        outputs=[
            gr.outputs.Label(label="Class"),
            "text",
            gr.outputs.Image(label="Output Face with Explainability", type="pil")
        ],
        title = title,
        description = description,
        article = article,
        examples = examples,
        allow_flagging = 'auto',
        allow_screenshot = False
).launch(enable_queue=True,cache_examples=True)
"""