Fix issue with dup seed argument in modulated noise sampler

* 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.
* Improve noise scaling (may change seeds).
* Add `KRestartSamplerCustomNoise` if the user has a recent enough version of ComfyUI_restart_sampling installed.

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.
@@ -100,6 +98,38 @@ From a usage perspective, using positive alpha will tend to create a colorful ef
 
 Noise from the `SonarCustomNoise` node and `SonarPowerNoise` can be freely mixed.
 
+### `SonarModulatedNoise`
+
+Experimental noise modulation based on code stolen from
+[ComfyUI-Extra-Samplers](https://github.com/Clybius/ComfyUI-Extra-Samplers). _Probably_ does not work correctly
+for normal sampling — I expect the modulation will be based on the tensor where the noise sampler was created
+rather than each step. However it may be useful for something like restart sampling noise
+(see `KRestartSamplerCustomNoise` below).
+
+*Note*: It's likely this node will be changed in the future.
+
+### `SonarRepeatedNoise`
+
+Experimental node to cache noise sampler results. Why would you want to do this? Some noise samplers are
+relatively slow (`pyramid` for example) or it may be slow to generate noise if you are mixing many types
+of noise. When `permute` is enabled, a random effect like flipping the noise or rolling it in some dimension
+will be chosen each time the noise sampler is called. I recommend leaving `permute` on. Note that repeated
+noise (especially with `permute` disabled) can be stronger than normal noise, so you may need to rescale to
+a value lower than `1.0` or decrease `s_noise` for the sampler.
+
+*Note*: It's likely this node will be changed in the future.
+
+### `KRestartSamplerCustomNoise`
+
+If you have a recent enough version of [ComfyUI_restart_sampling](https://github.com/ssitu/ComfyUI_restart_sampling/)
+installed, you'll also get the `KRestartSamplerCustomNoise` node which is exactly the same as `KRestartSamplerCustom`
+except for adding an optional custom noise input.
+See the restart sampling repo for more information: https://github.com/ssitu/ComfyUI_restart_sampling
+
+### `RestartSamplerCustomNoise`
+
+As above, except this is the custom sampler version.
+
 ## Sonar Sampler Parameters
 
 Very abbreviated section. The init type can make a big difference. If you use `RANDOM` you can get away with setting `direction` to high values (like up to `2.25` or so) and absurdly low values (like `-30.0`). It's also possible to set `momentum` and `momentum_hist` to negative values, although whether it's a good idea...
@@ -146,9 +176,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!
 

__init__.py

@@ -2,17 +2,9 @@ from .py import nodes, powernoise, sonar
 
 sonar.add_samplers()
 
-NODE_CLASS_MAPPINGS = {
-    "SamplerSonarEuler": nodes.SamplerNodeSonarEuler,
-    "SamplerSonarEulerA": nodes.SamplerNodeSonarEulerAncestral,
-    "SamplerSonarDPMPPSDE": nodes.SamplerNodeSonarDPMPPSDE,
-    "SamplerConfigOverride": nodes.SamplerNodeConfigOverride,
-    "NoisyLatentLike": nodes.NoisyLatentLikeNode,
-    "SonarCustomNoise": nodes.SonarCustomNoiseNode,
+NODE_CLASS_MAPPINGS = nodes.NODE_CLASS_MAPPINGS | {
     "SonarPowerNoise": powernoise.SonarPowerNoiseNode,
-    "SonarGuidanceConfig": nodes.GuidanceConfigNode,
 }
-
-NODE_DISPLAY_NAME_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = nodes.NODE_DISPLAY_NAME_MAPPINGS
 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

changelog.md

@@ -2,6 +2,20 @@
 
 Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top.
 
+## 20240506
+
+* Add `SonarModulatedNoise` and `SonarRepeatedNoise` nodes.
+
+## 20240327
+
+* 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.
+* Improve noise scaling (may change seeds).
+* Add `KRestartSamplerCustomNoise` if the user has a recent enough version of ComfyUI_restart_sampling installed.
+
 ## 20240320
 
 * `NoisyLatentLike` node improved to allow calculating strength with sigmas and injecting noise itself.

py/nodes.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 import abc
 import inspect
+from types import SimpleNamespace
 from typing import Any, Callable
 
 import torch
 from comfy import samplers
 
 from . import noise
+from .noise import NoiseType
 from .sonar import (
     GuidanceConfig,
     GuidanceType,
@@ -24,13 +26,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}),
@@ -69,7 +65,7 @@ class NoisyLatentLikeNode:
                 model = model.model
             latent_scale_factor = model.latent_format.scale_factor
             max_denoise = samplers.Sampler().max_denoise(
-                samplers.wrap_model(model),
+                SimpleNamespace(inner_model=model),
                 sigmas,
             )
             multiplier *= (
@@ -83,7 +79,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 +160,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
 
@@ -178,6 +168,65 @@ class SonarCustomNoiseNode(SonarCustomNoiseNodeBase):
         return noise.CustomNoiseItem
 
 
+class SonarModulatedNoiseNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sonar_custom_noise": ("SONAR_CUSTOM_NOISE",),
+                "modulation_type": (
+                    (
+                        "intensity",
+                        "frequency",
+                        "spectral_signum",
+                        "none",
+                    ),
+                ),
+                "dims": ("INT", {"default": 3, "min": 1, "max": 3}),
+                "strength": ("FLOAT", {"default": 2.0, "min": -100.0, "max": 100.0}),
+            },
+        }
+
+    RETURN_TYPES = ("SONAR_CUSTOM_NOISE",)
+    CATEGORY = "advanced/noise"
+    FUNCTION = "go"
+
+    def go(self, sonar_custom_noise, modulation_type, dims, strength):
+        return (
+            noise.ModulatedNoise(
+                sonar_custom_noise.make_noise_sampler,
+                modulation_type=modulation_type,
+                modulation_strength=strength,
+                modulation_dims=dims,
+            ),
+        )
+
+
+class SonarRepeatedNoiseNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sonar_custom_noise": ("SONAR_CUSTOM_NOISE",),
+                "repeat_length": ("INT", {"default": 8, "min": 1, "max": 100}),
+                "permute": ("BOOLEAN", {"default": True}),
+            },
+        }
+
+    RETURN_TYPES = ("SONAR_CUSTOM_NOISE",)
+    CATEGORY = "advanced/noise"
+    FUNCTION = "go"
+
+    def go(self, sonar_custom_noise, repeat_length, permute=True):
+        return (
+            noise.RepeatedNoise(
+                sonar_custom_noise.make_noise_sampler,
+                repeat_length,
+                permute=permute,
+            ),
+        )
+
+
 class GuidanceConfigNode:
     @classmethod
     def INPUT_TYPES(cls):
@@ -261,11 +310,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 +362,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 +392,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 +420,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 +453,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 +481,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 +542,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 +570,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,
@@ -598,3 +643,156 @@ class SamplerNodeConfigOverride:
             extra_args=extra_args,
             **kwargs,
         )
+
+
+NODE_CLASS_MAPPINGS = {
+    "SamplerSonarEuler": SamplerNodeSonarEuler,
+    "SamplerSonarEulerA": SamplerNodeSonarEulerAncestral,
+    "SamplerSonarDPMPPSDE": SamplerNodeSonarDPMPPSDE,
+    "SamplerConfigOverride": SamplerNodeConfigOverride,
+    "NoisyLatentLike": NoisyLatentLikeNode,
+    "SonarCustomNoise": SonarCustomNoiseNode,
+    "SonarModulatedNoise": SonarModulatedNoiseNode,
+    "SonarRepeatedNoise": SonarRepeatedNoiseNode,
+    "SonarGuidanceConfig": GuidanceConfigNode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {}
+
+try:
+    import custom_nodes.ComfyUI_restart_sampling as rs
+
+    if not hasattr(rs.restart_sampling, "DEFAULT_SEGMENTS"):
+        # Dumb test but this should only exist in restart sampling versions that
+        # support plugging in custom noise.
+        raise NotImplementedError  # noqa: TRY301
+
+    class KRestartSamplerCustomNoise:
+        @classmethod
+        def INPUT_TYPES(cls):
+            get_normal_schedulers = getattr(
+                rs.nodes,
+                "get_supported_normal_schedulers",
+                rs.nodes.get_supported_restart_schedulers,
+            )
+            return {
+                "required": {
+                    "model": ("MODEL",),
+                    "add_noise": (["enable", "disable"],),
+                    "noise_seed": (
+                        "INT",
+                        {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
+                    ),
+                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
+                    "sampler": ("SAMPLER",),
+                    "scheduler": (get_normal_schedulers(),),
+                    "positive": ("CONDITIONING",),
+                    "negative": ("CONDITIONING",),
+                    "latent_image": ("LATENT",),
+                    "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
+                    "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
+                    "return_with_leftover_noise": (["disable", "enable"],),
+                    "segments": (
+                        "STRING",
+                        {
+                            "default": rs.restart_sampling.DEFAULT_SEGMENTS,
+                            "multiline": False,
+                        },
+                    ),
+                    "restart_scheduler": (rs.nodes.get_supported_restart_schedulers(),),
+                    "chunked_mode": ("BOOLEAN", {"default": True}),
+                },
+                "optional": {
+                    "custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
+                },
+            }
+
+        RETURN_TYPES = ("LATENT", "LATENT")
+        RETURN_NAMES = ("output", "denoised_output")
+        FUNCTION = "sample"
+        CATEGORY = "sampling"
+
+        def sample(
+            self,
+            model,
+            add_noise,
+            noise_seed,
+            steps,
+            cfg,
+            sampler,
+            scheduler,
+            positive,
+            negative,
+            latent_image,
+            start_at_step,
+            end_at_step,
+            return_with_leftover_noise,
+            segments,
+            restart_scheduler,
+            chunked_mode=False,
+            custom_noise_opt=None,
+        ):
+            return rs.restart_sampling.restart_sampling(
+                model,
+                noise_seed,
+                steps,
+                cfg,
+                sampler,
+                scheduler,
+                positive,
+                negative,
+                latent_image,
+                segments,
+                restart_scheduler,
+                disable_noise=add_noise == "disable",
+                step_range=(start_at_step, end_at_step),
+                force_full_denoise=return_with_leftover_noise != "enable",
+                output_only=False,
+                chunked_mode=chunked_mode,
+                custom_noise=custom_noise_opt.make_noise_sampler
+                if custom_noise_opt
+                else None,
+            )
+
+    NODE_CLASS_MAPPINGS["KRestartSamplerCustomNoise"] = KRestartSamplerCustomNoise
+
+    if not hasattr(rs.restart_sampling, "RestartSampler"):
+        # Dumb test part II: The Dumbening
+        raise NotImplementedError  # noqa: TRY301
+
+    class RestartSamplerCustomNoise:
+        @classmethod
+        def INPUT_TYPES(cls):
+            return {
+                "required": {
+                    "sampler": ("SAMPLER",),
+                    "chunked_mode": ("BOOLEAN", {"default": True}),
+                },
+                "optional": {
+                    "custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
+                },
+            }
+
+        RETURN_TYPES = ("SAMPLER",)
+        FUNCTION = "go"
+        CATEGORY = "sampling/custom_sampling/samplers"
+
+        def go(self, sampler, chunked_mode, custom_noise_opt=None):
+            restart_options = {
+                "restart_chunked": chunked_mode,
+                "restart_wrapped_sampler": sampler,
+                "restart_custom_noise": None
+                if custom_noise_opt is None
+                else custom_noise_opt.make_noise_sampler,
+            }
+            restart_sampler = samplers.KSAMPLER(
+                rs.restart_sampling.RestartSampler.sampler_function,
+                extra_options=sampler.extra_options | restart_options,
+                inpaint_options=sampler.inpaint_options,
+            )
+            return (restart_sampler,)
+
+    NODE_CLASS_MAPPINGS["RestartSamplerCustomNoise"] = RestartSamplerCustomNoise
+except (ImportError, NotImplementedError):
+    pass

py/noise.py

@@ -11,13 +11,21 @@ from typing import Callable
 import torch
 from comfy.k_diffusion import sampling
 from torch import FloatTensor, Generator, Tensor
+from torch.distributions import StudentT
 
 # 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 +35,7 @@ class NoiseType(Enum):
     PERLIN = auto()
     STUDENTT = auto()
     HIGHRES_PYRAMID = auto()
+    PYRAMID = auto()
     PINK = auto()
     LAPLACIAN = auto()
     POWER = auto()
@@ -37,6 +46,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,9 +375,37 @@ def highres_pyramid_noise_like(x, discount=0.7):
     return noise / noise.std()  # Scaled back to roughly unit variance
 
 
-def studentt_noise_like(x):
-    from torch.distributions import StudentT
+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):
     noise = StudentT(loc=0, scale=0.2, df=1).rsample(x.size())
     s: FloatTensor = torch.quantile(noise.flatten(start_dim=1).abs(), 0.75, dim=-1)
     s = s.reshape(*s.shape, 1, 1, 1)
@@ -501,10 +547,276 @@ 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
 
 
+class RepeatedNoise:
+    def __init__(self, noise_sampler, repeat_length, permute=True):
+        self.noise_sampler = noise_sampler
+        self.repeat_length = repeat_length
+        self.permute = permute
+
+    def clone(self):
+        return RepeatedNoise(self.noise_sampler, self.repeat_length)
+
+    def make_noise_sampler(self, x, *args, **kwargs):
+        ns = self.noise_sampler(x, *args, **kwargs)
+        noise_items = []
+        permute_options = 2
+        u32_max = 0xFFFF_FFFF
+        seed = kwargs.get("seed")
+        if seed is None:
+            seed = torch.randint(
+                -u32_max,
+                u32_max,
+                (1,),
+                device="cpu",
+                dtype=torch.int64,
+            ).item()
+        gen = torch.Generator(device="cpu")
+        gen.manual_seed(seed)
+
+        def noise_sampler(s, sn):
+            rands = torch.randint(
+                u32_max,
+                (4,),
+                generator=gen,
+                dtype=torch.uint32,
+            ).tolist()
+            if len(noise_items) < self.repeat_length:
+                idx = len(noise_items)
+                noise_items.append(ns(s, sn))
+            else:
+                idx = rands[0] % self.repeat_length
+            noise = noise_items[idx]
+            if not self.permute:
+                return noise.clone()
+            noise_dims = len(noise.shape)
+            match rands[1] % permute_options:
+                case 0:
+                    if rands[2] <= u32_max // 10:
+                        # 10% of the time we return the original tensor instead of flipping
+                        noise = noise.clone()
+                    else:
+                        dim = -1 + (rands[2] % (noise_dims + 1))
+                        noise = torch.flip(noise, (dim,))
+                case 1:
+                    dim = rands[2] % noise_dims
+                    count = rands[3] % noise.shape[dim]
+                    noise = torch.roll(noise, count, dims=(dim,)).clone()
+            return noise
+
+        return noise_sampler
+
+
+# Modulated noise functions copied from https://github.com/Clybius/ComfyUI-Extra-Samplers
+# They probably don't work correctly for normal sampling.
+class ModulatedNoise:
+    MODULATION_DIMS = (-3, (-2, -1), (-3, -2, -1))
+
+    def __init__(
+        self,
+        noise_sampler,
+        modulation_type="none",
+        modulation_strength=2.0,
+        modulation_dims=3,
+    ):
+        self.noise_sampler = noise_sampler
+        self.dims = self.MODULATION_DIMS[modulation_dims - 1]
+        self.type = modulation_type
+        self.strength = modulation_strength
+        match self.type:
+            case "intensity":
+                self.modulation_function = self.intensity_based_multiplicative_noise
+            case "frequency":
+                self.modulation_function = self.frequency_based_noise
+            case "spectral_signum":
+                self.modulation_function = self.spectral_modulate_noise
+            case _:
+                self.modulation_function = None
+
+    def clone(self):
+        return ModulatedNoise(self.noise_sampler, self.type, self.strength, self.dims)
+
+    def make_noise_sampler(self, x, *args, **kwargs):
+        ns = self.noise_sampler(x, *args, **kwargs)
+        if not self.modulation_function:
+            return ns
+        s_noise = sigma_up = 1.0
+        return lambda s, sn: self.modulation_function(
+            x,
+            ns(s, sn),
+            s_noise,
+            sigma_up,
+            self.strength,
+            self.dims,
+        )
+
+    @staticmethod
+    def intensity_based_multiplicative_noise(
+        x,
+        noise,
+        s_noise,
+        sigma_up,
+        intensity,
+        dims,
+    ) -> torch.Tensor:
+        """Scales noise based on the intensities of the input tensor."""
+        std = torch.std(
+            x - x.mean(),
+            dim=dims,
+            keepdim=True,
+        )  # Average across channels to get intensity
+        scaling = (
+            1 / (std * abs(intensity) + 1.0)
+        )  # Scale std by intensity, as not doing this leads to more noise being left over, leading to crusty/preceivably extremely oversharpened images
+        additive_noise = noise * s_noise * sigma_up
+        scaled_noise = noise * s_noise * sigma_up * scaling + additive_noise
+
+        noise_norm = torch.norm(additive_noise)
+        scaled_noise_norm = torch.norm(scaled_noise)
+        scaled_noise *= noise_norm / scaled_noise_norm  # Scale to normal noise strength
+        return scaled_noise * intensity + additive_noise * (1 - intensity)
+
+    @staticmethod
+    def frequency_based_noise(
+        z_k,
+        noise,
+        s_noise,
+        sigma_up,
+        intensity,
+        channels,
+    ) -> torch.Tensor:
+        """Scales the high-frequency components of the noise based on the given intensity."""
+        additive_noise = noise * s_noise * sigma_up
+
+        std = torch.std(
+            z_k - z_k.mean(),
+            dim=channels,
+            keepdim=True,
+        )  # Average across channels to get intensity
+        scaling = 1 / (std * abs(intensity) + 1.0)
+        # Perform Fast Fourier Transform (FFT)
+        z_k_freq = torch.fft.fft2(scaling * additive_noise + additive_noise)
+
+        # Get the magnitudes of the frequency components
+        magnitudes = torch.abs(z_k_freq)
+
+        # Create a high-pass filter (emphasize high frequencies)
+        h, w = z_k.shape[-2:]
+        b = abs(
+            intensity,
+        )  # Controls the emphasis of the high pass (higher frequencies are boosted)
+        high_pass_filter = 1 - torch.exp(
+            -((torch.arange(h)[:, None] / h) ** 2 + (torch.arange(w)[None, :] / w) ** 2)
+            * b**2,
+        )
+        high_pass_filter = high_pass_filter.to(z_k.device)
+
+        # Apply the filter to the magnitudes
+        magnitudes_scaled = magnitudes * (1 + high_pass_filter)
+
+        # Reconstruct the complex tensor with scaled magnitudes
+        z_k_freq_scaled = magnitudes_scaled * torch.exp(1j * torch.angle(z_k_freq))
+
+        # Perform Inverse Fast Fourier Transform (IFFT)
+        z_k_scaled = torch.fft.ifft2(z_k_freq_scaled)
+
+        # Return the real part of the result
+        z_k_scaled = torch.real(z_k_scaled)
+
+        noise_norm = torch.norm(additive_noise)
+        scaled_noise_norm = torch.norm(z_k_scaled)
+
+        z_k_scaled *= noise_norm / scaled_noise_norm  # Scale to normal noise strength
+
+        return z_k_scaled * intensity + additive_noise * (1 - intensity)
+
+    @staticmethod
+    def spectral_modulate_noise(
+        _unused,
+        noise,
+        s_noise,
+        sigma_up,
+        intensity,
+        channels,
+        spectral_mod_percentile=5.0,
+    ) -> torch.Tensor:  # Modified for soft quantile adjustment using a novel:tm::c::r: method titled linalg.
+        additive_noise = noise * s_noise * sigma_up
+        # Convert image to Fourier domain
+        fourier = torch.fft.fftn(
+            additive_noise,
+            dim=channels,
+        )  # Apply FFT along Height and Width dimensions
+
+        log_amp = torch.log(torch.sqrt(fourier.real**2 + fourier.imag**2))
+
+        quantile_low = (
+            torch.quantile(
+                log_amp.abs().flatten(1),
+                spectral_mod_percentile * 0.01,
+                dim=1,
+            )
+            .unsqueeze(-1)
+            .unsqueeze(-1)
+            .expand(log_amp.shape)
+        )
+
+        quantile_high = (
+            torch.quantile(
+                log_amp.abs().flatten(1),
+                1 - (spectral_mod_percentile * 0.01),
+                dim=1,
+            )
+            .unsqueeze(-1)
+            .unsqueeze(-1)
+            .expand(log_amp.shape)
+        )
+
+        quantile_max = (
+            torch.quantile(log_amp.abs().flatten(1), 1, dim=1)
+            .unsqueeze(-1)
+            .unsqueeze(-1)
+            .expand(log_amp.shape)
+        )
+
+        # Decrease high-frequency components
+        mask_high = log_amp > quantile_high  # If we're larger than 95th percentile
+
+        additive_mult_high = torch.where(
+            mask_high,
+            1
+            - ((log_amp - quantile_high) / (quantile_max - quantile_high)).clamp_(
+                max=0.5,
+            ),  # (1) - (0-1), where 0 is 95th %ile and 1 is 100%ile
+            torch.tensor(1.0),
+        )
+
+        # Increase low-frequency components
+        mask_low = log_amp < quantile_low
+        additive_mult_low = torch.where(
+            mask_low,
+            1
+            + (1 - (log_amp / quantile_low)).clamp_(
+                max=0.5,
+            ),  # (1) + (0-1), where 0 is 5th %ile and 1 is 0%ile
+            torch.tensor(1.0),
+        )
+
+        mask_mult = (additive_mult_low * additive_mult_high) ** intensity
+        # print(mask_mult)
+        filtered_fourier = fourier * mask_mult
+
+        # Inverse transform back to spatial domain
+        inverse_transformed = torch.fft.ifftn(
+            filtered_fourier,
+            dim=channels,
+        )  # Apply IFFT along Height and Width dimensions
+
+        return inverse_transformed.real.to(additive_noise.device)
+
+
 NOISE_SAMPLERS: dict[NoiseType, Callable] = {
     NoiseType.BROWNIAN: NoiseSampler.wrap(sampling.BrownianTreeNoiseSampler),
     NoiseType.GAUSSIAN: NoiseSampler.simple(torch.randn_like),
@@ -513,6 +825,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 +835,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,
 }
 
 

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)

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,