Add RadianceShiftScheduler node and bump version from 1.0.4 to 1.0.5
This commit is contained in:
+89
-57
@@ -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
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user