Add multistep DPM-Solver discrete scheduler

[LuChengTHU]

Add the multistep version of DPM-Solver, accepting discrete time steps as DDPM, DDIM, and PNDM.

The solver can successfully run with the stable-diffusion pipeline for both pytorch and jax versions.

Currently, it supports the following algorithms:

• multistep DPM-Solver (for discretizing the integral w.r.t. noise prediction model) with order = 1, 2, 3
• multistep DPM-Solver++ (for discretizing the integral w.r.t. data prediction model) with order = 1, 2, 3

For more details of the algorithms, please refer to https://github.com/LuChengTHU/dpm-solver

#1101

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patil-suraj]

Great work @LuChengTHU, super cool PR!

The PR already looks super good, I just left some nits.

Also tried it out with stable diffusion and it works really well, and generates really good output with 20-25 steps, for some examples even 10-15 🤯 !

QQ: What are the recommended setting for stable diffusion, I used the following and it worked super well. We will need to update these in the config of sd checkpoints.

dpm = DPMSolverDiscreteScheduler.from_config("CompVis/stable-diffusion-v1-4", subfolder="scheduler", solver_order=2, predict_x0=False, denoise_final=True)

And is this scheduler different from the DPM version which requires two model evaluation per step ?

We are still thinking about an API for schedulers that require two model evaluations. Will have a draft PR next week. Then we can add the other version of DPM once the API is finalised.

Also, to differentiate it from the other scheduler could we call this DPMMultiStepScheduler ?

[LuChengTHU]

Great work @LuChengTHU, super cool PR!

The PR already looks super good, I just left some nits.

Also tried it out with stable diffusion and it works really well, and generates really good output with 20-25 steps, for some examples even 10-15 🤯 !

QQ: What are the recommended setting for stable diffusion, I used the following and it worked super well. We will need to update these in the config of sd checkpoints.

dpm = DPMSolverDiscreteScheduler.from_config("CompVis/stable-diffusion-v1-4", subfolder="scheduler", solver_order=2, predict_x0=False, denoise_final=True)

And is this scheduler different from the DPM version which requires two model evaluation per step ?

We are still thinking about an API for schedulers that require two model evaluations. Will have a draft PR next week. Then we can add the other version of DPM once the API is finalised.

Also, to differentiate it from the other scheduler could we call this DPMMultiStepScheduler ?

Thank you very much for the quick and high-quality review!

Q1: What are the recommended settings for stable diffusion?

I use the following setting, which works similarly to your settings:

dpm = DPMSolverMultistepScheduler.from_config(
    "CompVis/stable-diffusion-v1-4",  # or use the v1-5 version
    subfolder="scheduler",
    solver_order=2,
    predict_epsilon=True,
    thresholding=False,
    algorithm_type="dpmsolver++",
    solver_type="midpoint",
    denoise_final=True,  # the influence of this trick is effective for small (e.g. <=10) steps
)

Q2: is this scheduler different from the DPM version, which requires two model evaluations per step?

Yes, they are different. The DPM version you mentioned is actually the singlestep DPM-Solver in my repo, which is proposed by my paper "DPM-Solver". The currently committed PR version is the multistep DPM-Solver in my repo, which is proposed by my other paper "DPM-Solver++".

In the paper "DPM-Solver++", I carefully compare the different settings of the solvers for guided sampling (i.e. conditional sampling) by diffusion models, and I find that:

0. All previous high-order solvers are unstable and fail to generate samples for large guidance scales and small steps (<20).

1. We further proposed algorithms that discretize the integral w.r.t. the data prediction model, and can greatly stabilize the sample quality for small steps (<20). We use it as algorithm_type="dpmsolver++".

2. We find that for large guidance scales, higher-order solvers are extremely unstable. So we only use the 2nd-order solver for guided sampling.

3. We find that the multistep version is slightly better than the singlestep version.

Therefore, I suggest using the 2nd-order multistep DPM-Solver++, which is the above default settings.

However, the 3rd-order method is still useful for unconditional sampling, and it can achieve a better sample quality than the 2nd-order method for unconditional sampling. So I also provide a 3rd-order method in this PR.

References:
[1] DPM-Solver: https://arxiv.org/abs/2206.00927
[2] DPM-Solver++: https://arxiv.org/abs/2211.01095

Q3: We are still thinking about an API for schedulers that require two model evaluations. Will have a draft PR next week. Then we can add the other version of DPM once the API is finalised.

Great to know it! I can help to add the corresponding singlestep DPM-Solver at that time.

Q4: To differentiate it from the other scheduler could we call this DPMMultiStepScheduler ?

I've changed the name by "DPMSolverMultistepScheduler"

Q5: other modifications.

Thanks for your careful review! I will fix them in the latter commits.

[patil-suraj]

Thanks a lot for the detailed answer and addressing the review comments @LuChengTHU ! Good to merge for me :)

@patrickvonplaten or @anton-l would be nice if you could also take quick look :)

[patrickvonplaten]

Amazing addition @LuChengTHU!

The API choices are great and it's amazing that it fits so well with the existing API of Stable Diffusion :-)

Only left a couple of nits:

• suggestion to rename predict_x0 to predict_epsilon for consistency
• Add the DPM-Solver scheduler as a compatible scheduler to all other schedulers

[LuChengTHU]

Hi @patrickvonplaten, thank you for the quick reviews! I've fixed all the modifications. The main difference is that I've changed the APIs for predict_epsilon, algorithm_type, and solver_type to a clearer way (and also changed the docs).

Looking forward to your reply!

[patil-suraj]

Thank you for addressing the comments, let's go!