Upload comfy/custom_nodes/pingpongsampler_node.py with huggingface_hub

comfy/custom_nodes/pingpongsampler_node.py

@@ -0,0 +1,252 @@
+# By https://github.com/blepping
+# LICENSE: Apache2
+# Usage: Place this file in the custom_nodes directory and restart ComfyUI+refresh browser.
+#        It will add a PingPongSampler node that can be used with SamplerCustom, etc.
+
+import random
+import torch
+
+from tqdm.auto import trange
+
+from comfy import model_sampling
+from comfy.samplers import KSAMPLER
+import nodes
+
+
+BLEND_MODES = None
+
+def _ensure_blend_modes():
+    global BLEND_MODES
+    if BLEND_MODES is not None:
+        return
+    bleh = getattr(nodes, "_blepping_integrations", {}).get("bleh")
+    if bleh is not None:
+        BLEND_MODES = bleh.py.latent_utils.BLENDING_MODES
+    else:
+        BLEND_MODES = {"lerp": torch.lerp, "a_only": lambda a, _b, _t: a, "b_only": lambda _a, b, _t: b}
+
+class ModelProxy:
+    def __init__(self, model, last_x, last_sigma, last_denoised):
+        self.__model = model
+        self.__last_x = last_x
+        self.__last_sigma = last_sigma
+        self.__last_denoised = last_denoised
+
+    def __call__(self, x, sigma, *args, **kwargs):
+        if torch.allclose(sigma.to(self.__last_sigma), self.__last_sigma) and torch.allclose(x.to(self.__last_x), self.__last_x):
+            return self.__last_denoised.to(x, copy=True)
+        return self.__model(x, sigma, *args, **kwargs)
+
+    def __getattr__(self, k):
+        return getattr(self.__model, k)
+
+class PingPongSampler:
+    def __init__(self, model, x, sigmas, *args, extra_args=None, callback=None, disable=None, noise_sampler=None, s_noise=1.0, pingpong_options=None, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+        self.model_ = model
+        self.sigmas = sigmas
+        self.x = x
+        self.s_in = x.new_ones((x.shape[0],))
+        self.extra_args = extra_args.copy() if extra_args is not None else {}
+        self.seed = self.extra_args.pop("seed", 42)
+        self.disable = disable
+        self.callback_ = callback
+        if pingpong_options is None:
+            pingpong_options= {}
+        self.first_ancestral_step = pingpong_options.get("first_ancestral_step", 0)
+        self.last_ancestral_step = pingpong_options.get("last_ancestral_step", 0)
+
+        self.pingpong_blend = pingpong_options.get("pingpong_blend")
+        sampler_opt = pingpong_options.get("external_sampler")
+        if self.pingpong_blend != 1.0 and sampler_opt is None:
+            raise ValueError("Sampler input must be connect when pingpong_blend isn't 1.0")
+        self.external_sampler = sampler_opt
+        self.step_blend_function = pingpong_options.get("step_blend_function", torch.lerp)
+        self.blend_function = pingpong_options.get("blend_function", torch.lerp)
+        self.s_noise = s_noise
+        self.is_rf = isinstance(model.inner_model.inner_model.model_sampling, model_sampling.CONST)
+        if noise_sampler is None:
+            def noise_sampler(*_unused):
+                return torch.randn_like(x)
+        self.noise_sampler = noise_sampler
+
+    @classmethod
+    def go(cls, model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, s_noise=1.0, pingpong_options=None, **kwargs):
+        return cls(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler, s_noise=s_noise, pingpong_options=pingpong_options, **kwargs)()
+
+    def model(self, x, sigma, **kwargs):
+        return self.model_(x, sigma * self.s_in, **self.extra_args, **kwargs)
+
+    def callback(self, idx, x, sigma, denoised):
+        if self.callback_ is None:
+            return
+        self.callback_({
+            "i": idx,
+            "x": x,
+            "sigma": sigma,
+            "sigma_hat": sigma,
+            "denoised": denoised,
+        })
+
+    def __call__(self):
+        x = self.x
+        noise_sampler = self.noise_sampler
+        astart_step = self.first_ancestral_step
+        aend_step = self.last_ancestral_step
+        last_step_idx = len(self.sigmas) - 2
+        step_count = len(self.sigmas) - 1
+        if astart_step < 0:
+            astart_step = step_count + astart_step
+        if aend_step < 0:
+            aend_step = step_count + aend_step
+        astart_step = min(last_step_idx, max(0, astart_step))
+        aend_step = min(last_step_idx, max(0, aend_step))
+        s_noise = self.s_noise
+        seed_offset = 10
+        for idx in trange(step_count, disable=self.disable):
+            sigma, sigma_next = self.sigmas[idx:idx + 2]
+            orig_x = x
+            denoised = self.model(orig_x, sigma)
+            self.callback(idx, x, sigma, denoised)
+            use_ancestral = astart_step <= idx <= aend_step
+            if sigma_next <= 1e-06:
+                return denoised
+            if not use_ancestral:
+                x = self.step_blend_function(denoised, x, sigma_next / sigma)
+                continue
+            if self.pingpong_blend != 1.0:
+                alt_x = self.external_sampler.sampler_function(
+                    ModelProxy(self.model_, x, sigma, denoised),
+                    orig_x.clone(),
+                    self.sigmas[idx:idx + 2].clone(),
+                    *self.args,
+                    disable=True,
+                    callback=None,
+                    extra_args=self.extra_args | {"seed": self.seed + seed_offset},
+                    **self.external_sampler.extra_options,
+                    **self.kwargs,
+                )
+                seed_offset += 10
+                if self.pingpong_blend <= 0:
+                    x = alt_x
+                    continue
+            noise = noise_sampler(sigma, sigma_next).mul_(self.s_noise)
+            if self.is_rf:
+                x = self.step_blend_function(denoised, noise, sigma_next)
+            else:
+                x = denoised + noise * sigma_next
+            if self.pingpong_blend != 1.0:
+                x = self.blend_function(alt_x, x, self.pingpong_blend)
+                del alt_x
+        return x
+
+
+class PingPongSamplerNode:
+    CATEGORY = "sampling/custom_sampling/samplers"
+    RETURN_TYPES = ("SAMPLER",)
+    FUNCTION = "go"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        _ensure_blend_modes()
+        return {
+            "required": {
+                "s_noise": ("FLOAT", {"default": 1.0, "min": -1000.0, "max": 1000.0}),
+                "first_ancestral_step": ("INT", {"default": 0, "min": -10000, "max": 10000}),
+                "last_ancestral_step": ("INT", {"default": -1, "min": -10000, "max": 10000}),
+                "pingpong_blend": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001, "tooltip": "Allows blending pingpong sampling with a different sampler. Only has an effect during the ancestral_step range. If set to a value below 1.0 (100% pingpong) then sampler_opt must be attached."}),
+                "blend_mode": (tuple(BLEND_MODES), {"default": "lerp", "tooltip": "Blend mode to use when blending pingpong sampling with the external sampler. See tooltip for pingpong_blend. Can integrate with ComfyUI-bleh to add more blend modes."}),
+                "step_blend_mode": (tuple(BLEND_MODES), {"default": "lerp", "tooltip": "Blend mode to use for pingpong steps. Changing this is likely a bad idea. Does not apply for ancestral steps on non-flow models.  Can integrate with ComfyUI-bleh to add more blend modes."}),
+
+            },
+            "optional": {
+                "sampler_opt": ("SAMPLER", {"tooltip": "Optional when pingpong_blend is 1.0. Result of a pingpong step will be blended with output from this sampler with the configured ratio. Calls the sampler on a single step so will not work well with samplers that care about state (I.E. history samplers such as deis, res_multistep, etc)."}),
+            },
+        }
+
+    @classmethod
+    def go(cls, *, s_noise: float, first_ancestral_step: int, last_ancestral_step: int, pingpong_blend: float, blend_mode: str, step_blend_mode: str, sampler_opt = None):
+        options = {
+            "s_noise": s_noise,
+            "pingpong_options": {
+                "first_ancestral_step": first_ancestral_step,
+                "last_ancestral_step": last_ancestral_step,
+                "pingpong_blend": pingpong_blend,
+                "blend_function": BLEND_MODES[blend_mode],
+                "step_blend_function": BLEND_MODES[step_blend_mode],
+                "external_sampler": sampler_opt,
+            },
+        }
+        return (KSAMPLER(PingPongSampler.go, extra_options=options),)
+
+class RestlessSchedulerNode:
+    DESCRIPTION = "HACK: A weird scheduler that will randomly jump around a list of sigmas you input. Not recommended. Breaks most multi-step and history samplers. Works okay-ish with Pingpong."
+    CATEGORY = "sampling/custom_sampling/schedulers"
+    RETURN_TYPES = ("SIGMAS",)
+    FUNCTION = "go"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sigmas": ("SIGMAS",),
+                "seed": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "tooltip": "Seed to use for generating schedule.",
+                    },
+                ),
+                "shrink_factor": ("FLOAT", {
+                    "default": 0.3,
+                    "tooltip": "Amount the window for restless scheduling shrinks by per iteration.",
+                }),
+                "first_restless_step": ("INT", {
+                    "default": 3, "min": 1,
+                    "tooltip": "First step (0-based) to include for restless scheduling. Must be greater than 1 and less than last_restless_step.",
+                }),
+                "last_restless_step": ("INT", {
+                    "default": -4, "min": -10000, "max": 10000,
+                    "tooltip": "Last step (0-based) to include for restless scheduling. Can be negative to count from the end, but you cannot target the last sigma in the list.",
+                }),
+            },
+        }
+
+    @classmethod
+    def go(cls, *, sigmas: torch.Tensor, seed: int, shrink_factor: float, first_restless_step: int, last_restless_step: int) -> tuple:
+        n_sigmas = len(sigmas)
+        if n_sigmas < 3:
+            return (sigmas,)
+        if last_restless_step < 0:
+            last_restless_step = n_sigmas + last_restless_step
+        if last_restless_step <= first_restless_step:
+            raise ValueError("Last restless step <= first restless step!")
+        if last_restless_step >= n_sigmas - 1:
+            raise ValueError("Last restless step cannot include the final sigma")
+        orig_sigmas = sigmas
+        random.seed(seed)
+        result = sigmas[:first_restless_step].tolist()
+        end_chunk = sigmas[last_restless_step + 1:].tolist()
+        sigmas = sigmas[first_restless_step:last_restless_step + 1].tolist()
+        n_sigmas = len(sigmas)
+        shrinkage = 0.0
+        curr_idx = None
+        while (window_size := int((n_sigmas - 1) - shrinkage)) > 0:
+            next_idx = random.randint(0, window_size)
+            if next_idx == curr_idx:
+                next_idx += 1
+            result.append(sigmas[int(shrinkage) + next_idx])
+            curr_idx = next_idx
+            shrinkage += shrink_factor
+        result += end_chunk
+        return (torch.tensor(result, dtype=torch.float32, device="cpu").to(orig_sigmas),)
+
+
+
+NODE_CLASS_MAPPINGS = {
+    "PingPongSampler": PingPongSamplerNode,
+    "RestlessScheduler": RestlessSchedulerNode,
+}