Can we get more schedulers for flow based models such as SD3, SD3.5, and flux

[linjiapro]

It seems advanced schedulers such as DDIM, and the dpm++ 2m does work with flow based model such as SD3, SD3.5, and flux.
However, I only see 2 flow based schedulers in diffusers codebase:

FlowMatchEulerDiscreteScheduler, and'
FlowMatchHeunDiscreteScheduler

I tried to use DPMSolverMultistepScheduler, but it does not generate correct images with flow based models. Help?

[sayakpaul]

Cc: @asomoza @hlky @yiyixuxu

[linjiapro]

@ukaprch

From the import clauses txt file, it seems the classes are already under diffusers repo for the testing runs. It seems to be suitable to include the file to the official diffusers repo?

[ukaprch]

Some of you can update your diffusers library to incorporate the scheduler I proposed (which is a first start) for independent testing (as I have). The scheduler as written can be incorporated into the diffuser library w/o change and will work with the current FLUX and SD3 pipelines.

[linjiapro]

@sayakpaul @ukaprch

I have a pending PR which has not been merged:
#9758

Once above merges, I can work on sending another PR related to schedulers. In the meanwhile, anyone else are welcomed to raise a pull request.

[vladmandic]

also related #9607

[vladmandic]

@ukaprch why not create a pr to include this in diffusers?
i've integrated your code in sdnext, dpm flowmatch is working fine except any kind of karras/exponential results in full blur.

anyhow, here are some sample grids (including original euler and heun flowmatch for reference)

sd35:

flux.1

auraflow: also tried with this model since its also flow-match based and it works

[ukaprch]

Thanks for testing and your feedback. Please note that I had no such problems using the FLUX pipeline with the new scheduler using algorithms (karras, exponential, lambdas) see attached. I am using the quantized version (QINT8) FWIW for transformers and T5.

 image = pipe(
 prompt='a cat holding a sign that says hello world',
 width = 1024,
 height = 1024,
 guidance_scale=3.5,
 generator=generator, (seed = 114747598)
 max_sequence_length=512,
 num_inference_steps=30).images[0]

I ran the above pipeline using the basic FLUX dev prompt from Replicate and generated the following locally on my PC:

[vladmandic]

bit more testing, karras and exponential are working fine for flux, but not for sd35

flux:

sd35

actual params used:

10:50:38-967529 DEBUG Sampler: sampler="DPM2++ 2M FlowMatch" class="FlowMatchDPMSolverMultistepScheduler config={'num_train_timesteps': 1000, 'shift': 1, 'use_dynamic_shifting': False, 'solver_order': 2, 'sigma_schedule': 'karras', 'use_SD35_sigmas': True, 'algorithm_type': 'dpmsolver++2M', 'use_noise_sampler': True}

[ukaprch]

OK, I noticed you are using a shift: 1. This should be shift: 3 which is the std for SD3. I'll look into some more.

After running my test the only noticeable thing I see is a more pronounced 'bokeh' effect which I wouldn't necessarily characterize as blur per se. I do think the sigmas / timesteps for karras, exponential and lambdas for SD 3.5 could possibly use some further refinement.

image = pipe(
prompt='photorealistic image of a robot walking the streets of a city meeting passersby, 8K, dramatic lighting, textured skin, film grain, RAW photo, highly detailed',
negative_prompt=negative_prompt,
width = 1024,
height = 1024,
guidance_scale=3.5,
generator=generator, (seed = 580198207)
max_sequence_length=512,
num_inference_steps=30).images[0]

scheduler config:

◢ | config | OrderedDict([('num_train_timesteps', 1000), ('solver_order', 2), ('thresholding', False), ('dynamic_thresholding_ratio', 0.995), ('sample_max_value', 1.0), ('algorithm_type', 'dpmsolver++2Msde'), ('solver_type', 'midpoint'), ('sigma_schedule', None), ('shift', 3.0), ('midpoint_ratio', 0.5), ('s_noise', 1), ('use_noise_sampler', True), ('use_SD35_sigmas', True), ('use_dynamic_shifting', False), ('base_shift', 0.5), ('max_shift', 1.15), ('base_image_seq_len', 256), ('max_image_seq_len', 4096), ('_use_default_values', ['dynamic_thresholding_ratio', 'base_image_seq_len', 'midpoint_ratio', 'solver_type', 'thresholding', 'base_shift', 'max_image_seq_len', 'sample_max_value'])]) | FrozenDict

| algorithm_type | 'dpmsolver++2Msde' | str
| base_image_seq_len | 256 | int
| base_shift | 0.5 | float
| dynamic_thresholding_ratio | 0.995 | float
| max_image_seq_len | 4096 | int
| max_shift | 1.15 | float
| midpoint_ratio | 0.5 | float
| num_train_timesteps | 1000 | int
| s_noise | 1 | int
| sample_max_value | 1.0 | float
| shift | 3.0 | float
| sigma_schedule | None | NoneType
| solver_order | 2 | int
| solver_type | 'midpoint' | str
| thresholding | False | bool
| use_SD35_sigmas | True | bool
| use_dynamic_shifting | False | bool
| use_noise_sampler | True | bool

[vladmandic]

ahhh, found it. for testing purposes i was running with steps=15 only.
and that's fine for beta, but definitely not fine for other sigma methods.

btw, good catch on shift value, thats a user error on my side, had it left like that from flux, but it doesn't impact this.

[ukaprch]

I found with FLUX that you need generally 25 or more steps to get good images. Remember, we're not using the 3 step method as previously with SDXL. Your images look good BTW. It really comes into its own b/t 35 - 50 steps. BTW if you increase the shift factor from 3 to 3.5 you get a bit more contrast in the image which may or may not be to everyone's liking.
So what are you thinking? Are we almost good to proceed with this as a scheduler?

[vladmandic]

I found with FLUX that you need generally 25 or more steps to get good images.

actually, the worst-case scenario here is sd35-medium ;)
flux is behaving better in this regard.

re: shift - i've noticed. i have it as user-configurable item

So what are you thinking? Are we almost good to proceed with this as a scheduler?

definitely!

from a quick look at the code, 99% is clean.
del sigmas1 looks unsafe as it may not be defined in all code-paths. and in general, does it help at all to have all those del statements in set_timestep()?

i wish if use_SD35_sigmas could be auto-determined, but i don't see a clean way.
but it does need a better name, definitely not something with mixed-capitalization :)

[ukaprch]

vladmandic: As noted, this scheduler is not designed for flow match derived sampling.

I've attached a new FlowMatch version of the original: scheduling_dpmsolver_multistep.py (scheduling_flow_match_dpmsolver_multistep_orig.py)
and updated scheduling_flow_match_dpmsolver_multistep.py per some of your suggestions.
The powers that be TPTB if they decide to move forward on these can obviously rename / make changes, etc.

1. simags1 is gone. A holdover from my previous testing.
2. use_SD35_sigmas is now use_beta_sigmas and you'll see why when you look at the code.
3. both new scheduler's have a consistent variable interface so you'll be able to use both w/o much effort.
4. sigma_schedule (which replaces all those use_????_sigmas=True in the original version) include the following:
   [None, "karras", "exponential", "lambdas", "betas"].
5. Betas play an important role in determining sigmas and image rendering. Because of this, beta_start and beta_end become important variables and should be configurable. Currently I am using the default: beta_start = 0.0001; beta_end = 0.02 and custom: beta_start = 0.00085; beta_end=0.012. Values in between would also work well. Again, this is all subjective.
6. When using sigma_schedulers "karras", "exponential" and "lambdas" you will almost certainly have better outcomes with your images when using more inference steps (i.e. 40+ vs 25).
7. It's unclear with the current versions of FLUX DEV and SD 3.5 Large that you can run most images with these schedulers and get good results in under 20 steps. The current (Open) models were not designed for this. What is clear is that the dpm++ class of solvers are better at rendering realistic images (in my testing) than the current FlowMatchEulerDiscrete.py scheduler.

scheduling_flow_match_dpmsolver_multistep.txt
scheduling_flow_match_dpmsolver_multistep_orig.txt

Please do some more testing on both and state any concerns, etc. I think we are very close on this and the community would definitely benefit from both.

Example images using same generation data. Left is new proposed version of scheduling_flow_match_dpmsolver_multistep.py and right is current scheduling_flow_match_euler_discrete.py both used 25 steps: