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

@@ -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

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