From 6ff12ebed49cb2ec6277f61290a73bbb44bca834 Mon Sep 17 00:00:00 2001 From: Extraltodeus Date: Tue, 27 Aug 2024 09:33:57 +0200 Subject: [PATCH] Add files via upload --- __init__.py | 25 +++++++++++++ custom_samplers.py | 93 ++++++++++++++++++++++++++++++++++++++++++++++ presets_to_add.py | 37 ++++++++++++++++++ 3 files changed, 155 insertions(+) create mode 100644 __init__.py create mode 100644 custom_samplers.py create mode 100644 presets_to_add.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..95f6660 --- /dev/null +++ b/__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, +} \ No newline at end of file diff --git a/custom_samplers.py b/custom_samplers.py new file mode 100644 index 0000000..2378b49 --- /dev/null +++ b/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, ) \ No newline at end of file diff --git a/presets_to_add.py b/presets_to_add.py new file mode 100644 index 0000000..1fcd59a --- /dev/null +++ b/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 \ No newline at end of file