pub trait Scheduler {
    // Required methods
    fn timesteps(&self) -> &[usize];
    fn add_noise(
        &self,
        original: &Tensor,
        noise: Tensor,
        timestep: usize
    ) -> Result<Tensor>;
    fn init_noise_sigma(&self) -> f64;
    fn scale_model_input(
        &self,
        sample: Tensor,
        _timestep: usize
    ) -> Result<Tensor>;
    fn step(
        &self,
        model_output: &Tensor,
        timestep: usize,
        sample: &Tensor
    ) -> Result<Tensor>;
}
Expand description

This trait represents a scheduler for the diffusion process.

Required Methods§

source

fn timesteps(&self) -> &[usize]

source

fn add_noise( &self, original: &Tensor, noise: Tensor, timestep: usize ) -> Result<Tensor>

source

fn init_noise_sigma(&self) -> f64

source

fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> Result<Tensor>

source

fn step( &self, model_output: &Tensor, timestep: usize, sample: &Tensor ) -> Result<Tensor>

Implementors§