Diffusers

HeliosDMDScheduler

HeliosDMDScheduler is based on the pyramidal flow-matching sampling introduced in Helios.

HeliosDMDScheduler

class diffusers.HeliosDMDScheduler

( num_train_timesteps: int = 1000shift: float = 1.0stages: int = 3stage_range: list = [0, 0.3333333333333333, 0.6666666666666666, 1]gamma: float = 0.3333333333333333prediction_type: str = 'flow_prediction'use_flow_sigmas: bool = Trueuse_dynamic_shifting: bool = Falsetime_shift_type: typing.Literal['exponential', 'linear'] = 'linear' )

init_sigmas

< source >

( )

initialize the global timesteps and sigmas

init_sigmas_for_each_stage

< source >

( )

Init the timesteps for each stage

set_begin_index

< source >

( begin_index: int = 0 )

Parameters

begin_index (int) — The begin index for the scheduler.

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

set_timesteps

< source >

( num_inference_steps: intstage_index: int | None = Nonedevice: typing.Union[str, torch.device] = Nonesigmas: bool | None = Nonemu: bool | None = Noneis_amplify_first_chunk: bool = False )

Setting the timesteps and sigmas for each stage

time_shift

< source >

( mu: floatsigma: floatt: Tensor ) → torch.Tensor

Parameters

mu (float) — The mu parameter for the time shift.
sigma (float) — The sigma parameter for the time shift.
t (torch.Tensor) — The input timesteps.

Returns

torch.Tensor

The time-shifted timesteps.

Apply time shifting to the sigmas.

scheduling_helios_dmd

Update on GitHub