Ecosyste.ms: Awesome

An open API service indexing awesome lists of open source software.

Awesome Lists | Featured Topics | Projects

https://github.com/bes-dev/MobileStyleGAN.pytorch

An official implementation of MobileStyleGAN in PyTorch
https://github.com/bes-dev/MobileStyleGAN.pytorch

gan image-synthesis mobile-development openvino stylegan2 stylegan2-pytorch sylegan

Last synced: 3 months ago
JSON representation

An official implementation of MobileStyleGAN in PyTorch

Awesome Lists containing this project

README

        

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

Official PyTorch Implementation



The accompanying videos can be found on [YouTube](https://www.youtube.com/playlist?list=PLstKhmdpWBtwsvq_27ALmPbf_mBLmk0uI).
For more details, please refer to the [paper](https://arxiv.org/abs/2104.04767).

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

```bash
pip install -r requirements.txt
python train.py --cfg configs/mobile_stylegan_ffhq.json --gpus
```

## Convert checkpoint from rosinality/stylegan2-pytorch

Our framework supports StyleGAN2 checkpoints format from [rosinality/stylegan2-pytorch](https://github.com/rosinality/stylegan2-pytorch).
To convert ckpt your own checkpoint of StyleGAN2 to our framework:

```bash
python convert_rosinality_ckpt.py --ckpt --ckpt-mnet --ckpt-snet --cfg-path
```

### Check converted checkpoint

To check that your checkpoint is converted correctly, just run demo visualization:

```bash
python demo.py --cfg --ckpt "" --generator teacher
```

## Generate images using MobileStyleGAN

```bash
python generate.py --cfg configs/mobile_stylegan_ffhq.json --device cuda --ckpt --output-path --batch-size --n-batches
```

## Evaluate FID score

To evaluate the FID score we use a modified version of [pytorch-fid](https://github.com/mseitzer/pytorch-fid) library:

```bash
python evaluate_fid.py
```

## Demo

Run demo visualization using MobileStyleGAN:
```bash
python demo.py --cfg configs/mobile_stylegan_ffhq.json --ckpt
```

Run visual comparison using StyleGAN2 vs. MobileStyleGAN:
```bash
python compare.py --cfg configs/mobile_stylegan_ffhq.json --ckpt
```

## Convert to ONNX
```bash
python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt --export-model onnx --export-dir
```

## Convert to CoreML
```bash
python train.py --cfg configs/mobile_stylegan_ffhq.json --ckpt --export-model coreml --export-dir
```

## Deployment using OpenVINO

We provide external library [random_face](https://github.com/bes-dev/random_face) as an example of deploying our model at the edge devices using the [OpenVINO](https://github.com/openvinotoolkit/openvino) framework.

## Pretrained models

|Name|FID|
|:---|:--|
|[mobilestylegan_ffhq.ckpt](https://drive.google.com/uc?id=11Kja0XGE8liLb6R5slNZjF3j3v_6xydt)|7.75|

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

## Legacy license

|Code|Source|License|
|:---|:-----|:------|
|[Custom CUDA kernels](core/models/modules/ops/)|https://github.com/NVlabs/stylegan2|[Nvidia License](LICENSE-NVIDIA)|
|[StyleGAN2 blocks](core/models/modules/legacy.py)|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](https://github.com/rosinality/stylegan2-pytorch).
* Fergal Cotter for implementation of Discrete Wavelet Transforms and Inverse Discrete Wavelet Transforms in [PyTorch](https://github.com/fbcotter/pytorch_wavelets).
* Cyril Diagne for the excellent [demo of how to run MobileStyleGAN directly into the web-browser](https://github.com/cyrildiagne/mobilestylegan-web-demo).

## 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},
}
```