Add RadianceShiftScheduler node and bump version from 1.0.4 to 1.0.5

This commit is contained in:
silveroxides
2026-01-21 15:35:12 +01:00
parent b12ed2be23
commit 6c273bfed0
2 changed files with 90 additions and 58 deletions
+89 -57
View File
@@ -3,63 +3,6 @@ import numpy
import scipy.stats
from comfy.samplers import SchedulerHandler, SCHEDULER_HANDLERS, SCHEDULER_NAMES
def power_shift_scheduler(model_sampling, steps, power=2.0, midpoint_shift=1.0, discard_penultimate=False):
total_timesteps = (len(model_sampling.sigmas) - 1)
x = numpy.linspace(0, 1, steps, endpoint=False)
x = x**midpoint_shift
ts_normalized = (1 - x**power)**power
ts = numpy.rint(ts_normalized * total_timesteps)
sigs = []
last_t = -1
for t in ts:
t_int = min(int(t), total_timesteps)
if t_int != last_t:
sigs.append(float(model_sampling.sigmas[t_int]))
last_t = t_int
sigs.append(0.0)
if discard_penultimate is True:
sigmas = torch.FloatTensor(sigs)
return torch.cat((sigmas[:-2], sigmas[-1:]))
else:
return torch.FloatTensor(sigs)
class PowerShiftSchedulerNode:
@classmethod
def INPUT_TYPES(s):
return {"required":
{"model": ("MODEL",),
"steps": ("INT", {"default": 20, "min": 3, "max": 1000}),
"power": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 5.0, "step": 0.001}),
"midpoint_shift": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.001}),
"discard_penultimate": ("BOOLEAN", {"default": False}),
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
}
}
RETURN_TYPES = ("SIGMAS",)
CATEGORY = "sampling/custom_sampling/schedulers"
FUNCTION = "get_sigmas"
def get_sigmas(self, model, steps, power, midpoint_shift, discard_penultimate, denoise):
total_steps = steps
if denoise < 1.0:
total_steps = int(steps/denoise)
sigmas = power_shift_scheduler(model.get_model_object("model_sampling"), total_steps, power, midpoint_shift, discard_penultimate=discard_penultimate).cpu()
sigmas = sigmas[-(steps + 1):]
return (sigmas, )
scheduler_name = "power_shift"
if scheduler_name not in SCHEDULER_HANDLERS:
scheduler_handler = SchedulerHandler(handler=power_shift_scheduler, use_ms=True)
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
if scheduler_name not in SCHEDULER_NAMES:
SCHEDULER_NAMES.append(scheduler_name)
def modify_by_twenties(num: int) -> int:
return num + (num // 20)
@@ -177,6 +120,37 @@ if scheduler_name not in SCHEDULER_HANDLERS:
if scheduler_name not in SCHEDULER_NAMES:
SCHEDULER_NAMES.append(scheduler_name)
def power_shift_scheduler(model_sampling, steps, power=2.0, midpoint_shift=1.0, discard_penultimate=False):
total_timesteps = (len(model_sampling.sigmas) - 1)
x = numpy.linspace(0, 1, steps, endpoint=False)
x = x**midpoint_shift
ts_normalized = (1 - x**power)**power
ts = numpy.rint(ts_normalized * total_timesteps)
sigs = []
last_t = -1
for t in ts:
t_int = min(int(t), total_timesteps)
if t_int != last_t:
sigs.append(float(model_sampling.sigmas[t_int]))
last_t = t_int
sigs.append(0.0)
if discard_penultimate is True:
sigmas = torch.FloatTensor(sigs)
return torch.cat((sigmas[:-2], sigmas[-1:]))
else:
return torch.FloatTensor(sigs)
scheduler_name = "power_shift"
if scheduler_name not in SCHEDULER_HANDLERS:
scheduler_handler = SchedulerHandler(handler=power_shift_scheduler, use_ms=True)
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
if scheduler_name not in SCHEDULER_NAMES:
SCHEDULER_NAMES.append(scheduler_name)
def radiance_shift_scheduler(model_sampling, steps, power=2.4, midpoint_shift=0.98, discard_penultimate=True):
total_timesteps = (len(model_sampling.sigmas) - 1)
real_steps = steps + 1
@@ -208,9 +182,67 @@ if scheduler_name not in SCHEDULER_HANDLERS:
if scheduler_name not in SCHEDULER_NAMES:
SCHEDULER_NAMES.append(scheduler_name)
class PowerShiftSchedulerNode:
@classmethod
def INPUT_TYPES(s):
return {"required":
{"model": ("MODEL",),
"steps": ("INT", {"default": 20, "min": 3, "max": 1000}),
"power": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 5.0, "step": 0.001}),
"midpoint_shift": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.001}),
"discard_penultimate": ("BOOLEAN", {"default": False}),
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
}
}
RETURN_TYPES = ("SIGMAS",)
CATEGORY = "sampling/custom_sampling/schedulers"
FUNCTION = "get_sigmas"
def get_sigmas(self, model, steps, power, midpoint_shift, discard_penultimate, denoise):
total_steps = steps
if denoise < 1.0:
total_steps = int(steps/denoise)
sigmas = power_shift_scheduler(model.get_model_object("model_sampling"), total_steps, power, midpoint_shift, discard_penultimate=discard_penultimate).cpu()
sigmas = sigmas[-(steps + 1):]
return (sigmas, )
class RadianceShiftSchedulerNode:
@classmethod
def INPUT_TYPES(s):
return {"required":
{"model": ("MODEL",),
"steps": ("INT", {"default": 20, "min": 3, "max": 1000}),
"power": ("FLOAT", {"default": 2.4, "min": 0.0, "max": 5.0, "step": 0.001}),
"midpoint_shift": ("FLOAT", {"default": 0.98, "min": 0.0, "max": 5.0, "step": 0.001}),
"discard_penultimate": ("BOOLEAN", {"default": True}),
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
}
}
RETURN_TYPES = ("SIGMAS",)
CATEGORY = "sampling/custom_sampling/schedulers"
FUNCTION = "get_sigmas"
def get_sigmas(self, model, steps, power, midpoint_shift, discard_penultimate, denoise):
total_steps = steps
if denoise < 1.0:
total_steps = int(steps/denoise)
sigmas = power_shift_scheduler(model.get_model_object("model_sampling"), total_steps, power, midpoint_shift, discard_penultimate=discard_penultimate).cpu()
sigmas = sigmas[-(steps + 1):]
return (sigmas, )
NODE_CLASS_MAPPINGS = {
"PowerShiftScheduler": PowerShiftSchedulerNode,
"RadianceShiftScheduler": RadianceShiftSchedulerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"PowerShiftScheduler": "Power Shift Scheduler",
"RadianceShiftScheduler": "Radiance Shift Scheduler",
}
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI_PowerShiftScheduler"
description = 'Highly customizable Scheduler for ComfyUI.'
version = "1.0.4"
version = "1.0.5"
license = { file = "LICENSE" }
[project.urls]