diff --git a/examples/common/common.cpp b/examples/common/common.cpp index ceb78143..66d1d331 100644 --- a/examples/common/common.cpp +++ b/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}, {"", diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index bd568c68..cbc68647 100644 --- a/include/stable-diffusion.h +++ b/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 }; diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index 67aac538..f58f42b5 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -2578,6 +2578,88 @@ static sd::Tensor sample_tcd(denoise_cb_t model, return x; } +static sd::Tensor sample_lms(denoise_cb_t model, + sd::Tensor x, + const std::vector& 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> hist = {}; + + int steps = static_cast(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 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 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(max_order - 1)) { + hist.erase(hist.begin()); + } + hist.push_back(std::move(d_cur)); + } + return x; +} + static sd::Tensor sample_euler_cfg_pp(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas) { @@ -2739,6 +2821,8 @@ static sd::Tensor 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: diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index e5422052..9153b52b 100644 --- a/src/stable-diffusion.cpp +++ b/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) {