Upload webUI_ExtraSchedulers using SD-Hub

webUI_ExtraSchedulers/README.md

@@ -0,0 +1,46 @@
+## Extra Schedulers extension for Stable Diffusion webUI ##
+### built for new Forge, partial support for Automatic1111, and probably reForge ###
+#### (webUI must have split sampler/scheduler selection) ####
+
+>[!IMPORTANT]
+>not for old Forge. For some related stuff in old Forge, see my old [OverrideScheduler extension](https://github.com/DenOfEquity/SchedRide).
+
+### What do? ###
+Adds six new schedulers to the dropdown list:
+* cosine: follows a, you guessed it, cosine curve. Initial drop is relatively slow.
+* cosine-exponential blend: starts cosine, ends up exponential (long tail).
+* phi: (based on original by [Extraltodeus](https://github.com/Extraltodeus/sigmas_tools_and_the_golden_scheduler))
+* Laplace: (credit Tiankai et al. (2024), via Comfy)
+* Karras Dynamic: (via yoinked-h)
+* custom: either a list of sigmas [1.0, 0.6, 0.25, 0.1, 0.0] or an expression that will be evaluated for each sampling step. A list will be log-linear interpolated to the number of sampling steps. A list starting with 1.0 and ending with 0.0 will be scaled between sigma_max and sigma_min. Otherwise list will be interpreted as is.
+  * *m*: minimum sigma (adjustable in **Settings**, usually ~0.03)
+  * *M*: maximum sigma (adjustable in **Settings**, usually ~14.6)
+  * *n*: total steps
+  * *s*: this step
+  * *x*: step / (total steps - 1)
+  * *phi*: (1 + sqrt(5)) / 2
+
+Adds six new samplers:
+* Euler a CFG++ [Forge only]
+* Euler CFG++ [Forge only]
+* Euler Dy CFG++ (based on Euler Dy by Koishi-Star) [Forge only]
+* Euler SMEA Dy CFG++ (...) [Forge only]
+* Refined Exponential Solver (credit: Katherine Crowson, Birch-san, Clybius)
+* DPM++ 4M SDE (credit: Clybius)
+  
+### Why do? ###
+Different results, sometimes better. I tend to use cosine-exponential blend most of the time.
+
+### How do? ###
+*(schedulers)* It's just a calculation of different number sequences travelling from sigma_max to sigma_min over the set number of sampling steps, guiding the denoising process. Infinite possibilities, but few sweet spots.
+
+### Redo? ###
+Yes, custom scheduler is saved to image infotext and *params.txt*.
+
+### How install? ###
+Go to the **Extensions** tab, then **Install from URL**, use the URL for this repository.
+
+Then, go back to the **Installed** tab and hit **Apply and restart UI**.
+
+### more? ###
+Check the 'neg' branch instead.

webUI_ExtraSchedulers/old/extra_schedulers.py

@@ -0,0 +1,385 @@
+import gradio
+import math, numpy
+import torch
+from modules import scripts, shared
+
+# def get_sigmas_oss (n, sigma_min, sigma_max, device):
+    # # https://github.com/bebebe666/OptimalSteps
+    # def loglinear_interp(t_steps, num_steps):
+        # """
+        # Performs log-linear interpolation of a given array of decreasing numbers.
+        # """
+        # xs = numpy.linspace(0, 1, len(t_steps))
+        # ys = numpy.log(t_steps[::-1])
+
+        # new_xs = numpy.linspace(0, 1, num_steps)
+        # new_ys = numpy.interp(new_xs, xs, ys)
+
+        # interped_ys = numpy.exp(new_ys)[::-1].copy()
+        # return interped_ys
+
+    # if not shared.sd_model.is_webui_legacy_model():
+        # sigmas = [0.9968, 0.9886, 0.9819, 0.975, 0.966, 0.9471, 0.9158, 0.8287, 0.5512, 0.2808, 0.001]
+    # elif shared.sd_model.is_sd3:  #same as flux, but here for ease of changing later
+        # sigmas = [0.9968, 0.9886, 0.9819, 0.975, 0.966, 0.9471, 0.9158, 0.8287, 0.5512, 0.2808, 0.001]
+    # elif shared.sd_model.is_sdxl:   # fallback AYS11
+        # sigmas = [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.029]
+    # else:                           # fallback AYS11
+        # sigmas = [14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.029]
+
+    # if n != len(sigmas):
+        # sigmas = numpy.append(loglinear_interp(sigmas, n), [0.0])
+    # else:
+        # sigmas.append(0.0)
+
+    # return torch.FloatTensor(sigmas).to(device)
+
+
+def cosine_scheduler (n, sigma_min, sigma_max, device):
+    sigmas = torch.zeros(n, device=device)
+    if n == 1:
+        sigmas[0] = sigma_max ** 0.5
+    else:
+        for x in range(n):
+            p = x / (n-1)
+            C = sigma_min + 0.5*(sigma_max-sigma_min)*(1 - math.cos(math.pi*(1 - p**0.5)))
+            sigmas[x] = C
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+def cosexpblend_boost_scheduler (n, sigma_min, sigma_max, device):
+    sigmas = []
+    if n == 1:
+        sigmas.append(sigma_max ** 0.5)
+    else:
+        K = (sigma_min / sigma_max)**(1/(n-1))
+        E = sigma_max
+        detail = numpy.interp(numpy.linspace(0, 1, n), numpy.linspace(0, 1, 5), [1.0, 1.0, 1.27, 1.0, 1.0])
+        for x in range(n):
+            p = x / (n-1)
+            C = sigma_min + 0.5*(sigma_max-sigma_min)*(1 - math.cos(math.pi*(1 - p**0.5)))
+            sigmas.append(detail[x] * (C + p * (E - C)))
+            E *= K
+
+    sigmas += [0.0]
+    return torch.FloatTensor(sigmas).to(device)
+
+
+def cosexpblend_scheduler (n, sigma_min, sigma_max, device):
+    sigmas = []
+    if n == 1:
+        sigmas.append(sigma_max ** 0.5)
+    else:
+        K = (sigma_min / sigma_max)**(1/(n-1))
+        E = sigma_max
+        for x in range(n):
+            p = x / (n-1)
+            C = sigma_min + 0.5*(sigma_max-sigma_min)*(1 - math.cos(math.pi*(1 - p**0.5)))
+            sigmas.append(C + p * (E - C))
+            E *= K
+
+    sigmas += [0.0]
+    return torch.FloatTensor(sigmas).to(device)
+
+##  phi scheduler modified from original by @extraltodeus
+def phi_scheduler(n, sigma_min, sigma_max, device):
+    sigmas = torch.zeros(n, device=device)
+    if n == 1:
+        sigmas[0] = sigma_max ** 0.5
+    else:
+        phi = (1 + 5**0.5) / 2
+        for x in range(n):
+            sigmas[x] = sigma_min + (sigma_max-sigma_min)*((1-x/(n-1))**(phi*phi))
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+def get_sigmas_vp(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs a continuous VP noise schedule."""
+    
+    beta_d = 19.9
+    beta_min = 0.1
+    eps_s = 1e-3
+    
+    t = torch.linspace(1, eps_s, n, device=device)
+    sigmas = torch.sqrt(torch.exp(beta_d * t ** 2 / 2 + beta_min * t) - 1)
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+def get_sigmas_laplace(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs the noise schedule proposed by Tiankai et al. (2024). """
+    mu = 0.
+    beta = 0.5
+    epsilon = 1e-5 # avoid log(0)
+    x = torch.linspace(0, 1, n, device=device)
+    clamp = lambda x: torch.clamp(x, min=sigma_min, max=sigma_max)
+    lmb = mu - beta * torch.sign(0.5-x) * torch.log(1 - 2 * torch.abs(0.5-x) + epsilon)
+    sigmas = clamp(torch.exp(lmb))
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+
+
+def get_sigmas_sinusoidal_sf(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs a sinusoidal noise schedule."""
+    sf = 3.5
+    x = torch.linspace(0, 1, n, device=device)
+    sigmas = (sigma_min + (sigma_max - sigma_min) * (1 - torch.sin(torch.pi / 2 * x)))/sigma_max
+    sigmas = sigmas**sf
+    sigmas = sigmas * sigma_max
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_invcosinusoidal_sf(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs a sinusoidal noise schedule."""
+    sf = 3.5
+    x = torch.linspace(0, 1, n, device=device)
+    sigmas = (sigma_min + (sigma_max - sigma_min) * (0.5*(torch.cos(x * math.pi) + 1)))/sigma_max
+    sigmas = sigmas**sf
+    sigmas = sigmas * sigma_max
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_react_cosinusoidal_dynsf(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs a sinusoidal noise schedule."""
+    sf = 2.15
+    x = torch.linspace(0, 1, n, device=device)
+    sigmas = (sigma_min+(sigma_max-sigma_min)*(torch.cos(x*(torch.pi/2))))/sigma_max
+    sigmas = sigmas**(sf*(n*x/n))
+    sigmas = sigmas * sigma_max
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_karras_dynamic(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    rho = 7.
+    ramp = torch.linspace(0, 1, n, device=device)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = torch.zeros_like(ramp)
+    for i in range(n):
+        sigmas[i] = (max_inv_rho + ramp[i] * (min_inv_rho - max_inv_rho)) ** (math.cos(i*math.tau/n)*2+rho) 
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_karras_exponential_decay(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    rho = 7.
+    ramp = torch.linspace(0, 1, n, device=device)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = torch.zeros_like(ramp)
+    for i in range(n):
+        sigmas[i] = (max_inv_rho + ramp[i] * (min_inv_rho - max_inv_rho)) ** (rho-(3*i/n))
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_karras_exponential_increment(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    rho = 7.
+    ramp = torch.linspace(0, 1, n, device=device)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = torch.zeros_like(ramp)
+    for i in range(n):
+        sigmas[i] = (max_inv_rho + ramp[i] * (min_inv_rho - max_inv_rho)) ** (rho+3*i/n)
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+def custom_scheduler(n, sigma_min, sigma_max, device):
+    if 'import' in ExtraScheduler.customSigmas:
+        sigmas = torch.linspace(sigma_max, sigma_min, n, device=device)
+    elif 'eval' in ExtraScheduler.customSigmas:
+        sigmas = torch.linspace(sigma_max, sigma_min, n, device=device)
+    elif 'scripts' in ExtraScheduler.customSigmas:
+        sigmas = torch.linspace(sigma_max, sigma_min, n, device=device)
+
+    elif ExtraScheduler.customSigmas[0] == '[' and ExtraScheduler.customSigmas[-1] == ']':
+        sigmasList = [float(x) for x in ExtraScheduler.customSigmas.strip('[]').split(',')]
+
+        if sigmasList[0] == 1.0 and sigmasList[-1] == 0.0:
+            for x in range(len(sigmasList)):
+                sigmasList[x] *= (sigma_max - sigma_min)
+                sigmasList[x] += sigma_min
+        elif sigmasList[-1] == 0.0:
+            #don't interpolate to number of steps, use as is
+            return torch.tensor(sigmasList)
+
+        xs = numpy.linspace(0, 1, len(sigmasList))
+        ys = numpy.log(sigmasList[::-1])
+        
+        new_xs = numpy.linspace(0, 1, n)
+        new_ys = numpy.interp(new_xs, xs, ys)
+        
+        interpolated_ys = numpy.exp(new_ys)[::-1].copy()
+        sigmas = torch.tensor(interpolated_ys, device=device)
+    else:
+        sigmas = torch.linspace(sigma_max, sigma_min, n, device=device)
+        detail = numpy.interp(numpy.linspace(0, 1, n), numpy.linspace(0, 1, 5), [1.0, 1.0, 1.25, 1.0, 1.0])
+
+        phi = (1 + 5**0.5) / 2
+        pi = math.pi
+        
+        s = 0
+        while (s < n):
+            x = (s) / (n - 1)
+            M = sigma_max
+            m = sigma_min
+            d = detail[s]
+        
+            sigmas[s] = eval((ExtraScheduler.customSigmas))
+            s += 1
+
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+from scripts.simple_kes import get_sigmas_simple_kes
+
+from scripts.res_solver import sample_res_solver, sample_res_multistep, sample_res_multistep_cfgpp
+from scripts.clybius_dpmpp_4m_sde import sample_clyb_4m_sde_momentumized
+from scripts.gradient_estimation import sample_gradient_e, sample_gradient_e_cfgpp
+
+from modules import sd_samplers_common, sd_samplers
+from modules.sd_samplers_kdiffusion import sampler_extra_params, KDiffusionSampler
+
+class ExtraScheduler(scripts.Script):
+    sorting_priority = 99
+
+    installed = False
+    customSigmas = 'm + (M-m)*(1-x)**3'
+
+    def title(self):
+        return "Extra Schedulers (custom)"
+
+    def show(self, is_img2img):
+        # make this extension visible in both txt2img and img2img tab.
+        if ExtraScheduler.installed:
+            return scripts.AlwaysVisible
+        else:
+            return False
+
+    def ui(self, *args, **kwargs):
+        #with gradio.Accordion(open=False, label=self.title(), visible=ExtraScheduler.installed):
+        custom_sigmas = gradio.Textbox(value=ExtraScheduler.customSigmas, label='Extra Schedulers: custom function / list [n0, n1, n2, ...]', lines=1.01)
+
+        self.infotext_fields = [
+            (custom_sigmas, "es_custom"),
+        ]
+
+        return [custom_sigmas]
+
+    def process(self, params, *script_args, **kwargs):
+        if params.scheduler == 'custom':
+            custom_sigmas = script_args[0]
+            ExtraScheduler.customSigmas = custom_sigmas
+            params.extra_generation_params.update(dict(es_custom = ExtraScheduler.customSigmas, ))
+        elif params.scheduler == 'Simple KES':
+            params.extra_generation_params.update(dict(
+                es_KES_start_blend       = getattr(shared.opts, 'kes_start_blend'),
+                es_KES_end_blend         = getattr(shared.opts, 'kes_end_blend'),
+                es_KES_sharpness         = getattr(shared.opts, 'kes_sharpness'),
+                es_KES_initial_step_size = getattr(shared.opts, 'kes_initial_step_size'),
+                es_KES_final_step_size   = getattr(shared.opts, 'kes_final_step_size'),
+                es_KES_initial_noise     = getattr(shared.opts, 'kes_initial_noise'),
+                es_KES_final_noise       = getattr(shared.opts, 'kes_final_noise'),
+                es_KES_smooth_blend      = getattr(shared.opts, 'kes_smooth_blend'),
+                es_KES_step_size_factor  = getattr(shared.opts, 'kes_step_size_factor'),
+                es_KES_noise_scale       = getattr(shared.opts, 'kes_noise_scale'),
+            ))
+        return
+
+try:
+    import modules.sd_schedulers as schedulers
+
+    if "name='custom'" not in str(schedulers.schedulers[-1]):   # this is a bit lazy tbh
+        print ("Extension: Extra Schedulers: adding new schedulers")
+        CosineScheduler         = schedulers.Scheduler("cosine",        "Cosine",                   cosine_scheduler)
+        CosExpScheduler         = schedulers.Scheduler("cosexp",        "CosineExponential blend",  cosexpblend_scheduler)
+        CosExpBScheduler        = schedulers.Scheduler("cosprev",       "CosExp blend boost",       cosexpblend_boost_scheduler)
+        PhiScheduler            = schedulers.Scheduler("phi",           "Phi",                      phi_scheduler)
+        VPScheduler             = schedulers.Scheduler("vp",            "VP",                       get_sigmas_vp)
+        LaplaceScheduler        = schedulers.Scheduler("laplace",       "Laplace",                  get_sigmas_laplace)
+
+        SineScheduler           = schedulers.Scheduler("sine_sc",       "Sine scaled",              get_sigmas_sinusoidal_sf)
+        InvCosScheduler         = schedulers.Scheduler("inv_cos_sc",    "Inverse Cosine scaled",    get_sigmas_invcosinusoidal_sf)
+        CosDynScheduler         = schedulers.Scheduler("cosine_dyn",    "Cosine Dynamic",           get_sigmas_react_cosinusoidal_dynsf)
+        KarrasDynScheduler      = schedulers.Scheduler("karras_dyn",    "Karras Dynamic",           get_sigmas_karras_dynamic)
+        KarrasExpDecayScheduler = schedulers.Scheduler("karras_exp_d",  "Karras Exp Decay",         get_sigmas_karras_exponential_decay)
+        KarrasExpIncScheduler   = schedulers.Scheduler("karras_exp_i",  "Karras Exp Inc",           get_sigmas_karras_exponential_increment)
+
+        SimpleKEScheduler       = schedulers.Scheduler("simple_kes",    "Simple KES",               get_sigmas_simple_kes)
+
+        # OSSFlowScheduler        = schedulers.Scheduler("optimal_ss",    "Optimal Steps",            get_sigmas_oss)
+
+        CustomScheduler         = schedulers.Scheduler("custom",        "custom",                   custom_scheduler)
+
+
+        schedulers.schedulers.append(CosineScheduler)
+        schedulers.schedulers.append(CosExpScheduler)
+        schedulers.schedulers.append(CosExpBScheduler)
+        schedulers.schedulers.append(PhiScheduler)
+        schedulers.schedulers.append(VPScheduler)
+        schedulers.schedulers.append(LaplaceScheduler)
+
+        schedulers.schedulers.append(SineScheduler)
+        schedulers.schedulers.append(InvCosScheduler)
+        schedulers.schedulers.append(CosDynScheduler)
+        schedulers.schedulers.append(KarrasDynScheduler)
+        schedulers.schedulers.append(KarrasExpDecayScheduler)
+        schedulers.schedulers.append(KarrasExpIncScheduler)
+
+        schedulers.schedulers.append(SimpleKEScheduler)
+
+        # schedulers.schedulers.append(OSSFlowScheduler)
+
+        schedulers.schedulers.append(CustomScheduler)
+        schedulers.schedulers_map = {**{x.name: x for x in schedulers.schedulers}, **{x.label: x for x in schedulers.schedulers}}
+
+        try:
+            # CFG++ method is Forge only, not working in A1111
+            import modules_forge.forge_version
+            from scripts.samplers_cfgpp import sample_euler_ancestral_cfgpp, sample_euler_cfgpp, sample_euler_dy_cfgpp, sample_euler_smea_dy_cfgpp, sample_euler_negative_cfgpp, sample_euler_negative_dy_cfgpp
+            from scripts.forgeClassic_cfgpp import sample_dpmpp_sde_cfgpp, sample_dpmpp_2m_cfgpp, sample_dpmpp_2m_sde_cfgpp, sample_dpmpp_3m_sde_cfgpp, sample_dpmpp_2s_ancestral_cfgpp
+            samplers_cfgpp = [
+                ("Euler a CFG++",           sample_euler_ancestral_cfgpp,   ["k_euler_a_cfgpp"],            {"uses_ensd": True} ),
+                ("Euler CFG++",             sample_euler_cfgpp,             ["k_euler_cfgpp"],              {}                  ),
+                ("Euler Dy CFG++",          sample_euler_dy_cfgpp,          ["k_euler_dy_cfgpp"],           {}                  ),
+                ("Euler SMEA Dy CFG++",     sample_euler_smea_dy_cfgpp,     ["k_euler_smea_dy_cfgpp"],      {}                  ),
+                ("Euler Negative CFG++",    sample_euler_negative_cfgpp,    ["k_euler_negative_cfgpp"],     {}                  ),
+                ("Euler Negative Dy CFG++", sample_euler_negative_dy_cfgpp, ["k_euler_negative_dy_cfgpp"],  {}                  ),
+                ("RES multistep CFG++",     sample_res_multistep_cfgpp,     ["k_res_multi_cfgpp"],          {}                  ),
+                ("Gradient Estimation CFG++", sample_gradient_e_cfgpp,      ["k_grad_est_cfgpp"],           {}                  ),
+                ("DPM++ SDE CFG++",         sample_dpmpp_sde_cfgpp,         ["k_dpmpp_sde_cfgpp"],          {"brownian_noise": True, "second_order": True} ),
+                ("DPM++ 2M CFG++",          sample_dpmpp_2m_cfgpp,          ["k_dpmpp_2m_cfgpp"],           {}                  ),
+                ("DPM++ 2M SDE CFG++",      sample_dpmpp_2m_sde_cfgpp,      ["k_dpmpp_2m_sde_cfgpp"],       {"brownian_noise": True} ),
+                ("DPM++ 3M SDE CFG++",      sample_dpmpp_3m_sde_cfgpp,      ["k_dpmpp_3m_sde_cfgpp"],       {"brownian_noise": True, 'discard_next_to_last_sigma': True} ),
+                ("DPM++ 2S a CFG++",        sample_dpmpp_2s_ancestral_cfgpp,["k_dpmpp_2s_a_cfgpp"],         {"uses_ensd": True, "second_order": True} ),
+            ]
+            samplers_data_cfgpp = [
+                sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options)
+                for label, funcname, aliases, options in samplers_cfgpp
+                if callable(funcname)
+            ]
+            sampler_extra_params['sample_euler_cfgpp']             = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+            sampler_extra_params['sample_euler_negative_cfgpp']    = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+            sampler_extra_params['sample_euler_dy_cfgpp']          = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+            sampler_extra_params['sample_euler_negative_dy_cfgpp'] = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+            sampler_extra_params['sample_euler_smea_dy_cfgpp']     = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+
+            sampler_extra_params['sample_dpmpp_sde_cfgpp']         = ['s_noise']
+            sampler_extra_params['sample_dpmpp_2m_sde_cfgpp']      = ['s_noise']
+            sampler_extra_params['sample_dpmpp_3m_sde_cfgpp']      = ['s_noise']
+            sampler_extra_params['sample_dpmpp_2s_ancestral_cfgpp']= ['s_noise']
+
+            sd_samplers.all_samplers.extend(samplers_data_cfgpp)
+        except:
+            pass
+
+        samplers_extra = [
+            ("RES multistep",                sample_res_multistep,              ["k_res_multi"],        {}),
+            ("Refined Exponential Solver",   sample_res_solver,                 ["k_res"],              {}),
+            ("DPM++ 4M SDE",                 sample_clyb_4m_sde_momentumized,   ["k_dpmpp_4m_sde"],     {}),
+            ("Gradient Estimation",          sample_gradient_e,                 ["k_grad_est"],         {}),
+        ]
+        samplers_data_extra = [
+            sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options)
+            for label, funcname, aliases, options in samplers_extra
+            if callable(funcname)
+        ]
+
+        sd_samplers.all_samplers.extend(samplers_data_extra)
+        sd_samplers.all_samplers_map = {x.name: x for x in sd_samplers.all_samplers}
+        sd_samplers.set_samplers()
+
+    ExtraScheduler.installed = True
+except:
+    print ("Extension: Extra Schedulers: unsupported webUI")
+    ExtraScheduler.installed = False

webUI_ExtraSchedulers/old/forgeClassic_cfgpp.py

@@ -0,0 +1,269 @@
+# first 3 lifted from ForgeClassic (https://github.com/Haoming02/sd-webui-forge-classic/)
+# 4th is simple adaptation of 3M to 2M
+# 5th lifted from ReForge (https://github.com/Panchovix/stable-diffusion-webui-reForge)
+# all modified to work with Forge2
+
+import torch
+from tqdm.auto import trange
+from k_diffusion.sampling import (
+    default_noise_sampler,
+    BrownianTreeNoiseSampler,
+    get_ancestral_step,
+    to_d,
+)
+
+
+def _sigma_fn(t):
+    return t.neg().exp()
+
+
+def _t_fn(sigma):
+    return sigma.log().neg()
+
+
+@torch.no_grad()
+def sample_dpmpp_sde_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
+    eta = 1.0
+    s_noise = 1.0
+    r = 0.5
+
+    if len(sigmas) <= 1:
+        return x
+
+    seed = extra_args.get("seed", None)
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback(
+                {
+                    "x": x,
+                    "i": i,
+                    "sigma": sigmas[i],
+                    "sigma_hat": sigmas[i],
+                    "denoised": denoised,
+                }
+            )
+
+        if sigmas[i + 1] == 0:
+            d = model.last_noise_uncond
+            x = denoised + d * sigmas[i + 1]
+        else:
+            t, t_next = _t_fn(sigmas[i]), _t_fn(sigmas[i + 1])
+            h = t_next - t
+            s = t + h * r
+            fac = 1 / (2 * r)
+
+
+            sd, su = get_ancestral_step(_sigma_fn(t), _sigma_fn(s), eta)
+            s_ = _t_fn(sd)
+            x_2 = (_sigma_fn(s_) / _sigma_fn(t)) * x - (t - s_).expm1() * denoised
+            x_2 = x_2 + noise_sampler(_sigma_fn(t), _sigma_fn(s)) * s_noise * su
+            denoised_2 = model(x_2, _sigma_fn(s) * s_in, **extra_args)
+            u = x_2 - model.last_noise_uncond * _sigma_fn(s) * s_in                                             #d=(x-u)/sigma; d*sigma=x-u; u=x-d*sigma
+
+            sd, su = get_ancestral_step(_sigma_fn(t), _sigma_fn(t_next), eta)
+            denoised_d = (1 - fac) * u + fac * u
+            x = denoised_2 + to_d(x, sigmas[i], denoised_d) * sd
+            x = x + noise_sampler(_sigma_fn(t), _sigma_fn(t_next)) * s_noise * su
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_2m_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None):
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    old_uncond_denoised = None
+    uncond_denoised = None
+
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        uncond_denoised = x - model.last_noise_uncond * sigmas[i] * s_in
+        if callback is not None:
+            callback(
+                {
+                    "x": x,
+                    "i": i,
+                    "sigma": sigmas[i],
+                    "sigma_hat": sigmas[i],
+                    "denoised": denoised,
+                }
+            )
+        t, t_next = _t_fn(sigmas[i]), _t_fn(sigmas[i + 1])
+        h = t_next - t
+        if old_uncond_denoised is None or sigmas[i + 1] == 0:
+            denoised_mix = -torch.exp(-h) * uncond_denoised
+        else:
+            h_last = t - _t_fn(sigmas[i - 1])
+            r = h_last / h
+            denoised_mix = -torch.exp(-h) * uncond_denoised - torch.expm1(-h) * (1 / (2 * r)) * (denoised - old_uncond_denoised)
+        x = denoised + denoised_mix + torch.exp(-h) * x
+        old_uncond_denoised = uncond_denoised
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_3m_sde_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=None, s_noise=None, noise_sampler=None):
+    eta = 1.0 if eta is None else eta
+    s_noise = 1.0 if s_noise is None else s_noise
+
+    if len(sigmas) <= 1:
+        return x
+
+    seed = extra_args.get("seed", None)
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    denoised_1, denoised_2 = None, None
+    h, h_1, h_2 = None, None, None
+
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        u = x - model.last_noise_uncond * sigmas[i] * s_in                                            #d=(x-u)/sigma; d*sigma=x-u; u=x-d*sigma
+        if callback is not None:
+            callback(
+                {
+                    "x": x,
+                    "i": i,
+                    "sigma": sigmas[i],
+                    "sigma_hat": sigmas[i],
+                    "denoised": denoised,
+                }
+            )
+        if sigmas[i + 1] == 0:
+            x = denoised
+        else:
+            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = s - t
+            h_eta = h * (eta + 1)
+
+            x = torch.exp(-h_eta) * (x + (denoised - u)) + (-h_eta).expm1().neg() * denoised
+
+            if h_2 is not None:
+                r0 = h_1 / h
+                r1 = h_2 / h
+                d1_0 = (denoised - denoised_1) / r0
+                d1_1 = (denoised_1 - denoised_2) / r1
+                d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
+                d2 = (d1_0 - d1_1) / (r0 + r1)
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                phi_3 = phi_2 / h_eta - 0.5
+                x = x + phi_2 * d1 - phi_3 * d2
+            elif h_1 is not None:
+                r = h_1 / h
+                d = (denoised - denoised_1) / r
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                x = x + phi_2 * d
+
+            if eta:
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
+
+        denoised_1, denoised_2 = denoised, denoised_1
+        h_1, h_2 = h, h_1
+    return x
+
+
+## extra
+@torch.no_grad()
+def sample_dpmpp_2m_sde_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+# just cut down from 3m_sde version
+    seed = extra_args.get("seed", None)
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    denoised_1 = None
+    h_1 = None
+
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        u = x - model.last_noise_uncond * sigmas[i] * s_in
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            #Denoising step
+            x = denoised
+        else:
+            #DPM-Solver++(2M) SDE
+            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = s - t
+
+            h_eta = h * (eta + 1)
+            x = torch.exp(-h_eta) * (x + (denoised - u)) + (-h_eta).expm1().neg() * denoised
+
+            if denoised_1 is not None:
+                r = h_1 / h
+
+                d = (denoised - denoised_1) / r
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                x = x + phi_2 * d
+
+            if eta:
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
+
+            h_1 = h
+
+        denoised_1 = denoised
+    return x
+
+
+# via ReForge
+@torch.no_grad()
+def sample_dpmpp_2s_ancestral_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigma_down == 0:
+            # Euler method
+            d = model.last_noise_uncond 
+            dt = sigma_down - sigmas[i]
+            x = denoised + d * sigma_down
+        else:
+            u = x - model.last_noise_uncond * sigmas[i] * s_in
+
+            # DPM-Solver++(2S)
+            t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
+            # r = torch.sinh(1 + (2 - eta) * (t_next - t) / (t - t_fn(sigma_up))) works only on non-cfgpp, weird
+            r = 1 / 2
+            h = t_next - t
+            s = t + r * h
+            x_2 = (sigma_fn(s) / sigma_fn(t)) * (x + (denoised - u)) - (-h * r).expm1() * denoised
+            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
+            x = (sigma_fn(t_next) / sigma_fn(t)) * (x + (denoised - u)) - (-h).expm1() * denoised_2
+
+        # Noise addition
+        if sigmas[i + 1] > 0:
+            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+    return x

webUI_ExtraSchedulers/old/gradient_estimation.py

@@ -0,0 +1,70 @@
+## lifted from ReForge, original implementation from Comfy
+## CFG++ attempt by me
+
+import torch
+from tqdm.auto import trange
+
+
+#   copied from kdiffusion/sampling.py
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+def append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
+    return x[(...,) + (None,) * dims_to_append]
+
+
+@torch.no_grad()
+def sample_gradient_e(model, x, sigmas, extra_args=None, callback=None, disable=None, ge_gamma=2.):
+    """Gradient-estimation sampler. Paper: https://openreview.net/pdf?id=o2ND9v0CeK"""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    old_d = None
+
+    sigmas = sigmas.to(x.device)
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        
+        d = to_d(x, sigmas[i], denoised)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        dt = sigmas[i + 1] - sigmas[i]
+        if i == 0: # Euler method
+            x = x + d * dt
+        else:
+            # Gradient estimation
+            d_bar = ge_gamma * d + (1 - ge_gamma) * old_d
+            x = x + d_bar * dt
+        old_d = d
+    return x
+
+
+@torch.no_grad()
+def sample_gradient_e_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, ge_gamma=2.):
+    """Gradient-estimation sampler. Paper: https://openreview.net/pdf?id=o2ND9v0CeK"""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    old_d = None
+    
+    model.need_last_noise_uncond = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        dt = sigmas[i + 1] - sigmas[i]
+        if i == 0: # Euler method
+            x = denoised + d * sigmas[i+1]
+        else:
+            # Gradient estimation
+            d_bar = ge_gamma * d + (1 - ge_gamma) * old_d
+            x = denoised + d_bar * sigmas[i+1]
+        old_d = d
+    return x

webUI_ExtraSchedulers/old/res_solver.py

@@ -0,0 +1,396 @@
+import torch
+from torch import no_grad, FloatTensor
+from tqdm import tqdm
+from itertools import pairwise
+from typing import Protocol, Optional, Dict, Any, TypedDict, NamedTuple, Union, List
+import math
+
+from tqdm.auto import trange
+
+#   copied from kdiffusion/sampling.py and utils.py
+def default_noise_sampler(x):
+    return lambda sigma, sigma_next: torch.randn_like(x)
+def append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
+    return x[(...,) + (None,) * dims_to_append]
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+
+
+class DenoiserModel(Protocol):
+  def __call__(self, x: FloatTensor, t: FloatTensor, *args, **kwargs) -> FloatTensor: ...
+
+class RefinedExpCallbackPayload(TypedDict):
+  x: FloatTensor
+  i: int
+  sigma: FloatTensor
+  sigma_hat: FloatTensor
+
+class RefinedExpCallback(Protocol):
+  def __call__(self, payload: RefinedExpCallbackPayload) -> None: ...
+
+class NoiseSampler(Protocol):
+  def __call__(self, x: FloatTensor) -> FloatTensor: ...
+
+class StepOutput(NamedTuple):
+  x_next: FloatTensor
+  denoised: FloatTensor
+  denoised2: FloatTensor
+  vel: FloatTensor
+  vel_2: FloatTensor
+
+def _gamma(
+  n: int,
+) -> int:
+  """
+  https://en.wikipedia.org/wiki/Gamma_function
+  for every positive integer n,
+  Γ(n) = (n-1)!
+  """
+  return math.factorial(n-1)
+
+def _incomplete_gamma(
+  s: int,
+  x: float,
+  gamma_s: Optional[int] = None
+) -> float:
+  """
+  https://en.wikipedia.org/wiki/Incomplete_gamma_function#Special_values
+  if s is a positive integer,
+  Γ(s, x) = (s-1)!*∑{k=0..s-1}(x^k/k!)
+  """
+  if gamma_s is None:
+    gamma_s = _gamma(s)
+
+  sum_: float = 0
+  # {k=0..s-1} inclusive
+  for k in range(s):
+    numerator: float = x**k
+    denom: int = math.factorial(k)
+    quotient: float = numerator/denom
+    sum_ += quotient
+  incomplete_gamma_: float = sum_ * math.exp(-x) * gamma_s
+  return incomplete_gamma_
+
+# by Katherine Crowson
+def _phi_1(neg_h: FloatTensor):
+  return torch.nan_to_num(torch.expm1(neg_h) / neg_h, nan=1.0)
+
+# by Katherine Crowson
+def _phi_2(neg_h: FloatTensor):
+  return torch.nan_to_num((torch.expm1(neg_h) - neg_h) / neg_h**2, nan=0.5)
+
+# by Katherine Crowson
+def _phi_3(neg_h: FloatTensor):
+  return torch.nan_to_num((torch.expm1(neg_h) - neg_h - neg_h**2 / 2) / neg_h**3, nan=1 / 6)
+
+def _phi(
+  neg_h: float,
+  j: int,
+):
+  """
+  For j={1,2,3}: you could alternatively use Kat's phi_1, phi_2, phi_3 which perform fewer steps
+
+  Lemma 1
+  https://arxiv.org/abs/2308.02157
+  ϕj(-h) = 1/h^j*∫{0..h}(e^(τ-h)*(τ^(j-1))/((j-1)!)dτ)
+
+  https://www.wolframalpha.com/input?i=integrate+e%5E%28%CF%84-h%29*%28%CF%84%5E%28j-1%29%2F%28j-1%29%21%29d%CF%84
+  = 1/h^j*[(e^(-h)*(-τ)^(-j)*τ(j))/((j-1)!)]{0..h}
+  https://www.wolframalpha.com/input?i=integrate+e%5E%28%CF%84-h%29*%28%CF%84%5E%28j-1%29%2F%28j-1%29%21%29d%CF%84+between+0+and+h
+  = 1/h^j*((e^(-h)*(-h)^(-j)*h^j*(Γ(j)-Γ(j,-h)))/(j-1)!)
+  = (e^(-h)*(-h)^(-j)*h^j*(Γ(j)-Γ(j,-h))/((j-1)!*h^j)
+  = (e^(-h)*(-h)^(-j)*(Γ(j)-Γ(j,-h))/(j-1)!
+  = (e^(-h)*(-h)^(-j)*(Γ(j)-Γ(j,-h))/Γ(j)
+  = (e^(-h)*(-h)^(-j)*(1-Γ(j,-h)/Γ(j))
+
+  requires j>0
+  """
+  assert j > 0
+  gamma_: float = _gamma(j)
+  incomp_gamma_: float = _incomplete_gamma(j, neg_h, gamma_s=gamma_)
+
+  phi_: float = math.exp(neg_h) * neg_h**-j * (1-incomp_gamma_/gamma_)
+
+  return phi_
+
+class RESDECoeffsSecondOrder(NamedTuple):
+  a2_1: float
+  b1: float
+  b2: float
+
+def _de_second_order(
+  h: float,
+  c2: float,
+  simple_phi_calc = False,
+) -> RESDECoeffsSecondOrder:
+  """
+  Table 3
+  https://arxiv.org/abs/2308.02157
+  ϕi,j := ϕi,j(-h) = ϕi(-cj*h)
+  a2_1 = c2ϕ1,2
+       = c2ϕ1(-c2*h)
+  b1 = ϕ1 - ϕ2/c2
+  """
+  if simple_phi_calc:
+    # Kat computed simpler expressions for phi for cases j={1,2,3}
+    a2_1: float = c2 * _phi_1(-c2*h)
+    phi1: float = _phi_1(-h)
+    phi2: float = _phi_2(-h)
+  else:
+    # I computed general solution instead.
+    # they're close, but there are slight differences. not sure which would be more prone to numerical error.
+    a2_1: float = c2 * _phi(j=1, neg_h=-c2*h)
+    phi1: float = _phi(j=1, neg_h=-h)
+    phi2: float = _phi(j=2, neg_h=-h)
+  phi2_c2: float = phi2/c2
+  b1: float = phi1 - phi2_c2
+  b2: float = phi2_c2
+  return RESDECoeffsSecondOrder(
+    a2_1=a2_1,
+    b1=b1,
+    b2=b2,
+  )  
+
+def _refined_exp_sosu_step(
+  model: DenoiserModel,
+  x: FloatTensor,
+  sigma: FloatTensor,
+  sigma_next: FloatTensor,
+  c2 = 0.5,
+  extra_args: Dict[str, Any] = {},
+  pbar: Optional[tqdm] = None,
+  simple_phi_calc = False,
+  momentum = 0.0,
+  vel = None,
+  vel_2 = None,
+  time = None
+) -> StepOutput:
+  """
+  Algorithm 1 "RES Second order Single Update Step with c2"
+  https://arxiv.org/abs/2308.02157
+
+  Parameters:
+    model (`DenoiserModel`): a k-diffusion wrapped denoiser model (e.g. a subclass of DiscreteEpsDDPMDenoiser)
+    x (`FloatTensor`): noised latents (or RGB I suppose), e.g. torch.randn((B, C, H, W)) * sigma[0]
+    sigma (`FloatTensor`): timestep to denoise
+    sigma_next (`FloatTensor`): timestep+1 to denoise
+    c2 (`float`, *optional*, defaults to .5): partial step size for solving ODE. .5 = midpoint method
+    extra_args (`Dict[str, Any]`, *optional*, defaults to `{}`): kwargs to pass to `model#__call__()`
+    pbar (`tqdm`, *optional*, defaults to `None`): progress bar to update after each model call
+    simple_phi_calc (`bool`, *optional*, defaults to `True`): True = calculate phi_i,j(-h) via simplified formulae specific to j={1,2}. False = Use general solution that works for any j. Mathematically equivalent, but could be numeric differences.
+  """
+
+  def momentum_func(diff, velocity, timescale=1.0, offset=-momentum / 2.0): # Diff is current diff, vel is previous diff
+    if velocity is None:
+        momentum_vel = diff
+    else:
+        momentum_vel = momentum * (timescale + offset) * velocity + (1 - momentum * (timescale + offset)) * diff
+    return momentum_vel
+
+  lam_next, lam = (s.log().neg() for s in (sigma_next, sigma))
+
+  # type hints aren't strictly true regarding float vs FloatTensor.
+  # everything gets promoted to `FloatTensor` after interacting with `sigma: FloatTensor`.
+  # I will use float to indicate any variables which are scalars.
+  h: float = lam_next - lam
+  a2_1, b1, b2 = _de_second_order(h=h, c2=c2, simple_phi_calc=simple_phi_calc)
+  
+  denoised: FloatTensor = model(x, sigma.repeat(x.size(0)), **extra_args)
+  # if pbar is not None:
+    # pbar.update(0.5)
+
+  c2_h: float = c2*h
+
+  diff_2 = momentum_func(a2_1*h*denoised, vel_2, time)
+  vel_2 = diff_2
+  x_2: FloatTensor = math.exp(-c2_h)*x + diff_2
+  lam_2: float = lam + c2_h
+  sigma_2: float = lam_2.neg().exp()
+
+  denoised2: FloatTensor = model(x_2, sigma_2.repeat(x_2.size(0)), **extra_args)
+  if pbar is not None:
+    pbar.update()
+
+  diff = momentum_func(h*(b1*denoised + b2*denoised2), vel, time)
+  vel = diff
+
+  x_next: FloatTensor = math.exp(-h)*x + diff
+  
+  return StepOutput(
+    x_next=x_next,
+    denoised=denoised,
+    denoised2=denoised2,
+    vel=vel,
+    vel_2=vel_2,
+  )
+  
+
+@no_grad()
+def sample_refined_exp_s(
+  model: FloatTensor,
+  x: FloatTensor,
+  sigmas: FloatTensor,
+  denoise_to_zero: bool = True,
+  extra_args: Dict[str, Any] = {},
+  callback: Optional[RefinedExpCallback] = None,
+  disable: Optional[bool] = None,
+  ita: FloatTensor = torch.zeros((1,)),
+  c2 = .5,
+  noise_sampler: NoiseSampler = torch.randn_like,
+  simple_phi_calc = False,
+  momentum = 0.0,
+):
+  """
+  Refined Exponential Solver (S).
+  Algorithm 2 "RES Single-Step Sampler" with Algorithm 1 second-order step
+  https://arxiv.org/abs/2308.02157
+
+  Parameters:
+    model (`DenoiserModel`): a k-diffusion wrapped denoiser model (e.g. a subclass of DiscreteEpsDDPMDenoiser)
+    x (`FloatTensor`): noised latents (or RGB I suppose), e.g. torch.randn((B, C, H, W)) * sigma[0]
+    sigmas (`FloatTensor`): sigmas (ideally an exponential schedule!) e.g. get_sigmas_exponential(n=25, sigma_min=model.sigma_min, sigma_max=model.sigma_max)
+    denoise_to_zero (`bool`, *optional*, defaults to `True`): whether to finish with a first-order step down to 0 (rather than stopping at sigma_min). True = fully denoise image. False = match Algorithm 2 in paper
+    extra_args (`Dict[str, Any]`, *optional*, defaults to `{}`): kwargs to pass to `model#__call__()`
+    callback (`RefinedExpCallback`, *optional*, defaults to `None`): you can supply this callback to see the intermediate denoising results, e.g. to preview each step of the denoising process
+    disable (`bool`, *optional*, defaults to `False`): whether to hide `tqdm`'s progress bar animation from being printed
+    ita (`FloatTensor`, *optional*, defaults to 0.): degree of stochasticity, η, for each timestep. tensor shape must be broadcastable to 1-dimensional tensor with length `len(sigmas) if denoise_to_zero else len(sigmas)-1`. each element should be from 0 to 1.
+         - if used: batch noise doesn't match non-batch
+    c2 (`float`, *optional*, defaults to .5): partial step size for solving ODE. .5 = midpoint method
+    noise_sampler (`NoiseSampler`, *optional*, defaults to `torch.randn_like`): method used for adding noise
+    simple_phi_calc (`bool`, *optional*, defaults to `True`): True = calculate phi_i,j(-h) via simplified formulae specific to j={1,2}. False = Use general solution that works for any j. Mathematically equivalent, but could be numeric differences.
+  """
+  #assert sigmas[-1] == 0
+  device = x.device
+  ita = ita.to(device)
+  sigmas = sigmas.to(device)
+
+  sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+
+  vel, vel_2 = None, None
+  with tqdm(disable=disable, total=len(sigmas)-(1 if denoise_to_zero else 2)) as pbar:
+    for i, (sigma, sigma_next) in enumerate(pairwise(sigmas[:-1].split(1))):
+      time = sigmas[i] / sigma_max
+      if 'sigma' not in locals():
+        sigma = sigmas[i]
+      eps = torch.randn_like(x).float()
+      sigma_hat = sigma * (1 + ita)
+      x_hat = x + (sigma_hat ** 2 - sigma ** 2).sqrt() * eps
+      x_next, denoised, denoised2, vel, vel_2 = _refined_exp_sosu_step(
+        model,
+        x_hat,
+        sigma_hat,
+        sigma_next,
+        c2=c2,
+        extra_args=extra_args,
+        pbar=pbar,
+        simple_phi_calc=simple_phi_calc,
+        momentum = momentum,
+        vel = vel,
+        vel_2 = vel_2,
+        time = time
+      )
+      if callback is not None:
+        payload = RefinedExpCallbackPayload(
+          x=x,
+          i=i,
+          sigma=sigma,
+          sigma_hat=sigma_hat,
+          denoised=denoised,
+          denoised2=denoised2,
+        )
+        callback(payload)
+      x = x_next
+    if denoise_to_zero:
+      eps = torch.randn_like(x).float()
+      sigma_hat = sigma * (1 + ita)
+      x_hat = x + (sigma_hat ** 2 - sigma ** 2).sqrt() * eps
+      x_next: FloatTensor = model(x_hat, sigma.to(x_hat.device).repeat(x_hat.size(0)), **extra_args)
+      pbar.update()
+
+      if callback is not None:
+        payload = RefinedExpCallbackPayload(
+          x=x,
+          i=i,
+          sigma=sigma,
+          sigma_hat=sigma_hat,
+          denoised=denoised,
+          denoised2=denoised2,
+        )
+        callback(payload)
+
+
+      x = x_next
+  return x
+
+# Many thanks to Kat + Birch-San for this wonderful sampler implementation! https://github.com/Birch-san/sdxl-play/commits/res/
+def sample_res_solver(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler_type="gaussian", noise_sampler=None, denoise_to_zero=True, simple_phi_calc=False, c2=0.5, ita=torch.Tensor((0.0,)), momentum=0.0):
+    return sample_refined_exp_s(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler, denoise_to_zero=denoise_to_zero, simple_phi_calc=simple_phi_calc, c2=c2, ita=ita, momentum=momentum)
+
+
+##  modified from ReForge, original implementation ComfyUI
+@torch.no_grad()
+def res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., noise_sampler=None, cfgpp=False):
+    extra_args = {} if extra_args is None else extra_args
+    seed = extra_args.get("seed", None)
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+    phi1_fn = lambda t: torch.expm1(t) / t
+    phi2_fn = lambda t: (phi1_fn(t) - 1.0) / t
+    old_denoised = None
+
+    sigmas = sigmas.to(x.device)
+
+    if cfgpp:
+        model.need_last_noise_uncond = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        if s_churn > 0:
+            gamma = min(s_churn / (len(sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.0
+            sigma_hat = sigmas[i] * (gamma + 1)
+        else:
+            gamma = 0
+            sigma_hat = sigmas[i]
+        if gamma > 0:
+            eps = torch.randn_like(x) * s_noise
+            x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+
+        if callback is not None:
+            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigma_hat, "denoised": denoised})
+        if sigmas[i + 1] == 0 or old_denoised is None:
+            # Euler method
+            if cfgpp:
+                d = model.last_noise_uncond
+                x = denoised + d * sigmas[i + 1]
+            else:
+                d = to_d(x, sigma_hat, denoised)
+                dt = sigmas[i + 1] - sigma_hat
+                x = x + d * dt
+        else:
+            # Second order multistep method in https://arxiv.org/pdf/2308.02157
+            t, t_next, t_prev = t_fn(sigmas[i]), t_fn(sigmas[i + 1]), t_fn(sigmas[i - 1])
+            h = t_next - t
+            c2 = (t_prev - t) / h
+            phi1_val, phi2_val = phi1_fn(-h), phi2_fn(-h)
+            b1 = torch.nan_to_num(phi1_val - 1.0 / c2 * phi2_val, nan=0.0)
+            b2 = torch.nan_to_num(1.0 / c2 * phi2_val, nan=0.0)
+            if cfgpp:
+                d = model.last_noise_uncond
+                x = denoised + d * sigma_hat
+
+            x = (sigma_fn(t_next) / sigma_fn(t)) * x + h * (b1 * denoised + b2 * old_denoised)
+        old_denoised = denoised
+    return x
+@torch.no_grad()
+def sample_res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., noise_sampler=None):
+    return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise, noise_sampler=noise_sampler, cfgpp=False)
+@torch.no_grad()
+def sample_res_multistep_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., noise_sampler=None):
+    return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise, noise_sampler=noise_sampler, cfgpp=True)

webUI_ExtraSchedulers/old/samplers_cfgpp.py

@@ -0,0 +1,258 @@
+import torch
+from tqdm.auto import trange
+
+#   copied from kdiffusion/sampling.py and utils.py
+def default_noise_sampler(x):
+    return lambda sigma, sigma_next: torch.randn_like(x)
+def get_ancestral_step(sigma_from, sigma_to, eta=1.):
+    """Calculates the noise level (sigma_down) to step down to and the amount
+    of noise to add (sigma_up) when doing an ancestral sampling step."""
+    if not eta:
+        return sigma_to, 0.
+    sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
+    sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
+    return sigma_down, sigma_up
+def append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
+    return x[(...,) + (None,) * dims_to_append]
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+
+
+@torch.no_grad()
+def sample_euler_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        x = denoised + d * sigmas[i+1]
+    return x
+
+class _Rescaler:
+    def __init__(self, model, x, mode, **extra_args):
+        self.model = model
+        self.x = x
+        self.mode = mode
+        self.extra_args = extra_args
+        self.init_latent, self.mask, self.nmask = model.init_latent, model.mask, model.nmask
+
+    def __enter__(self):
+        if self.init_latent is not None:
+            self.model.init_latent = torch.nn.functional.interpolate(input=self.init_latent, size=self.x.shape[2:4], mode=self.mode)
+        if self.mask is not None:
+            self.model.mask = torch.nn.functional.interpolate(input=self.mask.unsqueeze(0), size=self.x.shape[2:4], mode=self.mode).squeeze(0)
+        if self.nmask is not None:
+            self.model.nmask = torch.nn.functional.interpolate(input=self.nmask.unsqueeze(0), size=self.x.shape[2:4], mode=self.mode).squeeze(0)
+
+        return self
+
+    def __exit__(self, type, value, traceback):
+        del self.model.init_latent, self.model.mask, self.model.nmask
+        self.model.init_latent, self.model.mask, self.model.nmask = self.init_latent, self.mask, self.nmask
+
+@torch.no_grad()
+def dy_sampling_step_cfgpp(x, model, sigma_hat, **extra_args):
+    original_shape = x.shape
+    batch_size, channels, m, n = original_shape[0], original_shape[1], original_shape[2] // 2, original_shape[3] // 2
+    extra_row = x.shape[2] % 2 == 1
+    extra_col = x.shape[3] % 2 == 1
+
+    if extra_row:
+        extra_row_content = x[:, :, -1:, :]
+        x = x[:, :, :-1, :]
+    if extra_col:
+        extra_col_content = x[:, :, :, -1:]
+        x = x[:, :, :, :-1]
+
+    a_list = x.unfold(2, 2, 2).unfold(3, 2, 2).contiguous().view(batch_size, channels, m * n, 2, 2)
+    c = a_list[:, :, :, 1, 1].view(batch_size, channels, m, n)
+
+    with _Rescaler(model, c, 'nearest-exact', **extra_args) as rescaler:
+        denoised = model(c, sigma_hat * c.new_ones([c.shape[0]]), **rescaler.extra_args)
+    d = model.last_noise_uncond
+    c = denoised + d * sigma_hat
+
+    d_list = c.view(batch_size, channels, m * n, 1, 1)
+    a_list[:, :, :, 1, 1] = d_list[:, :, :, 0, 0]
+    x = a_list.view(batch_size, channels, m, n, 2, 2).permute(0, 1, 2, 4, 3, 5).reshape(batch_size, channels, 2 * m, 2 * n)
+
+    if extra_row or extra_col:
+        x_expanded = torch.zeros(original_shape, dtype=x.dtype, device=x.device)
+        x_expanded[:, :, :2 * m, :2 * n] = x
+        if extra_row:
+            x_expanded[:, :, -1:, :2 * n + 1] = extra_row_content
+        if extra_col:
+            x_expanded[:, :, :2 * m, -1:] = extra_col_content
+        if extra_row and extra_col:
+            x_expanded[:, :, -1:, -1:] = extra_col_content[:, :, -1:, :]
+        x = x_expanded
+
+    return x
+
+@torch.no_grad()
+def smea_sampling_step_cfgpp(x, model, sigma_hat, **extra_args):
+    m, n = x.shape[2], x.shape[3]
+    x = torch.nn.functional.interpolate(input=x, scale_factor=(1.25, 1.25), mode='nearest-exact')
+    with _Rescaler(model, x, 'nearest-exact', **extra_args) as rescaler:
+        denoised = model(x, sigma_hat * x.new_ones([x.shape[0]]), **rescaler.extra_args)
+    d = model.last_noise_uncond
+    x = denoised + d * sigma_hat
+    x = torch.nn.functional.interpolate(input=x, size=(m,n), mode='nearest-exact')
+    return x
+
+
+@torch.no_grad()
+def sample_euler_dy_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """CFG++ version of Euler Dy by KoishiStar."""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        x = denoised + d * sigmas[i+1]
+
+        if sigmas[i + 1] > 0:
+            if i // 2 == 1:
+                x = dy_sampling_step_cfgpp(x, model, sigma_hat, **extra_args)        
+
+    return x
+
+@torch.no_grad()
+def sample_euler_negative_dy_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """CFG++ version of Euler Negative Dy by KoishiStar."""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        if sigmas[i + 1] > 0 and i // 2 == 1:
+            x = -denoised - d * sigmas[i+1]
+        else:
+            x = denoised + d * sigmas[i+1]
+
+        if sigmas[i + 1] > 0:
+            if i // 2 == 1:
+                x = dy_sampling_step_cfgpp(x, model, sigma_hat, **extra_args)        
+
+    return x
+
+@torch.no_grad()
+def sample_euler_negative_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """based on Euler Negative by KoishiStar"""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        if sigmas[i + 1] > 0 and i // 2 == 1:
+            x = -denoised - d * sigmas[i+1]
+        else:
+            x = denoised + d * sigmas[i+1]
+    return x
+
+
+@torch.no_grad()
+def sample_euler_smea_dy_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """CFG++ version of Euler SMEA Dy by KoishiStar."""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        x = denoised + d * sigmas[i+1]
+        
+        if sigmas[i + 1] > 0:
+            if i + 1 // 2 == 1:     #   ??  this is i == 1; why not if i // 2 == 1 same as Euler Dy
+                x = dy_sampling_step_cfgpp(x, model, sigma_hat, **extra_args)
+            if i + 1 // 2 == 0:     #   ??  this is i == 0
+                x = smea_sampling_step_cfgpp(x, model, sigma_hat, **extra_args)        
+    return x
+
+@torch.no_grad()
+def sample_euler_ancestral_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+    """Ancestral sampling with Euler method steps."""
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    model.need_last_noise_uncond = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
+        # Euler method
+        x = denoised + d * sigma_down
+        if sigmas[i + 1] > 0:
+            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+    return x

webUI_ExtraSchedulers/scripts/clybius_dpmpp_4m_sde.py

@@ -0,0 +1,124 @@
+# by Clybius : github.com/Clybius/ComfyUI-Extra-Samplers/
+
+import math
+
+import torch
+from torch import nn, FloatTensor
+import torchsde
+import kornia
+from tqdm.auto import trange, tqdm
+import numpy as np
+
+import sample
+
+from k_diffusion.sampling import BrownianTreeNoiseSampler, PIDStepSizeController, get_ancestral_step, to_d, default_noise_sampler, DPMSolver
+
+
+#   copied from kdiffusion/sampling.py and utils.py
+def default_noise_sampler(x):
+    return lambda sigma, sigma_next: torch.randn_like(x)
+
+
+@torch.no_grad()
+def sample_clyb_4m_sde_momentumized(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1., noise_sampler=None, momentum=0.0):
+    """DPM-Solver++(3M) SDE, modified with an extra SDE, and momentumized in both the SDE and ODE(?). 'its a first' - Clybius 2023
+    The expression for d1 is derived from the extrapolation formula given in the paper “Diffusion Monte Carlo with stochastic Hamiltonians” by M. Foulkes, L. Mitas, R. Needs, and G. Rajagopal. The formula is given as follows:
+    d1 = d1_0 + (d1_0 - d1_1) * r2 / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) * r2 / ((r2 + r1) * (r0 + r1))
+    (if this is an incorrect citing, we blame Google's Bard and OpenAI's ChatGPT for this and NOT me :^) )
+
+    where d1_0, d1_1, and d1_2 are defined as follows:
+    d1_0 = (denoised - denoised_1) / r2
+    d1_1 = (denoised_1 - denoised_2) / r1
+    d1_2 = (denoised_2 - denoised_3) / r0
+
+    The variables r0, r1, and r2 are defined as follows:
+    r0 = h_3 / h_2
+    r1 = h_2 / h
+    r2 = h / h_1
+    """
+
+    def momentum_func(diff, velocity, timescale=1.0, offset=-momentum / 2.0): # Diff is current diff, vel is previous diff
+        if velocity is None:
+            momentum_vel = diff
+        else:
+            momentum_vel = momentum * (timescale + offset) * velocity + (1 - momentum * (timescale + offset)) * diff
+        return momentum_vel
+
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    denoised_1, denoised_2, denoised_3 = None, None, None
+    h_1, h_2, h_3 = None, None, None
+    vel, vel_sde = None, None
+    for i in trange(len(sigmas) - 1, disable=disable):
+        time = sigmas[i] / sigma_max
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+
+        if sigmas[i + 1] == 0:
+            # Denoising step
+            x = denoised
+        else:
+            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = s - t
+            h_eta = h * (eta + 1)
+            x_diff = momentum_func((-h_eta).expm1().neg() * denoised, vel, time)
+            vel = x_diff
+            x = torch.exp(-h_eta) * x + vel
+
+            if h_3 is not None:
+                r0 = h_1 / h
+                r1 = h_2 / h
+                r2 = h_3 / h
+                d1_0 = (denoised   - denoised_1) / r0
+                d1_1 = (denoised_1 - denoised_2) / r1
+                d1_2 = (denoised_2 - denoised_3) / r2
+                # d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1) + ((d1_0 - d1_1) * r2 / (r1 + r2) - (d1_1 - d1_2) * r1 / (r0 + r1)) * r2 / ((r1 + r2) * (r0 + r1))
+                # d2 = (d1_0 - d1_1) / (r0 + r1) + ((d1_0 - d1_1) * r2 / (r1 + r2) - (d1_1 - d1_2) * r1 / (r0 + r1)) / ((r1 + r2) * (r0 + r1))
+
+                # r0 = h_3 / h_2
+                # r1 = h_2 / h
+                # r2 = h / h_1
+                # d1_0 = (denoised - denoised_1) / r2
+                # d1_1 = (denoised_1 - denoised_2) / r1
+                # d1_2 = (denoised_2 - denoised_3) / r0
+                d1 = d1_0 + (d1_0 - d1_1) * r2 / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) * r2 / ((r2 + r1) * (r0 + r1))
+                d2 = (d1_0 - d1_1) / (r2 + r1) + ((d1_0 - d1_1) * r2 / (r2 + r1) - (d1_1 - d1_2) * r1 / (r0 + r1)) / ((r2 + r1) * (r0 + r1))
+                phi_3 = h_eta.neg().expm1() / h_eta + 1
+                phi_4 = phi_3 / h_eta - 0.5
+                sde_diff = momentum_func(phi_3 * d1 - phi_4 * d2, vel_sde, time)
+                vel_sde = sde_diff
+                x = x + vel_sde
+            elif h_2 is not None:
+                r0 = h_1 / h
+                r1 = h_2 / h
+                d1_0 = (denoised - denoised_1) / r0
+                d1_1 = (denoised_1 - denoised_2) / r1
+                d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
+                d2 = (d1_0 - d1_1) / (r0 + r1)
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                phi_3 = phi_2 / h_eta - 0.5
+                sde_diff = momentum_func(phi_2 * d1 - phi_3 * d2, vel_sde, time)
+                vel_sde = sde_diff
+                x = x + vel_sde
+            elif h_1 is not None:
+                r = h_1 / h
+                d = (denoised - denoised_1) / r
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                sde_diff = momentum_func(phi_2 * d, vel_sde, time)
+                vel_sde = sde_diff
+                x = x + vel_sde
+
+            if eta:
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
+
+            denoised_1, denoised_2, denoised_3 = denoised, denoised_1, denoised_2
+            h_1, h_2, h_3 = h, h_1, h_2
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
+    return x

webUI_ExtraSchedulers/scripts/extra_schedulers.py

@@ -0,0 +1,432 @@
+import gradio
+import math, numpy
+import torch
+from modules import scripts, shared
+
+# Python 3.10+, PyTorch 2.1+, NumPy 1.24+
+def get_sigmas_oss(n, sigma_min, sigma_max, device):
+    """
+    Optimal Steps schedule (OSS).
+    Исправлено:
+      - Больше нет обращения к несуществующей переменной `sigmas`.
+      - Всегда возвращаем тензор float32 на переданном `device`.
+      - Порядок веток по типу модели: SD3/Flux → SDXL → общий (SD1/2).
+    Примечание: пресеты подобраны «в абсолютных» единицах под семейства моделей,
+    поэтому sigma_min/sigma_max здесь намеренно не используются.
+    """
+    import numpy
+    import torch
+    from modules import shared
+
+    def loglinear_interp(values: list[float], num_steps: int) -> numpy.ndarray:
+        """Лог-линейная интерполяция убывающей последовательности до num_steps."""
+        arr = numpy.asarray(values, dtype=float)
+        xs = numpy.linspace(0.0, 1.0, arr.shape[0])
+        ys = numpy.log(arr[::-1])                 # в возрастающую + логарифм
+        new_xs = numpy.linspace(0.0, 1.0, num_steps)
+        new_ys = numpy.interp(new_xs, xs, ys)     # интерполяция в лог-пространстве
+        out = numpy.exp(new_ys)[::-1].copy()      # обратно и снова убывающая
+        return out
+
+    m = shared.sd_model
+
+    # 1) Флоу-семейство (SD3/Flux) — нормализованный пресет ~[1..0]
+    if getattr(m, "is_sd3", False) or getattr(m, "is_flux", False):
+        base_sigmas = [0.9968, 0.9886, 0.9819, 0.975, 0.966, 0.9471, 0.9158, 0.8287, 0.5512, 0.2808, 0.001]
+
+    # 2) SDXL — свой AYS11 пресет
+    elif getattr(m, "is_sdxl", False):
+        base_sigmas = [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.029]
+
+    # 3) SD1.x/SD2.x и прочие — общий AYS11
+    else:
+        base_sigmas = [14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.029]
+
+    # Подгоняем длину к n и добавляем терминальный 0.0 (итого n+1 значений)
+    if n != len(base_sigmas):
+        sigmas_np = loglinear_interp(base_sigmas, n)
+        sigmas_np = numpy.append(sigmas_np, [0.0])
+    else:
+        sigmas_np = numpy.asarray(base_sigmas + [0.0], dtype=float)
+
+    # Единый путь возврата: float32 на переданном device
+    return torch.as_tensor(sigmas_np, dtype=torch.float32, device=device)
+
+
+
+    
+def cosine_scheduler (n, sigma_min, sigma_max, device):
+    sigmas = torch.zeros(n, device=device)
+    if n == 1:
+        sigmas[0] = sigma_max ** 0.5
+    else:
+        for x in range(n):
+            p = x / (n-1)
+            C = sigma_min + 0.5*(sigma_max-sigma_min)*(1 - math.cos(math.pi*(1 - p**0.5)))
+            sigmas[x] = C
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+def cosexpblend_boost_scheduler (n, sigma_min, sigma_max, device):
+    sigmas = []
+    if n == 1:
+        sigmas.append(sigma_max ** 0.5)
+    else:
+        K = (sigma_min / sigma_max)**(1/(n-1))
+        E = sigma_max
+        detail = numpy.interp(numpy.linspace(0, 1, n), numpy.linspace(0, 1, 5), [1.0, 1.0, 1.27, 1.0, 1.0])
+        for x in range(n):
+            p = x / (n-1)
+            C = sigma_min + 0.5*(sigma_max-sigma_min)*(1 - math.cos(math.pi*(1 - p**0.5)))
+            sigmas.append(detail[x] * (C + p * (E - C)))
+            E *= K
+
+    sigmas += [0.0]
+    return torch.FloatTensor(sigmas).to(device)
+
+def cosexpblend_scheduler (n, sigma_min, sigma_max, device):
+    sigmas = []
+    if n == 1:
+        sigmas.append(sigma_max ** 0.5)
+    else:
+        K = (sigma_min / sigma_max)**(1/(n-1))
+        E = sigma_max
+        for x in range(n):
+            p = x / (n-1)
+            C = sigma_min + 0.5*(sigma_max-sigma_min)*(1 - math.cos(math.pi*(1 - p**0.5)))
+            sigmas.append(C + p * (E - C))
+            E *= K
+    sigmas += [0.0]
+    return torch.FloatTensor(sigmas).to(device)
+
+##  phi scheduler modified from original by @extraltodeus
+def phi_scheduler(n, sigma_min, sigma_max, device):
+    sigmas = torch.zeros(n, device=device)
+    if n == 1:
+        sigmas[0] = sigma_max ** 0.5
+    else:
+        phi = (1 + 5**0.5) / 2
+        for x in range(n):
+            sigmas[x] = sigma_min + (sigma_max-sigma_min)*((1-x/(n-1))**(phi*phi))
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+def get_sigmas_vp(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs a continuous VP noise schedule."""
+    
+    beta_d = 19.9
+    beta_min = 0.1
+    eps_s = 1e-3
+    
+    t = torch.linspace(1, eps_s, n, device=device)
+    sigmas = torch.sqrt(torch.exp(beta_d * t ** 2 / 2 + beta_min * t) - 1)
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+def get_sigmas_laplace(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs the noise schedule proposed by Tiankai et al. (2024). """
+    mu = 0.
+    beta = 0.5
+    epsilon = 1e-5 # avoid log(0)
+    x = torch.linspace(0, 1, n, device=device)
+    clamp = lambda x: torch.clamp(x, min=sigma_min, max=sigma_max)
+    lmb = mu - beta * torch.sign(0.5-x) * torch.log(1 - 2 * torch.abs(0.5-x) + epsilon)
+    sigmas = clamp(torch.exp(lmb))
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+
+
+def get_sigmas_sinusoidal_sf(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs a sinusoidal noise schedule."""
+    sf = 3.5
+    x = torch.linspace(0, 1, n, device=device)
+    sigmas = (sigma_min + (sigma_max - sigma_min) * (1 - torch.sin(torch.pi / 2 * x)))/sigma_max
+    sigmas = sigmas**sf
+    sigmas = sigmas * sigma_max
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_invcosinusoidal_sf(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs a sinusoidal noise schedule."""
+    sf = 3.5
+    x = torch.linspace(0, 1, n, device=device)
+    sigmas = (sigma_min + (sigma_max - sigma_min) * (0.5*(torch.cos(x * math.pi) + 1)))/sigma_max
+    sigmas = sigmas**sf
+    sigmas = sigmas * sigma_max
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_react_cosinusoidal_dynsf(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs a sinusoidal noise schedule."""
+    sf = 2.15
+    x = torch.linspace(0, 1, n, device=device)
+    sigmas = (sigma_min+(sigma_max-sigma_min)*(torch.cos(x*(torch.pi/2))))/sigma_max
+    sigmas = sigmas**(sf*(n*x/n))
+    sigmas = sigmas * sigma_max
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_karras_dynamic(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    rho = 7.
+    ramp = torch.linspace(0, 1, n, device=device)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = torch.zeros_like(ramp)
+    for i in range(n):
+        sigmas[i] = (max_inv_rho + ramp[i] * (min_inv_rho - max_inv_rho)) ** (math.cos(i*math.tau/n)*2+rho) 
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_karras_exponential_decay(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    rho = 7.
+    ramp = torch.linspace(0, 1, n, device=device)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = torch.zeros_like(ramp)
+    for i in range(n):
+        sigmas[i] = (max_inv_rho + ramp[i] * (min_inv_rho - max_inv_rho)) ** (rho-(3*i/n))
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+ 
+def get_sigmas_karras_exponential_increment(n, sigma_min, sigma_max, device='cpu'):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    rho = 7.
+    ramp = torch.linspace(0, 1, n, device=device)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = torch.zeros_like(ramp)
+    for i in range(n):
+        sigmas[i] = (max_inv_rho + ramp[i] * (min_inv_rho - max_inv_rho)) ** (rho+3*i/n)
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+def custom_scheduler(n, sigma_min, sigma_max, device):
+    if 'import' in ExtraScheduler.customSigmas:
+        sigmas = torch.linspace(sigma_max, sigma_min, n, device=device)
+    elif 'eval' in ExtraScheduler.customSigmas:
+        sigmas = torch.linspace(sigma_max, sigma_min, n, device=device)
+    elif 'scripts' in ExtraScheduler.customSigmas:
+        sigmas = torch.linspace(sigma_max, sigma_min, n, device=device)
+
+    elif ExtraScheduler.customSigmas[0] == '[' and ExtraScheduler.customSigmas[-1] == ']':
+        sigmasList = [float(x) for x in ExtraScheduler.customSigmas.strip('[]').split(',')]
+
+        if sigmasList[0] == 1.0 and sigmasList[-1] == 0.0:
+            for x in range(len(sigmasList)):
+                sigmasList[x] *= (sigma_max - sigma_min)
+                sigmasList[x] += sigma_min
+        elif sigmasList[-1] == 0.0:
+            #don't interpolate to number of steps, use as is
+            return torch.tensor(sigmasList)
+
+        xs = numpy.linspace(0, 1, len(sigmasList))
+        ys = numpy.log(sigmasList[::-1])
+        
+        new_xs = numpy.linspace(0, 1, n)
+        new_ys = numpy.interp(new_xs, xs, ys)
+        
+        interpolated_ys = numpy.exp(new_ys)[::-1].copy()
+        sigmas = torch.tensor(interpolated_ys, device=device)
+    else:
+        sigmas = torch.linspace(sigma_max, sigma_min, n, device=device)
+        detail = numpy.interp(numpy.linspace(0, 1, n), numpy.linspace(0, 1, 5), [1.0, 1.0, 1.25, 1.0, 1.0])
+        
+        phi = (1 + 5**0.5) / 2
+        pi = math.pi
+        
+        s = 0
+        while (s < n):
+            x = (s) / (n - 1)
+            M = sigma_max
+            m = sigma_min
+            d = detail[s]
+        
+            sigmas[s] = eval((ExtraScheduler.customSigmas))
+            s += 1
+    return torch.cat([sigmas, sigmas.new_zeros([1])])
+
+from scripts.simple_kes import get_sigmas_simple_kes
+
+from scripts.res_solver import sample_res_solver, sample_res_multistep, sample_res_multistep_cfgpp
+from scripts.clybius_dpmpp_4m_sde import sample_clyb_4m_sde_momentumized
+from scripts.gradient_estimation import sample_gradient_e, sample_gradient_e_cfgpp
+from scripts.seeds import sample_seeds_2, sample_seeds_3
+
+from modules import sd_samplers_common, sd_samplers
+from modules.sd_samplers_kdiffusion import sampler_extra_params, KDiffusionSampler
+
+class ExtraScheduler(scripts.Script):
+    sorting_priority = 99
+
+    installed = False
+    customSigmas = 'm + (M-m)*(1-x)**3'
+
+    def title(self):
+        return "Extra Schedulers (custom)"
+
+    def show(self, is_img2img):
+        # make this extension visible in both txt2img and img2img tab.
+        if ExtraScheduler.installed:
+            return scripts.AlwaysVisible
+        else:
+            return False
+
+    def ui(self, *args, **kwargs):
+        #with gradio.Accordion(open=False, label=self.title(), visible=ExtraScheduler.installed):
+        custom_sigmas = gradio.Textbox(value=ExtraScheduler.customSigmas, label='Extra Schedulers: custom function / list [n0, n1, n2, ...]', lines=1.01)
+
+        self.infotext_fields = [
+            (custom_sigmas, "es_custom"),
+        ]
+
+        return [custom_sigmas]
+
+    def process(self, params, *script_args, **kwargs):
+        if params.scheduler == 'custom':
+            custom_sigmas = script_args[0]
+            ExtraScheduler.customSigmas = custom_sigmas
+            params.extra_generation_params.update(dict(es_custom = ExtraScheduler.customSigmas, ))
+        elif params.scheduler == 'Simple KES':
+            params.extra_generation_params.update(dict(
+                es_KES_start_blend       = getattr(shared.opts, 'kes_start_blend'),
+                es_KES_end_blend         = getattr(shared.opts, 'kes_end_blend'),
+                es_KES_sharpness         = getattr(shared.opts, 'kes_sharpness'),
+                es_KES_initial_step_size = getattr(shared.opts, 'kes_initial_step_size'),
+                es_KES_final_step_size   = getattr(shared.opts, 'kes_final_step_size'),
+                es_KES_initial_noise     = getattr(shared.opts, 'kes_initial_noise'),
+                es_KES_final_noise       = getattr(shared.opts, 'kes_final_noise'),
+                es_KES_smooth_blend      = getattr(shared.opts, 'kes_smooth_blend'),
+                es_KES_step_size_factor  = getattr(shared.opts, 'kes_step_size_factor'),
+                es_KES_noise_scale       = getattr(shared.opts, 'kes_noise_scale'),
+            ))
+        return
+
+try:
+    import modules.sd_schedulers as schedulers
+
+    if True:
+        # убираем уже зарегистрированные версии с тем же именем/лейблом
+        def _drop(name=None, label=None):
+            schedulers.schedulers = [
+                s for s in getattr(schedulers, "schedulers", [])
+                if (name is not None and getattr(s, "name", None) == name) is False
+                and (label is not None and getattr(s, "label", None) == label) is False
+            ]
+
+        _drop(name="optimal_ss");      _drop(label="Optimal Steps")
+        _drop(name="custom");          _drop(label="custom")
+
+        print("Extension: Extra Schedulers: (re)adding schedulers")
+
+        # далее — как у вас: создаём объекты Scheduler(...)
+
+        print("Extension: Extra Schedulers: adding new schedulers")
+        CosineScheduler         = schedulers.Scheduler("cosine",        "Cosine",                   cosine_scheduler)
+        CosExpScheduler         = schedulers.Scheduler("cosexp",        "CosineExponential blend",  cosexpblend_scheduler)
+        CosExpBScheduler        = schedulers.Scheduler("cosprev",       "CosExp blend boost",       cosexpblend_boost_scheduler)
+        PhiScheduler            = schedulers.Scheduler("phi",           "Phi",                      phi_scheduler)
+        VPScheduler             = schedulers.Scheduler("vp",            "VP",                       get_sigmas_vp)
+        LaplaceScheduler        = schedulers.Scheduler("laplace",       "Laplace",                  get_sigmas_laplace)
+
+        SineScheduler           = schedulers.Scheduler("sine_sc",       "Sine scaled",              get_sigmas_sinusoidal_sf)
+        InvCosScheduler         = schedulers.Scheduler("inv_cos_sc",    "Inverse Cosine scaled",    get_sigmas_invcosinusoidal_sf)
+        CosDynScheduler         = schedulers.Scheduler("cosine_dyn",    "Cosine Dynamic",           get_sigmas_react_cosinusoidal_dynsf)
+        KarrasDynScheduler      = schedulers.Scheduler("karras_dyn",    "Karras Dynamic",           get_sigmas_karras_dynamic)
+        KarrasExpDecayScheduler = schedulers.Scheduler("karras_exp_d",  "Karras Exp Decay",         get_sigmas_karras_exponential_decay)
+        KarrasExpIncScheduler   = schedulers.Scheduler("karras_exp_i",  "Karras Exp Inc",           get_sigmas_karras_exponential_increment)
+
+        SimpleKEScheduler       = schedulers.Scheduler("simple_kes",    "Simple KES",               get_sigmas_simple_kes)
+        OSSFlowScheduler        = schedulers.Scheduler("optimal_ss",    "Optimal Steps",            get_sigmas_oss)
+        CustomScheduler         = schedulers.Scheduler("custom",        "custom",                   custom_scheduler)
+
+        schedulers.schedulers.append(CosineScheduler)
+        schedulers.schedulers.append(CosExpScheduler)
+        schedulers.schedulers.append(CosExpBScheduler)
+        schedulers.schedulers.append(PhiScheduler)
+        schedulers.schedulers.append(VPScheduler)
+        schedulers.schedulers.append(LaplaceScheduler)
+        schedulers.schedulers.append(SineScheduler)
+        schedulers.schedulers.append(InvCosScheduler)
+        schedulers.schedulers.append(CosDynScheduler)
+        schedulers.schedulers.append(KarrasDynScheduler)
+        schedulers.schedulers.append(KarrasExpDecayScheduler)
+        schedulers.schedulers.append(KarrasExpIncScheduler)
+        schedulers.schedulers.append(SimpleKEScheduler)
+        schedulers.schedulers.append(OSSFlowScheduler)
+        schedulers.schedulers.append(CustomScheduler)
+
+        schedulers.schedulers_map = {
+            **{x.name: x for x in schedulers.schedulers},
+            **{x.label: x for x in schedulers.schedulers}
+        }
+
+        # CFG++ method is Forge only, not working in A1111
+        from modules import sd_samplers_common, sd_samplers
+        from modules.sd_samplers_kdiffusion import sampler_extra_params, KDiffusionSampler
+        from scripts.samplers_cfgpp import (
+            sample_euler_ancestral_cfgpp, sample_euler_cfgpp, sample_euler_dy_cfgpp,
+            sample_euler_smea_dy_cfgpp, sample_euler_negative_cfgpp, sample_euler_negative_dy_cfgpp
+        )
+        from scripts.forgeClassic_cfgpp import (
+            sample_dpmpp_sde_cfgpp, sample_dpmpp_2m_cfgpp,
+            sample_dpmpp_2m_sde_cfgpp, sample_dpmpp_3m_sde_cfgpp,
+            sample_dpmpp_2s_ancestral_cfgpp
+        )
+
+        samplers_cfgpp = [
+            ("Euler a CFG++",           sample_euler_ancestral_cfgpp,   ["k_euler_a_cfgpp"],            {"uses_ensd": True}),
+            ("Euler CFG++",             sample_euler_cfgpp,             ["k_euler_cfgpp"],              {}),
+            ("Euler Dy CFG++",          sample_euler_dy_cfgpp,          ["k_euler_dy_cfgpp"],           {}),
+            ("Euler SMEA Dy CFG++",     sample_euler_smea_dy_cfgpp,     ["k_euler_smea_dy_cfgpp"],      {}),
+            ("Euler Negative CFG++",    sample_euler_negative_cfgpp,    ["k_euler_negative_cfgpp"],     {}),
+            ("Euler Negative Dy CFG++", sample_euler_negative_dy_cfgpp, ["k_euler_negative_dy_cfgpp"],  {}),
+            ("RES multistep CFG++",     sample_res_multistep_cfgpp,     ["k_res_multi_cfgpp"],          {}),
+            ("Gradient Estimation CFG++", sample_gradient_e_cfgpp,      ["k_grad_est_cfgpp"],           {}),
+            # ("GE/DPM2 CFG++",           sample_ge_dpm2_cfgpp,           ["k_ge_dpm_cfgpp"],             {}),
+            ("DPM++ SDE CFG++",         sample_dpmpp_sde_cfgpp,         ["k_dpmpp_sde_cfgpp"],          {"brownian_noise": True, "second_order": True}),
+            ("DPM++ 2M CFG++",          sample_dpmpp_2m_cfgpp,          ["k_dpmpp_2m_cfgpp"],           {}),
+            ("DPM++ 2M SDE CFG++",      sample_dpmpp_2m_sde_cfgpp,      ["k_dpmpp_2m_sde_cfgpp"],       {"brownian_noise": True}),
+            ("DPM++ 3M SDE CFG++",      sample_dpmpp_3m_sde_cfgpp,      ["k_dpmpp_3m_sde_cfgpp"],       {"brownian_noise": True, 'discard_next_to_last_sigma': True}),
+            ("DPM++ 2S a CFG++",        sample_dpmpp_2s_ancestral_cfgpp,["k_dpmpp_2s_a_cfgpp"],         {"uses_ensd": True, "second_order": True}),
+        ]
+
+        samplers_data_cfgpp = [
+            sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options)
+            for label, funcname, aliases, options in samplers_cfgpp
+            if callable(funcname)
+        ]
+
+        sampler_extra_params['sample_euler_cfgpp']             = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+        sampler_extra_params['sample_euler_negative_cfgpp']    = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+        sampler_extra_params['sample_euler_dy_cfgpp']          = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+        sampler_extra_params['sample_euler_negative_dy_cfgpp'] = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+        sampler_extra_params['sample_euler_smea_dy_cfgpp']     = ['s_churn', 's_tmin', 's_tmax', 's_noise']
+        
+        sampler_extra_params['sample_dpmpp_sde_cfgpp']         = ['s_noise']
+        sampler_extra_params['sample_dpmpp_2m_sde_cfgpp']      = ['s_noise']
+        sampler_extra_params['sample_dpmpp_3m_sde_cfgpp']      = ['s_noise']
+        sampler_extra_params['sample_dpmpp_2s_ancestral_cfgpp']= ['s_noise']
+
+        sd_samplers.all_samplers.extend(samplers_data_cfgpp)
+        #except:
+            #pass
+
+        samplers_extra = [
+            ("RES multistep",                sample_res_multistep,              ["k_res_multi"],        {}),
+            ("Refined Exponential Solver",   sample_res_solver,                 ["k_res"],              {}),
+            ("DPM++ 4M SDE",                 sample_clyb_4m_sde_momentumized,   ["k_dpmpp_4m_sde"],     {}),
+            ("Gradient Estimation",          sample_gradient_e,                 ["k_grad_est"],         {}),
+            ("SEEDS-2",                      sample_seeds_2,                    ["k_seeds2"],         {}),
+            ("SEEDS-3",                      sample_seeds_3,                    ["k_seeds3"],         {}),
+        ]
+        sampler_extra_params['sample_seeds_2']         = ['s_noise']
+        sampler_extra_params['sample_seeds_3']         = ['s_noise']
+        
+        samplers_data_extra = [
+            sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options)
+            for label, funcname, aliases, options in samplers_extra
+            if callable(funcname)
+        ]
+
+        sd_samplers.all_samplers.extend(samplers_data_extra)
+        sd_samplers.all_samplers_map = {x.name: x for x in sd_samplers.all_samplers}
+        sd_samplers.set_samplers()
+
+
+    ExtraScheduler.installed = True
+except:
+    print ("Extension: Extra Schedulers: unsupported webUI")
+    ExtraScheduler.installed = False

webUI_ExtraSchedulers/scripts/forgeClassic_cfgpp.py

@@ -0,0 +1,269 @@
+# first 3 lifted from ForgeClassic (https://github.com/Haoming02/sd-webui-forge-classic/)
+# 4th is simple adaptation of 3M to 2M
+# 5th lifted from ReForge (https://github.com/Panchovix/stable-diffusion-webui-reForge)
+# all modified to work with Forge2
+
+import torch
+from tqdm.auto import trange
+from k_diffusion.sampling import (
+    default_noise_sampler,
+    BrownianTreeNoiseSampler,
+    get_ancestral_step,
+    to_d,
+)
+
+
+def _sigma_fn(t):
+    return t.neg().exp()
+
+
+def _t_fn(sigma):
+    return sigma.log().neg()
+
+
+@torch.no_grad()
+def sample_dpmpp_sde_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
+    eta = 1.0
+    s_noise = 1.0
+    r = 0.5
+
+    if len(sigmas) <= 1:
+        return x
+
+    seed = extra_args.get("seed", None)
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback(
+                {
+                    "x": x,
+                    "i": i,
+                    "sigma": sigmas[i],
+                    "sigma_hat": sigmas[i],
+                    "denoised": denoised,
+                }
+            )
+
+        if sigmas[i + 1] == 0:
+            d = model.last_noise_uncond
+            x = denoised + d * sigmas[i + 1]
+        else:
+            t, t_next = _t_fn(sigmas[i]), _t_fn(sigmas[i + 1])
+            h = t_next - t
+            s = t + h * r
+            fac = 1 / (2 * r)
+
+
+            sd, su = get_ancestral_step(_sigma_fn(t), _sigma_fn(s), eta)
+            s_ = _t_fn(sd)
+            x_2 = (_sigma_fn(s_) / _sigma_fn(t)) * x - (t - s_).expm1() * denoised
+            x_2 = x_2 + noise_sampler(_sigma_fn(t), _sigma_fn(s)) * s_noise * su
+            denoised_2 = model(x_2, _sigma_fn(s) * s_in, **extra_args)
+            u = x_2 - model.last_noise_uncond * _sigma_fn(s) * s_in[:1]
+
+            sd, su = get_ancestral_step(_sigma_fn(t), _sigma_fn(t_next), eta)
+            denoised_d = (1 - fac) * u + fac * u
+            x = denoised_2 + to_d(x, sigmas[i], denoised_d) * sd
+            x = x + noise_sampler(_sigma_fn(t), _sigma_fn(t_next)) * s_noise * su
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_2m_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None):
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    old_uncond_denoised = None
+    uncond_denoised = None
+
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        uncond_denoised = x - model.last_noise_uncond * sigmas[i] * s_in[:1]
+        if callback is not None:
+            callback(
+                {
+                    "x": x,
+                    "i": i,
+                    "sigma": sigmas[i],
+                    "sigma_hat": sigmas[i],
+                    "denoised": denoised,
+                }
+            )
+        t, t_next = _t_fn(sigmas[i]), _t_fn(sigmas[i + 1])
+        h = t_next - t
+        if old_uncond_denoised is None or sigmas[i + 1] == 0:
+            denoised_mix = -torch.exp(-h) * uncond_denoised
+        else:
+            h_last = t - _t_fn(sigmas[i - 1])
+            r = h_last / h
+            denoised_mix = -torch.exp(-h) * uncond_denoised - torch.expm1(-h) * (1 / (2 * r)) * (denoised - old_uncond_denoised)
+        x = denoised + denoised_mix + torch.exp(-h) * x
+        old_uncond_denoised = uncond_denoised
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_3m_sde_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=None, s_noise=None, noise_sampler=None):
+    eta = 1.0 if eta is None else eta
+    s_noise = 1.0 if s_noise is None else s_noise
+
+    if len(sigmas) <= 1:
+        return x
+
+    seed = extra_args.get("seed", None)
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    denoised_1, denoised_2 = None, None
+    h, h_1, h_2 = None, None, None
+
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        u = x - model.last_noise_uncond * sigmas[i] * s_in[:1]
+        if callback is not None:
+            callback(
+                {
+                    "x": x,
+                    "i": i,
+                    "sigma": sigmas[i],
+                    "sigma_hat": sigmas[i],
+                    "denoised": denoised,
+                }
+            )
+        if sigmas[i + 1] == 0:
+            x = denoised
+        else:
+            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = s - t
+            h_eta = h * (eta + 1)
+
+            x = torch.exp(-h_eta) * (x + (denoised - u)) + (-h_eta).expm1().neg() * denoised
+
+            if h_2 is not None:
+                r0 = h_1 / h
+                r1 = h_2 / h
+                d1_0 = (denoised - denoised_1) / r0
+                d1_1 = (denoised_1 - denoised_2) / r1
+                d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
+                d2 = (d1_0 - d1_1) / (r0 + r1)
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                phi_3 = phi_2 / h_eta - 0.5
+                x = x + phi_2 * d1 - phi_3 * d2
+            elif h_1 is not None:
+                r = h_1 / h
+                d = (denoised - denoised_1) / r
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                x = x + phi_2 * d
+
+            if eta:
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
+
+        denoised_1, denoised_2 = denoised, denoised_1
+        h_1, h_2 = h, h_1
+    return x
+
+
+## extra
+@torch.no_grad()
+def sample_dpmpp_2m_sde_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+# just cut down from 3m_sde version
+    seed = extra_args.get("seed", None)
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    denoised_1 = None
+    h_1 = None
+
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        u = x - model.last_noise_uncond * sigmas[i] * s_in[:1]
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            #Denoising step
+            x = denoised
+        else:
+            #DPM-Solver++(2M) SDE
+            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = s - t
+
+            h_eta = h * (eta + 1)
+            x = torch.exp(-h_eta) * (x + (denoised - u)) + (-h_eta).expm1().neg() * denoised
+
+            if denoised_1 is not None:
+                r = h_1 / h
+
+                d = (denoised - denoised_1) / r
+                phi_2 = h_eta.neg().expm1() / h_eta + 1
+                x = x + phi_2 * d
+
+            if eta:
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
+
+            h_1 = h
+
+        denoised_1 = denoised
+    return x
+
+
+# via ReForge
+@torch.no_grad()
+def sample_dpmpp_2s_ancestral_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigma_down == 0:
+            # Euler method
+            d = model.last_noise_uncond 
+            dt = sigma_down - sigmas[i]
+            x = denoised + d * sigma_down
+        else:
+            u = x - model.last_noise_uncond * sigmas[i] * s_in[:1]
+
+            # DPM-Solver++(2S)
+            t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
+            # r = torch.sinh(1 + (2 - eta) * (t_next - t) / (t - t_fn(sigma_up))) works only on non-cfgpp, weird
+            r = 1 / 2
+            h = t_next - t
+            s = t + r * h
+            x_2 = (sigma_fn(s) / sigma_fn(t)) * (x + (denoised - u)) - (-h * r).expm1() * denoised
+            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
+            x = (sigma_fn(t_next) / sigma_fn(t)) * (x + (denoised - u)) - (-h).expm1() * denoised_2
+
+        # Noise addition
+        if sigmas[i + 1] > 0:
+            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+    return x

webUI_ExtraSchedulers/scripts/gradient_estimation.py

@@ -0,0 +1,71 @@
+## lifted from ReForge, original implementation from Comfy
+## CFG++ attempt by me
+
+import torch
+from tqdm.auto import trange
+
+
+#   copied from kdiffusion/sampling.py
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+def append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
+    return x[(...,) + (None,) * dims_to_append]
+
+
+@torch.no_grad()
+def sample_gradient_e(model, x, sigmas, extra_args=None, callback=None, disable=None, ge_gamma=2.):
+    """Gradient-estimation sampler. Paper: https://openreview.net/pdf?id=o2ND9v0CeK"""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    old_d = None
+
+    sigmas = sigmas.to(x.device)
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        
+        d = to_d(x, sigmas[i], denoised)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        dt = sigmas[i + 1] - sigmas[i]
+        if i == 0: # Euler method
+            x = x + d * dt
+        else:
+            # Gradient estimation
+            d_bar = ge_gamma * d + (1 - ge_gamma) * old_d
+            x = x + d_bar * dt
+        old_d = d
+    return x
+
+
+@torch.no_grad()
+def sample_gradient_e_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, ge_gamma=2.):
+    """Gradient-estimation sampler. Paper: https://openreview.net/pdf?id=o2ND9v0CeK"""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    old_d = None
+    
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        dt = sigmas[i + 1] - sigmas[i]
+        if i == 0: # Euler method
+            x = denoised + d * sigmas[i+1]
+        else:
+            # Gradient estimation
+            d_bar = ge_gamma * d + (1 - ge_gamma) * old_d
+            x = denoised + d_bar * sigmas[i+1]
+        old_d = d
+    return x

webUI_ExtraSchedulers/scripts/res_solver.py

@@ -0,0 +1,398 @@
+import torch
+from torch import no_grad, FloatTensor
+from tqdm import tqdm
+from itertools import pairwise
+from typing import Protocol, Optional, Dict, Any, TypedDict, NamedTuple, Union, List
+import math
+
+from tqdm.auto import trange
+
+#   copied from kdiffusion/sampling.py and utils.py
+def default_noise_sampler(x):
+    return lambda sigma, sigma_next: torch.randn_like(x)
+def append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
+    return x[(...,) + (None,) * dims_to_append]
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+
+
+class DenoiserModel(Protocol):
+  def __call__(self, x: FloatTensor, t: FloatTensor, *args, **kwargs) -> FloatTensor: ...
+
+class RefinedExpCallbackPayload(TypedDict):
+  x: FloatTensor
+  i: int
+  sigma: FloatTensor
+  sigma_hat: FloatTensor
+
+class RefinedExpCallback(Protocol):
+  def __call__(self, payload: RefinedExpCallbackPayload) -> None: ...
+
+class NoiseSampler(Protocol):
+  def __call__(self, x: FloatTensor) -> FloatTensor: ...
+
+class StepOutput(NamedTuple):
+  x_next: FloatTensor
+  denoised: FloatTensor
+  denoised2: FloatTensor
+  vel: FloatTensor
+  vel_2: FloatTensor
+
+def _gamma(
+  n: int,
+) -> int:
+  """
+  https://en.wikipedia.org/wiki/Gamma_function
+  for every positive integer n,
+  Γ(n) = (n-1)!
+  """
+  return math.factorial(n-1)
+
+def _incomplete_gamma(
+  s: int,
+  x: float,
+  gamma_s: Optional[int] = None
+) -> float:
+  """
+  https://en.wikipedia.org/wiki/Incomplete_gamma_function#Special_values
+  if s is a positive integer,
+  Γ(s, x) = (s-1)!*∑{k=0..s-1}(x^k/k!)
+  """
+  if gamma_s is None:
+    gamma_s = _gamma(s)
+
+  sum_: float = 0
+  # {k=0..s-1} inclusive
+  for k in range(s):
+    numerator: float = x**k
+    denom: int = math.factorial(k)
+    quotient: float = numerator/denom
+    sum_ += quotient
+  incomplete_gamma_: float = sum_ * math.exp(-x) * gamma_s
+  return incomplete_gamma_
+
+# by Katherine Crowson
+def _phi_1(neg_h: FloatTensor):
+  return torch.nan_to_num(torch.expm1(neg_h) / neg_h, nan=1.0)
+
+# by Katherine Crowson
+def _phi_2(neg_h: FloatTensor):
+  return torch.nan_to_num((torch.expm1(neg_h) - neg_h) / neg_h**2, nan=0.5)
+
+# by Katherine Crowson
+def _phi_3(neg_h: FloatTensor):
+  return torch.nan_to_num((torch.expm1(neg_h) - neg_h - neg_h**2 / 2) / neg_h**3, nan=1 / 6)
+
+def _phi(
+  neg_h: float,
+  j: int,
+):
+  """
+  For j={1,2,3}: you could alternatively use Kat's phi_1, phi_2, phi_3 which perform fewer steps
+
+  Lemma 1
+  https://arxiv.org/abs/2308.02157
+  ϕj(-h) = 1/h^j*∫{0..h}(e^(τ-h)*(τ^(j-1))/((j-1)!)dτ)
+
+  https://www.wolframalpha.com/input?i=integrate+e%5E%28%CF%84-h%29*%28%CF%84%5E%28j-1%29%2F%28j-1%29%21%29d%CF%84
+  = 1/h^j*[(e^(-h)*(-τ)^(-j)*τ(j))/((j-1)!)]{0..h}
+  https://www.wolframalpha.com/input?i=integrate+e%5E%28%CF%84-h%29*%28%CF%84%5E%28j-1%29%2F%28j-1%29%21%29d%CF%84+between+0+and+h
+  = 1/h^j*((e^(-h)*(-h)^(-j)*h^j*(Γ(j)-Γ(j,-h)))/(j-1)!)
+  = (e^(-h)*(-h)^(-j)*h^j*(Γ(j)-Γ(j,-h))/((j-1)!*h^j)
+  = (e^(-h)*(-h)^(-j)*(Γ(j)-Γ(j,-h))/(j-1)!
+  = (e^(-h)*(-h)^(-j)*(Γ(j)-Γ(j,-h))/Γ(j)
+  = (e^(-h)*(-h)^(-j)*(1-Γ(j,-h)/Γ(j))
+
+  requires j>0
+  """
+  assert j > 0
+  gamma_: float = _gamma(j)
+  incomp_gamma_: float = _incomplete_gamma(j, neg_h, gamma_s=gamma_)
+
+  phi_: float = math.exp(neg_h) * neg_h**-j * (1-incomp_gamma_/gamma_)
+
+  return phi_
+
+class RESDECoeffsSecondOrder(NamedTuple):
+  a2_1: float
+  b1: float
+  b2: float
+
+def _de_second_order(
+  h: float,
+  c2: float,
+  simple_phi_calc = False,
+) -> RESDECoeffsSecondOrder:
+  """
+  Table 3
+  https://arxiv.org/abs/2308.02157
+  ϕi,j := ϕi,j(-h) = ϕi(-cj*h)
+  a2_1 = c2ϕ1,2
+       = c2ϕ1(-c2*h)
+  b1 = ϕ1 - ϕ2/c2
+  """
+  if simple_phi_calc:
+    # Kat computed simpler expressions for phi for cases j={1,2,3}
+    a2_1: float = c2 * _phi_1(-c2*h)
+    phi1: float = _phi_1(-h)
+    phi2: float = _phi_2(-h)
+  else:
+    # I computed general solution instead.
+    # they're close, but there are slight differences. not sure which would be more prone to numerical error.
+    a2_1: float = c2 * _phi(j=1, neg_h=-c2*h)
+    phi1: float = _phi(j=1, neg_h=-h)
+    phi2: float = _phi(j=2, neg_h=-h)
+  phi2_c2: float = phi2/c2
+  b1: float = phi1 - phi2_c2
+  b2: float = phi2_c2
+  return RESDECoeffsSecondOrder(
+    a2_1=a2_1,
+    b1=b1,
+    b2=b2,
+  )  
+
+def _refined_exp_sosu_step(
+  model: DenoiserModel,
+  x: FloatTensor,
+  sigma: FloatTensor,
+  sigma_next: FloatTensor,
+  c2 = 0.5,
+  extra_args: Dict[str, Any] = {},
+  pbar: Optional[tqdm] = None,
+  simple_phi_calc = False,
+  momentum = 0.0,
+  vel = None,
+  vel_2 = None,
+  time = None
+) -> StepOutput:
+  """
+  Algorithm 1 "RES Second order Single Update Step with c2"
+  https://arxiv.org/abs/2308.02157
+
+  Parameters:
+    model (`DenoiserModel`): a k-diffusion wrapped denoiser model (e.g. a subclass of DiscreteEpsDDPMDenoiser)
+    x (`FloatTensor`): noised latents (or RGB I suppose), e.g. torch.randn((B, C, H, W)) * sigma[0]
+    sigma (`FloatTensor`): timestep to denoise
+    sigma_next (`FloatTensor`): timestep+1 to denoise
+    c2 (`float`, *optional*, defaults to .5): partial step size for solving ODE. .5 = midpoint method
+    extra_args (`Dict[str, Any]`, *optional*, defaults to `{}`): kwargs to pass to `model#__call__()`
+    pbar (`tqdm`, *optional*, defaults to `None`): progress bar to update after each model call
+    simple_phi_calc (`bool`, *optional*, defaults to `True`): True = calculate phi_i,j(-h) via simplified formulae specific to j={1,2}. False = Use general solution that works for any j. Mathematically equivalent, but could be numeric differences.
+  """
+
+  def momentum_func(diff, velocity, timescale=1.0, offset=-momentum / 2.0): # Diff is current diff, vel is previous diff
+    if velocity is None:
+        momentum_vel = diff
+    else:
+        momentum_vel = momentum * (timescale + offset) * velocity + (1 - momentum * (timescale + offset)) * diff
+    return momentum_vel
+
+  lam_next, lam = (s.log().neg() for s in (sigma_next, sigma))
+
+  # type hints aren't strictly true regarding float vs FloatTensor.
+  # everything gets promoted to `FloatTensor` after interacting with `sigma: FloatTensor`.
+  # I will use float to indicate any variables which are scalars.
+  h: float = lam_next - lam
+  a2_1, b1, b2 = _de_second_order(h=h, c2=c2, simple_phi_calc=simple_phi_calc)
+  
+  denoised: FloatTensor = model(x, sigma.repeat(x.size(0)), **extra_args)
+  # if pbar is not None:
+    # pbar.update(0.5)
+
+  c2_h: float = c2*h
+
+  diff_2 = momentum_func(a2_1*h*denoised, vel_2, time)
+  vel_2 = diff_2
+  x_2: FloatTensor = math.exp(-c2_h)*x + diff_2
+  lam_2: float = lam + c2_h
+  sigma_2: float = lam_2.neg().exp()
+
+  denoised2: FloatTensor = model(x_2, sigma_2.repeat(x_2.size(0)), **extra_args)
+  if pbar is not None:
+    pbar.update()
+
+  diff = momentum_func(h*(b1*denoised + b2*denoised2), vel, time)
+  vel = diff
+
+  x_next: FloatTensor = math.exp(-h)*x + diff
+  
+  return StepOutput(
+    x_next=x_next,
+    denoised=denoised,
+    denoised2=denoised2,
+    vel=vel,
+    vel_2=vel_2,
+  )
+  
+
+@no_grad()
+def sample_refined_exp_s(
+  model: FloatTensor,
+  x: FloatTensor,
+  sigmas: FloatTensor,
+  denoise_to_zero: bool = True,
+  extra_args: Dict[str, Any] = {},
+  callback: Optional[RefinedExpCallback] = None,
+  disable: Optional[bool] = None,
+  ita: FloatTensor = torch.zeros((1,)),
+  c2 = .5,
+  noise_sampler: NoiseSampler = torch.randn_like,
+  simple_phi_calc = False,
+  momentum = 0.0,
+):
+  """
+  Refined Exponential Solver (S).
+  Algorithm 2 "RES Single-Step Sampler" with Algorithm 1 second-order step
+  https://arxiv.org/abs/2308.02157
+
+  Parameters:
+    model (`DenoiserModel`): a k-diffusion wrapped denoiser model (e.g. a subclass of DiscreteEpsDDPMDenoiser)
+    x (`FloatTensor`): noised latents (or RGB I suppose), e.g. torch.randn((B, C, H, W)) * sigma[0]
+    sigmas (`FloatTensor`): sigmas (ideally an exponential schedule!) e.g. get_sigmas_exponential(n=25, sigma_min=model.sigma_min, sigma_max=model.sigma_max)
+    denoise_to_zero (`bool`, *optional*, defaults to `True`): whether to finish with a first-order step down to 0 (rather than stopping at sigma_min). True = fully denoise image. False = match Algorithm 2 in paper
+    extra_args (`Dict[str, Any]`, *optional*, defaults to `{}`): kwargs to pass to `model#__call__()`
+    callback (`RefinedExpCallback`, *optional*, defaults to `None`): you can supply this callback to see the intermediate denoising results, e.g. to preview each step of the denoising process
+    disable (`bool`, *optional*, defaults to `False`): whether to hide `tqdm`'s progress bar animation from being printed
+    ita (`FloatTensor`, *optional*, defaults to 0.): degree of stochasticity, η, for each timestep. tensor shape must be broadcastable to 1-dimensional tensor with length `len(sigmas) if denoise_to_zero else len(sigmas)-1`. each element should be from 0 to 1.
+         - if used: batch noise doesn't match non-batch
+    c2 (`float`, *optional*, defaults to .5): partial step size for solving ODE. .5 = midpoint method
+    noise_sampler (`NoiseSampler`, *optional*, defaults to `torch.randn_like`): method used for adding noise
+    simple_phi_calc (`bool`, *optional*, defaults to `True`): True = calculate phi_i,j(-h) via simplified formulae specific to j={1,2}. False = Use general solution that works for any j. Mathematically equivalent, but could be numeric differences.
+  """
+  #assert sigmas[-1] == 0
+  device = x.device
+  ita = ita.to(device)
+  sigmas = sigmas.to(device)
+
+  sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+
+  vel, vel_2 = None, None
+  with tqdm(disable=disable, total=len(sigmas)-(1 if denoise_to_zero else 2)) as pbar:
+    for i, (sigma, sigma_next) in enumerate(pairwise(sigmas[:-1].split(1))):
+      time = sigmas[i] / sigma_max
+      if 'sigma' not in locals():
+        sigma = sigmas[i]
+      eps = torch.randn_like(x).float()
+      sigma_hat = sigma * (1 + ita)
+      x_hat = x + (sigma_hat ** 2 - sigma ** 2).sqrt() * eps
+      x_next, denoised, denoised2, vel, vel_2 = _refined_exp_sosu_step(
+        model,
+        x_hat,
+        sigma_hat,
+        sigma_next,
+        c2=c2,
+        extra_args=extra_args,
+        pbar=pbar,
+        simple_phi_calc=simple_phi_calc,
+        momentum = momentum,
+        vel = vel,
+        vel_2 = vel_2,
+        time = time
+      )
+      if callback is not None:
+        payload = RefinedExpCallbackPayload(
+          x=x,
+          i=i,
+          sigma=sigma,
+          sigma_hat=sigma_hat,
+          denoised=denoised,
+          denoised2=denoised2,
+        )
+        callback(payload)
+      x = x_next
+    if denoise_to_zero:
+      eps = torch.randn_like(x).float()
+      sigma_hat = sigma * (1 + ita)
+      x_hat = x + (sigma_hat ** 2 - sigma ** 2).sqrt() * eps
+      x_next: FloatTensor = model(x_hat, sigma.to(x_hat.device).repeat(x_hat.size(0)), **extra_args)
+      pbar.update()
+
+      if callback is not None:
+        payload = RefinedExpCallbackPayload(
+          x=x,
+          i=i,
+          sigma=sigma,
+          sigma_hat=sigma_hat,
+          denoised=denoised,
+          denoised2=denoised2,
+        )
+        callback(payload)
+
+
+      x = x_next
+  return x
+
+# Many thanks to Kat + Birch-San for this wonderful sampler implementation! https://github.com/Birch-san/sdxl-play/commits/res/
+def sample_res_solver(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler_type="gaussian", noise_sampler=None, denoise_to_zero=True, simple_phi_calc=False, c2=0.5, ita=torch.Tensor((0.0,)), momentum=0.0):
+    return sample_refined_exp_s(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler, denoise_to_zero=denoise_to_zero, simple_phi_calc=simple_phi_calc, c2=c2, ita=ita, momentum=momentum)
+
+
+##  modified from ReForge, original implementation ComfyUI
+@torch.no_grad()
+def res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., noise_sampler=None, cfgpp=False):
+    extra_args = {} if extra_args is None else extra_args
+    seed = extra_args.get("seed", None)
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+    phi1_fn = lambda t: torch.expm1(t) / t
+    phi2_fn = lambda t: (phi1_fn(t) - 1.0) / t
+    old_denoised = None
+
+    sigmas = sigmas.to(x.device)
+
+    if cfgpp:
+        model.need_last_noise_uncond = True
+        model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        if s_churn > 0:
+            gamma = min(s_churn / (len(sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.0
+            sigma_hat = sigmas[i] * (gamma + 1)
+        else:
+            gamma = 0
+            sigma_hat = sigmas[i]
+        if gamma > 0:
+            eps = torch.randn_like(x) * s_noise
+            x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+
+        if callback is not None:
+            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigma_hat, "denoised": denoised})
+        if sigmas[i + 1] == 0 or old_denoised is None:
+            # Euler method
+            if cfgpp:
+                d = model.last_noise_uncond
+                x = denoised + d * sigmas[i + 1]
+            else:
+                d = to_d(x, sigma_hat, denoised)
+                dt = sigmas[i + 1] - sigma_hat
+                x = x + d * dt
+        else:
+            # Second order multistep method in https://arxiv.org/pdf/2308.02157
+            t, t_next, t_prev = t_fn(sigmas[i]), t_fn(sigmas[i + 1]), t_fn(sigmas[i - 1])
+            h = t_next - t
+            c2 = (t_prev - t) / h
+            phi1_val, phi2_val = phi1_fn(-h), phi2_fn(-h)
+            b1 = torch.nan_to_num(phi1_val - 1.0 / c2 * phi2_val, nan=0.0)
+            b2 = torch.nan_to_num(1.0 / c2 * phi2_val, nan=0.0)
+            if cfgpp:
+                d = model.last_noise_uncond
+                x = denoised + d * sigma_hat
+
+            x = (sigma_fn(t_next) / sigma_fn(t)) * x + h * (b1 * denoised + b2 * old_denoised)
+        old_denoised = denoised
+    return x
+@torch.no_grad()
+def sample_res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., noise_sampler=None):
+    return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise, noise_sampler=noise_sampler, cfgpp=False)
+@torch.no_grad()
+def sample_res_multistep_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., noise_sampler=None):
+    return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise, noise_sampler=noise_sampler, cfgpp=True)
+

webUI_ExtraSchedulers/scripts/samplers_cfgpp.py

@@ -0,0 +1,264 @@
+import torch
+from tqdm.auto import trange
+
+#   copied from kdiffusion/sampling.py and utils.py
+def default_noise_sampler(x):
+    return lambda sigma, sigma_next: torch.randn_like(x)
+def get_ancestral_step(sigma_from, sigma_to, eta=1.):
+    """Calculates the noise level (sigma_down) to step down to and the amount
+    of noise to add (sigma_up) when doing an ancestral sampling step."""
+    if not eta:
+        return sigma_to, 0.
+    sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
+    sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
+    return sigma_down, sigma_up
+def append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
+    return x[(...,) + (None,) * dims_to_append]
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+
+
+@torch.no_grad()
+def sample_euler_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        x = denoised + d * sigmas[i+1]
+    return x
+
+class _Rescaler:
+    def __init__(self, model, x, mode, **extra_args):
+        self.model = model
+        self.x = x
+        self.mode = mode
+        self.extra_args = extra_args
+        self.init_latent, self.mask, self.nmask = model.init_latent, model.mask, model.nmask
+
+    def __enter__(self):
+        if self.init_latent is not None:
+            self.model.init_latent = torch.nn.functional.interpolate(input=self.init_latent, size=self.x.shape[2:4], mode=self.mode)
+        if self.mask is not None:
+            self.model.mask = torch.nn.functional.interpolate(input=self.mask.unsqueeze(0), size=self.x.shape[2:4], mode=self.mode).squeeze(0)
+        if self.nmask is not None:
+            self.model.nmask = torch.nn.functional.interpolate(input=self.nmask.unsqueeze(0), size=self.x.shape[2:4], mode=self.mode).squeeze(0)
+
+        return self
+
+    def __exit__(self, type, value, traceback):
+        del self.model.init_latent, self.model.mask, self.model.nmask
+        self.model.init_latent, self.model.mask, self.model.nmask = self.init_latent, self.mask, self.nmask
+
+@torch.no_grad()
+def dy_sampling_step_cfgpp(x, model, sigma_hat, **extra_args):
+    original_shape = x.shape
+    batch_size, channels, m, n = original_shape[0], original_shape[1], original_shape[2] // 2, original_shape[3] // 2
+    extra_row = x.shape[2] % 2 == 1
+    extra_col = x.shape[3] % 2 == 1
+
+    if extra_row:
+        extra_row_content = x[:, :, -1:, :]
+        x = x[:, :, :-1, :]
+    if extra_col:
+        extra_col_content = x[:, :, :, -1:]
+        x = x[:, :, :, :-1]
+
+    a_list = x.unfold(2, 2, 2).unfold(3, 2, 2).contiguous().view(batch_size, channels, m * n, 2, 2)
+    c = a_list[:, :, :, 1, 1].view(batch_size, channels, m, n)
+
+    with _Rescaler(model, c, 'nearest-exact', **extra_args) as rescaler:
+        denoised = model(c, sigma_hat * c.new_ones([c.shape[0]]), **rescaler.extra_args)
+    d = model.last_noise_uncond
+    c = denoised + d * sigma_hat
+
+    d_list = c.view(batch_size, channels, m * n, 1, 1)
+    a_list[:, :, :, 1, 1] = d_list[:, :, :, 0, 0]
+    x = a_list.view(batch_size, channels, m, n, 2, 2).permute(0, 1, 2, 4, 3, 5).reshape(batch_size, channels, 2 * m, 2 * n)
+
+    if extra_row or extra_col:
+        x_expanded = torch.zeros(original_shape, dtype=x.dtype, device=x.device)
+        x_expanded[:, :, :2 * m, :2 * n] = x
+        if extra_row:
+            x_expanded[:, :, -1:, :2 * n + 1] = extra_row_content
+        if extra_col:
+            x_expanded[:, :, :2 * m, -1:] = extra_col_content
+        if extra_row and extra_col:
+            x_expanded[:, :, -1:, -1:] = extra_col_content[:, :, -1:, :]
+        x = x_expanded
+
+    return x
+
+@torch.no_grad()
+def smea_sampling_step_cfgpp(x, model, sigma_hat, **extra_args):
+    m, n = x.shape[2], x.shape[3]
+    x = torch.nn.functional.interpolate(input=x, scale_factor=(1.25, 1.25), mode='nearest-exact')
+    with _Rescaler(model, x, 'nearest-exact', **extra_args) as rescaler:
+        denoised = model(x, sigma_hat * x.new_ones([x.shape[0]]), **rescaler.extra_args)
+    d = model.last_noise_uncond
+    x = denoised + d * sigma_hat
+    x = torch.nn.functional.interpolate(input=x, size=(m,n), mode='nearest-exact')
+    return x
+
+
+@torch.no_grad()
+def sample_euler_dy_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """CFG++ version of Euler Dy by KoishiStar."""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        x = denoised + d * sigmas[i+1]
+
+        if sigmas[i + 1] > 0:
+            if i // 2 == 1:
+                x = dy_sampling_step_cfgpp(x, model, sigma_hat, **extra_args)        
+
+    return x
+
+@torch.no_grad()
+def sample_euler_negative_dy_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """CFG++ version of Euler Negative Dy by KoishiStar."""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        if sigmas[i + 1] > 0 and i // 2 == 1:
+            x = -denoised - d * sigmas[i+1]
+        else:
+            x = denoised + d * sigmas[i+1]
+
+        if sigmas[i + 1] > 0:
+            if i // 2 == 1:
+                x = dy_sampling_step_cfgpp(x, model, sigma_hat, **extra_args)        
+
+    return x
+
+@torch.no_grad()
+def sample_euler_negative_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """based on Euler Negative by KoishiStar"""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        if sigmas[i + 1] > 0 and i // 2 == 1:
+            x = -denoised - d * sigmas[i+1]
+        else:
+            x = denoised + d * sigmas[i+1]
+    return x
+
+
+@torch.no_grad()
+def sample_euler_smea_dy_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
+    """CFG++ version of Euler SMEA Dy by KoishiStar."""
+    extra_args = {} if extra_args is None else extra_args
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
+        eps = torch.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+        denoised = model(x, sigma_hat * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+
+        # Euler method
+        x = denoised + d * sigmas[i+1]
+        
+        if sigmas[i + 1] > 0:
+            if i + 1 // 2 == 1:     #   ??  this is i == 1; why not if i // 2 == 1 same as Euler Dy
+                x = dy_sampling_step_cfgpp(x, model, sigma_hat, **extra_args)
+            if i + 1 // 2 == 0:     #   ??  this is i == 0
+                x = smea_sampling_step_cfgpp(x, model, sigma_hat, **extra_args)        
+    return x
+
+@torch.no_grad()
+def sample_euler_ancestral_cfgpp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+    """Ancestral sampling with Euler method steps."""
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    model.need_last_noise_uncond = True
+    model.inner_model.inner_model.forge_objects.unet.model_options["disable_cfg1_optimization"] = True
+    s_in = x.new_ones([x.shape[0]])
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        d = model.last_noise_uncond
+
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
+        # Euler method
+        x = denoised + d * sigma_down
+        if sigmas[i + 1] > 0:
+            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+    return x

webUI_ExtraSchedulers/scripts/seeds.py

@@ -0,0 +1,106 @@
+# SEEDS implementations by chaObserv : https://github.com/comfyanonymous/ComfyUI/pull/7580
+
+import torch
+from tqdm.auto import trange
+from k_diffusion.sampling import (
+    default_noise_sampler,
+)
+
+
+@torch.no_grad()
+def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=0.5):
+    '''
+    SEEDS-2 - Stochastic Explicit Exponential Derivative-free Solvers (VE Data Prediction) stage 2
+    Arxiv: https://arxiv.org/abs/2305.14267
+    '''
+
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+    sigmas = sigmas.to(x.device)
+
+    inject_noise = eta > 0 and s_noise > 0
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            x = denoised
+        else:
+            t, t_next = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = t_next - t
+            h_eta = h * (eta + 1)
+            s = t + r * h
+            fac = 1 / (2 * r)
+            sigma_s = s.neg().exp()
+
+            coeff_1, coeff_2 = (-r * h_eta).expm1(), (-h_eta).expm1()
+            if inject_noise:
+                noise_coeff_1 = (-2 * r * h * eta).expm1().neg().sqrt()
+                noise_coeff_2 = ((-2 * r * h * eta).expm1() - (-2 * h * eta).expm1()).sqrt()
+                noise_1, noise_2 = noise_sampler(sigmas[i], sigma_s), noise_sampler(sigma_s, sigmas[i + 1])
+
+            # Step 1
+            x_2 = (coeff_1 + 1) * x - coeff_1 * denoised
+            if inject_noise:
+                x_2 = x_2 + sigma_s * (noise_coeff_1 * noise_1) * s_noise
+            denoised_2 = model(x_2, sigma_s * s_in, **extra_args)
+
+            # Step 2
+            denoised_d = (1 - fac) * denoised + fac * denoised_2
+            x = (coeff_2 + 1) * x - coeff_2 * denoised_d
+            if inject_noise:
+                x = x + sigmas[i + 1] * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
+    return x
+
+@torch.no_grad()
+def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r_1=1./3, r_2=2./3):
+    '''
+    SEEDS-3 - Stochastic Explicit Exponential Derivative-free Solvers (VE Data Prediction) stage 3
+    Arxiv: https://arxiv.org/abs/2305.14267
+    '''
+
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+    sigmas = sigmas.to(x.device)
+
+    inject_noise = eta > 0 and s_noise > 0
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            x = denoised
+        else:
+            t, t_next = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = t_next - t
+            h_eta = h * (eta + 1)
+            s_1 = t + r_1 * h
+            s_2 = t + r_2 * h
+            sigma_s_1, sigma_s_2 = s_1.neg().exp(), s_2.neg().exp()
+
+            coeff_1, coeff_2, coeff_3 = (-r_1 * h_eta).expm1(), (-r_2 * h_eta).expm1(), (-h_eta).expm1()
+            if inject_noise:
+                noise_coeff_1 = (-2 * r_1 * h * eta).expm1().neg().sqrt()
+                noise_coeff_2 = ((-2 * r_1 * h * eta).expm1() - (-2 * r_2 * h * eta).expm1()).sqrt()
+                noise_coeff_3 = ((-2 * r_2 * h * eta).expm1() - (-2 * h * eta).expm1()).sqrt()
+                noise_1, noise_2, noise_3 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigma_s_2), noise_sampler(sigma_s_2, sigmas[i + 1])
+
+            # Step 1
+            x_2 = (coeff_1 + 1) * x - coeff_1 * denoised
+            if inject_noise:
+                x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
+            denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
+
+            # Step 2
+            x_3 = (coeff_2 + 1) * x - coeff_2 * denoised + (r_2 / r_1) * (coeff_2 / (r_2 * h_eta) + 1) * (denoised_2 - denoised)
+            if inject_noise:
+                x_3 = x_3 + sigma_s_2 * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
+            denoised_3 = model(x_3, sigma_s_2 * s_in, **extra_args)
+
+            # Step 3
+            x = (coeff_3 + 1) * x - coeff_3 * denoised + (1. / r_2) * (coeff_3 / h_eta + 1) * (denoised_3 - denoised)
+            if inject_noise:
+                x = x + sigmas[i + 1] * (noise_coeff_3 * noise_1 + noise_coeff_2 * noise_2 + noise_coeff_1 * noise_3) * s_noise
+    return x

webUI_ExtraSchedulers/scripts/simple_kes.py

@@ -0,0 +1,98 @@
+# based on -
+# Simple Karras-Exponential Scheduler, by Kittensx
+# https://github.com/Kittensx/Simple_KES
+
+import torch
+
+import gradio as gr
+from k_diffusion.sampling import get_sigmas_karras, get_sigmas_exponential
+
+from modules import shared
+
+def get_sigmas_simple_kes (n, sigma_min, sigma_max, device):
+    """
+    Scheduler function that blends sigma sequences using Karras and Exponential methods with adaptive parameters.
+
+    Parameters:
+        n (int): Number of steps.
+        sigma_min (float): Minimum sigma value.
+        sigma_max (float): Maximum sigma value.
+        device (torch.device): The device on which to perform computations (e.g., 'cuda' or 'cpu').
+        start_blend (float): Initial blend factor for dynamic blending.
+        end_blend (float): Final blend factor for dynamic blending.
+        sharpen_factor (float): Sharpening factor to be applied adaptively.
+        update_interval (int): Interval to update blend factors.
+        initial_step_size (float): Initial step size for adaptive step size calculation.
+        final_step_size (float): Final step size for adaptive step size calculation.
+        initial_noise_scale (float): Initial noise scale factor.
+        final_noise_scale (float): Final noise scale factor.
+        step_size_factor: Adjust to compensate for oversmoothing
+        noise_scale_factor: Adjust to provide more variation
+        
+    Returns:
+        torch.Tensor: A tensor of blended sigma values.
+    """
+
+    start_blend = getattr(shared.opts, 'kes_start_blend', 0.1)
+    end_blend = getattr(shared.opts, 'kes_end_blend', 0.5)
+    sharpness = getattr(shared.opts, 'kes_sharpness', 0.95)
+    initial_step_size = getattr(shared.opts, 'kes_initial_step_size', 0.9)
+    final_step_size = getattr(shared.opts, 'kes_final_step_size', 0.2)
+    initial_noise_scale = getattr(shared.opts, 'kes_initial_noise', 1.25)
+    final_noise_scale = getattr(shared.opts, 'kes_final_noise', 0.8)
+    smooth_blend_factor = getattr(shared.opts, 'kes_smooth_blend', 11)
+    step_size_factor = getattr(shared.opts, 'kes_step_size_factor', 0.8)
+    noise_scale_factor = getattr(shared.opts, 'kes_noise_scale', 0.9)
+    
+    # Expand sigma_max slightly to account for smoother transitions
+    # sigma_max = sigma_max * 1.1
+
+    # Generate sigma sequences using Karras and Exponential methods
+    sigmas_karras = get_sigmas_karras(n=n, sigma_min=sigma_min, sigma_max=sigma_max, device=device)
+    sigmas_exponential = get_sigmas_exponential(n=n, sigma_min=sigma_min, sigma_max=sigma_max, device=device)
+
+    # Define progress and initialize blend factor
+    progress = torch.linspace(0, 1, len(sigmas_karras)).to(device)
+    
+    sigs = torch.zeros_like(sigmas_karras).to(device)
+    
+    # Iterate through each step, dynamically adjust blend factor, step size, and noise scaling
+    for i in range(len(sigmas_karras)):
+        # Adaptive step size and blend factor calculations
+        step_size = initial_step_size * (1 - progress[i]) + final_step_size * progress[i] * step_size_factor  # 0.8 default value Adjusted to avoid over-smoothing
+
+        dynamic_blend_factor = start_blend * (1 - progress[i]) + end_blend * progress[i]
+
+        noise_scale = initial_noise_scale * (1 - progress[i]) + final_noise_scale * progress[i] * noise_scale_factor  # 0.9 default value Adjusted to keep more variation
+
+        # Calculate smooth blending between the two sigma sequences
+        smooth_blend = torch.sigmoid((dynamic_blend_factor - 0.5) * smooth_blend_factor) # Increase scaling factor to smooth transitions more
+        
+        # Compute blended sigma values
+        blended_sigma = sigmas_karras[i] * (1 - smooth_blend) + sigmas_exponential[i] * smooth_blend
+        
+        # Apply step size and noise scaling
+        sigs[i] = blended_sigma * step_size * noise_scale
+
+    # Optional: Adaptive sharpening based on sigma values
+    sharpen_mask = torch.where(sigs < sigma_min * 1.5, sharpness, 1.0).to(device)
+    sigs = sigs * sharpen_mask
+    
+    if torch.isnan(sigs).any() or torch.isinf(sigs).any():
+        raise ValueError("Invalid sigma values detected (NaN or Inf).")
+
+    return sigs.to(device)
+
+shared.options_templates.update(shared.options_section(('simple_kes', "Simple KES", ""), {
+    "kes_start_blend":       shared.OptionInfo(0.1,  "start blend factor",       gr.Slider, {"minimum": 0.0,  "maximum": 1.0,  "step": 0.01}),
+    "kes_end_blend":         shared.OptionInfo(0.5,  "end blend factor",         gr.Slider, {"minimum": 0.0,  "maximum": 1.0,  "step": 0.01}),
+    "kes_sharpness":         shared.OptionInfo(0.95, "sharpness",                gr.Slider, {"minimum": 0.0,  "maximum": 2.0,  "step": 0.01}),
+    "kes_initial_step_size": shared.OptionInfo(0.9,  "initial step size",        gr.Slider, {"minimum": 0.01, "maximum": 1.0,  "step": 0.01}), # larger max?
+    "kes_final_step_size":   shared.OptionInfo(0.2,  "final step size",          gr.Slider, {"minimum": 0.01, "maximum": 1.0,  "step": 0.01}), #larger max?
+    "kes_initial_noise":     shared.OptionInfo(1.25, "initial noise",            gr.Slider, {"minimum": 0.0,  "maximum": 4.0,  "step": 0.01}),
+    "kes_final_noise":       shared.OptionInfo(0.8,  "final noise",              gr.Slider, {"minimum": 0.0,  "maximum": 4.0,  "step": 0.01}),
+    "kes_smooth_blend":      shared.OptionInfo(11,   "smooth blend factor",      gr.Slider, {"minimum": 0.0,  "maximum": 50.0, "step": 0.1}),
+    "kes_step_size_factor":  shared.OptionInfo(0.8,  "step size factor",         gr.Slider, {"minimum": 0.0,  "maximum": 4.0,  "step": 0.01}),
+    "kes_noise_scale":       shared.OptionInfo(0.9,  "noise scale factor",       gr.Slider, {"minimum": 0.0,  "maximum": 4.0,  "step": 0.01}),
+}))
+