PingPong sampler for ComfyUI

pingpongsampler_node.py

# 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,
}