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.
This commit is contained in:
abhay-codes07 2026-07-12 18:01:43 +05:30
parent 8b099de36a
commit 8ed7447f6a
No known key found for this signature in database
2 changed files with 24 additions and 1 deletions

View File

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

View File

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