MobileStyleGAN.pytorch - GitLite

KO | EN

gitlite — search

#python #hacktoberfest #react #typescript #llm #go #deep-learning #docker #nodejs #windows #awesome #cli #awesome-list #openai

MobileStyleGAN.pytorch

(★ 697)

An official implementation of MobileStyleGAN in PyTorch

File Explorer

Download Latest Version (.zip)

- mobile_stylegan_ffhq.json
- model_zoo.json
- template_cfg.json
- - __init__.py
  - diffaug.py
  - distiller_loss.py
  - non_saturating_gan_loss.py
  - perceptual_loss.py
- - __init__.py
  - discriminator.py
  - inception_v3.py
  - mapping_network.py
  - mobile_synthesis_network.py
  - synthesis_network.py
  - utils.py
- __init__.py
- dataset.py
- distiller.py
- model_zoo.py
- utils.py
- faces.jpeg
.gitignore
__init__.py
CHANGELOG
compare.py
convert_rosinality_ckpt.py
demo.py
evaluate_fid.py
generate.py
LICENSE
LICENSE-NVIDIA
README.md
requirements.txt
train.py

# Installation Guide

1. Get the code

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

Downloads the entire project code from GitHub to your computer.

cd MobileStyleGAN.pytorch

Moves into the project folder you just downloaded.

2. Python

Easy Recommended

Prerequisites

Git Needed to download the project code from GitHub.
Python 3 On Windows, be sure to check 'Add Python to PATH' during installation.

pip install -r requirements.txt

Installs the Python libraries listed in requirements.txt (or similar).

✅ If it runs without errors and prints output in the terminal, it worked.

Pulled directly from this repo's README.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

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

Official PyTorch Implementation

<p align="center">
  <img src="res/faces.jpeg"/>
</p>

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 <n_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 <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:

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

## Generate images using MobileStyleGAN

```bash
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](https://github.com/mseitzer/pytorch-fid) library:

```bash
python evaluate_fid.py <path_to_ref_dataset> <path_to_generated_imgs>
```

## Demo

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

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

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

## Convert to CoreML
```bash
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](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 <pretrined_model_name>`.

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