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:
parent
8b099de36a
commit
8ed7447f6a
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Reference in New Issue