From f287b5aafbc8db2aff1752ef23550c4a83ce1b21 Mon Sep 17 00:00:00 2001
From: Daniele <57776841+daniandtheweb@users.noreply.github.com>
Date: Tue, 23 Dec 2025 18:13:28 +0100
Subject: [PATCH 1/2] fix: avoid issues when sigma_min is close to 0

---
 denoiser.hpp | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/denoiser.hpp b/denoiser.hpp
index 7a8242e7d..2df662a4d 100644
--- a/denoiser.hpp
+++ b/denoiser.hpp
@@ -347,7 +347,11 @@ struct SmoothStepScheduler : SigmaScheduler {
     }
 };
 
-// Implementation adapted from https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608
+/*
+* KL Optimal:
+*     Original work from https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608.
+*     Implemented using https://github.com/comfyanonymous/ComfyUI/pull/6206 as a reference.
+*/
 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;
@@ -355,27 +359,25 @@ struct KLOptimalScheduler : SigmaScheduler {
         if (n == 0) {
             return sigmas;
         }
+
         if (n == 1) {
             sigmas.push_back(sigma_max);
             sigmas.push_back(0.0f);
             return sigmas;
         }
 
+        sigmas.reserve(n + 1);
+
         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 t = static_cast<float>(i) / static_cast<float>(n - 1);
             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;
@@ -459,6 +461,10 @@ struct CompVisDenoiser : public Denoiser {
     }
 
     float sigma_to_t(float sigma) override {
+        if (sigma <= 1e-6f) {
+            return (float)(TIMESTEPS - 1);
+        }
+
         float log_sigma = std::log(sigma);
         std::vector<float> dists;
         dists.reserve(TIMESTEPS);
@@ -734,8 +740,12 @@ static bool sample_k_diffusion(sample_method_t method,
                     float* vec_x        = (float*)x->data;
                     float* vec_denoised = (float*)denoised->data;
 
-                    for (int i = 0; i < ggml_nelements(d); i++) {
-                        vec_d[i] = (vec_x[i] - vec_denoised[i]) / sigma;
+                    if (sigma < 1e-6f) {
+                        ggml_set_f32(d, 0.0f);
+                    } else {
+                        for (int i = 0; i < ggml_nelements(d); i++) {
+                            vec_d[i] = (vec_x[i] - vec_denoised[i]) / sigma;
+                        }
                     }
                 }
 

From f9022057deea1332c138b041ae031ac08fa3018d Mon Sep 17 00:00:00 2001
From: Daniele <57776841+daniandtheweb@users.noreply.github.com>
Date: Thu, 25 Dec 2025 20:04:06 +0100
Subject: [PATCH 2/2] fix: clamp sigma_min in Karras and KL Optimal to avoid
 NaNs

---
 denoiser.hpp | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/denoiser.hpp b/denoiser.hpp
index 2df662a4d..00bad2ec7 100644
--- a/denoiser.hpp
+++ b/denoiser.hpp
@@ -276,6 +276,10 @@ struct KarrasScheduler : SigmaScheduler {
         // but does anybody ever bother to touch them?
         float rho = 7.f;
 
+        if (sigma_min <= 1e-6f) {
+            sigma_min = 1e-6f;
+        }
+
         std::vector<float> result(n + 1);
 
         float min_inv_rho = pow(sigma_min, (1.f / rho));
@@ -366,6 +370,10 @@ struct KLOptimalScheduler : SigmaScheduler {
             return sigmas;
         }
 
+        if (sigma_min <= 1e-6f) {
+            sigma_min = 1e-6f;
+        }
+
         sigmas.reserve(n + 1);
 
         float alpha_min = std::atan(sigma_min);
@@ -461,10 +469,6 @@ struct CompVisDenoiser : public Denoiser {
     }
 
     float sigma_to_t(float sigma) override {
-        if (sigma <= 1e-6f) {
-            return (float)(TIMESTEPS - 1);
-        }
-
         float log_sigma = std::log(sigma);
         std::vector<float> dists;
         dists.reserve(TIMESTEPS);
@@ -740,12 +744,8 @@ static bool sample_k_diffusion(sample_method_t method,
                     float* vec_x        = (float*)x->data;
                     float* vec_denoised = (float*)denoised->data;
 
-                    if (sigma < 1e-6f) {
-                        ggml_set_f32(d, 0.0f);
-                    } else {
-                        for (int i = 0; i < ggml_nelements(d); i++) {
-                            vec_d[i] = (vec_x[i] - vec_denoised[i]) / sigma;
-                        }
+                    for (int i = 0; i < ggml_nelements(d); i++) {
+                        vec_d[i] = (vec_x[i] - vec_denoised[i]) / sigma;
                     }
                 }