search

crowsonkb / k-diffusion

Karras et al. (2022) diffusion models for PyTorch
MIT License
2.21k stars 371 forks source link

Feature Request: Support the "correct" SDE variant of DPM-Solver++ #59

Closed LuChengTHU closed 1 year ago

LuChengTHU commented 1 year ago

Hi crowsonkb, recently I'm studying the SDE solvers and I find that the current implementation of "DPM++2M SDE" is not exactly the solver of the reverse diffusion SDE.

I've implemented the correct version of SDE-variant DPM++2M, and its performance is quite amazing in DeepFloyd-IF. Could you please check this PR in huggingface diffusers: https://github.com/huggingface/diffusers/pull/3344 and implement it in your own repo? Thank you so much!

LuChengTHU commented 1 year ago

Here are some examples with deepfloyd-if, where the stage-1 only uses 10 steps! (under the same random seed)

Stage-1: 10-step "sde-dpmsolver++" or "dpmsolver++" or DDPM. Stage-2: 25-step "dpmsolver++". Stage-3: 25-step "dpmsolver++".

I recommend using "sde-dpmsolver++" in stage-1, and "dpmsolver++" in stage-2 and stage-3.

My implemented "sde-dpmsolver++":

2_sde-dpm++_10 3_sde-dpm++_10

Official DDPM:

2_ddpm_10 3_ddpm_10

The DPMSolverSDEScheduler which reproduces your implemented "DPM++2M SDE":

2_k-sde_10 3_k-sde_10

crowsonkb commented 1 year ago

Oh that's interesting, the reverse diffusion SDE does better than ancestral sampling based solvers like mine for higher order SDE solvers? I need to look at this again, I had stopped trying anything reverse diffusion SDE based.

crowsonkb commented 1 year ago

Ohh, this one is a linear multistep solver! Mine was a single step solver. Here's a WIP gist: https://gist.github.com/crowsonkb/3ed16fba35c73ece7cf4b9a2095f2b78

LuChengTHU commented 1 year ago

Amazing that you implement it so fast! So in you own experience, does this solver outperform the previous single step SDE solver?

crowsonkb commented 1 year ago

https://github.com/crowsonkb/k-diffusion/commit/962d62bc73444d005e5828f11ae5b053f287e12e :)