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+
- Ninja build:
sudo apt install -y ninja-build
- PyTorch < 1.11.0 & CUDA Toolkit
conda install pytorch==1.10.1 torchvision==0.11.2 torchaudio==0.10.1 cudatoolkit=11.3 -c pytorch -c conda-forge
orpip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/torch_stable.html
- 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>
Convert checkpoint from rosinality/stylegan2-pytorch
Our framework supports StyleGAN2 checkpoints format from rosinality/stylegan2-pytorch. To convert ckpt your own checkpoint of StyleGAN2 to our framework:
python convert_rosinality_ckpt.py --ckpt <path_to_rosinality_stylegan2_ckpt> --ckpt-mnet <path_to_output_mapping_network_ckpt> --ckpt-snet <path_to_output_synthesis_network_ckpt> --cfg-path <path_to_output_config_json>
Check converted checkpoint
To check that your checkpoint is converted correctly, just run demo visualization:
python demo.py --cfg <path_to_output_config_json> --ckpt "" --generator teacher
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> --export-model onnx --export-dir <output_dir>
Convert ONNX to OpenVINO
python3 -m venv vino_env=2021.4.2
source vino_env\=2021.4.2/bin/activate
python -m pip install --upgrade pip
pip install openvino-dev[pytorch,onnx]==2021.4.2
mo --input_model "logs/2022-10-04_17:09:07/MappingNetwork.onnx" --output_dir "logs/2022-10-04_17:09:07/"
mo --input_model "logs/2022-10-04_17:09:07/SynthesisNetwork.onnx" --output_dir "logs/2022-10-04_17:09:07/"
deactivate
ONNX MAC(=FLOPs/2) Counter
https://github.com/ThanatosShinji/onnx-tool
python3 -m venv onnx_tool
source onnx_tool/bin/activate
python -m pip install --upgrade pip
pip install onnx-tool
import onnx_tool
onnx_tool.model_profile("MappingNetwork.onnx", savenode="mnet_ops.txt")
onnx_tool.model_profile("SynthesisNetwork.onnx", savenode="snet_ops.txt")
Convert to CoreML
python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt <path_to_ckpt> --export-model coreml --export-dir <output_dir>
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 | 7.75 |
(*) 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.
- Cyril Diagne for the excellent demo of how to run MobileStyleGAN directly into the web-browser.
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}
}
@article{BELOUSOV2021100115,
title = {MobileStyleGAN.pytorch: PyTorch-based toolkit to compress StyleGAN2 model},
journal = {Software Impacts},
year = {2021},
issn = {2665-9638},
doi = {https://doi.org/10.1016/j.simpa.2021.100115},
url = {https://www.sciencedirect.com/science/article/pii/S2665963821000452},
author = {Sergei Belousov},
}