508 lines
17 KiB
Python
508 lines
17 KiB
Python
import torch
|
|
import numpy
|
|
import scipy.stats
|
|
from scipy.interpolate import PchipInterpolator
|
|
from comfy.samplers import SchedulerHandler, SCHEDULER_HANDLERS, SCHEDULER_NAMES
|
|
|
|
|
|
def modify_by_twenties(num: int) -> int:
|
|
return num + (num // 20)
|
|
|
|
def modify_by_tens(num: int) -> int:
|
|
return num + (num // 10)
|
|
|
|
def parse_float_list(s: str):
|
|
return [float(x.strip()) for x in s.split(",") if x.strip()]
|
|
|
|
BASE_SIGMA_POINTS = [
|
|
1.0,
|
|
0.99375,
|
|
0.9875,
|
|
0.98125,
|
|
0.975,
|
|
0.909375,
|
|
0.725,
|
|
0.421875,
|
|
]
|
|
|
|
|
|
def beta_33_scheduler(model_sampling, steps, alpha=0.3, beta=0.3):
|
|
total_timesteps = (len(model_sampling.sigmas) - 1)
|
|
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
|
|
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
|
|
|
|
sigs = []
|
|
last_t = -1
|
|
for t in ts:
|
|
if t != last_t:
|
|
sigs += [float(model_sampling.sigmas[int(t)])]
|
|
last_t = t
|
|
sigs += [0.0]
|
|
return torch.FloatTensor(sigs)
|
|
|
|
scheduler_name = "beta_33"
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=beta_33_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
if scheduler_name not in SCHEDULER_NAMES:
|
|
SCHEDULER_NAMES.append(scheduler_name)
|
|
|
|
|
|
def beta_44_scheduler(model_sampling, steps, alpha=0.4, beta=0.4):
|
|
total_timesteps = (len(model_sampling.sigmas) - 1)
|
|
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
|
|
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
|
|
|
|
sigs = []
|
|
last_t = -1
|
|
for t in ts:
|
|
if t != last_t:
|
|
sigs += [float(model_sampling.sigmas[int(t)])]
|
|
last_t = t
|
|
sigs += [0.0]
|
|
return torch.FloatTensor(sigs)
|
|
|
|
scheduler_name = "beta_44"
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=beta_44_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
if scheduler_name not in SCHEDULER_NAMES:
|
|
SCHEDULER_NAMES.append(scheduler_name)
|
|
|
|
|
|
def beta_53_scheduler(model_sampling, steps, alpha=0.5, beta=0.3):
|
|
total_timesteps = (len(model_sampling.sigmas) - 1)
|
|
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
|
|
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
|
|
|
|
sigs = []
|
|
last_t = -1
|
|
for t in ts:
|
|
if t != last_t:
|
|
sigs += [float(model_sampling.sigmas[int(t)])]
|
|
last_t = t
|
|
sigs += [0.0]
|
|
return torch.FloatTensor(sigs)
|
|
|
|
scheduler_name = "beta_53"
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=beta_53_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
if scheduler_name not in SCHEDULER_NAMES:
|
|
SCHEDULER_NAMES.append(scheduler_name)
|
|
|
|
|
|
def beta_54_scheduler(model_sampling, steps, alpha=0.5, beta=0.4):
|
|
total_timesteps = (len(model_sampling.sigmas) - 1)
|
|
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
|
|
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
|
|
|
|
sigs = []
|
|
last_t = -1
|
|
for t in ts:
|
|
if t != last_t:
|
|
sigs += [float(model_sampling.sigmas[int(t)])]
|
|
last_t = t
|
|
sigs += [0.0]
|
|
return torch.FloatTensor(sigs)
|
|
|
|
scheduler_name = "beta_54"
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=beta_54_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
if scheduler_name not in SCHEDULER_NAMES:
|
|
SCHEDULER_NAMES.append(scheduler_name)
|
|
|
|
|
|
def beta_57_scheduler(model_sampling, steps, alpha=0.5, beta=0.7):
|
|
total_timesteps = (len(model_sampling.sigmas) - 1)
|
|
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
|
|
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
|
|
|
|
sigs = []
|
|
last_t = -1
|
|
for t in ts:
|
|
if t != last_t:
|
|
sigs += [float(model_sampling.sigmas[int(t)])]
|
|
last_t = t
|
|
sigs += [0.0]
|
|
return torch.FloatTensor(sigs)
|
|
|
|
scheduler_name = "beta_57"
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=beta_57_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
if scheduler_name not in SCHEDULER_NAMES:
|
|
SCHEDULER_NAMES.append(scheduler_name)
|
|
|
|
|
|
def beta_32_scheduler(model_sampling, steps, alpha=0.3, beta=0.2):
|
|
steps = modify_by_twenties(steps)
|
|
total_timesteps = (len(model_sampling.sigmas) - 1)
|
|
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
|
|
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
|
|
|
|
sigs = []
|
|
last_t = -1
|
|
for t in ts:
|
|
if t != last_t:
|
|
sigs += [float(model_sampling.sigmas[int(t)])]
|
|
last_t = t
|
|
sigs += [0.0]
|
|
return torch.FloatTensor(sigs)
|
|
|
|
scheduler_name = "beta_32"
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=beta_32_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
if scheduler_name not in SCHEDULER_NAMES:
|
|
SCHEDULER_NAMES.append(scheduler_name)
|
|
|
|
|
|
def beta_43_scheduler(model_sampling, steps, alpha=0.4, beta=0.3):
|
|
steps = modify_by_twenties(steps)
|
|
total_timesteps = (len(model_sampling.sigmas) - 1)
|
|
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
|
|
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
|
|
|
|
sigs = []
|
|
last_t = -1
|
|
for t in ts:
|
|
if t != last_t:
|
|
sigs += [float(model_sampling.sigmas[int(t)])]
|
|
last_t = t
|
|
sigs += [0.0]
|
|
return torch.FloatTensor(sigs)
|
|
|
|
scheduler_name = "beta_43"
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=beta_43_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
if scheduler_name not in SCHEDULER_NAMES:
|
|
SCHEDULER_NAMES.append(scheduler_name)
|
|
|
|
def beta_42_scheduler(model_sampling, steps, alpha=0.4, beta=0.2):
|
|
steps = modify_by_tens(steps)
|
|
total_timesteps = (len(model_sampling.sigmas) - 1)
|
|
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
|
|
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
|
|
|
|
sigs = []
|
|
last_t = -1
|
|
for t in ts:
|
|
if t != last_t:
|
|
sigs += [float(model_sampling.sigmas[int(t)])]
|
|
last_t = t
|
|
sigs += [0.0]
|
|
return torch.FloatTensor(sigs)
|
|
|
|
scheduler_name = "beta_42"
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=beta_42_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
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
|
|
x = numpy.linspace(0, 1, real_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 = "radiance_shift"
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=radiance_shift_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
if scheduler_name not in SCHEDULER_NAMES:
|
|
SCHEDULER_NAMES.append(scheduler_name)
|
|
|
|
|
|
def sigma_curve_scheduler(model_sampling, steps, discard_penultimate=False, sigma_points=BASE_SIGMA_POINTS):
|
|
"""
|
|
Interpolates a sigma schedule from sigma_points to an arbitrary number
|
|
of steps, then appends a final 0.0 sigma.
|
|
|
|
For steps == 8, this reproduces exactly:
|
|
1.0
|
|
0.99375
|
|
0.9875
|
|
0.98125
|
|
0.975
|
|
0.909375
|
|
0.725
|
|
0.421875
|
|
0.0
|
|
"""
|
|
|
|
base = numpy.array(sigma_points, dtype=numpy.float32)
|
|
n_base = len(base) # 8 points
|
|
|
|
# Curve is defined over indices [0, n_base-1]
|
|
x_base = numpy.linspace(0.0, float(n_base - 1), n_base)
|
|
# We want exactly `steps` curve points (excluding the final 0.0)
|
|
x_new = numpy.linspace(0.0, float(n_base - 1), steps)
|
|
|
|
curve_sigs = numpy.interp(x_new, x_base, base).astype(numpy.float32)
|
|
|
|
# Append terminal 0.0 (not part of the curve shape)
|
|
sigs = numpy.concatenate([curve_sigs, numpy.array([0.0], dtype=numpy.float32)])
|
|
|
|
if discard_penultimate:
|
|
sigs_t = torch.from_numpy(sigs)
|
|
return torch.cat((sigs_t[:-2], sigs_t[-1:]))
|
|
|
|
return torch.from_numpy(sigs)
|
|
|
|
scheduler_name = "sigma_curve_from_points"
|
|
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=sigma_curve_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
if scheduler_name not in SCHEDULER_NAMES:
|
|
SCHEDULER_NAMES.append(scheduler_name)
|
|
|
|
def sigma_curve_pchip_scheduler(model_sampling, steps, discard_penultimate=False, sigma_points=BASE_SIGMA_POINTS):
|
|
"""
|
|
PCHIP monotonic cubic Hermite spline over sigma_points.
|
|
Produces `steps` curve points + final 0.0.
|
|
"""
|
|
|
|
base = numpy.array(sigma_points, dtype=numpy.float32)
|
|
n_base = len(base)
|
|
|
|
# Domain for the base curve: 0 .. n_base-1
|
|
x_base = numpy.linspace(0.0, float(n_base - 1), n_base)
|
|
|
|
# Domain for the new curve: steps points
|
|
x_new = numpy.linspace(0.0, float(n_base - 1), steps)
|
|
|
|
# PCHIP spline (monotonic, no overshoot)
|
|
pchip = PchipInterpolator(x_base, base)
|
|
curve_sigs = pchip(x_new).astype(numpy.float32)
|
|
|
|
# Append terminal 0.0
|
|
sigs = numpy.concatenate([curve_sigs, numpy.array([0.0], dtype=numpy.float32)])
|
|
|
|
if discard_penultimate:
|
|
sigs_t = torch.from_numpy(sigs)
|
|
return torch.cat((sigs_t[:-2], sigs_t[-1:]))
|
|
|
|
return torch.from_numpy(sigs)
|
|
|
|
|
|
scheduler_name = "sigma_curve_pchip"
|
|
|
|
if scheduler_name not in SCHEDULER_HANDLERS:
|
|
scheduler_handler = SchedulerHandler(handler=sigma_curve_pchip_scheduler, use_ms=True)
|
|
SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
|
|
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, )
|
|
|
|
|
|
class SigmaCurveFromPointsSchedulerNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"steps": ("INT", {"default": 8, "min": 1, "max": 1000}),
|
|
"discard_penultimate": ("BOOLEAN", {"default": False}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
},
|
|
"optional": {
|
|
"custom_points": ("STRING", {
|
|
"multiline": False,
|
|
"tooltip": "sigma curve points provided as comma separated list of floats. e.g. '1.0, 0.9, 0.8' etc."
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SIGMAS",)
|
|
CATEGORY = "sampling/custom_sampling/schedulers"
|
|
FUNCTION = "get_sigmas"
|
|
|
|
def get_sigmas(self, model, steps, discard_penultimate, denoise, custom_points=None):
|
|
total_steps = steps
|
|
if 0.0 < denoise < 1.0:
|
|
total_steps = int(steps / denoise)
|
|
|
|
if custom_points is not None:
|
|
sigma_points = parse_float_list(custom_points)
|
|
if len(sigma_points) < 2:
|
|
sigma_points = BASE_SIGMA_POINTS
|
|
else:
|
|
sigma_points = BASE_SIGMA_POINTS
|
|
|
|
sigmas = sigma_curve_scheduler(
|
|
model.get_model_object("model_sampling"),
|
|
total_steps,
|
|
discard_penultimate=discard_penultimate,
|
|
sigma_points=sigma_points,
|
|
).cpu()
|
|
|
|
sigmas = sigmas[-(steps + 1):]
|
|
return (sigmas,)
|
|
|
|
class SigmaCurvePchipSchedulerNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"steps": ("INT", {"default": 8, "min": 1, "max": 2000}),
|
|
"discard_penultimate": ("BOOLEAN", {"default": False}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
},
|
|
"optional": {
|
|
"custom_points": ("STRING", {
|
|
"multiline": False,
|
|
"tooltip": "sigma curve points provided as comma separated list of floats. e.g. '1.0, 0.9, 0.8'"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SIGMAS",)
|
|
CATEGORY = "sampling/custom_sampling/schedulers"
|
|
FUNCTION = "get_sigmas"
|
|
|
|
def get_sigmas(self, model, steps, discard_penultimate, denoise, custom_points=None):
|
|
total_steps = steps
|
|
if 0.0 < denoise < 1.0:
|
|
total_steps = int(steps / denoise)
|
|
|
|
if custom_points is not None:
|
|
sigma_points = parse_float_list(custom_points)
|
|
if len(sigma_points) < 2:
|
|
sigma_points = BASE_SIGMA_POINTS
|
|
else:
|
|
sigma_points = BASE_SIGMA_POINTS
|
|
|
|
sigmas = sigma_curve_pchip_scheduler(
|
|
model.get_model_object("model_sampling"),
|
|
total_steps,
|
|
discard_penultimate=discard_penultimate,
|
|
sigma_points=sigma_points,
|
|
).cpu()
|
|
|
|
sigmas = sigmas[-(steps + 1):]
|
|
return (sigmas,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PowerShiftScheduler": PowerShiftSchedulerNode,
|
|
"RadianceShiftScheduler": RadianceShiftSchedulerNode,
|
|
"SigmaCurveFromPointsScheduler": SigmaCurveFromPointsSchedulerNode,
|
|
"SigmaCurvePchipScheduler": SigmaCurvePchipSchedulerNode,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PowerShiftScheduler": "Power Shift Scheduler",
|
|
"RadianceShiftScheduler": "Radiance Shift Scheduler",
|
|
"SigmaCurveFromPointsScheduler": "From Points Scheduler",
|
|
"SigmaCurvePchipScheduler": "PCHIP Scheduler",
|
|
}
|