vllm_omni.diffusion.models.helios.scheduling_helios ¶

add_noise(
    original_samples, noise, timestep, sigmas, timesteps
)

convert_flow_pred_to_x0(
    flow_pred, xt, timestep, sigmas, timesteps
)

convert_model_output(
    model_output: Tensor,
    *args,
    sample: Tensor = None,
    sigma: Tensor = None,
    **kwargs,
) -> Tensor

index_for_timestep(timestep, schedule_timesteps=None)

init_sigmas()

init_sigmas_for_each_stage()

multistep_uni_c_bh_update(
    this_model_output: Tensor,
    *args,
    last_sample: Tensor = None,
    this_sample: Tensor = None,
    order: int = None,
    sigma_before: Tensor = None,
    sigma: Tensor = None,
    **kwargs,
) -> Tensor

multistep_uni_p_bh_update(
    model_output: Tensor,
    *args,
    sample: Tensor = None,
    order: int = None,
    sigma: Tensor = None,
    sigma_next: Tensor = None,
    **kwargs,
) -> Tensor

reset_scheduler_history()

set_begin_index(begin_index: int = 0)

set_timesteps(
    num_inference_steps: int,
    stage_index: int | None = None,
    device: str | device = None,
    sigmas: bool | None = None,
    mu: bool | None = None,
    is_amplify_first_chunk: bool = False,
)

step(
    model_output: FloatTensor,
    timestep: float | FloatTensor = None,
    sample: FloatTensor = None,
    generator: Generator | None = None,
    return_dict: bool = True,
    cur_sampling_step: int = 0,
    dmd_noisy_tensor: FloatTensor | None = None,
    dmd_sigmas: FloatTensor | None = None,
    dmd_timesteps: FloatTensor | None = None,
    all_timesteps: FloatTensor | None = None,
) -> HeliosSchedulerOutput | tuple

step_dmd(
    model_output: FloatTensor,
    timestep: float | FloatTensor = None,
    sample: FloatTensor = None,
    generator: Generator | None = None,
    return_dict: bool = True,
    cur_sampling_step: int = 0,
    dmd_noisy_tensor: FloatTensor | None = None,
    dmd_sigmas: FloatTensor | None = None,
    dmd_timesteps: FloatTensor | None = None,
    all_timesteps: FloatTensor | None = None,
)

step_euler(
    model_output: FloatTensor,
    timestep: float | FloatTensor = None,
    sample: FloatTensor = None,
    generator: Generator | None = None,
    sigma: FloatTensor | None = None,
    sigma_next: FloatTensor | None = None,
    return_dict: bool = True,
) -> HeliosSchedulerOutput | tuple

step_unipc(
    model_output: Tensor,
    timestep: int | Tensor = None,
    sample: Tensor = None,
    return_dict: bool = True,
    model_outputs: list = None,
    timestep_list: list = None,
    sigma_before: Tensor = None,
    sigma: Tensor = None,
    sigma_next: Tensor = None,
    cus_step_index: int = None,
    cus_lower_order_num: int = None,
    cus_this_order: int = None,
    cus_last_sample: Tensor = None,
) -> HeliosSchedulerOutput | tuple

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