94 lines
4.2 KiB
Python
94 lines
4.2 KiB
Python
import torch
|
|
from comfy.k_diffusion.sampling import trange, to_d
|
|
import comfy.model_patcher
|
|
import comfy.samplers
|
|
|
|
@torch.no_grad()
|
|
def fast_distance_weights(t,p):
|
|
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))
|
|
# progression = lambda x: max(0,min(1,((x - s_min) / (s_max - s_min))))
|
|
|
|
if resample == -1:
|
|
current_resample = min(10, sigmas.shape[0] // 2)
|
|
else:
|
|
current_resample = resample
|
|
|
|
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(current_resample,resample_end),min(max(current_resample,resample_end),int(current_resample * res_mul + resample_end * (1 - res_mul))))
|
|
else:
|
|
resample_steps = current_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_n.append(d)
|
|
x = x + d * dt
|
|
return x
|
|
return sample_distance_advanced
|
|
|
|
class SamplerDistanceAdvanced:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"resample": ("INT", {"default": 3, "min": -1, "max": 32, "step": 1,
|
|
"tooltip":"0 all along gives Euler. 1 gives Heun.\nAnything starting from 2 will use the distance method.\n-1 will do remaining steps + 1 as the resample value. This can be pretty slow."}),
|
|
"resample_end": ("INT", {"default": 1, "min": -1, "max": 32, "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, )
|