From 951043fca9aa1033d293ef8d11a01adc3c6a8155 Mon Sep 17 00:00:00 2001 From: Daniele <57776841+daniandtheweb@users.noreply.github.com> Date: Mon, 15 Dec 2025 18:46:07 +0100 Subject: [PATCH] feat: add KL Optimal scheduler --- denoiser.hpp | 39 ++++++++++++++++++++++++++++++++++++++ examples/cli/README.md | 2 +- examples/common/common.hpp | 4 ++-- examples/server/README.md | 4 ++-- stable-diffusion.cpp | 3 ++- stable-diffusion.h | 1 + 6 files changed, 47 insertions(+), 6 deletions(-) diff --git a/denoiser.hpp b/denoiser.hpp index 32f402786..fc5230d7b 100644 --- a/denoiser.hpp +++ b/denoiser.hpp @@ -347,6 +347,41 @@ struct SmoothStepScheduler : SigmaScheduler { } }; +// Implementation adapted from https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608 +struct KLOptimalScheduler : SigmaScheduler { + std::vector get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override { + std::vector sigmas; + + if (n == 0) { + return sigmas; + } + if (n == 1) { + sigmas.push_back(sigma_max); + sigmas.push_back(0.0f); + return sigmas; + } + + float alpha_min = std::atan(sigma_min); + float alpha_max = std::atan(sigma_max); + + for (uint32_t i = 0; i < n; ++i) { + // t goes from 0.0 to 1.0 + float t = static_cast(i) / static_cast(n-1); + + // Interpolate in the angle domain + float angle = t * alpha_min + (1.0f - t) * alpha_max; + + // Convert back to sigma + sigmas.push_back(std::tan(angle)); + } + + // Append the final zero to sigma + sigmas.push_back(0.0f); + + return sigmas; + } +}; + struct Denoiser { virtual float sigma_min() = 0; virtual float sigma_max() = 0; @@ -392,6 +427,10 @@ struct Denoiser { LOG_INFO("get_sigmas with SmoothStep scheduler"); scheduler = std::make_shared(); break; + case KL_OPTIMAL_SCHEDULER: + LOG_INFO("get_sigmas with KL Optimal scheduler"); + scheduler = std::make_shared(); + break; case LCM_SCHEDULER: LOG_INFO("get_sigmas with LCM scheduler"); scheduler = std::make_shared(); diff --git a/examples/cli/README.md b/examples/cli/README.md index 8531b2aed..da7734d31 100644 --- a/examples/cli/README.md +++ b/examples/cli/README.md @@ -120,7 +120,7 @@ Generation Options: tcd] (default: euler for Flux/SD3/Wan, euler_a otherwise) --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] default: euler for Flux/SD3/Wan, euler_a otherwise - --scheduler denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, lcm], + --scheduler denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm], default: discrete --sigmas custom sigma values for the sampler, comma-separated (e.g., "14.61,7.8,3.5,0.0"). --skip-layers layers to skip for SLG steps (default: [7,8,9]) diff --git a/examples/common/common.hpp b/examples/common/common.hpp index f3a561367..b81dab784 100644 --- a/examples/common/common.hpp +++ b/examples/common/common.hpp @@ -1409,7 +1409,7 @@ struct SDGenerationParams { on_high_noise_sample_method_arg}, {"", "--scheduler", - "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, lcm], default: discrete", + "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm], default: discrete", on_scheduler_arg}, {"", "--sigmas", @@ -1911,4 +1911,4 @@ uint8_t* load_image_from_memory(const char* image_bytes, int expected_height = 0, int expected_channel = 3) { return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel); -} \ No newline at end of file +} diff --git a/examples/server/README.md b/examples/server/README.md index a475856fa..ae1049680 100644 --- a/examples/server/README.md +++ b/examples/server/README.md @@ -114,11 +114,11 @@ Default Generation Options: tcd] (default: euler for Flux/SD3/Wan, euler_a otherwise) --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] default: euler for Flux/SD3/Wan, euler_a otherwise - --scheduler denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, lcm], + --scheduler denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm], default: discrete --sigmas custom sigma values for the sampler, comma-separated (e.g., "14.61,7.8,3.5,0.0"). --skip-layers layers to skip for SLG steps (default: [7,8,9]) --high-noise-skip-layers (high noise) layers to skip for SLG steps (default: [7,8,9]) -r, --ref-image reference image for Flux Kontext models (can be used multiple times) --easycache enable EasyCache for DiT models with optional "threshold,start_percent,end_percent" (default: 0.2,0.15,0.95) -``` \ No newline at end of file +``` diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index 44bd3ccac..d58c7dd26 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -2409,6 +2409,7 @@ const char* scheduler_to_str[] = { "sgm_uniform", "simple", "smoothstep", + "kl_optimal", "lcm", }; @@ -3885,4 +3886,4 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s LOG_INFO("generate_video completed in %.2fs", (t5 - t0) * 1.0f / 1000); return result_images; -} \ No newline at end of file +} diff --git a/stable-diffusion.h b/stable-diffusion.h index 9266ba437..110120e97 100644 --- a/stable-diffusion.h +++ b/stable-diffusion.h @@ -60,6 +60,7 @@ enum scheduler_t { SGM_UNIFORM_SCHEDULER, SIMPLE_SCHEDULER, SMOOTHSTEP_SCHEDULER, + KL_OPTIMAL_SCHEDULER, LCM_SCHEDULER, SCHEDULER_COUNT };