/PmSFC

We use self-expressive layer to learn the affinity matrix of the hidden space of a generator, and we found subspace in GAN. Also we propsed a subspace-based high-fidelity GAN inversion model.

Primary LanguagePythonApache License 2.0Apache-2.0

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.