Add custom k_diffusion scheduler

[KohakuBlueleaf]

Describe what this pull request is trying to achieve.

#10621 implement this feature request.
I use it as hires fix, only make the options visible if it is enabled.
I didn't remove the "XXX karras" sampler since user may not know what is the sigma_max/sigma_min/rho for them.
This custom scheduler can be used on all the k_diffusion sampler

Additional notes and description of your changes

I add some new arguments for txt2img and img2img, I don't know if I should use this way or another api/hook.

Environment this was tested in

List the environment you have developed / tested this on. As per the contributing page, changes should be able to work on Windows out of the box.

• OS: Win11, MacOS
• Browser: Edge, Chrome
• Graphics card: RTX 3050, M1Max

Screenshots or videos of your changes

[catboxanon]

imo, if you're doing this, you may want to go ahead and just add every available scheduler in kdiffusion. It was pointed out here that for DPM-Solver++(2M) SDE, exponential is technically better than Karras.

[KohakuBlueleaf]

imo, if you're doing this, you may want to go ahead and just add every available scheduler in kdiffusion. It was pointed out here that for DPM-Solver++(2M) SDE, exponential is technically better than Karras.

Actually I also point this out in my issue but I am still struggling on it
You may need to wait one more hour for what you said(

[KohakuBlueleaf]

The problem here for "every scheduler" is that not all the scheduler need same parameters

If you think ignore vp is ok, I can make it work in 10min

[catboxanon]

VP is the default iirc, so that might be fine.

[KohakuBlueleaf]

oh that's good, ok, I will add a dropdown for it

[catboxanon]

Actually on doing a sanity check, that isn't true. Default noise scheduler for kdiffusion is found here.

VP can be addressed later though, I think what this is mostly trying to solve, exposing Karras and other scheduler parameters, is more important.

[KohakuBlueleaf]

Actually on doing a sanity check, that isn't true. Default noise scheduler for kdiffusion is found here.

VP can be addressed later though, I think what this is mostly trying to solve, exposing Karras and other scheduler parameters, is more important.

ok, The default noise scheduler is as same as what SD have in their model.
So yes, not vp I think

[gesen2egee]

great work!

[KohakuBlueleaf]

DPM++ SDE 2M, 10step

Karras 0.1/20/8

Exponential 0.1/20

Polyexponential 0.1/20/1.5

UI be like:

[catboxanon]

Probably want to rename everywhere where applicable Karras to kdiffusion or similar since it's allowing more than just Karras now too.

[KohakuBlueleaf]

Probably want to rename Karras to kdiffusion or similar since it's allowing more than just Karras now too.

oh yes, make sense

[catboxanon]

Should only ever be adding the extra params to kdiffusion samplers, and only when the custom kdiffusion scheduler is checked. For example, it currently adds it for DDIM and when the custom scheduler is unchecked.

[KohakuBlueleaf]

Should only ever be adding the extra params to kdiffusion samplers, and only when the custom kdiffusion scheduler is checked. For example, it currently adds it for DDIM and when the custom scheduler is unchecked.

Like this?

[KohakuBlueleaf]

Should I just add a check for sampler name so if user choose the non-kdiffusion sampler, p.enable_karras will automatically be False?

[catboxanon]

I think it can be safe to assume the custom scheduler is not in use if it's not contained in the metadata. That logic would put in False for it which seems unnecessary. It might be worth adding a check if p.sampler is also an instance of KDiffusionSampler or not (haven't double checked but I seem to remember p.sampler being stored), for cases where the user tries to enable it for a non-kdiffusion sampler (which will affect nothing).

[KohakuBlueleaf]

I think it can be safe to assume the custom scheduler is not in use if it's not contained in the metadata. That logic would put in False for it which seems unnecessary. It might be worth adding a check if p.sampler is also an instance of KDiffusionSampler or not (haven't double checked but I seem to remember p.sampler being stored).

I can just add a samplers map for k_diffusion (and this also may be useful in the future)
And, (False or None) actually will produce None, not False.

[KohakuBlueleaf]

oh wait I pushed some not-related code(
let me fix it

[catboxanon]

Minor naming fixes

[catboxanon]

In my UI, enabling it causes it to appear above CFG scale instead of below the checkbox. I assume this isn't intentional since your screenshots show otherwise. This is with only built-in extensions.

Also, rho is potentially dangerous when certain samplers are used (like DPM++ 2M SDE) and the type is polyexponential. At a high value, I get a RecursionError during inference. Doing some brief tests this seems to start at around 6. May be worth adding an additional except to this when that's the case to return the last latent and also inform the user in a print statement of what happened.

stable-diffusion-webui/modules/sd_samplers_kdiffusion.py

Lines 250 to 253 in 89f9faa

	try:
	return func()
	except sd_samplers_common.InterruptedException:
	return self.last_latent

[KohakuBlueleaf]

@catboxanon for your screen shot, you may want to add "kdiffusion_scheduler" into your ui settings (the order part)

And for high rho for polyexponential
I will add a expect for RecursionError and add a print for it

[catboxanon]

Ah, I overlooked that, thanks! For most users that shouldn't matter then, it only did for me since I had set that manually prior.

VP scheduler would still be nice to have, but again, that could wait until a follow-up PR. This otherwise seems to be working with everything else I've tested thus far.

[KohakuBlueleaf]

@catboxanon cool!
And I will add another features that:
if you set the value to 0.0, it will use the sigmas that used in "built-in karras scheduler"
so user can just change the rho and scheduler type
no need to struggle on sigmas

[catboxanon]

A couple more comments. Title of this PR should be changed too.

[AUTOMATIC1111]

I don't believe those belong in the main image generation interface. We have a settings page with settings for other sampler parameters, and if the user wants to have those available to him in main interface, he can use quicksettings. Have them added as settings.

[catboxanon]

I would generally agree to that but I think this is a bit more debatable. As mentioned in the issue (#10621) Karras samplers shouldn't even be a separate option in the first place. And this impacts all kdiffusion samplers, not just a minor few like sampler_extra_params does. Ideally this would depracate the separate Karras samplers at some point.

[catboxanon]

At minimum, I would say the scheduler should be selectable in the main interface, and everything else (sigma, rho) could potentially be relegated to settings.

[AUTOMATIC1111]

Here's a simple solution: keep the list as it is, add those overrides as settings, do not deprecate anything. If you really want to remove some entries from the list, you'd also have to write some code to be compatible with infotexts of already generated images. Is it worth it?

You can make scheduler selectable in the main interface via quicksettings.

[catboxanon]

I think we need a better way to display quicksettings then because it's going to continually get more cluttered as more things are relegated there. Controls like this and CFG rescale don't feel like one-and-done type settings to me.

That can be addressed in a separate PR though, so I otherwise am OK with your judgement.

[KohakuBlueleaf]

I don't believe those belong in the main image generation interface. We have a settings page with settings for other sampler parameters, and if the user wants to have those available to him in main interface, he can use quicksettings. Have them added as settings.

The problem is, this kind of settings is not something you can set it once and done.
Like, for img2img, User may want to use denoise=1.0 with special sigma_max and rho to make a schedule that stay on sigma=0.7~1.2 more time.

Actually in my opinion, if k_diffusion scheduler settings should not in the main interface, CFG scale also should not in the main interface.

And for deprecate "karras" sampler, I agree that we can keep the list of it. Or we will need to make some convertion to accept it and automatically fill the sigma_max/sigma_min for them

[catboxanon]

@KohakuBlueleaf @AUTOMATIC1111 Here, take a look at this. I tried my hand at implementing an option for a second quicksettings area that should be more sufficient for controls like this. #10648

[KohakuBlueleaf]

@catboxanon it looks good!
Should I just make custom k_diffusion scheduler as settings now?
(BTW, for me, make custom scheduler in main interface still make more sense)

[catboxanon]

I would open a separate PR incase auto changes his mind, but it seems like putting it in settings is where this is going to end up.

[KohakuBlueleaf]

I would open a separate PR incase auto changes his mind, but it seems like putting it in settings is where this is going to end up.

Makes sense.
I will open seperate PR later.

[KohakuBlueleaf]

@catboxanon I find a problem here.
How should I set the value range for this thing?
The range of sigma_max actually can be 0~inf.
And it looks like the value range of settings is only in the source code.

rho also has this problem, the value range should not be restricted.

[catboxanon]

Well s_tmax is currently not exposed for the same reason.

stable-diffusion-webui/modules/processing.py

Line 147 in 89f9faa

self.s_tmax = s_tmax or float('inf') # not representable as a standard ui option

I think gr.Number might be your only passable option.

[KohakuBlueleaf]

@catboxanon so something like this?

[catboxanon]

LGTM

[KohakuBlueleaf]

@catboxanon Another problem here. How can I let the "send to txt2img/img2img" in image browser to send the sigma_min/sigma_max if they are in settings?

edit: I found it