Add beta, exponential and karras sigmas to FlowMatchEulerDiscreteScheduler

[hlky]

What does this PR do?

Add beta, exponential and karras sigmas to FlowMatchEulerDiscreteScheduler.

from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
import numpy as np


def calculate_shift(
    image_seq_len,
    base_seq_len: int = 256,
    max_seq_len: int = 4096,
    base_shift: float = 0.5,
    max_shift: float = 1.16,
):
    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
    b = base_shift - m * base_seq_len
    mu = image_seq_len * m + b
    return mu

image_seq_len = 4096
mu = calculate_shift(image_seq_len)
num_inference_steps = 8
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)

flow_match = FlowMatchEulerDiscreteScheduler.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="scheduler", use_beta_sigmas=True)
flow_match.set_timesteps(sigmas=sigmas, mu=mu)
print(f"flow_match beta (sigmas) {flow_match.sigmas} {flow_match.sigmas.shape[0]}")
print(f"flow_match beta (sigmas) {flow_match.timesteps} {flow_match.timesteps.shape[0]}")

flow_match.set_timesteps(num_inference_steps, mu=mu)
print(f"flow_match beta ({num_inference_steps}) {flow_match.sigmas} {flow_match.sigmas.shape[0]}")
print(f"flow_match beta ({num_inference_steps}) {flow_match.timesteps} {flow_match.timesteps.shape[0]}")

flow_match = FlowMatchEulerDiscreteScheduler.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="scheduler", use_exponential_sigmas=True)
flow_match.set_timesteps(sigmas=sigmas, mu=mu)
print(f"flow_match exponential (sigmas) {flow_match.sigmas} {flow_match.sigmas.shape[0]}")
print(f"flow_match exponential (sigmas) {flow_match.timesteps} {flow_match.timesteps.shape[0]}")

flow_match.set_timesteps(num_inference_steps, mu=mu)
print(f"flow_match exponential ({num_inference_steps}) {flow_match.sigmas} {flow_match.sigmas.shape[0]}")
print(f"flow_match exponential ({num_inference_steps}) {flow_match.timesteps} {flow_match.timesteps.shape[0]}")

flow_match = FlowMatchEulerDiscreteScheduler.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="scheduler", use_beta_sigmas=True)
flow_match.set_timesteps(sigmas=sigmas, mu=mu)
print(f"flow_match karras (sigmas) {flow_match.sigmas} {flow_match.sigmas.shape[0]}")
print(f"flow_match karras (sigmas) {flow_match.timesteps} {flow_match.timesteps.shape[0]}")

flow_match.set_timesteps(num_inference_steps, mu=mu)
print(f"flow_match karras ({num_inference_steps}) {flow_match.sigmas} {flow_match.sigmas.shape[0]}")
print(f"flow_match karras ({num_inference_steps}) {flow_match.timesteps} {flow_match.timesteps.shape[0]}")

flow_match beta (sigmas) tensor([1.0000, 0.9511, 0.8510, 0.7242, 0.5888, 0.4620, 0.3619, 0.3130, 0.0000]) 9
flow_match beta (sigmas) tensor([1000.0000,  951.1288,  851.0334,  724.2367,  588.8108,  462.0141,
         361.9187,  313.0475]) 8
flow_match beta (8) tensor([1.0000, 0.9291, 0.7838, 0.5998, 0.4033, 0.2193, 0.0741, 0.0032, 0.0000]) 9
flow_match beta (8) tensor([1000.0000,  929.0844,  783.8388,  599.8478,  403.3352,  219.3442,
          74.0985,    3.1830]) 8
flow_match exponential (sigmas) tensor([1.0000, 0.8471, 0.7176, 0.6079, 0.5150, 0.4362, 0.3695, 0.3130, 0.0000]) 9
flow_match exponential (sigmas) tensor([1000.0000,  847.1188,  717.6103,  607.9012,  514.9645,  436.2361,
         369.5438,  313.0475]) 8
flow_match exponential (8) tensor([1.0000, 0.4398, 0.1934, 0.0851, 0.0374, 0.0165, 0.0072, 0.0032, 0.0000]) 9
flow_match exponential (8) tensor([1000.0000,  439.8065,  193.4298,   85.0717,   37.4151,   16.4554,
           7.2372,    3.1830]) 8
flow_match karras (sigmas) tensor([1.0000, 0.9511, 0.8510, 0.7242, 0.5888, 0.4620, 0.3619, 0.3130, 0.0000]) 9
flow_match karras (sigmas) tensor([1000.0000,  951.1288,  851.0334,  724.2367,  588.8108,  462.0141,
         361.9187,  313.0475]) 8
flow_match karras (8) tensor([1.0000, 0.9291, 0.7838, 0.5998, 0.4033, 0.2193, 0.0741, 0.0032, 0.0000]) 9
flow_match karras (8) tensor([1000.0000,  929.0844,  783.8388,  599.8478,  403.3352,  219.3442,
          74.0985,    3.1830]) 8

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

cc @yiyixuxu @asomoza

[yiyixuxu]

thanks for the initiative!! @hlky

I left some comment as I was going through the the PR; however, now after going through all the code changes, I actually think it becomes pretty clear to me these two things do not naturally combine: they share very little common logic together, and most of the code should be moved inside a big if use_flow_match: ... else: ... block if they have not already in there

I think we should:

1. add the karras/beta/exponential to flow matching scheduler
2. in a future PR, we should think after separating the sigmas schedule as its own abstraction, like you suggested to me:)

let me know what you think!

[yiyixuxu]

this whole section to calculation betas -> alphas_cumprod is not relevant to flow matching, no? since it's only used to calculate sigmas when not use_flow_match; but it isn't clear from the code because is it outside of the if use_flow_match ... else ... block

[yiyixuxu]

is v_prediction relevant to flow matching? can people configure use_flow_match + v_prediction? based on the code, it is possible. But does this make sense?

[yiyixuxu]

we have this argument sigmas is to accept a custom sigmas schedule from user, i.e. they should either pass a custom sigmas or choose to use one of the pre-set sigma schedules (e.g. karras, beta, exponential) ;

so we should not have this logic here

if sigma is not None and self.config.use_flow_match:
    ...
    if self.config.use_karras_sigmas:
          ...

[hlky]

Yes this was due to how pipelines are currently set up and I wanted to test the noise schedules without modifying the pipelines.

Some of the logic is the same as existing FlowMatchEuler which applies shifting etc to the supplied sigmas.

Although I'm not sure the context of why Flux pipelines pass sigmas and why those are calculated differently to the sigmas is None path. Karras etc might actually work better when applied before the scaling so the future refactor will be really useful here, I'll run some tests to check that and I'll check pipelines modified to pass number of steps when using karras etc

[hlky]

Thanks for the review. This was mainly to see whether it could be combined, there are only a few key differences, it should be possible to refactor Euler in a way that FlowMatchEuler only overrides a few things instead, it would be great to use copied from more when creating other FlowMatch variants and make the differences clear. For now I'll change this PR to add karras/beta/exponential to flow match euler and think about that future PR with an abstraction of noise schedulers.

[ukaprch]

@hlky @yiyixuxu @asomoza
Before you folks implement this PR I'd like to offer another approach which I used in my (hopefully) soon to be implemented:
scheduling_flow_match_dpmsolver_multistep.py
which I think is more straightforward, makes more sense and makes use of betas to calculate sigmas as I feel they do make a difference. I will come up with this scheduler hopefully today / tomorrow which you can test for yourselves. The user will only have to make one decision whether to use the betas or not. The betas will be governed by the variables beta_start and beta_end as is currently the case for SDXL.

[yiyixuxu]

@hlky is this PR ready for review? can we see some images? :)

@ukaprch thanks! Happy to take a look!

[hlky]

Yes it's ready for review

import torch
from diffusers import FluxPipeline, FlowMatchEulerDiscreteScheduler
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.enable_vae_tiling()
pipe = pipe.to("cuda")
config = pipe.scheduler.config
euler_flow_beta = FlowMatchEulerDiscreteScheduler.from_config(config, use_beta_sigmas=True)

euler_flow_exponential = FlowMatchEulerDiscreteScheduler.from_config(config, use_exponential_sigmas=True)

euler_flow_karras = FlowMatchEulerDiscreteScheduler.from_config(config, use_karras_sigmas=True)

pipe.scheduler = euler_flow_beta
generator = torch.Generator("cuda").manual_seed(0)
prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=30, guidance_scale=3.5, generator=generator).images[0]
image.save("flow_beta.png")

pipe.scheduler = euler_flow_exponential
generator = torch.Generator("cuda").manual_seed(0)
prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=30, guidance_scale=3.5, generator=generator).images[0]
image.save("flow_exponential.png")

pipe.scheduler = euler_flow_karras
generator = torch.Generator("cuda").manual_seed(0)
prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=30, guidance_scale=3.5, generator=generator).images[0]
image.save("flow_karras.png")

beta

exponential

karras

[ukaprch]

The proposed scheduler can handle SD 3.5 Medium/Large & Flux equally well. Working out final details.
Here's a taste of what the proposed FlowMatchEulerDiscrete scheduler can do using SD 3.5 Medium:
sigma scheduler = karras
use beta sigmas = True
30 steps
seed = 114747598
prompt = a cat holding a sign that says "hello world"

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

thanks!

[yiyixuxu]

@ukaprch thanks!

I'm going to merge this PR now since it is consistent with the current design. We realize that our current scheduler has many limitations and are very eager to refactor and improve.

I look forward to seeing your PR!! We can discuss it from there. And if we need to change the current design, we will apply it across all schedulers :)

[ukaprch]

One of the limitations of FlowMatch for sigma schedules like 'karras', 'exponential' and 'lambdas' is that the current scheme of creating sigmas for them is not biased (weighted) enough in the early steps delaying convergence as opposed to what was used in SDXL. The use of betas is instrumental in building the necessary timesteps / sigmas to take this into account. I felt this feature should not be ignored in our new FlowMatch schedulers. The problem was how to implement for FlowMatch.
The image outputs generated during preliminary testing justify using beta/sigmas as a variable option for Flux and especially for SD 3.5.
Below are (2) images using the same generational data with the exception that 1 uses std flux/karras sigma generation and the other uses beta/sigma generation. The bottom image uses beta/sigmas Custom Betas: start=0.00085; end=0.012. I also prefer using scaled linear betas as they seem to have a marginal improvement.