diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 4bddaaff3..4192aa64c 100644 --- a/examples/common/common.cpp +++ b/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}, {"", diff --git a/examples/server/routes_sdapi.cpp b/examples/server/routes_sdapi.cpp index b87b57105..7da456f18 100644 --- a/examples/server/routes_sdapi.cpp +++ b/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}, diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 5eacdddfa..c6092e57d 100644 --- a/include/stable-diffusion.h +++ b/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 }; diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index f6e4e7e2f..96c84269d 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -4,7 +4,10 @@ #include #include #include +#include #include +#include +#include #include #include @@ -1874,6 +1877,158 @@ static sd::Tensor 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& 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(); + rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL)); + w_at_tmax_ = sd::Tensor::randn(shape_, rng) * std::sqrt(static_cast(t_max_ - t_min_)); + } + + sd::Tensor 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(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 w(double t) { + auto it = cache_.find(t); + if (it != cache_.end()) { + return it->second; + } + sd::Tensor zero = sd::Tensor::zeros(shape_); + sd::Tensor out = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth); + cache_.emplace(t, out); + return out; + } + + sd::Tensor bridge(double a, + double c, + const sd::Tensor& w_a, + const sd::Tensor& w_c, + double t, + uint64_t node_seed, + int depth) { + if (depth <= 0 || c - a < 1e-9) { + float alpha = (c > a) ? static_cast((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(); + rng->manual_seed(node_seed); + auto z = sd::Tensor::randn(shape_, rng); + auto w_m = 0.5f * (w_a + w_c) + static_cast(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 shape_; + uint64_t root_seed_; + sd::Tensor w_at_tmax_; + std::map> 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 sample_dpmpp_2m_sde_bt(denoise_cb_t model, + sd::Tensor x, + const std::vector& sigmas, + std::shared_ptr rng, + float eta) { + double sigma_max = 0.0; + double sigma_min = std::numeric_limits::infinity(); + for (float s : sigmas) { + if (s > 0.0f) { + sigma_max = std::max(sigma_max, static_cast(s)); + sigma_min = std::min(sigma_min, static_cast(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 old_denoised; + bool have_old_denoised = false; + float h_last = 0.f; + + int steps = static_cast(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 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 sample_lcm(denoise_cb_t model, @@ -2552,6 +2707,8 @@ static sd::Tensor 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); diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index ae1224da9..8252640e1 100644 --- a/src/stable-diffusion.cpp +++ b/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:; }