ECT: Consistency Models Made Easy

Pytorch implementation for Easy Consistency Tuning (ECT).

ECT unlocks SoTA few-step generative abilities through a simple yet principled approach. With a negligible tuning cost, ECT demonstrates promising early results while benefiting from the scaling to improve its few-step generation capability.

Try your own Consistency Models! You only need to fine-tune a bit. :D

Introduction

v This branch provides a minimal setup to tune CMs on the ImageNet 64x64 dataset. The dataset preparation follows the format used by ADM and EDM2.

Environment

You can run the following command to set up the Python environment through conda. Pytorch 2.3.0 and Python 3.9.18 will be installed.

conda env create -f env.yml

Datasets

To prepare the ImageNet 64x64 dataset, please follow the steps below:

1. Download the the ILSVRC2012 data archive from Kaggle and extract it.
2. Clone the EDM2 repository

   • git clone https://github.com/NVlabs/edm2.git
     mv data_prep.sh edm2

3. Configure paths in data_prep.sh.
   • DATA_DIR: Set this to the path where you saved the extracted ImageNet train set.
   • DEST_DIR: Set this to the directory where you want to store the processed ImageNet 64x64 dataset.
4. Run bash data_prep.sh.
5. Link your zip dataset to the local datasets directory.

   mkdir datasets
   ln -s $DEST_DIR/edm2-imagenet-64x64.zip ./datasets/

Training

Run the following command to run ECT at batch size 128 and 100k iterations. NGPUs=4/8 is recommended.

bash run_imgnet_ecm.sh <NGPUs> <PORT> --desc bs128.100k

• Replace <NGPUs> and <PORT> with the number of GPUs used for training and the port number for DDP sync.
• Modify --preset to the model configs to launch.

Evaluation

Run the following command to calculate FID of a pretrained checkpoint.

bash eval_imgnet_ecm.sh <NGPUs> <PORT> --resume <CKPT_PATH> 

• Use --resume_pkl to specify a snapshot (.pkl) or the directory containing multiple snapshots for evaluation.
• Use --resume to specify a checkpoint (.pt) or the directory containing multiple checkpoints for evaluation.

Checkpoints

• ImgNet 64x64 ECM-XL checkpoints trained at bs1024 for 100k iterations.

Contact

Feel free to drop me an email at [email protected] if you have additional questions or are interested in collaboration. You can also find me on Twitter or WeChat.

Citation

@article{ect,
  title={Consistency Models Made Easy},
  author={Geng, Zhengyang and Pokle, Ashwini and Luo, William and Lin, Justin and Kolter, J Zico},
  journal={arXiv preprint arXiv:2406.14548},
  year={2024}
}