diff --git a/custom_samplers.py b/custom_samplers.py index 177bf16..2f7fed2 100644 --- a/custom_samplers.py +++ b/custom_samplers.py @@ -4,14 +4,15 @@ 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) +def fast_distance_weights(t): + z = t / t.norm() + distances = (z.unsqueeze(0) - z.unsqueeze(1)).abs().sum(dim=0) + distances = 1 - (distances - distances.min(dim=0).values) / (distances.max(dim=0).values - distances.min(dim=0).values) + distances[~torch.isfinite(distances)] = 1 + distances = distances.pow(t.shape[0]) + distances = distances / distances.sum(dim=0) + res = (t * distances).sum(dim=0) + return res # Euler and CFGpp part taken from comfy_extras/nodes_advanced_samplers def distance_wrap(resample,resample_end=-1,cfgpp=False): @@ -69,7 +70,7 @@ def distance_wrap(resample,resample_end=-1,cfgpp=False): if re_step == 0: d = (new_d + d) / 2 else: - d = fast_distance_weights(torch.stack(x_n), re_step + 2) + d = fast_distance_weights(torch.stack(x_n)) x_n.append(d) x = x + d * dt return x @@ -80,7 +81,7 @@ class SamplerDistanceAdvanced: 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."}), + "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",)