Skip to content

vllm_omni.diffusion.models.schedulers.base

Base scheduler class for diffusion models.

BaseScheduler

Bases: ABC

Abstract base class for schedulers.

Subclasses must define
  • timesteps: torch.Tensor
  • order: int
  • num_train_timesteps: int

num_train_timesteps instance-attribute 

num_train_timesteps: int

order instance-attribute 

order: int

timesteps instance-attribute 

timesteps: Tensor

scale_model_input abstractmethod 

scale_model_input(
    sample: Tensor, timestep: int | None = None
) -> Tensor

Scale the model input.

set_shift abstractmethod 

set_shift(shift: float) -> None

Set the shift parameter for the scheduler.

set_timesteps abstractmethod 

set_timesteps(*args, **kwargs) -> None

Set the timesteps for the scheduler.