feat: add linear multi-step sampling method (#1843)

examples/common/common.cpp

@@ -1008,7 +1008,7 @@ ArgOptions SDGenerationParams::get_options() {
          &hires_upscaler},
         {"",
          "--extra-sample-args",
-         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
+         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_divisions",
          (int)',',
          &extra_sample_args},
         {"",
@@ -1538,12 +1538,12 @@ ArgOptions SDGenerationParams::get_options() {
          on_seed_arg},
         {"",
          "--sampling-method",
-         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
+         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
          "(default: euler for Flux/SD3/Wan, euler_a otherwise)",
          on_sample_method_arg},
         {"",
          "--high-noise-sampling-method",
-         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
+         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
          " default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_high_noise_sample_method_arg},
         {"",

include/stable-diffusion.h

@@ -56,6 +56,7 @@ enum sample_method_t {
     EULER_GE_SAMPLE_METHOD,
     DPMPP2M_SDE_SAMPLE_METHOD,
     DPMPP2M_SDE_BT_SAMPLE_METHOD,
+    LMS_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 

src/runtime/denoiser.hpp

@@ -2578,6 +2578,88 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
     return x;
 }
 
+static sd::Tensor<float> sample_lms(denoise_cb_t model,
+                                    sd::Tensor<float> x,
+                                    const std::vector<float>& sigmas,
+                                    const SamplerExtraArgs& extra_sample_args) {
+    // Linear Multi-Step from https://github.com/crowsonkb/k-diffusion
+
+    int divisions = 1000;
+    for (const auto& [key, value] : extra_sample_args) {
+        int parsed = 0;
+        if (key == "lms_divisions") {
+            if (!parse_strict_int(value, parsed)) {
+                LOG_WARN("ignoring invalid lms extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                continue;
+            }
+            divisions = parsed;  // std::max(1, parsed);
+            // values above 35M produce noise, can be fixed by double precision
+            // values < 1 always produce noise
+        }
+    }
+    LOG_DEBUG("linear multi-step sampler: integrating using %i division%s", divisions, (divisions == 1) ? "" : "s");
+
+    auto linear_multistep_coeff = [=](const int order, const int m, const int j) -> float {
+        if (!divisions)
+            return sigmas[m + 1] - sigmas[m];  // delta / 0 * 0
+#define LMS_PRECISION float                    // double
+        const LMS_PRECISION a = sigmas[m], dx = (sigmas[m + 1] - a) / divisions, s = sigmas[m - j];
+        const LMS_PRECISION b0 = a + 0.5f * dx;  // using Riemann middle integral
+        LMS_PRECISION sum      = 0.0f;
+        for (int h = 0; h < divisions; h++) {
+            const LMS_PRECISION b = h * dx + b0;
+            LMS_PRECISION prod    = 1.0f;
+            for (int k = 0; k < j; k++) {
+                const LMS_PRECISION t = sigmas[m - k];
+                prod *= (b - t) / (s - t);
+            }
+            for (int k = j + 1; k < order; k++) {
+                const LMS_PRECISION t = sigmas[m - k];
+                prod *= (b - t) / (s - t);
+            }
+            sum += prod;
+        }
+        return sum * dx;
+    };
+
+    const int max_order = 4;
+    float lms_coeff[max_order];
+    std::vector<sd::Tensor<float>> hist = {};
+
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        const float sigma = sigmas[i];
+
+        auto denoised_opt = model(x, sigma, i + 1);
+        if (denoised_opt.pred.empty()) {
+            return {};
+        }
+        sd::Tensor<float> denoised = std::move(denoised_opt.pred);
+
+        const int order = std::min(max_order, i + 1);
+        for (int c = 0; c < order; c++)  // computing coefficients
+            lms_coeff[c] = linear_multistep_coeff(order, i, c);
+
+        sd::Tensor<float> d_cur = (x - denoised) / sigma;
+        switch (order) {
+            case 4:  // derivative + 3 history points
+                x += hist[hist.size() - 2] * lms_coeff[3];
+            case 3:
+                x += hist[hist.size() - 1] * lms_coeff[2];
+            case 2:
+                x += hist.back() * lms_coeff[1];
+            case 1:
+                x += d_cur * lms_coeff[0];
+        }
+
+        if (hist.size() == static_cast<size_t>(max_order - 1)) {
+            hist.erase(hist.begin());
+        }
+        hist.push_back(std::move(d_cur));
+    }
+    return x;
+}
+
 static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
                                              sd::Tensor<float> x,
                                              const std::vector<float>& sigmas) {
@@ -2739,6 +2821,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
             return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case TCD_SAMPLE_METHOD:
             return sample_tcd(model, std::move(x), sigmas, rng, eta);
+        case LMS_SAMPLE_METHOD:
+            return sample_lms(model, std::move(x), sigmas, extra_args);
         case EULER_CFG_PP_SAMPLE_METHOD:
             return sample_euler_cfg_pp(model, std::move(x), sigmas);
         case EULER_A_CFG_PP_SAMPLE_METHOD:

src/stable-diffusion.cpp

@@ -143,6 +143,7 @@ const char* sampling_methods_str[] = {
     "Euler GE",
     "DPM++ (2M) SDE",
     "DPM++ (2M) SDE BT",
+    "LMS",
 };
 
 /*================================================== Helper Functions ================================================*/
@@ -3208,6 +3209,7 @@ const char* sample_method_to_str[] = {
     "euler_ge",
     "dpm++2m_sde",
     "dpm++2m_sde_bt",
+    "lms",
 };
 
 const char* sd_sample_method_name(enum sample_method_t sample_method) {