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

This commit is contained in:
vmobilis 2026-08-19 17:29:29 +03:00 committed by GitHub
parent 88b044be7f
commit 760717a060
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 47 additions and 23 deletions

View File

@ -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},
{"",

View File

@ -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,26 +2660,27 @@ 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 (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() == static_cast<size_t>(max_order - 1)) {
if (hist_size_p1 == max_order) {
hist.erase(hist.begin());
}
hist.push_back(std::move(d_cur));
}
}
return x;
}