Diffusers documentation
HeliosScheduler
HeliosScheduler
HeliosScheduler is based on the pyramidal flow-matching sampling introduced in Helios.
HeliosScheduler
class diffusers.HeliosScheduler
< source >( num_train_timesteps: int = 1000shift: float = 1.0stages: int = 3stage_range: list = [0, 0.3333333333333333, 0.6666666666666666, 1]gamma: float = 0.3333333333333333thresholding: bool = Falseprediction_type: str = 'flow_prediction'solver_order: int = 2predict_x0: bool = Truesolver_type: str = 'bh2'lower_order_final: bool = Truedisable_corrector: list = []solver_p: SchedulerMixin = Noneuse_flow_sigmas: bool = Truescheduler_type: str = 'unipc'use_dynamic_shifting: bool = Falsetime_shift_type: typing.Literal['exponential', 'linear'] = 'exponential' )
convert_model_output
< source >( model_output: Tensor*argssample: Tensor = Nonesigma: Tensor = None**kwargs ) → torch.Tensor
Parameters
- model_output (
torch.Tensor) — The direct output from the learned diffusion model. - timestep (
int) — The current discrete timestep in the diffusion chain. - sample (
torch.Tensor) — A current instance of a sample created by the diffusion process. - sigma (
torch.Tensor, optional) — The sigma of the current step in the noise schedule.
Returns
torch.Tensor
The converted model output.
Convert the model output to the corresponding type the UniPC algorithm needs.
initialize the global timesteps and sigmas
Init the timesteps for each stage
multistep_uni_c_bh_update
< source >( this_model_output: Tensor*argslast_sample: Tensor = Nonethis_sample: Tensor = Noneorder: int = Nonesigma_before: Tensor = Nonesigma: Tensor = None**kwargs ) → torch.Tensor
Parameters
- this_model_output (
torch.Tensor) — The model outputs atx_t. - this_timestep (
int) — The current timestept. - last_sample (
torch.Tensor) — The generated sample before the last predictorx_{t-1}. - this_sample (
torch.Tensor) — The generated sample after the last predictorx_{t}. - order (
int) — Thepof UniC-p at this step. The effective order of accuracy should beorder + 1. - sigma_before (
torch.Tensor, optional) — The sigma of the previous step in the noise schedule. - sigma (
torch.Tensor, optional) — The sigma of the current step in the noise schedule.
Returns
torch.Tensor
The corrected sample tensor at the current timestep.
One step for the UniC (B(h) version).
multistep_uni_p_bh_update
< source >( model_output: Tensor*argssample: Tensor = Noneorder: int = Nonesigma: Tensor = Nonesigma_next: Tensor = None**kwargs ) → torch.Tensor
Parameters
- model_output (
torch.Tensor) — The direct output from the learned diffusion model at the current timestep. - prev_timestep (
int) — The previous discrete timestep in the diffusion chain. - sample (
torch.Tensor) — A current instance of a sample created by the diffusion process. - order (
int) — The order of UniP at this timestep (corresponds to the p in UniPC-p). - sigma (
torch.Tensor, optional) — The sigma of the current step in the noise schedule. - sigma_next (
torch.Tensor, optional) — The sigma of the next step in the noise schedule.
Returns
torch.Tensor
The sample tensor at the previous timestep.
One step for the UniP (B(h) version). Alternatively, self.solver_p is used if is specified.
set_begin_index
< source >( begin_index: int = 0 )
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
Apply time shifting to the sigmas.
scheduling_helios
Update on GitHub