genforce/sefa

[CVPR 2021] Closed-Form Factorization of Latent Semantics in GANs

GenForce Lib for Generative Modeling

An efficient PyTorch library for deep generative modeling. May the Generative Force (GenForce) be with You.

Updates

• Encoder Training: We support training encoders on top of pre-trained GANs for GAN inversion.
• Model Converters: You can easily migrate your already started projects to this repository. Please check here for more details.

Highlights

• Distributed training framework.
• Fast training speed.
• Modular design for prototyping new models.
• Model zoo containing a rich set of pretrained GAN models, with Colab live demo to play.

Installation

1. Create a virtual environment via conda.

   conda create -n genforce python=3.7
   conda activate genforce

2. Install cuda and cudnn. (We use CUDA 10.0 in case you would like to use TensorFlow 1.15 for model conversion.)

   conda install cudatoolkit=10.0 cudnn=7.6.5

3. Install torch and torchvision.

   pip install torch==1.7 torchvision==0.8

4. Install requirements

   pip install -r requirements.txt

Quick Demo

We provide a quick training demo, scripts/stylegan_training_demo.py, which allows to train StyleGAN on a toy dataset (500 animeface images with 64 x 64 resolution). Try it via

./scripts/stylegan_training_demo.sh

We also provide an inference demo, synthesize.py, which allows to synthesize images with pre-trained models. Generated images can be found at work_dirs/synthesis_results/. Try it via

python synthesize.py stylegan_ffhq1024

You can also play the demo at Colab.

Play with GANs

Test

Pre-trained models can be found at model zoo.

• On local machine:

  GPUS=8
  CONFIG=configs/stylegan_ffhq256_val.py
  WORK_DIR=work_dirs/stylegan_ffhq256_val
  CHECKPOINT=checkpoints/stylegan_ffhq256.pth
  ./scripts/dist_test.sh ${GPUS} ${CONFIG} ${WORK_DIR} ${CHECKPOINT}

• Using slurm:

  CONFIG=configs/stylegan_ffhq256_val.py
  WORK_DIR=work_dirs/stylegan_ffhq256_val
  CHECKPOINT=checkpoints/stylegan_ffhq256.pth
  GPUS=8 ./scripts/slurm_test.sh ${PARTITION} ${JOB_NAME} \
      ${CONFIG} ${WORK_DIR} ${CHECKPOINT}

Train

All log files in the training process, such as log message, checkpoints, synthesis snapshots, etc, will be saved to the work directory.

• On local machine:

  GPUS=8
  CONFIG=configs/stylegan_ffhq256.py
  WORK_DIR=work_dirs/stylegan_ffhq256_train
  ./scripts/dist_train.sh ${GPUS} ${CONFIG} ${WORK_DIR} \
      [--options additional_arguments]

• Using slurm:

  CONFIG=configs/stylegan_ffhq256.py
  WORK_DIR=work_dirs/stylegan_ffhq256_train
  GPUS=8 ./scripts/slurm_train.sh ${PARTITION} ${JOB_NAME} \
      ${CONFIG} ${WORK_DIR} \
      [--options additional_arguments]

Play with Encoders for GAN Inversion

Train

• On local machine:

  GPUS=8
  CONFIG=configs/stylegan_ffhq256_encoder_y.py
  WORK_DIR=work_dirs/stylegan_ffhq256_encoder_y
  ./scripts/dist_train.sh ${GPUS} ${CONFIG} ${WORK_DIR} \
      [--options additional_arguments]

• Using slurm:

  CONFIG=configs/stylegan_ffhq256_encoder_y.py
  WORK_DIR=work_dirs/stylegan_ffhq256_encoder_y
  GPUS=8 ./scripts/slurm_train.sh ${PARTITION} ${JOB_NAME} \
      ${CONFIG} ${WORK_DIR} \
      [--options additional_arguments]

Contributors

Member	Module
Yujun Shen	models and running controllers
Yinghao Xu	runner and loss functions
Ceyuan Yang	data loader
Jiapeng Zhu	evaluation metrics
Bolei Zhou	cheerleader

NOTE: The above form only lists the person in charge for each module. We help each other a lot and develop as a TEAM.

We welcome external contributors to join us for improving this library.

License

The project is under the MIT License.

Acknowledgement

We thank PGGAN, StyleGAN, StyleGAN2, StyleGAN2-ADA for their work on high-quality image synthesis. We thank IDInvert and GHFeat for their contribution to GAN inversion. We also thank MMCV for the inspiration on the design of controllers.

BibTex

We open source this library to the community to facilitate the research of generative modeling. If you do like our work and use the codebase or models for your research, please cite our work as follows.

@misc{genforce2020,
  title =        {GenForce},
  author =       {Shen, Yujun and Xu, Yinghao and Yang, Ceyuan and Zhu, Jiapeng and Zhou, Bolei},
  howpublished = {\url{https://github.com/genforce/genforce}},
  year =         {2020}
}

sefa's Issues

About re-scoring analysis

Hi, I want to do some experiment on disentanglement, but i am confused about the metric.
Suppose G(z) is the origin image and d is the dirction, how to select k for re-scoring G(z+kd)?

Annotations corresponding to eigen vectors

Hi,

I couldn't find the post labeling for the semantic indices.
Could you share the post annotation for StyleGAN and StyleGAN2 on celeba/ffhq?

Thanks

Some questions about your code

sefa/utils.py

Line 185 in c60c536

weight = weight.flip(2, 3).permute(1, 0, 2, 3).flatten(1)

In this line, I wonder why you used torch.flip to flip the weight. I think there is no need to flip the weights before implementing convolution in CNN.

EOFError: Ran out of input

EOFError: Ran out of input
Traceback:
File "c:\users\a\anaconda3\lib\site-packages\streamlit\script_runner.py", line 332, in _run_script
exec(code, module.dict)
File "C:\Users\A\Desktop\sefa-master\interface.py", line 128, in
main()
File "C:\Users\A\Desktop\sefa-master\interface.py", line 69, in main
model = get_model(model_name)
File "c:\users\a\anaconda3\lib\site-packages\streamlit\caching.py", line 606, in wrapped_func
return get_or_create_cached_value()
File "c:\users\a\anaconda3\lib\site-packages\streamlit\caching.py", line 588, in get_or_create_cached_value
return_value = func(*args, **kwargs)
File "C:\Users\A\Desktop\sefa-master\interface.py", line 19, in get_model
return load_generator(model_name)
File "C:\Users\A\Desktop\sefa-master\utils.py", line 85, in load_generator
checkpoint = torch.load(checkpoint_path, map_location='cpu')
File "c:\users\a\anaconda3\lib\site-packages\torch\serialization.py", line 585, in load
return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
File "c:\users\a\anaconda3\lib\site-packages\torch\serialization.py", line 755, in _legacy_load
magic_number = pickle_module.load(f, **pickle_load_args)

How??

StyleGAN 2 ADA transfer learning sourcenet error

I trained a couple of models with the pytorch version of StyleGAN 2 ADA and want to use them with SeFa. Most of them run fine, but using models that are based on the transfer learning sourcenet models provided in the official ADA repo results in the following error:

RuntimeError: Error(s) in loading state_dict for StyleGAN2Generator:
Missing key(s) in state_dict: "mapping.dense2.weight", "mapping.dense2.bias", "mapping.dense3.weight", "mapping.dense3.bias", "mapping.dense4.weight", "mapping.dense4.bias", "mapping.dense5.weight", "mapping.dense5.bias", "mapping.dense6.weight", "mapping.dense6.bias", "mapping.dense7.weight", "mapping.dense7.bias".
Unexpected key(s) in state_dict: "synthesis.layer15.weight", "synthesis.layer15.bias", "synthesis.layer15.noise_strength", "synthesis.layer15.noise", "synthesis.layer15.filter.kernel", "synthesis.layer15.style.weight", "synthesis.layer15.style.bias", "synthesis.layer16.weight", "synthesis.layer16.bias", "synthesis.layer16.noise_strength", "synthesis.layer16.noise", "synthesis.layer16.style.weight", "synthesis.layer16.style.bias", "synthesis.output8.weight", "synthesis.output8.bias", "synthesis.output8.style.weight", "synthesis.output8.style.bias".

I know that those models have a different layer structure, but is there any way to run them with this code anyways?

Compability issue for .pkl extension

We have a model stored as .pkl file which is trained on NVIDIA's StyleGAN. Since its extension is .pkl, and your structure asks for a .pth file, is there any way to convert our .pkl file to .pth file? Or is there any way to work on your system with a .pkl file?

Help -- factorize_weight --

Hi,
Please in factorize_weight, what does generator.synthesis.getattr(layer_name).style.weight correspond to exactly ??
i am trying to reimplement your paper and my stylegan is based on this version and I can't figure out the exact correspondances
thanks in advance for your help!!

Using a custom model

Hi thanks for the release, I've been waiting for it.
So I have a stylegan2 model converted from tf to pytorch, how would I go about using it?
I am getting a keyerror when loading the converted .pt model

Traceback (most recent call last):
  File "sefa.py", line 145, in <module>
    main()
  File "sefa.py", line 70, in main
    generator = load_generator(args.model_name)
  File "J:\sefa-master\utils.py", line 101, in load_generator
    generator.load_state_dict(checkpoint['generator'])
KeyError: 'generator'

Printing the keys of my model yields dict_keys(['g_ema', 'latent_avg']) but your pretrained models yield dict_keys(['generator', 'discriminator', 'generator_smooth'])
What should I still need to do? I am able to use a closed-form-factorization implementation using my converted model from this stylegan2-pytorch repo but how to do it in this one?

I can't run this scripts

I run streamlit run interface.py, but I got Error like this
Traceback (most recent call last):
File "d:\program files\python36\lib\site-packages\streamlit\script_runner.py", line 354, in _run_script
exec(code, module.dict)
File "D:\code\sefa-master\interface.py", line 129, in
main()
File "D:\code\sefa-master\interface.py", line 75, in main
layers, boundaries, eigen_values = factorize_model(model, layer_idx)
File "d:\program files\python36\lib\site-packages\streamlit\caching.py", line 545, in wrapped_func
return get_or_create_cached_value()
File "d:\program files\python36\lib\site-packages\streamlit\caching.py", line 527, in get_or_create_cached_value
return_value = func(*args, **kwargs)
File "D:\code\sefa-master\interface.py", line 26, in factorize_model
return factorize_weight(model, layer_idx)
File "D:\code\sefa-master\utils.py", line 187, in factorize_weight
weight = generator.synthesis.getattr(layer_name).style.weight.T
AttributeError: 'Parameter' object has no attribute 'T'

I change this code to weight = generator.synthesis.getattr(layer_name).style.weight.t()
then, I still got Error like this
Traceback (most recent call last):
File "d:\program files\python36\lib\site-packages\streamlit\script_runner.py", line 354, in _run_script
exec(code, module.dict)
File "D:\code\sefa-master\interface.py", line 129, in
main()
File "D:\code\sefa-master\interface.py", line 124, in main
image = synthesize(model, gan_type, code)
File "d:\program files\python36\lib\site-packages\streamlit\caching.py", line 545, in wrapped_func
return get_or_create_cached_value()
File "d:\program files\python36\lib\site-packages\streamlit\caching.py", line 527, in get_or_create_cached_value
return_value = func(*args, **kwargs)
File "D:\code\sefa-master\interface.py", line 54, in synthesize
image = model.synthesis(to_tensor(code))['image']
File "d:\program files\python36\lib\site-packages\torch\nn\modules\module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "D:\code\sefa-master\models\stylegan_generator.py", line 471, in forward
if wp.ndim != 3 or wp.shape[1:] != (self.num_layers, self.w_space_dim):
AttributeError: 'Tensor' object has no attribute 'ndim'

Is this something wrong with my pytorch vesion or settings?

Size mismatch error

When I load weight from genforce format, it reports error as follow:

RuntimeError: Error(s) in loading state_dict for StyleGAN2Generator:
        size mismatch for synthesis.layer7.weight: copying a param with shape torch.Size([256, 512, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for synthesis.layer7.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for synthesis.layer8.weight: copying a param with shape torch.Size([256, 256, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for synthesis.layer8.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for synthesis.layer8.style.weight: copying a param with shape torch.Size([256, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for synthesis.layer8.style.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for synthesis.output4.weight: copying a param with shape torch.Size([3, 256, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 512, 1, 1]).
        size mismatch for synthesis.output4.style.weight: copying a param with shape torch.Size([256, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for synthesis.output4.style.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for synthesis.layer9.weight: copying a param with shape torch.Size([128, 256, 3, 3]) from checkpoint, the shape in current model is torch.Size([256, 512, 3, 3]).
        size mismatch for synthesis.layer9.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for synthesis.layer9.style.weight: copying a param with shape torch.Size([256, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for synthesis.layer9.style.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for synthesis.layer10.weight: copying a param with shape torch.Size([128, 128, 3, 3]) from checkpoint, the shape in current model is torch.Size([256, 256, 3, 3]).
        size mismatch for synthesis.layer10.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for synthesis.layer10.style.weight: copying a param with shape torch.Size([128, 512]) from checkpoint, the shape in current model is torch.Size([256, 512]).
        size mismatch for synthesis.layer10.style.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for synthesis.output5.weight: copying a param with shape torch.Size([3, 128, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 256, 1, 1]).
        size mismatch for synthesis.output5.style.weight: copying a param with shape torch.Size([128, 512]) from checkpoint, the shape in current model is torch.Size([256, 512]).
        size mismatch for synthesis.output5.style.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for synthesis.layer11.weight: copying a param with shape torch.Size([64, 128, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 256, 3, 3]).
        size mismatch for synthesis.layer11.bias: copying a param with shape torch.Size([64]) from checkpoint, the shape in current model is torch.Size([128]).
        size mismatch for synthesis.layer11.style.weight: copying a param with shape torch.Size([128, 512]) from checkpoint, the shape in current model is torch.Size([256, 512]).
        size mismatch for synthesis.layer11.style.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for synthesis.layer12.weight: copying a param with shape torch.Size([64, 64, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 128, 3, 3]).
        size mismatch for synthesis.layer12.bias: copying a param with shape torch.Size([64]) from checkpoint, the shape in current model is torch.Size([128]).
        size mismatch for synthesis.layer12.style.weight: copying a param with shape torch.Size([64, 512]) from checkpoint, the shape in current model is torch.Size([128, 512]).
        size mismatch for synthesis.layer12.style.bias: copying a param with shape torch.Size([64]) from checkpoint, the shape in current model is torch.Size([128]).
        size mismatch for synthesis.output6.weight: copying a param with shape torch.Size([3, 64, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 128, 1, 1]).
        size mismatch for synthesis.output6.style.weight: copying a param with shape torch.Size([64, 512]) from checkpoint, the shape in current model is torch.Size([128, 512]).
        size mismatch for synthesis.output6.style.bias: copying a param with shape torch.Size([64]) from checkpoint, the shape in current model is torch.Size([128]).

After I check the stylegan2_generator.py, I found that the default value of 'fmaps_base' is said to be 16 << 10, but actually it is 32 << 10 in the python code. When I change the 'fmaps_base' into 16 << 10, the code works?

However, in the genforce (where I obtain the model weight), the 'fmaps_base' is 32 << 10! It seems that the same parameter does not work for the same model??

Can anyone tell me what happens? Thank you very much!

About attribute predictor and re-scoring analysis.

Dear @ShenYujun :

Thanks for your excellent works. It is very useful for my research and I will cite your work in my study. However, I have some issues with re-scoring analysis since the attribute predictor is lost (The same issue on InterfaceGAN). Would you mind releasing your pre-trained attribute predictor's code and model?

Thank you again.
Best wishes.

BigGAN

Hi,

Would it be possible to release the model used to produce the BigGAN results in the paper?

Thanks,
Sagie

A question about the assumption of the method

Hi,

I have a question that whether we can apply this method for any layers of generator? Do you try any experiments to explore it?
In the paper, you only show theoretical proof on the mapping from the latent code to the first linear layer of the generator, but ignoring some following non-linear and linear layers.

Thanks!

StyleGAN3 Support

This is really interesting work, thanks for publishing.

Do you have plans to support StyleGAN3 in this repository?

About ignore the non-linear layers of the mapping net.

Thanks for your excellent work! but the mapping net consists of 8 FC layers with 'leak_relu' in between. But the code seems to concatenate the weight of each layer as the final weight, but ignores the 'leak_relu' layers. Why ignore them?

Direction indices for Closed Form applied on PGGAN

Thanks for the great work. Could you please let us know the indices of the closed form directions corresponding to Smiling, Pose, Eyeglasses, Age and Gender obtained from PGGAN on CelebA-HQ dataset?

About code release schedule

I found that paper is really innovative. I want to try building something on top of it.
Can I ask about the estimated released date of the code?

Thank you a lot.

how to directly use checkpoint from stylegan2-ada-pytorch

Hi, thanks for your great work and easy to use streamlit app. The current available checkpoint has impressed me a lot, however, i find it hard to use other checkpoint from nvlab official repositorys such as stylegan2-ada-pytorch. Detail list following:

different checkpoint type

genfporce checkpoint

I try to download some checkpoint from 'model zoo', the checkpoint end with .pth and load through torch.load()

stylegan2-ada-checkpoint

My trained checkpoint end with .pkl. It seems to be saved in nvlab's unique method.

my effort

• Directly use torch.load() to load checkpoint

• Search for some method to convert the pkl checkpoint in genforce repository issues.
  But I failed to find a resolution.

StyleGAN2 LSUN bedroom

Dear GenForce team,

I really like your work . It is a very nice paper.

In figure 1, you show manipulation in StyleGAN2 LSUN bedroom, but this model is not in model zoo. Would you mind sharing it with us?

I know that in model zoo, there are bedroom model for pggan and StyleGAN1, but I want to try it on StyleGAN2.

Thank you very much for your help.

Best Wishes,

Alex

Edit face in Eyeglasses, Gender, Hair Color, Pose , Smile ... direction failed

I edited face in Eyeglasses, Gender, Hair Color, Pose , Smile ... direction by MODEL_ZOO/StyleGAN2/ffhq-1024x1024

it shown that some attributes/directions are highly correlated with each other .

reproduce results matched figure 5 in paper

Hi, thanks for your great research first. May I have the seed to generate the source StyleGAN image which used in figure 5 of the paper? And also the index of sorted eigen value corresponding to "Pose", "Eyeglasses" and "Smile" manipulation so that I can reproduce results in figure 5.

By the way, the InterFaceGAN result in qualitative comparison were produced in W space or W+ space?

Thanks.

Question about factorize_weight()

Hi, I first like to thank you for your great research, and generous release of code.
I'm using your approach on a research using styleGAN.
while looking at your code, I realized that on concatenating weights to get A from the paper,
last layer is from output layer(like synthesis.output7.style.weight),
unlike others which are from style weight layers (like synthesis.layer14.style.weight)

I was wondering why last output layer was included in constructing A,
as some other unofficial implements which do not include output layer seems to work as well.
If there is specific reason why output layer was included, I'd like to know.

(I understand such curiosity may have come from my lack of knowledge of GAN structure.
On such case, I would be also grateful even if you point it out )

Thank you for your kind response!

streetscapes stylegan 2 model

Is there a pretrained stylegan 2 model on streettscape dataset