peterzhousz / mobilestylegan.pytorch

An official implementation of MobileStyleGAN in PyTorch

MobileStyleGAN: A Lightweight Convolutional Neural Network for High-Fidelity Image Synthesis

Official PyTorch Implementation

The accompanying videos can be found on YouTube. For more details, please refer to the paper.

Requirements

• Python 3.8+
• 1–8 high-end NVIDIA GPUs with at least 12 GB of memory. We have done all testing and development using DL Workstation with 4x2080Ti

Training

pip install -r requirements.txt
python train.py --cfg configs/mobile_stylegan_ffhq.json --gpus <n_gpus>

Generate images using MobileStyleGAN

python generate.py --cfg configs/mobile_stylegan_ffhq.json --device cuda --ckpt <path_to_ckpt> --output-path <path_to_store_imgs> --batch-size <batch_size> --n-batches <n_batches>

Evaluate FID score

To evaluate the FID score we use a modified version of pytorch-fid library:

python evaluate_fid.py <path_to_ref_dataset> <path_to_generated_imgs>

Demo

Run demo visualization using MobileStyleGAN:

python demo.py --cfg configs/mobile_stylegan_ffhq.json --ckpt <path_to_ckpt>

Run visual comparison using StyleGAN2 vs. MobileStyleGAN:

python compare.py --cfg configs/mobile_stylegan_ffhq.json --ckpt <path_to_ckpt>

Convert to ONNX

python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt <path_to_ckpt> --to-onnx <onnx_prefix_name>

Deployment using OpenVINO

We provide external library random_face as an example of deploying our model at the edge devices using the OpenVINO framework.

Pretrained models

Name	FID
mobilestylegan_ffhq.ckpt	12.38

(*) Our framework supports automatic download pretrained models, just use --ckpt <pretrined_model_name>.

Legacy license

Code	Source	License
Custom CUDA kernels	https://github.com/NVlabs/stylegan2	Nvidia License
StyleGAN2 blocks	https://github.com/rosinality/stylegan2-pytorch	MIT

Acknowledgements

We want to thank the people whose works contributed to our project::

• Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, Timo Aila for research related to style based generative models.
• Kim Seonghyeon for implementation of StyleGAN2 in PyTorch.
• Fergal Cotter for implementation of Discrete Wavelet Transforms and Inverse Discrete Wavelet Transforms in PyTorch.

Citation

If you are using the results and code of this work, please cite it as:

@misc{belousov2021mobilestylegan,
      title={MobileStyleGAN: A Lightweight Convolutional Neural Network for High-Fidelity Image Synthesis},
      author={Sergei Belousov},
      year={2021},
      eprint={2104.04767},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}