feat: add DPM++ 2M SDE sampler (#1742)

examples/common/common.cpp

@@ -1106,7 +1106,7 @@ ArgOptions SDGenerationParams::get_options() {
          &sample_params.guidance.slg.layer_end},
         {"",
          "--eta",
-         "noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
+         "noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a and dpm++2m_sde)",
          &sample_params.eta},
         {"",
          "--flow-shift",
@@ -1138,7 +1138,7 @@ ArgOptions SDGenerationParams::get_options() {
          &high_noise_sample_params.guidance.slg.layer_end},
         {"",
          "--high-noise-eta",
-         "(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
+         "(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a and dpm++2m_sde)",
          &high_noise_sample_params.eta},
         {"",
          "--strength",
@@ -1509,12 +1509,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, 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, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
          "(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, 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, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
          " default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_high_noise_sample_method_arg},
         {"",

examples/server/routes_sdapi.cpp

@@ -63,6 +63,8 @@ static enum sample_method_t get_sdapi_sample_method(std::string name) {
         {"ddim", DDIM_TRAILING_SAMPLE_METHOD},
         {"dpm++ 2m", DPMPP2M_SAMPLE_METHOD},
         {"k_dpmpp_2m", DPMPP2M_SAMPLE_METHOD},
+        {"dpm++ 2m sde", DPMPP2M_SDE_SAMPLE_METHOD},
+        {"k_dpmpp_2m_sde", DPMPP2M_SDE_SAMPLE_METHOD},
         {"res multistep", RES_MULTISTEP_SAMPLE_METHOD},
         {"k_res_multistep", RES_MULTISTEP_SAMPLE_METHOD},
         {"res 2s", RES_2S_SAMPLE_METHOD},

include/stable-diffusion.h

@@ -54,6 +54,7 @@ enum sample_method_t {
     EULER_CFG_PP_SAMPLE_METHOD,
     EULER_A_CFG_PP_SAMPLE_METHOD,
     EULER_GE_SAMPLE_METHOD,
+    DPMPP2M_SDE_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 

src/runtime/denoiser.hpp

@@ -1814,6 +1814,52 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
     return x;
 }
 
+// DPM-Solver++(2M) SDE, midpoint variant. Ref: Lu et al. arXiv:2211.01095;
+// k-diffusion sample_dpmpp_2m_sde.
+static sd::Tensor<float> sample_dpmpp_2m_sde(denoise_cb_t model,
+                                             sd::Tensor<float> x,
+                                             const std::vector<float>& sigmas,
+                                             std::shared_ptr<RNG> rng,
+                                             float eta) {
+    sd::Tensor<float> old_denoised;
+    bool have_old_denoised = false;
+    float h_last           = 0.f;
+
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        auto denoised_opt = model(x, sigmas[i], i + 1);
+        if (denoised_opt.pred.empty()) {
+            return {};
+        }
+        sd::Tensor<float> denoised = std::move(denoised_opt.pred);
+
+        if (sigmas[i + 1] == 0.f) {
+            x = denoised;
+        } else {
+            float t     = -std::log(sigmas[i]);
+            float s     = -std::log(sigmas[i + 1]);
+            float h     = s - t;
+            float eta_h = eta * h;
+            float a     = sigmas[i + 1] / sigmas[i] * std::exp(-eta_h);
+            float b     = -std::expm1(-h - eta_h);
+
+            x = a * x + b * denoised;
+
+            if (have_old_denoised) {
+                float r = h_last / h;
+                x += (0.5f * b / r) * (denoised - old_denoised);
+            }
+            if (eta > 0.f) {
+                x += sd::Tensor<float>::randn_like(x, rng) * (sigmas[i + 1] * std::sqrt(-std::expm1(-2.f * eta_h)));
+            }
+            h_last = h;
+        }
+        old_denoised      = denoised;
+        have_old_denoised = true;
+    }
+    return x;
+}
+
 using SamplerExtraArgs = KeyValueArgs;
 
 static sd::Tensor<float> sample_lcm(denoise_cb_t model,
@@ -2490,6 +2536,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
             return sample_res_2s(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case ER_SDE_SAMPLE_METHOD:
             return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
+        case DPMPP2M_SDE_SAMPLE_METHOD:
+            return sample_dpmpp_2m_sde(model, std::move(x), sigmas, rng, eta);
         case DDIM_TRAILING_SAMPLE_METHOD:
             // DDIM is equivalent to Euler Ancestral with the Simple scheduler
             return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);

src/stable-diffusion.cpp

@@ -2818,6 +2818,7 @@ const char* sample_method_to_str[] = {
     "euler_cfg_pp",
     "euler_a_cfg_pp",
     "euler_ge",
+    "dpm++2m_sde",
 };
 
 const char* sd_sample_method_name(enum sample_method_t sample_method) {
@@ -3519,6 +3520,7 @@ static float resolve_eta(sd_ctx_t* sd_ctx,
             case DPMPP2S_A_SAMPLE_METHOD:
             case ER_SDE_SAMPLE_METHOD:
             case EULER_A_CFG_PP_SAMPLE_METHOD:
+            case DPMPP2M_SDE_SAMPLE_METHOD:
                 return 1.0f;
             default:;
         }