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

This commit is contained in:
fszontagh
2026-07-05 10:13:25 +02:00
committed by GitHub
parent 45714b1069
commit 38a51f8105
5 changed files with 57 additions and 4 deletions

View File

@@ -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},
{"",

View File

@@ -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},

View File

@@ -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
};

View File

@@ -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);

View File

@@ -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:;
}