feat: add gits scheduler

[zhyzhouu]

The geometry-inspired time scheduling (GITS, ICML'24) provides a series of time schedules (with a hyperparameter 'coeff') for fast sampling of diffusion models. The implementation follows that of nodes_align_your_steps.py. For steps <= 20 we provide GITS schedules and for steps > 20 we follow nodes_align_your_steps.py to perform a log-linear interpolation based on the 20-step GITS schedule.

Below is our workflow for comparison of these three schedules:

[ltdrdata]

This seems interesting.

[comfyanonymous]

can you tell me where these values are from in the reference code or how you calculated them?

[zhyzhouu]

can you tell me where these values are from in the reference code or how you calculated them?

TL;DR

To avoid extra computational overhead for users, in the submitted script, we provide the sigmas pre-calculated on our device (requiring only several minutes using 4 A100 GPUs). Our proposed GITS utilizes the regularity of the sampling trajectories of diffusion models (they all look like a “boomerang” curve) and finds optimal time schedules using dynamic programming with a hyperparameter 'coeff'. Dynamic programming is performed on a cost matrix, where each cost evalutes the difference between an Euler step and the ground-truth prediction from the current position to the next position.

Derivation of GITS

In our ICML'24 paper On the Trajectory Regularity of ODE-based Diffusion Sampling, we illustrate that different sampling trajectories of diffusion models share almost the same shape structure (a “boomerang” curve), which inspires us to search a better time schedule for ODE-based diffusion sampling. Our approach is named as geometry-inspired time scheduling (GITS).

We first define a searching space denoted as $\Gamma_{g}$, which is a fine-grained grid including all possible intermediate timestamps. We use $\Gamma_{g}$ to generate a 'ground-truth' sampling trajectory and then measure the trajectory curvature by the local truncation errors. Specifically, we define the cost from the current position $x_{t_i}$ to the next position $x_{t_j}$ as the difference between an Euler step and the ground-truth prediction. Given all computed pair-wise costs, which form a cost matrix, this becomes a standard minimum-cost path problem and can be solved with dynamic programming. Additionally, as different trajectories share almost the same shape, we can use just a few “warmup” samples (256 by defualt) to estimate the appropriate schedule. We further introduce a hyperparameter $\gamma$ to control the schedule. The algorithm is shown below (see our paper for more details).

How to obtain these schedules

We run the following command in our repo which takes several minutes using 4 A100 GPUs:

SOLVER_FLAGS="--solver=dpmpp --num_steps=21 --afs=False"
SCHEDULE_FLAGS="--schedule_type=discrete --schedule_rho=1"
ADDITIONAL_FLAGS="--max_order=2 --predict_x0=False --lower_order_final=True"
GUIDANCE_FLAGS="--guidance_type=cfg --guidance_rate=7.5"
GITS_FLAGS="--dp=True --metric=dev --coeff=1.2 --num_steps_tea=101 --solver_tea=dpmpp"
torchrun --standalone --nproc_per_node=4 --master_port=11111 \
sample.py --dataset_name="ms_coco" --batch=16 --seeds="0-15" $SOLVER_FLAGS $SCHEDULE_FLAGS $ADDITIONAL_FLAGS $GUIDANCE_FLAGS $GITS_FLAGS

This command is used to generate GITS schedules for Stable Diffusion v1-5 using DPM-Solver++(2M) for sampling. The GITS_FLAGS includes the settings for GITS where we set 100-step trajectory as the ground-truth. This command for Stable Diffusion will by default generate all time schedules at one time (including indices and sigmas), for each coefficient $\gamma \in [0.80, 1.20]$ with a step size of $0.05$ and different number of sampling steps $\in [3,20]$ (thanks to the optimal substructure property of dynamic programming). The generated schedules will be saved to ./dp_record.txt.

To avoid extra computational overhead for users, in the submitted script, we provide the sigmas generated for $\gamma \in [0.80, 1.50]$ and steps $\in [3,20]$. Other values can be similarly obtained by modifying the above command.

[not-ski]

I see that the examples all use SD1.5; is SDXL supported? And if so, do you still recommend the default "coeff" value of 1.20? Thanks! <3

[zhyzhouu]

I see that the examples all use SD1.5; is SDXL supported? And if so, do you still recommend the default "coeff" value of 1.20? Thanks! <3

Thanks for your attention! Due to limited resources, we have not yet conducted experiments on SDXL. We will test our method on SDXL in the next few days. We speculate that the schedule will not differ significantly, and the current schedule for SD1.5 should be applicable to SDXL.

The larger the "coeff", the more biased the schedule is towards 0. So according to the schedules provided in nodes_align_your_steps.py, a default value of 1.20 or 1.25 should be fine.