UniPC for diffusion sampling

[nicksenger]

Hi, thanks for the awesome library! It's great having these tools available in the Rust ecosystem.

I was interested in low step-count inference for some experiments, so ported over the UniPC scheduler. I figured I'd open a PR here in case this work is useful to others. Here is a comparison with Euler A and DDIM at 5 steps which demonstrates UniPC's benefits for quick convergence:

[LaurentMazare]

Would you have some examples with good quality where this makes a difference?

[nicksenger]

Sure, here is a comparison of DDIM (default cfg) and UniPC (corrector enabled past step 2, Bh2 solver type, otherwise defaults) on sd1.5 across 50 steps with the following prompt with seed 1984 and guidance scale 15:

a rusty robot holding a fire torch in its hand,android scrapyard,oxidized,intricate exposed wires,arcing,neon lights,sci-fi,dystopia,futuristic city background,silhouettes,strange illuminated mannequins,geisha billboard,flying cars,4k,cinematic lighting,photo-realistic,extremely detailed,high quality,epic,lfg

negative prompt:

reduction,reducing agents,electron rich,drawing,illustrated,low quality,out of focus,blurry,cut off,simple,clean,organized,daytime,utopia,peaceful,low contrast,lowres,worst quality

6 steps is about where the output starts becoming clear from both schedulers, but the composition is quite different:

At 50 steps both have converged to a similar output with some differences. Which one looks better is a bit subjective, but I think most would agree the composition is closer to the 6-step output from UniPC than that from DDIM:

Just for fun here's the 50-step output from UniPC with the same settings as above but using 3 for the solver order:

Somewhat different output again, but overall the composition pretty much resembles what it produces at 5 steps:

[LaurentMazare]

Could you give some insights on the ordering that this provides?

[nicksenger]

I added a comment in fedce6287405fe04826344c3b5537375b317760a, the ordering is:

NaN | -Infinity | x < 0 | -0 | +0 | x > 0 | +Infinity | NaN

It's the same strategy used by float-ord, which is in turn a dependency of the average crate where this quantile computation comes from. These are only used for the dynamic thresholding, so I thought it'd be better to hide the logic within this module versus introducing additional dependencies crate-wide.

[LaurentMazare]

Agreed that it's better not to introduce a dependency but why not use the total ordering f64::total_cmp? Do you expect some differences with it that would be helpful here?

[nicksenger]

Ah, nope. This was just an oversight on my part, I thought the msrv for the project was lower for some reason, but now I'm seeing total_cmp is used in quite a few other places. It looks like the standard lib uses the same method: https://doc.rust-lang.org/src/core/num/f64.rs.html#1350

Updated in f0384fa

[LaurentMazare]

Thanks for the detailed analysis, I've put a bunch of comments inline, please also look at the clippy and rustfmt failures in the CI.

[nicksenger]

Thanks for the detailed analysis, I've put a bunch of comments inline, please also look at the clippy and rustfmt failures in the CI.

Thanks, these should be addressed now.

[LaurentMazare]

Thanks!

[super-fun-surf]

stoked!