import gradio as gr
import torch
from torchvision import transforms
from archs import CycleGenerator
import warnings

warnings.filterwarnings("ignore")

# instantiate generators
G_XtoY = CycleGenerator(conv_dim=64) # Apple -> Windows
G_YtoX = CycleGenerator(conv_dim=64) # Windows -> Apple

# load weights (on CPU because Huggingface does not provide free GPU computing)
device = torch.device('cpu')
G_XtoY.load_state_dict(torch.load('G_XtoY.pth', map_location=device)); G_XtoY.eval()
G_YtoX.load_state_dict(torch.load('G_YtoX.pth', map_location=device)); G_YtoX.eval()

def generate(input_image, radio):
    transform = transforms.Compose([
                    transforms.Resize(32),
                    transforms.ToTensor(),
                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                ])
    input_image = transform(input_image).unsqueeze(0)
    with torch.no_grad():
        if radio == 'Apple to Windows':
            out = G_XtoY(input_image).squeeze().numpy()
        else:
            out = G_YtoX(input_image).squeeze().numpy()
    return out.transpose(1, 2, 0)


input_image = gr.inputs.Image(source="upload", type="pil", label="Input Image")
radio = gr.inputs.Radio(
    choices=['Apple to Windows', 'Windows to Apple'],
    label="Choose Conversion"
)

output_image = gr.outputs.Image(type="numpy", label="Converted Image")

iface = gr.Interface(
    generate,
    [input_image, radio],
    output_image,
    title = "Apple/Windows Style Emoji Conversion Using Cycle-GAN",
    article = "By: Arian Tashakkor"
)

iface.launch()