refactor: extract model loader initialization (#1844 )

feat: add linear multi-step sampling method (#1843 )
2026-09-22 22:17:53 -05:00 · 2026-08-02 17:24:10 +08:00 · 2026-08-02 16:19:47 +08:00
4 changed files with 156 additions and 60 deletions
@@ -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},
        {"",
@@ -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
 };

@@ -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:
@@ -143,6 +143,7 @@ const char* sampling_methods_str[] = {
    "Euler GE",
    "DPM++ (2M) SDE",
    "DPM++ (2M) SDE BT",
+    "LMS",
 };

 /*================================================== Helper Functions ================================================*/
@@ -695,45 +696,11 @@ public:
        LOG_DEBUG("loaded alphas_cumprod from model file");
    }

-    bool init(const sd_ctx_params_t* sd_ctx_params) {
-        n_threads           = sd_ctx_params->n_threads;
-        enable_mmap         = sd_ctx_params->enable_mmap;
-        stream_layers       = sd_ctx_params->stream_layers;
-        eager_load          = sd_ctx_params->eager_load;
-        backend_spec        = SAFE_STR(sd_ctx_params->backend);
-        params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
-        split_mode_spec     = SAFE_STR(sd_ctx_params->split_mode);
-        auto_fit_enabled    = sd_ctx_params->auto_fit;
-        max_vram_assignment.reset(0.f);
-        {
-            std::string error;
-            if (!max_vram_assignment.parse(SAFE_STR(sd_ctx_params->max_vram), &error)) {
-                LOG_ERROR("%s", error.c_str());
-                return false;
-            }
-        }
-
-        std::string rpc_servers_spec = SAFE_STR(sd_ctx_params->rpc_servers);
-        add_rpc_devices(rpc_servers_spec);
-
-        bool use_tae         = false;
-        bool use_audio_vae   = false;
-        bool use_control_net = false;
-
-        rng = get_rng(sd_ctx_params->rng_type);
-        if (sd_ctx_params->sampler_rng_type != RNG_TYPE_COUNT && sd_ctx_params->sampler_rng_type != sd_ctx_params->rng_type) {
-            sampler_rng = get_rng(sd_ctx_params->sampler_rng_type);
-        } else {
-            sampler_rng = rng;
-        }
-
-        ggml_log_set(ggml_log_callback_default, nullptr);
-
-        model_manager = std::make_shared<ModelManager>();
-        model_manager->set_n_threads(n_threads);
-        model_manager->set_enable_mmap(enable_mmap);
-        ModelLoader& model_loader = model_manager->loader();
-
+    bool init_model_loader(ModelLoader& model_loader,
+                           const sd_ctx_params_t* sd_ctx_params,
+                           bool& use_tae,
+                           bool& use_audio_vae,
+                           bool& use_control_net) {
        if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) {
            LOG_INFO("loading model from '%s'", sd_ctx_params->model_path);
            if (!model_loader.init_from_file(sd_ctx_params->model_path)) {
@@ -873,24 +840,69 @@ public:

        model_loader.convert_tensors_name();

-        version = model_loader.get_sd_version();
-        if (version == VERSION_COUNT) {
-            LOG_ERROR("get sd version from file failed: '%s'", SAFE_STR(sd_ctx_params->model_path));
-            return false;
-        }
-
-        auto& tensor_storage_map = model_loader.get_tensor_storage_map();
-
-        LOG_INFO("Version: %s ", model_version_to_str[version]);
        ggml_type wtype               = sd_type_to_ggml_type(sd_ctx_params->wtype);
        std::string tensor_type_rules = SAFE_STR(sd_ctx_params->tensor_type_rules);
        if (wtype != GGML_TYPE_COUNT || tensor_type_rules.size() > 0) {
            model_loader.set_wtype_override(wtype, tensor_type_rules);
        }

+        return true;
+    }
+
+    bool init(const sd_ctx_params_t* sd_ctx_params) {
+        n_threads           = sd_ctx_params->n_threads;
+        enable_mmap         = sd_ctx_params->enable_mmap;
+        stream_layers       = sd_ctx_params->stream_layers;
+        eager_load          = sd_ctx_params->eager_load;
+        backend_spec        = SAFE_STR(sd_ctx_params->backend);
+        params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
+        split_mode_spec     = SAFE_STR(sd_ctx_params->split_mode);
+        auto_fit_enabled    = sd_ctx_params->auto_fit;
+        max_vram_assignment.reset(0.f);
+        {
+            std::string error;
+            if (!max_vram_assignment.parse(SAFE_STR(sd_ctx_params->max_vram), &error)) {
+                LOG_ERROR("%s", error.c_str());
+                return false;
+            }
+        }
+
+        std::string rpc_servers_spec = SAFE_STR(sd_ctx_params->rpc_servers);
+        add_rpc_devices(rpc_servers_spec);
+
+        bool use_tae         = false;
+        bool use_audio_vae   = false;
+        bool use_control_net = false;
+
+        rng = get_rng(sd_ctx_params->rng_type);
+        if (sd_ctx_params->sampler_rng_type != RNG_TYPE_COUNT && sd_ctx_params->sampler_rng_type != sd_ctx_params->rng_type) {
+            sampler_rng = get_rng(sd_ctx_params->sampler_rng_type);
+        } else {
+            sampler_rng = rng;
+        }
+
+        ggml_log_set(ggml_log_callback_default, nullptr);
+
+        model_manager = std::make_shared<ModelManager>();
+        model_manager->set_n_threads(n_threads);
+        model_manager->set_enable_mmap(enable_mmap);
+        ModelLoader& model_loader = model_manager->loader();
+
+        if (!init_model_loader(model_loader, sd_ctx_params, use_tae, use_audio_vae, use_control_net)) {
+            return false;
+        }
+
+        version = model_loader.get_sd_version();
+        if (version == VERSION_COUNT) {
+            LOG_ERROR("get sd version from file failed: '%s'", SAFE_STR(sd_ctx_params->model_path));
+            return false;
+        } else {
+            LOG_INFO("Version: %s ", model_version_to_str[version]);
+        }
+
        if (auto_fit_enabled) {
            if (!sd::backend_fit::derive_backend_specs(model_loader,
-                                                       wtype,
+                                                       sd_type_to_ggml_type(sd_ctx_params->wtype),
                                                       max_vram_assignment,
                                                       backend_spec,
                                                       params_backend_spec)) {
@@ -945,14 +957,10 @@ public:

        if (sd_ctx_params->lora_apply_mode == LORA_APPLY_AUTO) {
            bool have_quantized_weight = false;
-            if (wtype != GGML_TYPE_COUNT && ggml_is_quantized(wtype)) {
-                have_quantized_weight = true;
-            } else {
-                for (const auto& [type, _] : wtype_stat) {
-                    if (ggml_is_quantized(type)) {
-                        have_quantized_weight = true;
-                        break;
-                    }
+            for (const auto& [type, _] : wtype_stat) {
+                if (ggml_is_quantized(type)) {
+                    have_quantized_weight = true;
+                    break;
                }
            }
            // Avoid full-model LoRA merge buffers on constrained setups.
@@ -996,6 +1004,8 @@ public:
            use_tae          = true;
        }

+        auto& tensor_storage_map = model_loader.get_tensor_storage_map();
+
        {
            if (!ensure_backend_pair(SDBackendModule::TE) ||
                !ensure_backend_pair(SDBackendModule::DIFFUSION)) {
@@ -3208,6 +3218,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) {
Author	SHA1	Message	Date
leejet	db99efdd6d	refactor: extract model loader initialization (#1844 )	2026-08-02 17:24:10 +08:00
vmobilis	eb7f35ca49	feat: add linear multi-step sampling method (#1843 )	2026-08-02 16:19:47 +08:00