Asymmetric VQGAN

Designing a Better Asymmetric VQGAN for StableDiffusion

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:

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

model	pretrain	resolution	fid	lpips	pre_error
StableDiffusion + vanilla VQGAN	ImageNet-1K	224x224	9.57	0.255	1082.8e^-5
StableDiffusion + asymmetric VQGAN (base)	ImageNet-1K	224x224	7.60	0.137	5.7e^-5
StableDiffusion + asymmetric VQGAN (Large 1.5x)	ImageNet-1k	224x224	7.55	0.136	2.6e^-5
StableDiffusion + asymmetric VQGAN (Large 2x)	ImageNet-1k	224x224	7.49	0.134	2.1e^-5

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

model	fid	is
StableDiffusion + vanilla VQGAN	19.88	37.55
StableDiffusion + asymmetric VQGAN (base) w/o mask	19.92	37.52
StableDiffusion + asymmetric VQGAN (Large 1.5x) w/o mask	19.75	37.64
StableDiffusion + asymmetric VQGAN (Large 2x) w/o mask	19.68	37.73

Comments

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

treksis / asymmetric_vqgan Goto Github PK