feat: add DPM++ 2M SDE (Brownian tree) sampler

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, dpm++2s_a and dpm++2m_sde)",
+         "noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a, dpm++2m_sde and dpm++2m_sde_bt)",
          &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, dpm++2s_a and dpm++2m_sde)",
+         "(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a, dpm++2m_sde and dpm++2m_sde_bt)",
          &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, dpm++2m_sde, 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]"
          "(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, 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]"
          " default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_high_noise_sample_method_arg},
         {"",

examples/server/routes_sdapi.cpp

@@ -65,6 +65,8 @@ static enum sample_method_t get_sdapi_sample_method(std::string name) {
         {"k_dpmpp_2m", DPMPP2M_SAMPLE_METHOD},
         {"dpm++ 2m sde", DPMPP2M_SDE_SAMPLE_METHOD},
         {"k_dpmpp_2m_sde", DPMPP2M_SDE_SAMPLE_METHOD},
+        {"dpm++ 2m sde gpu", DPMPP2M_SDE_BT_SAMPLE_METHOD},
+        {"k_dpmpp_2m_sde_gpu", DPMPP2M_SDE_BT_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

@@ -55,6 +55,7 @@ enum sample_method_t {
     EULER_A_CFG_PP_SAMPLE_METHOD,
     EULER_GE_SAMPLE_METHOD,
     DPMPP2M_SDE_SAMPLE_METHOD,
+    DPMPP2M_SDE_BT_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 

src/runtime/denoiser.hpp

@@ -4,7 +4,10 @@
 #include <algorithm>
 #include <cctype>
 #include <cmath>
+#include <cstring>
 #include <functional>
+#include <limits>
+#include <map>
 #include <string>
 #include <utility>
 
@@ -1874,6 +1877,158 @@ static sd::Tensor<float> sample_dpmpp_2m_sde(denoise_cb_t model,
     return x;
 }
 
+// Seeded Brownian tree providing deterministic, step-count-stable Gaussian
+// increments for stochastic samplers. Constructed once per generation; each
+// call returns unit-variance noise for interval [sigma_a, sigma_b].
+// Reference: torchsde BrownianTree; k-diffusion BatchedBrownianTree.
+class BrownianTreeNoiseSampler {
+public:
+    BrownianTreeNoiseSampler(const sd::Tensor<float>& x_template,
+                             double sigma_min,
+                             double sigma_max,
+                             uint64_t seed)
+        : t_min_(sigma_min),
+          t_max_(sigma_max),
+          shape_(x_template.shape()),
+          root_seed_(mix64(seed, 0x9E3779B97F4A7C15ULL)) {
+        auto rng = std::make_shared<STDDefaultRNG>();
+        rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL));
+        w_at_tmax_ = sd::Tensor<float>::randn(shape_, rng) * std::sqrt(static_cast<float>(t_max_ - t_min_));
+    }
+
+    sd::Tensor<float> operator()(double sigma_a, double sigma_b) {
+        double a          = clamp(std::min(sigma_a, sigma_b));
+        double b          = clamp(std::max(sigma_a, sigma_b));
+        auto dW           = w(b) - w(a);
+        float span        = static_cast<float>(std::max(std::abs(sigma_b - sigma_a), 1e-12));
+        return dW * (1.0f / std::sqrt(span));
+    }
+
+private:
+    static constexpr int kMaxDepth = 24;
+
+    static uint64_t mix64(uint64_t v, uint64_t salt) {
+        uint64_t z = v + salt;
+        z          = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
+        z          = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
+        return z ^ (z >> 31);
+    }
+
+    double clamp(double t) const {
+        return std::min(std::max(t, t_min_), t_max_);
+    }
+
+    sd::Tensor<float> w(double t) {
+        auto it = cache_.find(t);
+        if (it != cache_.end()) {
+            return it->second;
+        }
+        sd::Tensor<float> zero = sd::Tensor<float>::zeros(shape_);
+        sd::Tensor<float> out  = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth);
+        cache_.emplace(t, out);
+        return out;
+    }
+
+    sd::Tensor<float> bridge(double a,
+                             double c,
+                             const sd::Tensor<float>& w_a,
+                             const sd::Tensor<float>& w_c,
+                             double t,
+                             uint64_t node_seed,
+                             int depth) {
+        if (depth <= 0 || c - a < 1e-9) {
+            float alpha = (c > a) ? static_cast<float>((t - a) / (c - a)) : 0.5f;
+            return (1.0f - alpha) * w_a + alpha * w_c;
+        }
+        double m       = 0.5 * (a + c);
+        double std_dev = std::sqrt((c - m) * (m - a) / (c - a));
+        auto rng       = std::make_shared<STDDefaultRNG>();
+        rng->manual_seed(node_seed);
+        auto z   = sd::Tensor<float>::randn(shape_, rng);
+        auto w_m = 0.5f * (w_a + w_c) + static_cast<float>(std_dev) * z;
+        if (t == m) {
+            return w_m;
+        }
+        if (t < m) {
+            return bridge(a, m, w_a, w_m, t, mix64(node_seed, 1), depth - 1);
+        }
+        return bridge(m, c, w_m, w_c, t, mix64(node_seed, 2), depth - 1);
+    }
+
+    double t_min_;
+    double t_max_;
+    std::vector<int64_t> shape_;
+    uint64_t root_seed_;
+    sd::Tensor<float> w_at_tmax_;
+    std::map<double, sd::Tensor<float>> cache_;
+};
+
+// DPM-Solver++(2M) SDE, midpoint variant, with step-count-stable Brownian-tree
+// noise. Same trajectory shape at any step count for a given seed. Aliased in
+// k-diffusion / ComfyUI as sample_dpmpp_2m_sde_gpu.
+// Ref: Lu et al. arXiv:2211.01095; torchsde BrownianTree.
+static sd::Tensor<float> sample_dpmpp_2m_sde_bt(denoise_cb_t model,
+                                                sd::Tensor<float> x,
+                                                const std::vector<float>& sigmas,
+                                                std::shared_ptr<RNG> rng,
+                                                float eta) {
+    double sigma_max = 0.0;
+    double sigma_min = std::numeric_limits<double>::infinity();
+    for (float s : sigmas) {
+        if (s > 0.0f) {
+            sigma_max = std::max(sigma_max, static_cast<double>(s));
+            sigma_min = std::min(sigma_min, static_cast<double>(s));
+        }
+    }
+    if (sigma_max <= sigma_min) {
+        return x;
+    }
+    uint64_t tree_seed = 0;
+    {
+        auto draw = rng->randn(2);
+        std::memcpy(&tree_seed, draw.data(), sizeof(tree_seed));
+    }
+    BrownianTreeNoiseSampler noise_sampler(x, sigma_min, sigma_max, tree_seed);
+
+    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 += noise_sampler(sigmas[i], sigmas[i + 1]) * (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,
@@ -2552,6 +2707,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t 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 DPMPP2M_SDE_BT_SAMPLE_METHOD:
+            return sample_dpmpp_2m_sde_bt(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

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