|
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() |
|
|
|
|
|
encode = transforms.Compose([ |
|
transforms.ToTensor(), |
|
transforms.Resize((256, 256)) |
|
]) |
|
transform = transforms.ToPILImage() |
|
|
|
def generate_pokemon(pet_img): |
|
|
|
encoded_img = encode(pet_img) |
|
|
|
|
|
with torch.no_grad(): |
|
generated_img = generator(encoded_img) |
|
|
|
|
|
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() |
|
|
