zju-pi / diff-sampler

An open-source toolbox for fast sampling of diffusion models. Official implementations for our [CVPR-2024, ICML-2024] papers

diff-sampler

diff-sampler is an open-source toolbox for fast sampling of diffusion models, to provide a fair comparison of existing approaches and help researchers to develp better approaches. diff-sampler contains various model implementations, numerical-based solvers, time schedules, and other features.

This repository also includes (or will include) the official implementations of our following works:

• [arXiv 2023] A Geometric Perspective on Diffusion Models
• [CVPR 2024] Fast ODE-based Sampling for Diffusion Models in Around 5 Steps
• [ICML 2024] On the Trajectory Regularity of ODE-based Diffusion Sampling

News

• 2024-06-20. GITS is integrated into ComfyUI as GITSScheduler! We provide a ComfyUI workflow to compare GITS with the baseline time schedule and the concurrent ICML 2024 paper Align Your Steps: Optimizing Sampling Schedules in Diffusion Models. More details are included in our pull request here.
  
• 2024-06-18. The official implementation of the paper On the Trajectory Regularity of ODE-based Diffusion Sampling is released in gits-main.
  • This new repo is built on the basis of diff-solvers-main, supporting baseline sampling as well as customized time schedule. Our method (GITS) offers a more efficient time schedule for fast sampling of diffusion models.
  • For text-to-image generation, we provide a colab script using Diffusers 🧨 where we also compare our methods with the concurrent work Align Your Steps. Below is a 10-step comparison.
  
• 2024-06-03. The repo is made easier to use. Now the pre-trained models will be automatically downloaded to ./src/dataset_name. Some errors and typos are fixed. Detailed running scripts are provided in launch.sh, where we also add new scripts for evaluation of CLIP score for Stable Diffusion.
• 2024-05-02. Our work On the Trajectory Regularity of ODE-based Diffusion Sampling is accepted by ICML 2024.
• 2024-04-05. Our work Fast ODE-based Sampling for Diffusion Models in Around 5 Steps is selected as a Highlight paper (Top 2.8%) by CVPR 2024!
• 2024-03-25. The official implementation of the paper Fast ODE-based Sampling for Diffusion Models in Around 5 Steps (amed-solver-main) and a toolbox (diff-solvers-main) for fast sampling of diffusion models is released. We upload the reference statistics for FID evaluation here. Hope that this repo can facilitate researchers on fast sampling of diffusion models!
• 2024-02-27. Our work Fast ODE-based Sampling for Diffusion Models in Around 5 Steps is accepted by CVPR 2024.

TODO

• Give an example of AMED-Plugin on Stable Diffusion using Diffusers 🧨.
• Integrate the GITS time schedule into ComfyUI and provide a ComfyUI workflow.
• Release the code of A Geometric Perspective on Diffusion Models for easier analysis of Diffusion Models.

Supported Fast Samplers for Diffusion Models

Name	Max Order	Source	Location
Euler	1	Denoising Diffusion Implicit Models	diff-solvers-main
Heun	2	Elucidating the Design Space of Diffusion-Based Generative Models	diff-solvers-main
DPM-Solver-2	2	DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps	diff-solvers-main
DPM-Solver++	3	DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models	diff-solvers-main
UniPC	3	UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models	diff-solvers-main
DEIS	4	Fast Sampling of Diffusion Models with Exponential Integrator	diff-solvers-main
iPNDM	4	Fast Sampling of Diffusion Models with Exponential Integrator	diff-solvers-main
iPNDM_v	4	The variable-step version of the Adams–Bashforth methods	diff-solvers-main
AMED-Solver	2	Fast ODE-based Sampling for Diffusion Models in Around 5 Steps	amed-solver-main
AMED-Plugin	-	Fast ODE-based Sampling for Diffusion Models in Around 5 Steps	amed-solver-main
GITS	-	On the Trajectory Regularity of ODE-based Diffusion Sampling	gits-main

Citation

If you find this repository useful, please consider citing the following paper (reverse chronological order):

@article{chen2024trajectory,
  title={On the Trajectory Regularity of ODE-based Diffusion Sampling},
  author={Chen, Defang and Zhou, Zhenyu and Wang, Can and Shen, Chunhua and Lyu, Siwei},
  journal={arXiv preprint arXiv:2405.11326},
  year={2024}
}

@article{zhou2023fast,
  title={Fast ODE-based Sampling for Diffusion Models in Around 5 Steps},
  author={Zhou, Zhenyu and Chen, Defang and Wang, Can and Chen, Chun},
  journal={arXiv preprint arXiv:2312.00094},
  year={2023}
}

@article{chen2023geometric,
  title={A geometric perspective on diffusion models},
  author={Chen, Defang and Zhou, Zhenyu and Mei, Jian-Ping and Shen, Chunhua and Chen, Chun and Wang, Can},
  journal={arXiv preprint arXiv:2305.19947},
  year={2023}
}