feat: add KL Optimal scheduler (#1098)

Author: daniandtheweb

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<float> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override {
+        std::vector<float> 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<float>(i) / static_cast<float>(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<SmoothStepScheduler>();
                 break;
+            case KL_OPTIMAL_SCHEDULER:
+                LOG_INFO("get_sigmas with KL Optimal scheduler");
+                scheduler = std::make_shared<KLOptimalScheduler>();
+                break;
             case LCM_SCHEDULER:
                 LOG_INFO("get_sigmas with LCM scheduler");
                 scheduler = std::make_shared<LCMScheduler>();

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])

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);
-}
+}

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)
-```
+```

stable-diffusion.cpp

@@ -2412,6 +2412,7 @@ const char* scheduler_to_str[] = {
     "sgm_uniform",
     "simple",
     "smoothstep",
+    "kl_optimal",
     "lcm",
 };
 
@@ -3888,4 +3889,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;
-}
+}

stable-diffusion.h

@@ -60,6 +60,7 @@ enum scheduler_t {
     SGM_UNIFORM_SCHEDULER,
     SIMPLE_SCHEDULER,
     SMOOTHSTEP_SCHEDULER,
+    KL_OPTIMAL_SCHEDULER,
     LCM_SCHEDULER,
     SCHEDULER_COUNT
 };