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