For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Python class

UniPCMultistepScheduler

`UniPCMultistepScheduler`

class max.pipelines.diffusion.schedulers.UniPCMultistepScheduler(num_train_timesteps=1000, solver_order=2, prediction_type='flow_prediction', predict_x0=True, solver_type='bh2', lower_order_final=True, disable_corrector=None, thresholding=False, dynamic_thresholding_ratio=0.995, sample_max_value=1.0, use_flow_sigmas=False, flow_shift=1.0, time_shift_type='exponential', final_sigmas_type='zero', order=1, **unused_kwargs)

source

Bases: object

NumPy-only UniPC multistep scheduler for diffusion models.

Implements the UniPC (Unified Predictor-Corrector) framework with B(h) updates for fast sampling of diffusion models. Supports flow-matching prediction type and the BH2 solver variant.

This scheduler is designed for use with the Wan 2.2 T2V pipeline.

Parameters:

num_train_timesteps (int)
solver_order (int)
prediction_type (str)
predict_x0 (bool)
solver_type (str)
lower_order_final (bool)
disable_corrector (list[int] | None)
thresholding (bool)
dynamic_thresholding_ratio (float)
sample_max_value (float)
use_flow_sigmas (bool)
flow_shift (float)
time_shift_type (str)
final_sigmas_type (str)
order (int)

`begin_index`

property begin_index: int | None

source

The index for the first timestep.

`build_step_coefficients()`

build_step_coefficients()

source

Pre-compute all UniPC step coefficients as a dense numpy array.

Returns array of shape [num_steps, 9] with columns [sigma, corrected_input_scale, corrector_sample_scale, corrector_m0_scale, corrector_m1_scale, corrector_mt_scale, predictor_sample_scale, predictor_m0_scale, predictor_m1_scale].

Must be called AFTER set_timesteps().

Return type:: ndarray[tuple[Any, …], dtype[float32]]

`convert_model_output()`

convert_model_output(model_output, sample)

source

Convert the raw model output to x0 prediction.

For flow_prediction: x0 = sample - sigma_t * model_output.

Parameters:

model_output (ndarray) – Direct output from the diffusion model.
sample (ndarray) – Current noisy sample.

Returns:

Converted output (x0 prediction when predict_x0=True).

Return type:

ndarray

`index_for_timestep()`

index_for_timestep(timestep)

source

Find the index for a given timestep in the schedule.

Parameters:: timestep (float)
Return type:: int

`multistep_uni_c_bh_update()`

multistep_uni_c_bh_update(this_model_output, last_sample, this_sample, order)

source

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

Parameters:

this_model_output (ndarray) – Converted model output at current step.
last_sample (ndarray) – Sample before the last predictor step.
this_sample (ndarray) – Sample after the last predictor step.
order (int) – Corrector order.

Returns:

Corrected sample.

Return type:

ndarray

`multistep_uni_p_bh_update()`

multistep_uni_p_bh_update(model_output, sample, order)

source

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

Parameters:

model_output (ndarray) – Direct (non-converted) model output.
sample (ndarray) – Current sample.
order (int) – Solver order for this step.

Returns:

Predicted sample at the next timestep.

Return type:

ndarray

`retrieve_timesteps_and_sigmas()`

retrieve_timesteps_and_sigmas(image_seq_len, num_inference_steps, reverse=False, sigma_min=None)

source

Build timestep/sigma schedule, with flow-matching support.

Also calls set_timesteps so the scheduler is ready for stepping.

Parameters:

image_seq_len (int) – Sequence length (unused for flow-matching).
num_inference_steps (int) – Number of denoising steps.
reverse (bool) – Whether to reverse timesteps (non-flow only).
sigma_min (float | None) – Unused. Accepted for interface compatibility with other schedulers.

Returns:

Tuple of (timesteps, sigmas) as float32 arrays.

Return type:

tuple[ndarray[tuple[Any, …], dtype[float32]], ndarray[tuple[Any, …], dtype[float32]]]

`set_begin_index()`

set_begin_index(begin_index=0)

source

Sets the begin index for the scheduler.

Parameters:: begin_index (int)
Return type:: None

`set_timesteps()`

set_timesteps(num_inference_steps)

source

Initialize internal state for a denoising run.

Must be called before the first step() call.

Parameters:: num_inference_steps (int) – Number of denoising steps.
Return type:: None

`step()`

step(model_output, timestep, sample)

source

Predict the sample at the previous timestep using UniPC.

Orchestrates the corrector (UniC) and predictor (UniP) updates.

Parameters:

model_output (ndarray) – Direct output from the learned diffusion model.
timestep (float | int) – Current discrete timestep.
sample (ndarray) – Current noisy sample.

Returns:

Denoised sample at the next timestep.

Return type:

ndarray

`step_index`

property step_index: int | None

source

The index counter for the current timestep.

UniPCMultistepScheduler​

begin_index​

build_step_coefficients()​

convert_model_output()​

index_for_timestep()​

multistep_uni_c_bh_update()​

multistep_uni_p_bh_update()​

retrieve_timesteps_and_sigmas()​

set_begin_index()​

set_timesteps()​

step()​

step_index​

`UniPCMultistepScheduler`

`begin_index`

`build_step_coefficients()`

`convert_model_output()`

`index_for_timestep()`

`multistep_uni_c_bh_update()`

`multistep_uni_p_bh_update()`

`retrieve_timesteps_and_sigmas()`

`set_begin_index()`

`set_timesteps()`

`step()`

`step_index`