flow matching lcm scheduler

[quickjkee]

Why?

Most cutting-edge image generative models (SD3.5,
FLUX) are formulated using flow matching. Moreover, some state-of-the-art distillation methods (DMD2, CD) use stochastic multistep sampling (LCM scheduler). However, there is no realization of the LCM scheduler for the flow matching formulation.

What does this PR do?

Thus, this PR provides a scheduling_flow_match_lcm.py that is in the same spirit as the scheduling_lcm.py. The goal is for distilled models (using, e.g SD3.5 or FLUX) to use this sampler.

Moreover, as an additional tool, this PR provides the ability to generate on different resolution scales.
Which could significantly accelerate the generation process. This concept was introduced in SwD.

How it works?

For example, one of the most prominent distillation approaches, DMD2, uses LCM. It is distilled from SDXL. But there is also an implementation for SD3.5 Large and with PR it works as follows:

import torch
from diffusers.schedulers.scheduling_flow_match_lcm import FlowMatchLCMScheduler
from diffusers import StableDiffusion3Pipeline
from peft import PeftModel

# Loading models
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-large", 
                                                torch_dtype=torch.float16)
pipe = pipe.to("cuda")
lora_path = 'yresearch/stable-diffusion-3.5-large-dmd2'
pipe.transformer = PeftModel.from_pretrained(
    pipe.transformer,
    lora_path,
)

# LCM scheduler 
pipe.scheduler = FlowMatchLCMScheduler.from_config(pipe.scheduler.config, 
                                                   shift=1.0)

# Generation
image = pipe(
    "a cat reading a newspaper",
    sigmas=[1.0000, 0.9454, 0.8959, 0.7904, 0.7371, 0.6022],
    guidance_scale=0.0,
    generator=torch.Generator().manual_seed(0),
).images[0]

To generate in a scale-wise regime, it is necessary to provide scale factors through the pipe.scheduler.set_scale_factors.
This modification is optional.

import torch
from diffusers.schedulers.scheduling_flow_match_lcm import FlowMatchLCMScheduler
from diffusers import StableDiffusion3Pipeline
from peft import PeftModel

# Loading models
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-large", 
                                                torch_dtype=torch.float16)
pipe = pipe.to("cuda")
lora_path = 'yresearch/swd-large-6-steps'
pipe.transformer = PeftModel.from_pretrained(
    pipe.transformer,
    lora_path,
)

# LCM scheduler
pipe.scheduler = FlowMatchLCMScheduler.from_config(pipe.scheduler.config, 
                                                   shift=1.0)

# Setting up the scale factors
# We define it wrt input size 
# In this case: 32->48->64->80->96->128
pipe.scheduler.set_scale_factors(
    scale_factors=[1.5, 2., 2.5, 3., 4.], 
    upscale_type='bicubic'
)

# Generation
image = pipe(
    "a cat reading a newspaper",
    sigmas=[1.0000, 0.9454, 0.8959, 0.7904, 0.7371, 0.6022],
    guidance_scale=0.0,
    generator=torch.Generator().manual_seed(0),
    width=256,
    height=256
).images[0]

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

[hlky]

@bot /style

[github-actions]

Style fixes have been applied. View the workflow run here.

[yiyixuxu]

thanks @quickjkee
do we have a training script using this scheduler?

[quickjkee]

thanks @quickjkee do we have a training script using this scheduler?

thanks for the question and reviewing @yiyixuxu!
I don't think so. This scheduler is only needed for inference, as the basic LCM scheduler.

[yiyixuxu]

@quickjkee
do you have a training script for these checkpoints that can be used with this scheduler? (even though the training itself does not require it)

[ShuyUSTC]

Thank you for your great contribution! Is there any plan to merge this pr? Our HiDream-I1 fast and dev model use this flow_match_lcm scheduler by default. We are waiting for this pr to be merged. Thanks!

[hlky]

Hi @quickjkee. Left a few suggestions on unused parameters but LGTM, thanks!