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)

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

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")