app.py

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
                      <br/>
                      **Examples**: [Oxford-IIIT Pet Dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/)
                      """
                      )
    gr.Interface(fn=generate_pokemon, 
                 inputs=gr.Image(), 
                 outputs="image",
                 examples=["data/german_shorthaired_164.jpg", "data/samoyed_189.jpg", "data/shiba_inu_139.jpg"]
    )

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