https://github.com/kohya-ss/sd-scripts/tree/gradual_latent_hires_fix

gradual_latent.py

import comfy
from comfy.samplers import KSAMPLER
import torch
from comfy.k_diffusion.sampling import default_noise_sampler, get_ancestral_step, to_d, BrownianTreeNoiseSampler
from tqdm.auto import trange

@torch.no_grad()
def sample_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, upscale_ratio=2.0, start_step=5, end_step=15, upscale_n_step=3):
    """Ancestral sampling with Euler method steps."""
    extra_args = {} if extra_args is None else extra_args
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
    s_in = x.new_ones([x.shape[0]])
        
    # make upscale info
    upscale_steps = []
    step = start_step-1
    while step < end_step-1:
        upscale_steps.append(step)
        step += upscale_n_step
    height, width = x.shape[2:]
    upscale_shapes = [(int(height * (((upscale_ratio-1) / i) + 1)), int(width * (((upscale_ratio-1) / i) + 1))) for i in reversed(range(1, len(upscale_steps)+1))]
    upscale_info = {k: v for k, v in zip(upscale_steps, upscale_shapes)}


    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        d = to_d(x, sigmas[i], denoised)
        # Euler method
        dt = sigma_down - sigmas[i]
        x = x + d * dt
        if sigmas[i + 1] > 0:
            # Resize
            if i in upscale_info:
                x = torch.nn.functional.interpolate(x, size=upscale_info[i], mode='bicubic', align_corners=False)
            noise_sampler = default_noise_sampler(x)
            noise = noise_sampler(sigmas[i], sigmas[i + 1])
            x = x + noise * sigma_up * s_noise
    return x

@torch.no_grad()
def sample_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, upscale_ratio=2.0, start_step=5, end_step=15, upscale_n_step=3):
    """Ancestral sampling with DPM-Solver++(2S) second-order steps."""
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    sigma_fn = lambda t: t.neg().exp()
    t_fn = lambda sigma: sigma.log().neg()

    # make upscale info
    upscale_steps = []
    step = start_step-1
    while step < end_step-1:
        upscale_steps.append(step)
        step += upscale_n_step
    height, width = x.shape[2:]
    upscale_shapes = [(int(height * (((upscale_ratio-1) / i) + 1)), int(width * (((upscale_ratio-1) / i) + 1))) for i in reversed(range(1, len(upscale_steps)+1))]
    upscale_info = {k: v for k, v in zip(upscale_steps, upscale_shapes)}

    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        if sigma_down == 0:
            # Euler method
            d = to_d(x, sigmas[i], denoised)
            dt = sigma_down - sigmas[i]
            x = x + d * dt
        else:
            # DPM-Solver++(2S)
            t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
            r = 1 / 2
            h = t_next - t
            s = t + r * h
            x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * denoised
            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_2
        # Noise addition
        if sigmas[i + 1] > 0:
            # Resize
            if i in upscale_info:
                x = torch.nn.functional.interpolate(x, size=upscale_info[i], mode='bicubic', align_corners=False)
            noise_sampler = default_noise_sampler(x)
            noise = noise_sampler(sigmas[i], sigmas[i + 1])
            x = x + noise * sigma_up * s_noise
    return x

@torch.no_grad()
def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint', upscale_ratio=2.0, start_step=5, end_step=15, upscale_n_step=3):
    """DPM-Solver++(2M) SDE."""

    if solver_type not in {'heun', 'midpoint'}:
        raise ValueError('solver_type must be \'heun\' or \'midpoint\'')

    seed = extra_args.get("seed", None)
    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])

    old_denoised = None
    h_last = None
    h = None

    # make upscale info
    upscale_steps = []
    step = start_step-1
    while step < end_step-1:
        upscale_steps.append(step)
        step += upscale_n_step
    height, width = x.shape[2:]
    upscale_shapes = [(int(height * (((upscale_ratio-1) / i) + 1)), int(width * (((upscale_ratio-1) / i) + 1))) for i in reversed(range(1, len(upscale_steps)+1))]
    upscale_info = {k: v for k, v in zip(upscale_steps, upscale_shapes)}

    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
        if sigmas[i + 1] == 0:
            # Denoising step
            x = denoised
        else:
            # DPM-Solver++(2M) SDE
            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
            h = s - t
            eta_h = eta * h

            x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised

            if old_denoised is not None:
                r = h_last / h
                if solver_type == 'heun':
                    x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised)
                elif solver_type == 'midpoint':
                    x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised)

            if eta:
                # Resize
                if i in upscale_info:
                    x = torch.nn.functional.interpolate(x, size=upscale_info[i], mode='bicubic', align_corners=False)
                    denoised = None # 次ステップとサイズがあわないのでとりあえずNoneにしておく。
                noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True)
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise

        old_denoised = denoised
        h_last = h
    return x

@torch.no_grad()
def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, eta=None, s_noise=None, upscale_ratio=2.0, start_step=5, end_step=15, upscale_n_step=3):
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])

    # make upscale info
    upscale_steps = []
    step = start_step-1
    while step < end_step-1:
        upscale_steps.append(step)
        step += upscale_n_step
    height, width = x.shape[2:]
    upscale_shapes = [(int(height * (((upscale_ratio-1) / i) + 1)), int(width * (((upscale_ratio-1) / i) + 1))) for i in reversed(range(1, len(upscale_steps)+1))]
    upscale_info = {k: v for k, v in zip(upscale_steps, upscale_shapes)}

    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})

        x = denoised
        if sigmas[i + 1] > 0:
            # Resize
            if i in upscale_info:
                x = torch.nn.functional.interpolate(x, size=upscale_info[i], mode='bicubic', align_corners=False)
            noise_sampler = default_noise_sampler(x)
            x += sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1])
            
    return x

class GradualLatentSampler:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {
                        "sampler_name": (["euler_ancestral", "dpmpp_2s_ancestral", "dpmpp_2m_sde", "lcm"], ),
                        "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step":0.01, "round": False}),
                        "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step":0.01, "round": False}),
                        "upscale_ratio": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 16.0, "step":0.01, "round": False}),
                        "start_step": ("INT", {"default": 5, "min": 0, "max": 1000, "step": 1}),
                        "end_step": ("INT", {"default": 15, "min": 0, "max": 1000, "step": 1}),
                        "upscale_n_step": ("INT", {"default": 3, "min": 0, "max": 1000, "step": 1}),
                    }
               }
    RETURN_TYPES = ("SAMPLER",)
    CATEGORY = "sampling/custom_sampling/samplers"

    FUNCTION = "get_sampler"

    def get_sampler(self, sampler_name, eta, s_noise, upscale_ratio, start_step, end_step, upscale_n_step):
        if sampler_name == "euler_ancestral":
            sample_function = sample_euler_ancestral
        elif sampler_name == "dpmpp_2s_ancestral":
            sample_function = sample_dpmpp_2s_ancestral
        elif sampler_name == "dpmpp_2m_sde":
            sample_function = sample_dpmpp_2m_sde
        elif sampler_name == "lcm":
            sample_function = sample_lcm
        else:
            raise ValueError("Unknown sampler name")

        sampler = KSAMPLER(sample_function, {"eta":eta, "s_noise":s_noise, "upscale_ratio": upscale_ratio, "start_step": start_step, "end_step": end_step, "upscale_n_step": upscale_n_step})
        return (sampler, )
    
NODE_CLASS_MAPPINGS = {
    "GradualLatentSampler": GradualLatentSampler,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    # Sampling
    "GradualLatentSampler": "GradualLatentSampler",
}