fix: make max_order of lms sampler configurable (#1885)

Author: vmobilis

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; lms supports lms_divisions",
+         "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_max_order, lms_shift, lms_divisions",
          (int)',',
          &extra_sample_args},
         {"",

src/runtime/denoiser.hpp

@@ -2582,27 +2582,49 @@ 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
-
+    // Linear Multi-Step from https://github.com/crowsonkb/k-diffusion,
+    // modified with "history shift" value, which seemingly needs less steps
     int divisions = 1000;
+    int max_order = 4;
+    int shift = 1;  // 4, 0 - original; 4, 1 - PR #1843; 3, 1 - smoother image
     for (const auto& [key, value] : extra_sample_args) {
         int parsed = 0;
+        if (key == "lms_max_order") {
+            if (!parse_strict_int(value, parsed)) {
+                LOG_WARN("ignoring invalid lms extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                continue;
+            }
+            max_order = std::max(1, parsed);
+            // smaller values make the result softer, closer to Euler
+            // higher values need more steps
+            // values above 12 can produce NaNs, depending on steps and scheduler
+        }
+        if (key == "lms_shift") {
+            if (!parse_strict_int(value, parsed)) {
+                LOG_WARN("ignoring invalid lms extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                continue;
+            }
+            shift = std::max(0, parsed);
+            // for a low number of steps, the value 1 works best
+        }
         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
+            // values above 30M require double precision in the integrator
+            // (they are needless and just slow the integration down, but
+            //  with single precision they softly produce noise
+            //  near the 35M, it can be used for distorted generations)
         }
     }
-    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
+#define LMS_PRECISION float  // when divisions > 30 millions, the double precision fixes noise
         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;
@@ -2622,11 +2644,12 @@ static sd::Tensor<float> sample_lms(denoise_cb_t model,
         return sum * dx;
     };
 
-    const int max_order = 4;
-    float lms_coeff[max_order];
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    max_order = std::min(max_order, steps);  // history can not be larger than steps
+    LOG_DEBUG("linear multi-step sampler: lms_max_order = %i, lms_shift = %i, lms_divisions = %i", max_order, shift, divisions);
+    std::vector<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];
 
@@ -2637,25 +2660,26 @@ static sd::Tensor<float> sample_lms(denoise_cb_t model,
         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());
+        x += d_cur * lms_coeff[0];
+        if (max_order > 1) {  // if max_order == 1, the history is not used (order always < 2)
+            int hist_size_p1 = hist.size() + 1;
+            if (i) {  // history does not exist at 1st step
+                int hist_max = hist.size() - 1;
+                for (int c = 2; c <= order; c++)
+                    x += hist[std::min(hist_max, hist_size_p1 - c + shift)] * lms_coeff[c - 1];
+                    // max_order == 4  =>  hist[] index = 2, 1, 0
+                    // shift == 1      =>  hist[] index = 2, 2, 1
+            }
+            if (hist_size_p1 == max_order) {
+                hist.erase(hist.begin());
+            }
+            hist.push_back(std::move(d_cur));
         }
-        hist.push_back(std::move(d_cur));
     }
     return x;
 }