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

@@ -4,6 +4,8 @@ A janky implementation of Sonar sampling (momentum-based sampling) for [ComfyUI]
 
 Currently supports Euler, Euler Ancestral, and DPM++ SDE sampling.
 
+See the [ChangeLog](changelog.md) for recent user-visible changes.
+
 ## Description
 
 See https://github.com/Kahsolt/stable-diffusion-webui-sonar for a more in-depth explanation.
@@ -22,13 +24,113 @@ You can also just choose `sonar_euler`, `sonar_euler_ancestral` or `sonar_dpmpp_
 
 ## Nodes
 
-1. `SamplerSonarEuler` — Custom sampler node that combines Euler sampling and momentum and optionally guidance. A bit boring compared to the ancestral version but it has predictability going for it. You can possibly try setting init type to `RAND` and using different noise types, however this sampler seems _very_ sensitive to that init type. You may want to set direction to a very low value like `0.05` or `-0.15` when using the `RAND` init type.
-2. `SamplerSonarEulerAncestral` — Ancestral version of the above. Same features, just with ancestral Euler.
-4. `SonarGuidanceConfig` — You can optionally plug this into the Sonar sampler nodes. See the [Guidance](#guidance) section below.
-5. `NoisyLatentLike` — If you give it a latent (or latent batch) it'll return a noisy latent of the same shape. Allows specifying all the custom noise types except `brownian` which has some special requirements. Provided just because the noise generation functions are conveniently available. You can also use this as a reference latent with `SonarGuidanceConfig` node and depending on the strength it can act like variation seed (you'd change the seed in the `NoisyLatentLike` node). *Note*: The seed stuff may or may not work correctly.
-6. `SamplerSonarDPMPPSDE` — This one is extra experimental but it is an attempt to add moment and guidance to the DPM++ SDE sampler. It may not work correctly but you can sample stuff with it and get interesting results. I actually really like this one, and you can get away with more extreme stuff like `green_test` noise and still produce reasonable results. You may want to use the `BlehDiscardPenultimateSigma` node from my [ComfyUI-bleh](https://github.com/blepping/ComfyUI-bleh) collection if you find the result seems a bit washed out and b lurry.
+### `SamplerSonarEuler`
 
-## Parameters
+Custom sampler node that combines Euler sampling and momentum and optionally guidance. A bit boring compared to the ancestral version but it has predictability going for it. You can possibly try setting init type to `RAND` and using different noise types, however this sampler seems _very_ sensitive to that init type. You may want to set direction to a very low value like `0.05` or `-0.15` when using the `RAND` init type. Setting `momentum=1` is the same as disabling momentum, so this sampler with `momentum=1` is basically the same as the basic `euler` sampler.
+
+### `SamplerSonarEulerAncestral`
+
+Ancestral version of the above. Same features, just with ancestral Euler.
+
+### `SamplerSonarDPMPPSDE`
+
+Attempt to add momentum and guidance to the DPM++ SDE sampler. It may not work correctly but you can sample stuff with it and get interesting results. I actually really like this one, and you can get away with more extreme stuff like `green_test` noise and still produce reasonable results. You may want to use the `BlehDiscardPenultimateSigma` node from my [ComfyUI-bleh](https://github.com/blepping/ComfyUI-bleh) collection if you find the result seems a bit washed out and blurry.
+
+### `SonarGuidanceConfig`
+
+You can optionally plug this into the Sonar sampler nodes. See the [Guidance](#guidance) section below.
+
+### `NoisyLatentLike`
+
+This node takes a reference latent and generates noise of the same shape. The one required input is `latent`.
+
+You can connect a `SonarCustomNoise` or `SonerPowerNoise` node to the `custom_noise_opt` input: if that is attached, the built in noise type selector is ignored. The generated noise will be multiplied by the `multiplier` value. Note that you cannot use `brownian` noise whether specified directly or via custom noise nodes.
+
+The node has two main modes: simply generate and scale the noise by the multiplier and return or add it to the input latent. In this mode, you don't connect anything to the `mul_by_sigmas_opt` or `model_opt` inputs and you would use other nodes to calculate the correct strength.
+
+In the second mode you must connect sigmas (for example from a `BasicScheduler` node) to the `mul_by_sigmas_opt` input and connect a model to the `model_opt` input. It will calculate the strength based on the first item in the list of sigmas (so you could use something like a `SplitSigmas` node to slice them as needed). Note that `multiplier` still applies: the calculated strength will be scaled by it. This second mode is generally this is the most convenient way to use the node since the two main uses cases are: making a latent with initial noise or adding noise to a latent (for img2img type stuff).
+
+If you want to create noise for initial sampling, connect model and sigmas to the node, connect an empty latent (or one of the appropriate size) to it and that is basically all you need to do (aside from configuring the noise types). For img2img (upscaling, etc), either slice the sigmas at the appropriate or set a denoise in something like the `BasicScheduler` node. **Note**: You also need to turn on the `add_to_latent` toggle. Turning this on doesn't matter for initial noise since an empty latent is all zeros.
+
+
+### `SamplerConfigOverride`
+
+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.
+
+### `SonarCustomNoise`
+
+See the [Noise](#noise) section below for information on noise types.
+
+### `SonarPowerNoise`
+
+This node generates [fractional Brownian motion (fBm) noise](https://en.wikipedia.org/wiki/Fractional_Brownian_motion#Frequency-domain_interpretation). It offers versatility in producing various types of noise including gaussian, pink, 2D brownian noise, and all intermediates.
+
+By default, the node generates normal gaussian noise.
+
+<details>
+
+<summary>Expand detailed explanation</summary>
+
+
+Here's an overview of its parameters:
+
+- `factor` and `rescale` operate similarly to `SonarCustomNoise`, enabling the addition of multiple sources of noises.
+- `time_brownian` introduces correlation across sampler timesteps for SDE solvers.
+- `alpha` is the main parameter. `alpha > 0` amplifies low frequencies; `alpha = 1` yields pink noise, and `alpha = 2` produces brownian noise. Conversely, for `alpha < 0`, it amplifies high frequencies.
+- `min_freq` and `max_freq` determine the range of frequencies allowed through. Setting `max_freq = `$\sqrt{1/2} \simeq 0.7071$ enables the passage of the highest frequencies. In cases where `alpha < 0`, setting `max_freq = 0.5` is advisable to diminish the power of diagonally oriented frequencies.
+- `stretch`, `rotate`, and `pnorm` alter the filter's shape by stretching, rotating, or cushioning the band-pass region.
+- Lowering `mix` moderates the filter's effect by blending back unfiltered gaussian noise from the same sample.
+- `common_mode` is an attempt to desaturate the latent by injecting the average across channels into every latent channel. However, this may result in a specific color due to the encoding of the unit vector by the latent space. Note that this is done _after_ the `mix`ing of unfiltered gaussian noise.
+- Enabling `preview` provides a visual representation of the filter. `no_mix` sets `mix = 1` for the preview. The preview includes, from left to right:
+  - Fourier domain visualization: Low frequencies at the center, with black indicating filtered-out frequencies.
+  - Spatial visualization of the 2D kernel: The filtering can be interpreted as convolution with the displayed kernel.
+  - Sample: Gaussian sample with shaped frequency spectrum. A single latent channel will look like this.
+
+**Frequency-domain Interpretation**: The Fourier transform decomposes a 2D latent into sinusoids covering all spatial orientations and frequencies. For an independent and identically distributed gaussian sample, energy is evenly distributed across all frequencies and orientations. Scaling the power spectrum by $1 / f^\alpha$, where $\alpha>0$, boosts low frequencies, introducing spatial correlations.
+
+**Spatial Domain Interpretation**: A gaussian latent sample comprises independently sampled pixels, exhibiting no spatial correlations. Conversely, a requirement that each pixel value differs from its neighbors by a $\epsilon \sim \mathcal{N}(0, 1)$ results in 2D brownian noise ($\alpha=2$).
+
+**Seed Considerations**: While the node defaults to outputting gaussian noise, a given seed produce a different sample than the one produced by other gaussian noise sources. This stems from sampling the noise directly in the frequency domain to avoid the cost of a FFT. When `time_brownian = true`, noise sampling occurs in the spatial domain, ensuring that default parameters yield output equivalent to `SonarCustomNoise` set to `brownian`.
+
+</details>
+
+From a usage perspective, using positive alpha will tend to create a colorful effect, using negative alpha will create line/streak like artifacts sort of like an oil painting canvas. Start with small values at first (`-0.1`, `0.1`) and adjust as necessary. `time_brownian` makes the effect of power noise (and alpha) stronger - also note that it can only be used when sampling and not for `NoisyLatentLike`. Setting `common_mode` also generally seems to intensify these effects. Different types of models (normal EPS models, v-prediction models, SDXL) generally react differently to these exotic noise types so my advice is to experiment! Lowering `mix` uses normal gaussian noise for part of the generated noise. For example, `mix=1.0` means 100% power noise, `mix=0.5` means 50/50 power noise and normal gaussian noise. This also is about the same as setting factor to `0.5` and plugging in a `SonarCustomNoise` node with factor at `0.5` also and the type set to `guassian`.
+
+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...
 
@@ -51,30 +153,39 @@ I basically just copied a bunch of noise functions without really knowing what t
 3. `brownian`: This is the noise type SDE samplers use.
 4. `perlin`
 5. `studentt`: There's a comment that says it may enhance subject details. It seemed to produce a fairly dark result.
-6. `studentt_test`: An experiment that may be removed, it doesn't seem to be adding enough noise. You can possibly compensate by increasing `s_noise`.
-7. `pink`
-8. `highres_pyramid`: Not extensively tested, but it is slower than the other noise types. I would guess it does something like enhance details.
-9. `laplacian`
-10. `power`
-11. `rainbow_mild` and `rainbow_intense`: A combination of green (-ish, the implementation may be broken) noise plus perlin noise. Very colorful results.
-12. `green_test`: Even more rainbow-y than the rainbow noise types. It _probably_ isn't working correctly, but the results are very interesting and colorful. Depending on the model, it may not work well for an initial generation but may be worth trying with img2img type workflows.
+6. `pink`
+7. `highres_pyramid`: Not extensively tested, but it is slower than the other noise types. I would guess it does something like enhance details.
+8. `laplacian`
+9. `power`
+10. `rainbow_mild` and `rainbow_intense`: A combination of green (-ish, the implementation may be broken) noise plus perlin noise. Very colorful results.
+11. `green_test`: Even more rainbow-y than the rainbow noise types. It _probably_ isn't working correctly, but the results are very interesting and colorful. Depending on the model, it may not work well for an initial generation but may be worth trying with img2img type workflows.
 
 You can scroll down to the the [Examples](#examples) section near the bottom to see some example generations with different noise types.
 
 The sampler and `NoisyLatentLike` nodes now take an optional `SonarCustomNoise` input. You can chain `SonarCustomNoise` nodes together to mix different types of noise, similar to how some of the built in ones. It shouldn't matter what order the noise types are chained. If `rescale` is set to `0.0` no rescaling will occur. `factor` is the proportion of that type of noise you want. If you want to use `rescale` it should be on the node that you are plugging into a sampler. Just for example if you had two `SonarCustomNoise` nodes both with `factor=0.7` and `rescale=1.0` on the last one, it would be effectively the same as if you'd used `factor=0.5` and `rescale=1.0` doesn't actually do anything. You can also rescale to values above `1.0` — the result is more noise, similar to increasing `s_noise` above `1.0` on a sampler. The simple explanation is `rescale` means you don't have to make sure the `factor`s add up to the scale you want (which normally would be `1.0`).
 
-**Note**: If you connect the optional `SonarCustomNoise` node to a Sonar sampler or the `NoisyLatentLike` node it will override the noise type selected in the node.
+**Note**: If you connect the optional `SonarCustomNoise` node to a Sonar sampler, the `NoisyLatentLike` node or the `SamplerConfigOverride` node, it will override the noise type selected in the node.
+
+
+
+## Related
+
+I also have some other ComfyUI nodes here: https://github.com/blepping/ComfyUI-bleh/
 
 ## Credits
 
 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!
 
 ## Examples
 
+Unfortunately, right now these examples are somewhat incomplete and out of date. I hope to update them when I get the time.
+
 ### Guidance
 
 <details>
@@ -141,11 +252,15 @@ Normal (non-sonar) Eular A. Not really a comparison with noise (think it would u
 
 #### StudentT
 
+**outdated**
+
 ![StudentT](assets/example_images/noise/renoise_studentt.png)
 
 
 #### StudentT_test
 
+**outdated**
+
 ![StudentT_test](assets/example_images/noise/renoise_studentt_test.png)
 
 #### Laplacian
@@ -164,7 +279,7 @@ Normal (non-sonar) Eular A. Not really a comparison with noise (think it would u
 
 ![Rainbow Intense](assets/example_images/noise/renoise_rainbow_intense.png)
 
-#### Green_test_
+#### Green_test
 
 ![Green_test](assets/example_images/noise/renoise_green_test.png)
 
@@ -203,10 +318,14 @@ These were generated with `s_noise=1.1` to make the noise effect more pronounced
 
 #### StudentT
 
+**outdated**
+
 ![StudentT](assets/example_images/noise/noise_studentt.png)
 
 #### StudentT_test
 
+**outdated**
+
 ![StudentT_test](assets/example_images/noise/noise_studentt_test.png)
 
 #### Laplacian

__init__.py

@@ -1,16 +1,10 @@
-from .py import nodes, sonar
+from .py import nodes, powernoise, sonar
 
 sonar.add_samplers()
 
-NODE_CLASS_MAPPINGS = {
-    "SamplerSonarEuler": nodes.SamplerNodeSonarEuler,
-    "SamplerSonarEulerA": nodes.SamplerNodeSonarEulerAncestral,
-    "SamplerSonarDPMPPSDE": nodes.SamplerNodeSonarDPMPPSDE,
-    "NoisyLatentLike": nodes.NoisyLatentLikeNode,
-    "SonarCustomNoise": nodes.SonarCustomNoiseNode,
-    "SonarGuidanceConfig": nodes.GuidanceConfigNode,
+NODE_CLASS_MAPPINGS = nodes.NODE_CLASS_MAPPINGS | {
+    "SonarPowerNoise": powernoise.SonarPowerNoiseNode,
 }
-
-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,34 @@
 
 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.
+
+## 20240314
+
+* `SonarPowerNoise` node added.
+
+## 20240227
+
+* Refactored noise generation functions (will break seeds).
+* Added `SamplerOverride` node.
+* `studentt` noise type replaced with `studentt_test` (the more correct version).
+
 ## 20240210
 
 * Added `SonarCustomNoise` node.

py/nodes.py

@@ -1,9 +1,15 @@
 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,
@@ -20,18 +26,16 @@ 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}),
+                "add_to_latent": ("BOOLEAN", {"default": False}),
             },
             "optional": {
                 "custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
+                "mul_by_sigmas_opt": ("SIGMAS",),
+                "model_opt": ("MODEL",),
             },
         }
 
@@ -42,21 +46,45 @@ class NoisyLatentLikeNode:
 
     def go(
         self,
-        noise_type,
-        seed,
-        latent,
-        custom_noise_opt=None,
+        noise_type: str,
+        seed: None | int,
+        latent: dict,
+        multiplier: float = 1.0,
+        add_to_latent=False,
+        custom_noise_opt: object | None = None,
+        mul_by_sigmas_opt: None | torch.Tensor = None,
+        model_opt: object | None = None,
     ):
+        model, sigmas = model_opt, mul_by_sigmas_opt
+        if sigmas is not None and len(sigmas) > 0:
+            if model is None:
+                raise ValueError(
+                    "NoisyLatentLike requires a model when sigmas are connected!",
+                )
+            while hasattr(model, "model"):
+                model = model.model
+            latent_scale_factor = model.latent_format.scale_factor
+            max_denoise = samplers.Sampler().max_denoise(
+                SimpleNamespace(inner_model=model),
+                sigmas,
+            )
+            multiplier *= (
+                float(
+                    torch.sqrt(1.0 + sigmas[0] ** 2.0) if max_denoise else sigmas[0],
+                )
+                / latent_scale_factor
+            )
+        latent_samples = latent["samples"]
         if custom_noise_opt is not None:
-            ns = custom_noise_opt.make_noise_sampler(latent["samples"])
+            ns = custom_noise_opt.make_noise_sampler(latent_samples)
         else:
             ns = noise.get_noise_sampler(
-                noise.NoiseType[noise_type.upper()],
-                latent["samples"],
+                NoiseType[noise_type.upper()],
+                latent_samples,
                 None,
                 None,
-                seed=None,
-                use_cpu=True,
+                seed=seed,
+                cpu=True,
             )
         randst = torch.random.get_rng_state()
         try:
@@ -64,10 +92,18 @@ class NoisyLatentLikeNode:
             result = ns(None, None)
         finally:
             torch.random.set_rng_state(randst)
+        if multiplier != 1.0:
+            result *= multiplier
+        if add_to_latent:
+            result += latent_samples.to(result.device)
         return ({"samples": result},)
 
 
-class SonarCustomNoiseNode:
+class SonarCustomNoiseNodeBase(abc.ABC):
+    @abc.abstractmethod
+    def get_item_class(self):
+        raise NotImplementedError
+
     @classmethod
     def INPUT_TYPES(cls):
         return {
@@ -92,13 +128,6 @@ class SonarCustomNoiseNode:
                         "round": False,
                     },
                 ),
-                "noise_type": (
-                    tuple(
-                        t.name.lower()
-                        for t in noise.NoiseType
-                        if t is not noise.NoiseType.BROWNIAN
-                    ),
-                ),
             },
             "optional": {
                 "sonar_custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
@@ -109,17 +138,95 @@ class SonarCustomNoiseNode:
     CATEGORY = "advanced/noise"
     FUNCTION = "go"
 
-    def go(self, factor, rescale, noise_type, sonar_custom_noise_opt=None):
+    def go(
+        self,
+        factor,
+        rescale,
+        sonar_custom_noise_opt=None,
+        **kwargs: dict[str, Any],
+    ):
         nis = (
             sonar_custom_noise_opt.clone()
             if sonar_custom_noise_opt
-            else noise.CustomNoise()
+            else noise.CustomNoiseChain()
         )
         if factor != 0:
-            nis.add(noise.CustomNoiseItem(factor, noise_type))
+            nis.add(self.get_item_class()(factor, **kwargs))
         return (nis if rescale == 0 else nis.rescaled(rescale),)
 
 
+class SonarCustomNoiseNode(SonarCustomNoiseNodeBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        result = super().INPUT_TYPES()
+        result["required"] |= {
+            "noise_type": (tuple(NoiseType.get_names()),),
+        }
+        return result
+
+    def get_item_class(self):
+        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):
@@ -203,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": {
@@ -259,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 (
@@ -289,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(
@@ -317,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,
         )
@@ -350,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(
@@ -378,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,
         )
@@ -393,3 +496,303 @@ class SamplerNodeSonarDPMPPSDE(SamplerNodeSonarEuler):
                 },
             ),
         )
+
+
+class SamplerNodeConfigOverride:
+    KWARG_OVERRIDES = ("s_noise", "eta", "s_churn", "r", "solver_type")
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sampler": ("SAMPLER",),
+                "eta": (
+                    "FLOAT",
+                    {
+                        "default": 1.0,
+                        "step": 0.01,
+                        "round": False,
+                    },
+                ),
+                "s_noise": (
+                    "FLOAT",
+                    {
+                        "default": 1.0,
+                        "step": 0.01,
+                        "round": False,
+                    },
+                ),
+                "s_churn": (
+                    "FLOAT",
+                    {
+                        "default": 0.0,
+                        "min": 0.0,
+                        "step": 0.01,
+                        "round": False,
+                    },
+                ),
+                "r": (
+                    "FLOAT",
+                    {
+                        "default": 0.5,
+                        "step": 0.01,
+                        "round": False,
+                    },
+                ),
+                "sde_solver": (("midpoint", "heun"),),
+            },
+            "optional": {
+                "noise_type": (tuple(NoiseType.get_names()),),
+                "custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
+            },
+        }
+
+    RETURN_TYPES = ("SAMPLER",)
+    CATEGORY = "sampling/custom_sampling/samplers"
+
+    FUNCTION = "get_sampler"
+
+    def get_sampler(
+        self,
+        sampler,
+        eta,
+        s_noise,
+        s_churn,
+        r,
+        sde_solver,
+        noise_type=None,
+        custom_noise_opt=None,
+    ):
+        return (
+            samplers.KSAMPLER(
+                self.sampler_function,
+                extra_options=sampler.extra_options
+                | {
+                    "override_sampler_cfg": {
+                        "sampler": sampler,
+                        "noise_type": NoiseType[noise_type.upper()]
+                        if noise_type is not None
+                        else None,
+                        "custom_noise": custom_noise_opt,
+                        "s_noise": s_noise,
+                        "eta": eta,
+                        "s_churn": s_churn,
+                        "r": r,
+                        "solver_type": sde_solver,
+                    },
+                },
+                inpaint_options=sampler.inpaint_options | {},
+            ),
+        )
+
+    @classmethod
+    @torch.no_grad()
+    def sampler_function(
+        cls,
+        model,
+        x,
+        sigmas,
+        *args: list[Any],
+        override_sampler_cfg: dict[str, Any] | None = None,
+        noise_sampler: Callable | None = None,
+        extra_args: dict[str, Any] | None = None,
+        **kwargs: dict[str, Any],
+    ):
+        if not override_sampler_cfg:
+            raise ValueError("Override sampler config missing!")
+        if extra_args is None:
+            extra_args = {}
+        cfg = override_sampler_cfg
+        sampler, noise_type, custom_noise = (
+            cfg["sampler"],
+            cfg.get("noise_type"),
+            cfg.get("custom_noise"),
+        )
+        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+        seed = extra_args.get("seed")
+        if custom_noise is not None:
+            noise_sampler = custom_noise.make_noise_sampler(
+                x,
+                sigma_min,
+                sigma_max,
+                seed=seed,
+            )
+        elif noise_type is not None:
+            noise_sampler = noise.get_noise_sampler(
+                noise_type,
+                x,
+                sigma_min,
+                sigma_max,
+                seed=seed,
+                cpu=True,
+            )
+        sig = inspect.signature(sampler.sampler_function)
+        params = sig.parameters
+        kwargs = kwargs | {}
+        if "noise_sampler" in params:
+            kwargs["noise_sampler"] = noise_sampler
+        for k in cls.KWARG_OVERRIDES:
+            if k not in params or cfg.get(k) is None:
+                continue
+            kwargs[k] = cfg[k]
+        return sampler.sampler_function(
+            model,
+            x,
+            sigmas,
+            *args,
+            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

@@ -1,25 +1,41 @@
 # Noise generation functions shamelessly yoinked from https://github.com/Clybius/ComfyUI-Extra-Samplers
 from __future__ import annotations
 
+import abc
+import functools as fun
 import math
+import operator as op
 from enum import Enum, auto
 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, 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):
     GAUSSIAN = auto()
     UNIFORM = auto()
     BROWNIAN = auto()
     PERLIN = auto()
     STUDENTT = auto()
-    STUDENTT_TEST = auto()
     HIGHRES_PYRAMID = auto()
+    PYRAMID = auto()
     PINK = auto()
     LAPLACIAN = auto()
     POWER = auto()
@@ -30,24 +46,79 @@ 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
 
 
-class CustomNoiseItem:
-    def __init__(self, factor, noise_type):
+class CustomNoiseItemBase(abc.ABC):
+    def __init__(self, factor, **kwargs):
         self.factor = factor
-        self.noise_type = noise_type
+        self.keys = set(kwargs.keys())
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def clone(self):
+        return self.__class__(self.factor, **{k: getattr(self, k) for k in self.keys})
+
+    def set_factor(self, factor):
+        self.factor = factor
+        return self
+
+    @abc.abstractmethod
+    def make_noise_sampler(
+        self,
+        x: Tensor,
+        sigma_min=None,
+        sigma_max=None,
+        seed=None,
+        cpu=True,
+    ):
+        raise NotImplementedError
 
 
-class CustomNoise:
+class CustomNoiseItem(CustomNoiseItemBase):
+    def __init__(self, factor, **kwargs):
+        super().__init__(factor, **kwargs)
+        if getattr(self, "noise_type", None) is None:
+            raise ValueError("Noise type required!")
+
+    @torch.no_grad()
+    def make_noise_sampler(
+        self,
+        x: Tensor,
+        sigma_min=None,
+        sigma_max=None,
+        seed=None,
+        cpu=True,
+    ):
+        return get_noise_sampler(
+            self.noise_type,
+            x,
+            sigma_min,
+            sigma_max,
+            seed=seed,
+            cpu=cpu,
+            factor=self.factor,
+        )
+
+
+class CustomNoiseChain:
     def __init__(self, items=None):
         self.items = items if items is not None else []
 
     def clone(self):
-        return CustomNoise(
-            [CustomNoiseItem(i.factor, i.noise_type) for i in self.items],
+        return CustomNoiseChain(
+            [i.clone() for i in self.items],
         )
 
     def add(self, item):
@@ -56,27 +127,40 @@ class CustomNoise:
     def rescaled(self, scale=1.0):
         total = sum(i.factor for i in self.items)
         divisor = total / scale
-        return CustomNoise(
-            [CustomNoiseItem(i.factor / divisor, i.noise_type) for i in self.items],
+        divisor = divisor if divisor != 0 else 1.0
+        return CustomNoiseChain(
+            [i.clone().set_factor(i.factor / divisor) for i in self.items],
         )
 
     @torch.no_grad()
-    def make_noise_sampler(self, x: Tensor) -> Callable:
-        items = tuple(
-            (get_noise_sampler(i.noise_type, x, None, None), i.factor)
+    def make_noise_sampler(
+        self,
+        x: Tensor,
+        sigma_min=None,
+        sigma_max=None,
+        seed=None,
+        cpu=True,
+    ) -> Callable:
+        noise_samplers = tuple(
+            i.make_noise_sampler(
+                x,
+                sigma_min,
+                sigma_max,
+                seed=seed,
+                cpu=cpu,
+            )
             for i in self.items
         )
-        if not items or not all(i[0] for i in items):
+        if not noise_samplers or not all(noise_samplers):
             raise ValueError("Failed to get noise sampler")
+        scale = sum(i.factor for i in self.items)
 
-        def noise_sampler(s, sn):
-            nonlocal items
-            result = items[0][0](s, sn) * items[0][1]
-            for ns, factor in items[1:]:
-                result += ns(s, sn) * factor
-            result /= result.std()
-            scale = sum(i[1] for i in items)
-            return result * scale
+        def noise_sampler(sigma, sigma_next):
+            result = fun.reduce(
+                op.add,
+                (ns(sigma, sigma_next) for ns in noise_samplers),
+            )
+            return scale_noise(result, scale)
 
         return noise_sampler
 
@@ -291,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)
@@ -369,8 +481,9 @@ def power_noise_like(tensor, alpha=2, k=1):  # This doesn't work properly right
     """
     tensor = torch.randn_like(tensor)
     fft = torch.fft.fft2(tensor)
-    freq = torch.arange(1, len(fft) + 1, dtype=torch.float)
-    freq = freq.reshape(freq.shape + (1,) * (len(tensor.shape) - 1))
+    freq = torch.arange(1, len(fft) + 1, dtype=torch.float).reshape(
+        (len(fft),) + (1,) * (tensor.dim() - 1),
+    )
     spectral_density = k / freq**alpha
     noise = torch.rand(tensor.shape) * spectral_density
     mean = torch.mean(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
@@ -378,40 +491,350 @@ def power_noise_like(tensor, alpha=2, k=1):  # This doesn't work properly right
     return noise.to(tensor.device).sub_(mean).div_(std)
 
 
+class NoiseSampler:
+    def __init__(
+        self,
+        x: Tensor,
+        sigma_min: float | None = None,
+        sigma_max: float | None = None,
+        seed: int | None = None,
+        cpu: bool = False,
+        transform: Callable = lambda t: t,
+        make_noise_sampler: Callable | None = None,
+        normalize_noise=False,
+        factor: float = 1.0,
+    ):
+        try:
+            self.noise_sampler = make_noise_sampler(
+                x,
+                transform(torch.as_tensor(sigma_min))
+                if sigma_min is not None
+                else None,
+                transform(torch.as_tensor(sigma_max))
+                if sigma_max is not None
+                else None,
+                seed=seed,
+                cpu=cpu,
+            )
+        except TypeError:
+            self.noise_sampler = make_noise_sampler(x)
+        self.factor = factor
+        self.normalize_noise = normalize_noise
+        self.transform = transform
+        self.device = x.device
+        self.dtype = x.dtype
+
+    @classmethod
+    def simple(cls, f):
+        return lambda *args, **kwargs: cls(
+            *args,
+            **kwargs,
+            make_noise_sampler=lambda x, *_args, **_kwargs: lambda _s, _sn: f(x),
+        )
+
+    @classmethod
+    def wrap(cls, f):
+        return lambda *args, **kwargs: cls(*args, **kwargs, make_noise_sampler=f)
+
+    def __call__(self, *args, **kwargs):
+        args = (
+            self.transform(torch.as_tensor(s)) if s is not None else s for s in args
+        )
+        noise = self.noise_sampler(*args, **kwargs)
+        noise = (
+            scale_noise(noise, self.factor)
+            if self.normalize_noise
+            else noise.mul_(self.factor)
+        )
+        if hasattr(noise, "to"):
+            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] = {
-    # No brownian as it is a special case that requires extra stuff like seed.
-    NoiseType.GAUSSIAN: sampling.default_noise_sampler,
-    NoiseType.UNIFORM: lambda x: lambda _s, _sn: uniform_noise_like(x),
-    NoiseType.PERLIN: lambda x: lambda _s, _sn: rand_perlin_like(x),
-    NoiseType.STUDENTT: studentt_noise_sampler,
-    NoiseType.STUDENTT_TEST: lambda x: lambda _s, _sn: studentt_noise_like(x).to(
-        x.device,
+    NoiseType.BROWNIAN: NoiseSampler.wrap(sampling.BrownianTreeNoiseSampler),
+    NoiseType.GAUSSIAN: NoiseSampler.simple(torch.randn_like),
+    NoiseType.UNIFORM: NoiseSampler.simple(uniform_noise_like),
+    NoiseType.PERLIN: NoiseSampler.simple(rand_perlin_like),
+    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,
     ),
-    NoiseType.PINK: lambda x: lambda _s, _sn: pink_noise_like(x),
-    NoiseType.HIGHRES_PYRAMID: lambda x: lambda _s, _sn: highres_pyramid_noise_like(x),
-    NoiseType.RAINBOW_MILD: lambda x: lambda _s, _sn: (
-        green_noise_like(x) * 0.55 + rand_perlin_like(x) * 0.7
-    )
-    * 1.15,
-    NoiseType.RAINBOW_INTENSE: lambda x: lambda _s, _sn: (
-        green_noise_like(x) * 0.75 + rand_perlin_like(x) * 0.5
-    )
-    * 1.15,
-    NoiseType.LAPLACIAN: lambda x: lambda _s, _sn: laplacian_noise_like(x),
-    NoiseType.POWER: lambda x: lambda _s, _sn: power_noise_like(x),
-    NoiseType.GREEN_TEST: lambda x: lambda _s, _sn: green_noise_like(x),
-    # 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,
+    NoiseType.RAINBOW_INTENSE: NoiseSampler.simple(
+        lambda x: (green_noise_like(x) * 0.75 + rand_perlin_like(x) * 0.5) * 1.15,
+    ),
+    NoiseType.LAPLACIAN: NoiseSampler.simple(laplacian_noise_like),
+    NoiseType.POWER: NoiseSampler.simple(power_noise_like),
+    NoiseType.GREEN_TEST: NoiseSampler.simple(green_noise_like),
 }
 
 
@@ -421,23 +844,25 @@ def get_noise_sampler(
     sigma_min: float | None,
     sigma_max: float | None,
     seed: int | None = None,
-    use_cpu: bool = True,
+    cpu: bool = True,
+    factor: float = 1.0,
+    normalize_noise=True,
 ) -> Callable:
     if noise_type is None:
         noise_type = NoiseType.GAUSSIAN
     elif isinstance(noise_type, str):
         noise_type = NoiseType[noise_type.upper()]
-    if noise_type == NoiseType.BROWNIAN:
-        if sigma_min is None or sigma_max is None:
-            raise ValueError("Must pass sigma min/max when using brownian noise")
-        return sampling.BrownianTreeNoiseSampler(
-            x,
-            sigma_min,
-            sigma_max,
-            seed=seed,
-            cpu=use_cpu,
-        )
-    ns = NOISE_SAMPLERS.get(noise_type)
-    if ns is None:
+    if noise_type == NoiseType.BROWNIAN and (sigma_min is None or sigma_max is None):
+        raise ValueError("Must pass sigma min/max when using brownian noise")
+    mkns = NOISE_SAMPLERS.get(noise_type)
+    if mkns is None:
         raise ValueError("Unknown noise sampler")
-    return ns(x)
+    return mkns(
+        x,
+        sigma_min,
+        sigma_max,
+        seed=seed,
+        cpu=cpu,
+        factor=factor,
+        normalize_noise=normalize_noise,
+    )

py/powernoise.py

@@ -0,0 +1,334 @@
+from __future__ import annotations
+
+import math
+import os
+import random
+
+import folder_paths
+import torch
+from comfy.k_diffusion.sampling import BrownianTreeNoiseSampler
+from PIL import Image
+from torch import Tensor
+
+from .nodes import SonarCustomNoiseNodeBase
+from .noise import CustomNoiseItemBase
+
+# ruff: noqa: ANN003, FBT001, FBT002
+
+
+class PowerNoiseItem(CustomNoiseItemBase):
+    def __init__(self, factor, **kwargs):
+        super().__init__(factor, **kwargs)
+        self.max_freq = max(self.max_freq, self.min_freq)
+
+    def make_filter(self, shape, oversample=4, rel_bw=0.125):
+        """Construct a band-pass * 1/f^alpha filter in rfft space."""
+        height, width = shape[-2:]
+        hfreq_bins = width // 2 + 1
+
+        # Flat unit gain frequency response
+        if self.mix < 1.0:
+            flat = torch.ones(1, 1, height, hfreq_bins)
+        if self.mix <= 0.0:
+            return flat
+
+        # Start with an over-sampled fftshift(rfft2freq()) grid. uses complex
+        # numbers for convenient 2d rotation (unrelated to the fft complex phase
+        # space)
+        fc = torch.complex(
+            # real-fftfreq
+            torch.linspace(0, 0.5, oversample * hfreq_bins),
+            # normal fftfreq
+            torch.linspace(
+                -(height // 2) / height,
+                ((height - 1) // 2) / height,
+                oversample * height,
+            ).unsqueeze(1),
+        )
+        # Rotate, stretch and p-norm
+        if abs(self.rotate) >= 1e-3:
+            fc *= torch.exp(1.0j * torch.deg2rad(torch.scalar_tensor(self.rotate)))
+        if self.stretch > 1.0:
+            fc.real *= self.stretch
+        else:
+            fc.imag *= 1.0 / self.stretch
+        if abs(self.pnorm - 2.0) < 1e-3:
+            d = fc.abs()
+        else:
+            d = (
+                torch.view_as_real(fc)
+                .abs()
+                .pow(self.pnorm)
+                .sum(-1)
+                .pow(1.0 / self.pnorm)
+            )
+
+        # filter gain function
+        op = torch.empty_like(d)
+        m_highpass = d >= self.min_freq
+        m_lowpass = d < self.max_freq
+        m_band = m_highpass & m_lowpass
+        # 1 / f^alpha for the band-pass region
+        op[m_band] = d[m_band].pow(-self.alpha)
+        # easing gaussian (TODO: try cosine windows)
+        m_lowpass = ~m_lowpass
+        op[m_lowpass] = math.pow(self.max_freq, -self.alpha) * torch.exp(
+            -(d[m_lowpass] - self.max_freq).square() / (rel_bw * self.max_freq) ** 2,
+        )
+        if self.min_freq > 0.0:
+            m_highpass = ~m_highpass
+            op[m_highpass] = math.pow(self.min_freq, -self.alpha) * torch.exp(
+                -(d[m_highpass] - self.min_freq).square()
+                / (rel_bw * self.min_freq) ** 2,
+            )
+        op = torch.nn.functional.interpolate(
+            op[None, None, ...],
+            (height, hfreq_bins),
+            mode="bilinear",
+            align_corners=True,
+        )
+        op = op.roll(-(height // 2), -2)  # ifftshift
+        if self.alpha > 0:
+            # In general, the mean offset should be kept as is, sampled from
+            # N(0, 1 / sqrt(H*W) ). However, gain goes to inf when alpha>0.
+            op[..., 0, 0] = 0
+
+        # Scale to unit power gain, then mix flat filter
+        mean_pow_gain = op.mean()
+        if mean_pow_gain <= 0.0:
+            # don't fail catastrophically when something broke
+            return flat
+        op *= 1.0 / mean_pow_gain
+        if self.mix < 1.0:
+            op = torch.lerp(flat, op, self.mix, out=op)
+        return op.sqrt_()
+
+    def make_noise_sampler(
+        self,
+        x: Tensor,
+        sigma_min: float | None,
+        sigma_max: float | None,
+        seed: int | None,
+        cpu: bool = True,
+    ):
+        shape = x.shape
+        device = x.device
+        time_brownian = self.time_brownian
+        if self.time_brownian:
+            if sigma_min is None:
+                raise ValueError(
+                    "time correlated brownian mode is valid only for stochastic samplers",
+                )
+            brownian_tree = BrownianTreeNoiseSampler(
+                x,
+                sigma_min,
+                sigma_max,
+                seed=seed,
+                cpu=cpu,
+            )
+
+        common_mode = self.common_mode
+        if common_mode > 0.0:
+            b, c, h, w = shape
+            torch.eye(c, c)
+            channel_mixer = torch.lerp(
+                torch.eye(c, c),
+                torch.ones(c, c) / c,
+                common_mode,
+            )
+            channel_mixer = channel_mixer.sqrt().to(device, non_blocking=True)
+
+        filter_rfft = self.make_filter(shape).to(device, non_blocking=True)
+
+        def sampler(sigma, sigma_next):
+            if time_brownian:
+                noise = brownian_tree(sigma, sigma_next).to(device)
+                noise_rfft = torch.fft.rfft2(noise, norm="ortho")
+            else:
+                noise_rfft = torch.randn(
+                    (*shape[:-1], filter_rfft.shape[-1]),
+                    dtype=torch.complex64,
+                    device=device,
+                )
+            noise = torch.fft.irfft2(
+                noise_rfft.mul_(filter_rfft),
+                s=shape[-2:],
+                norm="ortho",
+            )
+
+            if common_mode > 0.0:
+                noise = channel_mixer @ noise.swapaxes(0, 1).reshape(c, -1)
+                noise = noise.reshape(c, b, h, w).swapaxes(1, 0)
+            return noise.mul_(self.factor)
+
+        return sampler
+
+    def preview(self, size=(128, 128)):
+        filter_rfft = self.make_filter(size, oversample=1)
+        filter_fft = rfft2_to_fft2(filter_rfft)
+        noise = torch.fft.irfft2(
+            filter_rfft
+            * torch.randn(
+                filter_rfft.shape,
+                dtype=torch.complex64,
+                generator=torch.Generator().manual_seed(0),
+            ),
+            s=size,
+            norm="ortho",
+        )
+        kernel = torch.fft.irfft2(filter_rfft, s=size, norm="ortho")
+        kernel = kernel.roll((size[0] // 2, size[1] // 2), (-2, -1))
+        img = (
+            torch.cat(
+                [
+                    filter_fft.mul_(1 / 3).tanh_().mul_(256.0),
+                    kernel.mul_(1 / 3).tanh_().add_(1.0).mul_(128.0),
+                    noise.mul_(1 / 3).tanh_().add_(1.0).mul_(128.0),
+                ],
+                dim=-1,
+            )
+            .clamp(0, 255)
+            .to(torch.uint8)
+        )
+        return Image.fromarray(img[0, 0].numpy())
+
+
+def rfft2_to_fft2(x):
+    """Apply hermitian-summetry to reconstruct the second half of a fft.
+
+    Only for previews.
+    """
+    height, width = x.shape[-2:]
+    x_r = x.roll(height // 2, -2)  # torch.fft.fftshift(x, -2)
+    x_l = x_r[..., 1 : -1 if width & 1 else None]
+    x_l = torch.flip(x_l.conj(), dims=(-2, -1))
+    if height & 1 == 0:
+        x_l = x_l.roll(1, -2)
+    return torch.cat([x_l, x_r], dim=-1)
+
+
+class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        result = super().INPUT_TYPES()
+        result["required"] |= {
+            "time_brownian": ("BOOLEAN", {"default": False}),
+            "alpha": (
+                "FLOAT",
+                {
+                    "default": 0.0,
+                    "min": -5.0,
+                    "max": 5.0,
+                    "step": 0.001,
+                    "round": False,
+                },
+            ),
+            "max_freq": (
+                "FLOAT",
+                {
+                    "default": 0.7071,
+                    "min": 0.0,
+                    "max": 0.7071,
+                    "step": 0.001,
+                    "round": False,
+                },
+            ),
+            "min_freq": (
+                "FLOAT",
+                {
+                    "default": 0.0,
+                    "min": 0.0,
+                    "max": 0.7071,
+                    "step": 0.001,
+                    "round": False,
+                },
+            ),
+            "stretch": (
+                "FLOAT",
+                {
+                    "default": 1.0,
+                    "min": 0.01,
+                    "max": 100,
+                    "step": 0.1,
+                    "round": False,
+                },
+            ),
+            "rotate": (
+                "FLOAT",
+                {
+                    "default": 0,
+                    "min": -90,
+                    "max": 90,
+                    "step": 5,
+                    "round": False,
+                },
+            ),
+            "pnorm": (
+                "FLOAT",
+                {
+                    "default": 2,
+                    "min": 0.125,
+                    "max": 100,
+                    "step": 0.1,
+                    "round": False,
+                },
+            ),
+            "mix": (
+                "FLOAT",
+                {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.001,
+                    "round": False,
+                },
+            ),
+            "common_mode": (
+                "FLOAT",
+                {
+                    "default": 0.0,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.001,
+                    "round": False,
+                },
+            ),
+            "preview": (["none", "no_mix", "mix"],),
+        }
+        return result
+
+    def get_item_class(self):
+        return PowerNoiseItem
+
+    def go(
+        self,
+        preview="none",
+        **kwargs,
+    ):
+        result = super().go(**kwargs)
+        if preview == "none":
+            return result
+        if preview == "no_mix":
+            kwargs["mix"] = 1.0
+        img = PowerNoiseItem(**kwargs).preview()
+
+        output_dir = folder_paths.get_temp_directory()
+        prefix_append = "sonar_temp_" + "".join(
+            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)
+        )
+        filename = f"{filename}_{counter:05}_.png"
+        file_path = os.path.join(full_output_folder, filename)  # noqa: PTH118
+        img.save(file_path, compress_level=1)
+
+        return {
+            "ui": {
+                "images": [
+                    {"filename": filename, "subfolder": subfolder, "type": "temp"},
+                ],
+            },
+            "result": result,
+        }

py/sonar.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 from enum import Enum, auto
-from typing import Any, NamedTuple
+from sys import stderr
+from typing import Any, Callable, NamedTuple
 
 import torch
 from comfy.k_diffusion import sampling
@@ -44,14 +45,48 @@ class SonarConfig(NamedTuple):
 
 
 class SonarBase:
-    def __init__(
-        self,
-        cfg: SonarConfig,
-    ) -> None:
+    DEFAULT_NOISE_TYPE = noise.NoiseType.GAUSSIAN
+
+    def __init__(self, cfg: SonarConfig) -> None:
         self.history_d = None
         self.cfg = cfg
         self.noise_sampler = None
 
+    def set_noise_sampler(
+        self,
+        x: Tensor,
+        sigmas,
+        noise_sampler: Callable | None,
+        seed: int | None = None,
+    ):
+        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+        if noise_sampler is not None and self.cfg.noise_type not in (
+            None,
+            self.DEFAULT_NOISE_TYPE,
+        ):
+            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(
+                x,
+                sigma_min,
+                sigma_max,
+                seed=seed,
+            )
+        elif noise_sampler is None:
+            noise_sampler = noise.get_noise_sampler(
+                self.cfg.noise_type or self.DEFAULT_NOISE_TYPE,
+                x,
+                sigma_min,
+                sigma_max,
+                seed=seed,
+                cpu=True,
+            )
+        self.noise_sampler = noise_sampler
+        return noise_sampler
+
     def init_hist_d(self, x: Tensor) -> None:
         if self.history_d is not None:
             return
@@ -67,7 +102,7 @@ class SonarBase:
                 None,
                 None,
                 seed=self.extra_args.get("seed"),
-                use_cpu=True,
+                cpu=True,
             )
             self.history_d = ns(None, None)
         else:
@@ -201,7 +236,7 @@ class SonarEuler(SonarSampler):
     ):
         self.init_hist_d(sample)
 
-        sigma = self.sigmas[step_index]
+        sigma, sigma_to = self.sigmas[step_index], self.sigmas[step_index + 1]
 
         gamma = (
             min(self.s_churn / (len(self.sigmas) - 1), 2**0.5 - 1)
@@ -212,8 +247,11 @@ class SonarEuler(SonarSampler):
         sigma_hat = sigma * (gamma + 1)
 
         if gamma > 0:
-            noise = torch.randn_like(sample.shape)
-
+            noise = (
+                self.noise_sampler(sigma, sigma_to)
+                if self.noise_sampler
+                else torch.randn_like(sample)
+            )
             eps = noise * self.s_noise
             sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
 
@@ -243,6 +281,7 @@ class SonarEuler(SonarSampler):
         extra_args=None,
         callback=None,
         disable=None,
+        noise_sampler: Callable | None = None,
         sonar_config=None,
         s_churn=0.0,
         s_tmin=0.0,
@@ -263,6 +302,12 @@ class SonarEuler(SonarSampler):
             {} if extra_args is None else extra_args,
             sonar_config,
         )
+        sonar.set_noise_sampler(
+            x,
+            sigmas,
+            noise_sampler,
+            seed=extra_args.get("seed"),
+        )
 
         for i in trange(len(sigmas) - 1, disable=disable):
             x, sigma, sigma_hat, denoised = sonar.step(
@@ -285,14 +330,12 @@ class SonarEuler(SonarSampler):
 class SonarEulerAncestral(SonarSampler):
     def __init__(
         self,
-        noise_sampler,
         eta: float = 1.0,
         s_noise: float = 1.0,
         *args: list[Any],
         **kwargs: dict[str, Any],
     ):
         super().__init__(*args, **kwargs)
-        self.noise_sampler = noise_sampler
         self.eta = eta
         self.s_noise = s_noise
 
@@ -342,33 +385,12 @@ class SonarEulerAncestral(SonarSampler):
         sonar_config=None,
         eta=1.0,
         s_noise=1.0,
-        noise_sampler=None,
+        noise_sampler: Callable | None = None,
     ):
         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",
-            )
-        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-        if sonar_config.custom_noise:
-            noise_sampler = sonar_config.custom_noise.make_noise_sampler(x)
-        else:
-            noise_sampler = noise.get_noise_sampler(
-                sonar_config.noise_type,
-                x,
-                sigma_min,
-                sigma_max,
-                seed=extra_args.get("seed"),
-                use_cpu=True,
-            )
         s_in = x.new_ones([x.shape[0]])
         sonar = cls(
-            noise_sampler,
             eta,
             s_noise,
             model,
@@ -377,6 +399,12 @@ class SonarEulerAncestral(SonarSampler):
             {} if extra_args is None else extra_args,
             sonar_config,
         )
+        sonar.set_noise_sampler(
+            x,
+            sigmas,
+            noise_sampler,
+            seed=extra_args.get("seed"),
+        )
 
         for i in trange(len(sigmas) - 1, disable=disable):
             x, sigma, sigma_hat, denoised = sonar.step(
@@ -397,16 +425,16 @@ class SonarEulerAncestral(SonarSampler):
 
 
 class SonarDPMPPSDE(SonarSampler):
+    DEFAULT_NOISE_TYPE = noise.NoiseType.BROWNIAN
+
     def __init__(
         self,
-        noise_sampler,
         eta: float = 1.0,
         s_noise: float = 1.0,
         *args: list[Any],
         **kwargs: dict[str, Any],
     ):
         super().__init__(*args, **kwargs)
-        self.noise_sampler = noise_sampler
         self.eta = eta
         self.s_noise = s_noise
 
@@ -529,29 +557,8 @@ 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",
-            )
-        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-        if sonar_config.custom_noise:
-            noise_sampler = sonar_config.custom_noise.make_noise_sampler(x)
-        else:
-            noise_sampler = noise.get_noise_sampler(
-                sonar_config.noise_type,
-                x,
-                sigma_min,
-                sigma_max,
-                seed=extra_args.get("seed"),
-                use_cpu=True,
-            )
         s_in = x.new_ones([x.shape[0]])
         sonar = cls(
-            noise_sampler,
             eta,
             s_noise,
             model,
@@ -560,6 +567,12 @@ class SonarDPMPPSDE(SonarSampler):
             {} if extra_args is None else extra_args,
             sonar_config,
         )
+        sonar.set_noise_sampler(
+            x,
+            sigmas,
+            noise_sampler,
+            seed=extra_args.get("seed"),
+        )
 
         for i in trange(len(sigmas) - 1, disable=disable):
             x, sigma, sigma_hat, denoised = sonar.step(

ruff.toml

@@ -22,6 +22,7 @@ ignore = [
     "ERA001",
     "F403",
     "F405",
+    "FBT002",
     "PLR0912",
     "PLR0913",
     "PLR0915",