Stable Diffusion Turbo Support (#1395)

* Add support for SD Turbo

* Set Leading as default in euler_ancestral discrete

* Use the appropriate default values for n_steps and guidance_scale.

---------

Co-authored-by: Laurent <laurent.mazare@gmail.com>

candle-transformers/src/models/stable_diffusion/ddim.rs

@@ -7,7 +7,9 @@
 //!
 //! Denoising Diffusion Implicit Models, J. Song et al, 2020.
 //! https://arxiv.org/abs/2010.02502
-use super::schedulers::{betas_for_alpha_bar, BetaSchedule, PredictionType, TimestepSpacing};
+use super::schedulers::{
+    betas_for_alpha_bar, BetaSchedule, PredictionType, Scheduler, SchedulerConfig, TimestepSpacing,
+};
 use candle::{Result, Tensor};
 
 /// The configuration for the DDIM scheduler.
@@ -48,6 +50,12 @@ impl Default for DDIMSchedulerConfig {
     }
 }
 
+impl SchedulerConfig for DDIMSchedulerConfig {
+    fn build(&self, inference_steps: usize) -> Result<Box<dyn Scheduler>> {
+        Ok(Box::new(DDIMScheduler::new(inference_steps, *self)?))
+    }
+}
+
 /// The DDIM scheduler.
 #[derive(Debug, Clone)]
 pub struct DDIMScheduler {
@@ -63,7 +71,7 @@ impl DDIMScheduler {
     /// Creates a new DDIM scheduler given the number of steps to be
     /// used for inference as well as the number of steps that was used
     /// during training.
-    pub fn new(inference_steps: usize, config: DDIMSchedulerConfig) -> Result<Self> {
+    fn new(inference_steps: usize, config: DDIMSchedulerConfig) -> Result<Self> {
         let step_ratio = config.train_timesteps / inference_steps;
         let timesteps: Vec<usize> = match config.timestep_spacing {
             TimestepSpacing::Leading => (0..(inference_steps))
@@ -115,19 +123,11 @@ impl DDIMScheduler {
             config,
         })
     }
+}
 
-    pub fn timesteps(&self) -> &[usize] {
-        self.timesteps.as_slice()
-    }
-
-    ///  Ensures interchangeability with schedulers that need to scale the denoising model input
-    /// depending on the current timestep.
-    pub fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> Result<Tensor> {
-        Ok(sample)
-    }
-
+impl Scheduler for DDIMScheduler {
     /// Performs a backward step during inference.
-    pub fn step(&self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> Result<Tensor> {
+    fn step(&self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> Result<Tensor> {
         let timestep = if timestep >= self.alphas_cumprod.len() {
             timestep - 1
         } else {
@@ -186,7 +186,17 @@ impl DDIMScheduler {
         }
     }
 
-    pub fn add_noise(&self, original: &Tensor, noise: Tensor, timestep: usize) -> Result<Tensor> {
+    ///  Ensures interchangeability with schedulers that need to scale the denoising model input
+    /// depending on the current timestep.
+    fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> Result<Tensor> {
+        Ok(sample)
+    }
+
+    fn timesteps(&self) -> &[usize] {
+        self.timesteps.as_slice()
+    }
+
+    fn add_noise(&self, original: &Tensor, noise: Tensor, timestep: usize) -> Result<Tensor> {
         let timestep = if timestep >= self.alphas_cumprod.len() {
             timestep - 1
         } else {
@@ -197,7 +207,7 @@ impl DDIMScheduler {
         (original * sqrt_alpha_prod)? + (noise * sqrt_one_minus_alpha_prod)?
     }
 
-    pub fn init_noise_sigma(&self) -> f64 {
+    fn init_noise_sigma(&self) -> f64 {
         self.init_noise_sigma
     }
 }