Add dynamic_shifting to SD3

[hlky]

What does this PR do?

Adds dynamic_shifting to SD3.

importtorchfromdiffusersimportStableDiffusion3Pipeline, FlowMatchEulerDiscreteSchedulerpipe=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.

importtorchfromdiffusersimportStableDiffusion3Pipeline, FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteSchedulerprompt="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.configpipe.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:

fromdiffusersimportFlowMatchEulerDiscreteSchedulerscheduler_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 functionimage=pipe(prompt, ..., mu=mu).images[0]

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

defcalculate_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_lenmu=image_seq_len*m+breturnmu

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!!