Add `beta`, `exponential` and `karras` sigmas to `FlowMatchEulerDiscreteScheduler` #10001

hlky · 2024-11-23T13:24:40Z

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:

add the karras/beta/exponential to flow matching scheduler
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 · 2024-11-23T20:11:51Z

src/diffusers/schedulers/scheduling_euler_discrete.py

+        max_shift: Optional[float] = 1.15,
+        base_image_seq_len: Optional[int] = 256,
+        max_image_seq_len: Optional[int] = 4096,
+        invert_sigmas: bool = False,
    ):
        if self.config.use_beta_sigmas and not is_scipy_available():


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 · 2024-11-23T20:13:29Z

src/diffusers/schedulers/scheduling_euler_discrete.py

+                sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+        else:
+            sigmas = (((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5).flip(0)
+
        # setable values
        self.num_inference_steps = None

        # TODO: Support the full EDM scalings for all prediction types and timestep types
        if timestep_type == "continuous" and prediction_type == "v_prediction":


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 · 2024-11-23T20:28:25Z

src/diffusers/schedulers/scheduling_euler_discrete.py


-        else:
+            if self.config.use_karras_sigmas:


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: ...

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 · 2024-11-23T21:44:25Z

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.

hlky · 2024-11-25T14:30:56Z

src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

            timesteps = np.linspace(
                self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
            )

            sigmas = timesteps / self.config.num_train_timesteps
+        else:
+            num_inference_steps = len(sigmas)


Just to note passing sigmas to Euler vs FlowMatchEuler is different

diffusers/docs/source/en/using-diffusers/scheduler_features.md

Line 133 in 074e123

sigmas = [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.0]

diffusers/src/diffusers/schedulers/scheduling_euler_discrete.py

Line 367 in 074e123

num_inference_steps = len(timesteps) if timesteps is not None else len(sigmas) - 1

diffusers/src/diffusers/pipelines/flux/pipeline_flux.py

Line 702 in 074e123

sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)

yes, it is a little bit different (even though I think it should be the same); that's why we created it in the pipeline for convenience (and had planned to move the logic to inside scheduler too, but just haven't done that)

The current logic of set_timesteps (for sigma is None path) in FlowMatchEuler is implemented based on SD3( it was the first flow-matching model we have), and it is somehow a little bit different from all the flow-matching models following it; I think the main difference is it applied the shift twice, first time here

diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

Line 82 in 074e123

sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)

and then

diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

Line 195 in 074e123

self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps

BTW, it is very clear the flux method is the "standard" way by now

We had planned to move the flux sigmas to the scheduler, too, but since we want to separate the noise scheduler logic way from the scheduler; maybe we don't have to spend time on that for now because it might be the opposite direction :)

oh sorry this here is just explain your question here

why Flux pipelines pass sigmas and why those are calculated differently to the sigmas is None path

I think i might put it in the wrong thread here, sorry!

ukaprch · 2024-11-26T13:25:13Z

@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 · 2024-11-26T17:11:05Z

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

@ukaprch thanks! Happy to take a look!

hlky · 2024-11-26T20:18:48Z

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 · 2024-11-27T14:15:54Z

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"

hlky force-pushed the combine-flow-match-euler branch 3 times, most recently from 634cb90 to f41b2f1 Compare November 23, 2024 17:14

hlky marked this pull request as ready for review November 23, 2024 17:23

yiyixuxu reviewed Nov 23, 2024

View reviewed changes

hlky force-pushed the combine-flow-match-euler branch from f41b2f1 to 3319bc5 Compare November 24, 2024 10:02

hlky changed the title ~~Combine Flow Match Euler into Euler~~ Add beta, exponential and karras sigmas to FlowMatchEulerDiscreteScheduler Nov 24, 2024

Add beta, exponential and karras sigmas to FlowMatchEuler

39f634f

hlky force-pushed the combine-flow-match-euler branch from 3319bc5 to 39f634f Compare November 25, 2024 13:26

hlky commented Nov 25, 2024

View reviewed changes

yiyixuxu mentioned this pull request Nov 27, 2024

[Sana] Add Sana, including SanaPipeline, SanaPAGPipeline, LinearAttentionProcessor, Flow-based DPM-sovler and so on. #9982

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `beta`, `exponential` and `karras` sigmas to `FlowMatchEulerDiscreteScheduler` #10001

Add `beta`, `exponential` and `karras` sigmas to `FlowMatchEulerDiscreteScheduler` #10001

hlky commented Nov 23, 2024 •

edited

Loading

yiyixuxu left a comment

yiyixuxu Nov 23, 2024

yiyixuxu Nov 23, 2024

yiyixuxu Nov 23, 2024

hlky Nov 23, 2024 •

edited

Loading

hlky commented Nov 23, 2024 •

edited

Loading

hlky Nov 25, 2024

yiyixuxu Nov 25, 2024

yiyixuxu Nov 25, 2024

yiyixuxu Nov 25, 2024

ukaprch commented Nov 26, 2024

yiyixuxu commented Nov 26, 2024

hlky commented Nov 26, 2024

ukaprch commented Nov 27, 2024

Add beta, exponential and karras sigmas to FlowMatchEulerDiscreteScheduler #10001

Are you sure you want to change the base?

Add beta, exponential and karras sigmas to FlowMatchEulerDiscreteScheduler #10001

Conversation

hlky commented Nov 23, 2024 • edited Loading

What does this PR do?

Who can review?

yiyixuxu left a comment

Choose a reason for hiding this comment

yiyixuxu Nov 23, 2024

Choose a reason for hiding this comment

yiyixuxu Nov 23, 2024

Choose a reason for hiding this comment

yiyixuxu Nov 23, 2024

Choose a reason for hiding this comment

hlky Nov 23, 2024 • edited Loading

Choose a reason for hiding this comment

hlky commented Nov 23, 2024 • edited Loading

hlky Nov 25, 2024

Choose a reason for hiding this comment

yiyixuxu Nov 25, 2024

Choose a reason for hiding this comment

yiyixuxu Nov 25, 2024

Choose a reason for hiding this comment

yiyixuxu Nov 25, 2024

Choose a reason for hiding this comment

ukaprch commented Nov 26, 2024

yiyixuxu commented Nov 26, 2024

hlky commented Nov 26, 2024

ukaprch commented Nov 27, 2024

Add `beta`, `exponential` and `karras` sigmas to `FlowMatchEulerDiscreteScheduler` #10001

Add `beta`, `exponential` and `karras` sigmas to `FlowMatchEulerDiscreteScheduler` #10001

hlky commented Nov 23, 2024 •

edited

Loading

hlky Nov 23, 2024 •

edited

Loading

hlky commented Nov 23, 2024 •

edited

Loading