Awesome

Asymmetric VQGAN

Designing a Better Asymmetric VQGAN for StableDiffusion<br/>

Introduction

We propose the Asymmetric VQGAN, to preserve the information of conditional image input. Asymmetric VQGAN involves two core designs compared with the original VQGAN as shown in the figure. First, we introduce a conditional branch into the decoder of the VQGAN which aims to handle the conditional input for image manipulation tasks. Second, we design a larger decoder for VQGAN to better recover the losing details of the quantized codes.

Top: The inference process of our symmetric VQGAN. Bottom: The inference process of vanilla VQGAN.

• Our pre-trained models are available:
  • A large 2x deocder
  • A large 1.5x decoder
  • A base decoder

Visualization Results

• results on inpainting task

• results on text2image task

Requirements

pip install -r requirements.txt
pip install wandb
pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
pip install -e git+https://github.com/openai/CLIP.git@main#egg=clip

Pretrained diffusion Models

The inpainting model sd-v1-5-inpainting.ckpt of StableDiffusion is here

The text2image model v1-5-pruned-emaonly.ckpt of StableDiffusion is here

Inpainting task

Download our images and masks .

python inpaint_st.py  --config {config_spec}

where config_spec is one of {autoencoder_kl_32x32x4.yaml(base decoder), autoencoder_kl_32x32x4_large.yaml(large decode 1.5x),

autoencoder_kl_32x32x4_large2.yaml(large decoder 2x).

Main Results on ImageNet

Text2image task

python txt2img.py --plms --config_c {config_spec}

where config_spec is one of {autoencoder_kl_woc_32x32x4.yaml(base decoder), autoencoder_kl_woc_32x32x4_large.yaml(large decode 1.5x),

autoencoder_kl_woc_32x32x4_large2.yaml(large decoder 2x).

Main Results on MSCOCO

Train your own asymmetric vqgan

Data preparation

ImageNet

The code will try to download (through Academic Torrents) and prepare ImageNet the first time it is used. However, since ImageNet is quite large, this requires a lot of disk space and time. If you already have ImageNet on your disk, you can speed things up by putting the data into ./datasets/ImageNet/train. It should have the following structure:

./datasets/ImageNet/train/
├── n01440764
│   ├── n01440764_10026.JPEG
│   ├── n01440764_10027.JPEG
│   ├── ...
├── n01443537
│   ├── n01443537_10007.JPEG
│   ├── n01443537_10014.JPEG
│   ├── ...
├── ...

Training autoencoder models

First, download weights of the autoencoder stable_vqgan.ckpt obtained from StableDiffusion.

Second, input your own key of wandb in main.py (line 679).

Configs for training a KL-regularized autoencoder on ImageNet are provided at configs/autoencoder. Training can be started by running

python main.py --base configs/autoencoder/{config_spec} -t --gpus 0,1,2,3,4,5,6,7 --tag <yourtag>   

where config_spec is one of {autoencoder_kl_32x32x4_train.yaml(base decoder), autoencoder_kl_32x32x4_large_train.yaml(large decode 1.5x),

autoencoder_kl_32x32x4_large2_train.yaml(large decoder 2x). It is worth noting that the parameter num_gpus in config_spec is still needed to be set as the same as the number of gpus which you use.

Comments

• Our codebase for the diffusion models builds heavily on StableDiffusion. Thanks for open-sourcing!

BibTeX

@misc{zhu2023designing,
      title={Designing a Better Asymmetric VQGAN for StableDiffusion}, 
      author={Zixin Zhu and Xuelu Feng and Dongdong Chen and Jianmin Bao and Le Wang and Yinpeng Chen and Lu Yuan and Gang Hua},
      year={2023},
      eprint={2306.04632},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}