From 8ed7447f6a2b1b25184a528c86b105750684178e Mon Sep 17 00:00:00 2001 From: abhay-codes07 Date: Sun, 12 Jul 2026 18:01:43 +0530 Subject: [PATCH] Fix SplitSigmasDenoise returning empty high_sigmas at low denoise When round(steps * denoise) is 0 (e.g. a small denoise), sigmas[:-(total_steps)] becomes sigmas[:0] because -0 == 0, collapsing high_sigmas to an empty tensor instead of the full schedule. Slice with len(sigmas) - total_steps so the zero case keeps the whole schedule; the behavior is identical for total_steps > 0. --- comfy_extras/nodes_custom_sampler.py | 2 +- .../split_sigmas_denoise_test.py | 23 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) create mode 100644 tests-unit/comfy_extras_test/split_sigmas_denoise_test.py diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 56ef5f526..809cfc6ee 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -245,7 +245,7 @@ class SplitSigmasDenoise(io.ComfyNode): def execute(cls, sigmas, denoise) -> io.NodeOutput: steps = max(sigmas.shape[-1] - 1, 0) total_steps = round(steps * denoise) - sigmas1 = sigmas[:-(total_steps)] + sigmas1 = sigmas[:len(sigmas) - total_steps] sigmas2 = sigmas[-(total_steps + 1):] return io.NodeOutput(sigmas1, sigmas2) diff --git a/tests-unit/comfy_extras_test/split_sigmas_denoise_test.py b/tests-unit/comfy_extras_test/split_sigmas_denoise_test.py new file mode 100644 index 000000000..f5f8b68d8 --- /dev/null +++ b/tests-unit/comfy_extras_test/split_sigmas_denoise_test.py @@ -0,0 +1,23 @@ +import torch +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + +from comfy_extras.nodes_custom_sampler import SplitSigmasDenoise + + +class TestSplitSigmasDenoise: + def test_low_denoise_keeps_full_high_sigmas(self): + # denoise small enough that round(steps * denoise) == 0. sigmas[:-0] == sigmas[:0] + # collapsed high_sigmas to empty; it should stay the full schedule. + sigmas = torch.linspace(10, 0, 21) # 20 steps + high, low = SplitSigmasDenoise.execute(sigmas, 0.02) + assert high.shape[-1] == 21 + assert low.shape[-1] == 1 + + def test_normal_denoise_split(self): + sigmas = torch.linspace(10, 0, 21) + high, low = SplitSigmasDenoise.execute(sigmas, 0.5) # total_steps = 10 + assert high.shape[-1] == 11 + assert low.shape[-1] == 11