Accelerating Diffusion Transformer via Error-Optimized Cache

Requirement

With pytorch(>2.0) installed, execute the following command to install necessary packages

pip install accelerate diffusers timm torchvision wandb

Prior Knowledge Extraction

Change the cache matrix in ~/EOC/DiT/models/models.py to matrix1

python 1-blocks.py
python 2-block_pth_average.py

Sample Image

For learning-to-cache DDIM-20 + EOC:

python sample.py --model DiT-XL/2 --num-sampling-steps 20 --ddim-sample --accelerate-method dynamiclayer --path ckpt/DDIM20_router.pt --thres 0.1

For Rule-based Method + EOC Change the cache matrix in ~/EOC/DiT/models/models.py to matrix3

python sample.py --model DiT-XL/2 --num-sampling-steps 20 --ddim-sample 

Sample 50k images for Evaluation

If you want to reproduce the FID results from the paper, you can use the following command to sample 50k images:

For learning-to-cache DDIM-20 + EOC:

torchrun --nnodes=1 --nproc_per_node=8 --master_port 12345 sample_ddp.py --model DiT-XL/2 --num-sampling-steps NUM_STEPS --ddim-sample --accelerate-method dynamiclayer --path PATH_TO_TRAINED_ROUTER --thres 0.1

For Rule-based Method + EOC:

torchrun --nnodes=1 --nproc_per_node=8 --master_port 12345 sample_ddp.py --model DiT-XL/2 --num-sampling-steps NUM_STEPS --ddim-sample --path PATH_TO_TRAINED_ROUTER --thres 0.1

Be sure to modify NUM_STEPS and PATH_TO_TRAINED_ROUTER to correspond to the respective NFE steps and the location of the router.

Calculate FID

We follow DiT to evaluate FID by the code. Please install the required packages, download the pre-computed sample batches, and then run the following command:

python evaluator.py ~/ckpt/VIRTUAL_imagenet256_labeled.npz PATH_TO_NPZ

BibTeX

@article{qiu2025accelerating,
  title={Accelerating Diffusion Transformer via Error-Optimized Cache},
  author={Qiu, Junxiang and Wang, Shuo and Lu, Jinda and Liu, Lin and Jiang, Houcheng and Zhu, Xingyu and Hao, Yanbin},
  journal={arXiv preprint arXiv:2501.19243},
  year={2025}
}

Acknowledgement

This implementation is based on DiT and Learning-to-cache.