Rename, add interpolation functions

Co-authored-by: liesen <liesen.dev@gmail.com>
This commit is contained in:
catboxanon 2025-05-01 18:14:27 -04:00
parent 1cf0b79fd1
commit ea613ed2e3

View File

@ -1,3 +1,4 @@
import math
import comfy.samplers import comfy.samplers
import comfy.sample import comfy.sample
from comfy.k_diffusion import sampling as k_diffusion_sampling from comfy.k_diffusion import sampling as k_diffusion_sampling
@ -249,45 +250,53 @@ class SetFirstSigma:
sigmas[0] = sigma sigmas[0] = sigma
return (sigmas, ) return (sigmas, )
class ExpandSigmas: class ExtendIntermediateSigmas:
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
return {"required": return {"required":
{"sigmas": ("SIGMAS", ), {"sigmas": ("SIGMAS", ),
"steps": ("INT", {"default": 2, "min": 1, "max": 100}), "steps": ("INT", {"default": 2, "min": 1, "max": 100}),
"sigma_min": ("FLOAT", {"default": 12.0, "min": 0.0, "max": 20000.0, "step": 0.01, "round": False}), "start_at_sigma": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20000.0, "step": 0.01, "round": False}),
"sigma_max": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20000.0, "step": 0.01, "round": False}), "end_at_sigma": ("FLOAT", {"default": 12.0, "min": 0.0, "max": 20000.0, "step": 0.01, "round": False}),
"spacing": (['linear', 'cosine', 'sine'],),
} }
} }
RETURN_TYPES = ("SIGMAS",) RETURN_TYPES = ("SIGMAS",)
CATEGORY = "sampling/custom_sampling/sigmas" CATEGORY = "sampling/custom_sampling/sigmas"
FUNCTION = "expand" FUNCTION = "extend"
def expand(self, sigmas, steps, sigma_max, sigma_min): def extend(self, sigmas, steps, start_at_sigma, end_at_sigma, spacing):
if sigma_max < 0: if start_at_sigma < 0:
sigma_max = float("inf") start_at_sigma = float("inf")
expanded_sigmas = [] interpolator = {
'linear': lambda x: x,
'cosine': lambda x: torch.sin(x*math.pi/2),
'sine': lambda x: 1 - torch.cos(x*math.pi/2)
}[spacing]
# linear space for our interpolation function
x = torch.linspace(0, 1, steps + 1, device=sigmas.device)[1:-1]
extended_sigmas = []
for i in range(len(sigmas) - 1): for i in range(len(sigmas) - 1):
sigma_current = sigmas[i] sigma_current = sigmas[i]
sigma_next = sigmas[i+1] sigma_next = sigmas[i+1]
expanded_sigmas.append(sigma_current) extended_sigmas.append(sigma_current)
if sigma_min <= sigma_current <= sigma_max: if end_at_sigma <= sigma_current <= start_at_sigma:
#XXX: Might be nice to interpolate the sigmas with different methods. interpolated_steps = interpolator(x) * (sigma_next - sigma_current) + sigma_current
# Would require writing this not as a loop. extended_sigmas.extend(interpolated_steps.tolist())
interpolated_steps = torch.linspace(sigma_current, sigma_next, steps + 1, device=sigmas.device)[1:-1]
expanded_sigmas.extend(interpolated_steps.tolist())
# Add the last sigma value # Add the last sigma value
if len(sigmas) > 0: if len(sigmas) > 0:
expanded_sigmas.append(sigmas[-1]) extended_sigmas.append(sigmas[-1])
expanded_sigmas = torch.FloatTensor(expanded_sigmas) extended_sigmas = torch.FloatTensor(extended_sigmas)
return (expanded_sigmas,) return (extended_sigmas,)
class KSamplerSelect: class KSamplerSelect:
@classmethod @classmethod
@ -775,7 +784,7 @@ NODE_CLASS_MAPPINGS = {
"SplitSigmasDenoise": SplitSigmasDenoise, "SplitSigmasDenoise": SplitSigmasDenoise,
"FlipSigmas": FlipSigmas, "FlipSigmas": FlipSigmas,
"SetFirstSigma": SetFirstSigma, "SetFirstSigma": SetFirstSigma,
"ExpandSigmas": ExpandSigmas, "ExtendIntermediateSigmas": ExtendIntermediateSigmas,
"CFGGuider": CFGGuider, "CFGGuider": CFGGuider,
"DualCFGGuider": DualCFGGuider, "DualCFGGuider": DualCFGGuider,