From 02c7fdf34991c42012c61ac594fc6b23e849a7a4 Mon Sep 17 00:00:00 2001
From: abhay-codes07 <abhaysingh0293@gmail.com>
Date: Sun, 12 Jul 2026 17:56:25 +0530
Subject: [PATCH] Fix kl_optimal scheduler producing NaN at steps=1

kl_optimal_scheduler divided torch.arange(n) by (n - 1), which is a
division by zero when n == 1, making the first sigma NaN and corrupting
the whole generation. steps=1 is a valid KSampler/BasicScheduler input.
Clamp the divisor to at least 1 so a single step starts at sigma_max.
---
 comfy/samplers.py                             |  2 +-
 .../comfy_test/kl_optimal_scheduler_test.py   | 25 +++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)
 create mode 100644 tests-unit/comfy_test/kl_optimal_scheduler_test.py

diff --git a/comfy/samplers.py b/comfy/samplers.py
index 25c5a855f..a67b86f65 100755
--- a/comfy/samplers.py
+++ b/comfy/samplers.py
@@ -729,7 +729,7 @@ def linear_quadratic_schedule(model_sampling, steps, threshold_noise=0.025, line
 
 # Referenced from https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15608
 def kl_optimal_scheduler(n: int, sigma_min: float, sigma_max: float) -> torch.Tensor:
-    adj_idxs = torch.arange(n, dtype=torch.float).div_(n - 1)
+    adj_idxs = torch.arange(n, dtype=torch.float).div_(max(n - 1, 1))
     sigmas = adj_idxs.new_zeros(n + 1)
     sigmas[:-1] = (adj_idxs * math.atan(sigma_min) + (1 - adj_idxs) * math.atan(sigma_max)).tan_()
     return sigmas
diff --git a/tests-unit/comfy_test/kl_optimal_scheduler_test.py b/tests-unit/comfy_test/kl_optimal_scheduler_test.py
new file mode 100644
index 000000000..0035bcf0b
--- /dev/null
+++ b/tests-unit/comfy_test/kl_optimal_scheduler_test.py
@@ -0,0 +1,25 @@
+import torch
+from comfy.cli_args import args as cli_args
+
+if not torch.cuda.is_available():
+    cli_args.cpu = True
+
+from comfy.samplers import kl_optimal_scheduler
+
+
+class TestKLOptimalScheduler:
+    def test_single_step_is_not_nan(self):
+        # steps=1 is a valid KSampler input; div by (n - 1) made the first sigma NaN.
+        sigmas = kl_optimal_scheduler(1, 0.0291675, 14.614642)
+        assert not torch.isnan(sigmas).any()
+        assert sigmas.shape == (2,)
+        # the single step should start at sigma_max and end at 0
+        assert torch.isclose(sigmas[0], torch.tensor(14.614642), atol=1e-3)
+        assert sigmas[-1] == 0
+
+    def test_multi_step_unchanged(self):
+        sigmas = kl_optimal_scheduler(20, 0.0291675, 14.614642)
+        assert not torch.isnan(sigmas).any()
+        assert sigmas.shape == (21,)
+        assert torch.isclose(sigmas[0], torch.tensor(14.614642), atol=1e-3)
+        assert sigmas[-1] == 0