Refactor common parts of multistep/singlestep into dpmsolver

Author: rockerBOO

src/schedulers/dpmsolver.rs

@@ -0,0 +1,67 @@
+use crate::schedulers::BetaSchedule;
+use crate::schedulers::PredictionType;
+
+/// The algorithm type for the solver.
+///
+#[derive(Default, Debug, Clone, PartialEq, Eq)]
+pub enum DPMSolverAlgorithmType {
+    /// Implements the algorithms defined in <https://arxiv.org/abs/2211.01095>.
+    #[default]
+    DPMSolverPlusPlus,
+    /// Implements the algorithms defined in <https://arxiv.org/abs/2206.00927>.
+    DPMSolver,
+}
+
+/// The solver type for the second-order solver.
+/// The solver type slightly affects the sample quality, especially for
+/// small number of steps.
+#[derive(Default, Debug, Clone, PartialEq, Eq)]
+pub enum DPMSolverType {
+    #[default]
+    Midpoint,
+    Heun,
+}
+
+#[derive(Debug, Clone)]
+pub struct DPMSolverSchedulerConfig {
+    /// The value of beta at the beginning of training.
+    pub beta_start: f64,
+    /// The value of beta at the end of training.
+    pub beta_end: f64,
+    /// How beta evolved during training.
+    pub beta_schedule: BetaSchedule,
+    /// number of diffusion steps used to train the model.
+    pub train_timesteps: usize,
+    /// the order of DPM-Solver; can be `1` or `2` or `3`. We recommend to use `solver_order=2` for guided
+    /// sampling, and `solver_order=3` for unconditional sampling.
+    pub solver_order: usize,
+    /// prediction type of the scheduler function
+    pub prediction_type: PredictionType,
+    /// The threshold value for dynamic thresholding. Valid only when `thresholding: true` and
+    /// `algorithm_type: DPMSolverAlgorithmType::DPMSolverPlusPlus`.
+    pub sample_max_value: f32,
+    /// The algorithm type for the solver
+    pub algorithm_type: DPMSolverAlgorithmType,
+    /// The solver type for the second-order solver.
+    pub solver_type: DPMSolverType,
+    /// Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. We empirically
+    /// find this can stabilize the sampling of DPM-Solver for `steps < 15`, especially for steps <= 10.
+    pub lower_order_final: bool,
+}
+
+impl Default for DPMSolverSchedulerConfig {
+    fn default() -> Self {
+        Self {
+            beta_start: 0.0001,
+            beta_end: 0.02,
+            beta_schedule: BetaSchedule::Linear,
+            train_timesteps: 1000,
+            solver_order: 2,
+            prediction_type: PredictionType::Epsilon,
+            sample_max_value: 1.0,
+            algorithm_type: DPMSolverAlgorithmType::DPMSolverPlusPlus,
+            solver_type: DPMSolverType::Midpoint,
+            lower_order_final: true,
+        }
+    }
+}

src/schedulers/dpmsolver_multistep.rs

@@ -1,86 +1,22 @@
-use super::{betas_for_alpha_bar, BetaSchedule, PredictionType};
-use std::iter;
+use super::{betas_for_alpha_bar, BetaSchedule, PredictionType, dpmsolver::{DPMSolverSchedulerConfig, DPMSolverAlgorithmType, DPMSolverType}};
 use tch::{kind, Kind, Tensor};
 
-/// The algorithm type for the solver.
-///
-#[derive(Default, Debug, Clone, PartialEq, Eq)]
-pub enum DPMSolverAlgorithmType {
-    /// Implements the algorithms defined in <https://arxiv.org/abs/2211.01095>.
-    #[default]
-    DPMSolverPlusPlus,
-    /// Implements the algorithms defined in <https://arxiv.org/abs/2206.00927>.
-    DPMSolver,
-}
-
-/// The solver type for the second-order solver.
-/// The solver type slightly affects the sample quality, especially for
-/// small number of steps.
-#[derive(Default, Debug, Clone, PartialEq, Eq)]
-pub enum DPMSolverType {
-    #[default]
-    Midpoint,
-    Heun,
-}
-
-#[derive(Debug, Clone)]
-pub struct DPMSolverMultistepSchedulerConfig {
-    /// The value of beta at the beginning of training.
-    pub beta_start: f64,
-    /// The value of beta at the end of training.
-    pub beta_end: f64,
-    /// How beta evolved during training.
-    pub beta_schedule: BetaSchedule,
-    /// number of diffusion steps used to train the model.
-    pub train_timesteps: usize,
-    /// the order of DPM-Solver; can be `1` or `2` or `3`. We recommend to use `solver_order=2` for guided
-    /// sampling, and `solver_order=3` for unconditional sampling.
-    pub solver_order: usize,
-    /// prediction type of the scheduler function
-    pub prediction_type: PredictionType,
-    /// The threshold value for dynamic thresholding. Valid only when `thresholding: true` and
-    /// `algorithm_type: DPMSolverAlgorithmType::DPMSolverPlusPlus`.
-    pub sample_max_value: f32,
-    /// The algorithm type for the solver
-    pub algorithm_type: DPMSolverAlgorithmType,
-    /// The solver type for the second-order solver.
-    pub solver_type: DPMSolverType,
-    /// Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. We empirically
-    /// find this can stabilize the sampling of DPM-Solver for `steps < 15`, especially for steps <= 10.
-    pub lower_order_final: bool,
-}
-
-impl Default for DPMSolverMultistepSchedulerConfig {
-    fn default() -> Self {
-        Self {
-            beta_start: 0.00085,
-            beta_end: 0.012,
-            beta_schedule: BetaSchedule::ScaledLinear,
-            train_timesteps: 1000,
-            solver_order: 2,
-            prediction_type: PredictionType::Epsilon,
-            sample_max_value: 1.0,
-            algorithm_type: DPMSolverAlgorithmType::DPMSolverPlusPlus,
-            solver_type: DPMSolverType::Midpoint,
-            lower_order_final: true,
-        }
-    }
-}
-
 pub struct DPMSolverMultistepScheduler {
     alphas_cumprod: Vec<f64>,
     alpha_t: Vec<f64>,
     sigma_t: Vec<f64>,
     lambda_t: Vec<f64>,
     init_noise_sigma: f64,
     lower_order_nums: usize,
+    /// Direct outputs from learned diffusion model at current and latter timesteps
     model_outputs: Vec<Tensor>,
+    /// List of current discrete timesteps in the diffusion chain
     timesteps: Vec<usize>,
-    pub config: DPMSolverMultistepSchedulerConfig,
+    pub config: DPMSolverSchedulerConfig,
 }
 
 impl DPMSolverMultistepScheduler {
-    pub fn new(inference_steps: usize, config: DPMSolverMultistepSchedulerConfig) -> Self {
+    pub fn new(inference_steps: usize, config: DPMSolverSchedulerConfig) -> Self {
         let betas = match config.beta_schedule {
             BetaSchedule::ScaledLinear => Tensor::linspace(
                 config.beta_start.sqrt(),

src/schedulers/dpmsolver_singlestep.rs

@@ -1,88 +1,30 @@
 use std::iter::repeat;
 
-use super::{betas_for_alpha_bar, BetaSchedule, PredictionType};
+use super::{
+    betas_for_alpha_bar,
+    dpmsolver::{DPMSolverAlgorithmType, DPMSolverSchedulerConfig, DPMSolverType},
+    BetaSchedule, PredictionType,
+};
 use tch::{kind, Kind, Tensor};
 
-/// The algorithm type for the solver.
-///
-#[derive(Default, Debug, Clone, PartialEq, Eq)]
-pub enum DPMSolverAlgorithmType {
-    /// Implements the algorithms defined in <https://arxiv.org/abs/2211.01095>.
-    #[default]
-    DPMSolverPlusPlus,
-    /// Implements the algorithms defined in <https://arxiv.org/abs/2206.00927>.
-    DPMSolver,
-}
-
-/// The solver type for the second-order solver.
-/// The solver type slightly affects the sample quality, especially for
-/// small number of steps.
-#[derive(Default, Debug, Clone, PartialEq, Eq)]
-pub enum DPMSolverType {
-    #[default]
-    Midpoint,
-    Heun,
-}
-
-#[derive(Debug, Clone)]
-pub struct DPMSolverSinglestepSchedulerConfig {
-    /// The value of beta at the beginning of training.
-    pub beta_start: f64,
-    /// The value of beta at the end of training.
-    pub beta_end: f64,
-    /// How beta evolved during training.
-    pub beta_schedule: BetaSchedule,
-    /// number of diffusion steps used to train the model.
-    pub train_timesteps: usize,
-    /// the order of DPM-Solver; can be `1` or `2` or `3`. We recommend to use `solver_order=2` for guided
-    /// sampling, and `solver_order=3` for unconditional sampling.
-    pub solver_order: usize,
-    /// prediction type of the scheduler function
-    pub prediction_type: PredictionType,
-    /// The threshold value for dynamic thresholding. Valid only when `thresholding: true` and
-    /// `algorithm_type: DPMSolverAlgorithmType::DPMSolverPlusPlus`.
-    pub sample_max_value: f32,
-    /// The algorithm type for the solver
-    pub algorithm_type: DPMSolverAlgorithmType,
-    /// The solver type for the second-order solver.
-    pub solver_type: DPMSolverType,
-    /// Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. We empirically
-    /// find this can stabilize the sampling of DPM-Solver for `steps < 15`, especially for steps <= 10.
-    pub lower_order_final: bool,
-}
-
-impl Default for DPMSolverSinglestepSchedulerConfig {
-    fn default() -> Self {
-        Self {
-            beta_start: 0.0001,
-            beta_end: 0.02,
-            train_timesteps: 1000,
-            beta_schedule: BetaSchedule::Linear,
-            solver_order: 2,
-            prediction_type: PredictionType::Epsilon,
-            sample_max_value: 1.0,
-            algorithm_type: DPMSolverAlgorithmType::DPMSolverPlusPlus,
-            solver_type: DPMSolverType::Midpoint,
-            lower_order_final: true,
-        }
-    }
-}
-
 pub struct DPMSolverSinglestepScheduler {
     alphas_cumprod: Vec<f64>,
     alpha_t: Vec<f64>,
     sigma_t: Vec<f64>,
     lambda_t: Vec<f64>,
     init_noise_sigma: f64,
     order_list: Vec<usize>,
+    /// Direct outputs from learned diffusion model at current and latter timesteps
     model_outputs: Vec<Tensor>,
+    /// List of current discrete timesteps in the diffusion chain
     timesteps: Vec<usize>,
+    /// Current instance of sample being created by diffusion process
     sample: Option<Tensor>,
-    pub config: DPMSolverSinglestepSchedulerConfig,
+    pub config: DPMSolverSchedulerConfig,
 }
 
 impl DPMSolverSinglestepScheduler {
-    pub fn new(inference_steps: usize, config: DPMSolverSinglestepSchedulerConfig) -> Self {
+    pub fn new(inference_steps: usize, config: DPMSolverSchedulerConfig) -> Self {
         let betas = match config.beta_schedule {
             BetaSchedule::ScaledLinear => Tensor::linspace(
                 config.beta_start.sqrt(),
@@ -462,8 +404,6 @@ fn get_order_list(steps: usize, solver_order: usize, lower_order_final: bool) ->
                     .chain([&[1][..]])
                     .flatten()
                     .map(|v| *v)
-                    .collect()
-            }
         } else if solver_order == 1 {
             repeat(&[1][..]).take(steps).flatten().map(|v| *v).collect()
         } else {
@@ -473,11 +413,9 @@ fn get_order_list(steps: usize, solver_order: usize, lower_order_final: bool) ->
         if solver_order == 3 {
             repeat(&[1, 2, 3][..]).take(steps / 3).flatten().map(|v| *v).collect()
         } else if solver_order == 2 {
-            repeat(&[1, 2][..]).take(steps / 2).flatten().map(|v| *v).collect()
+            repeat(dbg!(&[1, 2][..])).take(dbg!(steps / 2)).flatten().map(|v| dbg!(*v)).collect()
         } else if solver_order == 1 {
             repeat(&[1][..]).take(steps).flatten().map(|v| *v).collect()
-        } else {
-            panic!("invalid solver_order");
         }
     }
 }

src/schedulers/mod.rs

@@ -7,6 +7,7 @@ use tch::{Kind, Tensor};
 
 pub mod ddim;
 pub mod ddpm;
+pub mod dpmsolver;
 pub mod dpmsolver_multistep;
 pub mod dpmsolver_singlestep;
 pub mod euler_ancestral_discrete;