November 2024 mega update, see changelog for details

Add normalize and cpu_noise parameters to NoisyLatentLike

README.md

@@ -107,11 +107,15 @@ Original Sonar Sampler implementation (for A1111): https://github.com/Kahsolt/st
 
 My version was initially based on this Sonar sampler implementation for Diffusers: https://github.com/alexblattner/modified-euler-samplers-for-sonar-diffusers/
 
-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.
+* 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.
+* New pyramid noise based on implementation in [Jonathan Whitaker](https://wandb.ai/johnowhitaker/multires_noise/reports/Multi-Resolution-Noise-for-Diffusion-Model-Training--VmlldzozNjYyOTU2)'s article on multi-resolution noise.
+* Original `SonarPowerNoise` contributed by [elias-gaeros](https://github.com/elias-gaeros/). Additionally, he provided a lot of guidance with refactoring it to allow separate filtering and other enhancements and answered a multitude of dumb questions. To say those changes are only co-authored is probably giving myself too much credit. Thank you! Your patience and help is very much appreciated.
+* New 1/f (onef) and power law (white, grey, violet, velvet) noise types referenced from https://github.com/WASasquatch/PowerNoiseSuite
 
-New pyramid noise based on implementation in [Jonathan Whitaker](https://wandb.ai/johnowhitaker/multires_noise/reports/Multi-Resolution-Noise-for-Diffusion-Model-Training--VmlldzozNjYyOTU2)'s article on multi-resolution noise.
+## Errata
 
-Original `SonarPowerNoise` contributed by [elias-gaeros](https://github.com/elias-gaeros/). Additionally, he provided a lot of guidance with refactoring it to allow separate filtering and other enhancements and answered a multitude of dumb questions. To say those changes are only co-authored is probably giving myself too much credit. Thank you! Your patience and help is very much appreciated.
+* The noise types might not actually do what they claim. In that, I mean something I called "pink" noise might not be what is technically known as "pink noise". My implementations are best-effort. Bug reports and contributions to improve this repo are always welcome!
+* Whether noise gets generated on GPU or CPU is probably inconsistent. This means changing GPU types may change seeds, also when this eventually gets fixed it will probably also change seeds.
 
 ## Sonar Examples
 

changelog.md

@@ -2,6 +2,28 @@
 
 Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top.
 
+## 20241129
+
+*Note*: Contains some potentially workflow-breaking changes.
+
+* `pink` noise type renamed to `pink_old` - the implementation was incorrect.
+* `power` noise type renamed to `power_old` - the implementation was incorrect.
+* Added `onef_pinkish` (higher frequencye) and `onef_greenish` (lower frequency) noise types.
+* Added `SonarAdvanced1fNoise` node and `onef_pinkish`, `onef_greenish`, `onef_pinkish_mix`, `onef_greenish_mix`, and `onef_pinkishgreenish` noise types.
+* Added `SonarAdvancedPowerLawNoise` node and `grey`, `white`, `violet` and `velvet` noise types.
+* The `SonarAdvancedPyramidNoise` node can now use upscale methods from my [ComfyUI-bleh](https://github.com/blepping/ComfyUI-bleh) node pack if it is available.
+* Added the `SonarChannelNoise` and `SonarBlendedNoise` nodes.
+* Added the `SonarBlehOpsNoise` node.
+* Added advanced parameter input to the SampleConfigOverride node, you can now pass options directly to the wrapped sampler function.
+* Custom noise inputs now are semi-wildcard and will accept `OCS_NOISE` or `SONAR_CUSTOM_NOISE` interchangeably.
+
+## 20240823
+
+* Added descriptions and tooltips for most nodes.
+* Added `repeat_batch` parameter to `NoisyLatentLike` node.
+* Added a `SONAR_CUSTOM_NOISE to NOISE` node to allow converting from Sonar's custom noise type to the built in ComfyUI `NOISE` (used by `SamplerCustomAdvanced` and possibly other nodes).
+* Added a `SonarAdvancedPyramidNoise` node that allows setting parameters for the pyramid noise variants.
+
 ## 20240521
 
 Mega update! Many new features, documentation reorganized.

docs/advanced_noise_nodes.md

@@ -72,6 +72,40 @@ If you want to create noise for initial sampling, connect model and sigmas to th
 
 This node 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.
 
+You can enter YAML parameters in the text input, these arguments are passed directly to the sampler function without any error checking. If the same key exists in the node itself (i.e. `s_noise`) the one in the text input will take precedence. Note that these are based on the internal sampler function so the names of the arguments won't necessarily be the same as the sampler node (but they often are). You may need to check the source code for the sampler.
+
+***
+
+### `SONAR_CUSTOM_NOISE to NOISE`
+
+This node can be used to convert Sonar custom noise to the `NOISE` type used by the builtin `SamplerCustomAdvanced` (and any other nodes that take a `NOISE` input).
+
+***
+
+### `SonarAdvancedPyramidNoise`
+
+Allows setting some parameters for the pyramid noise variants (`pyramid`, `highres_pyramid` and `pyramid_old`). `discount` further from zero generally results in a more extreme colorful effect (can also be set to negative values). Higher `iterations` also tends to make the effect more extreme - zero iterations will just return normal Gaussian noise. You can also experiment with the `upscale_mode` for different effects.
+
+### `SonarAdvanced1fNoise`
+
+More extensive documentation TBD (hopefully). For now, a few recipes:
+
+These differ differ only in alpha. For the other parameters, use `k=1, vf=1, hf=1, use_sqrt=true` to start.
+
+* `blue`: `alpha=1`
+* `green`: `alpha=0.75`
+* `pink`: `alpha=0.5`
+*
+
+### `SonarAdvancedPowerLawNoise`
+
+More extensive documentation TBD (hopefully). For now, a few recipes:
+
+* `white`: `alpha=0, use_sign=true, div_max_dims=none`
+* `grey`: `alpha=0, use_sign=false, div_max_dims=none`
+* `velvet`: `alpha=1, use_sign=true, div_max_dims=all, use_div_max_abs=true`
+* `violet`: `alpha=0.5, use_sign=true, div_max_dims=all, use_div_max_abs=true`
+
 ***
 
 ### `SonarModulatedNoise`
@@ -314,3 +348,29 @@ Light to dark (negative strength):
 Randomly chooses between the noise types in the chain connected to it each time the noise sampler is called.
 You generally do not want to use `rescale` here. You can also set `mix_count` to choose and combine multiple
 types.
+
+### `SonarChannelNoise`
+
+Allows using a different noise generator per channel. The custom noise items attached to this node are treated as a list where the furthest item from the node will correspond to channel 0. For example where CN is a custom noise node and SCN is the `SonarChannelNoise` node:
+
+```plaintext
+CN (channel 0) -> CN (channel 1) -> SCN
+```
+
+Don't enable `rescale` in the custom noise nodes attached to `SonarChannelNoise`. If you want a blend of noise types for a channel, you can use something like `SonarBlendedNoise`.
+
+### `SonarBlendedNoise`
+
+Allows blending two noise generators. If [ComfyUI-bleh](https://github.com/blepping/ComfyUI-bleh) is available, you will have access to many more blending modes.
+
+### `SonarBlehOpsNoise`
+
+Only provided if [ComfyUI-bleh](https://github.com/blepping/ComfyUI-bleh) is available. Allows transforming/manipulating noise with bleh blockops expressions. For instance, you can do something like:
+
+```yaml
+- ops:
+  - [multiply, -1]
+  - [roll, -2, 0.5]
+```
+
+to flip the sign on the noise and then roll dimension -2 (height) by 50%.

docs/base_noise_types.md

@@ -8,6 +8,18 @@ noise of that type. However you can either schedule the noise type to kick in at
 (as in these examples) and/or mix it with something a bit more run of the mill. See
 [advanced_noise_nodes](advanced_noise_nodes.md).
 
+## Documentation TBD
+
+* `grey`
+* `onef_greenish_mix` (50/50 mix of positive/negative noise.)
+* `onef_greenish`
+* `onef_pinkish_mix` (50/50 mix of positive/negative noise.)
+* `onef_pinkish`
+* `onef_pinkishgreenish` (50/50 mix of `onef_pinkish` and `onef_greenish`.)
+* `velvet`
+* `violet`
+* `white`
+
 ## Brownian
 
 This is the default noise type for SDE samplers.
@@ -62,13 +74,17 @@ Variation using bislerp scaling:
 
 ***
 
-## Pink
+## Pink Old
+
+Previously known as `pink`. The implementation isn't correct, though in terms of results it's fine.
 
 ![Pink](../assets/example_images/noise_base_types/noise_pink.png)
 
 ***
 
-## Power Builtin
+## Power Old
+
+Previously known as `power`. The implementation isn't correct, though in terms of results it's fine.
 
 ![PowerBuiltin](../assets/example_images/noise_base_types/noise_power_builtin.png)
 

py/freeu_extreme.py

@@ -36,6 +36,7 @@ BLEND_OPS = (
 
 
 class FreeUExtremeConfigNode:
+    DESCRIPTION = "Allows setting configuration for FreeU Extreme."
     RETURN_TYPES = ("FRUX_CONFIG",)
     FUNCTION = "go"
     CATEGORY = "model_patches"
@@ -44,10 +45,33 @@ class FreeUExtremeConfigNode:
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "stage_1": ("BOOLEAN", {"default": True}),
-                "stage_2": ("BOOLEAN", {"default": False}),
-                "stage_3": ("BOOLEAN", {"default": False}),
-                "target": (("backbone", "skip", "both"),),
+                "stage_1": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": "Controls whether this configuration applies to stage 1.",
+                    },
+                ),
+                "stage_2": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Controls whether this configuration applies to stage 2.",
+                    },
+                ),
+                "stage_3": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Controls whether this configuration applies to stage 3.",
+                    },
+                ),
+                "target": (
+                    ("backbone", "skip", "both"),
+                    {
+                        "tooltip": "Controls whether this filter applies to backbone or skip layers (or both).",
+                    },
+                ),
                 "start": (
                     "FLOAT",
                     {
@@ -56,6 +80,7 @@ class FreeUExtremeConfigNode:
                         "max": 1.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Start time as percentage of sampling this configuration applies to. Inclusive.",
                     },
                 ),
                 "end": (
@@ -66,6 +91,7 @@ class FreeUExtremeConfigNode:
                         "max": 1.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "End time as percentage of sampling this configuration applies to. Inclusive.",
                     },
                 ),
                 "slice": (
@@ -76,6 +102,7 @@ class FreeUExtremeConfigNode:
                         "max": 1.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Percentage of the layer the FreeU effect is applied to.",
                     },
                 ),
                 "slice_offset": (
@@ -86,6 +113,7 @@ class FreeUExtremeConfigNode:
                         "max": 1.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Offset as a percentage the layer is applied to. For example if slice is 0.25 and slice_offset is 0.25 then the filter will apply to the range 25% through 50%.",
                     },
                 ),
                 "filter_norm": (
@@ -96,6 +124,7 @@ class FreeUExtremeConfigNode:
                         "max": 10.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Normalization factor applied to the filter. 1.0 means 100% normalized.",
                     },
                 ),
                 "scale": (
@@ -106,6 +135,7 @@ class FreeUExtremeConfigNode:
                         "max": 100.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Strength of the effects applied by this configuration.",
                     },
                 ),
                 "blend": (
@@ -116,19 +146,48 @@ class FreeUExtremeConfigNode:
                         "max": 10.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Blends the filtered result based on the specified strength where 1.0 means 100% filtered.",
+                    },
+                ),
+                "blend_mode": (
+                    tuple(BLEND_OPS.keys()),
+                    {
+                        "tooltip": "Mode used when blending. Generally only has an effect when blend is set to values other than 0 or 1",
+                    },
+                ),
+                "hidden_mean": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": "You can think of this as FreeU V2 mode.",
+                    },
+                ),
+                "final": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": "When enabled, other configurations won't be considered if this one matched. Otherwise, multiple configurations/filter effects can be stacked.",
                     },
                 ),
-                "blend_mode": (tuple(BLEND_OPS.keys()),),
-                "hidden_mean": ("BOOLEAN", {"default": True}),
-                "final": ("BOOLEAN", {"default": True}),
             },
             "optional": {
-                "sonar_power_filter_opt": ("SONAR_POWER_FILTER",),
-                "frux_config_opt": ("FRUX_CONFIG",),
+                "sonar_power_filter_opt": (
+                    "SONAR_POWER_FILTER",
+                    {
+                        "tooltip": "Optionally attach a Power Filter here to set filtering parameters.",
+                    },
+                ),
+                "frux_config_opt": (
+                    "FRUX_CONFIG",
+                    {
+                        "tooltip": "Optionally attach another configuration node here.",
+                    },
+                ),
             },
         }
 
-    def go(self, **kwargs: dict):
+    @classmethod
+    def go(cls, **kwargs: dict):
         return (FreeUExtremeConfig(**kwargs),)
 
 
@@ -223,12 +282,10 @@ class FreeUExtremeConfig:
             return False
         if not getattr(self, f"stage_{stage}"):
             return False
-        if self.target not in ("skip" if is_skip else "backbone", "both"):
-            return False
-        return True
+        return not self.target not in {"skip" if is_skip else "backbone", "both"}
 
     def apply(self, idx, x, filter_cache, cpu_fft=False):
-        batch, features, height, width = x.shape
+        _batch, features, _height, _width = x.shape
         scale = self.get_scale(x)
         slice_size = int(features * self.slice)
         slice_offs = int(features * self.slice_offset)
@@ -280,6 +337,7 @@ class FreeUExtremeConfig:
 
 
 class FreeUExtremeNode:
+    DESCRIPTION = "Main FreeU Extreme node. Allows patching a model with the FreeU (V2) effect with more control."
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "go"
     CATEGORY = "model_patches"
@@ -288,18 +346,45 @@ class FreeUExtremeNode:
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "model": ("MODEL",),
-                "cpu_fft": ("BOOLEAN", {"default": False}),
+                "model": (
+                    "MODEL",
+                    {
+                        "tooltip": "Model to patch.",
+                    },
+                ),
+                "cpu_fft": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Controls whether to perform FFT calculations on the CPU. May be necessary for some GPUs that don't have native support for FFT operations at the cost of performance.",
+                    },
+                ),
             },
             "optional": {
-                "input_config": ("FRUX_CONFIG",),
-                "middle_config": ("FRUX_CONFIG",),
-                "output_config": ("FRUX_CONFIG",),
+                "input_config": (
+                    "FRUX_CONFIG",
+                    {
+                        "tooltip": "Allows specifying configuration for input blocks.",
+                    },
+                ),
+                "middle_config": (
+                    "FRUX_CONFIG",
+                    {
+                        "tooltip": "Allows specifying configuration for middle blocks.",
+                    },
+                ),
+                "output_config": (
+                    "FRUX_CONFIG",
+                    {
+                        "tooltip": "Allows specifying configuration for output blocks.",
+                    },
+                ),
             },
         }
 
+    @classmethod
     def go(
-        self,
+        cls,
         model,
         cpu_fft,
         input_config=None,

py/noise.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+from functools import partial
 from typing import Callable
 
 import comfy
@@ -10,6 +11,7 @@ from torch import Tensor
 
 from . import external
 from .noise_generation import *
+from .sonar import SonarGuidanceMixin
 
 # ruff: noqa: D412, D413, D417, D212, D407, ANN002, ANN003, FBT001, FBT002, S311
 
@@ -197,6 +199,62 @@ class NoiseSampler:
         return noise
 
 
+class AdvancedNoiseBase(CustomNoiseItemBase):
+    ns_factory_arg_keys = ()
+
+    # This has to be done as a property for some reason.
+    @property
+    def ns_factory(self):
+        raise NotImplementedError
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.ns_factory is None:
+            raise NotImplementedError("ns_factory not implemented")
+        noise_sampler_kwargs = {}
+        for k in self.ns_factory_arg_keys:
+            v = getattr(self, k, None)
+            if v is not None:
+                noise_sampler_kwargs[k] = v
+        self.sampler_factory = NoiseSampler.simple(
+            partial(self.ns_factory, **noise_sampler_kwargs),
+        )
+
+    @torch.no_grad()
+    def make_noise_sampler(self, *args, **kwargs):
+        return self.sampler_factory(*args, factor=self.factor, **kwargs)
+
+
+class AdvancedPyramidNoise(AdvancedNoiseBase):
+    ns_factory_arg_keys = ("discount", "iterations", "upscale_mode")
+
+    pyramid_variants_map = {  # noqa: RUF012
+        "pyramid": pyramid_noise_like,
+        "pyramid_old": pyramid_old_noise_like,
+        "highres_pyramid": highres_pyramid_noise_like,
+    }
+
+    @property
+    def ns_factory(self):
+        return self.pyramid_variants_map[self.variant]
+
+
+class Advanced1fNoise(AdvancedNoiseBase):
+    ns_factory_arg_keys = ("alpha", "hfac", "wfac", "k", "use_sqrt", "base_power")
+
+    @property
+    def ns_factory(self):
+        return onef_noise_like
+
+
+class AdvancedPowerLawNoise(AdvancedNoiseBase):
+    ns_factory_arg_keys = ("alpha", "div_max_dims", "use_sign")
+
+    @property
+    def ns_factory(self):
+        return powerlaw_noise_like
+
+
 class CompositeNoise(CustomNoiseItemBase):
     def __init__(
         self,
@@ -220,7 +278,7 @@ class CompositeNoise(CustomNoiseItemBase):
         )
 
     def clone_key(self, k):
-        if k in ("mask", "src_noise", "dst_noise"):
+        if k in {"mask", "src_noise", "dst_noise"}:
             return getattr(self, k).clone()
         return super().clone_key(k)
 
@@ -286,13 +344,11 @@ class GuidedNoise(CustomNoiseItemBase):
         )
 
     def clone_key(self, k):
-        if k in ("noise", "ref_latent"):
+        if k in {"noise", "ref_latent"}:
             return getattr(self, k).clone()
         return super().clone_key(k)
 
     def make_noise_sampler(self, x, *args, normalized=True, **kwargs):
-        from .sonar import SonarGuidanceMixin
-
         factor, guidance_factor = self.factor, self.guidance_factor
         normalize_noise, normalize_result = (
             self.get_normalize(f"normalize_{k}", normalized)
@@ -325,6 +381,7 @@ class GuidedNoise(CustomNoiseItemBase):
                         factor,
                         normalized=normalize_result,
                     )
+
             case "euler":
 
                 def noise_sampler(s, sn):
@@ -808,9 +865,142 @@ class RandomNoise(CustomNoiseItemBase):
         return noise_sampler
 
 
+class ChannelNoise(CustomNoiseItemBase):
+    def __init__(self, factor, *, noise, insufficient_channels_mode, normalize):
+        if len(noise.items) == 0:
+            raise ValueError("ChannelNoise requires at least one noise item")
+        if insufficient_channels_mode not in {"wrap", "repeat", "zero"}:
+            raise ValueError("Bad insufficient_channels_mode")
+        super().__init__(
+            factor,
+            noise=noise.clone(),
+            insufficient_channels_mode=insufficient_channels_mode,
+            normalize=normalize,
+        )
+
+    def clone_key(self, k):
+        if k == "noise":
+            return self.noise.clone()
+        return super().clone_key(k)
+
+    def make_noise_sampler(self, x, *args, normalized=True, **kwargs):
+        factor = self.factor
+        icmode = self.insufficient_channels_mode
+        c = x.shape[1]
+        noise_items = self.noise.items[:c]
+        num_samplers = len(noise_items)
+
+        def make_zero_noise_sampler(x, *_args, **_kwargs):
+            return lambda *_args, **_kwargs: torch.zeros_like(x)
+
+        make_zero_noise_sampler.make_noise_sampler = make_zero_noise_sampler
+
+        while len(noise_items) < c:
+            if icmode == "wrap":
+                item = noise_items[len(noise_items) % num_samplers]
+            elif icmode == "repeat":
+                item = noise_items[num_samplers - 1]
+            elif icmode == "zero":
+                item = make_zero_noise_sampler
+            else:
+                raise ValueError("Bad insufficient_channels_mode")
+            noise_items.append(item)
+        noise_samplers = tuple(
+            ni.make_noise_sampler(
+                x[:, ni_channel : ni_channel + 1, ...],
+                *args,
+                normalized=False,
+                **kwargs,
+            )
+            for ni_channel, ni in enumerate(noise_items)
+        )
+        normalize = self.get_normalize("normalize", normalized)
+
+        def noise_sampler(s, sn):
+            noise = torch.cat(tuple(ns(s, sn) for ns in noise_samplers), dim=1)
+            return scale_noise(noise, factor, normalized=normalize)
+
+        return noise_sampler
+
+
+class BlendedNoise(CustomNoiseItemBase):
+    def __init__(
+        self,
+        factor,
+        *,
+        normalize,
+        blend_function,
+        custom_noise_1=None,
+        custom_noise_2=None,
+        noise_2_percent=0.5,
+    ):
+        if custom_noise_1 is None and noise_2_percent != 1:
+            raise ValueError(
+                "When custom_noise_1 is not attached noise_2_percent must be set to 1",
+            )
+        if custom_noise_2 is None and noise_2_percent != 0:
+            raise ValueError(
+                "When custom_noise_2 is not attached noise_2_percent must be set to 0",
+            )
+        if noise_2_percent == 1:
+            custom_noise_1, custom_noise_2 = custom_noise_2, None
+            noise_2_percent = 0.0
+
+        super().__init__(
+            factor,
+            noise_2_percent=noise_2_percent,
+            blend_function=blend_function,
+            custom_noise_1=custom_noise_1.clone(),
+            custom_noise_2=None if custom_noise_2 is None else custom_noise_2.clone(),
+            normalize=normalize,
+        )
+
+    def clone_key(self, k):
+        if k == "custom_noise_1":
+            return self.custom_noise_1.clone()
+        if k == "custom_noise_2":
+            return None if self.custom_noise_2 is None else self.custom_noise_2.clone()
+        return super().clone_key(k)
+
+    def make_noise_sampler(self, x, *args, normalized=True, **kwargs):
+        factor = self.factor
+        normalize = self.get_normalize("normalize", normalized)
+        blend_function = self.blend_function
+        n2_blend = self.noise_2_percent
+        n2_blend_tensor = x.new_full((1,), n2_blend)
+        ns_1 = self.custom_noise_1.make_noise_sampler(
+            x,
+            *args,
+            normalized=False,
+            **kwargs,
+        )
+        ns_2 = (
+            None
+            if self.custom_noise_2 is None
+            else self.custom_noise_2.make_noise_sampler(
+                x,
+                *args,
+                normalized=False,
+                **kwargs,
+            )
+        )
+
+        def noise_sampler(s, sn):
+            noise_1 = ns_1(s, sn)
+            noise = (
+                noise_1
+                if n2_blend == 0 or ns_2 is None
+                else blend_function(noise_1, ns_2(s, sn), n2_blend_tensor)
+            )
+            return scale_noise(noise, factor, normalized=normalize)
+
+        return noise_sampler
+
+
 if "bleh" in external.MODULES:
     bleh = external.MODULES["bleh"]
     BLU = bleh.py.latent_utils
+    BOPS = bleh.py.nodes.ops
 
     class BlendFilterNoise(CustomNoiseItemBase):
         def __init__(
@@ -830,7 +1020,7 @@ if "bleh" in external.MODULES:
             normalize_noise,
         ):
             if len(noise.items) == 0:
-                raise ValueError("BlendFilterNoise requires ta least one noise item")
+                raise ValueError("BlendFilterNoise requires at least one noise item")
             super().__init__(
                 factor,
                 noise=noise.clone(),
@@ -884,8 +1074,8 @@ if "bleh" in external.MODULES:
                 normalized or num_samplers > 1,
             )
             normalize_result = self.get_normalize("normalize_result", normalized)
-            noise_effects = self.affect in ("noise", "both")
-            result_effects = self.affect in ("result", "both")
+            noise_effects = self.affect in {"noise", "both"}
+            result_effects = self.affect in {"result", "both"}
             noise_init = torch.zeros_like(x)
 
             def noise_sampler(s, sn):
@@ -910,6 +1100,58 @@ if "bleh" in external.MODULES:
 
             return noise_sampler
 
+    class BlehOpsNoise(CustomNoiseItemBase):
+        def __init__(
+            self,
+            factor,
+            *,
+            noise,
+            rules,
+            normalize,
+        ):
+            if len(noise.items) == 0:
+                raise ValueError("BlehOpsNoise requires at least one noise item")
+            super().__init__(
+                factor,
+                noise=noise.clone(),
+                rules=rules,
+                normalize=normalize,
+            )
+
+        def clone_key(self, k):
+            if k == "noise":
+                return self.noise.clone()
+            return super().clone_key(k)
+
+        def make_noise_sampler(self, x, *args, normalized=True, **kwargs):
+            factor = self.factor
+            normalize = self.get_normalize("normalize", normalized)
+            rulegroup = self.rules
+            internal_ns = self.noise.make_noise_sampler(
+                x,
+                *args,
+                normalized=False,
+                **kwargs,
+            )
+
+            def noise_sampler(s, sn):
+                noise = internal_ns(s, sn)
+                if len(rulegroup.rules):
+                    state = {
+                        BOPS.CondType.TYPE: BOPS.PatchType.LATENT,
+                        BOPS.CondType.PERCENT: 0.0,
+                        BOPS.CondType.BLOCK: -1,
+                        BOPS.CondType.STAGE: -1,
+                        "sigma": None if s is None else s,
+                        "h": noise,
+                        "hsp": x.detach().clone(),
+                        "target": "h",
+                    }
+                    noise = rulegroup.eval(state, toplevel=True)["h"]
+                return scale_noise(noise, factor, normalized=normalize)
+
+            return noise_sampler
+
 
 NOISE_SAMPLERS: dict[NoiseType, Callable] = {
     NoiseType.BROWNIAN: NoiseSampler.wrap(sampling.BrownianTreeNoiseSampler),
@@ -917,39 +1159,94 @@ NOISE_SAMPLERS: dict[NoiseType, Callable] = {
     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.ONEF_PINKISH: NoiseSampler.simple(partial(onef_noise_like, alpha=-0.5)),
+    NoiseType.ONEF_GREENISH: NoiseSampler.simple(partial(onef_noise_like, alpha=0.5)),
+    NoiseType.ONEF_PINKISHGREENISH: NoiseSampler.simple(
+        lambda x: onef_noise_like(x, alpha=0.5)
+        .add_(onef_noise_like(x, alpha=-0.5))
+        .mul_(0.5),
+    ),
+    NoiseType.ONEF_PINKISH_MIX: NoiseSampler.simple(
+        lambda x: onef_noise_like(x, alpha=-0.5)
+        .mul_(-1.0)
+        .add_(onef_noise_like(x, alpha=-0.5))
+        .mul_(0.5),
+    ),
+    NoiseType.ONEF_GREENISH_MIX: NoiseSampler.simple(
+        lambda x: onef_noise_like(x, alpha=0.5)
+        .mul_(-1.0)
+        .add_(onef_noise_like(x, alpha=0.5))
+        .mul_(0.5),
+    ),
+    NoiseType.WHITE: NoiseSampler.simple(
+        partial(
+            powerlaw_noise_like,
+            alpha=0.0,
+            use_sign=True,
+        ),
+    ),
+    NoiseType.GREY: NoiseSampler.simple(
+        partial(
+            powerlaw_noise_like,
+            alpha=0.0,
+            use_sign=False,
+        ),
+    ),
+    NoiseType.VELVET: NoiseSampler.simple(
+        partial(
+            powerlaw_noise_like,
+            alpha=1.0,
+            use_sign=True,
+            div_max_dims=(-3, -2, -1),
+        ),
+    ),
+    NoiseType.VIOLET: NoiseSampler.simple(
+        partial(
+            powerlaw_noise_like,
+            alpha=0.5,
+            use_sign=True,
+            div_max_dims=(-3, -2, -1),
+        ),
+    ),
+    NoiseType.PINK_OLD: NoiseSampler.simple(pink_noise_old_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,
+        lambda x: green_noise_like(x)
+        .mul_(0.55)
+        .add_(rand_perlin_like(x).mul_(0.7))
+        .mul_(1.15),
     ),
     NoiseType.RAINBOW_INTENSE: NoiseSampler.simple(
-        lambda x: (green_noise_like(x) * 0.75 + rand_perlin_like(x) * 0.5) * 1.15,
+        lambda x: green_noise_like(x)
+        .mul_(0.75)
+        .add_(rand_perlin_like(x).mul_(0.5))
+        .mul_(1.15),
     ),
     NoiseType.LAPLACIAN: NoiseSampler.simple(laplacian_noise_like),
-    NoiseType.POWER: NoiseSampler.simple(power_noise_like),
+    NoiseType.POWER_OLD: NoiseSampler.simple(power_noise_old_like),
     NoiseType.GREEN_TEST: NoiseSampler.simple(green_noise_like),
     NoiseType.PYRAMID_OLD: NoiseSampler.simple(pyramid_old_noise_like),
     NoiseType.PYRAMID_BISLERP: NoiseSampler.simple(
-        lambda x: pyramid_noise_like(x, upscale_mode="bislerp"),
+        partial(pyramid_noise_like, upscale_mode="bislerp"),
     ),
     NoiseType.HIGHRES_PYRAMID_BISLERP: NoiseSampler.simple(
-        lambda x: highres_pyramid_noise_like(x, upscale_mode="bislerp"),
+        partial(highres_pyramid_noise_like, upscale_mode="bislerp"),
     ),
     NoiseType.PYRAMID_AREA: NoiseSampler.simple(
-        lambda x: pyramid_noise_like(x, upscale_mode="area"),
+        partial(pyramid_noise_like, upscale_mode="area"),
     ),
     NoiseType.HIGHRES_PYRAMID_AREA: NoiseSampler.simple(
-        lambda x: highres_pyramid_noise_like(x, upscale_mode="area"),
+        partial(highres_pyramid_noise_like, upscale_mode="area"),
     ),
     NoiseType.PYRAMID_OLD_BISLERP: NoiseSampler.simple(
-        lambda x: pyramid_old_noise_like(x, upscale_mode="bislerp"),
+        partial(pyramid_old_noise_like, upscale_mode="bislerp"),
     ),
     NoiseType.PYRAMID_OLD_AREA: NoiseSampler.simple(
-        lambda x: pyramid_old_noise_like(x, upscale_mode="area"),
+        partial(pyramid_old_noise_like, upscale_mode="area"),
     ),
     NoiseType.PYRAMID_DISCOUNT5: NoiseSampler.simple(
-        lambda x: pyramid_noise_like(x, discount=0.5),
+        partial(pyramid_noise_like, discount=0.5),
     ),
     NoiseType.PYRAMID_MIX: NoiseSampler.simple(
         lambda x: pyramid_noise_like(x, discount=0.6)

py/noise_generation.py

@@ -6,42 +6,56 @@ from enum import Enum, auto
 from typing import Callable
 
 import torch
+from comfy.model_management import device_supports_non_blocking
 from comfy.utils import common_upscale
 from torch import FloatTensor, Generator, Tensor
 from torch.distributions import Laplace, StudentT
 
+from .external import MODULES as EXT
+
 # ruff: noqa: D412, D413, D417, D212, D407, ANN002, ANN003, FBT001, FBT002, S311
 
 
 class NoiseType(Enum):
-    GAUSSIAN = auto()
-    UNIFORM = auto()
     BROWNIAN = auto()
-    PERLIN = auto()
-    STUDENTT = auto()
-    HIGHRES_PYRAMID = auto()
-    PYRAMID = auto()
-    PYRAMID_MIX = auto()
-    PINK = auto()
-    LAPLACIAN = auto()
-    POWER = auto()
-    RAINBOW_MILD = auto()
-    RAINBOW_INTENSE = auto()
+    GAUSSIAN = auto()
     GREEN_TEST = auto()
-    PYRAMID_OLD = auto()
-    PYRAMID_BISLERP = auto()
-    HIGHRES_PYRAMID_BISLERP = auto()
-    PYRAMID_OLD_BISLERP = auto()
-    PYRAMID_OLD_AREA = auto()
-    PYRAMID_AREA = auto()
+    GREY = auto()
+    HIGHRES_PYRAMID = auto()
     HIGHRES_PYRAMID_AREA = auto()
+    HIGHRES_PYRAMID_BISLERP = auto()
+    LAPLACIAN = auto()
+    ONEF_GREENISH = auto()
+    ONEF_GREENISH_MIX = auto()
+    ONEF_PINKISH = auto()
+    ONEF_PINKISH_MIX = auto()
+    ONEF_PINKISHGREENISH = auto()
+    PERLIN = auto()
+    PINK_OLD = auto()
+    POWER_OLD = auto()
+    PYRAMID = auto()
+    PYRAMID_AREA = auto()
+    PYRAMID_BISLERP = auto()
     PYRAMID_DISCOUNT5 = auto()
-    PYRAMID_MIX_BISLERP = auto()
+    PYRAMID_MIX = auto()
     PYRAMID_MIX_AREA = auto()
+    PYRAMID_MIX_BISLERP = auto()
+    PYRAMID_OLD = auto()
+    PYRAMID_OLD_AREA = auto()
+    PYRAMID_OLD_BISLERP = auto()
+    RAINBOW_INTENSE = auto()
+    RAINBOW_MILD = auto()
+    STUDENTT = auto()
+    UNIFORM = auto()
+    VELVET = auto()
+    VIOLET = auto()
+    WHITE = auto()
 
     @classmethod
-    def get_names(cls, default=None, skip=None):
+    def get_names(cls, default=GAUSSIAN, skip=None):
         if default is not None:
+            if isinstance(default, int):
+                default = cls(default)
             yield default.name.lower()
         for nt in cls:
             if nt == default or (skip and nt in skip):
@@ -65,6 +79,32 @@ def scale_noise(noise, factor=1.0, *, normalized=True, threshold_std_devs=2.5):
     return noise.mul_(factor) if factor != 1 else noise
 
 
+if "bleh" in EXT:
+    scale_samples = EXT["bleh"].py.latent_utils.scale_samples
+else:
+
+    def scale_samples(
+        samples,
+        width,
+        height,
+        *,
+        mode="bicubic",
+    ):
+        return common_upscale(samples, width, height, mode, None)
+
+
+CAN_NONBLOCK = {}
+
+
+def tensor_to(tensor, dest):
+    device = dest.device if isinstance(dest, torch.Tensor) else dest
+    non_blocking = CAN_NONBLOCK.get(device)
+    if non_blocking is None:
+        non_blocking = device_supports_non_blocking(device)
+        CAN_NONBLOCK[device] = non_blocking
+    return tensor.to(dest, non_blocking=non_blocking)
+
+
 def get_positions(block_shape: tuple[int, int]) -> Tensor:
     """
     Generate position tensor.
@@ -95,7 +135,7 @@ def unfold_grid(vectors: Tensor) -> Tensor:
     Returns:
         batched grid vectors
     """
-    batch_size, _, gpy, gpx = vectors.shape
+    batch_size, _channels, gpy, gpx = vectors.shape
     return (
         torch.nn.functional.unfold(vectors, (2, 2))
         .view(batch_size, 2, 4, -1)
@@ -133,7 +173,7 @@ def perlin_noise_tensor(
         step -- smooth step function [0, 1] -> [0, 1] (default: `smooth_step`)
 
     Raises:
-        Exception: if position and vector shapes do not match
+        NoiseError: if position and vector shapes do not match
 
     Returns:
         (batch_size, block_height * grid_height, block_width * grid_width)
@@ -148,11 +188,11 @@ def perlin_noise_tensor(
     bh, bw = positions.shape[1:3]
 
     for i in range(2):
-        if positions.shape[i + 3] not in (1, vectors.shape[i + 2]):
+        if positions.shape[i + 3] not in {1, vectors.shape[i + 2]}:
             msg = f"Blocks shapes do not match: vectors ({vectors.shape[1]}, {vectors.shape[2]}), positions {gh}, {gw})"
             raise NoiseError(msg)
 
-    if positions.shape[0] not in (1, batch_size):
+    if positions.shape[0] not in {1, batch_size}:
         msg = f"Batch sizes do not match: vectors ({vectors.shape[0]}), positions ({positions.shape[0]})"
         raise NoiseError(msg)
 
@@ -206,7 +246,7 @@ def perlin_noise(
         generator -- random generator used for grid vectors (default: {None})
 
     Raises:
-        Exception: if grid and out shapes do not match
+        NoiseError: if grid and out shapes do not match
 
     Returns:
         Noise image shaped (batch_size, height, width)
@@ -233,28 +273,55 @@ def perlin_noise(
     # random vectors on grid points
     vectors = unfold_grid(torch.stack((torch.cos(angle), torch.sin(angle)), dim=1))
     # positions inside grid cells [0, 1)
-    positions = get_positions((bh, bw)).to(vectors)
+    positions = tensor_to(get_positions((bh, bw)), vectors)
     return perlin_noise_tensor(vectors, positions).squeeze(0)
 
 
-def rand_perlin_like(x):
-    noise = torch.randn_like(x) / 2.0
+def rand_perlin_like(x, *, generator=None):
+    noise = (
+        torch.rand(
+            x.shape,
+            dtype=x.dtype,
+            device=x.device,
+            layout=x.layout,
+            generator=generator,
+        )
+        / 2.0
+    )
     noise_height = noise.size(dim=2)
     noise_width = noise.size(dim=3)
     for _ in range(2):
-        noise += perlin_noise(
-            (noise_height, noise_width),
-            (noise_height, noise_width),
-            batch_size=x.shape[1],  # This should be the number of channels.
-        ).to(x.device)
+        noise += tensor_to(
+            perlin_noise(
+                (noise_height, noise_width),
+                (noise_height, noise_width),
+                batch_size=x.shape[1],  # This should be the number of channels.
+            ),
+            x.device,
+        )
     return scale_noise(noise)
 
 
-def uniform_noise_like(x):
-    return (torch.rand_like(x) - 0.5) * 3.46
+def uniform_noise_like(x, *, generator=None):
+    return (
+        torch.rand(
+            x.shape,
+            dtype=x.dtype,
+            device=x.device,
+            layout=x.layout,
+            generator=generator,
+        ).sub_(0.5)
+    ).mul_(3.46)
 
 
-def highres_pyramid_noise_like(x, discount=0.7, upscale_mode="bilinear"):
+def highres_pyramid_noise_like(
+    x,
+    *,
+    discount=0.7,
+    upscale_mode="bilinear",
+    iterations=4,
+    generator=None,
+):
     (
         b,
         c,
@@ -262,17 +329,16 @@ def highres_pyramid_noise_like(x, discount=0.7, upscale_mode="bilinear"):
         w,
     ) = x.shape  # EDIT: w and h get over-written, rename for a different variant!
     orig_w, orig_h = w, h
-    noise = uniform_noise_like(x)
-    rs = torch.rand(4, dtype=torch.float32) * 2 + 2
-    for i in range(4):
-        r = rs[i]
+    noise = uniform_noise_like(x, generator=generator)
+    rs = torch.rand(iterations, dtype=torch.float32, generator=generator).cpu() * 2 + 2
+    for i in range(iterations):
+        r = rs[i].item()
         h, w = min(orig_h * 15, int(h * (r**i))), min(orig_w * 15, int(w * (r**i)))
-        noise += common_upscale(
-            torch.randn(b, c, h, w).to(x),
+        noise += scale_samples(
+            tensor_to(torch.randn(b, c, h, w, generator=generator), x),
             orig_w,
             orig_h,
-            upscale_mode,
-            None,
+            mode=upscale_mode,
         ).mul_(discount**i)
         if h >= orig_h * 15 or w >= orig_w * 15:
             break  # Lowest resolution is 1x1
@@ -281,19 +347,21 @@ def highres_pyramid_noise_like(x, discount=0.7, upscale_mode="bilinear"):
 
 def pyramid_old_noise_like(
     x,
+    *,
     generator=None,
     device="cpu",
     discount=0.8,
+    iterations=5,
     upscale_mode="nearest-exact",
 ):
-    size = x.size()
+    size = x.shape
     b, c, h, w = size
     orig_h, orig_w = h, w
     noise = torch.zeros(size=size, dtype=x.dtype, layout=x.layout, device=device)
     r = 1
-    for i in range(5):
+    for i in range(iterations):
         r *= 2
-        noise += common_upscale(
+        noise += scale_samples(
             torch.normal(
                 mean=0,
                 std=0.5**i,
@@ -305,28 +373,35 @@ def pyramid_old_noise_like(
             ),
             orig_w,
             orig_h,
-            upscale_mode,
-            None,
+            mode=upscale_mode,
         ).mul_(discount**i)
-    return noise.to(device=x.device)
+    return tensor_to(noise, x.device)
 
 
 # Copied from https://wandb.ai/johnowhitaker/multires_noise/reports/Multi-Resolution-Noise-for-Diffusion-Model-Training--VmlldzozNjYyOTU2
-def pyramid_noise_like(x, discount=0.7, upscale_mode="bilinear"):
+def pyramid_noise_like(
+    x,
+    *,
+    discount=0.7,
+    upscale_mode="bilinear",
+    iterations=10,
+    generator=None,
+):
     b, c, w, h = (
         x.shape
     )  # NOTE: w and h get over-written, rename for a different variant!
     orig_w, orig_h = w, h
     noise = torch.randn_like(x)
-    for i in range(10):
-        r = torch.rand(1, device="cpu").item() * 2 + 2  # Rather than always going 2x,
+    for i in range(iterations):
+        r = (
+            torch.rand(1, generator=generator).cpu().item() * 2 + 2
+        )  # Rather than always going 2x,
         w, h = max(1, int(w / (r**i))), max(1, int(h / (r**i)))
-        noise += common_upscale(
-            torch.randn(b, c, w, h).to(x),
+        noise += scale_samples(
+            tensor_to(torch.randn(b, c, w, h), x),
             orig_h,
             orig_w,
-            upscale_mode,
-            None,
+            mode=upscale_mode,
         ).mul_(
             discount**i,
         )
@@ -336,57 +411,119 @@ def pyramid_noise_like(x, discount=0.7, upscale_mode="bilinear"):
 
 
 def studentt_noise_like(x):
-    noise = StudentT(loc=0, scale=0.2, df=1).rsample(x.size())
+    noise = StudentT(loc=0, scale=0.2, df=1).rsample(x.shape)
     s: FloatTensor = torch.quantile(noise.flatten(start_dim=1).abs(), 0.75, dim=-1)
     s = s.reshape(*s.shape, 1, 1, 1)
     noise = noise.clamp(-s, s)
     return torch.copysign(torch.pow(torch.abs(noise), 0.5), noise)
 
 
-def green_noise_like(x):
+def green_noise_like(x, *, generator=None):  # noqa: ARG001
     # The comments said this didn't work and I had to learn the hard way. Turns out it's true!
-    width, height = x.size(dim=2), x.size(dim=3)
+    height, width = x.shape[-2:]
     noise = torch.randn_like(x)
     scale = 1.0 / (width * height)
-    fy = torch.fft.fftfreq(width, device=x.device)[:, None] ** 2
-    fx = torch.fft.fftfreq(height, device=x.device) ** 2
+    fy = torch.fft.fftfreq(height, device=x.device)[:, None] ** 2
+    fx = torch.fft.fftfreq(width, device=x.device) ** 2
     f = fy + fx
     power = torch.sqrt(f)
     power[0, 0] = 1
     noise = torch.fft.ifft2(torch.fft.fft2(noise) / torch.sqrt(power))
     noise *= scale / noise.std()
-    noise = torch.real(noise).to(x.device)
+    noise = tensor_to(torch.real(noise), x.device)
     return scale_noise(noise)
 
 
-def generate_1f_noise(tensor, alpha, k, generator=None):
-    """Generate 1/f noise for a given tensor.
-
-    Args:
-        tensor: The tensor to add noise to.
-        alpha: The parameter that determines the slope of the spectrum.
-        k: A constant.
-
-    Returns:
-        A tensor with the same shape as `tensor` containing 1/f noise.
-    """
-    fft = torch.fft.fft2(tensor)
-    freq = torch.arange(1, len(fft) + 1, dtype=torch.float)
+# Completely wrong implementation here.
+def generate_1f_noise_old(tensor, alpha, k, generator=None):
+    freq = 1.0
     spectral_density = k / freq**alpha
     return torch.randn(tensor.shape, generator=generator) * spectral_density
 
 
-def pink_noise_like(x):
-    return scale_noise(generate_1f_noise(x, 2.0, 1.0)).to(x.device)
+def pink_noise_old_like(x, *, generator=None):
+    return tensor_to(
+        scale_noise(generate_1f_noise_old(x, 2.0, 1.0, generator=generator)),
+        x.device,
+    )
+
+
+# Referenced from: https://github.com/WASasquatch/PowerNoiseSuite
+def generate_1f_noise(
+    tensor,
+    *,
+    alpha=-2.0,
+    k=1.0,
+    hfac=1.0,
+    wfac=1.0,
+    base_power=1.0,
+    use_sqrt=True,
+    generator=None,
+):
+    batch, _channels, height, width = tensor.shape
+    noise = torch.randn(tensor.shape, generator=generator)
+
+    freq_x = torch.fft.fftfreq(height, hfac)
+    freq_y = torch.fft.fftfreq(width, wfac)
+    fx, fy = torch.meshgrid(freq_x, freq_y, indexing="ij")
+
+    power = (fx**2 + fy**2) ** (-alpha / 2.0)
+    if k != 0:
+        power = k / power
+    power[0, 0] = base_power
+    power = power.unsqueeze(0).expand(batch, 1, height, width)
+
+    noise_fft = torch.fft.fftn(noise)
+    noise_fft /= (
+        torch.sqrt(power.to(noise_fft.dtype)) if use_sqrt else power.to(noise_fft.dtype)
+    )
+
+    return torch.fft.ifftn(noise_fft).real
+
+
+def onef_noise_like(x, *, generator=None, **kwargs):
+    return tensor_to(
+        scale_noise(generate_1f_noise(x, generator=generator, **kwargs)),
+        x.device,
+    )
+
+
+# Referenced from: https://github.com/WASasquatch/PowerNoiseSuite
+def generate_powerlaw_noise(
+    tensor: torch.Tensor,
+    *,
+    alpha=1.0,
+    div_max_dims=None,
+    use_sign=False,
+    use_div_max_abs=True,
+    generator=None,
+) -> torch.Tensor:
+    noise = torch.randn(tensor.shape, generator=generator)
+    modulation = torch.abs(noise) ** alpha
+    noise = (torch.sign(noise) if use_sign else noise).mul_(modulation)
+    if div_max_dims is not None:
+        noise /= torch.amax(
+            torch.abs(noise) if use_div_max_abs else noise,
+            keepdim=True,
+            dim=div_max_dims,
+        )
+    return noise
+
+
+def powerlaw_noise_like(x, *, generator=None, **kwargs):
+    return tensor_to(
+        scale_noise(generate_powerlaw_noise(x, generator=generator, **kwargs)),
+        x.device,
+    )
 
 
 def laplacian_noise_like(x):
     noise = torch.randn_like(x).div_(4.0)
-    noise += Laplace(loc=0, scale=1.0).rsample(x.size()).to(noise.device)
+    noise += tensor_to(Laplace(loc=0, scale=1.0).rsample(x.shape), noise.device)
     return scale_noise(noise)
 
 
-def power_noise_like(tensor, alpha=2, k=1):  # This doesn't work properly right now
+def power_noise_old_like(tensor, alpha=2, k=1):  # This doesn't work properly right now
     """Generate 1/f noise for a given tensor.
 
     Args:
@@ -403,24 +540,26 @@ def power_noise_like(tensor, alpha=2, k=1):  # This doesn't work properly right
         (len(fft),) + (1,) * (tensor.dim() - 1),
     )
     spectral_density = k / freq**alpha
-    noise = torch.rand(tensor.shape).mul_(spectral_density)
-    mean = torch.mean(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
-    std = torch.std(noise, dim=(-2, -1), keepdim=True).to(tensor.device)
-    return noise.to(tensor.device).sub_(mean).div_(std)
+    noise = tensor_to(torch.rand(tensor.shape).mul_(spectral_density), tensor.device)
+    mean = torch.mean(noise, dim=(-2, -1), keepdim=True)
+    std = torch.std(noise, dim=(-2, -1), keepdim=True)
+    return noise.sub_(mean).div_(std)
 
 
 __all__ = (
-    "NoiseType",
     "NoiseError",
-    "scale_noise",
+    "NoiseType",
     "green_noise_like",
     "highres_pyramid_noise_like",
     "laplacian_noise_like",
-    "pink_noise_like",
-    "power_noise_like",
+    "onef_noise_like",
+    "pink_noise_old_like",
+    "power_noise_old_like",
+    "powerlaw_noise_like",
     "pyramid_noise_like",
     "pyramid_old_noise_like",
     "rand_perlin_like",
+    "scale_noise",
     "studentt_noise_like",
     "uniform_noise_like",
 )

py/powernoise.py

@@ -16,7 +16,11 @@ from comfy.k_diffusion.sampling import BrownianTreeNoiseSampler
 from PIL import Image
 from torch import Tensor
 
-from .nodes import SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin
+from .nodes import (
+    WILDCARD_NOISE,
+    SonarCustomNoiseNodeBase,
+    SonarNormalizeNoiseNodeMixin,
+)
 from .noise import CustomNoiseItemBase
 from .noise_generation import scale_noise
 
@@ -69,8 +73,10 @@ class ChannelMixer:
                 ),
             ),
         )
-        channel_mixer = torch.eye(c)
-        channel_mixer[*torch.tril_indices(c, c, offset=-1)] = channel_correlation
+        channel_mixer = torch.eye(c).index_put_(
+            tuple(torch.tril_indices(c, c, offset=-1)),
+            channel_correlation,
+        )
         channel_mixer += channel_mixer.tril(-1).mT
         channel_mixer = torch.linalg.ldl_factor(channel_mixer).LD
         dc = torch.diagonal_copy(channel_mixer)
@@ -110,7 +116,7 @@ class PowerFilter:
         scale=1.0,
         rel_bw=0.125,
         oversample=4,
-        compose_with: None | PowerFilter = None,
+        compose_with: PowerFilter | None = None,
         compose_mode="max",
     ):
         self.min_freq = min_freq
@@ -514,7 +520,7 @@ class PowerFilterNoiseItem(PowerNoiseItem):
             x,
             ns,
             self.make_filter(x.shape),
-            self.normalize_result in (True, None),
+            self.normalize_result in {True, None},
         )
         filtered_noise = filtered_ns(
             torch.scalar_tensor(14.0),
@@ -531,11 +537,19 @@ class PowerFilterNoiseItem(PowerNoiseItem):
 
 
 class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
+    DESCRIPTION = "Custom noise type that applies a filter to generated noise."
+
     @classmethod
     def INPUT_TYPES(cls, *args: list, **kwargs: dict):
         result = super().INPUT_TYPES(*args, **kwargs)
         result["required"] |= {
-            "time_brownian": ("BOOLEAN", {"default": False}),
+            "time_brownian": (
+                "BOOLEAN",
+                {
+                    "default": False,
+                    "tooltip": "Controls whether brownian noise is used when mix isn't 1.0.",
+                },
+            ),
             "alpha": (
                 "FLOAT",
                 {
@@ -544,6 +558,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
                     "max": 5.0,
                     "step": 0.001,
                     "round": False,
+                    "tooltip": "Values above 0 will amplify low frequencies, negative values will amplify high frequencies.",
                 },
             ),
             "max_freq": (
@@ -554,6 +569,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
                     "max": 0.7071,
                     "step": 0.001,
                     "round": False,
+                    "tooltip": "Maximum frequency to pass through the filter.",
                 },
             ),
             "min_freq": (
@@ -564,6 +580,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
                     "max": 0.7071,
                     "step": 0.001,
                     "round": False,
+                    "tooltip": "Minimum frequency to pass through the filter.",
                 },
             ),
             "stretch": (
@@ -574,6 +591,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
                     "max": 100,
                     "step": 0.1,
                     "round": False,
+                    "tooltip": "Stretches the filter's shape by the specified factor.",
                 },
             ),
             "rotate": (
@@ -584,6 +602,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
                     "max": 90,
                     "step": 5,
                     "round": False,
+                    "tooltip": "Rotates the filter.",
                 },
             ),
             "pnorm": (
@@ -594,6 +613,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
                     "max": 100,
                     "step": 0.1,
                     "round": False,
+                    "tooltip": "Factor used for cushioning the band-pass region.",
                 },
             ),
             "mix": (
@@ -604,6 +624,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
                     "max": 1.0,
                     "step": 0.001,
                     "round": False,
+                    "tooltip": "Controls the ratio of filtered noise. For example, 0.75 means 75% noise with the filter effects applied, 25% raw noise.",
                 },
             ),
             "common_mode": (
@@ -614,6 +635,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
                     "max": 100.0,
                     "step": 0.001,
                     "round": False,
+                    "tooltip": "Attempts to desaturate thelatent by injecting the average across channels (controlled by channel_correction). Applied after mix.",
                 },
             ),
             "channel_correlation": (
@@ -622,13 +644,20 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
                     "default": "1, 1, 1, 1, 1, 1",
                     "multiline": False,
                     "dynamicPrompts": False,
+                    "tooltip": "Comma-separated list of channel correlation strengths.",
+                },
+            ),
+            "preview": (
+                ("none", "no_mix", "mix"),
+                {
+                    "tooltip": "When enabled, displays a preview of the filter shape and a sample of noise. Mix - previews noise after mix is applied. no_mix - only previews the filtered noise.",
                 },
             ),
-            "preview": (("none", "no_mix", "mix"),),
         }
         return result
 
-    def get_item_class(self):
+    @classmethod
+    def get_item_class(cls):
         return PowerNoiseItem
 
     def go(
@@ -646,6 +675,8 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase):
 
 
 class SonarPowerFilterNoiseNode(SonarPowerNoiseNode, SonarNormalizeNoiseNodeMixin):
+    DESCRIPTION = "Custom noise type that allows applying a Power Filter to another custom noise generator."
+
     @classmethod
     def INPUT_TYPES(cls):
         result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
@@ -660,8 +691,18 @@ class SonarPowerFilterNoiseNode(SonarPowerNoiseNode, SonarNormalizeNoiseNodeMixi
         ):
             del result["required"][k]
         result["required"] |= {
-            "sonar_custom_noise": ("SONAR_CUSTOM_NOISE",),
-            "sonar_power_filter": ("SONAR_POWER_FILTER",),
+            "sonar_custom_noise": (
+                WILDCARD_NOISE,
+                {
+                    "tooltip": "Custom noise type to filter.",
+                },
+            ),
+            "sonar_power_filter": (
+                "SONAR_POWER_FILTER",
+                {
+                    "tooltip": "Filter to use.",
+                },
+            ),
             "filter_norm_factor": (
                 "FLOAT",
                 {
@@ -670,15 +711,32 @@ class SonarPowerFilterNoiseNode(SonarPowerNoiseNode, SonarNormalizeNoiseNodeMixi
                     "max": 1.0,
                     "step": 0.1,
                     "round": False,
+                    "tooltip": "Normalization factor applied to the specified filter. 1.0 means 100% normalized.",
+                },
+            ),
+            "normalize_result": (
+                ("default", "forced", "disabled"),
+                {
+                    "tooltip": "Controls whether the final result is normalized to 1.0 strength.",
+                },
+            ),
+            "normalize_noise": (
+                ("default", "forced", "disabled"),
+                {
+                    "tooltip": "Controls whether the generated noise is normalized to 1.0 strength.",
                 },
             ),
-            "normalize_result": (("default", "forced", "disabled"),),
-            "normalize_noise": (("default", "forced", "disabled"),),
         }
-        result["required"]["preview"] = ((*result["required"]["preview"][0], "custom"),)
+        result["required"]["preview"] = (
+            (*result["required"]["preview"][0], "custom"),
+            {
+                "tooltip": "When enabled, displays a preview of the filter shape and a sample of noise. Mix - previews noise after mix is applied. no_mix - only previews the filtered noise. custom - Like no_mix, but will use a latent previewer to display a color preview of the generated noise. Works best when previewer is set to TAESD.",
+            },
+        )
         return result
 
-    def get_item_class(self):
+    @classmethod
+    def get_item_class(cls):
         return PowerFilterNoiseItem
 
     def go(
@@ -712,69 +770,23 @@ class SonarPowerFilterNode:
 
     @classmethod
     def INPUT_TYPES(cls):
+        include_keys = {"alpha", "max_freq", "min_freq", "stretch", "rotate", "pnorm"}
         return {
             "required": {
-                "alpha": (
-                    "FLOAT",
+                k: v
+                for k, v in SonarPowerNoiseNode.INPUT_TYPES()["required"].items()
+                if k in include_keys
+            }
+            | {
+                "oversample": (
+                    "INT",
                     {
-                        "default": 0.0,
-                        "min": -5.0,
-                        "max": 5.0,
-                        "step": 0.001,
-                        "round": False,
+                        "default": 4,
+                        "min": 1,
+                        "max": 128,
+                        "tooltip": "Oversampling factor used for the filter size.",
                     },
                 ),
-                "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,
-                    },
-                ),
-                "oversample": ("INT", {"default": 4, "min": 1, "max": 128}),
                 "blur": (
                     "FLOAT",
                     {
@@ -783,6 +795,7 @@ class SonarPowerFilterNode:
                         "max": 10.0,
                         "step": 0.01,
                         "round": False,
+                        "tooltip": "Slightly blurs the filter to reduce artifacts.",
                     },
                 ),
                 "scale": (
@@ -793,17 +806,24 @@ class SonarPowerFilterNode:
                         "max": 100.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Scales the filter to the specified strength. May be negative.",
+                    },
+                ),
+                "compose_mode": (
+                    ("max", "min", "add", "sub", "mul"),
+                    {
+                        "tooltip": "Controls composition of the option attached filter. For example, when set to MUL the result will be this filter multiplied by the attached filter. No effect if the optional filter input is not attached.",
                     },
                 ),
-                "compose_mode": (("max", "min", "add", "sub", "mul"),),
             },
             "optional": {
                 "power_filter_opt": ("SONAR_POWER_FILTER",),
             },
         }
 
+    @classmethod
     def go(
-        self,
+        cls,
         min_freq=0.0,
         max_freq=0.7071,
         stretch=1.0,
@@ -834,6 +854,7 @@ class SonarPowerFilterNode:
 
 
 class SonarPreviewFilterNode:
+    DESCRIPTION = "Allows previewing a Power Filter."
     RETURN_TYPES = ("SONAR_POWER_FILTER",)
     CATEGORY = "advanced/noise"
     FUNCTION = "go"
@@ -843,7 +864,12 @@ class SonarPreviewFilterNode:
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "sonar_power_filter": ("SONAR_POWER_FILTER",),
+                "sonar_power_filter": (
+                    "SONAR_POWER_FILTER",
+                    {
+                        "tooltip": "Power Filter to preview.",
+                    },
+                ),
                 "filter_gain": (
                     "FLOAT",
                     {
@@ -852,6 +878,7 @@ class SonarPreviewFilterNode:
                         "max": 1000000.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Gain factor applied to the filter part of the preview.",
                     },
                 ),
                 "kernel_gain": (
@@ -862,6 +889,7 @@ class SonarPreviewFilterNode:
                         "max": 1000000.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Gain factor applied to the kernel part of the preview.",
                     },
                 ),
                 "norm_factor": (
@@ -872,6 +900,7 @@ class SonarPreviewFilterNode:
                         "max": 1.0,
                         "step": 0.1,
                         "round": False,
+                        "tooltip": "Normalization factor applied to the filter before previewing. 1.0 means 100% normalized.",
                     },
                 ),
                 "preview_size": (
@@ -886,12 +915,16 @@ class SonarPreviewFilterNode:
                         "128x127",
                         "127x128",
                     ),
+                    {
+                        "tooltip": "Controls the size of the generated preview. Note: Sizes are in latent pixels. For most models, one latent pixel equals eight pixels",
+                    },
                 ),
             },
         }
 
+    @classmethod
     def go(
-        self,
+        cls,
         sonar_power_filter,
         filter_gain=1 / 3,
         kernel_gain=1 / 3,

py/sonar.py

@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
+import importlib
 from enum import Enum, auto
 from sys import stderr
 from typing import Any, Callable, NamedTuple
 
 import torch
 from comfy.k_diffusion import sampling
+from comfy.samplers import KSampler, k_diffusion_sampling
 from torch import Tensor
 from tqdm.auto import trange
 
@@ -60,10 +62,10 @@ class SonarBase:
         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 (
+        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,
@@ -127,7 +129,7 @@ class SonarBase:
         momentum_d = (1.0 - p) * d + p * hd
 
         # Euler method with momentum
-        x = x + momentum_d * dt
+        x = x + momentum_d * dt  # noqa: PLR6104
 
         self.update_hist(momentum_d)
 
@@ -155,9 +157,7 @@ class SonarGuidanceMixin:
         return ((latent - avg_s) / std_s).to(latent.dtype)
 
     def guidance_step(self, step_index: int, x: Tensor, denoised: Tensor):
-        if (self.guidance is None or self.guidance.factor == 0.0) or not (
-            self.guidance.start_step <= (step_index + 1) <= self.guidance.end_step
-        ):
+        if self.guidance is None or self.guidance.factor == 0.0 or not self.guidance.start_step <= step_index <= self.guidance.end_step:
             return x
         if self.ref_latent.device != x.device:
             self.ref_latent = self.ref_latent.to(device=x.device)
@@ -263,7 +263,7 @@ class SonarEuler(SonarSampler):
                 else torch.randn_like(sample)
             )
             eps = noise * self.s_noise
-            sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
+            sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5  # noqa: PLR6104
 
         denoised = self.model(sample, sigma_hat * self.s_in, **self.extra_args)
         derivative = sampling.to_d(sample, sigma, denoised)
@@ -320,7 +320,7 @@ class SonarEuler(SonarSampler):
         )
 
         for i in trange(len(sigmas) - 1, disable=disable):
-            x, sigma, sigma_hat, denoised = sonar.step(
+            x, _sigma, sigma_hat, denoised = sonar.step(
                 i,
                 x,
             )
@@ -370,7 +370,7 @@ class SonarEulerAncestral(SonarSampler):
         result_sample = self.momentum_step(sample, derivative, dt)
         if sigma_to > 0:
             result_sample = self.guidance_step(step_index, result_sample, denoised)
-            result_sample = (
+            result_sample = (  # noqa: PLR6104
                 result_sample
                 + self.noise_sampler(sigma_from, sigma_to) * self.s_noise * sigma_up
             )
@@ -417,7 +417,7 @@ class SonarEulerAncestral(SonarSampler):
         )
 
         for i in trange(len(sigmas) - 1, disable=disable):
-            x, sigma, sigma_hat, denoised = sonar.step(
+            x, _sigma, sigma_hat, denoised = sonar.step(
                 i,
                 x,
             )
@@ -457,7 +457,7 @@ class SonarDPMPPSDE(SonarSampler):
         return sigma.log.neg()
 
     # DPM++ solver algorithm copied from ComfyUI source.
-    def momentum_step(
+    def momentum_step(  # noqa: PLR0914
         self,
         step_index,
         x: Tensor,
@@ -495,7 +495,7 @@ class SonarDPMPPSDE(SonarSampler):
         self.update_hist(momentum_d)
         hd = self.history_d
         x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - momentum_d
-        x_2 = x_2 + self.noise_sampler(sigma_fn(t), sigma_fn(s)) * self.s_noise * su
+        x_2 += self.noise_sampler(sigma_fn(t), sigma_fn(s)) * self.s_noise * su
         denoised_2 = self.model(x_2, sigma_fn(s) * self.s_in, **self.extra_args)
 
         # Step 2
@@ -527,7 +527,7 @@ class SonarDPMPPSDE(SonarSampler):
         self.init_hist_d(sample)
 
         sigma_from, sigma_to = self.sigmas[step_index], self.sigmas[step_index + 1]
-        sigma_down, sigma_up = sampling.get_ancestral_step(
+        sigma_down, _sigma_up = sampling.get_ancestral_step(
             sigma_from,
             sigma_to,
             eta=self.eta,
@@ -585,7 +585,7 @@ class SonarDPMPPSDE(SonarSampler):
         )
 
         for i in trange(len(sigmas) - 1, disable=disable):
-            x, sigma, sigma_hat, denoised = sonar.step(
+            x, _sigma, sigma_hat, denoised = sonar.step(
                 i,
                 x,
             )
@@ -603,10 +603,6 @@ class SonarDPMPPSDE(SonarSampler):
 
 
 def add_samplers():
-    import importlib
-
-    from comfy.samplers import KSampler, k_diffusion_sampling
-
     extra_samplers = {
         "sonar_euler": SonarEuler.sampler,
         "sonar_euler_ancestral": SonarEulerAncestral.sampler,

ruff.toml

@@ -8,6 +8,8 @@ ignore = [
     "ANN204",
     "ANN206",
     "C901",
+    "CPY001",
+    "DOC201",
     "D100",
     "D101",
     "D102",
@@ -26,6 +28,7 @@ ignore = [
     "PLR0912",
     "PLR0913",
     "PLR0915",
+    "PLR0917",
     "PLR2004",
     "T201",
     "TRY003",