Add files via upload

__init__.py

@@ -0,0 +1,25 @@
+from .custom_samplers import SamplerDistanceAdvanced
+from .presets_to_add import extra_samplers
+
+def add_samplers():
+    from comfy.samplers import KSampler, k_diffusion_sampling
+    added = 0
+    samplers_names = [n for n in extra_samplers][::-1]
+    for sampler in samplers_names:
+        if sampler not in KSampler.SAMPLERS:
+            try:
+                idx = KSampler.SAMPLERS.index("uni_pc_bh2") # Last item in the samplers list
+                KSampler.SAMPLERS.insert(idx+1, sampler) # Add our custom samplers
+                setattr(k_diffusion_sampling, "sample_{}".format(sampler), extra_samplers[sampler])
+                added += 1
+            except ValueError as _err:
+                pass
+    if added > 0:
+        import importlib
+        importlib.reload(k_diffusion_sampling)
+
+add_samplers()
+
+NODE_CLASS_MAPPINGS = {
+    "SamplerDistance": SamplerDistanceAdvanced,
+}

custom_samplers.py

@@ -0,0 +1,93 @@
+import torch
+from comfy.k_diffusion.sampling import trange, to_d
+import comfy.model_patcher
+import comfy.samplers
+
+"""
+I wrote this logic initially to merge models.
+If you want to use this for that, I do not recommand to subtract
+the entire batch at once but to iterate manually unless you have
+a ton of memory.
+"""
+
+@torch.no_grad()
+def fast_distance_weights(t,p=2):
+    d = torch.zeros_like(t,device=t.device)
+    for i in range(t.shape[0]):
+        d[i] = (t - t[i]).abs().sum(dim=0)
+    d = (1 - (d - d.min()) / (d.max() - d.min())).pow(p)
+    d = torch.nan_to_num(d,nan=1,neginf=1,posinf=1)
+    d = (d / d.sum(dim=0))
+    return (d * t).sum(dim=0)
+
+# Euler and CFGpp part taken from comfy_extras/nodes_advanced_samplers
+def distance_wrap(resample,resample_end=-1,cfgpp=False):
+    @torch.no_grad()
+    def sample_distance_advanced(model, x, sigmas, extra_args=None, callback=None, disable=None):
+        extra_args = {} if extra_args is None else extra_args
+        if cfgpp:
+            uncond = None
+            def post_cfg_function(args):
+                nonlocal uncond
+                uncond = args["uncond_denoised"]
+                return args["denoised"]
+            model_options = extra_args.get("model_options", {}).copy()
+            extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function)
+
+        s_min, s_max = sigmas[sigmas > 0].min(), sigmas.max()
+        progression = lambda x: max(0,min(1,((x - s_min) / (s_max - s_min)) ** 0.5))
+
+        s_in = x.new_ones([x.shape[0]])
+        for i in trange(len(sigmas) - 1, disable=disable):
+            sigma_hat = sigmas[i]
+            denoised = model(x, sigma_hat * s_in, **extra_args)
+
+            if cfgpp and torch.any(uncond):
+                d = to_d(x - denoised + uncond, sigmas[i], denoised)
+            else:
+                d = to_d(x, sigma_hat, denoised)
+
+            if callback is not None:
+                callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+            dt = sigmas[i + 1] - sigma_hat
+
+            if resample_end >= 0:
+                res_mul = progression(sigma_hat)
+                resample_steps = max(min(resample,resample_end),min(max(resample,resample_end),int(resample * res_mul + resample_end * (1 - res_mul))))
+            else:
+                resample_steps = resample
+
+            if sigmas[i + 1] == 0 or resample_steps == 0:
+                # Euler method
+                x = x + d * dt
+            else:
+                x_n = [d]
+                for re_step in range(resample_steps):
+                    x_new = x + d * dt
+                    new_denoised = model(x_new, sigmas[i + 1] * s_in, **extra_args)
+                    new_d = to_d(x_new, sigmas[i + 1], new_denoised)
+                    x_n.append(new_d)
+                    if re_step == 0:
+                        d = (new_d + d) / 2
+                    else:
+                        d = fast_distance_weights(torch.stack(x_n), re_step + 2)
+                x = x + d * dt
+        return x
+    return sample_distance_advanced
+
+class SamplerDistanceAdvanced:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"resample": ("INT", {"default": 3, "min": 0, "max": 128, "step": 1,
+                                                  "tooltip":"0 all along gives Euler. 1 gives Heun.\nAnything starting from 2 will use the distance method."}),
+                             "resample_end": ("INT", {"default": -1, "min": 0, "max": 128, "step": 1, "tooltip":"How many resamples for the end. -1 means constant."}),
+                             "cfgpp" : ("BOOLEAN", {"default": True}),
+                             }}
+    RETURN_TYPES = ("SAMPLER",)
+    CATEGORY = "sampling/custom_sampling/samplers"
+    FUNCTION = "get_sampler"
+
+    def get_sampler(self,resample,resample_end,cfgpp):
+        sampler = comfy.samplers.KSAMPLER(
+            distance_wrap(resample=resample,cfgpp=cfgpp,resample_end=resample_end))
+        return (sampler, )

presets_to_add.py

@@ -0,0 +1,37 @@
+from .custom_samplers import distance_wrap
+
+extra_samplers = {}
+extra_samplers["heun_cfg_pp"] = distance_wrap(resample=1,cfgpp=True)
+
+"""
+To add a sampler to the list of samplers you can do it this way (outside of this comment):
+
+extra_samplers["the_name_that_you_want"] = distance_wrap(resample=3,resample_end=-1,cfgpp=False)
+
+"resample" is the starting value, how many more inferences it will use.
+For resample_end "-1" means constant.
+Constant resample at 0 gives Euler, 1 gives Heun.
+cfgpp will determin if you want it or not. True or False.
+
+You can remove the part below if you prefer to clean the list from the preset that I added.
+"""
+
+def make_preset(cfgpp,start,end):
+    ppname = "_cfg_pp" if cfgpp else ""
+    stepsn = "constant" if end == -1 else "fast"
+    preset_name = f"distance_{stepsn}_{start}{ppname}"
+    preset_sampler = distance_wrap(resample=start,resample_end=end,cfgpp=cfgpp)
+    return preset_name, preset_sampler
+
+ispp = [False,True] # CFGpp
+resample_start = [3,4]
+resample_const = [2,3,4]
+for ipp in ispp:
+    distance_p = resample_start
+    for ep in distance_p:
+        name, p_sampler = make_preset(ipp,ep,1)
+        extra_samplers[name] = p_sampler
+    distance_p = resample_const
+    for ep in distance_p:
+        name, p_sampler = make_preset(ipp,ep,-1)
+        extra_samplers[name] = p_sampler