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.