Skip to content

vllm_omni.diffusion.models.schedulers

Modules:

Name Description
base

Base scheduler class for diffusion models.
scheduling_dmd2_euler
scheduling_flow_unipc_multistep

FlowUniPCMultistepScheduler - A training-free framework for fast sampling of flow-matching diffusion models.

DMD2EulerScheduler

Bases: FlowMatchEulerDiscreteScheduler

Euler scheduler that always uses the fixed DMD2 training timestep schedule.

set_timesteps 

set_timesteps(
    num_inference_steps: int | None = None,
    device: str | device | None = None,
    timesteps: list[int] | None = None,
    sigmas: list[float] | None = None,
    **kwargs,
) -> None

FlowUniPCMultistepScheduler

Bases: SchedulerMixin, ConfigMixin, BaseScheduler

FlowUniPCMultistepScheduler is a training-free framework designed for the fast sampling of flow-matching diffusion models.

This scheduler implements the UniPC (Unified Predictor-Corrector) algorithm adapted for flow matching, which can achieve the same quality as Euler methods in fewer steps (typically 20-30 steps vs 40-50).

Parameters:

Name Type Description Default
num_train_timesteps `int`, defaults to 1000

The number of diffusion steps to train the model.

 1000
solver_order `int`, default `2`

The UniPC order which can be any positive integer. The effective order of accuracy is solver_order + 1 due to the UniC. It is recommended to use solver_order=2 for guided sampling, and solver_order=3 for unconditional sampling.

 2
prediction_type `str`, defaults to "flow_prediction"

Prediction type of the scheduler function; must be flow_prediction for this scheduler.

 'flow_prediction'
shift `float`, defaults to 1.0

The shift parameter for the noise schedule. For Wan2.2: use 5.0 for 720p, 12.0 for 480p.

 1.0
use_dynamic_shifting `bool`, defaults to False

Whether to use dynamic shifting based on image resolution.

 False
thresholding `bool`, defaults to `False`

Whether to use the "dynamic thresholding" method.

 False
dynamic_thresholding_ratio `float`, defaults to 0.995

The ratio for the dynamic thresholding method.

 0.995
sample_max_value `float`, defaults to 1.0

The threshold value for dynamic thresholding.

 1.0
predict_x0 `bool`, defaults to `True`

Whether to use the updating algorithm on the predicted x0.

 True
solver_type `str`, default `bh2`

Solver type for UniPC. Use bh1 for unconditional sampling when steps < 10, bh2 otherwise.

 'bh2'
lower_order_final `bool`, default `True`

Whether to use lower-order solvers in the final steps. Stabilizes sampling for steps < 15.

 True
disable_corrector `list`, default `[]`

Steps to disable the corrector to mitigate misalignment with large guidance scales.

 ()
timestep_spacing `str`, defaults to `"linspace"`

The way the timesteps should be scaled.

 'linspace'
final_sigmas_type `str`, defaults to `"zero"`

The final sigma value for the noise schedule. Either "zero" or "sigma_min".

 'zero'

begin_index property 

begin_index: int | None

The index for the first timestep. Should be set from pipeline with set_begin_index method.

disable_corrector instance-attribute 

disable_corrector = list(disable_corrector)

last_sample instance-attribute 

last_sample: Tensor | None = None

lower_order_nums instance-attribute 

lower_order_nums = 0

model_outputs instance-attribute 

model_outputs: list[Tensor | None] = [None] * solver_order

num_inference_steps instance-attribute 

num_inference_steps: int | None = None

num_train_timesteps instance-attribute 

num_train_timesteps = num_train_timesteps

order class-attribute instance-attribute 

order = 1

predict_x0 instance-attribute 

predict_x0 = predict_x0

sigma_max instance-attribute 

sigma_max = item()

sigma_min instance-attribute 

sigma_min = item()

sigmas instance-attribute 

sigmas = to('cpu')

solver_p instance-attribute 

solver_p = solver_p

step_index property 

step_index: int | None

The index counter for current timestep. Increases by 1 after each scheduler step.

this_order instance-attribute 

this_order: int = 1

timestep_list instance-attribute 

timestep_list: list[Any | None] = [None] * solver_order

timesteps instance-attribute 

timesteps = sigmas * num_train_timesteps

add_noise 

add_noise(
    original_samples: Tensor,
    noise: Tensor,
    timesteps: IntTensor,
) -> Tensor

Add noise to the original samples.

Parameters:

Name Type Description Default
original_samples `torch.Tensor`

Original samples.

 required
noise `torch.Tensor`

Noise to add.

 required
timesteps `torch.IntTensor`

Timesteps for noise addition.

 required

Returns:

Type Description
Tensor

torch.Tensor: Noisy samples.

convert_model_output 

convert_model_output(
    model_output: Tensor,
    *args,
    sample: Tensor | None = None,
    **kwargs,
) -> Tensor

Convert the model output to the format needed by the UniPC algorithm.

Parameters:

Name Type Description Default
model_output `torch.Tensor`

Direct output from the diffusion model.

 required
sample `torch.Tensor`

Current sample in the diffusion process.

 None

Returns:

Type Description
Tensor

torch.Tensor: Converted model output.

index_for_timestep 

index_for_timestep(
    timestep: Tensor,
    schedule_timesteps: Tensor | None = None,
) -> int

Get the index for a given timestep.

multistep_uni_c_bh_update 

multistep_uni_c_bh_update(
    this_model_output: Tensor,
    *args,
    last_sample: Tensor | None = None,
    this_sample: Tensor | None = None,
    order: int | None = None,
    **kwargs,
) -> Tensor

One step for the UniC (B(h) version) corrector.

Parameters:

Name Type Description Default
this_model_output `torch.Tensor`

Model outputs at x_t.

 required
last_sample `torch.Tensor`

Sample before the last predictor x_{t-1}.

 None
this_sample `torch.Tensor`

Sample after the last predictor x_{t}.

 None
order `int`

The order of UniC-p. Effective accuracy is order + 1.

 None

Returns:

Type Description
Tensor

torch.Tensor: The corrected sample tensor.

multistep_uni_p_bh_update 

multistep_uni_p_bh_update(
    model_output: Tensor,
    *args,
    sample: Tensor | None = None,
    order: int | None = None,
    **kwargs,
) -> Tensor

One step for the UniP (B(h) version) predictor.

Parameters:

Name Type Description Default
model_output `torch.Tensor`

Direct output from the diffusion model.

 required
sample `torch.Tensor`

Current sample.

 None
order `int`

The order of UniP at this timestep.

 None

Returns:

Type Description
Tensor

torch.Tensor: The sample tensor at the previous timestep.

scale_model_input 

scale_model_input(
    sample: Tensor, *args, **kwargs
) -> Tensor

Ensures interchangeability with schedulers that need to scale the denoising model input.

Parameters:

Name Type Description Default
sample `torch.Tensor`

The input sample.

 required

Returns:

Type Description
Tensor

torch.Tensor: A scaled input sample (unchanged for this scheduler).

set_begin_index 

set_begin_index(begin_index: int = 0) -> None

Sets the begin index for the scheduler. Run from pipeline before inference.

Parameters:

Name Type Description Default
begin_index `int`

The begin index for the scheduler.

 0

set_shift 

set_shift(shift: float) -> None

Set the shift parameter for the scheduler.

set_timesteps 

set_timesteps(
    num_inference_steps: int | None = None,
    device: str | device | None = None,
    sigmas: list[float] | None = None,
    mu: float | None = None,
    shift: float | None = None,
) -> None

Sets the discrete timesteps used for the diffusion chain (run before inference).

Parameters:

Name Type Description Default
num_inference_steps `int`

Total number of timesteps.

 None
device `str` or `torch.device`, *optional*

The device to move timesteps to.

 None
sigmas `list[float]`, *optional*

Custom sigma schedule.

 None
mu `float`, *optional*

Parameter for dynamic shifting.

 None
shift `float`, *optional*

Override shift parameter.

 None

step 

step(
    model_output: Tensor,
    timestep: int | Tensor,
    sample: Tensor,
    return_dict: bool = True,
    generator: Generator | None = None,
) -> SchedulerOutput | tuple

Predict the sample from the previous timestep by reversing the SDE using multistep UniPC.

Parameters:

Name Type Description Default
model_output `torch.Tensor`

Direct output from the diffusion model.

 required
timestep `int`

Current discrete timestep in the diffusion chain.

 required
sample `torch.Tensor`

Current sample created by the diffusion process.

 required
return_dict `bool`

Whether to return a SchedulerOutput or tuple.

 True

Returns:

Type Description
SchedulerOutput | tuple

SchedulerOutput or tuple: The sample tensor at the previous timestep.

time_shift 

time_shift(mu: float, sigma: float, t: ndarray) -> ndarray

Apply time shift transformation.