From 81a163fcf033458cecebdcae4d982170f9809c7a Mon Sep 17 00:00:00 2001 From: blepping Date: Mon, 25 Mar 2024 06:47:15 -0600 Subject: [PATCH] Noise improvements phase 1 --- README.md | 6 ++--- changelog.md | 8 ++++++ py/nodes.py | 43 ++++++++++-------------------- py/noise.py | 68 ++++++++++++++++++++++++++++++++++++------------ py/powernoise.py | 3 ++- py/sonar.py | 34 ++++++++---------------- 6 files changed, 89 insertions(+), 73 deletions(-) diff --git a/README.md b/README.md index 58f6941..75e73f7 100644 --- a/README.md +++ b/README.md @@ -57,8 +57,6 @@ If you want to create noise for initial sampling, connect model and sigmas to th can be used to override configuration settings for other samplers, including the noise type. For example, you could force `euler_ancestral` to use a different noise type. It's also possible to override other settings like `s_noise`, etc. *Note*: The wrapper inspects the sampling function's arguments to see what it supports, so you should connect the sampler directly to this rather than having other nodes (like a different sampler wrapper) in between. -**Note**: If you are using this with Sonar samplers, make sure you set the noise type in the sampler to `gaussian` as the Sonar samplers only allow overriding noise types in that case. - ### `SonarCustomNoise` See the [Noise](#noise) section below for information on noise types. @@ -146,9 +144,9 @@ I also have some other ComfyUI nodes here: https://github.com/blepping/ComfyUI-b Original Sonar Sampler implementation (for A1111): https://github.com/Kahsolt/stable-diffusion-webui-sonar -My version basically just rips off this Sonar sampler implementation for Diffusers: https://github.com/alexblattner/modified-euler-samplers-for-sonar-diffusers/ +My version was initially based on this Sonar sampler implementation for Diffusers: https://github.com/alexblattner/modified-euler-samplers-for-sonar-diffusers/ -Noise generation functions copied from https://github.com/Clybius/ComfyUI-Extra-Samplers with only minor modifications. I may have broken some of them in the process _or_ they may not have been suitable for use and I took them anyway. If they don't work it is not a reflection on the original source. +Many noise generation functions copied from https://github.com/Clybius/ComfyUI-Extra-Samplers with only minor modifications. I may have broken some of them in the process _or_ they may not have been suitable for use and I took them anyway. If they don't work it is not a reflection on the original source. `SonarPowerNoise` contributed by [elias-gaeros](https://github.com/elias-gaeros/). Thanks! diff --git a/changelog.md b/changelog.md index 7476bb1..5e19a75 100644 --- a/changelog.md +++ b/changelog.md @@ -2,6 +2,14 @@ Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top. +## 20240325 + +* Fixed issue when using Sonar samplers in normal sampling nodes/via stuff like `KSamplerSelect`. +* Add `pyramid` (non-high-res) noise type. +* Allow selecting `brownian` noise in custom noise nodes (but it won't work with `NoisyLatentLike`). +* Use `brownian` as the default noise type for `SamplerSonarDPMPP`. +* Make overriding the selected noise type in Sonar samplers a warning instead of a hard error. + ## 20240320 * `NoisyLatentLike` node improved to allow calculating strength with sigmas and injecting noise itself. diff --git a/py/nodes.py b/py/nodes.py index 9e1247d..c76b0bc 100644 --- a/py/nodes.py +++ b/py/nodes.py @@ -8,6 +8,7 @@ import torch from comfy import samplers from . import noise +from .noise import NoiseType from .sonar import ( GuidanceConfig, GuidanceType, @@ -24,13 +25,7 @@ class NoisyLatentLikeNode: def INPUT_TYPES(cls): return { "required": { - "noise_type": ( - tuple( - t.name.lower() - for t in noise.NoiseType - if t is not noise.NoiseType.BROWNIAN - ), - ), + "noise_type": (tuple(NoiseType.get_names(skip=(NoiseType.BROWNIAN,))),), "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), "latent": ("LATENT",), "multiplier": ("FLOAT", {"default": 1.0}), @@ -83,7 +78,7 @@ class NoisyLatentLikeNode: ns = custom_noise_opt.make_noise_sampler(latent_samples) else: ns = noise.get_noise_sampler( - noise.NoiseType[noise_type.upper()], + NoiseType[noise_type.upper()], latent_samples, None, None, @@ -164,13 +159,7 @@ class SonarCustomNoiseNode(SonarCustomNoiseNodeBase): def INPUT_TYPES(cls): result = super().INPUT_TYPES() result["required"] |= { - "noise_type": ( - tuple( - t.name.lower() - for t in noise.NoiseType - if t is not noise.NoiseType.BROWNIAN - ), - ), + "noise_type": (tuple(NoiseType.get_names()),), } return result @@ -261,11 +250,7 @@ class SamplerNodeSonarBase: }, ), "rand_init_noise_type": ( - tuple( - t.name.lower() - for t in noise.NoiseType - if t is not noise.NoiseType.BROWNIAN - ), + tuple(NoiseType.get_names(skip=(NoiseType.BROWNIAN,))), ), }, "optional": { @@ -317,7 +302,7 @@ class SamplerNodeSonarEuler(SamplerNodeSonarBase): init=HistoryType[momentum_init.upper()], momentum_hist=momentum_hist, direction=direction, - rand_init_noise_type=noise.NoiseType[rand_init_noise_type.upper()], + rand_init_noise_type=NoiseType[rand_init_noise_type.upper()], guidance=guidance_cfg_opt, ) return ( @@ -347,7 +332,7 @@ class SamplerNodeSonarEulerAncestral(SamplerNodeSonarEuler): "round": False, }, ), - "noise_type": (tuple(t.name.lower() for t in noise.NoiseType),), + "noise_type": (tuple(NoiseType.get_names()),), }, ) result["optional"].update( @@ -375,8 +360,8 @@ class SamplerNodeSonarEulerAncestral(SamplerNodeSonarEuler): init=HistoryType[momentum_init.upper()], momentum_hist=momentum_hist, direction=direction, - rand_init_noise_type=noise.NoiseType[rand_init_noise_type.upper()], - noise_type=noise.NoiseType[noise_type.upper()], + rand_init_noise_type=NoiseType[rand_init_noise_type.upper()], + noise_type=NoiseType[noise_type.upper()], custom_noise=custom_noise_opt.clone() if custom_noise_opt else None, guidance=guidance_cfg_opt, ) @@ -408,7 +393,7 @@ class SamplerNodeSonarDPMPPSDE(SamplerNodeSonarEuler): "round": False, }, ), - "noise_type": (tuple(t.name.lower() for t in noise.NoiseType),), + "noise_type": (tuple(NoiseType.get_names(default=NoiseType.BROWNIAN)),), }, ) result["optional"].update( @@ -436,8 +421,8 @@ class SamplerNodeSonarDPMPPSDE(SamplerNodeSonarEuler): init=HistoryType[momentum_init.upper()], momentum_hist=momentum_hist, direction=direction, - rand_init_noise_type=noise.NoiseType[rand_init_noise_type.upper()], - noise_type=noise.NoiseType[noise_type.upper()], + rand_init_noise_type=NoiseType[rand_init_noise_type.upper()], + noise_type=NoiseType[noise_type.upper()], custom_noise=custom_noise_opt.clone() if custom_noise_opt else None, guidance=guidance_cfg_opt, ) @@ -497,7 +482,7 @@ class SamplerNodeConfigOverride: "sde_solver": (("midpoint", "heun"),), }, "optional": { - "noise_type": (tuple(t.name.lower() for t in noise.NoiseType),), + "noise_type": (tuple(NoiseType.get_names()),), "custom_noise_opt": ("SONAR_CUSTOM_NOISE",), }, } @@ -525,7 +510,7 @@ class SamplerNodeConfigOverride: | { "override_sampler_cfg": { "sampler": sampler, - "noise_type": noise.NoiseType[noise_type.upper()] + "noise_type": NoiseType[noise_type.upper()] if noise_type is not None else None, "custom_noise": custom_noise_opt, diff --git a/py/noise.py b/py/noise.py index e80d6b2..c4a4a79 100644 --- a/py/noise.py +++ b/py/noise.py @@ -15,9 +15,16 @@ from torch import FloatTensor, Generator, Tensor # ruff: noqa: D412, D413, D417, D212, D407, ANN002, ANN003, FBT001, FBT002, S311 -def scale_noise(noise, factor=1.0): - mean, std = noise.mean(), noise.std() - return (noise - mean).div_(std).mul_(factor) +def scale_noise(noise, factor=1.0, threshold_std_devs=2.5): + mean, std = noise.mean().item(), noise.std().item() + threshold = threshold_std_devs / math.sqrt(noise.numel()) + if abs(mean) > threshold: + noise -= mean + if abs(1.0 - std) > threshold: + noise /= std + if factor != 1.0: + noise *= factor + return noise class NoiseType(Enum): @@ -27,6 +34,7 @@ class NoiseType(Enum): PERLIN = auto() STUDENTT = auto() HIGHRES_PYRAMID = auto() + PYRAMID = auto() PINK = auto() LAPLACIAN = auto() POWER = auto() @@ -37,6 +45,15 @@ class NoiseType(Enum): # RAINBOW_INTENSE3 = auto() GREEN_TEST = auto() + @classmethod + def get_names(cls, default=None, skip=None): + if default is not None: + yield default.name.lower() + for nt in cls: + if nt == default or (skip and nt in skip): + continue + yield nt.name.lower() + class NoiseError(Exception): pass @@ -357,6 +374,36 @@ def highres_pyramid_noise_like(x, discount=0.7): return noise / noise.std() # Scaled back to roughly unit variance +def pyramid_noise_like(x, generator=None, device="cpu", discount=0.8): + size = x.size() + b, c, h, w = size + orig_h = h + orig_w = w + noise = torch.zeros(size=size, dtype=x.dtype, layout=x.layout, device=device) + r = 1 + for i in range(5): + r *= 2 # Rather than always going 2x, + noise += ( + torch.nn.functional.interpolate( + ( + torch.normal( + mean=0, + std=0.5**i, + size=(b, c, h * r, w * r), + dtype=x.dtype, + layout=x.layout, + generator=generator, + device=device, + ) + ), + size=(orig_h, orig_w), + mode="nearest-exact", + ) + * discount**i + ) + return noise.to(device=x.device) + + def studentt_noise_like(x): from torch.distributions import StudentT @@ -501,7 +548,7 @@ class NoiseSampler: else noise.mul_(self.factor) ) if hasattr(noise, "to"): - return noise.to(dtype=self.dtype, device=self.device) + noise = noise.to(dtype=self.dtype, device=self.device) return noise @@ -513,6 +560,7 @@ NOISE_SAMPLERS: dict[NoiseType, Callable] = { NoiseType.STUDENTT: NoiseSampler.simple(studentt_noise_like), NoiseType.PINK: NoiseSampler.simple(pink_noise_like), NoiseType.HIGHRES_PYRAMID: NoiseSampler.simple(highres_pyramid_noise_like), + NoiseType.PYRAMID: NoiseSampler.simple(pyramid_noise_like), NoiseType.RAINBOW_MILD: NoiseSampler.simple( lambda x: (green_noise_like(x) * 0.55 + rand_perlin_like(x) * 0.7) * 1.15, ), @@ -522,18 +570,6 @@ NOISE_SAMPLERS: dict[NoiseType, Callable] = { NoiseType.LAPLACIAN: NoiseSampler.simple(laplacian_noise_like), NoiseType.POWER: NoiseSampler.simple(power_noise_like), NoiseType.GREEN_TEST: NoiseSampler.simple(green_noise_like), - # NoiseType.RAINBOW_MILD2: lambda x: lambda _s, _sn: ( - # green_noise_like(x) * 0.55 + uniform_noise_like(x) * 0.7 - # ) - # * 1.15, - # NoiseType.RAINBOW_INTENSE2: lambda x: lambda _s, _sn: ( - # green_noise_like(x) * 0.75 + uniform_noise_like(x) * 0.5 - # ) - # * 1.15, - # NoiseType.RAINBOW_INTENSE3: lambda x: lambda _s, _sn: ( - # green_noise_like(x) * 0.75 + highres_pyramid_noise_like(x) * 0.5 - # ) - # * 1.15, } diff --git a/py/powernoise.py b/py/powernoise.py index f30a8b8..298af8c 100644 --- a/py/powernoise.py +++ b/py/powernoise.py @@ -314,7 +314,8 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): output_dir = folder_paths.get_temp_directory() prefix_append = "sonar_temp_" + "".join( - random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5) # noqa: S311 + random.choice("abcdefghijklmnopqrstupvxyz") # noqa: S311 + for x in range(5) ) full_output_folder, filename, counter, subfolder, _ = ( folder_paths.get_save_image_path(prefix_append, output_dir) diff --git a/py/sonar.py b/py/sonar.py index 8ee7ab8..b183ec0 100644 --- a/py/sonar.py +++ b/py/sonar.py @@ -3,6 +3,7 @@ from __future__ import annotations from enum import Enum, auto +from sys import stderr from typing import Any, Callable, NamedTuple import torch @@ -44,6 +45,8 @@ class SonarConfig(NamedTuple): class SonarBase: + DEFAULT_NOISE_TYPE = noise.NoiseType.GAUSSIAN + def __init__(self, cfg: SonarConfig) -> None: self.history_d = None self.cfg = cfg @@ -59,11 +62,11 @@ class SonarBase: sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() if noise_sampler is not None and self.cfg.noise_type not in ( None, - noise.NoiseType.GAUSSIAN, + self.DEFAULT_NOISE_TYPE, ): - # Possibly we should just use the supplied already-created noise sampler here. - raise ValueError( - "Unexpected noise_sampler presence with non-default noise type requested", + print( + "Sonar: Warning: Noise sampler supplied, overriding noise type from settings", + file=stderr, ) if self.cfg.custom_noise: noise_sampler = self.cfg.custom_noise.make_noise_sampler( @@ -72,9 +75,9 @@ class SonarBase: sigma_max, seed=seed, ) - elif noise_sampler is None and self.cfg.noise_type: + elif noise_sampler is None: noise_sampler = noise.get_noise_sampler( - self.cfg.noise_type, + self.cfg.noise_type or self.DEFAULT_NOISE_TYPE, x, sigma_min, sigma_max, @@ -386,14 +389,6 @@ class SonarEulerAncestral(SonarSampler): ): if sonar_config is None: sonar_config = SonarConfig() - if ( - noise_sampler is not None - and sonar_config.noise_type != noise.NoiseType.GAUSSIAN - ): - # Possibly we should just use the supplied already-created noise sampler here. - raise ValueError( - "Unexpected noise_sampler presence with non-default noise type requested", - ) s_in = x.new_ones([x.shape[0]]) sonar = cls( eta, @@ -430,6 +425,8 @@ class SonarEulerAncestral(SonarSampler): class SonarDPMPPSDE(SonarSampler): + DEFAULT_NOISE_TYPE = noise.NoiseType.BROWNIAN + def __init__( self, eta: float = 1.0, @@ -560,15 +557,6 @@ class SonarDPMPPSDE(SonarSampler): ): if sonar_config is None: sonar_config = SonarConfig() - if ( - noise_sampler is not None - and sonar_config.noise_type != noise.NoiseType.GAUSSIAN - ): - # Possibly we should just use the supplied already-created noise sampler here. - raise ValueError( - "Unexpected noise_sampler presence with non-default noise type requested", - ) - s_in = x.new_ones([x.shape[0]]) sonar = cls( eta,