Home

Awesome

StyleGAN2 — Encoder/Projector for Official TensorFlow Implementation

Python 3.6 TensorFlow 1.10 cuDNN 7.3.1 License CC BY-NC

This is a port of Puzer/stylegan-encoder for NVlabs/stylegan2, plus a modified StyleGAN2 projector.

Teaser image

Generating latent representation of your images, using the original encoder

pip install tensorflow-gpu==1.14

git clone https://github.com/rolux/stylegan2encoder.git

cd stylegan2encoder

You can generate latent representations of your own images using two scripts:

  1. Extract and align faces from images

python align_images.py raw_images/ aligned_images/

  1. Find latent representation of aligned images

python encode_images.py aligned_images/ generated_images/ latent_representations/

Generating latent representation of your images, using the modified projector

Replace step 2 with:

python project_images.py aligned_images/ generated_images/

This is usually preferable. It also allows you to render a video of the optimization process. To see all available options, type:

python project_images.py -h

Original Readme

Teaser image

Analyzing and Improving the Image Quality of StyleGAN<br> Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, Timo Aila<br>

Paper: http://arxiv.org/abs/1912.04958<br> Video: https://youtu.be/c-NJtV9Jvp0<br>

Abstract: The style-based GAN architecture (StyleGAN) yields state-of-the-art results in data-driven unconditional generative image modeling. We expose and analyze several of its characteristic artifacts, and propose changes in both model architecture and training methods to address them. In particular, we redesign generator normalization, revisit progressive growing, and regularize the generator to encourage good conditioning in the mapping from latent vectors to images. In addition to improving image quality, this path length regularizer yields the additional benefit that the generator becomes significantly easier to invert. This makes it possible to reliably detect if an image is generated by a particular network. We furthermore visualize how well the generator utilizes its output resolution, and identify a capacity problem, motivating us to train larger models for additional quality improvements. Overall, our improved model redefines the state of the art in unconditional image modeling, both in terms of existing distribution quality metrics as well as perceived image quality.

For business inquiries, please contact researchinquiries@nvidia.com<br> For press and other inquiries, please contact Hector Marinez at hmarinez@nvidia.com<br>

Additional material 
StyleGAN2Main Google Drive folder
├  stylegan2-paper.pdfHigh-quality version of the paper
├  stylegan2-video.mp4High-quality version of the video
├  imagesExample images produced using our method
│  ├  curated-imagesHand-picked images showcasing our results
│  └  100k-generated-imagesRandom images with and without truncation
├  videosIndividual clips of the video as high-quality MP4
└  networksPre-trained networks
   ├  stylegan2-ffhq-config-f.pklStyleGAN2 for <span style="font-variant:small-caps">FFHQ</span> dataset at 1024×1024
   ├  stylegan2-car-config-f.pklStyleGAN2 for <span style="font-variant:small-caps">LSUN Car</span> dataset at 512×384
   ├  stylegan2-cat-config-f.pklStyleGAN2 for <span style="font-variant:small-caps">LSUN Cat</span> dataset at 256×256
   ├  stylegan2-church-config-f.pklStyleGAN2 for <span style="font-variant:small-caps">LSUN Church</span> dataset at 256×256
   ├  stylegan2-horse-config-f.pklStyleGAN2 for <span style="font-variant:small-caps">LSUN Horse</span> dataset at 256×256
   └ ⋯Other training configurations used in the paper

Requirements

StyleGAN2 relies on custom TensorFlow ops that are compiled on the fly using NVCC. To test that your NVCC installation is working correctly, run:

nvcc test_nvcc.cu -o test_nvcc -run
| CPU says hello.
| GPU says hello.

On Windows, the compilation requires Microsoft Visual Studio to be in PATH. We recommend installing Visual Studio Community Edition and adding into PATH using "C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Auxiliary\Build\vcvars64.bat".

Preparing datasets

Datasets are stored as multi-resolution TFRecords, similar to the original StyleGAN. Each dataset consists of multiple *.tfrecords files stored under a common directory, e.g., ~/datasets/ffhq/ffhq-r*.tfrecords. In the following sections, the datasets are referenced using a combination of --dataset and --data-dir arguments, e.g., --dataset=ffhq --data-dir=~/datasets.

FFHQ. To download the Flickr-Faces-HQ dataset as multi-resolution TFRecords, run:

pushd ~
git clone https://github.com/NVlabs/ffhq-dataset.git
cd ffhq-dataset
python download_ffhq.py --tfrecords
popd
python dataset_tool.py display ~/ffhq-dataset/tfrecords/ffhq

LSUN. Download the desired LSUN categories in LMDB format from the LSUN project page. To convert the data to multi-resolution TFRecords, run:

python dataset_tool.py create_lsun_wide ~/datasets/car ~/lsun/car_lmdb --width=512 --height=384
python dataset_tool.py create_lsun ~/datasets/cat ~/lsun/cat_lmdb --resolution=256
python dataset_tool.py create_lsun ~/datasets/church ~/lsun/church_outdoor_train_lmdb --resolution=256
python dataset_tool.py create_lsun ~/datasets/horse ~/lsun/horse_lmdb --resolution=256

Custom. Create custom datasets by placing all training images under a single directory. The images must be square-shaped and they must all have the same power-of-two dimensions. To convert the images to multi-resolution TFRecords, run:

python dataset_tool.py create_from_images ~/datasets/my-custom-dataset ~/my-custom-images
python dataset_tool.py display ~/datasets/my-custom-dataset

Using pre-trained networks

Pre-trained networks are stored as *.pkl files on the StyleGAN2 Google Drive folder. Below, you can either reference them directly using the syntax gdrive:networks/<filename>.pkl, or download them manually and reference by filename.

Generating images:

# Generate uncurated ffhq images (matches paper Figure 12)
python run_generator.py generate-images --network=gdrive:networks/stylegan2-ffhq-config-f.pkl \
  --seeds=6600-6625 --truncation-psi=0.5

# Generate curated ffhq images (matches paper Figure 11)
python run_generator.py generate-images --network=gdrive:networks/stylegan2-ffhq-config-f.pkl \
  --seeds=66,230,389,1518 --truncation-psi=1.0

# Generate uncurated car images
python run_generator.py generate-images --network=gdrive:networks/stylegan2-car-config-f.pkl \
  --seeds=6000-6025 --truncation-psi=0.5

# Example of style mixing (matches the corresponding video clip)
python run_generator.py style-mixing-example --network=gdrive:networks/stylegan2-ffhq-config-f.pkl \
  --row-seeds=85,100,75,458,1500 --col-seeds=55,821,1789,293 --truncation-psi=1.0

The results are placed in results/<RUNNING_ID>/*.png. You can change the location with --result-dir. For example, --result-dir=~/my-stylegan2-results.

Projecting images to latent space:

# Project generated images
python run_projector.py project-generated-images --network=gdrive:networks/stylegan2-car-config-f.pkl \
  --seeds=0,1,5

# Project real images
python run_projector.py project-real-images --network=gdrive:networks/stylegan2-car-config-f.pkl \
  --dataset=car --data-dir=~/datasets

You can import the networks in your own Python code using pickle.load(). For this to work, you need to include the dnnlib source directory in PYTHONPATH and create a default TensorFlow session by calling dnnlib.tflib.init_tf(). See run_generator.py and pretrained_networks.py for examples.

Training networks

To reproduce the training runs for config F in Tables 1 and 3, run:

python run_training.py --num-gpus=8 --data-dir=~/datasets --config=config-f \
  --dataset=ffhq --mirror-augment=true
python run_training.py --num-gpus=8 --data-dir=~/datasets --config=config-f \
  --dataset=car --total-kimg=57000
python run_training.py --num-gpus=8 --data-dir=~/datasets --config=config-f \
  --dataset=cat --total-kimg=88000
python run_training.py --num-gpus=8 --data-dir=~/datasets --config=config-f \
  --dataset=church --total-kimg 88000 --gamma=100
python run_training.py --num-gpus=8 --data-dir=~/datasets --config=config-f \
  --dataset=horse --total-kimg 100000 --gamma=100

For other configurations, see python run_training.py --help.

We have verified that the results match the paper when training with 1, 2, 4, or 8 GPUs. Note that training FFHQ at 1024×1024 resolution requires GPU(s) with at least 16 GB of memory. The following table lists typical training times using NVIDIA DGX-1 with 8 Tesla V100 GPUs:

ConfigurationResolutionTotal kimg1 GPU2 GPUs4 GPUs8 GPUsGPU mem
config-f1024×10242500069d 23h36d 4h18d 14h9d 18h13.3 GB
config-f1024×10241000027d 23h14d 11h7d 10h3d 22h13.3 GB
config-e1024×10242500035d 11h18d 15h9d 15h5d 6h8.6 GB
config-e1024×10241000014d 4h7d 11h3d 20h2d 3h8.6 GB
config-f256×2562500032d 13h16d 23h8d 21h4d 18h6.4 GB
config-f256×2561000013d 0h6d 19h3d 13h1d 22h6.4 GB

Training curves for FFHQ config F (StyleGAN2) compared to original StyleGAN using 8 GPUs:

Training curves

After training, the resulting networks can be used the same way as the official pre-trained networks:

# Generate 1000 random images without truncation
python run_generator.py generate-images --seeds=0-999 --truncation-psi=1.0 \
  --network=results/00006-stylegan2-ffhq-8gpu-config-f/networks-final.pkl

Evaluation metrics

To reproduce the numbers for config F in Tables 1 and 3, run:

python run_metrics.py --data-dir=~/datasets --network=gdrive:networks/stylegan2-ffhq-config-f.pkl \
  --metrics=fid50k,ppl_wend --dataset=ffhq --mirror-augment=true
python run_metrics.py --data-dir=~/datasets --network=gdrive:networks/stylegan2-car-config-f.pkl \
  --metrics=fid50k,ppl2_wend --dataset=car
python run_metrics.py --data-dir=~/datasets --network=gdrive:networks/stylegan2-cat-config-f.pkl \
  --metrics=fid50k,ppl2_wend --dataset=cat
python run_metrics.py --data-dir=~/datasets --network=gdrive:networks/stylegan2-church-config-f.pkl \
  --metrics=fid50k,ppl2_wend --dataset=church
python run_metrics.py --data-dir=~/datasets --network=gdrive:networks/stylegan2-horse-config-f.pkl \
  --metrics=fid50k,ppl2_wend --dataset=horse

For other configurations, see the StyleGAN2 Google Drive folder.

Note that the metrics are evaluated using a different random seed each time, so the results will vary between runs. In the paper, we reported the average result of running each metric 10 times. The following table lists the available metrics along with their expected runtimes and random variation:

MetricFFHQ config F1 GPU2 GPUs4 GPUsDescription
fid50k2.84 ± 0.0322 min14 min10 minFréchet Inception Distance
is50k5.13 ± 0.0223 min14 min8 minInception Score
ppl_zfull348.0 ± 3.841 min22 min14 minPerceptual Path Length in Z, full paths
ppl_wfull126.9 ± 0.242 min22 min13 minPerceptual Path Length in W, full paths
ppl_zend348.6 ± 3.041 min22 min14 minPerceptual Path Length in Z, path endpoints
ppl_wend129.4 ± 0.840 min23 min13 minPerceptual Path Length in W, path endpoints
ppl2_wend145.0 ± 0.541 min23 min14 minPerceptual Path Length without center crop
ls154.2 / 4.2710 hrs6 hrs4 hrsLinear Separability
pr50k30.689 / 0.49226 min17 min12 minPrecision and Recall

Note that some of the metrics cache dataset-specific data on the disk, and they will take somewhat longer when run for the first time.

License

Copyright © 2019, NVIDIA Corporation. All rights reserved.

This work is made available under the Nvidia Source Code License-NC. To view a copy of this license, visit https://nvlabs.github.io/stylegan2/license.html

Citation

@article{Karras2019stylegan2,
  title   = {Analyzing and Improving the Image Quality of {StyleGAN}},
  author  = {Tero Karras and Samuli Laine and Miika Aittala and Janne Hellsten and Jaakko Lehtinen and Timo Aila},
  journal = {CoRR},
  volume  = {abs/1912.04958},
  year    = {2019},
}

Acknowledgements

We thank Ming-Yu Liu for an early review, Timo Viitanen for his help with code release, and Tero Kuosmanen for compute infrastructure.