Dogyun Park, Sojin Lee, Sihyeon Kim, Taehoon Lee, Youngjoon Hong†, Hyunwoo J. Kim†.

NeurIPS 2024

This repository is an official implementation of "Constant Acceleration Flow" accepted at NeurIPS 2024.

Constant Acceleration Flow (CAF) introduces a novel ODE framework that models the transformation between two distributions using constant acceleration. This approach enables high-quality image generation in a single step. Notably, CAF also achieves high-accuracy single-step inversion, demonstrating its capability to learn precise forward and reverse flows.

• Release training code.
• Release inference (sampling and inversion) code.
• Release the pre-trained model.

FID=1.69

N=1

To install requirements, run:

git clone https://github.com/mlvlab/CAF.git
cd CAF
conda create -n caf python==3.9
conda activate caf
conda install pytorch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 pytorch-cuda=11.8 -c pytorch -c nvidia

pip install -e .

(RECOMMENDED, linux) Install PyTorch 2.2.0 with CUDA 11.8 for xformers, recommended for memory-efficient computation.

We have utilized pre-trained EDM models to generate deterministic couplings for training our models. To construct deterministic couplings for CIFAR-10 and ImageNet, we selected $N=18$ and $N=40$, respectively, employing deterministic sampling following the protocol outlined in the EDM repository. For CIFAR-10 and ImageNet, we generated 1M and 3M pairs, respectively, and saved them in data/ folder as NumPy (.npy) format. Each dataset should be organized in the following directory structure:

data
|-- cifar10_npy (imagenet_npy)
    |-- z0
        |-- data1.npy
        |-- data2.npy
        |-- ...
    |-- z1
        |-- data1.npy
        |-- data2.npy
        |-- ...
    |-- class
        |-- data1.npy
        |-- data2.npy
        |-- ...

It is important to ensure that the z0, z1, and class files correspond to each other by sharing the same filenames. Each triplet (data1.npy in z0, z1, and class) should be aligned to represent a single data pair.

We utilized real datasets for adversarial training. You can download ImageNet 64x64 following the download script here. Place the downloaded dataset in the data/imagenet folder.

You can update the training config in flow/script_util.py.

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --multi_gpu --num_processes=8 \
    script/caf_train_vel.py

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --multi_gpu --num_processes=8 \
    script/caf_train_acc.py

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --multi_gpu --num_processes=8 \
    script/caf_train_gan.py 

You can update the inference config in flow/script_util.py.

python script/inference.py --sample_step [N] --save_image True \
     --save_dir [path to model]

python scripts/evaluate_reconstruction.py --sample_step [N] --save_image True \
     --save_pth [path to model]

This repo is built upon Consistency Models and CTM.

If you have any questions, please create an issue on this repository or contact at [email protected].

If you find our work interesting, please consider giving a ⭐ and citation.

@article{park2024constant,
  title={Constant Acceleration Flow},
  author={Park, Dogyun and Lee, Sojin and Kim, Sihyeon and Lee, Taehoon and Hong, Youngjoon and Kim, Hyunwoo J},
  journal={arXiv preprint arXiv:2411.00322},
  year={2024}
}