Pytorch implementation of the paper Face Identity Disentanglement via Latent Space Mapping for both training and evaluation, with StyleGAN 2.
- Paper: https://arxiv.org/abs/2005.07728
- Official TensorFlow Implementation: https://github.com/YotamNitzan/ID-disentanglement
- Important: this implementation does not follow the exact same architecture of the original paper. See changes below
- instead of using a Discriminator loss for the mapper. We have used several other losses such as:
- LPIPS Loss (The Unreasonable Effectiveness of Deep Features as a Perceptual Metric, Zhang el al, 2018)
- MSE Loss
- Different ID Loss
- Different landmark detector
- The reason for those changes resides in the fact that the training procedure with Discriminator is often hard and does not converge. We have found that replacing the Discriminator with LPIPS and MSE losses we can achieve the same result. Nevertheless, our code supports training with a discriminator which can be activated using the configuration.
- The other changes are due to better Recognition models that have developed since the original paper was published
We used several pretrained models:
- StyleGan2 Generator for image size 256 - 550000.pt
- ID Encoder - model_ir_se50.pth
- Landmarks Detection - mobilefacenet_model_best.pth.tar
Weight files attached at this Drive folder.
You can also find at the above link our environment.yml file to create a relevant conda environment.
If your model is trained by stylegan2-ada-pytorch, please change file and use Training/train_script2.py instead of Training/train_script.py to Train
The dataset is comprised of StyleGAN 2 generated images and W latent codes. see Utils/data_creator.py.
if you want to use model trained by stylegan2-ada-pytorch, please use Utils/data_creator2.py
Examples of our generated dataset attached at this Drive folder.
To train the model run train_script.py, you can change parameters in Configs/ folder.
Try Inference.ipynb notebook to disentangle identity from attributes by yourself
Our pretrained checkpoint attached at this Drive folder.
