/CIPS-3D

3D-aware GANs based on NeRF (arXiv).

Primary LanguagePythonMIT LicenseMIT

This repository contains the code of the paper CIPS-3D: A 3D-Aware Generator of GANs Based on Conditionally-Independent Pixel Synthesis. For 3D GAN inversion and editing of real images, please refer to this repository: https://github.com/PeterouZh/CIPS-3Dplusplus (CIPS-3D++).

Updates

✔️ (2022-1-5) The code has been refactored. Please refer to the scripts in exp/cips3d/bash. Please upgrade the tl2 package with pip install -I tl2.

✔️ (2021-11-26) The configuration files (yaml files) for training are being released.

✔️ (2021-10-27) All the code files have been released. The configuration files (yaml files) for training will be released next. Now I have provided a GUI script and models to facilitate the experiment of network interpolation (see below). If you find any problems, please open an issue. Have fun with it.

✔️ (2021-10-25) Thank you for your kind attention. The github star has reached two hundred. I will open source the training code in the near future.

✔️ (2021-10-20) We are planning to publish the training code here in December. But if the github star reaches two hundred, I will advance the date. Stay tuned 🕙.

Demo videos

demo1.mp4
demo2.mp4
demo_animal_finetuned.mp4
demo3.mp4
demo4.mp4
demo5.mp4
Mirror symmetry problem

The problem of mirror symmetry refers to the sudden change of the direction of the bangs near the yaw angle of pi/2. We propose to use an auxiliary discriminator to solve this problem (please see the paper).

Note that in the initial stage of training, the auxiliary discriminator must dominate the generator more than the main discriminator does. Otherwise, if the main discriminator dominates the generator, the mirror symmetry problem will still occur. In practice, progressive training is able to guarantee this. We have trained many times from scratch. Adding an auxiliary discriminator stably solves the mirror symmetry problem. If you find any problems with this idea, please open an issue.

Prepare environment 
git clone --recursive https://github.com/PeterouZh/CIPS-3D.git
cd CIPS-3D

# Create virtual environment
conda create -y --name cips3d python=3.6.7
conda activate cips3d

pip install torch==1.8.2+cu102 torchvision==0.9.2+cu102 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html

pip install --no-cache-dir -r requirements.txt
pip install -I tl2

pip install -e torch_fidelity_lib
pip install -e pytorch_ema_lib

Model interpolation (web demo)

Download the pre-trained checkpoints.

Execute this command:

streamlit run --server.port 8650 -- scripts/web_demo.py  \
  --outdir results/model_interpolation \
  --cfg_file configs/web_demo.yaml \
  --command model_interpolation

Then open the browser: http://your_ip_address:8650.

You can debug this script with this command:

python scripts/web_demo.py  \
  --outdir results/model_interpolation \
  --cfg_file configs/web_demo.yaml \
  --command model_interpolation \
  --debug True

Pre-trained checkpoints

ffhq_exp
FFHQ_r256 train_ffhq_high-20220105_143314_190
AFHQ_r256 finetune_afhq-20220124_193407_473
CartoonFaces_r256 finetune_photo2cartoon-20220107_172255_454
Prepare dataset

FFHQ: Download FFHQ dataset images1024x1024 (89.1 GB)

# Downsampling images in advance to speed up training
python scripts/dataset_tool.py \
    --source=datasets/ffhq/images1024x1024 \
    --dest=datasets/ffhq/downsample_ffhq_256x256.zip \
    --width=256 --height=256

CartoonFaces Download photo2cartoon dataset

# Prepare training dataset.
python scripts/dataset_tool.py \
    --source=datasets/photo2cartoon/photo2cartoon \
    --dest=datasets/photo2cartoon/photo2cartoon_stylegan2.zip

AFHQ Download afhq dataset

# Prepare training dataset.
python scripts/dataset_tool.py \
    --source=datasets/AFHQv2/AFHQv2 \
    --dest=datasets/AFHQv2/AFHQv2_stylegan2.zip

Training

Please refer to the scripts in exp/cips3d/bash. I will release all the pre-trained models when the reproducing is over.

running order:

  • exp/cips3d/bash/ffhq_exp:

    • train_ffhq_r32.sh -> train_ffhq_r64.sh -> train_ffhq_r128.sh -> train_ffhq_r256.sh
    • eval_fid.sh

  • exp/cips3d/bash/finetuning_exp: (require pre-trained models from the above step)

    • finetune_photo2cartoon.sh

developing:

  • exp/cips3d/bash/ffhq_exp_v1:

  • exp/cips3d/bash/afhq_exp:

Bug fixed

  • If the training process is blocked when training with multi GPUs, please upgrade the tl2 via pip install -I tl2
Old readme

Note:

  • In order to ensure that this code is consistent with my original dirty code, please follow me to reproduce the results using this code step by step.
  • The training script train_v16.py is dirty, but I'm not going to refactor it. After all, it still works stably.

Start training at 64x64

Training:

export CUDA_HOME=/usr/local/cuda-10.2/
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export PYTHONPATH=.
python exp/dev/nerf_inr/scripts/train_v16.py \
    --port 8888 \
    --tl_config_file configs/train_ffhq.yaml \
    --tl_command train_ffhq \
    --tl_outdir results/train_ffhq \
    --tl_opts curriculum.new_attrs.image_list_file datasets/ffhq/images256x256_image_list.txt \
      D_first_layer_warmup True

Dummy training (for debug):

export CUDA_HOME=/usr/local/cuda-10.2/
export CUDA_VISIBLE_DEVICES=1
python exp/dev/nerf_inr/scripts/train_v16.py \
    --port 8888 \ 
    --tl_config_file configs/train_ffhq.yaml \
    --tl_command train_ffhq \
    --tl_outdir results/train_ffhq_debug \ 
    --tl_debug \
    --tl_opts curriculum.new_attrs.image_list_file datasets/ffhq/images256x256_image_list.txt \
      num_workers 0 num_images_real_eval 10 num_images_gen_eval 2

When the FID of the 64x64 model reaches about 16, we start the next step: resume training at 128x128. Let's wait for the training (about 2 days or less).

Reproduced results: best_FID=15.27

Resume training at 128x128 from the 64x64 models

Training:

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export PYTHONPATH=.
python exp/dev/nerf_inr/scripts/train_v16.py \
    --port 8888 \
    --tl_config_file configs/train_ffhq.yaml \
    --tl_command train_ffhq_r128 \
    --tl_outdir results/train_ffhq \
    --tl_resume \
    --tl_resumedir results/train_ffhq \
    --tl_opts curriculum.new_attrs.image_list_file datasets/ffhq/images256x256_image_list.txt \
      D_first_layer_warmup True reset_best_fid True update_aux_every 16 d_reg_every 1 train_aux_img True

Dummy training (for debug):

export CUDA_HOME=/usr/local/cuda-10.2/
export CUDA_VISIBLE_DEVICES=1
python exp/dev/nerf_inr/scripts/train_v16.py \
    --port 8888 \ 
    --tl_config_file configs/train_ffhq.yaml \
    --tl_command train_ffhq_r128 \
    --tl_outdir results/train_ffhq \ 
    --tl_resume \
    --tl_resumedir results/train_ffhq \
    --tl_debug \
    --tl_opts curriculum.new_attrs.image_list_file datasets/ffhq/images256x256_image_list.txt \
      num_workers 0 num_images_real_eval 10 num_images_gen_eval 2 reset_best_fid True

When the FID of the 128x128 model reaches about 16, we start the next step.

Some hyperparameters may be different from the original experiment. Hope it works normally. Let's wait for the training (maybe longer).

Resume training at 256x256 from the 128x128 models



Finetune INR Net






Citation


If you find our work useful in your research, please cite:



@article{zhou2021CIPS3D,
  title = {{{CIPS}}-{{3D}}: A {{3D}}-{{Aware Generator}} of {{GANs Based}} on {{Conditionally}}-{{Independent Pixel Synthesis}}},
  shorttitle = {{{CIPS}}-{{3D}}},
  author = {Zhou, Peng and Xie, Lingxi and Ni, Bingbing and Tian, Qi},
  year = {2021},
  eprint = {2110.09788},
  eprinttype = {arxiv},
}




Acknowledgments