Add dynamic_shifting to SD3

[hlky]

What does this PR do?

Adds dynamic_shifting to SD3.

import torch
from diffusers import StableDiffusion3Pipeline, FlowMatchEulerDiscreteScheduler
pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
)
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config, use_dynamic_shifting=True)

We can pass a custom mu value

image = pipe(prompt, ..., mu=...).images[0]

When not provided mu is calculated the same as in Flux.

import torch
from diffusers import StableDiffusion3Pipeline, FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler

prompt = "A cat holding a sign that says hello world"

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
).to("cuda")

original_config = pipe.scheduler.config

pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(original_config)
print(pipe.scheduler.config)
image = pipe(prompt, generator=torch.Generator().manual_seed(4)).images[0]
image.save("euler.png")

pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(original_config, use_dynamic_shifting=True)
print(pipe.scheduler.config)
image = pipe(prompt, generator=torch.Generator().manual_seed(4)).images[0]
image.save("euler_dynamic_shift.png")

pipe.scheduler = FlowMatchHeunDiscreteScheduler.from_config(original_config)
print(pipe.scheduler.config)
image = pipe(prompt, generator=torch.Generator().manual_seed(4)).images[0]
image.save("heun.png")

pipe.scheduler, unused_kwargs = FlowMatchHeunDiscreteScheduler.from_config(original_config, use_dynamic_shifting=True, return_unused_kwargs=True)
print(unused_kwargs)
image = pipe(prompt, generator=torch.Generator().manual_seed(4)).images[0]
image.save("heun_dynamic_shift.png")

Euler

No shift	Dynamic shift

Img2Img

No shift	Dynamic shift

Inpaint

No shift	Dynamic shift

Heun

As expected Heun is unaffected by use_dynamic_shift.

No shift	Dynamic shift

Img2Img

No shift	Dynamic shift

Inpaint

No shift	Dynamic shift

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.

@yiyixuxu @cubiq

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

[cubiq]

since there's not an integrated way in the sd35 pipeline to calculate mu the use of dynamic shifting is not straight forward.

to expand on @hlky example, with the suggested implementation this is what would be needed:

from diffusers import FlowMatchEulerDiscreteScheduler
scheduler_config = {
    'num_train_timesteps': 1000,
    'shift': 3.0,
    'use_dynamic_shifting': True, # <- note this
    'base_shift': 0.5,
    'max_shift': 1.15,
    'base_image_seq_len': 256,
    'max_image_seq_len': 4096,
    'invert_sigmas': False,
}
scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
pipe = StableDiffusion3Pipeline.from_pretrained(
    ...
    scheduler=scheduler,
}
mu = calculate_shift(...) # your own shift function
image = pipe(prompt, ..., mu=mu).images[0]

To calculate the shift you could use a strategy similar to what Flux does:

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

But there might be better strategies (maybe logarithmic?), that also raises the question if mu should be a configurable parameter for Flux too.

As a side note maybe a more robust solution would be to have a scheduler_args option and send them as **kwargs to retrieve_timesteps? If done in all pipelines that would make it future proof.

[yiyixuxu]

thanks @hlky @cubiq
amazing!!