fix: fix kl_optimal when sigma_min is 0

Author: daniandtheweb

denoiser.hpp

@@ -347,36 +347,60 @@ struct SmoothStepScheduler : SigmaScheduler {
     }
 };
 
-// Implementation adapted from https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608
+// KL Optimal: 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);
+        if (sigma_min <= 1e-6f) {
+            /* sigma_min = 0:
+             * implemented using
+             * https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608
+             * as a reference.
+             */
+            
+            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);
+            for (uint32_t i = 0; i <= n; ++i) {
 
-            // Interpolate in the angle domain
-            float angle = t * alpha_min + (1.0f - t) * alpha_max;
+                float t = static_cast<float>(i) / static_cast<float>(n);
 
-            // Convert back to sigma
-            sigmas.push_back(std::tan(angle));
-        }
+                float angle = t * alpha_min + (1.0f - t) * alpha_max;
 
-        // Append the final zero to sigma
-        sigmas.push_back(0.0f);
+                sigmas.push_back(std::tan(angle));
+            }
+            
+        } else {
+            /* sigma_min != 0:
+             * implemented using
+             * https://github.com/comfyanonymous/ComfyUI/pull/6206
+             * as a reference.
+             */
+             
+            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) {
+
+                float t = static_cast<float>(i) / static_cast<float>(n - 1);
+
+                float angle = t * alpha_min + (1.0f - t) * alpha_max;
+                sigmas.push_back(std::tan(angle));
+            }
+
+            sigmas.push_back(0.0f);
+        }
 
         return sigmas;
     }