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}
}