vllm_omni.diffusion.models.schedulers ¶

Modules:

Name	Description
`base`	Base scheduler class for diffusion models.
`scheduling_dmd2_euler`
`scheduling_flow_match_euler_discrete`	Flow Match Euler scheduler copied from Diffusers v0.39.0.
`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

FlowMatchEulerDiscreteScheduler ¶

Bases: SchedulerMixin, ConfigMixin

Euler scheduler.

This model inherits from [SchedulerMixin] and [ConfigMixin]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving.

Parameters:

Name	Type	Description	Default
`num_train_timesteps`	`int`, defaults to 1000	The number of diffusion steps to train the model.	`1000`
`shift`	`float`, defaults to 1.0	The shift value for the timestep schedule.	`1.0`
`use_dynamic_shifting`	`bool`, defaults to False	Whether to apply timestep shifting on-the-fly based on the image resolution.	`False`
`base_shift`	`float`, defaults to 0.5	Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent with desired output.	`0.5`
`max_shift`	`float`, defaults to 1.15	Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be more exaggerated or stylized.	`1.15`
`base_image_seq_len`	`int`, defaults to 256	The base image sequence length.	`256`
`max_image_seq_len`	`int`, defaults to 4096	The maximum image sequence length.	`4096`
`invert_sigmas`	`bool`, defaults to False	Whether to invert the sigmas.	`False`
`shift_terminal`	`float`, defaults to None	The end value of the shifted timestep schedule.	`None`
`use_karras_sigmas`	`bool`, defaults to False	Whether to use Karras sigmas for step sizes in the noise schedule during sampling.	`False`
`use_exponential_sigmas`	`bool`, defaults to False	Whether to use exponential sigmas for step sizes in the noise schedule during sampling.	`False`
`use_beta_sigmas`	`bool`, defaults to False	Whether to use beta sigmas for step sizes in the noise schedule during sampling.	`False`
`time_shift_type`	`str`, defaults to "exponential"	The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".	`'exponential'`
`stochastic_sampling`	`bool`, defaults to False	Whether to use stochastic sampling.	`False`

begin_index `property` ¶

begin_index

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

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

order = 1

shift `property` ¶

shift

The value used for shifting.

sigma_max `instance-attribute` ¶

sigma_max = self.sigmas[0].item()

sigma_min `instance-attribute` ¶

sigma_min = self.sigmas[-1].item()

sigmas `instance-attribute` ¶

sigmas = sigmas.to('cpu')

step_index `property` ¶

step_index

The index counter for current timestep. It will increase 1 after each scheduler step.

timesteps `instance-attribute` ¶

timesteps = sigmas * num_train_timesteps

index_for_timestep ¶

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

Get the index for the given timestep.

Parameters:

Name	Type	Description	Default
`timestep`	`float` or `torch.FloatTensor`	The timestep to find the index for.	required
`schedule_timesteps`	`torch.FloatTensor`, optional	The schedule timesteps to validate against. If `None`, the scheduler's timesteps are used.	`None`

Returns:

Type	Description
`int`	`int`: The index of the timestep.

scale_noise ¶

scale_noise(
    sample: FloatTensor,
    timestep: float | FloatTensor,
    noise: FloatTensor | None = None,
) -> FloatTensor

Forward process in flow-matching

Parameters:

Name	Type	Description	Default
`sample`	`torch.FloatTensor`	The input sample.	required
`timestep`	`torch.FloatTensor`	The current timestep in the diffusion chain.	required
`noise`	`torch.FloatTensor`	The noise tensor.	`None`

Returns:

Type	Description
`FloatTensor`	`torch.FloatTensor`: A scaled input sample.

set_begin_index ¶

set_begin_index(begin_index: int = 0)

Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

Parameters:

Name	Type	Description	Default
`begin_index`	`int`, defaults to `0`	The begin index for the scheduler.	`0`

set_shift ¶

set_shift(shift: float)

Sets the shift value for the scheduler.

Parameters:

Name	Type	Description	Default
`shift`	`float`	The shift value to be set.	required

set_timesteps ¶

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

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

Parameters:

Name	Type	Description	Default
`num_inference_steps`	`int`, optional	The number of diffusion steps used when generating samples with a pre-trained model.	`None`
`device`	`str` or `torch.device`, optional	The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.	`None`
`sigmas`	`list[float]`, optional	Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed automatically.	`None`
`mu`	`float`, optional	Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep shifting.	`None`
`timesteps`	`list[float]`, optional	Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed automatically.	`None`

step ¶

step(
    model_output: FloatTensor,
    timestep: float | FloatTensor,
    sample: FloatTensor,
    s_churn: float = 0.0,
    s_tmin: float = 0.0,
    s_tmax: float = float("inf"),
    s_noise: float = 1.0,
    generator: Generator | None = None,
    per_token_timesteps: Tensor | None = None,
    return_dict: bool = True,
) -> FlowMatchEulerDiscreteSchedulerOutput | tuple

Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise).

Parameters:

Name	Type	Description	Default
`model_output`	`torch.FloatTensor`	The direct output from learned diffusion model.	required
`timestep`	`float`	The current discrete timestep in the diffusion chain.	required
`sample`	`torch.FloatTensor`	A current instance of a sample created by the diffusion process.	required
`s_churn`	`float`		`0.0`
`s_tmin`	`float`		`0.0`
`s_tmax`	`float`		`float('inf')`
`s_noise`	`float`, defaults to 1.0	Scaling factor for noise added to the sample.	`1.0`
`generator`	`torch.Generator`, optional	A random number generator.	`None`
`per_token_timesteps`	`torch.Tensor`, optional	The timesteps for each token in the sample.	`None`
`return_dict`	`bool`, defaults to `True`	Whether or not to return a [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.	`True`

Returns:

Type	Description
`FlowMatchEulerDiscreteSchedulerOutput \| tuple`	[`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: If return_dict is `True`, [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor.

stretch_shift_to_terminal ¶

stretch_shift_to_terminal(t: Tensor) -> Tensor

Stretches and shifts the timestep schedule to ensure it terminates at the configured shift_terminal config value.

Reference: https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51

Parameters:

Name	Type	Description	Default
`t`	`torch.Tensor`	A tensor of timesteps to be stretched and shifted.	required

Returns:

Type	Description
`Tensor`	`torch.Tensor`: A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`.

time_shift ¶

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

Apply time shifting to the sigmas.

Parameters:

Name	Type	Description	Default
`mu`	`float`	The mu parameter for the time shift.	required
`sigma`	`float`	The sigma parameter for the time shift.	required
`t`	`torch.Tensor`	The input timesteps.	required

Returns:

Type	Description
`Tensor`	`torch.Tensor`: The time-shifted timesteps.

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 = self.sigmas[0].item()

sigma_min `instance-attribute` ¶

sigma_min = self.sigmas[-1].item()

sigmas `instance-attribute` ¶

sigmas = self.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 ¶

DMD2EulerScheduler ¶

set_timesteps ¶

FlowMatchEulerDiscreteScheduler ¶

begin_index property ¶

order class-attribute instance-attribute ¶

shift property ¶

sigma_max instance-attribute ¶

sigma_min instance-attribute ¶

sigmas instance-attribute ¶

step_index property ¶

timesteps instance-attribute ¶

index_for_timestep ¶

scale_noise ¶

set_begin_index ¶

set_shift ¶

set_timesteps ¶

step ¶

stretch_shift_to_terminal ¶

time_shift ¶

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` ¶

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

shift `property` ¶

sigma_max `instance-attribute` ¶

sigma_min `instance-attribute` ¶

sigmas `instance-attribute` ¶

step_index `property` ¶

timesteps `instance-attribute` ¶

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` ¶