PmSFC

Implementation for paper G. Li, Q. Fen, S. Qian, and et al. High Fidelity GAN Inversion via Prior Multi-Subspace Feature Composition, accepted for Association for Advancements of Artificial Intelligence (AAAI) 2021. For more details, please refer to my personal website or read the PDF.

Teaser: The flowchart of the PmSFC. The top one is about excavating useful subspaces from the pre-trained GAN. The bottom one is to cooperate subspace in GAN inversion for high-fidelity inversion.

Description

• /models: This directory contains most of the SOTA generators, e.g., PGGAN, StyleGAN2 and BigGAN. The original repo is from mGANprior. Please download necessary checkpoints from the following links and place them in the /models/pretrain/tensorflow, the project will convert tensorflow-based checkpoints to pytorch version, and store them in /models/pretrain/pytorch. For a complete link, please refer to /models/README.MD.

PGGAN Official
Face
celebahq-1024x1024
Indoor Scene
bedroom-256x256	livingroom-256x256	diningroom-256x256	kitchen-256x256
Outdoor Scene
churchoutdoor-256x256	tower-256x256	bridge-256x256
Animal
person-256x256	cat-256x256	dog-256x256	bird-256x256
horse-256x256	sheep-256x256	cow-256x256

• /bin: contains some testing generated images and pretrained affinity matrixs.

• /examples: contains some target images.

• Note: The original repo has some mistakes in loading StyleGAN, and I have corrected them. I also add many personal utilities to the repo, e.g., getting intermediate features for calculating subspaces. It could be messy, so you can pull requests if you are interested.

Environment

Please use following command to install the environment.

python -m pip install -r requirements.txt

I also provide the conda environment in file conda_env.yaml if you are interested.

conda env create -f conda_env.yaml

Run the code

Take CelebaHQ as example, you can find useful subspaces by running the command.

python discover_subspace.py \
--outputs=./TRAIN --inversion_type=PGGAN-Layerwise \
--gan_model=pggan_celebahq --layer=3 --iterations=200 \
--optimization=Adam --lr=0.0001 --report_image=5 \
--report_model=5 --batch_size=4 --exp_id=SelfExpress \
--init_type=Zeros --beta=1 --cluster_numbers=6 --alpha=0.2 \
--subspace_dimension=6 --sparse_type=L1 --power=2.0

Place the pretrained affinity matrix to args.matrix_dir and run the following command.

python PmSFC.py --outputs=./TRAIN \
--inversion_type=PGGAN-Layerwise --gan_model=pggan_celebahq \
--layer=3 --iterations=20000 --optimization=Adam --lr=0.0001 \
--report_image=5 --report_model=5 --batch_size=4 --exp_id=SelfExpress \
--beta0=1 --beta1=1 --n_subs=6 --alpha=0.2 --d_subs=6 \
--sparse_type=L1 --power=2.0 \
--matrix_dir=./bin/pggan/celebahq/matrix/layer3/value50.pkl

You should get similar results like this:

Subspace visualization by switching the subspace features.