refactor: split and simplify sample_k_diffusion samplers (#1377)

This commit is contained in:
leejet 2026-03-31 00:32:14 +08:00 committed by GitHub
parent 83e8f6f0af
commit 1d6cb0f8c3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -2,6 +2,7 @@
#define __DENOISER_HPP__ #define __DENOISER_HPP__
#include <cmath> #include <cmath>
#include <utility>
#include "ggml_extend.hpp" #include "ggml_extend.hpp"
#include "gits_noise.inl" #include "gits_noise.inl"
@ -763,16 +764,33 @@ struct Flux2FlowDenoiser : public FluxFlowDenoiser {
typedef std::function<sd::Tensor<float>(const sd::Tensor<float>&, float, int)> denoise_cb_t; typedef std::function<sd::Tensor<float>(const sd::Tensor<float>&, float, int)> denoise_cb_t;
// k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t static std::pair<float, float> get_ancestral_step(float sigma_from,
static sd::Tensor<float> sample_k_diffusion(sample_method_t method, float sigma_to,
denoise_cb_t model, float eta = 1.0f) {
float sigma_up = 0.0f;
float sigma_down = sigma_to;
if (eta <= 0.0f) {
return {sigma_down, sigma_up};
}
float sigma_from_sq = sigma_from * sigma_from;
float sigma_to_sq = sigma_to * sigma_to;
if (sigma_from_sq > 0.0f) {
float term = sigma_to_sq * (sigma_from_sq - sigma_to_sq) / sigma_from_sq;
sigma_up = std::min(sigma_to, eta * std::sqrt(std::max(term, 0.0f)));
}
float sigma_down_sq = sigma_to_sq - sigma_up * sigma_up;
sigma_down = sigma_down_sq > 0.0f ? std::sqrt(sigma_down_sq) : 0.0f;
return {sigma_down, sigma_up};
}
static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
sd::Tensor<float> x, sd::Tensor<float> x,
std::vector<float> sigmas, const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng, std::shared_ptr<RNG> rng) {
float eta) { int steps = static_cast<int>(sigmas.size()) - 1;
size_t steps = sigmas.size() - 1;
switch (method) {
case EULER_A_SAMPLE_METHOD: {
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
float sigma = sigmas[i]; float sigma = sigmas[i];
auto denoised_opt = model(x, sigma, i + 1); auto denoised_opt = model(x, sigma, i + 1);
@ -781,18 +799,19 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
} }
sd::Tensor<float> denoised = std::move(denoised_opt); sd::Tensor<float> denoised = std::move(denoised_opt);
sd::Tensor<float> d = (x - denoised) / sigma; sd::Tensor<float> d = (x - denoised) / sigma;
float sigma_up = std::min(sigmas[i + 1], auto [sigma_down, sigma_up] = get_ancestral_step(sigmas[i], sigmas[i + 1]);
std::sqrt(sigmas[i + 1] * sigmas[i + 1] * (sigmas[i] * sigmas[i] - sigmas[i + 1] * sigmas[i + 1]) / (sigmas[i] * sigmas[i]))); x += d * (sigma_down - sigmas[i]);
float sigma_down = std::sqrt(sigmas[i + 1] * sigmas[i + 1] - sigma_up * sigma_up);
float dt = sigma_down - sigmas[i];
x += d * dt;
if (sigmas[i + 1] > 0) { if (sigmas[i + 1] > 0) {
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up; x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
} }
} }
return x; return x;
} }
case EULER_SAMPLE_METHOD: {
static sd::Tensor<float> sample_euler(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
float sigma = sigmas[i]; float sigma = sigmas[i];
auto denoised_opt = model(x, sigma, i + 1); auto denoised_opt = model(x, sigma, i + 1);
@ -801,12 +820,15 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
} }
sd::Tensor<float> denoised = std::move(denoised_opt); sd::Tensor<float> denoised = std::move(denoised_opt);
sd::Tensor<float> d = (x - denoised) / sigma; sd::Tensor<float> d = (x - denoised) / sigma;
float dt = sigmas[i + 1] - sigma; x += d * (sigmas[i + 1] - sigma);
x += d * dt;
} }
return x; return x;
} }
case HEUN_SAMPLE_METHOD: {
static sd::Tensor<float> sample_heun(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], -(i + 1)); auto denoised_opt = model(x, sigmas[i], -(i + 1));
if (denoised_opt.empty()) { if (denoised_opt.empty()) {
@ -829,8 +851,12 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
} }
} }
return x; return x;
} }
case DPM2_SAMPLE_METHOD: {
static sd::Tensor<float> sample_dpm2(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], -(i + 1)); auto denoised_opt = model(x, sigmas[i], -(i + 1));
if (denoised_opt.empty()) { if (denoised_opt.empty()) {
@ -839,8 +865,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
sd::Tensor<float> denoised = std::move(denoised_opt); sd::Tensor<float> denoised = std::move(denoised_opt);
sd::Tensor<float> d = (x - denoised) / sigmas[i]; sd::Tensor<float> d = (x - denoised) / sigmas[i];
if (sigmas[i + 1] == 0) { if (sigmas[i + 1] == 0) {
float dt = sigmas[i + 1] - sigmas[i]; x += d * (sigmas[i + 1] - sigmas[i]);
x += d * dt;
} else { } else {
float sigma_mid = exp(0.5f * (log(sigmas[i]) + log(sigmas[i + 1]))); float sigma_mid = exp(0.5f * (log(sigmas[i]) + log(sigmas[i + 1])));
float dt_1 = sigma_mid - sigmas[i]; float dt_1 = sigma_mid - sigmas[i];
@ -855,19 +880,23 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
} }
} }
return x; return x;
} }
case DPMPP2S_A_SAMPLE_METHOD: {
static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng) {
auto t_fn = [](float sigma) -> float { return -log(sigma); };
auto sigma_fn = [](float t) -> float { return exp(-t); };
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], -(i + 1)); auto denoised_opt = model(x, sigmas[i], -(i + 1));
if (denoised_opt.empty()) { if (denoised_opt.empty()) {
return {}; return {};
} }
sd::Tensor<float> denoised = std::move(denoised_opt); sd::Tensor<float> denoised = std::move(denoised_opt);
float sigma_up = std::min(sigmas[i + 1], auto [sigma_down, sigma_up] = get_ancestral_step(sigmas[i], sigmas[i + 1]);
std::sqrt(sigmas[i + 1] * sigmas[i + 1] * (sigmas[i] * sigmas[i] - sigmas[i + 1] * sigmas[i + 1]) / (sigmas[i] * sigmas[i])));
float sigma_down = std::sqrt(sigmas[i + 1] * sigmas[i + 1] - sigma_up * sigma_up);
auto t_fn = [](float sigma) -> float { return -log(sigma); };
auto sigma_fn = [](float t) -> float { return exp(-t); };
if (sigma_down == 0) { if (sigma_down == 0) {
x = denoised; x = denoised;
@ -882,7 +911,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
return {}; return {};
} }
sd::Tensor<float> denoised2 = std::move(denoised2_opt); sd::Tensor<float> denoised2 = std::move(denoised2_opt);
x = (sigma_fn(t_next) / sigma_fn(t)) * (x) - (exp(-h) - 1) * denoised2; x = (sigma_fn(t_next) / sigma_fn(t)) * x - (exp(-h) - 1) * denoised2;
} }
if (sigmas[i + 1] > 0) { if (sigmas[i + 1] > 0) {
@ -890,10 +919,15 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
} }
} }
return x; return x;
} }
case DPMPP2M_SAMPLE_METHOD: {
static sd::Tensor<float> sample_dpmpp_2m(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
sd::Tensor<float> old_denoised = x; sd::Tensor<float> old_denoised = x;
auto t_fn = [](float sigma) -> float { return -log(sigma); }; auto t_fn = [](float sigma) -> float { return -log(sigma); };
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], i + 1); auto denoised_opt = model(x, sigmas[i], i + 1);
if (denoised_opt.empty()) { if (denoised_opt.empty()) {
@ -907,20 +941,25 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
float b = exp(-h) - 1.f; float b = exp(-h) - 1.f;
if (i == 0 || sigmas[i + 1] == 0) { if (i == 0 || sigmas[i + 1] == 0) {
x = a * (x)-b * denoised; x = a * x - b * denoised;
} else { } else {
float h_last = t - t_fn(sigmas[i - 1]); float h_last = t - t_fn(sigmas[i - 1]);
float r = h_last / h; float r = h_last / h;
sd::Tensor<float> denoised_d = (1.f + 1.f / (2.f * r)) * denoised - (1.f / (2.f * r)) * old_denoised; sd::Tensor<float> denoised_d = (1.f + 1.f / (2.f * r)) * denoised - (1.f / (2.f * r)) * old_denoised;
x = a * (x)-b * denoised_d; x = a * x - b * denoised_d;
} }
old_denoised = denoised; old_denoised = denoised;
} }
return x; return x;
} }
case DPMPP2Mv2_SAMPLE_METHOD: {
static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
sd::Tensor<float> old_denoised = x; sd::Tensor<float> old_denoised = x;
auto t_fn = [](float sigma) -> float { return -log(sigma); }; auto t_fn = [](float sigma) -> float { return -log(sigma); };
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], i + 1); auto denoised_opt = model(x, sigmas[i], i + 1);
if (denoised_opt.empty()) { if (denoised_opt.empty()) {
@ -931,9 +970,10 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
float t_next = t_fn(sigmas[i + 1]); float t_next = t_fn(sigmas[i + 1]);
float h = t_next - t; float h = t_next - t;
float a = sigmas[i + 1] / sigmas[i]; float a = sigmas[i + 1] / sigmas[i];
if (i == 0 || sigmas[i + 1] == 0) { if (i == 0 || sigmas[i + 1] == 0) {
float b = exp(-h) - 1.f; float b = exp(-h) - 1.f;
x = a * (x)-b * denoised; x = a * x - b * denoised;
} else { } else {
float h_last = t - t_fn(sigmas[i - 1]); float h_last = t - t_fn(sigmas[i - 1]);
float h_min = std::min(h_last, h); float h_min = std::min(h_last, h);
@ -942,30 +982,38 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
float h_d = (h_max + h_min) / 2.f; float h_d = (h_max + h_min) / 2.f;
float b = exp(-h_d) - 1.f; float b = exp(-h_d) - 1.f;
sd::Tensor<float> denoised_d = (1.f + 1.f / (2.f * r)) * denoised - (1.f / (2.f * r)) * old_denoised; sd::Tensor<float> denoised_d = (1.f + 1.f / (2.f * r)) * denoised - (1.f / (2.f * r)) * old_denoised;
x = a * (x)-b * denoised_d; x = a * x - b * denoised_d;
} }
old_denoised = denoised; old_denoised = denoised;
} }
return x; return x;
} }
case LCM_SAMPLE_METHOD: {
static sd::Tensor<float> sample_lcm(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng) {
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], i + 1); auto denoised_opt = model(x, sigmas[i], i + 1);
if (denoised_opt.empty()) { if (denoised_opt.empty()) {
return {}; return {};
} }
sd::Tensor<float> denoised = std::move(denoised_opt); x = std::move(denoised_opt);
x = denoised;
if (sigmas[i + 1] > 0) { if (sigmas[i + 1] > 0) {
x += sd::Tensor<float>::randn_like(x, rng) * sigmas[i + 1]; x += sd::Tensor<float>::randn_like(x, rng) * sigmas[i + 1];
} }
} }
return x; return x;
} }
case IPNDM_SAMPLE_METHOD: {
int max_order = 4; static sd::Tensor<float> sample_ipndm(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
const int max_order = 4;
std::vector<sd::Tensor<float>> hist = {}; std::vector<sd::Tensor<float>> hist = {};
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
float sigma = sigmas[i]; float sigma = sigmas[i];
float sigma_next = sigmas[i + 1]; float sigma_next = sigmas[i + 1];
@ -1001,10 +1049,15 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
hist.push_back(std::move(d_cur)); hist.push_back(std::move(d_cur));
} }
return x; return x;
} }
case IPNDM_V_SAMPLE_METHOD: {
int max_order = 4; static sd::Tensor<float> sample_ipndm_v(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
const int max_order = 4;
std::vector<sd::Tensor<float>> hist = {}; std::vector<sd::Tensor<float>> hist = {};
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
float sigma = sigmas[i]; float sigma = sigmas[i];
float t_next = sigmas[i + 1]; float t_next = sigmas[i + 1];
@ -1041,10 +1094,14 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
hist.push_back(std::move(d_cur)); hist.push_back(std::move(d_cur));
} }
return x; return x;
} }
case RES_MULTISTEP_SAMPLE_METHOD: {
sd::Tensor<float> old_denoised = x;
static sd::Tensor<float> sample_res_multistep(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 = x;
bool have_old_sigma = false; bool have_old_sigma = false;
float old_sigma_down = 0.0f; float old_sigma_down = 0.0f;
@ -1064,6 +1121,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
return (phi1_val - 1.0f) / t; return (phi1_val - 1.0f) / t;
}; };
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], i + 1); auto denoised_opt = model(x, sigmas[i], i + 1);
if (denoised_opt.empty()) { if (denoised_opt.empty()) {
@ -1073,22 +1131,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
float sigma_from = sigmas[i]; float sigma_from = sigmas[i];
float sigma_to = sigmas[i + 1]; float sigma_to = sigmas[i + 1];
float sigma_up = 0.0f; auto [sigma_down, sigma_up] = get_ancestral_step(sigma_from, sigma_to, eta);
float sigma_down = sigma_to;
if (eta > 0.0f) {
float sigma_from_sq = sigma_from * sigma_from;
float sigma_to_sq = sigma_to * sigma_to;
if (sigma_from_sq > 0.0f) {
float term = sigma_to_sq * (sigma_from_sq - sigma_to_sq) / sigma_from_sq;
if (term > 0.0f) {
sigma_up = eta * std::sqrt(term);
}
}
sigma_up = std::min(sigma_up, sigma_to);
float sigma_down_sq = sigma_to_sq - sigma_up * sigma_up;
sigma_down = sigma_down_sq > 0.0f ? std::sqrt(sigma_down_sq) : 0.0f;
}
if (sigma_down == 0.0f || !have_old_sigma) { if (sigma_down == 0.0f || !have_old_sigma) {
x += ((x - denoised) / sigma_from) * (sigma_down - sigma_from); x += ((x - denoised) / sigma_from) * (sigma_down - sigma_from);
@ -1112,7 +1155,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
b2 = 0.0f; b2 = 0.0f;
} }
x = sigma_fn(h) * (x) + h * (b1 * denoised + b2 * old_denoised); x = sigma_fn(h) * x + h * (b1 * denoised + b2 * old_denoised);
} }
if (sigmas[i + 1] > 0 && sigma_up > 0.0f) { if (sigmas[i + 1] > 0 && sigma_up > 0.0f) {
@ -1124,8 +1167,13 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
have_old_sigma = true; have_old_sigma = true;
} }
return x; return x;
} }
case RES_2S_SAMPLE_METHOD: {
static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
float eta) {
const float c2 = 0.5f; const float c2 = 0.5f;
auto t_fn = [](float sigma) -> float { return -logf(sigma); }; auto t_fn = [](float sigma) -> float { return -logf(sigma); };
auto phi1_fn = [](float t) -> float { auto phi1_fn = [](float t) -> float {
@ -1142,6 +1190,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
return (phi1_val - 1.0f) / t; return (phi1_val - 1.0f) / t;
}; };
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
float sigma_from = sigmas[i]; float sigma_from = sigmas[i];
float sigma_to = sigmas[i + 1]; float sigma_to = sigmas[i + 1];
@ -1152,21 +1201,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
} }
sd::Tensor<float> denoised = std::move(denoised_opt); sd::Tensor<float> denoised = std::move(denoised_opt);
float sigma_up = 0.0f; auto [sigma_down, sigma_up] = get_ancestral_step(sigma_from, sigma_to, eta);
float sigma_down = sigma_to;
if (eta > 0.0f) {
float sigma_from_sq = sigma_from * sigma_from;
float sigma_to_sq = sigma_to * sigma_to;
if (sigma_from_sq > 0.0f) {
float term = sigma_to_sq * (sigma_from_sq - sigma_to_sq) / sigma_from_sq;
if (term > 0.0f) {
sigma_up = eta * std::sqrt(term);
}
}
sigma_up = std::min(sigma_up, sigma_to);
float sigma_down_sq = sigma_to_sq - sigma_up * sigma_up;
sigma_down = sigma_down_sq > 0.0f ? std::sqrt(sigma_down_sq) : 0.0f;
}
sd::Tensor<float> x0 = x; sd::Tensor<float> x0 = x;
if (sigma_down == 0.0f || sigma_from == 0.0f) { if (sigma_down == 0.0f || sigma_from == 0.0f) {
@ -1200,8 +1235,13 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
} }
} }
return x; return x;
} }
case DDIM_TRAILING_SAMPLE_METHOD: {
static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
float eta) {
float beta_start = 0.00085f; float beta_start = 0.00085f;
float beta_end = 0.0120f; float beta_end = 0.0120f;
std::vector<double> alphas_cumprod(TIMESTEPS); std::vector<double> alphas_cumprod(TIMESTEPS);
@ -1218,9 +1258,10 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]); std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
} }
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
int timestep = static_cast<int>(roundf(TIMESTEPS - i * ((float)TIMESTEPS / steps))) - 1; int timestep = static_cast<int>(roundf(TIMESTEPS - i * ((float)TIMESTEPS / steps))) - 1;
int prev_timestep = timestep - TIMESTEPS / static_cast<int>(steps); int prev_timestep = timestep - TIMESTEPS / steps;
float sigma = static_cast<float>(compvis_sigmas[timestep]); float sigma = static_cast<float>(compvis_sigmas[timestep]);
if (i == 0) { if (i == 0) {
x *= std::sqrt(sigma * sigma + 1) / sigma; x *= std::sqrt(sigma * sigma + 1) / sigma;
@ -1256,8 +1297,13 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
} }
} }
return x; return x;
} }
case TCD_SAMPLE_METHOD: {
static sd::Tensor<float> sample_tcd(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
float eta) {
float beta_start = 0.00085f; float beta_start = 0.00085f;
float beta_end = 0.0120f; float beta_end = 0.0120f;
std::vector<double> alphas_cumprod(TIMESTEPS); std::vector<double> alphas_cumprod(TIMESTEPS);
@ -1273,7 +1319,9 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
compvis_sigmas[i] = compvis_sigmas[i] =
std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]); std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
} }
int original_steps = 50; int original_steps = 50;
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) { for (int i = 0; i < steps; i++) {
int timestep = TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor(i * ((float)original_steps / steps)); int timestep = TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor(i * ((float)original_steps / steps));
int prev_timestep = i >= steps - 1 ? 0 : TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor((i + 1) * ((float)original_steps / steps)); int prev_timestep = i >= steps - 1 ? 0 : TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor((i + 1) * ((float)original_steps / steps));
@ -1307,12 +1355,49 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
std::sqrt(beta_prod_s) * model_output; std::sqrt(beta_prod_s) * model_output;
if (eta > 0 && i != steps - 1) { if (eta > 0 && i != steps - 1) {
x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * (x) + x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * x +
std::sqrt(1.0f - alpha_prod_t_prev / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng); std::sqrt(1.0f - alpha_prod_t_prev / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
} }
} }
return x; return x;
} }
// k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t
static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
denoise_cb_t model,
sd::Tensor<float> x,
std::vector<float> sigmas,
std::shared_ptr<RNG> rng,
float eta) {
switch (method) {
case EULER_A_SAMPLE_METHOD:
return sample_euler_ancestral(model, std::move(x), sigmas, rng);
case EULER_SAMPLE_METHOD:
return sample_euler(model, std::move(x), sigmas);
case HEUN_SAMPLE_METHOD:
return sample_heun(model, std::move(x), sigmas);
case DPM2_SAMPLE_METHOD:
return sample_dpm2(model, std::move(x), sigmas);
case DPMPP2S_A_SAMPLE_METHOD:
return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, rng);
case DPMPP2M_SAMPLE_METHOD:
return sample_dpmpp_2m(model, std::move(x), sigmas);
case DPMPP2Mv2_SAMPLE_METHOD:
return sample_dpmpp_2m_v2(model, std::move(x), sigmas);
case LCM_SAMPLE_METHOD:
return sample_lcm(model, std::move(x), sigmas, rng);
case IPNDM_SAMPLE_METHOD:
return sample_ipndm(model, std::move(x), sigmas);
case IPNDM_V_SAMPLE_METHOD:
return sample_ipndm_v(model, std::move(x), sigmas);
case RES_MULTISTEP_SAMPLE_METHOD:
return sample_res_multistep(model, std::move(x), sigmas, rng, eta);
case RES_2S_SAMPLE_METHOD:
return sample_res_2s(model, std::move(x), sigmas, rng, eta);
case DDIM_TRAILING_SAMPLE_METHOD:
return sample_ddim_trailing(model, std::move(x), sigmas, rng, eta);
case TCD_SAMPLE_METHOD:
return sample_tcd(model, std::move(x), sigmas, rng, eta);
default: default:
return {}; return {};
} }