From 5c0145dd69aa1cb22344212a1895f9d4a63452c8 Mon Sep 17 00:00:00 2001 From: Clybius Date: Thu, 15 Feb 2024 12:58:35 -0600 Subject: [PATCH] Add experimental sampler, remove extraneous import --- __init__.py | 1 + extra_samplers.py | 147 +++++++++++++++++++++++++++++++++++++++++++++- nodes.py | 26 +++++++- 3 files changed, 170 insertions(+), 4 deletions(-) diff --git a/__init__.py b/__init__.py index 1ff303d..d726b5c 100644 --- a/__init__.py +++ b/__init__.py @@ -13,5 +13,6 @@ NODE_CLASS_MAPPINGS = { "SamplerCLYB_4M_SDE_Momentumized": nodes.SamplerCLYB_4M_SDE_MOMENTUMIZED, "SamplerTTM": nodes.SamplerTTM, "SamplerLCMCustom": nodes.SamplerLCMCustom, + "SamplerEulerAncestralDancing_Experimental": nodes.SamplerEULER_ANCESTRAL_DANCING, } __all__ = ['NODE_CLASS_MAPPINGS'] diff --git a/extra_samplers.py b/extra_samplers.py index a60472a..8b92e4a 100644 --- a/extra_samplers.py +++ b/extra_samplers.py @@ -1,6 +1,5 @@ import math -from scipy import integrate import torch from torch import nn import torchsde @@ -244,6 +243,16 @@ def studentt_noise_sampler(x): # Produces more subject-focused outputs due to di print(s) return lambda sigma, sigma_next: noise.to(x.device) / (7/3) +from torch.distributions import Laplace +def rand_laplacian_like(x): + noise = torch.randn_like(x) / 4.0 + noise_size_H = noise.size(dim=2) + noise_size_W = noise.size(dim=3) + noise += Laplace(loc=0, scale=1.0).rsample(x.size()).to(noise.device) + #noise += perlin + #print(noise) + return noise / noise.std() + def highres_pyramid_noise_like(x, discount=0.7): b, c, h, w = x.shape # EDIT: w and h get over-written, rename for a different variant! orig_h = h @@ -461,6 +470,8 @@ def sample_res_solver(model, x, sigmas, extra_args=None, callback=None, disable= noise_sampler = lambda sigma, sigma_next: highres_pyramid_noise_like(x) case "perlin": noise_sampler = lambda sigma, sigma_next: rand_perlin_like(x) + case "laplacian": + noise_sampler = lambda sigma, sigma_next: rand_laplacian_like(x) case _: noise_sampler = lambda sigma, sigma_next: torch.randn_like(x) return sample_refined_exp_s(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, noise_sampler=noise_sampler, denoise_to_zero=denoise_to_zero, simple_phi_calc=simple_phi_calc, c2=c2, ita=ita, momentum=momentum) @@ -577,6 +588,8 @@ def sample_dpmpp_dualsdemomentum(model, x, sigmas, extra_args=None, callback=Non noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73 case "perlin": noise_sampler = lambda sigma, sigma_next: rand_perlin_like(x) + case "laplacian": + noise_sampler = lambda sigma, sigma_next: rand_laplacian_like(x) case _: noise_sampler = lambda sigma, sigma_next: torch.randn_like(x) return sample_dpmpp_dualsde_momentum(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, r=r, momentum=momentum) @@ -625,10 +638,141 @@ def sample_clyb_4m_sde(model, x, sigmas, extra_args=None, callback=None, disable noise_sampler = lambda sigma, sigma_next: highres_pyramid_noise_like(x) case "perlin": noise_sampler = lambda sigma, sigma_next: rand_perlin_like(x) + case "laplacian": + noise_sampler = lambda sigma, sigma_next: rand_laplacian_like(x) case _: noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73 return sample_clyb_4m_sde_momentumized(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, momentum=momentum) +""" +# This code works, but I'm currently experimenting with different methods +@torch.no_grad() +def sampler_euler_ancestral_dancing(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, leap=2, eta_dance=1.0): + #Ancestral sampling with Euler method steps, dancing 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 + unsample_noise_sampler = lambda sigma, sigma_next: torch.randn_like(x) + s_in = x.new_ones([x.shape[0]]) + for i in trange(len(sigmas) - 1, disable=disable): + if i < len(sigmas) - leap: + is_danceable = sigmas[i + leap] > 0 + else: + is_danceable = False + denoised = model(x, sigmas[i] * s_in, **extra_args) + sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + leap] if is_danceable else 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: + if is_danceable: + x = x + noise_sampler(sigmas[i], sigmas[i + leap]) * s_noise * sigma_up + #x = x + noise_sampler(sigmas[i + 2], sigmas[i + 1]) * s_noise * sigma_up + #denoised2 = model(x, sigmas[i + 2] * s_in, **extra_args) + sigma_down2, sigma_up2 = get_ancestral_step(sigmas[i + leap], sigmas[i + 1], eta=eta_dance) + d_2 = to_d(x, sigmas[i + leap], denoised) + dt_2 = sigma_down2 - sigmas[i + leap] + x = x + d_2 * dt_2 + x = x + noise_sampler(sigmas[i + leap], sigmas[i + 1]) * s_noise * sigma_up2 + + #sigma_down3, sigma_up3 = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) + #x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up3 + + #denoised2 = model(x, sigmas[i] * s_in, **extra_args) + #d_3 = to_d(x, sigmas[i], denoised2) + #dt_3 = sigma_down3 - sigmas[i] + #x = x + d_3 * dt_3 + d_2 * dt_2 + #print(dt_3, dt_2) + #x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up3 + #x = x + d * dt + else: + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up + + return x +""" +def rej(a, b): + """ + Implements the rejection function for alternative diffusion. + + Args: + a: Tensor of shape (B, H, W, C), where B is batch size, H and W are spatial dimensions, and C is number of channels. + b: Tensor of the same shape as a. + + Returns: + Tensor of the same shape as a and b, containing the rejection output. + """ + return (b * torch.tensordot(a, b, dims=len(a.shape)) / torch.tensordot(b, b, dims=len(a.shape))) - a + +@torch.no_grad() +def sampler_euler_ancestral_dancing(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, leap=2, eta_dance=1.0): + #Ancestral sampling with Euler method steps, dancing 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 + unsample_noise_sampler = lambda sigma, sigma_next: torch.randn_like(x) + s_in = x.new_ones([x.shape[0]]) + for i in trange(len(sigmas) - 1, disable=disable): + if i < len(sigmas) - leap: + is_danceable = sigmas[i + leap] > 0 + else: + is_danceable = False + orig_x = x + denoised = model(x, sigmas[i] * s_in, **extra_args) + sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + leap] if is_danceable else 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: + if is_danceable: + #x = x + noise_sampler(sigmas[i], sigmas[i + leap]) * s_noise * sigma_up + #x = x + noise_sampler(sigmas[i + 2], sigmas[i + 1]) * s_noise * sigma_up + denoised2 = model(x, sigmas[i + leap] * s_in, **extra_args) + sigma_down2, sigma_up2 = get_ancestral_step(sigmas[i + leap], sigmas[i + 1], eta=eta_dance) + d_2 = to_d(x, sigmas[i + leap], denoised2) + dt_2 = sigma_down2 - sigmas[i + leap] + x_2 = x + d_2 * dt_2 + #x_2 = x_2 + noise_sampler(sigmas[i + leap], sigmas[i]) * s_noise * sigma_up2 + + sigma_down3, sigma_up3 = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) + #x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up3 + + denoised3 = model(x_2, sigmas[i] * s_in, **extra_args) + d_3 = to_d(orig_x, sigmas[i], denoised3 + rej(denoised3 - denoised, denoised2 - denoised)) + #d_3 = to_d(x_2, sigmas[i], denoised3) + dt_3 = sigma_down3 - sigmas[i] + x = orig_x + d_3 * dt_3 # Very denoised, slightly denoised + #print(dt_3, dt_2) + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up3 + #x = x + d * dt + else: + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up + + return x + +def sample_euler_ancestral_dancing(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler="gaussian", leap=2, eta_dance=1.0): + sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() + seed = extra_args.get("seed", None) + match noise_sampler: + case "brownian": + noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=False) + case "gaussian": + noise_sampler = lambda sigma, sigma_next: torch.randn_like(x) + case "uniform": + noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73 + case "highres-pyramid": + noise_sampler = lambda sigma, sigma_next: highres_pyramid_noise_like(x) + case "perlin": + noise_sampler = lambda sigma, sigma_next: rand_perlin_like(x) + case "laplacian": + noise_sampler = lambda sigma, sigma_next: rand_laplacian_like(x) + case _: + noise_sampler = lambda sigma, sigma_next: (torch.rand_like(x) - 0.5) * 2 * 1.73 + return sampler_euler_ancestral_dancing(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, leap=leap, eta_dance=eta_dance) + # Add your personal samplers below here, just for formatting purposes ;3 # Add any extra samplers to the following dictionary @@ -638,4 +782,5 @@ extra_samplers = { "clyb_4m_sde_momentumized": sample_clyb_4m_sde, "ttm": sample_ttmcustom, "lcm_custom_noise": sample_lcmcustom, + "euler_ancestral_dancing": sample_euler_ancestral_dancing, } \ No newline at end of file diff --git a/nodes.py b/nodes.py index 7185afb..b9899dc 100644 --- a/nodes.py +++ b/nodes.py @@ -78,7 +78,7 @@ class SamplerRES_MOMENTUMIZED: @classmethod def INPUT_TYPES(s): return {"required": - {"noise_sampler_type": (["gaussian", "uniform", "brownian", "highres-pyramid", "perlin"], ), + {"noise_sampler_type": (["gaussian", "uniform", "brownian", "highres-pyramid", "perlin", "laplacian"], ), "momentum": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 1.0, "step":0.01}), "denoise_to_zero": ("BOOLEAN", {"default": True}), "simple_phi_calc": ("BOOLEAN", {"default": False}), @@ -99,7 +99,7 @@ class SamplerDPMPP_DUALSDE_MOMENTUMIZED: @classmethod def INPUT_TYPES(s): return {"required": - {"noise_sampler_type": (["gaussian", "uniform", "brownian", "perlin"], ), + {"noise_sampler_type": (["gaussian", "uniform", "brownian", "perlin", "laplacian"], ), "momentum": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 1.0, "step":0.01}), "eta": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}), "s_noise": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}), @@ -154,7 +154,7 @@ class SamplerCLYB_4M_SDE_MOMENTUMIZED: @classmethod def INPUT_TYPES(s): return {"required": - {"noise_sampler_type": (["gaussian", "uniform", "brownian", "highres-pyramid", "perlin"], ), + {"noise_sampler_type": (["gaussian", "uniform", "brownian", "highres-pyramid", "perlin", "laplacian"], ), "momentum": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 1.0, "step":0.01}), "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}), "s_noise": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}), @@ -169,6 +169,26 @@ class SamplerCLYB_4M_SDE_MOMENTUMIZED: sampler = comfy.samplers.ksampler("clyb_4m_sde_momentumized", {"noise_sampler": noise_sampler_type, "eta": eta, "s_noise": s_noise, "momentum": momentum}) return (sampler, ) +class SamplerEULER_ANCESTRAL_DANCING: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"noise_sampler_type": (["gaussian", "uniform", "brownian", "highres-pyramid", "perlin", "laplacian"], ), + "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}), + "eta_dance": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01}), + "s_noise": ("FLOAT", {"default": 1, "min": 0.0, "max": 100.0, "step":0.01}), + "leap": ("INT", {"default": 2, "min": 1, "max": 16, "step":1}), + } + } + RETURN_TYPES = ("SAMPLER",) + CATEGORY = "sampling/custom_sampling" + + FUNCTION = "get_sampler" + + def get_sampler(self, noise_sampler_type, eta, s_noise, leap, eta_dance): + sampler = comfy.samplers.ksampler("euler_ancestral_dancing", {"noise_sampler": noise_sampler_type, "eta": eta, "s_noise": s_noise, "leap": leap, "eta_dance": eta_dance}) + return (sampler, ) + from comfy import model_management import comfy.utils import comfy.conds