import gradio as gr
from huggingface_hub import hf_hub_download

import torch
from networks import define_G
from torchvision import transforms

REPO_ID = "Launchpad/ditto"
FILENAME = "model.pth"

model_dict = torch.load(
    hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
)
generator = define_G(input_nc=3, output_nc=3, ngf=64, netG="resnet_9blocks", norm="instance")
generator.load_state_dict(model_dict)
generator.eval()

# set up transforms for model
encode = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((256, 256))
])
transform = transforms.ToPILImage()

def generate_pokemon(pet_img):
    # encode image
    encoded_img = encode(pet_img)

    # evaluate model on pet image
    with torch.no_grad():
        generated_img = generator(encoded_img)
        
    # transform to PIL image
    return transform(generated_img)

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Image("https://www.ocf.berkeley.edu/~launchpad/media/uploads/project_logos/Ditto.png", elem_id="logo-img", show_label=False, show_share_button=False, show_download_button=False)
            
        with gr.Column(scale=3):
          gr.Markdown("""Ditto is a [Launchpad](https://launchpad.studentorg.berkeley.edu/) project (Fall 2022) that transfers styles of Pokemon sprites onto pet images using GANs and contrastive learning.
                      <br/><br/>
                      **Model**: [ditto](https://huggingface.co/Launchpad/ditto)
                      <br/>
                      **Developed by**: Kiran Suresh, Annie Lee, Chloe Wong, Tony Xin, Sebastian Zhao
                      """
                      )
    with gr.Row():  
        gr.Interface(generate_pokemon, gr.Image(), "image")

if __name__ == '__main__':
    demo.launch()