animegan2-pytorch: A Jupyter Notebook repository from chu108

PyTorch Implementation of AnimeGANv2

Updates

2021-10-17 Add weights for FacePortraitV2.
2021-11-07 Thanks to ak92501, a web demo is integrated to Huggingface Spaces with Gradio.
2021-11-07 Thanks to xhlulu, the torch.hub model is now available. See Torch Hub Usage.
2021-11-07 Add FacePortraitV2 style demo to a telegram bot. See @face2stickerbot by sxela

Basic Usage

Weight Conversion from the Original Repo (Requires TensorFlow 1.x)

git clone https://github.com/TachibanaYoshino/AnimeGANv2
python convert_weights.py

Inference

python test.py --input_dir [image_folder_path] --device [cpu/cuda]

samples

Results from converted `Paprika` style model (input image, original tensorflow result, pytorch result from left to right)

Note: Training code not included / Results from converted weights slightly different due to the bilinear upsample issue

Torch Hub Usage

You can load the model via torch.hub:

import torch
model = torch.hub.load("bryandlee/animegan2-pytorch", "generator").eval()
out = model(img_tensor)  # BCHW tensor

Currently, the following pretrained shorthands are available:

model = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="celeba_distill")
model = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v1")
model = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v2")
model = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="paprika")

You can also load the face2paint util function:

from PIL import Image

face2paint = torch.hub.load("bryandlee/animegan2-pytorch:main", "face2paint", size=512)

img = Image.open(...).convert("RGB")
out = face2paint(model, img)