vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep ¶

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

This scheduler implements the UniPC (Unified Predictor-Corrector) algorithm adapted for flow matching, providing faster convergence than simple Euler methods while maintaining quality.

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.

vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep ¶

FlowUniPCMultistepScheduler ¶

begin_index property ¶

disable_corrector instance-attribute ¶

last_sample instance-attribute ¶

lower_order_nums instance-attribute ¶

model_outputs instance-attribute ¶

num_inference_steps instance-attribute ¶

num_train_timesteps instance-attribute ¶

order class-attribute instance-attribute ¶

predict_x0 instance-attribute ¶

sigma_max instance-attribute ¶

sigma_min instance-attribute ¶

sigmas instance-attribute ¶

solver_p instance-attribute ¶

step_index property ¶

this_order instance-attribute ¶

timestep_list instance-attribute ¶

timesteps instance-attribute ¶

add_noise ¶

convert_model_output ¶

index_for_timestep ¶

multistep_uni_c_bh_update ¶

multistep_uni_p_bh_update ¶

scale_model_input ¶

set_begin_index ¶

set_shift ¶

set_timesteps ¶

step ¶

time_shift ¶

begin_index `property` ¶

disable_corrector `instance-attribute` ¶

last_sample `instance-attribute` ¶

lower_order_nums `instance-attribute` ¶

model_outputs `instance-attribute` ¶

num_inference_steps `instance-attribute` ¶

num_train_timesteps `instance-attribute` ¶

order `class-attribute` `instance-attribute` ¶

predict_x0 `instance-attribute` ¶

sigma_max `instance-attribute` ¶

sigma_min `instance-attribute` ¶

sigmas `instance-attribute` ¶

solver_p `instance-attribute` ¶

step_index `property` ¶

this_order `instance-attribute` ¶

timestep_list `instance-attribute` ¶

timesteps `instance-attribute` ¶