diff --git a/README.md b/README.md
index 183419d..2768471 100644
--- a/README.md
+++ b/README.md
@@ -64,6 +64,47 @@ You can optionally plug this into the Sonar sampler nodes. See the [Guidance](#g
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...
+
+Click to expand advanced parameters info
+
+There are some extra advanced parameters that may be passed by YAML/JSON using `SamplerConfigOVerride`'s `yaml_parameters`. Defaults:
+
+```yaml
+sonar_params:
+ # One of: classic, new, denoised
+ # classic: Should be the same as the way it works in the A1111 extension.
+ # new: Possibly improved version that doesn't blend in the history again.
+ # denoised: Instead of using the noise prediction, we do momentum on denoised instead.
+ momentum_mode: new
+
+ # The following two parameters may be used to control when
+ # momentum sampling is active. Steps are 0-based with 0 being the first step.
+ momentum_start_step: 0
+ momentum_end_step: 9999
+
+ # Controls whether history always gets updated, whether or not within the
+ # start/end step range or only in that range. Can be used to affect the initial
+ # history value.
+ always_update_history: true
+
+ # Only applies when the init type is RAND.
+ rand_init_noise_multiplier: 1.0
+
+ # If you have ComfyUI-bleh installed, you can use any blend mode it provides.
+ # Otherwise you can have your blend mode in any color you want as long as it's lerp.
+ blend_mode: lerp
+
+ # Defaultss to blend_mode if unset.
+ momentum_blend_mode: null
+
+ # Defaults to blend_mode if unset. Only applies to linear guidance mode.
+ guidance_blend_mode: null
+```
+
+Additionally, it's possible to override the normal Sonar parameters here as well. If they exist in the `sonar_params` block, they will overwrite the values in the node.
+
+
+
## Guidance
You can try the `SamplerSonarNaive` sampler which has an optional latent input. The guidance _probably_ isn't working correctly and the implementation definitely isn't exactly the same as the original A1111 version but it still might be fun to play with. The `linear` guidance type is a lot more sensitive to the `guidance_factor` than the `euler` type. For `euler`, reasonable values are around `0.01` to `0.1`, for `linear` reasonable values are more like `0.001` to `0.02`. It is also possible to set guidance factor to a negative value, I've found this results in high contrast and very vivid colors.
@@ -112,6 +153,7 @@ My version was initially based on this Sonar sampler implementation for Diffuser
* 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
+* Wavelet noise idea (and some of the default settings) from https://github.com/ClownsharkBatwing/RES4LYF
## Errata
diff --git a/changelog.md b/changelog.md
index 220123a..94a68bf 100644
--- a/changelog.md
+++ b/changelog.md
@@ -2,6 +2,22 @@
Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top.
+## 20250130
+
+*Note*: May change seeds.
+
+This set of changes includes some pretty major internal refactoring. Definitely possible that I broke something, so please create an issue if you run into problems.
+
+* Noise generation should now respect whether generating on CPU vs GPU is selected. Previously it likely was defaulting to generating on GPU. This may change seeds.
+* Refactored momentum samplers, this may change seeds especially if you were using weird parameters like negative direction.
+* Added some new parameters for momentum samplers.
+* Removed the `s_noise` and churn parameters from the normal Sonar Euler sampler. May break workflows. (Churn was the predecessor to ancestral samplers and is basically obsolete.)
+* Added `wavelet` and `distro` noise types.
+* Added `SonarCustomNoiseAdv` node that allows passing parameters via YAML.
+* Added `SonarResizedNoise` node that allows you to generate noise at a fixed size and then crop/resize it to match the generation.
+* Added `SonarAdvancedDistroNoise` node that allows generating noise with basically all the distributions PyTorch supports.
+* Added `SonarWaveletFilteredNoise` node that lets you filter another noise generator using wavelets.
+
## 20241129
*Note*: Contains some potentially workflow-breaking changes.
diff --git a/docs/advanced_noise_nodes.md b/docs/advanced_noise_nodes.md
index 079f4a3..b7b8c11 100644
--- a/docs/advanced_noise_nodes.md
+++ b/docs/advanced_noise_nodes.md
@@ -52,6 +52,21 @@ Parameters:
***
+### `SonarCustomNoiseAdv`
+
+Same as the `SonarCustomNoise` except it also includes a text widget for passing parameters by YAML or JSON (JSON is valid YAML).
+
+Just for example, instead of using the absurdly large `SonarAdvancedDistroNoise` node, you could do something like:
+
+```yaml
+distro: wishart
+quantile_norm: 0.5
+wishart_cov_size: 4
+wishart_df: 3.5
+```
+
+***
+
### `NoisyLatentLike`
This node takes a reference latent and generates noise of the same shape. The one required input is `latent`.
@@ -108,6 +123,59 @@ More extensive documentation TBD (hopefully). For now, a few recipes:
***
+## `SonarWaveletFilteredNoise`
+
+You will need [pytorch_wavelets](https://github.com/fbcotter/pytorch_wavelets) installed in your Python environment to use this one.
+
+Allows filtering another noise source using wavelets. Parameters are specified using YAML (or JSON) in the text widget. The defaults are:
+
+```yaml
+use_dtcwt: false
+mode: periodization
+level: 3
+wave: haar
+
+# Only used in DTCWT mode.
+qshift: qshift_a
+# Only used in DTCWT mode.
+biort: near_sym_a
+
+# Additional parameters for the inverse wavelet operation
+# are null by default and will use whatever the
+# forward parameter is set to:
+# inv_mode, inv_wave, inv_biort, inv_qshift
+# Note: Using different parameters for the inverse wavelet
+# operation may not work well (or just fail entirely).
+
+# Scale for the lowpass filter.
+yl_scale: 1.0
+
+# Scales for the highpass filter. Can be a single value (null is basically 1.0).
+yh_scales: null
+```
+
+***
+
+### `SonarAdvancedDistroNoise`
+
+See: https://pytorch.org/docs/stable/distributions.html
+
+For the most part, we just pass parameters directly to PyTorch's distribution classes. Some of them have specific requirements so it is possible to set invalid parameters.
+
+It may be more convenient to specify parameters using the `SonarCustomNoiseAdv` node than this gigantic monstrosity of a node. **Note**: In that case, pass the distribution name using `distro`, i.e. `distro: laplacian`.
+
+Common parameters:
+
+* `quantile_norm`: When enabled, will normalize generated noise to this quantile (i.e. 0.75 means outliers >75% will be clipped). Set to 1.0 or 0.0 to disable quantile normalization. A value like 0.75 or 0.85 should be reasonable, it really depends on the distribution and how many of the values are extreme. Some actually work better with quantile normalization disabled.
+* `quantile_norm_mode`: Controls what dimensions quantile normalization uses. By default, the noise is flattened first. You can try the nonflat versions but they may have a very strong row/column influence. Only applies when quantile_norm is active.
+* `result_index`: When noise generation returns a batch of items, it will select the specified index. Negative indexes count from the end. Values outside the valid range will be automatically adjusted. You may enter a space-separated list of values for the case where there might be multiple added batch dimensions. Excess batch dimensions are removed from the end, indexe from result_index are used in order so you may want to enter the indexes in reverse order. Example: If your noise has shape `(1, 4, 3, 3)` and two 2-sized batch dims are added resulting in `(1, 4, 3, 3, 2, 2)` and you wanted index 0 from the first additional batch dimension and 1 from the second you would use result_index: `1 0`
+
+Individual distributions have parameters beginning with their name, i.e. `laplacian_loc`. Parameters that are string inputs usually allow entering multiple space-separated items. This will usually result in the output noise being a batch, which can be selected with the `result_index` parameter.
+
+Suggestions for fun distributions to try: Wishart and VonMises can produce some interesting results.
+
+***
+
### `SonarModulatedNoise`
Experimental noise modulation based on code stolen from
diff --git a/py/external.py b/py/external.py
index 4d2474c..1faed51 100644
--- a/py/external.py
+++ b/py/external.py
@@ -1,22 +1,125 @@
+from __future__ import annotations
+
import contextlib
import importlib
+import sys
+from functools import partial
+from typing import TYPE_CHECKING, Callable, NamedTuple
-MODULES = {}
+if TYPE_CHECKING:
+ from types import ModuleType
-with contextlib.suppress(ImportError, NotImplementedError):
- bleh = importlib.import_module("custom_nodes.ComfyUI-bleh")
- bleh_version = getattr(bleh, "BLEH_VERSION", -1)
- if bleh_version < 1:
- raise NotImplementedError
- MODULES["bleh"] = bleh
-with contextlib.suppress(ImportError, NotImplementedError):
- import custom_nodes.ComfyUI_restart_sampling as rs
+class Integrations:
+ class Integration(NamedTuple):
+ key: str
+ module_name: str
+ handler: Callable | None = None
+
+ def __init__(self):
+ self.initialized = False
+ self.modules = {}
+ self.init_handlers = []
+ self.handlers = []
+
+ def __getitem__(self, key):
+ return self.modules[key]
+
+ def __contains__(self, key):
+ return key in self.modules
+
+ def __getattr__(self, key):
+ return self.modules.get(key)
+
+ @staticmethod
+ def get_custom_node(name: str) -> ModuleType | None:
+ module_key = f"custom_nodes.{name}"
+ with contextlib.suppress(StopIteration):
+ spec = importlib.util.find_spec(module_key)
+ if spec is None:
+ return None
+ return next(
+ v
+ for v in sys.modules.copy().values()
+ if hasattr(v, "__spec__")
+ and v.__spec__ is not None
+ and v.__spec__.origin == spec.origin
+ )
+ return None
+
+ def register_init_handler(self, handler):
+ self.init_handlers.append(handler)
+
+ def register_integration(self, key: str, module_name: str, handler=None) -> None:
+ if self.initialized:
+ raise ValueError(
+ "Internal error: Cannot register integration after initialization",
+ )
+ if any(item[0] == key or item[1] == module_name for item in self.handlers):
+ errstr = (
+ f"Module {module_name} ({key}) already in integration handlers list!"
+ )
+ raise ValueError(errstr)
+ self.handlers.append(self.Integration(key, module_name, handler))
+
+ def initialize(self) -> None:
+ if self.initialized:
+ return
+ self.initialized = True
+ for ih in self.handlers:
+ module = self.get_custom_node(ih.module_name)
+ if module is None:
+ continue
+ if ih.handler is not None:
+ module = ih.handler(module)
+ if module is not None:
+ self.modules[ih.key] = module
+
+ for init_handler in self.init_handlers:
+ init_handler(self)
+
+
+class SonarIntegrations(Integrations):
+ def __init__(self, *args: list, **kwargs: dict):
+ super().__init__(*args, **kwargs)
+ self.register_integration("bleh", "ComfyUI-bleh", self.bleh_integration)
+ self.register_integration(
+ "restart",
+ "ComfyUI_restart_sampling",
+ self.restart_integration,
+ )
+
+ @classmethod
+ def bleh_integration(cls, module: ModuleType) -> ModuleType | None:
+ bleh_version = getattr(module, "BLEH_VERSION", -1)
+ if bleh_version < 1:
+ return None
+ return module
+
+ @classmethod
+ def restart_integration(cls, module: ModuleType) -> ModuleType | None:
+ if hasattr(module, "restart_sampling") and hasattr(
+ module.restart_sampling,
+ "DEFAULT_SEGMENTS",
+ ):
+ return module
+ return None
+
+
+MODULES = SonarIntegrations()
+
+
+class IntegratedNode(type):
+ @staticmethod
+ def wrap_INPUT_TYPES(orig_method: Callable, *args: list, **kwargs: dict) -> dict:
+ MODULES.initialize()
+ return orig_method(*args, **kwargs)
+
+ def __new__(cls: type, name: str, bases: tuple, attrs: dict) -> object:
+ obj = type.__new__(cls, name, bases, attrs)
+ if hasattr(obj, "INPUT_TYPES"):
+ obj.INPUT_TYPES = partial(cls.wrap_INPUT_TYPES, obj.INPUT_TYPES)
+ return obj
- 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
- MODULES["restart"] = rs
__all__ = ("MODULES",)
diff --git a/py/freeu_extreme.py b/py/freeu_extreme.py
index d9b401d..d94ea40 100644
--- a/py/freeu_extreme.py
+++ b/py/freeu_extreme.py
@@ -2,7 +2,8 @@ from __future__ import annotations
import torch
-from .external import MODULES as EXTERNAL_MODULES
+from . import utils
+from .external import IntegratedNode
from .powernoise import PowerFilter
@@ -28,14 +29,7 @@ def ffilter(x, pfilter, normalization_factor=1.0, cfg_idx=None, filter_cache=Non
return x_filt.to(x.dtype, non_blocking=True)
-BLEND_OPS = (
- {"lerp": torch.lerp}
- if "bleh" not in EXTERNAL_MODULES
- else EXTERNAL_MODULES["bleh"].py.latent_utils.BLENDING_MODES
-)
-
-
-class FreeUExtremeConfigNode:
+class FreeUExtremeConfigNode(metaclass=IntegratedNode):
DESCRIPTION = "Allows setting configuration for FreeU Extreme."
RETURN_TYPES = ("FRUX_CONFIG",)
FUNCTION = "go"
@@ -150,7 +144,7 @@ class FreeUExtremeConfigNode:
},
),
"blend_mode": (
- tuple(BLEND_OPS.keys()),
+ tuple(utils.BLENDING_MODES.keys()),
{
"tooltip": "Mode used when blending. Generally only has an effect when blend is set to values other than 0 or 1",
},
@@ -302,7 +296,7 @@ class FreeUExtremeConfig:
x[:, slice_offs : slice_offs + slice_size] = (
xslice
if self.blend == 1.0
- else BLEND_OPS[self.blend_mode](
+ else utils.BLENDING_MODES[self.blend_mode](
x[:, slice_offs : slice_offs + slice_size],
xslice,
self.blend,
@@ -331,12 +325,12 @@ class FreeUExtremeConfig:
def clone(self):
return self.__class__(**{k: getattr(self, k) for k in self._keys})
- def __repr__(self): # noqa: D105
+ def __repr__(self):
meh = {k: getattr(self, k) for k in self._keys}
return f""
-class FreeUExtremeNode:
+class FreeUExtremeNode(metaclass=IntegratedNode):
DESCRIPTION = "Main FreeU Extreme node. Allows patching a model with the FreeU (V2) effect with more control."
RETURN_TYPES = ("MODEL",)
FUNCTION = "go"
diff --git a/py/nodes.py b/py/nodes.py
index bcea240..a682ba2 100644
--- a/py/nodes.py
+++ b/py/nodes.py
@@ -2,18 +2,18 @@ from __future__ import annotations
import abc
import inspect
+import math
import random
-from types import SimpleNamespace
from typing import Any, Callable
import numpy as np
import torch
import yaml
-from comfy import samplers
+from comfy import model_management, samplers
-from . import external, noise
+from . import external, noise, utils
+from .external import IntegratedNode
from .noise import NoiseType
-from .noise_generation import scale_noise
from .sonar import (
GuidanceConfig,
GuidanceType,
@@ -52,7 +52,7 @@ if not HAVE_COMFY_UNION_TYPE:
result.whitelist = whitelist
return result
- def __ne__(self, other): # noqa: D105
+ def __ne__(self, other):
return False if self.whitelist is None else other not in self.whitelist
WILDCARD_NOISE = Wildcard("*", whitelist=NOISE_INPUT_TYPES)
@@ -65,24 +65,7 @@ NOISE_INPUT_TYPES_HINT = (
)
-if "bleh" in external.MODULES:
- bleh_latent_utils = external.MODULES["bleh"].py.latent_utils
- BLEND_MODES = bleh_latent_utils.BLENDING_MODES
- UPSCALE_METHODS = bleh_latent_utils.UPSCALE_METHODS
- del bleh_latent_utils
-else:
- BLEND_MODES = {"lerp": torch.lerp}
- UPSCALE_METHODS = (
- "bilinear",
- "nearest-exact",
- "nearest",
- "area",
- "bicubic",
- "bislerp",
- )
-
-
-class NoisyLatentLikeNode:
+class NoisyLatentLikeNode(metaclass=IntegratedNode):
DESCRIPTION = "Allows generating noise (and optionally adding it) based on a reference latent. Note: For img2img workflows, you will generally want to enable add_to_latent as well as connecting the model and sigmas inputs."
RETURN_TYPES = ("LATENT",)
OUTPUT_TOOLTIPS = ("The noisy latent image.",)
@@ -120,6 +103,8 @@ class NoisyLatentLikeNode:
"FLOAT",
{
"default": 1.0,
+ "step": 0.001,
+ "round": False,
"tooltip": "Multiplier for the strength of the generated noise. Performed after mul_by_sigmas_opt.",
},
),
@@ -175,7 +160,7 @@ class NoisyLatentLikeNode:
}
@classmethod
- def go(
+ def go( # noqa: PLR0914
cls,
*,
noise_type: str,
@@ -199,9 +184,11 @@ class NoisyLatentLikeNode:
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,
+ model_sigma_max = float(model.model_sampling.sigma_max)
+ first_sigma = float(sigmas[0])
+ max_denoise = (
+ math.isclose(model_sigma_max, first_sigma, rel_tol=1e-05)
+ or first_sigma > model_sigma_max
)
multiplier *= (
float(
@@ -210,11 +197,17 @@ class NoisyLatentLikeNode:
/ latent_scale_factor
)
if sigmas is not None and sigmas.numel() > 1:
- sigma_min, sigma_max = sigmas[0], sigmas[-1]
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
sigma, sigma_next = sigmas[0], sigmas[1]
else:
sigma_min, sigma_max, sigma, sigma_next = (None,) * 4
latent_samples = latent["samples"]
+ orig_device = latent_samples.device
+ want_device = (
+ torch.device("cpu") if cpu_noise else model_management.get_torch_device()
+ )
+ if latent_samples.device != want_device:
+ latent_samples = latent_samples.detach().clone().to(want_device)
if custom_noise_opt is not None:
ns = custom_noise_opt.make_noise_sampler(
latent_samples,
@@ -243,15 +236,16 @@ class NoisyLatentLikeNode:
)
finally:
torch.random.set_rng_state(randst)
- result = scale_noise(result, multiplier, normalized=True)
+ result = utils.scale_noise(result, multiplier, normalized=True)
if add_to_latent:
result += latent_samples.repeat(
*(repeat_batch if i == 0 else 1 for i in range(latent_samples.ndim)),
).to(result)
+ result = result.to(orig_device)
return ({"samples": result},)
-class SonarCustomNoiseNodeBase(abc.ABC):
+class SonarCustomNoiseNodeBase(metaclass=IntegratedNode):
DESCRIPTION = "A custom noise item."
RETURN_TYPES = ("SONAR_CUSTOM_NOISE",)
OUTPUT_TOOLTIPS = ("A custom noise chain.",)
@@ -341,6 +335,26 @@ class SonarCustomNoiseNode(SonarCustomNoiseNodeBase):
return noise.CustomNoiseItem
+class SonarCustomNoiseAdvNode(SonarCustomNoiseNode):
+ DESCRIPTION = "A custom noise item allowing advanced YAML parameter input."
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ result = super().INPUT_TYPES()
+ result["optional"] |= {
+ "yaml_parameters": (
+ "STRING",
+ {
+ "tooltip": "Allows specifying custom parameters via YAML. Note: When specifying paramaters this way, there is no error checking.",
+ "placeholder": "# YAML or JSON here",
+ "dynamicPrompts": False,
+ "multiline": True,
+ },
+ ),
+ }
+ return result
+
+
class SonarNormalizeNoiseNodeMixin:
@staticmethod
def get_normalize(val: str) -> bool | None:
@@ -386,6 +400,8 @@ class SonarModulatedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeM
"default": 2.0,
"min": -100.0,
"max": 100.0,
+ "step": 0.001,
+ "round": False,
"tooltip": "Controls the strength of the modulation effect.",
},
),
@@ -515,7 +531,7 @@ class SonarRepeatedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMi
class SonarScheduledNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin):
- DESCRIPTION = "Custom noise type that allows scheduling the output of other custom noise generators. NOTE: If you don't connect the fallback custom noise input, no noise will be generated outside of the start_percent, end_percent range. Recommend connecting a 1.0 strength Gaussian custom noise node as the fallback."
+ DESCRIPTION = "Custom noise type that allows scheduling the output of other custom noise generators. NOTE: If you don't connect the fallback custom noise input, no noise will be generated outside of the start_percent, end_percent range. I recommend connecting a 1.0 strength Gaussian custom noise node as the fallback."
@classmethod
def INPUT_TYPES(cls):
@@ -539,6 +555,8 @@ class SonarScheduledNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeM
"default": 0.0,
"min": 0.0,
"max": 1.0,
+ "step": 0.001,
+ "round": False,
"tooltip": "Time the custom noise becomes active. Note: Sampling percentage where 1.0 indicates 100%, not based on steps.",
},
),
@@ -548,6 +566,8 @@ class SonarScheduledNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeM
"default": 1.0,
"min": 0.0,
"max": 1.0,
+ "step": 0.001,
+ "round": False,
"tooltip": "Time the custom noise effect ends - inclusive, so only sampling percentages greater than this will be excluded. Note: Sampling percentage where 1.0 indicates 100%, not based on steps.",
},
),
@@ -744,7 +764,7 @@ class SonarGuidedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixi
):
return super().go(
factor,
- ref_latent=scale_noise(
+ ref_latent=utils.scale_noise(
SonarGuidanceMixin.prepare_ref_latent(latent["samples"].clone()),
normalized=normalize_ref,
),
@@ -868,11 +888,13 @@ class SonarBlendedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMix
"FLOAT",
{
"default": 0.5,
+ "step": 0.001,
+ "round": False,
"tooltip": "Blend strength for custom_noise_2. Note that if set to 0 then custom_noise_2 is optional (and will not be called to generate noise) and if set to 1 then custom_noise_1 will not be called to generate noise. This is worth mentioning since going from a strength of 0.000000001 to 0 could make a big difference.",
},
),
"blend_mode": (
- tuple(BLEND_MODES.keys()),
+ tuple(utils.BLENDING_MODES.keys()),
{
"default": "lerp",
"tooltip": "Mode used for blending the two noise types. More modes will be available if ComfyUI-bleh is installed.",
@@ -917,7 +939,7 @@ class SonarBlendedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMix
custom_noise_2=None,
blend_mode="lerp",
):
- blend_function = BLEND_MODES.get(blend_mode)
+ blend_function = utils.BLENDING_MODES.get(blend_mode)
if blend_function is None:
raise ValueError("Unknown blend mode")
return super().go(
@@ -932,6 +954,143 @@ class SonarBlendedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMix
)
+class SonarResizedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin):
+ DESCRIPTION = "Custom noise type that allows resizing another noise item."
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
+ result["required"] |= {
+ "width": (
+ "INT",
+ {
+ "default": 1152,
+ "min": 16,
+ "step": 8,
+ "tooltip": "Note: This should almost always be set to a higher value than the image you're actually sampling.",
+ },
+ ),
+ "height": (
+ "INT",
+ {
+ "default": 1152,
+ "min": 16,
+ "step": 8,
+ "tooltip": "Note: This should almost always be set to a higher value than the image you're actually sampling.",
+ },
+ ),
+ "downscale_strategy": (
+ ("crop", "scale"),
+ {
+ "default": "crop",
+ "tooltip": "Scaling noise is something you'd pretty much only use to create weird effects. For normal workflows, leave this on crop.",
+ },
+ ),
+ "initial_reference": (
+ ("prefer_crop", "prefer_scale"),
+ {
+ "default": "prefer_crop",
+ "tooltip": "The initial latent the noise sampler uses as a reference may not match the requested width/height. This setting controls whether to crop or scale. Note: Cropping can only occur when the initial reference is larger than width/height in both dimensions which is unlikely (and not recommended).",
+ },
+ ),
+ "crop_mode": (
+ (
+ "center",
+ "top_left",
+ "top_center",
+ "top_right",
+ "center_left",
+ "center_right",
+ "bottom_left",
+ "bottom_center",
+ "bottom_right",
+ ),
+ {
+ "default": "center",
+ "tooltip": "Note: Crops will have a bias toward the lower number when the size isn't divisible by two. For example, a center crop of size 3 from (0, 1, 2, 3, 4, 5) will result in (1, 2, 3).",
+ },
+ ),
+ "crop_offset_horizontal": (
+ "INT",
+ {
+ "default": 0,
+ "step": 8,
+ "tooltip": "This offsets the cropped view by the specified size. Positive values will move it toward the right, negative values will move it toward the left. The offsets will be adjusted to to fit in the available space. For example, if you have crop_mode set to top_right then setting a positive offset isn't going to do anything: it's already as far right as it can go.",
+ },
+ ),
+ "crop_offset_vertical": (
+ "INT",
+ {
+ "default": 0,
+ "step": 8,
+ "tooltip": "This offsets the cropped view by the specified size. Positive values will move it toward the bottom, negative values will move it toward the top. The offsets will be adjusted to to fit in the available space. For example, if you have crop_mode set to bottom_right then setting a positive offset isn't going to do anything: it's already as far down as it can go.",
+ },
+ ),
+ "upscale_mode": (
+ utils.UPSCALE_METHODS,
+ {
+ "tooltip": "Allows setting the scaling mode when width/height is smaller than the requested size.",
+ "default": "nearest-exact",
+ },
+ ),
+ "downscale_mode": (
+ utils.UPSCALE_METHODS,
+ {
+ "tooltip": "Allows setting the scaling mode when width/height is larger than the requested size and downscale_strategy is set to 'scale'.",
+ "default": "nearest-exact",
+ },
+ ),
+ "normalize": (
+ ("default", "forced", "disabled"),
+ {
+ "tooltip": "Controls whether the generated noise is normalized to 1.0 strength. For weird blend modes, you may want to set this to forced.",
+ },
+ ),
+ "custom_noise": (
+ WILDCARD_NOISE,
+ {
+ "tooltip": f"Custom noise.\n{NOISE_INPUT_TYPES_HINT}",
+ },
+ ),
+ }
+ return result
+
+ @classmethod
+ def get_item_class(cls):
+ return noise.ResizedNoise
+
+ def go(
+ self,
+ *,
+ factor,
+ width,
+ height,
+ downscale_strategy,
+ initial_reference,
+ crop_offset_horizontal,
+ crop_offset_vertical,
+ crop_mode,
+ upscale_mode,
+ downscale_mode,
+ normalize,
+ custom_noise,
+ ):
+ return super().go(
+ factor,
+ width=width,
+ height=height,
+ downscale_strategy=downscale_strategy,
+ initial_reference=initial_reference,
+ crop_offset_horizontal=crop_offset_horizontal,
+ crop_offset_vertical=crop_offset_vertical,
+ crop_mode=crop_mode,
+ upscale_mode=upscale_mode,
+ downscale_mode=downscale_mode,
+ normalize=normalize,
+ custom_noise=custom_noise,
+ )
+
+
class SonarAdvancedPyramidNoiseNode(SonarCustomNoiseNodeBase):
DESCRIPTION = (
"Custom noise type that allows specifying parameters for Pyramid variants."
@@ -965,11 +1124,13 @@ class SonarAdvancedPyramidNoiseNode(SonarCustomNoiseNodeBase):
"FLOAT",
{
"default": 0.0,
+ "step": 0.001,
+ "round": False,
"tooltip": "When set to 0 will use the variant default.",
},
),
"upscale_mode": (
- ("default", *UPSCALE_METHODS),
+ ("default", *utils.UPSCALE_METHODS),
{
"tooltip": "Allows setting the scaling mode for Pyramid noise. Leave on default to use the variant default.",
"default": "default",
@@ -1015,6 +1176,8 @@ class SonarAdvanced1fNoiseNode(SonarCustomNoiseNodeBase):
"FLOAT",
{
"default": 0.25,
+ "step": 0.001,
+ "round": False,
"tooltip": "Similar to the advanced power noise node, positive values increase low frequencies (with colorful effects), negative values increase high frequencies.",
},
),
@@ -1022,6 +1185,8 @@ class SonarAdvanced1fNoiseNode(SonarCustomNoiseNodeBase):
"FLOAT",
{
"default": 1.0,
+ "step": 0.001,
+ "round": False,
"tooltip": "Currently no description of exactly what it does, it's just another knob you can try turning for a different effect.",
},
),
@@ -1029,6 +1194,8 @@ class SonarAdvanced1fNoiseNode(SonarCustomNoiseNodeBase):
"FLOAT",
{
"default": 1.0,
+ "step": 0.001,
+ "round": False,
"tooltip": "Vertical frequency scaling factor.",
},
),
@@ -1036,6 +1203,8 @@ class SonarAdvanced1fNoiseNode(SonarCustomNoiseNodeBase):
"FLOAT",
{
"default": 1.0,
+ "step": 0.001,
+ "round": False,
"tooltip": "Horizontal frequency scaling factor.",
},
),
@@ -1088,6 +1257,8 @@ class SonarAdvancedPowerLawNoiseNode(SonarCustomNoiseNodeBase):
"FLOAT",
{
"default": 0.5,
+ "step": 0.001,
+ "round": False,
"tooltip": "Alpha parameter of the generated noise. Positive values (low frequency noise) tend to produce colorful results.",
},
),
@@ -1161,6 +1332,173 @@ class SonarAdvancedPowerLawNoiseNode(SonarCustomNoiseNodeBase):
)
+class SonarAdvancedDistroNoiseNode(SonarCustomNoiseNodeBase):
+ DESCRIPTION = "Custom noise type that allows specifying parameters for Distro variants. See: https://pytorch.org/docs/stable/distributions.html"
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ngcls = cls.get_item_class().ns_factory
+ distro_params = ngcls.distro_params
+ variants = tuple(sorted(distro_params.keys()))
+ combined_params = ngcls.build_params()
+
+ result = super().INPUT_TYPES()
+ result["required"] |= {
+ "distribution": (
+ variants,
+ {
+ "tooltip": "Sets the distribution used for noise generation. See: https://pytorch.org/docs/stable/distributions.html",
+ "default": "uniform",
+ },
+ ),
+ "quantile_norm": (
+ "FLOAT",
+ {
+ "default": 0.85,
+ "min": 0.0,
+ "max": 1.0,
+ "step": 0.001,
+ "round": False,
+ "tooltip": "When enabled, will normalize generated noise to this quantile (i.e. 0.75 means outliers >75% will be clipped). Set to 1.0 or 0.0 to disable quantile normalization. A value like 0.75 or 0.85 should be reasonable, it really depends on the distribution and how many of the values are extreme.",
+ },
+ ),
+ "quantile_norm_mode": (
+ (
+ "global",
+ "batch",
+ "channel",
+ "batch_row",
+ "batch_col",
+ "nonflat_row",
+ "nonflat_col",
+ ),
+ {
+ "default": "batch",
+ "tooltip": "Controls what dimensions quantile normalization uses. By default, the noise is flattened first. You can try the nonflat versions but they may have a very strong row/column influence. Only applies when quantile_norm is active.",
+ },
+ ),
+ "result_index": (
+ "STRING",
+ {
+ "default": "-1",
+ "tooltip": "When noise generation returns a batch of items, it will select the specified index. Negative indexes count from the end. Values outside the valid range will be automatically adjusted. You may enter a space-separated list of values for the case where there might be multiple added batch dimensions. Excess batch dimensions are removed from the end, indexe from result_index are used in order so you may want to enter the indexes in reverse order.\nExample: If your noise has shape (1, 4, 3, 3) and two 2-sized batch dims are added resulting in (1, 4, 3, 3, 2, 2) and you wanted index 0 from the first additional batch dimension and 1 from the second you would use result_index: 1 0",
+ },
+ ),
+ } | {
+ k: ("STRING" if isinstance(v["default"], str) else v.get("_ty", "FLOAT"), v)
+ for k, v in combined_params.items()
+ }
+ # print("RESULT:", result)
+ return result
+
+ @classmethod
+ def get_item_class(cls):
+ return noise.AdvancedDistroNoise
+
+ def go(
+ self,
+ *,
+ factor,
+ rescale,
+ distribution,
+ quantile_norm,
+ quantile_norm_mode,
+ result_index,
+ sonar_custom_noise_opt=None,
+ **kwargs: dict[str],
+ ):
+ normdim, normflat = {
+ "global": (None, True),
+ "batch": (0, True),
+ "channel": (1, True),
+ "batch_row": (2, True),
+ "batch_col": (3, True),
+ "nonflat_row": (2, False),
+ "nonflat_col": (3, False),
+ }.get(quantile_norm_mode, (1, True))
+ result_index = tuple(int(v) for v in result_index.split(None))
+ return super().go(
+ factor,
+ rescale=rescale,
+ sonar_custom_noise_opt=sonar_custom_noise_opt,
+ distro=distribution,
+ quantile_norm=quantile_norm,
+ quantile_norm_dim=normdim,
+ quantile_norm_flatten=normflat,
+ result_index=result_index,
+ **kwargs,
+ )
+
+
+class SonarWaveletFilteredNoiseNode(
+ SonarCustomNoiseNodeBase,
+ SonarNormalizeNoiseNodeMixin,
+):
+ DESCRIPTION = "Custom noise type that allows filtering another custom noise source with wavelets."
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ result = super().INPUT_TYPES()
+ result["required"] |= {
+ "normalize_noise": (
+ "BOOLEAN",
+ {
+ "default": False,
+ "tooltip": "Controls whether the noise source is normalized before wavelet filtering occurs.",
+ },
+ ),
+ "normalize": (
+ ("default", "forced", "disabled"),
+ {
+ "tooltip": "Controls whether the generated noise is normalized to 1.0 strength. For weird blend modes, you may want to set this to forced.",
+ },
+ ),
+ }
+ result["optional"] |= {
+ "custom_noise": (
+ WILDCARD_NOISE,
+ {
+ "tooltip": f"Optional: Custom noise input. If unconnected will default to Gaussian noise.\n{NOISE_INPUT_TYPES_HINT}",
+ },
+ ),
+ "yaml_parameters": (
+ "STRING",
+ {
+ "tooltip": "Allows specifying custom parameters via YAML. Note: When specifying paramaters this way, there is no error checking.",
+ "placeholder": "# YAML or JSON here",
+ "dynamicPrompts": False,
+ "multiline": True,
+ },
+ ),
+ }
+ return result
+
+ @classmethod
+ def get_item_class(cls):
+ return noise.WaveletFilteredNoise
+
+ def go(
+ self,
+ *,
+ factor,
+ rescale,
+ normalize,
+ normalize_noise,
+ custom_noise=None,
+ yaml_parameters=None,
+ sonar_custom_noise_opt=None,
+ ):
+ return super().go(
+ factor,
+ rescale=rescale,
+ sonar_custom_noise_opt=sonar_custom_noise_opt,
+ normalize=self.get_normalize(normalize),
+ normalize_noise=normalize_noise,
+ noise=custom_noise,
+ yaml_parameters=yaml_parameters,
+ )
+
+
class CustomNOISE:
def __init__(
self,
@@ -1223,7 +1561,7 @@ class CustomNOISE:
return torch.cat(tuple(result[i] for i in inverse_inds), axis=0)
-class SonarToComfyNOISENode:
+class SonarToComfyNOISENode(metaclass=IntegratedNode):
DESCRIPTION = "Allows converting SONAR_CUSTOM_NOISE to NOISE (used by SamplerCustomAdvanced and possibly other custom samplers). NOTE: Does not work with noise types that depend on sigma (Brownian, ScheduledNoise, etc)."
RETURN_TYPES = ("NOISE",)
CATEGORY = "sampling/custom_sampling/noise"
@@ -1359,7 +1697,7 @@ class GuidanceConfigNode:
)
-class SamplerNodeSonarBase:
+class SamplerNodeSonarBase(metaclass=IntegratedNode):
DESCRIPTION = "Sonar - momentum based sampler node."
@classmethod
@@ -1374,7 +1712,7 @@ class SamplerNodeSonarBase:
"max": 2.5,
"step": 0.01,
"round": False,
- "tooltip": "Strength of the output from normal sampling. When set to 1.0 effectively disables momentum.",
+ "tooltip": "How much of the normal result to keep during sampling. 0.95 means 95% normal, 5% from history. When set to 1.0 effectively disables momentum.",
},
),
"momentum_hist": (
@@ -1385,7 +1723,7 @@ class SamplerNodeSonarBase:
"max": 1.5,
"step": 0.01,
"round": False,
- "tooltip": "Strength of momentum history",
+ "tooltip": "How much of the existing history to leave at each update. 0.75 means keep 75%, mix in 25% of the new result.",
},
),
"momentum_init": (
@@ -1427,6 +1765,34 @@ class SamplerNodeSonarBase:
class SamplerNodeSonarEuler(SamplerNodeSonarBase):
+ RETURN_TYPES = ("SAMPLER",)
+ CATEGORY = "sampling/custom_sampling/samplers"
+
+ FUNCTION = "get_sampler"
+
+ @classmethod
+ def get_sampler(
+ cls,
+ *,
+ momentum,
+ momentum_hist,
+ momentum_init,
+ direction,
+ rand_init_noise_type,
+ guidance_cfg_opt=None,
+ ):
+ cfg = SonarConfig(
+ momentum=momentum,
+ init=HistoryType[momentum_init.upper()],
+ momentum_hist=momentum_hist,
+ direction=direction,
+ rand_init_noise_type=NoiseType[rand_init_noise_type.upper()],
+ guidance=guidance_cfg_opt,
+ )
+ return (samplers.KSAMPLER(SonarEuler.sampler, {"sonar_config": cfg}),)
+
+
+class SamplerNodeSonarEulerAncestral(SamplerNodeSonarEuler):
@classmethod
def INPUT_TYPES(cls):
result = super().INPUT_TYPES()
@@ -1443,52 +1809,6 @@ class SamplerNodeSonarEuler(SamplerNodeSonarBase):
"tooltip": "Multiplier for noise added during ancestral or SDE sampling.",
},
),
- },
- )
- return result
-
- RETURN_TYPES = ("SAMPLER",)
- CATEGORY = "sampling/custom_sampling/samplers"
-
- FUNCTION = "get_sampler"
-
- @classmethod
- def get_sampler(
- cls,
- *,
- momentum,
- momentum_hist,
- momentum_init,
- direction,
- rand_init_noise_type,
- s_noise,
- guidance_cfg_opt=None,
- ):
- cfg = SonarConfig(
- momentum=momentum,
- init=HistoryType[momentum_init.upper()],
- momentum_hist=momentum_hist,
- direction=direction,
- rand_init_noise_type=NoiseType[rand_init_noise_type.upper()],
- guidance=guidance_cfg_opt,
- )
- return (
- samplers.KSAMPLER(
- SonarEuler.sampler,
- {
- "s_noise": s_noise,
- "sonar_config": cfg,
- },
- ),
- )
-
-
-class SamplerNodeSonarEulerAncestral(SamplerNodeSonarEuler):
- @classmethod
- def INPUT_TYPES(cls):
- result = super().INPUT_TYPES()
- result["required"].update(
- {
"eta": (
"FLOAT",
{
@@ -1563,6 +1883,17 @@ class SamplerNodeSonarDPMPPSDE(SamplerNodeSonarEuler):
result = super().INPUT_TYPES()
result["required"].update(
{
+ "s_noise": (
+ "FLOAT",
+ {
+ "default": 1.0,
+ "min": 0.0,
+ "max": 100.0,
+ "step": 0.01,
+ "round": False,
+ "tooltip": "Multiplier for noise added during ancestral or SDE sampling.",
+ },
+ ),
"eta": (
"FLOAT",
{
@@ -1631,9 +1962,8 @@ class SamplerNodeSonarDPMPPSDE(SamplerNodeSonarEuler):
)
-class SamplerNodeConfigOverride:
+class SamplerNodeConfigOverride(metaclass=IntegratedNode):
DESCRIPTION = "Allows overriding paramaters for a SAMPLER. Only parameters that particular sampler supports will be applied, so for example setting ETA will have no effect for non-ancestral Euler."
- # KWARG_OVERRIDES = ("s_noise", "eta", "s_churn", "r", "solver_type")
@classmethod
def INPUT_TYPES(cls):
@@ -1845,6 +2175,8 @@ class SamplerNodeConfigOverride:
)
+NODE_DISPLAY_NAME_MAPPINGS = {}
+
NODE_CLASS_MAPPINGS = {
"SamplerSonarEuler": SamplerNodeSonarEuler,
"SamplerSonarEulerA": SamplerNodeSonarEulerAncestral,
@@ -1855,7 +2187,9 @@ NODE_CLASS_MAPPINGS = {
"SonarAdvancedPyramidNoise": SonarAdvancedPyramidNoiseNode,
"SonarAdvanced1fNoise": SonarAdvanced1fNoiseNode,
"SonarAdvancedPowerLawNoise": SonarAdvancedPowerLawNoiseNode,
+ "SonarAdvancedDistroNoise": SonarAdvancedDistroNoiseNode,
"SonarCustomNoise": SonarCustomNoiseNode,
+ "SonarCustomNoiseAdv": SonarCustomNoiseAdvNode,
"SonarCompositeNoise": SonarCompositeNoiseNode,
"SonarModulatedNoise": SonarModulatedNoiseNode,
"SonarRepeatedNoise": SonarRepeatedNoiseNode,
@@ -1864,232 +2198,278 @@ NODE_CLASS_MAPPINGS = {
"SonarRandomNoise": SonarRandomNoiseNode,
"SonarChannelNoise": SonarChannelNoiseNode,
"SonarBlendedNoise": SonarBlendedNoiseNode,
+ "SonarResizedNoise": SonarResizedNoiseNode,
+ "SonarWaveletFilteredNoise": SonarWaveletFilteredNoiseNode,
"SONAR_CUSTOM_NOISE to NOISE": SonarToComfyNOISENode,
}
-NODE_DISPLAY_NAME_MAPPINGS = {}
+
+bleh = None
-if "bleh" in external.MODULES:
- import ast
+class SonarBlendFilterNoiseNode(
+ SonarCustomNoiseNodeBase,
+ SonarNormalizeNoiseNodeMixin,
+):
+ DESCRIPTION = "Custom noise type that allows blending and filtering the output of another noise generator using ComfyUI-bleh."
- bleh = external.MODULES["bleh"]
- bleh_latentutils = bleh.py.latent_utils
- bleh_ops = bleh.py.nodes.ops
+ @classmethod
+ def INPUT_TYPES(cls):
+ result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
+ result["required"] |= {
+ "sonar_custom_noise": (
+ WILDCARD_NOISE,
+ {
+ "tooltip": f"Custom noise input.\n{NOISE_INPUT_TYPES_HINT}",
+ },
+ ),
+ "blend_mode": (("simple_add", *utils.BLENDING_MODES.keys()),),
+ "ffilter": (tuple(bleh.py.latent_utils.FILTER_PRESETS.keys()),),
+ "ffilter_custom": ("STRING", {"default": ""}),
+ "ffilter_scale": (
+ "FLOAT",
+ {
+ "default": 1.0,
+ "min": -100.0,
+ "max": 100.0,
+ "step": 0.001,
+ "round": False,
+ },
+ ),
+ "ffilter_strength": (
+ "FLOAT",
+ {
+ "default": 0.0,
+ "min": -100.0,
+ "max": 100.0,
+ "step": 0.001,
+ "round": False,
+ },
+ ),
+ "ffilter_threshold": (
+ "INT",
+ {"default": 1, "min": 1, "max": 32},
+ ),
+ "enhance_mode": (("none", *bleh.py.latent_utils.ENHANCE_METHODS),),
+ "enhance_strength": (
+ "FLOAT",
+ {
+ "default": 0.0,
+ "min": -100.0,
+ "max": 100.0,
+ "step": 0.001,
+ "round": False,
+ },
+ ),
+ "affect": (("result", "noise", "both"),),
+ "normalize_result": (("default", "forced", "disabled"),),
+ "normalize_noise": (("default", "forced", "disabled"),),
+ }
+ return result
- class SonarBlendFilterNoiseNode(
- SonarCustomNoiseNodeBase,
- SonarNormalizeNoiseNodeMixin,
+ @classmethod
+ def get_item_class(cls):
+ return noise.BlendFilterNoise
+
+ def go(
+ self,
+ *,
+ factor,
+ sonar_custom_noise,
+ blend_mode,
+ ffilter,
+ ffilter_custom,
+ ffilter_scale,
+ ffilter_strength,
+ ffilter_threshold,
+ enhance_mode,
+ enhance_strength,
+ affect,
+ normalize_result,
+ normalize_noise,
):
- DESCRIPTION = "Custom noise type that allows blending and filtering the output of another noise generator using ComfyUI-bleh."
+ if bleh is None:
+ raise RuntimeError("bleh not available")
+ import ast # noqa: PLC0415
- @classmethod
- def INPUT_TYPES(cls):
- result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
- result["required"] |= {
- "sonar_custom_noise": (
- WILDCARD_NOISE,
- {
- "tooltip": f"Custom noise input.\n{NOISE_INPUT_TYPES_HINT}",
- },
- ),
- "blend_mode": (
- ("simple_add", *bleh_latentutils.BLENDING_MODES.keys()),
- ),
- "ffilter": (tuple(bleh_latentutils.FILTER_PRESETS.keys()),),
- "ffilter_custom": ("STRING", {"default": ""}),
- "ffilter_scale": (
- "FLOAT",
- {"default": 1.0, "min": -100.0, "max": 100.0},
- ),
- "ffilter_strength": (
- "FLOAT",
- {"default": 0.0, "min": -100.0, "max": 100.0},
- ),
- "ffilter_threshold": (
- "INT",
- {"default": 1, "min": 1, "max": 32},
- ),
- "enhance_mode": (("none", *bleh_latentutils.ENHANCE_METHODS),),
- "enhance_strength": (
- "FLOAT",
- {"default": 0.0, "min": -100.0, "max": 100.0},
- ),
- "affect": (("result", "noise", "both"),),
- "normalize_result": (("default", "forced", "disabled"),),
- "normalize_noise": (("default", "forced", "disabled"),),
- }
- return result
-
- @classmethod
- def get_item_class(cls):
- return noise.BlendFilterNoise
-
- def go(
- self,
- *,
+ ffilter_custom = ffilter_custom.strip()
+ normalize_result = (
+ None if normalize_result == "default" else normalize_result == "forced"
+ )
+ normalize_noise = (
+ None if normalize_noise == "default" else normalize_noise == "forced"
+ )
+ if ffilter_custom:
+ ffilter = ast.literal_eval(f"[{ffilter_custom}]")
+ else:
+ ffilter = bleh.py.latent_utils.FILTER_PRESETS[ffilter]
+ return super().go(
factor,
- sonar_custom_noise,
- blend_mode,
- ffilter,
- ffilter_custom,
- ffilter_scale,
- ffilter_strength,
- ffilter_threshold,
- enhance_mode,
- enhance_strength,
- affect,
- normalize_result,
- normalize_noise,
- ):
- ffilter_custom = ffilter_custom.strip()
- normalize_result = (
- None if normalize_result == "default" else normalize_result == "forced"
- )
- normalize_noise = (
- None if normalize_noise == "default" else normalize_noise == "forced"
- )
- if ffilter_custom:
- ffilter = ast.literal_eval(f"[{ffilter_custom}]")
- else:
- ffilter = bleh_latentutils.FILTER_PRESETS[ffilter]
- return super().go(
- factor,
- noise=sonar_custom_noise.clone(),
- blend_mode=blend_mode,
- ffilter=ffilter,
- ffilter_scale=ffilter_scale,
- ffilter_strength=ffilter_strength,
- ffilter_threshold=ffilter_threshold,
- enhance_mode=enhance_mode,
- enhance_strength=enhance_strength,
- affect=affect,
- normalize_noise=self.get_normalize(normalize_noise),
- normalize_result=self.get_normalize(normalize_result),
- )
-
- class SonarBlehOpsNoiseNode(
- SonarCustomNoiseNodeBase,
- SonarNormalizeNoiseNodeMixin,
- ):
- DESCRIPTION = (
- "Custom noise type that allows manipulating noise with ComfyUI-bleh ops."
+ noise=sonar_custom_noise.clone(),
+ blend_mode=blend_mode,
+ ffilter=ffilter,
+ ffilter_scale=ffilter_scale,
+ ffilter_strength=ffilter_strength,
+ ffilter_threshold=ffilter_threshold,
+ enhance_mode=enhance_mode,
+ enhance_strength=enhance_strength,
+ affect=affect,
+ normalize_noise=self.get_normalize(normalize_noise),
+ normalize_result=self.get_normalize(normalize_result),
)
- @classmethod
- def INPUT_TYPES(cls):
- result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
- result["required"] |= {
- "sonar_custom_noise": (
- WILDCARD_NOISE,
+
+class SonarBlehOpsNoiseNode(
+ SonarCustomNoiseNodeBase,
+ SonarNormalizeNoiseNodeMixin,
+):
+ DESCRIPTION = (
+ "Custom noise type that allows manipulating noise with ComfyUI-bleh ops."
+ )
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
+ result["required"] |= {
+ "sonar_custom_noise": (
+ WILDCARD_NOISE,
+ {
+ "tooltip": f"Custom noise input.\n{NOISE_INPUT_TYPES_HINT}",
+ },
+ ),
+ "normalize": (
+ ("default", "forced", "disabled"),
+ {
+ "tooltip": "Controls whether the generated noise is normalized to 1.0 strength.",
+ },
+ ),
+ "rules": (
+ "STRING",
+ {
+ "tooltip": "Enter rules in the bleh block ops format here.",
+ "placeholder": "# YAML ops here",
+ "dynamicPrompts": False,
+ "multiline": True,
+ },
+ ),
+ }
+ return result
+
+ @classmethod
+ def get_item_class(cls):
+ return noise.BlehOpsNoise
+
+ def go(
+ self,
+ *,
+ factor,
+ sonar_custom_noise,
+ rules,
+ normalize,
+ ):
+ if bleh is None:
+ raise RuntimeError("bleh not available")
+ return super().go(
+ factor,
+ noise=sonar_custom_noise.clone(),
+ rules=bleh.py.nodes.ops.RuleGroup.from_yaml(rules),
+ normalize=normalize,
+ )
+
+
+restart = None
+
+
+class KRestartSamplerCustomNoise(metaclass=IntegratedNode):
+ DESCRIPTION = "Restart sampler variant that allows specifying a custom noise type for noise added by restarts."
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ get_normal_schedulers = getattr(
+ restart.nodes,
+ "get_supported_normal_schedulers",
+ restart.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",
{
- "tooltip": f"Custom noise input.\n{NOISE_INPUT_TYPES_HINT}",
+ "default": 8.0,
+ "min": 0.0,
+ "max": 100.0,
+ "step": 0.001,
+ "round": False,
},
),
- "normalize": (
- ("default", "forced", "disabled"),
- {
- "tooltip": "Controls whether the generated noise is normalized to 1.0 strength.",
- },
- ),
- "rules": (
+ "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",
{
- "tooltip": "Enter rules in the bleh block ops format here.",
- "placeholder": "# YAML ops here",
- "dynamicPrompts": False,
- "multiline": True,
+ "default": restart.restart_sampling.DEFAULT_SEGMENTS,
+ "multiline": False,
},
),
- }
- return result
+ "restart_scheduler": (
+ restart.nodes.get_supported_restart_schedulers(),
+ ),
+ "chunked_mode": ("BOOLEAN", {"default": True}),
+ },
+ "optional": {
+ "custom_noise_opt": (
+ WILDCARD_NOISE,
+ {
+ "tooltip": f"Optional custom noise input.\n{NOISE_INPUT_TYPES_HINT}",
+ },
+ ),
+ },
+ }
- @classmethod
- def get_item_class(cls):
- return noise.BlehOpsNoise
+ RETURN_TYPES = ("LATENT", "LATENT")
+ RETURN_NAMES = ("output", "denoised_output")
+ FUNCTION = "go"
+ CATEGORY = "sampling"
- def go(
- self,
- *,
- factor,
- sonar_custom_noise,
- rules,
- normalize,
- ):
- return super().go(
- factor,
- noise=sonar_custom_noise.clone(),
- rules=bleh_ops.RuleGroup.from_yaml(rules),
- normalize=normalize,
- )
-
- NODE_CLASS_MAPPINGS |= {
- "SonarBlendFilterNoise": SonarBlendFilterNoiseNode,
- "SonarBlehOpsNoise": SonarBlehOpsNoiseNode,
- }
-
-if "restart" in external.MODULES:
- rs = external.MODULES["restart"]
-
- class KRestartSamplerCustomNoise:
- DESCRIPTION = "Restart sampler variant that allows specifying a custom noise type for noise added by restarts."
-
- @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": (
- WILDCARD_NOISE,
- {
- "tooltip": f"Optional custom noise input.\n{NOISE_INPUT_TYPES_HINT}",
- },
- ),
- },
- }
-
- RETURN_TYPES = ("LATENT", "LATENT")
- RETURN_NAMES = ("output", "denoised_output")
- FUNCTION = "sample"
- CATEGORY = "sampling"
-
- @classmethod
- def sample(
- cls,
- *,
+ @classmethod
+ def go(
+ cls,
+ *,
+ 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,
+ ):
+ if restart is None:
+ raise RuntimeError("Restart not available")
+ return restart.restart_sampling.restart_sampling(
model,
- add_noise,
noise_seed,
steps,
cfg,
@@ -2098,78 +2478,74 @@ if "restart" in external.MODULES:
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,
- )
+ 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 hasattr(rs.restart_sampling, "RestartSampler"):
+class RestartSamplerCustomNoise(metaclass=IntegratedNode):
+ DESCRIPTION = "Wrapper used to make another sampler Restart compatible. Allows specifying a custom type for noise added by restarts."
- class RestartSamplerCustomNoise:
- DESCRIPTION = "Wrapper used to make another sampler Restart compatible. Allows specifying a custom type for noise added by restarts."
-
- @classmethod
- def INPUT_TYPES(cls):
- return {
- "required": {
- "sampler": ("SAMPLER",),
- "chunked_mode": ("BOOLEAN", {"default": True}),
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "sampler": ("SAMPLER",),
+ "chunked_mode": ("BOOLEAN", {"default": True}),
+ },
+ "optional": {
+ "custom_noise_opt": (
+ WILDCARD_NOISE,
+ {
+ "tooltip": f"Optional custom noise input.\n{NOISE_INPUT_TYPES_HINT}",
},
- "optional": {
- "custom_noise_opt": (
- WILDCARD_NOISE,
- {
- "tooltip": f"Optional custom noise input.\n{NOISE_INPUT_TYPES_HINT}",
- },
- ),
- },
- }
+ ),
+ },
+ }
- RETURN_TYPES = ("SAMPLER",)
- FUNCTION = "go"
- CATEGORY = "sampling/custom_sampling/samplers"
+ RETURN_TYPES = ("SAMPLER",)
+ FUNCTION = "go"
+ CATEGORY = "sampling/custom_sampling/samplers"
- @classmethod
- def go(cls, 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,)
+ @classmethod
+ def go(cls, sampler, chunked_mode, custom_noise_opt=None):
+ if restart is None or not hasattr(restart.restart_sampling, "RestartSampler"):
+ raise RuntimeError("Restart not available")
+ 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(
+ restart.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
+
+NODE_CLASS_MAPPINGS |= {
+ "KRestartSamplerCustomNoise": KRestartSamplerCustomNoise,
+ "RestartSamplerCustomNoise": RestartSamplerCustomNoise,
+ "SonarBlendFilterNoise": SonarBlendFilterNoiseNode,
+ "SonarBlehOpsNoise": SonarBlehOpsNoiseNode,
+}
+
+
+def init_integrations(integrations):
+ global restart, bleh # noqa: PLW0603
+ restart = integrations.restart
+ bleh = integrations.bleh
+
+
+external.MODULES.register_init_handler(init_integrations)
diff --git a/py/noise.py b/py/noise.py
index 534e748..54a2180 100644
--- a/py/noise.py
+++ b/py/noise.py
@@ -6,18 +6,30 @@ from typing import Callable
import comfy
import torch
+import yaml
from comfy.k_diffusion import sampling
from torch import Tensor
-from . import external
+from . import external, utils
from .noise_generation import *
from .sonar import SonarGuidanceMixin
+from .utils import crop_samples, scale_noise
# ruff: noqa: D412, D413, D417, D212, D407, ANN002, ANN003, FBT001, FBT002, S311
class CustomNoiseItemBase(abc.ABC):
- def __init__(self, factor, **kwargs):
+ def __init__(self, factor, *, yaml_parameters=None, **kwargs):
+ if yaml_parameters:
+ extra_params = yaml.safe_load(yaml_parameters)
+ if extra_params is None:
+ pass
+ elif not isinstance(extra_params, dict):
+ raise ValueError(
+ "CustomNoiseItem: yaml_parameters must either be null or an object",
+ )
+ else:
+ kwargs["ns_kwargs"] = extra_params
self.factor = factor
self.keys = set(kwargs.keys())
for k, v in kwargs.items():
@@ -46,6 +58,7 @@ class CustomNoiseItemBase(abc.ABC):
seed=None,
cpu=True,
normalized=True,
+ **kwargs,
):
raise NotImplementedError
@@ -65,16 +78,25 @@ class CustomNoiseItem(CustomNoiseItemBase):
seed=None,
cpu=True,
normalized=True,
+ **kwargs,
):
+ ns_kwargs = getattr(self, "ns_kwargs", {}).copy()
+ # print("NS KWARGS", ns_kwargs)
+
return get_noise_sampler(
self.noise_type,
x,
sigma_min,
sigma_max,
- seed=seed,
- cpu=cpu,
+ seed=ns_kwargs.pop("seed", seed),
+ cpu=ns_kwargs.pop("cpu", cpu),
factor=self.factor,
- normalized=self.get_normalize("normalize", normalized),
+ normalized=ns_kwargs.pop(
+ "normalized",
+ self.get_normalize("normalize", normalized),
+ ),
+ **ns_kwargs,
+ **kwargs,
)
@@ -155,6 +177,7 @@ class NoiseSampler:
make_noise_sampler: Callable | None = None,
normalized=False,
factor: float = 1.0,
+ **kwargs,
):
self.factor = factor
self.normalized = normalized
@@ -164,16 +187,18 @@ class NoiseSampler:
try:
self.noise_sampler = make_noise_sampler(
x,
- transform(torch.as_tensor(sigma_min))
+ sigma_min=transform(torch.as_tensor(sigma_min))
if sigma_min is not None
else None,
- transform(torch.as_tensor(sigma_max))
+ sigma_max=transform(torch.as_tensor(sigma_max))
if sigma_max is not None
else None,
seed=seed,
cpu=cpu,
+ **kwargs,
)
except TypeError as _exc:
+ print("GOT EXC", _exc)
self.noise_sampler = make_noise_sampler(x)
@classmethod
@@ -216,7 +241,7 @@ class AdvancedNoiseBase(CustomNoiseItemBase):
v = getattr(self, k, None)
if v is not None:
noise_sampler_kwargs[k] = v
- self.sampler_factory = NoiseSampler.simple(
+ self.sampler_factory = NoiseSampler.wrap(
partial(self.ns_factory, **noise_sampler_kwargs),
)
@@ -229,9 +254,9 @@ 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,
+ "pyramid": PyramidNoiseGenerator,
+ "pyramid_old": PyramidOldNoiseGenerator,
+ "highres_pyramid": HighresPyramidNoiseGenerator,
}
@property
@@ -244,15 +269,33 @@ class Advanced1fNoise(AdvancedNoiseBase):
@property
def ns_factory(self):
- return onef_noise_like
+ return OneFNoiseGenerator
class AdvancedPowerLawNoise(AdvancedNoiseBase):
ns_factory_arg_keys = ("alpha", "div_max_dims", "use_sign")
+ @classmethod
@property
- def ns_factory(self):
- return powerlaw_noise_like
+ def ns_factory(cls):
+ return PowerLawNoiseGenerator
+
+
+class AdvancedDistroNoise(AdvancedNoiseBase):
+ distro_params = DistroNoiseGenerator.build_params()
+ ns_factory_arg_keys = (
+ "distro",
+ "quantile_norm",
+ "quantile_norm_dim",
+ "quantile_norm_flatten",
+ "result_index",
+ *distro_params.keys(),
+ )
+
+ @classmethod
+ @property
+ def ns_factory(cls):
+ return DistroNoiseGenerator
class CompositeNoise(CustomNoiseItemBase):
@@ -446,6 +489,10 @@ class ScheduledNoise(CustomNoiseItemBase):
return torch.zeros_like(x)
def noise_sampler(s, sn):
+ if s is None or sn is None:
+ raise ValueError(
+ "ScheduledNoise requires sigma, sigma_next to be passed",
+ )
noise = (ns if end_sigma <= s <= start_sigma else nsa)(s, sn)
return scale_noise(noise, factor, normalized=normalize)
@@ -997,271 +1044,508 @@ class BlendedNoise(CustomNoiseItemBase):
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__(
- self,
+class ResizedNoise(CustomNoiseItemBase):
+ def __init__(
+ self,
+ factor,
+ *,
+ width,
+ height,
+ downscale_strategy,
+ initial_reference,
+ crop_offset_horizontal,
+ crop_offset_vertical,
+ crop_mode,
+ upscale_mode,
+ downscale_mode,
+ normalize,
+ custom_noise,
+ ):
+ if len(custom_noise.items) == 0:
+ raise ValueError("ResizedNoise requires at least one noise item")
+ super().__init__(
factor,
- *,
- noise,
- blend_mode,
- ffilter,
- ffilter_scale,
- ffilter_strength,
- ffilter_threshold,
- enhance_mode,
- enhance_strength,
- affect,
- normalize_result,
- normalize_noise,
- ):
- if len(noise.items) == 0:
- raise ValueError("BlendFilterNoise requires at least one noise item")
- super().__init__(
- factor,
- noise=noise.clone(),
- blend_mode=blend_mode,
- ffilter=ffilter,
- ffilter_scale=ffilter_scale,
- ffilter_strength=ffilter_strength,
- ffilter_threshold=ffilter_threshold,
- enhance_mode=enhance_mode,
- enhance_strength=enhance_strength,
- affect=affect,
- normalize_result=normalize_result,
- normalize_noise=normalize_noise,
- )
+ width=width,
+ height=height,
+ downscale_strategy=downscale_strategy,
+ initial_reference=initial_reference,
+ crop_offset_horizontal=crop_offset_horizontal,
+ crop_offset_vertical=crop_offset_vertical,
+ crop_mode=crop_mode,
+ upscale_mode=upscale_mode,
+ downscale_mode=downscale_mode,
+ custom_noise=custom_noise.clone(),
+ normalize=normalize,
+ )
- def clone_key(self, k):
- if k == "noise":
- return self.noise.clone()
- return super().clone_key(k)
+ def clone_key(self, k):
+ if k == "custom_noise":
+ return self.custom_noise.clone()
+ return super().clone_key(k)
- def apply_effects(self, noise, sigma):
- if self.ffilter:
- noise = BLU.ffilter(
- noise,
- self.ffilter_threshold,
- self.ffilter_scale,
- self.ffilter,
- self.ffilter_strength,
- )
- if self.enhance_mode != "none" and self.enhance_strength != 0:
- noise = BLU.enhance_tensor(
- noise,
- self.enhance_mode,
- self.enhance_strength,
- sigma=sigma,
- skip_multiplier=0,
- adjust_scale=False,
- )
- return noise
-
- def make_noise_sampler(self, x, *args, normalized=True, **kwargs):
- factor = self.factor
- noise_items = self.noise.items
- noise_samplers = tuple(
- ni.make_noise_sampler(x, *args, normalized=False, **kwargs)
- for ni in noise_items
- )
- num_samplers = len(noise_samplers)
- normalize_noise = self.get_normalize(
- "normalize_noise",
- 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_init = torch.zeros_like(x)
-
- def noise_sampler(s, sn):
- noise = noise_init.clone()
- for ni, ns in zip(noise_items, noise_samplers):
- curr_noise = scale_noise(ns(s, sn), normalized=normalize_noise)
- if noise_effects:
- curr_noise = self.apply_effects(curr_noise, s)
- if self.blend_mode == "simple_add":
- noise += curr_noise.mul_(ni.factor)
- else:
- noise = BLU.BLENDING_MODES[self.blend_mode](
- noise,
- curr_noise,
- ni.factor,
- )
- del curr_noise
- noise = scale_noise(noise, factor, normalized=normalize_result)
- if result_effects:
- noise = self.apply_effects(noise, s)
- return noise
-
- 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(
+ def make_noise_sampler(self, x, *args, normalized=True, **kwargs):
+ if x.ndim < 3:
+ raise ValueError("ResizedNoise can only handle 3+ dimensional latents")
+ factor = self.factor
+ normalize = self.get_normalize("normalize", normalized)
+ xh, xw = x.shape[-2:]
+ nh, nw = self.height // 8, self.width // 8
+ offsh, offsw = self.crop_offset_vertical // 8, self.crop_offset_horizontal // 8
+ if xh == nh and xw == nw:
+ ns = self.custom_noise.make_noise_sampler(
x,
*args,
- normalized=False,
+ normalized=normalize,
**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)
+ def noise_sampler(*args, **kwargs):
+ return ns(*args, **kwargs).mul_(factor)
return noise_sampler
+ upscale_mode = self.upscale_mode
+ downscale_mode = self.downscale_mode
+ crop_mode = self.crop_mode
+ x_all_bigger = xh >= nh and xw >= nw
+ x_any_bigger = xh >= nh or xw >= nw
+ if x_all_bigger:
+ if self.initial_reference == "prefer_crop":
+ x = crop_samples(
+ x,
+ nw,
+ nh,
+ mode=self.crop_mode,
+ offset_width=offsw,
+ offset_height=offsh,
+ )
+ else:
+ x = utils.scale_samples(x, nw, nh, mode=self.downscale_mode)
+ output = partial(
+ utils.scale_samples,
+ width=xw,
+ height=xh,
+ mode=upscale_mode,
+ )
+ else:
+ x = utils.scale_samples(x, nw, nh, mode=self.upscale_mode)
+ if x_any_bigger:
+ output = partial(
+ utils.scale_samples,
+ width=xw,
+ height=xh,
+ mode=upscale_mode,
+ )
+ elif self.downscale_strategy == "scale":
+ output = partial(
+ utils.scale_samples,
+ width=xw,
+ height=xh,
+ mode=downscale_mode,
+ )
+ else:
+ output = partial(
+ crop_samples,
+ width=xw,
+ height=xh,
+ mode=crop_mode,
+ offset_width=offsw,
+ offset_height=offsh,
+ )
+ ns = self.custom_noise.make_noise_sampler(x, *args, normalized=False, **kwargs)
+ del x
+
+ def noise_sampler(*args, **kwargs):
+ return output(
+ scale_noise(ns(*args, **kwargs), factor, normalized=normalize),
+ )
+
+ return noise_sampler
+
+
+class WaveletFilteredNoise(CustomNoiseItemBase):
+ def __init__(self, factor, *, normalize, noise, normalize_noise=False, **kwargs):
+ super().__init__(
+ factor,
+ noise=noise,
+ normalize=normalize,
+ normalize_noise=normalize_noise,
+ **kwargs,
+ )
+
+ def clone_key(self, k):
+ if k == "noise":
+ return self.noise.clone()
+ return super().clone_key(k)
+
+ def make_noise_sampler(
+ self,
+ x,
+ sigma_min,
+ sigma_max,
+ *args,
+ normalized=True,
+ **kwargs,
+ ):
+ factor = self.factor
+ normalize = self.get_normalize("normalize", normalized)
+ internal_ns = self.noise.make_noise_sampler(
+ x,
+ *args,
+ sigma_min=sigma_min,
+ sigma_max=sigma_max,
+ normalized=self.normalize_noise,
+ **kwargs,
+ )
+ ns_kwargs = getattr(self, "ns_kwargs", {}).copy()
+ # print("WF:NS KWARGS", ns_kwargs)
+ kwargs |= ns_kwargs
+ ns = WaveletNoiseGenerator(
+ x,
+ *args,
+ sigma_min=sigma_min,
+ sigma_max=sigma_max,
+ normalized=False,
+ noise_sampler=internal_ns,
+ **kwargs,
+ )
+
+ def noise_sampler(sigma, sigma_next):
+ return scale_noise(
+ ns(sigma, sigma_next),
+ factor,
+ normalized=normalize,
+ )
+
+ return noise_sampler
+
+
+class BlendFilterNoise(CustomNoiseItemBase):
+ def __init__(
+ self,
+ factor,
+ *,
+ noise,
+ blend_mode,
+ ffilter,
+ ffilter_scale,
+ ffilter_strength,
+ ffilter_threshold,
+ enhance_mode,
+ enhance_strength,
+ affect,
+ normalize_result,
+ normalize_noise,
+ ):
+ if len(noise.items) == 0:
+ raise ValueError("BlendFilterNoise requires at least one noise item")
+ super().__init__(
+ factor,
+ noise=noise.clone(),
+ blend_mode=blend_mode,
+ ffilter=ffilter,
+ ffilter_scale=ffilter_scale,
+ ffilter_strength=ffilter_strength,
+ ffilter_threshold=ffilter_threshold,
+ enhance_mode=enhance_mode,
+ enhance_strength=enhance_strength,
+ affect=affect,
+ normalize_result=normalize_result,
+ normalize_noise=normalize_noise,
+ )
+
+ def clone_key(self, k):
+ if k == "noise":
+ return self.noise.clone()
+ return super().clone_key(k)
+
+ def apply_effects(self, noise, sigma):
+ blu = external.MODULES.bleh.py.latent_utils
+ if self.ffilter:
+ noise = blu.ffilter(
+ noise,
+ self.ffilter_threshold,
+ self.ffilter_scale,
+ self.ffilter,
+ self.ffilter_strength,
+ )
+ if self.enhance_mode != "none" and self.enhance_strength != 0:
+ noise = blu.enhance_tensor(
+ noise,
+ self.enhance_mode,
+ self.enhance_strength,
+ sigma=sigma,
+ skip_multiplier=0,
+ adjust_scale=False,
+ )
+ return noise
+
+ def make_noise_sampler(self, x, *args, normalized=True, **kwargs):
+ factor = self.factor
+ noise_items = self.noise.items
+ noise_samplers = tuple(
+ ni.make_noise_sampler(x, *args, normalized=False, **kwargs)
+ for ni in noise_items
+ )
+ num_samplers = len(noise_samplers)
+ normalize_noise = self.get_normalize(
+ "normalize_noise",
+ 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_init = torch.zeros_like(x)
+
+ def noise_sampler(s, sn):
+ noise = noise_init.clone()
+ for ni, ns in zip(noise_items, noise_samplers):
+ curr_noise = scale_noise(ns(s, sn), normalized=normalize_noise)
+ if noise_effects:
+ curr_noise = self.apply_effects(curr_noise, s)
+ if self.blend_mode == "simple_add":
+ noise += curr_noise.mul_(ni.factor)
+ else:
+ noise = utils.BLENDING_MODES[self.blend_mode](
+ noise,
+ curr_noise,
+ ni.factor,
+ )
+ del curr_noise
+ noise = scale_noise(noise, factor, normalized=normalize_result)
+ if result_effects:
+ noise = self.apply_effects(noise, s)
+ return noise
+
+ 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):
+ bops = external.MODULES.bleh.py.nodes.ops
+ 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),
- 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.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(
+ NoiseType.BROWNIAN: NoiseSampler.wrap(BrownianNoiseGenerator),
+ NoiseType.DISTRO: NoiseSampler.wrap(DistroNoiseGenerator),
+ NoiseType.GAUSSIAN: NoiseSampler.wrap(GaussianNoiseGenerator),
+ NoiseType.UNIFORM: NoiseSampler.wrap(UniformNoiseGenerator),
+ NoiseType.PERLIN: NoiseSampler.wrap(PerlinOldNoiseGenerator),
+ NoiseType.STUDENTT: NoiseSampler.wrap(StudentTNoiseGenerator),
+ NoiseType.ONEF_PINKISH: NoiseSampler.wrap(partial(OneFNoiseGenerator, alpha=-0.5)),
+ NoiseType.ONEF_GREENISH: NoiseSampler.wrap(partial(OneFNoiseGenerator, alpha=0.5)),
+ NoiseType.ONEF_PINKISHGREENISH: NoiseSampler.wrap(
partial(
- powerlaw_noise_like,
+ MixedNoiseGenerator,
+ name="onef_pinkishgreenish",
+ noise_mix=(
+ partial(OneFNoiseGenerator, alpha=0.5),
+ partial(OneFNoiseGenerator, alpha=-0.5),
+ ),
+ output_fun=lambda t: t.mul_(0.5),
+ ),
+ ),
+ NoiseType.ONEF_PINKISH_MIX: NoiseSampler.wrap(
+ partial(
+ MixedNoiseGenerator,
+ name="onef_pinkish_mix",
+ noise_mix=(
+ (partial(OneFNoiseGenerator, alpha=0.5), lambda t: t.mul_(-1.0)),
+ partial(OneFNoiseGenerator, alpha=0.5),
+ ),
+ output_fun=lambda t: t.mul_(0.5),
+ ),
+ ),
+ NoiseType.ONEF_GREENISH_MIX: NoiseSampler.wrap(
+ partial(
+ MixedNoiseGenerator,
+ name="onef_greenish_mix",
+ noise_mix=(
+ (partial(OneFNoiseGenerator, alpha=0.5), lambda t: t.mul_(-1.0)),
+ partial(OneFNoiseGenerator, alpha=0.5),
+ ),
+ output_fun=lambda t: t.mul_(0.5),
+ ),
+ ),
+ NoiseType.WHITE: NoiseSampler.wrap(
+ partial(
+ PowerLawNoiseGenerator,
alpha=0.0,
use_sign=True,
),
),
- NoiseType.GREY: NoiseSampler.simple(
+ NoiseType.GREY: NoiseSampler.wrap(
partial(
- powerlaw_noise_like,
+ PowerLawNoiseGenerator,
alpha=0.0,
use_sign=False,
),
),
- NoiseType.VELVET: NoiseSampler.simple(
+ NoiseType.VELVET: NoiseSampler.wrap(
partial(
- powerlaw_noise_like,
+ PowerLawNoiseGenerator,
alpha=1.0,
use_sign=True,
div_max_dims=(-3, -2, -1),
),
),
- NoiseType.VIOLET: NoiseSampler.simple(
+ NoiseType.VIOLET: NoiseSampler.wrap(
partial(
- powerlaw_noise_like,
+ PowerLawNoiseGenerator,
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)
- .mul_(0.55)
- .add_(rand_perlin_like(x).mul_(0.7))
- .mul_(1.15),
+ NoiseType.WAVELET: NoiseSampler.wrap(WaveletNoiseGenerator),
+ NoiseType.PINK_OLD: NoiseSampler.wrap(PinkOldNoiseGenerator),
+ NoiseType.HIGHRES_PYRAMID: NoiseSampler.wrap(HighresPyramidNoiseGenerator),
+ NoiseType.PYRAMID: NoiseSampler.wrap(PyramidNoiseGenerator),
+ NoiseType.RAINBOW_MILD: NoiseSampler.wrap(
+ partial(
+ MixedNoiseGenerator,
+ name="rainbow_mild",
+ noise_mix=(
+ (GreenTestNoiseGenerator, lambda t: t.mul_(0.55)),
+ (GreenTestNoiseGenerator, lambda t: t.mul_(0.7)),
+ ),
+ output_fun=lambda t: t.mul_(1.15),
+ ),
),
- NoiseType.RAINBOW_INTENSE: NoiseSampler.simple(
- lambda x: green_noise_like(x)
- .mul_(0.75)
- .add_(rand_perlin_like(x).mul_(0.5))
- .mul_(1.15),
+ NoiseType.RAINBOW_INTENSE: NoiseSampler.wrap(
+ partial(
+ MixedNoiseGenerator,
+ name="rainbow_intense",
+ noise_mix=(
+ (GreenTestNoiseGenerator, lambda t: t.mul_(0.75)),
+ (GreenTestNoiseGenerator, lambda t: t.mul_(0.5)),
+ ),
+ output_fun=lambda t: t.mul_(1.15),
+ ),
),
- NoiseType.LAPLACIAN: NoiseSampler.simple(laplacian_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(
- partial(pyramid_noise_like, upscale_mode="bislerp"),
+ NoiseType.LAPLACIAN: NoiseSampler.wrap(LaplacianNoiseGenerator),
+ NoiseType.POWER_OLD: NoiseSampler.wrap(PowerOldNoiseGenerator),
+ NoiseType.GREEN_TEST: NoiseSampler.wrap(GreenTestNoiseGenerator),
+ NoiseType.PYRAMID_OLD: NoiseSampler.wrap(PyramidOldNoiseGenerator),
+ NoiseType.PYRAMID_BISLERP: NoiseSampler.wrap(
+ partial(PyramidNoiseGenerator, upscale_mode="bislerp"),
),
- NoiseType.HIGHRES_PYRAMID_BISLERP: NoiseSampler.simple(
- partial(highres_pyramid_noise_like, upscale_mode="bislerp"),
+ NoiseType.HIGHRES_PYRAMID_BISLERP: NoiseSampler.wrap(
+ partial(HighresPyramidNoiseGenerator, upscale_mode="bislerp"),
),
- NoiseType.PYRAMID_AREA: NoiseSampler.simple(
- partial(pyramid_noise_like, upscale_mode="area"),
+ NoiseType.PYRAMID_AREA: NoiseSampler.wrap(
+ partial(PyramidNoiseGenerator, upscale_mode="area"),
),
- NoiseType.HIGHRES_PYRAMID_AREA: NoiseSampler.simple(
- partial(highres_pyramid_noise_like, upscale_mode="area"),
+ NoiseType.HIGHRES_PYRAMID_AREA: NoiseSampler.wrap(
+ partial(HighresPyramidNoiseGenerator, upscale_mode="area"),
),
- NoiseType.PYRAMID_OLD_BISLERP: NoiseSampler.simple(
- partial(pyramid_old_noise_like, upscale_mode="bislerp"),
+ NoiseType.PYRAMID_OLD_BISLERP: NoiseSampler.wrap(
+ partial(PyramidOldNoiseGenerator, upscale_mode="bislerp"),
),
- NoiseType.PYRAMID_OLD_AREA: NoiseSampler.simple(
- partial(pyramid_old_noise_like, upscale_mode="area"),
+ NoiseType.PYRAMID_OLD_AREA: NoiseSampler.wrap(
+ partial(PyramidOldNoiseGenerator, upscale_mode="area"),
),
- NoiseType.PYRAMID_DISCOUNT5: NoiseSampler.simple(
- partial(pyramid_noise_like, discount=0.5),
+ NoiseType.PYRAMID_DISCOUNT5: NoiseSampler.wrap(
+ partial(PyramidNoiseGenerator, discount=0.5),
),
- NoiseType.PYRAMID_MIX: NoiseSampler.simple(
- lambda x: pyramid_noise_like(x, discount=0.6)
- .mul_(0.2)
- .add_(pyramid_noise_like(x, discount=0.6).mul_(-0.8)),
+ NoiseType.PYRAMID_MIX: NoiseSampler.wrap(
+ partial(
+ MixedNoiseGenerator,
+ name="pyramid_mix",
+ noise_mix=(
+ (partial(PyramidNoiseGenerator, discount=0.6), lambda t: t.mul_(0.2)),
+ (partial(PyramidNoiseGenerator, discount=0.6), lambda t: t.mul_(-0.8)),
+ ),
+ ),
),
- NoiseType.PYRAMID_MIX_AREA: NoiseSampler.simple(
- lambda x: pyramid_noise_like(x, discount=0.5, upscale_mode="area")
- .mul_(0.2)
- .add_(pyramid_noise_like(x, discount=0.5, upscale_mode="area").mul_(-0.8)),
+ NoiseType.PYRAMID_MIX_AREA: NoiseSampler.wrap(
+ partial(
+ MixedNoiseGenerator,
+ name="pyramid_mix_area",
+ noise_mix=(
+ (
+ partial(PyramidNoiseGenerator, discount=0.5, upscale_mode="area"),
+ lambda t: t.mul_(0.2),
+ ),
+ (
+ partial(PyramidNoiseGenerator, discount=0.5, upscale_mode="area"),
+ lambda t: t.mul_(-0.8),
+ ),
+ ),
+ ),
),
- NoiseType.PYRAMID_MIX_BISLERP: NoiseSampler.simple(
- lambda x: pyramid_noise_like(x, discount=0.5, upscale_mode="bislerp")
- .mul_(0.2)
- .add_(pyramid_noise_like(x, discount=0.5, upscale_mode="bislerp").mul_(-0.8)),
+ NoiseType.PYRAMID_MIX_BISLERP: NoiseSampler.wrap(
+ partial(
+ MixedNoiseGenerator,
+ name="pyramid_mix_bislerp",
+ noise_mix=(
+ (
+ partial(
+ PyramidNoiseGenerator,
+ discount=0.5,
+ upscale_mode="bislerp",
+ ),
+ lambda t: t.mul_(0.2),
+ ),
+ (
+ partial(
+ PyramidNoiseGenerator,
+ discount=0.5,
+ upscale_mode="bislerp",
+ ),
+ lambda t: t.mul_(-0.8),
+ ),
+ ),
+ ),
),
}
@@ -1275,6 +1559,7 @@ def get_noise_sampler(
cpu: bool = True,
factor: float = 1.0,
normalized=False,
+ **kwargs,
) -> Callable:
if noise_type is None:
noise_type = NoiseType.GAUSSIAN
@@ -1293,4 +1578,5 @@ def get_noise_sampler(
cpu=cpu,
factor=factor,
normalized=normalized,
+ **kwargs,
)
diff --git a/py/noise_generation.py b/py/noise_generation.py
index 40a9e90..2fe42e1 100644
--- a/py/noise_generation.py
+++ b/py/noise_generation.py
@@ -3,21 +3,30 @@ from __future__ import annotations
import math
from enum import Enum, auto
+from functools import lru_cache
from typing import Callable
import torch
-from comfy.model_management import device_supports_non_blocking
-from comfy.utils import common_upscale
+from comfy.k_diffusion import sampling
from torch import FloatTensor, Generator, Tensor
from torch.distributions import Laplace, StudentT
-from .external import MODULES as EXT
+try:
+ import pytorch_wavelets as ptwav
+
+ HAVE_WAVELETS = True
+except ImportError:
+ HAVE_WAVELETS = False
+
+from . import utils
+from .utils import quantile_normalize, scale_noise, tensor_to
# ruff: noqa: D412, D413, D417, D212, D407, ANN002, ANN003, FBT001, FBT002, S311
class NoiseType(Enum):
BROWNIAN = auto()
+ DISTRO = auto()
GAUSSIAN = auto()
GREEN_TEST = auto()
GREY = auto()
@@ -49,6 +58,7 @@ class NoiseType(Enum):
UNIFORM = auto()
VELVET = auto()
VIOLET = auto()
+ WAVELET = auto()
WHITE = auto()
@classmethod
@@ -67,499 +77,1282 @@ class NoiseError(Exception):
pass
-def scale_noise(noise, factor=1.0, *, normalized=True, threshold_std_devs=2.5):
- if not normalized or noise.numel() == 0:
- return noise.mul_(factor) if factor != 1 else noise
- 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
- return noise.mul_(factor) if factor != 1 else noise
+class NoiseGenerator:
+ name = "unknown"
+ SAVE_X = False
+ MIN_DIMS = 1
+ MAX_DIMS = 0
-
-if "bleh" in EXT:
- scale_samples = EXT["bleh"].py.latent_utils.scale_samples
-else:
-
- def scale_samples(
- samples,
- width,
- height,
- *,
- mode="bicubic",
+ def __init__(
+ self,
+ x,
+ **kwargs,
):
- return common_upscale(samples, width, height, mode, None)
+ if x.ndim < self.MIN_DIMS:
+ errstr = f"Noise generator {self.name} requires at least {self.MIN_DIMS} dimension(s) but got input with shape {x.shape}"
+ raise ValueError(errstr)
+ if self.MAX_DIMS > 0 and x.ndim > self.MAX_DIMS:
+ errstr = f"Noise generator {self.name} requires at most {self.MAX_DIMS} dimension(s) but got input with shape {x.shape}"
+ raise ValueError(errstr)
+ params = self.ng_params
+ kwarg_params = params | kwargs
+ for k in params:
+ setattr(self, k, kwarg_params.pop(k))
+ self.options = kwarg_params
+ self.update_x(x)
+ print("CREATE NG", self, kwargs)
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return {
+ "normalized": True,
+ "force_normalize": None,
+ "normalize_dims": None,
+ "cpu": True,
+ "generator": None,
+ }
+
+ def update_x(self, x):
+ self.x = None if not self.SAVE_X else x.detach().clone()
+ self.shape = x.shape
+ self.batch, self.channels, self.height, self.width = (
+ x.shape if x.ndim == 4 else (None,) * 4
+ )
+ self.device = x.device
+ self.gen_device = torch.device("cpu") if self.cpu else self.device
+ self.layout = x.layout
+ self.dtype = x.dtype
+
+ def rand_like(self, *, fun=torch.randn, cpu=None, to_device=True):
+ cpu = cpu if cpu is not None else self.cpu
+ noise = fun(
+ *self.shape,
+ generator=self.generator,
+ dtype=self.dtype,
+ layout=self.layout,
+ device=self.gen_device,
+ )
+ # print("GEN NOISE", noise)
+ if to_device and noise.device != self.device:
+ noise = tensor_to(noise, self.device)
+ # print("MADE NOISE", noise)
+ return noise
+
+ def output_hook(self, noise):
+ # print("NOISE OUT1", noise)
+ if noise.device != self.device:
+ noise = tensor_to(noise, self.device)
+ return scale_noise(
+ noise,
+ normalized=self.normalized
+ and (self.force_normalize is None or self.force_normalize is True),
+ normalize_dims=self.normalize_dims,
+ )
+
+ def pre_hook(self):
+ pass
+
+ def generate(self):
+ raise NotImplementedError
+
+ def __call__(self, *args, **kwargs):
+ self.pre_hook()
+ return self.output_hook(self.generate(*args, **kwargs))
+
+ def __str__(self):
+ pretty_params = ", ".join(f"{k}={getattr(self, k)!s}" for k in self.ng_params)
+ return f""
-CAN_NONBLOCK = {}
+class MixedNoiseGenerator(NoiseGenerator):
+ MIN_DIMS = MAX_DIMS = 4
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "name": "mixed_noise",
+ "normalized": True,
+ "pass_args": frozenset(("cpu",)),
+ "noise_mix": (),
+ "output_fun": None,
+ }
+
+ def __init__(self, x, *args, **kwargs):
+ super().__init__(x, *args, **kwargs)
+ ng_list = []
+ for item in self.noise_mix:
+ if isinstance(item, (tuple, list)):
+ ng_class, transform_fun = item
+ else:
+ ng_class, transform_fun = item, None
+ ng_kwargs = {k: v for k, v in kwargs.items() if k in self.pass_args}
+ ng_list.append((ng_class(x, **ng_kwargs), transform_fun))
+ self.ng_list = ng_list
+
+ def generate(self, *args):
+ noise = None
+ for ng, transform_fun in self.ng_list:
+ new_noise = ng(*args)
+ if transform_fun is not None:
+ new_noise = transform_fun(new_noise)
+ noise = new_noise if noise is None else noise.add_(new_noise)
+ if self.output_fun is not None:
+ noise = self.output_fun(noise)
+ return noise
-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)
+class GaussianNoiseGenerator(NoiseGenerator):
+ name = "gaussian"
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {"normalized": False}
+
+ def generate(self, *_args):
+ return self.rand_like()
-def get_positions(block_shape: tuple[int, int]) -> Tensor:
- """
- Generate position tensor.
+class BrownianNoiseGenerator(NoiseGenerator):
+ name = "brownian"
- Arguments:
- block_shape -- (height, width) of position tensor
+ def __init__(self, x, *args, **kwargs):
+ super().__init__(x, *args, **kwargs)
+ seed = self.options.get("seed")
+ sigma_min = self.options.get("sigma_min")
+ sigma_max = self.options.get("sigma_max")
+ if sigma_min is None or sigma_max is None:
+ raise ValueError("Brownian noise requires sigma_min and sigma_max")
+ self.brownian_tree_ns = sampling.BrownianTreeNoiseSampler(
+ x,
+ sigma_min,
+ sigma_max,
+ seed=seed,
+ cpu=self.cpu,
+ )
- Returns:
- position vector shaped (1, height, width, 1, 1, 2)
- """
- bh, bw = block_shape
- return torch.stack(
- torch.meshgrid(
- [(torch.arange(b) + 0.5) / b for b in (bw, bh)],
- indexing="xy",
- ),
- -1,
- ).view(1, bh, bw, 1, 1, 2)
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {"normalized": False}
+
+ def generate(self, *args):
+ return self.brownian_tree_ns(*args)
-def unfold_grid(vectors: Tensor) -> Tensor:
- """
- Unfold vector grid to batched vectors.
+class PerlinOldNoiseGenerator(NoiseGenerator):
+ name = "perlin_old"
+ MIN_DIMS = MAX_DIMS = 4
- Arguments:
- vectors -- grid vectors
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "div_fac": 2.0,
+ "iterations": 2,
+ "blend_mode": "lerp",
+ }
- Returns:
- batched grid vectors
- """
- batch_size, _channels, gpy, gpx = vectors.shape
- return (
- torch.nn.functional.unfold(vectors, (2, 2))
- .view(batch_size, 2, 4, -1)
- .permute(0, 2, 3, 1)
- .view(batch_size, 4, gpy - 1, gpx - 1, 2)
- )
+ @staticmethod
+ def get_positions(block_shape: tuple[int, int]) -> Tensor:
+ """
+ Generate position tensor.
+ Arguments:
+ block_shape -- (height, width) of position tensor
-def smooth_step(t: Tensor) -> Tensor:
- """
- Smooth step function [0, 1] -> [0, 1].
+ Returns:
+ position vector shaped (1, height, width, 1, 1, 2)
+ """
+ bh, bw = block_shape
+ return torch.stack(
+ torch.meshgrid(
+ [(torch.arange(b) + 0.5) / b for b in (bw, bh)],
+ indexing="xy",
+ ),
+ -1,
+ ).view(1, bh, bw, 1, 1, 2)
- Arguments:
- t -- input values (any shape)
+ @staticmethod
+ def unfold_grid(vectors: Tensor) -> Tensor:
+ """
+ Unfold vector grid to batched vectors.
- Returns:
- output values (same shape as input values)
- """
- return t * t * (3.0 - 2.0 * t)
+ Arguments:
+ vectors -- grid vectors
+ Returns:
+ batched grid vectors
+ """
+ batch_size, _channels, gpy, gpx = vectors.shape
+ return (
+ torch.nn.functional.unfold(vectors, (2, 2))
+ .view(batch_size, 2, 4, -1)
+ .permute(0, 2, 3, 1)
+ .view(batch_size, 4, gpy - 1, gpx - 1, 2)
+ )
-def perlin_noise_tensor(
- vectors: Tensor,
- positions: Tensor,
- step: Callable | None = None,
-) -> Tensor:
- """
- Generate perlin noise from batched vectors and positions.
+ @staticmethod
+ def smooth_step(t: Tensor) -> Tensor:
+ """
+ Smooth step function [0, 1] -> [0, 1].
- Arguments:
- vectors -- batched grid vectors shaped (batch_size, 4, grid_height, grid_width, 2)
- positions -- batched grid positions shaped (batch_size or 1, block_height, block_width, grid_height or 1, grid_width or 1, 2)
+ Arguments:
+ t -- input values (any shape)
- Keyword Arguments:
- step -- smooth step function [0, 1] -> [0, 1] (default: `smooth_step`)
+ Returns:
+ output values (same shape as input values)
+ """
+ return t * t * (3.0 - 2.0 * t)
- Raises:
- NoiseError: if position and vector shapes do not match
+ @classmethod
+ def perlin_noise_tensor(
+ cls,
+ vectors: Tensor,
+ positions: Tensor,
+ step: Callable | None = None,
+ blend=torch.lerp,
+ ) -> Tensor:
+ """
+ Generate perlin noise from batched vectors and positions.
- Returns:
- (batch_size, block_height * grid_height, block_width * grid_width)
- """
- if step is None:
- step = smooth_step
+ Arguments:
+ vectors -- batched grid vectors shaped (batch_size, 4, grid_height, grid_width, 2)
+ positions -- batched grid positions shaped (batch_size or 1, block_height, block_width, grid_height or 1, grid_width or 1, 2)
- batch_size = vectors.shape[0]
- # grid height, grid width
- gh, gw = vectors.shape[2:4]
- # block height, block width
- bh, bw = positions.shape[1:3]
+ Keyword Arguments:
+ step -- smooth step function [0, 1] -> [0, 1] (default: `smooth_step`)
- for i in range(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})"
+ Raises:
+ NoiseError: if position and vector shapes do not match
+
+ Returns:
+ (batch_size, block_height * grid_height, block_width * grid_width)
+ """
+ if step is None:
+ step = cls.smooth_step
+
+ batch_size = vectors.shape[0]
+ # grid height, grid width
+ gh, gw = vectors.shape[2:4]
+ # block height, block width
+ bh, bw = positions.shape[1:3]
+
+ for i in range(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}:
+ msg = f"Batch sizes do not match: vectors ({vectors.shape[0]}), positions ({positions.shape[0]})"
raise NoiseError(msg)
- 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)
+ vectors = vectors.view(batch_size, 4, 1, gh * gw, 2)
+ positions = positions.view(positions.shape[0], bh * bw, -1, 2)
- vectors = vectors.view(batch_size, 4, 1, gh * gw, 2)
- positions = positions.view(positions.shape[0], bh * bw, -1, 2)
+ step_x = step(positions[..., 0])
+ step_y = step(positions[..., 1])
- step_x = step(positions[..., 0])
- step_y = step(positions[..., 1])
-
- row0 = torch.lerp(
- (vectors[:, 0] * positions).sum(dim=-1),
- (vectors[:, 1] * (positions - positions.new_tensor((1, 0)))).sum(dim=-1),
- step_x,
- )
- row1 = torch.lerp(
- (vectors[:, 2] * (positions - positions.new_tensor((0, 1)))).sum(dim=-1),
- (vectors[:, 3] * (positions - positions.new_tensor((1, 1)))).sum(dim=-1),
- step_x,
- )
- noise = torch.lerp(row0, row1, step_y)
- return (
- noise.view(
- batch_size,
- bh,
- bw,
- gh,
- gw,
+ row0 = blend(
+ (vectors[:, 0] * positions).sum(dim=-1),
+ (vectors[:, 1] * (positions - positions.new_tensor((1, 0)))).sum(dim=-1),
+ step_x,
+ )
+ row1 = blend(
+ (vectors[:, 2] * (positions - positions.new_tensor((0, 1)))).sum(dim=-1),
+ (vectors[:, 3] * (positions - positions.new_tensor((1, 1)))).sum(dim=-1),
+ step_x,
+ )
+ noise = blend(row0, row1, step_y)
+ return (
+ noise.view(
+ batch_size,
+ bh,
+ bw,
+ gh,
+ gw,
+ )
+ .permute(0, 3, 1, 4, 2)
+ .reshape(batch_size, gh * bh, gw * bw)
)
- .permute(0, 3, 1, 4, 2)
- .reshape(batch_size, gh * bh, gw * bw)
- )
-
-def perlin_noise(
- grid_shape: tuple[int, int],
- out_shape: tuple[int, int],
- batch_size: int = 1,
- generator: Generator | None = None,
- *args,
- **kwargs,
-) -> Tensor:
- """
- Generate perlin noise with given shape. `*args` and `**kwargs` are forwarded to `Tensor` creation.
-
- Arguments:
- grid_shape -- Shape of grid (height, width).
- out_shape -- Shape of output noise image (height, width).
-
- Keyword Arguments:
- batch_size -- (default: {1})
- generator -- random generator used for grid vectors (default: {None})
-
- Raises:
- NoiseError: if grid and out shapes do not match
-
- Returns:
- Noise image shaped (batch_size, height, width)
- """
- # grid height and width
- gh, gw = grid_shape
- # output height and width
- oh, ow = out_shape
- # block height and width
- bh, bw = oh // gh, ow // gw
-
- if oh != bh * gh:
- msg = f"Output height {oh} must be divisible by grid height {gh}"
- raise NoiseError(msg)
- if ow != bw * gw != 0:
- msg = f"Output width {ow} must be divisible by grid width {gw}"
- raise NoiseError(msg)
-
- angle = torch.empty(
- [batch_size] + [s + 1 for s in grid_shape],
+ @classmethod
+ def perlin_noise(
+ cls,
+ grid_shape: tuple[int, int],
+ out_shape: tuple[int, int],
+ batch_size: int = 1,
+ blend=torch.lerp,
+ generator: Generator | None = None,
*args,
**kwargs,
- ).uniform_(to=2.0 * math.pi, generator=generator)
- # 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 = tensor_to(get_positions((bh, bw)), vectors)
- return perlin_noise_tensor(vectors, positions).squeeze(0)
+ ) -> Tensor:
+ """
+ Generate perlin noise with given shape. `*args` and `**kwargs` are forwarded to `Tensor` creation.
+ Arguments:
+ grid_shape -- Shape of grid (height, width).
+ out_shape -- Shape of output noise image (height, width).
-def rand_perlin_like(x, *, generator=None):
- noise = (
- torch.rand(
- x.shape,
- dtype=x.dtype,
- device=x.device,
- layout=x.layout,
- generator=generator,
+ Keyword Arguments:
+ batch_size -- (default: {1})
+ generator -- random generator used for grid vectors (default: {None})
+
+ Raises:
+ NoiseError: if grid and out shapes do not match
+
+ Returns:
+ Noise image shaped (batch_size, height, width)
+ """
+ # grid height and width
+ gh, gw = grid_shape
+ # output height and width
+ oh, ow = out_shape
+ # block height and width
+ bh, bw = oh // gh, ow // gw
+
+ if oh != bh * gh:
+ msg = f"Output height {oh} must be divisible by grid height {gh}"
+ raise NoiseError(msg)
+ if ow != bw * gw != 0:
+ msg = f"Output width {ow} must be divisible by grid width {gw}"
+ raise NoiseError(msg)
+
+ angle = torch.empty(
+ [batch_size] + [s + 1 for s in grid_shape],
+ *args,
+ **kwargs,
+ ).uniform_(to=2.0 * math.pi, generator=generator)
+ # random vectors on grid points
+ vectors = cls.unfold_grid(
+ torch.stack((torch.cos(angle), torch.sin(angle)), dim=1),
)
- / 2.0
- )
- noise_height = noise.size(dim=2)
- noise_width = noise.size(dim=3)
- for _ in range(2):
+ # positions inside grid cells [0, 1)
+ positions = tensor_to(cls.get_positions((bh, bw)), vectors)
+ return cls.perlin_noise_tensor(vectors, positions, blend=blend).squeeze(0)
+
+ def generate(self, *_args):
+ blend = utils.BLENDING_MODES[self.blend_mode]
+ noise = self.rand_like(fun=torch.rand).div_(self.div_fac)
+
+ _batch, channels, noise_height, noise_width = noise.shape
+ for _ in range(self.iterations):
+ noise += self.perlin_noise(
+ (noise_height, noise_width),
+ (noise_height, noise_width),
+ batch_size=channels, # This should be the number of channels.
+ blend=blend,
+ dtype=noise.dtype,
+ layout=noise.layout,
+ device=noise.device,
+ )
+ return noise
+
+
+class UniformNoiseGenerator(NoiseGenerator):
+ name = "uniform"
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "normalized": False,
+ "sub_fac": 0.5,
+ "mul_fac": 3.46,
+ "mean_fac": 0.0,
+ }
+
+ def generate(self, *_args):
+ return (
+ self.rand_like(fun=torch.rand)
+ .sub_(self.sub_fac)
+ .mul_(self.mul_fac)
+ .add_(self.mean_fac)
+ )
+
+
+class HighresPyramidNoiseGenerator(NoiseGenerator):
+ name = "highres_pyramid"
+ MIN_DIMS = MAX_DIMS = 4
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.uniform_ng = UniformNoiseGenerator(
+ *args,
+ **(
+ kwargs
+ | {
+ "normalized": self.uniform_normalized,
+ "normalize_dims": self.options.get("uniform_normalize_dims"),
+ }
+ ),
+ )
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "normalized": True,
+ "uniform_normalized": False,
+ "discount": 0.7,
+ "upscale_mode": "bilinear",
+ "iterations": 4,
+ }
+
+ def generate(self, s, sn):
+ (
+ b,
+ c,
+ h,
+ w,
+ ) = self.shape # EDIT: w and h get over-written, rename for a different variant!
+
+ orig_w, orig_h = w, h
+ noise = self.uniform_ng(s, sn)
+ rs = (
+ torch.rand(
+ self.iterations,
+ dtype=torch.float32,
+ generator=self.generator,
+ ).cpu()
+ * 2
+ + 2
+ )
+ for i in range(self.iterations):
+ r = rs[i].item()
+ h, w = min(orig_h * 15, int(h * (r**i))), min(orig_w * 15, int(w * (r**i)))
+ noise += utils.scale_samples(
+ tensor_to(torch.randn(b, c, h, w, generator=self.generator), noise),
+ orig_w,
+ orig_h,
+ mode=self.upscale_mode,
+ ).mul_(self.discount**i)
+ if h >= orig_h * 15 or w >= orig_w * 15:
+ break # Lowest resolution is 1x1
+ return noise
+
+
+class PyramidOldNoiseGenerator(NoiseGenerator):
+ name = "pyramid_old"
+ MIN_DIMS = MAX_DIMS = 4
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "discount": 0.8,
+ "iterations": 5,
+ "upscale_mode": "nearest-exact",
+ "normalized": False,
+ }
+
+ def generate(self, *_args):
+ b, c, h, w = self.shape
+ orig_h, orig_w = h, w
+ noise = torch.zeros(
+ size=self.shape,
+ dtype=self.dtype,
+ layout=self.layout,
+ device=self.gen_device,
+ )
+ r = 1
+ for i in range(self.iterations):
+ r *= 2
+ noise += utils.scale_samples(
+ torch.normal(
+ mean=0,
+ std=0.5**i,
+ size=(b, c, h * r, w * r),
+ dtype=noise.dtype,
+ layout=noise.layout,
+ generator=self.generator,
+ device=noise.device,
+ ),
+ orig_w,
+ orig_h,
+ mode=self.upscale_mode,
+ ).mul_(self.discount**i)
+ return noise
+
+
+class PyramidNoiseGenerator(NoiseGenerator):
+ name = "pyramid"
+ MIN_DIMS = MAX_DIMS = 4
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "discount": 0.7,
+ "upscale_mode": "bilinear",
+ "iterations": 10,
+ }
+
+ # Modified from https://wandb.ai/johnowhitaker/multires_noise/reports/Multi-Resolution-Noise-for-Diffusion-Model-Training--VmlldzozNjYyOTU2
+ def generate(self, *_args):
+ b, c, w, h = (
+ self.shape
+ ) # NOTE: w and h get over-written, rename for a different variant!
+
+ orig_w, orig_h = w, h
+ noise = self.rand_like()
+ for i in range(self.iterations):
+ r = (
+ torch.rand(1, generator=self.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 += utils.scale_samples(
+ torch.randn(
+ b,
+ c,
+ w,
+ h,
+ device=noise.device,
+ layout=noise.layout,
+ dtype=noise.dtype,
+ ),
+ orig_h,
+ orig_w,
+ mode=self.upscale_mode,
+ ).mul_(
+ self.discount**i,
+ )
+ if w == 1 or h == 1:
+ break # Lowest resolution is 1x1
+ return noise
+
+
+class StudentTNoiseGenerator(NoiseGenerator):
+ name = "studentt"
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "loc": 0,
+ "scale": 0.2,
+ "df": 1,
+ "quantile_fac": 0.75,
+ "pow_fac": 0.5,
+ "nq_fac": 1.0,
+ "normalized": False,
+ }
+
+ def generate(self, *_args):
+ noise = StudentT(loc=self.loc, scale=self.scale, df=self.df).rsample(self.shape)
+ nq: FloatTensor = torch.quantile(
+ noise.flatten(start_dim=1).abs(),
+ self.quantile_fac,
+ dim=-1,
+ )
+ nq_shape = tuple(nq.shape) + (1,) * (noise.ndim - nq.ndim)
+ nq = nq.mul_(self.nq_fac).reshape(*nq_shape)
+ noise = noise.clamp(-nq, nq)
+ return torch.copysign(torch.pow(torch.abs(noise), self.pow_fac), noise)
+
+
+class GreenTestNoiseGenerator(NoiseGenerator):
+ name = "green_test"
+ MIN_DIMS = MAX_DIMS = 4
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "scale_fac": 1.0,
+ "x_pow": 2,
+ "y_pow": 2,
+ "power_base": 1,
+ }
+
+ def generate(self, *_args):
+ noise = self.rand_like()
+ scale = self.scale_fac / (self.width * self.height)
+ fy = torch.fft.fftfreq(self.height, device=noise.device)[:, None] ** self.y_pow
+ fx = torch.fft.fftfreq(self.width, device=noise.device) ** self.x_pow
+ f = fy + fx
+ power = torch.sqrt(f)
+ power[0, 0] = self.power_base
+ noise = torch.fft.ifft2(torch.fft.fft2(noise) / torch.sqrt(power))
+ noise *= scale / noise.std()
+ return torch.real(noise)
+
+
+class PinkOldNoiseGenerator(NoiseGenerator):
+ name = "pink_old"
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {"alpha": 2.0, "k": 1.0, "freq": 1.0}
+
+ # Completely wrong implementation here.
+ def generate(self, *_args):
+ spectral_density = self.k / self.freq**self.alpha
+ return self.rand_like() * spectral_density
+
+
+class OneFNoiseGenerator(NoiseGenerator):
+ name = "onef"
+ MIN_DIMS = MAX_DIMS = 4
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "alpha": 2.0,
+ "k": 1.0,
+ "hfac": 1.0,
+ "wfac": 1.0,
+ "base_power": 1.0,
+ "use_sqrt": True,
+ }
+
+ # Referenced from: https://github.com/WASasquatch/PowerNoiseSuite
+ def generate(self, *_args):
+ batch, _channels, height, width = self.shape
+
+ noise = self.rand_like()
+
+ freq_x = tensor_to(torch.fft.fftfreq(height, self.hfac), noise)
+ freq_y = tensor_to(torch.fft.fftfreq(width, self.wfac), noise)
+ fx, fy = torch.meshgrid(freq_x, freq_y, indexing="ij")
+
+ power = (fx**2 + fy**2) ** (-self.alpha / 2.0)
+ if self.k != 0:
+ power = self.k / power
+ power[0, 0] = self.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 self.use_sqrt
+ else power.to(noise_fft.dtype)
+ )
+
+ return torch.fft.ifftn(noise_fft).real
+
+
+class PowerLawNoiseGenerator(NoiseGenerator):
+ name = "powerlaw"
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "alpha": 2.0,
+ "div_max_dims": None,
+ "use_sign": False,
+ "use_div_max_abs": True,
+ }
+
+ # Referenced from: https://github.com/WASasquatch/PowerNoiseSuite
+ def generate(self, *_args):
+ noise = self.rand_like()
+
+ modulation = torch.abs(noise) ** self.alpha
+ noise = (torch.sign(noise) if self.use_sign else noise).mul_(modulation)
+ if self.div_max_dims is not None:
+ noise /= torch.amax(
+ torch.abs(noise) if self.use_div_max_abs else noise,
+ keepdim=True,
+ dim=self.div_max_dims,
+ )
+ return noise
+
+
+class LaplacianNoiseGenerator(NoiseGenerator):
+ name = "laplacian"
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {"loc": 0, "scale": 1.0, "div_fac": 4.0}
+
+ def generate(self, *_args):
+ noise = self.rand_like().div_(self.div_fac)
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.
+ Laplace(loc=self.loc, scale=self.scale).rsample(self.shape),
+ noise.device,
+ )
+ return noise
+
+
+class DistroNoiseGenerator(NoiseGenerator):
+ name = "distro"
+
+ simple_distros = frozenset((
+ "cauchy",
+ "exponential",
+ "geometric",
+ "log_normal",
+ "normal",
+ ))
+
+ def __init__(self, x, *args, **kwargs):
+ super().__init__(x, *args, **kwargs)
+ if self.distro not in self.distro_params:
+ raise ValueError("Bad distro")
+
+ @classmethod
+ @property
+ @lru_cache(maxsize=1)
+ def distro_params(cls):
+ td = torch.distributions
+ tt = torch.Tensor
+ return {
+ # Simple
+ "exponential": (
+ tt.exponential_,
+ {
+ "lambd": {
+ "default": 1.0,
+ },
+ },
),
- x.device,
- )
- return scale_noise(noise)
-
-
-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",
- iterations=4,
- generator=None,
-):
- (
- b,
- c,
- h,
- 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, 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 += scale_samples(
- tensor_to(torch.randn(b, c, h, w, generator=generator), x),
- orig_w,
- orig_h,
- mode=upscale_mode,
- ).mul_(discount**i)
- if h >= orig_h * 15 or w >= orig_w * 15:
- break # Lowest resolution is 1x1
- return scale_noise(noise)
-
-
-def pyramid_old_noise_like(
- x,
- *,
- generator=None,
- device="cpu",
- discount=0.8,
- iterations=5,
- upscale_mode="nearest-exact",
-):
- 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(iterations):
- r *= 2
- noise += scale_samples(
- 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,
+ "cauchy": (
+ tt.cauchy_,
+ {
+ "median": {
+ "default": "0.0",
+ },
+ "sigma": {
+ "default": 1.0,
+ "min": 0.0,
+ },
+ },
),
- orig_w,
- orig_h,
- mode=upscale_mode,
- ).mul_(discount**i)
- return tensor_to(noise, x.device)
+ "geometric": (
+ tt.geometric_,
+ {
+ "p": {
+ "default": 0.25,
+ },
+ },
+ ),
+ "log_normal": (
+ tt.log_normal_,
+ {
+ "mean": {
+ "default": 1.0,
+ },
+ "std": {
+ "default": 2.0,
+ },
+ },
+ ),
+ "normal": (
+ tt.normal_,
+ {
+ "mean": {
+ "default": 0.0,
+ },
+ "std": {
+ "default": 1.0,
+ },
+ },
+ ),
+ # Complex distros
+ "beta": (
+ td.Beta,
+ {
+ "concentration0": {
+ "default": "0.5",
+ },
+ "concentration1": {
+ "default": "0.5",
+ },
+ },
+ ),
+ "continuous_bernoulli": (
+ td.ContinuousBernoulli,
+ {
+ "probs": {
+ "default": "0.5",
+ },
+ },
+ ),
+ "dirichlet": (
+ td.Dirichlet,
+ {
+ "concentration": {
+ "default": "0.5 0.5",
+ },
+ },
+ ),
+ "fisher_snedecor": (
+ td.FisherSnedecor,
+ {
+ "df1": {
+ "default": "1.0",
+ },
+ "df2": {
+ "default": "2.0",
+ },
+ },
+ ),
+ "gamma": (
+ td.Gamma,
+ {
+ "concentration": {
+ "default": "1.0",
+ },
+ "rate": {
+ "default": "1.0",
+ },
+ },
+ ),
+ "gumbel": (
+ td.Gumbel,
+ {
+ "loc": {
+ "default": "1.0",
+ },
+ "scale": {
+ "default": "2.0",
+ },
+ },
+ ),
+ "inverse_gamma": (
+ td.InverseGamma,
+ {
+ "concentration": {
+ "default": "1.0",
+ },
+ "rate": {
+ "default": "1.0",
+ },
+ },
+ ),
+ "kumaraswamy": (
+ td.Kumaraswamy,
+ {
+ "concentration0": {
+ "default": "1.0",
+ },
+ "concentration1": {
+ "default": "1.0",
+ },
+ },
+ ),
+ "laplacian": (
+ td.Laplace,
+ {
+ "loc": {
+ "default": "0.0",
+ },
+ "scale": {
+ "default": "1.0",
+ },
+ },
+ ),
+ "lkjcholesky": (
+ td.LKJCholesky,
+ {
+ "dim": {
+ "_ty": "INT",
+ "default": 3,
+ },
+ "concentration": {
+ "default": "1.0",
+ },
+ },
+ ),
+ "lrmvariate_normal": (
+ lambda loc, cov_factor, cov_diag: td.LowRankMultivariateNormal(
+ loc=loc,
+ cov_factor=cov_factor.reshape(loc.numel(), -1),
+ cov_diag=cov_diag,
+ ),
+ {
+ "loc": {
+ "default": "0.0 0.0",
+ },
+ "cov_factor": {
+ "default": "1.0 0.0",
+ },
+ "cov_diag": {
+ "default": "1.0 1.0",
+ },
+ },
+ ),
+ "mvariate_normal": (
+ lambda loc, cov_multiplier=1.0: td.MultivariateNormal(
+ loc=loc,
+ covariance_matrix=torch.eye(
+ loc.numel(),
+ dtype=loc.dtype,
+ device=loc.device,
+ ).mul_(cov_multiplier),
+ ),
+ {
+ "loc": {
+ "default": "0.0 0.0",
+ },
+ "cov_multiplier": {
+ "default": 1.0,
+ },
+ },
+ ),
+ "pareto": (
+ td.Pareto,
+ {
+ "scale": {
+ "default": "1.0",
+ },
+ "alpha": {
+ "default": "1.0",
+ },
+ },
+ ),
+ "poisson": (
+ td.Poisson,
+ {
+ "rate": {
+ "default": "1.5",
+ },
+ },
+ ),
+ "relaxed_bernoulli": (
+ td.RelaxedBernoulli,
+ {
+ "temperature": {
+ "default": 0.75,
+ },
+ "probs": {
+ "default": "0.66",
+ },
+ },
+ ),
+ "relaxed_onehotcategorical": (
+ td.RelaxedOneHotCategorical,
+ {
+ "temperature": {
+ "default": 1.5,
+ },
+ "probs": {
+ "default": "0.33 0.66",
+ },
+ },
+ ),
+ "studentt": (
+ td.StudentT,
+ {
+ "loc": {
+ "default": "0.0",
+ },
+ "scale": {
+ "default": "1.0",
+ },
+ "df": {
+ "default": "1.0",
+ },
+ },
+ ),
+ "uniform": (
+ td.Uniform,
+ {
+ "low": {
+ "default": 0.0,
+ },
+ "high": {
+ "default": 1.0,
+ },
+ },
+ ),
+ "vonmises": (
+ td.VonMises,
+ {
+ "loc": {
+ "default": "1.0",
+ },
+ "concentration": {
+ "default": "1.0",
+ },
+ },
+ ),
+ "weibull": (
+ td.Weibull,
+ {
+ "scale": {
+ "default": "1.0",
+ },
+ "concentration": {
+ "default": "1.0",
+ },
+ },
+ ),
+ "wishart": (
+ lambda df, cov_size=2, cov_multiplier=1.0: td.Wishart(
+ df=df,
+ covariance_matrix=torch.eye(
+ int(cov_size),
+ dtype=df.dtype,
+ device=df.device,
+ ).mul_(cov_multiplier),
+ ),
+ {
+ "df": {
+ "default": "2.0",
+ },
+ "cov_size": {
+ "_ty": "INT",
+ "default": 2,
+ },
+ "cov_multiplier": {
+ "default": 1.0,
+ },
+ },
+ ),
+ }
+ @classmethod
+ @lru_cache(maxsize=1)
+ def build_params(cls):
+ return {
+ f"{tykey}_{pkey}": pval
+ for tykey, tyval in cls.distro_params.items()
+ for pkey, pval in tyval[1].items()
+ if not pkey.startswith("_")
+ }
-# 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",
- 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(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 += scale_samples(
- tensor_to(torch.randn(b, c, w, h), x),
- orig_h,
- orig_w,
- mode=upscale_mode,
- ).mul_(
- discount**i,
+ @classmethod
+ @property
+ @lru_cache(maxsize=1)
+ def ng_params(cls):
+ dparams = {
+ k: v["default"]
+ for k, v in cls.build_params().items()
+ if not k.startswith("_")
+ }
+ return (
+ super().ng_params
+ | {
+ "distro": "normal",
+ "quantile_norm": 0.85,
+ "quantile_norm_flatten": True,
+ "quantile_norm_dim": 1,
+ "quantile_norm_pow": 0.5,
+ "quantile_norm_fac": 1.0,
+ "result_index": "-1",
+ }
+ | dparams
)
- if w == 1 or h == 1:
- break # Lowest resolution is 1x1
- return scale_noise(noise)
-
-def studentt_noise_like(x):
- 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, *, generator=None): # noqa: ARG001
- # The comments said this didn't work and I had to learn the hard way. Turns out it's true!
- height, width = x.shape[-2:]
- noise = torch.randn_like(x)
- scale = 1.0 / (width * height)
- 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 = tensor_to(torch.real(noise), x.device)
- return scale_noise(noise)
-
-
-# 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_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,
+ def norm_output(self, noise):
+ if noise.ndim > len(self.shape):
+ if noise.shape[: len(self.shape)] != self.shape:
+ errstr = f"Unexpected shape when normalizing distro({self.distro}) noise! Output shape={self.shape}, noise shape={noise.shape}, generator dump: {self}"
+ raise RuntimeError(errstr)
+ selfdims = len(self.shape)
+ result_index = self.result_index
+ if not isinstance(result_index, (tuple, list)):
+ result_index = (result_index,)
+ ri_len = len(result_index)
+ if ri_len == 0:
+ raise ValueError("When result_index is a list, it must not be empty")
+ trim_count = 0
+ while noise.ndim > selfdims:
+ idx = result_index[trim_count % ri_len]
+ if idx < 0:
+ idx = noise.shape[-1] + idx
+ noise = noise[..., max(0, min(noise.shape[-1] - 1, idx))]
+ trim_count += 1
+ return (
+ quantile_normalize(
+ noise,
+ quantile=self.quantile_norm,
+ dim=self.quantile_norm_dim,
+ flatten=self.quantile_norm_flatten,
+ nq_fac=self.quantile_norm_fac,
+ pow_fac=self.quantile_norm_pow,
+ )
+ .reshape(self.shape)
+ .contiguous()
)
- return noise
+
+ def distro_param(self, val, *, simple_fun=None):
+ if isinstance(val, torch.Tensor):
+ return simple_fun(val) if simple_fun is not None else val
+ if isinstance(val, str):
+ val = tuple(float(v) for v in val.split(None))
+ if simple_fun is not None:
+ if isinstance(val, (float, int)):
+ return simple_fun(val)
+ if len(val) > 1:
+ raise ValueError("Couldn't return result as float")
+ return simple_fun(val[0])
+ if not isinstance(val, (tuple, list)):
+ val = (val,)
+ return torch.tensor(
+ val,
+ dtype=self.dtype,
+ device=self.gen_device,
+ )
+
+ def get_distro_kwargs(self, distro, ddef, *, simple=False):
+ return {
+ k: self.distro_param(
+ getattr(self, f"{distro}_{k}"),
+ simple_fun=None
+ if not simple and k != "dim"
+ else (int if k == "dim" else float),
+ )
+ for k in ddef
+ }
+
+ def generate(self, *_args):
+ distro = self.distro
+ dfun, ddef = self.distro_params[distro]
+ is_simple = distro in self.simple_distros
+ dkwargs = self.get_distro_kwargs(distro, ddef, simple=is_simple)
+ if is_simple:
+ noise = torch.empty(
+ *self.shape,
+ device=self.gen_device,
+ dtype=self.dtype,
+ layout=self.layout,
+ )
+ noise = dfun(noise, **dkwargs)
+ else:
+ dobj = dfun(**dkwargs)
+ noise = (
+ dobj.rsample if getattr(dobj, "has_rsample", False) else dobj.sample
+ )(self.shape)
+ return self.norm_output(noise)
-def powerlaw_noise_like(x, *, generator=None, **kwargs):
- return tensor_to(
- scale_noise(generate_powerlaw_noise(x, generator=generator, **kwargs)),
- x.device,
- )
+class PowerOldNoiseGenerator(NoiseGenerator):
+ name = "power_old"
+
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {"alpha": 2, "k": 1, "normalized": False}
+
+ def generate(self, *_args):
+ tensor = self.rand_like()
+ fft = torch.fft.fft2(tensor)
+ freq = torch.arange(
+ 1,
+ len(fft) + 1,
+ dtype=tensor.dtype,
+ layout=tensor.layout,
+ device=tensor.device,
+ ).reshape(
+ (len(fft),) + (1,) * (tensor.dim() - 1),
+ )
+ spectral_density = self.k / freq**self.alpha
+ noise = torch.rand(
+ tensor.shape,
+ device=tensor.device,
+ layout=tensor.layout,
+ dtype=tensor.dtype,
+ ).mul_(spectral_density)
+ mean = torch.mean(noise, dim=(-2, -1), keepdim=True)
+ std = torch.std(noise, dim=(-2, -1), keepdim=True)
+ return noise.sub_(mean).div_(std)
-def laplacian_noise_like(x):
- noise = torch.randn_like(x).div_(4.0)
- noise += tensor_to(Laplace(loc=0, scale=1.0).rsample(x.shape), noise.device)
- return scale_noise(noise)
+# Idea from https://github.com/ClownsharkBatwing/RES4LYF/ (wave and mode defaults also from that source)
+class WaveletNoiseGenerator(NoiseGenerator):
+ name = "wavelet"
+ MIN_DIMS = MAX_DIMS = 4
+ def __init__(self, *args, **kwargs):
+ if not HAVE_WAVELETS:
+ raise RuntimeError(
+ "Wavelet noise requires the pytorch_wavelets package installed in your environment",
+ )
+ super().__init__(*args, **kwargs)
+ if self.use_dtcwt:
+ self.wavelet_forward = ptwav.DTCWTForward(
+ J=self.level,
+ mode=self.mode,
+ biort=self.biort,
+ qshift=self.qshift,
+ ).to(self.gen_device)
+ self.wavelet_inverse = ptwav.DTCWTInverse(
+ mode=self.options.get("inv_mode", self.mode),
+ biort=self.options.get("inv_biort", self.biort),
+ qshift=self.options.get("inv_qshift", self.qshift),
+ ).to(
+ self.gen_device,
+ )
+ else:
+ self.wavelet_forward = ptwav.DWTForward(
+ J=self.level,
+ wave=self.wave,
+ mode=self.mode,
+ ).to(self.gen_device)
+ self.wavelet_inverse = ptwav.DWTInverse(
+ wave=self.options.get("inv_wave", self.wave),
+ mode=self.options.get("inv_mode", self.mode),
+ ).to(
+ self.gen_device,
+ )
-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.
+ @classmethod
+ @property
+ def ng_params(cls):
+ return super().ng_params | {
+ "mode": "periodization",
+ "level": 3,
+ "wave": "haar",
+ "use_dtcwt": False,
+ "qshift": "qshift_a",
+ "biort": "near_sym_a",
+ "yl_scale": 1.0,
+ "yh_scales": None,
+ "noise_sampler": None,
+ }
- 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.
- """
- tensor = torch.randn_like(tensor)
- fft = torch.fft.fft2(tensor)
- freq = torch.arange(1, len(fft) + 1, dtype=torch.float).reshape(
- (len(fft),) + (1,) * (tensor.dim() - 1),
- )
- spectral_density = k / freq**alpha
- 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)
+ def generate(self, *args):
+ noise = (
+ self.rand_like()
+ if self.noise_sampler is None
+ else self.noise_sampler(*args)
+ )
+ yl, yh = self.wavelet_forward(noise)
+ if self.yl_scale != 1:
+ yl *= self.yl_scale
+ if self.yh_scales is not None:
+ yh_scales = self.yh_scales
+ if isinstance(yh_scales, (int, float)):
+ yh_scales = (yh_scales,) * len(yh)
+ # print("SCALES", self.yl_scale, yh_scales)
+ for hscale, ht in zip(yh_scales, yh):
+ # print(">> SCALING", hscale)
+ if isinstance(hscale, (int, float)):
+ ht *= hscale # noqa: PLW2901
+ continue
+ for lidx in range(min(ht.shape[2], len(hscale))):
+ # print(">> SCALE IDX", lidx)
+ ht[:, :, lidx, :, :] *= hscale[lidx]
+ return self.wavelet_inverse((yl, yh))
__all__ = (
+ "BrownianNoiseGenerator",
+ "DistroNoiseGenerator",
+ "GaussianNoiseGenerator",
+ "GreenTestNoiseGenerator",
+ "HighresPyramidNoiseGenerator",
+ "LaplacianNoiseGenerator",
+ "MixedNoiseGenerator",
"NoiseError",
"NoiseType",
- "green_noise_like",
- "highres_pyramid_noise_like",
- "laplacian_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",
+ "OneFNoiseGenerator",
+ "PerlinOldNoiseGenerator",
+ "PinkOldNoiseGenerator",
+ "PowerLawNoiseGenerator",
+ "PowerOldNoiseGenerator",
+ "PyramidNoiseGenerator",
+ "PyramidOldNoiseGenerator",
+ "StudentTNoiseGenerator",
+ "UniformNoiseGenerator",
+ "WaveletNoiseGenerator",
)
diff --git a/py/powernoise.py b/py/powernoise.py
index 81c9d38..bfff7bb 100644
--- a/py/powernoise.py
+++ b/py/powernoise.py
@@ -23,7 +23,7 @@ from .nodes import (
SonarNormalizeNoiseNodeMixin,
)
from .noise import CustomNoiseItemBase
-from .noise_generation import scale_noise
+from .utils import scale_noise
# ruff: noqa: ANN003, FBT001, FBT002
diff --git a/py/sonar.py b/py/sonar.py
index c05b7d6..e30bad0 100644
--- a/py/sonar.py
+++ b/py/sonar.py
@@ -4,22 +4,24 @@ from __future__ import annotations
import importlib
from enum import Enum, auto
+from functools import lru_cache
from sys import stderr
from typing import Any, Callable, NamedTuple
import torch
-from comfy.k_diffusion import sampling
+from comfy.k_diffusion.sampling import get_ancestral_step, to_d
from comfy.samplers import KSampler, k_diffusion_sampling
from torch import Tensor
from tqdm.auto import trange
-from . import noise
+from . import noise, utils
class HistoryType(Enum):
ZERO = auto()
RAND = auto()
SAMPLE = auto()
+ SAMPLE_NORM = auto()
class GuidanceType(Enum):
@@ -35,15 +37,34 @@ class GuidanceConfig(NamedTuple):
latent: Tensor | None = None
+class MomentumMode(Enum):
+ CLASSIC = auto()
+ NEW = auto()
+ DENOISED = auto()
+
+
class SonarConfig(NamedTuple):
momentum: float = 0.95
momentum_hist: float = 0.75
direction: float = 1.0
+ momentum_start_step: int = 0
+ momentum_end_step: int = 9999
+ always_update_history: bool = True
+ momentum_mode: MomentumMode = MomentumMode.NEW
init: HistoryType = HistoryType.ZERO
noise_type: noise.NoiseType | None = None
custom_noise: noise.CustomNoise | None = None
rand_init_noise_type: noise.NoiseType | None = None
+ rand_init_noise_multiplier: float | int = 1.0
guidance: GuidanceConfig | None = None
+ blend_mode: str = "lerp"
+ momentum_blend_mode: str | None = None
+ history_blend_mode: str | None = None
+ guidance_blend_mode: str | None = None
+
+ def get_with_default(self, k: str, default: Any) -> Any: # noqa: ANN401
+ val = getattr(self, k)
+ return val if val is not None else default
class SonarBase:
@@ -53,14 +74,69 @@ class SonarBase:
self.history_d = None
self.cfg = cfg
self.noise_sampler = None
+ blend_mode = cfg.blend_mode
+ momentum_blend_mode = cfg.get_with_default("momentum_blend_mode", blend_mode)
+ history_blend_mode = cfg.get_with_default("history_blend_mode", blend_mode)
+ guidance_blend_mode = cfg.get_with_default("guidance_blend_mode", blend_mode)
+ bf = self.blend = utils.BLENDING_MODES[blend_mode]
+ self.momentum_blend = (
+ bf
+ if momentum_blend_mode == blend_mode
+ else utils.BLENDING_MODES[momentum_blend_mode]
+ )
+ self.history_blend = (
+ bf
+ if history_blend_mode == blend_mode
+ else utils.BLENDING_MODES[history_blend_mode]
+ )
+ self.guidance_blend = (
+ bf
+ if guidance_blend_mode == blend_mode
+ else utils.BLENDING_MODES[guidance_blend_mode]
+ )
+
+ _cfg_fixups = (
+ ("momentum_mode", MomentumMode),
+ ("init", HistoryType),
+ ("noise_type", noise.NoiseType),
+ )
+
+ @classmethod
+ def get_config(
+ cls,
+ cfg: SonarConfig | None = None,
+ ext: dict | None = None,
+ ) -> SonarConfig:
+ cfgdict = ext.copy() if ext is not None else {}
+ empty = object()
+ for k, enum_class in cls._cfg_fixups:
+ val = cfgdict.get(k, empty)
+ if val is empty:
+ continue
+ if isinstance(val, str):
+ val = getattr(enum_class, val.strip().upper(), empty)
+ if val is empty:
+ validstr = ", ".join(enum_class.__members__.keys())
+ errstr = f"Bad value for {k} of type enum {enum_class.__name__}, must be one of the following: {validstr}"
+ raise ValueError(errstr)
+ cfgdict[k] = val
+ continue
+ if not isinstance(val, enum_class):
+ errstr = f"Bad parameter type for {k}: Must be valid string or instance of {enum_class.__name__}"
+ raise TypeError(errstr)
+
+ if cfg is None:
+ return SonarConfig(**cfgdict)
+ cfgdict = cfg._asdict() | cfgdict
+ return SonarConfig(**cfgdict)
def set_noise_sampler(
self,
x: Tensor,
- sigmas,
+ sigmas: Tensor,
noise_sampler: Callable | None,
seed: int | None = None,
- ):
+ ) -> Callable:
sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
if noise_sampler is not None and self.cfg.noise_type not in {
None,
@@ -90,17 +166,32 @@ class SonarBase:
self.noise_sampler = noise_sampler
return noise_sampler
- def init_hist_d(self, x: Tensor) -> None:
- if self.history_d is not None:
+ def init_hist_d(
+ self,
+ x: Tensor,
+ denoised: Tensor,
+ sigma: Tensor,
+ *,
+ step: int,
+ ) -> None:
+ if self.history_d is not None or not self.check_step(step, is_history=True):
return
+ cfg = self.cfg
+ init = cfg.init
# memorize delta momentum
- if self.cfg.init == HistoryType.ZERO:
- self.history_d = 0
- elif self.cfg.init == HistoryType.SAMPLE:
- self.history_d = x
- elif self.cfg.init == HistoryType.RAND:
+ if init == HistoryType.ZERO:
+ self.history_d = None
+ elif init == HistoryType.SAMPLE:
+ self.history_d = (
+ x if cfg.momentum_mode != MomentumMode.DENOISED else denoised
+ )
+ elif init == HistoryType.SAMPLE_NORM:
+ self.history_d = (
+ x if cfg.momentum_mode != MomentumMode.DENOISED else denoised
+ ) / sigma
+ elif init == HistoryType.RAND:
ns = noise.get_noise_sampler(
- self.cfg.rand_init_noise_type,
+ cfg.rand_init_noise_type,
x,
None,
None,
@@ -109,31 +200,124 @@ class SonarBase:
normalized=True,
)
self.history_d = ns(None, None)
+ if cfg.rand_init_noise_multiplier != 1:
+ self.history_d *= cfg.rand_init_noise_multiplier
else:
raise ValueError("Sonar sampler: bad history type")
- def update_hist(self, momentum_d):
- q = 1.0 - self.cfg.momentum_hist
+ @property
+ @lru_cache(maxsize=1) # noqa: B019
+ def history_ratios(self):
+ direction = self.cfg.direction
+ momentum_hist = self.cfg.momentum_hist
+ return (
+ momentum_hist,
+ 1.0 + abs(direction) * (1 - momentum_hist)
+ if direction < 0
+ else 2.0 - direction,
+ direction,
+ )
+
+ def check_step(self, step: int, *, is_history: bool = False):
+ cfg = self.cfg
+ if is_history and cfg.always_update_history:
+ return True
+ return cfg.momentum_start_step <= step <= cfg.momentum_end_step
+
+ def update_hist(self, momentum_d: torch.Tensor, step: int) -> None:
+ hd, cfg = self.history_d, self.cfg
+ if cfg.momentum_hist == 1 or not self.check_step(step, is_history=True):
+ return
+ hd_ratio, hd_scale, md_scale = self.history_ratios
+ self.history_d = (
+ momentum_d
+ if hd is None
+ else self.history_blend(momentum_d * md_scale, hd * hd_scale, hd_ratio)
+ )
+
+ def momentum_mix(
+ self,
+ history: Tensor | None,
+ item: Tensor,
+ sigma: Tensor,
+ *,
+ is_denoised: bool = False,
+ momentum=None,
+ ) -> Tensor:
+ momentum = self.cfg.momentum if momentum is None else momentum
+ mode = self.cfg.momentum_mode
+ if (
+ momentum == 1 # noqa: PLR0916
+ or history is None
+ or (mode == MomentumMode.DENOISED and not is_denoised)
+ or (mode != MomentumMode.DENOISED and is_denoised)
+ ):
+ return item
+ return self.momentum_blend(
+ history * sigma if is_denoised else history,
+ item,
+ momentum,
+ )
+
+ def get_momentum_denoised(
+ self,
+ x: Tensor,
+ denoised: Tensor,
+ sigma: Tensor,
+ *,
+ step: int,
+ momentum: float | None = None,
+ update_history=True,
+ ) -> Tensor:
hd = self.history_d
- if isinstance(hd, int) and hd == 0:
- self.history_d = momentum_d
- else:
- self.history_d = (1.0 - q) * hd + q * momentum_d
+ momentum_denoised = self.momentum_mix(
+ hd,
+ denoised,
+ sigma,
+ is_denoised=True,
+ momentum=momentum,
+ )
+ if update_history:
+ self.init_hist_d(x, denoised, sigma, step=step)
+ self.update_hist(denoised / sigma, step=step)
+ return momentum_denoised if self.check_step(step) else denoised
- def momentum_step(self, x: Tensor, d: Tensor, dt: Tensor):
- if self.cfg.momentum == 1.0:
- return x + d * dt
+ def get_momentum_d(
+ self,
+ x: Tensor,
+ denoised: Tensor,
+ sigma: Tensor,
+ *,
+ step: int,
+ momentum: float | None = None,
+ d: Tensor | None = None,
+ update_history=True,
+ ) -> Tensor:
hd = self.history_d
- # correct current `d` with momentum
- p = (1.0 - self.cfg.momentum) * self.cfg.direction
- momentum_d = (1.0 - p) * d + p * hd
+ cfg = self.cfg
+ momentum = cfg.momentum if momentum is None else momentum
+ mode = cfg.momentum_mode
+ d = to_d(x, sigma, denoised) if d is None else d
+ if momentum == 1 or mode == MomentumMode.DENOISED:
+ return d
+ momentum_d = self.momentum_mix(hd, d, sigma)
+ if update_history:
+ self.init_hist_d(x, denoised, sigma, step=step)
+ self.update_hist(d if mode == MomentumMode.NEW else momentum_d, step=step)
+ return momentum_d if self.check_step(step) else d
- # Euler method with momentum
- x = x + momentum_d * dt # noqa: PLR6104
-
- self.update_hist(momentum_d)
-
- return x
+ def momentum_step(
+ self,
+ step: int,
+ x: Tensor,
+ denoised: Tensor,
+ sigma: Tensor,
+ sigma_down: Tensor,
+ ) -> Tensor:
+ dt = sigma_down - sigma
+ denoised = self.get_momentum_denoised(x, denoised, sigma, step=step)
+ momentum_d = self.get_momentum_d(x, denoised, sigma, step=step)
+ return (momentum_d * dt).add_(x)
class SonarGuidanceMixin:
@@ -152,17 +336,26 @@ class SonarGuidanceMixin:
def prepare_ref_latent(latent: Tensor | None) -> Tensor:
if latent is None:
return None
- avg_s = latent.mean(dim=[2, 3], keepdim=True)
- std_s = latent.std(dim=[2, 3], keepdim=True)
- return ((latent - avg_s) / std_s).to(latent.dtype)
+ avg_s = latent.mean(dim=(-2, -1), keepdim=True)
+ std_s = latent.std(dim=(-2, -1), keepdim=True)
+ return (latent - avg_s).div_(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 <= self.guidance.end_step:
+ def guidance_step(self, step_index: int, x: Tensor, denoised: Tensor) -> Tensor:
+ 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)
if self.guidance.guidance_type == GuidanceType.LINEAR:
- return self.guidance_linear(x, self.ref_latent, self.guidance.factor)
+ return self.guidance_linear(
+ x,
+ self.ref_latent,
+ self.guidance.factor,
+ blend=self.guidance_blend,
+ )
if self.guidance.guidance_type == GuidanceType.EULER:
sigma, sigma_next = self.sigmas[step_index], self.sigmas[step_index + 1]
return self.guidance_euler(
@@ -184,20 +377,26 @@ class SonarGuidanceMixin:
ref_latent: Tensor,
factor: float = 0.2,
) -> Tensor:
- avg_t = denoised.mean(dim=[1, 2, 3], keepdim=True)
- std_t = denoised.std(dim=[1, 2, 3], keepdim=True)
+ avg_t = denoised.mean(dim=(-3, -2, -1), keepdim=True)
+ std_t = denoised.std(dim=(-3, -2, -1), keepdim=True)
ref_img_shift = ref_latent * std_t + avg_t
- d = sampling.to_d(x, sigma, ref_img_shift)
+ d = to_d(x, sigma, ref_img_shift)
dt = (sigma_next - sigma) * factor
- return x + d * dt
+ return (d * dt).add_(x)
@staticmethod
- def guidance_linear(x: Tensor, ref_latent: Tensor, factor: float = 0.2) -> Tensor:
- avg_t = x.mean(dim=[1, 2, 3], keepdim=True)
- std_t = x.std(dim=[1, 2, 3], keepdim=True)
- ref_img_shift = ref_latent * std_t + avg_t
- return (1.0 - factor) * x + factor * ref_img_shift
+ def guidance_linear(
+ x: Tensor,
+ ref_latent: Tensor,
+ factor: float = 0.2,
+ *,
+ blend=torch.lerp,
+ ) -> Tensor:
+ avg_t = x.mean(dim=(-3, -2, -1), keepdim=True)
+ std_t = x.std(dim=(-3, -2, -1), keepdim=True)
+ ref_img_shift = (ref_latent * std_t).add_(avg_t)
+ return blend(x, ref_img_shift, factor)
class SonarWithGuidance(SonarBase, SonarGuidanceMixin):
@@ -210,9 +409,9 @@ class SonarSampler(SonarWithGuidance):
def __init__(
self,
model,
- sigmas,
- s_in,
- extra_args,
+ sigmas: Tensor,
+ s_in: Tensor,
+ extra_args: dict[str, Any],
*args: list[Any],
**kwargs: dict[str, Any],
):
@@ -222,62 +421,49 @@ class SonarSampler(SonarWithGuidance):
self.s_in = s_in
self.extra_args = extra_args
+ def call_model(
+ self,
+ x: Tensor,
+ sigma: Tensor,
+ *args: list[Any],
+ s_in=None,
+ extra_args=None,
+ ) -> Tensor:
+ if s_in is None:
+ s_in = self.s_in
+ extra_args = (
+ self.extra_args if extra_args is None else self.extra_args | extra_args
+ )
+ return self.model(x, sigma * s_in, *args, **extra_args)
+
class SonarEuler(SonarSampler):
def __init__(
self,
- s_churn: float = 0.0,
- s_tmin: float = 0.0,
- s_tmax: float = float("inf"),
- s_noise: float = 1.0,
*args: list[Any],
**kwargs: dict[str, Any],
):
super().__init__(*args, **kwargs)
- self.s_churn = s_churn
- self.s_tmin = s_tmin
- self.s_tmax = s_tmax
- self.s_noise = s_noise
- def step(
- self,
- step_index: int,
- sample: torch.FloatTensor,
- ):
- self.init_hist_d(sample)
+ def step(self, step_index: int, sample: torch.FloatTensor):
+ sigma, sigma_next = self.sigmas[step_index], self.sigmas[step_index + 1]
- 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)
- if self.s_tmin <= sigma <= self.s_tmax
- else 0.0
+ denoised = self.call_model(sample, sigma)
+ result_sample = self.momentum_step(
+ step_index,
+ sample,
+ denoised,
+ sigma,
+ sigma_next,
)
- sigma_hat = sigma * (gamma + 1)
-
- if gamma > 0:
- 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 # noqa: PLR6104
-
- denoised = self.model(sample, sigma_hat * self.s_in, **self.extra_args)
- derivative = sampling.to_d(sample, sigma, denoised)
- dt = self.sigmas[step_index + 1] - sigma_hat
-
- result_sample = self.momentum_step(sample, derivative, dt)
-
- if self.sigmas[step_index + 1] > 0:
+ if sigma_next > 0:
result_sample = self.guidance_step(step_index, result_sample, denoised)
return (
result_sample,
sigma,
- sigma_hat,
+ sigma,
denoised,
)
@@ -286,26 +472,18 @@ class SonarEuler(SonarSampler):
def sampler(
cls,
model,
- x,
- sigmas,
- extra_args=None,
+ x: Tensor,
+ sigmas: Tensor,
+ extra_args: dict | None = None,
callback=None,
- disable=None,
+ disable: bool | None = None, # noqa: FBT001
noise_sampler: Callable | None = None,
- sonar_config=None,
- s_churn=0.0,
- s_tmin=0.0,
- s_tmax=float("inf"),
- s_noise=1.0,
- ):
- if sonar_config is None:
- sonar_config = SonarConfig()
- s_in = x.new_ones([x.shape[0]])
+ sonar_config: SonarConfig | None = None,
+ sonar_params: dict | None = None,
+ ) -> Tensor:
+ sonar_config = cls.get_config(sonar_config, sonar_params)
+ s_in = x.new_ones((x.shape[0],))
sonar = cls(
- s_churn,
- s_tmin,
- s_tmax,
- s_noise,
model,
sigmas,
s_in,
@@ -320,7 +498,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,
)
@@ -329,7 +507,7 @@ class SonarEuler(SonarSampler):
{
"x": x,
"i": i,
- "sigma": sigmas[i],
+ "sigma": sigma,
"sigma_hat": sigma_hat,
"denoised": denoised,
},
@@ -354,31 +532,32 @@ class SonarEulerAncestral(SonarSampler):
step_index: int,
sample: torch.FloatTensor,
):
- 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_from,
- sigma_to,
+ sigma, sigma_next = self.sigmas[step_index], self.sigmas[step_index + 1]
+ sigma_down, sigma_up = get_ancestral_step(
+ sigma,
+ sigma_next,
eta=self.eta,
)
- denoised = self.model(sample, sigma_from * self.s_in, **self.extra_args)
- derivative = sampling.to_d(sample, sigma_from, denoised)
- dt = sigma_down - sigma_from
-
- result_sample = self.momentum_step(sample, derivative, dt)
- if sigma_to > 0:
+ denoised = self.call_model(sample, sigma)
+ result_sample = self.momentum_step(
+ step_index,
+ sample,
+ denoised,
+ sigma,
+ sigma_down,
+ )
+ if sigma_next > 0:
result_sample = self.guidance_step(step_index, result_sample, denoised)
result_sample = ( # noqa: PLR6104
result_sample
- + self.noise_sampler(sigma_from, sigma_to) * self.s_noise * sigma_up
+ + self.noise_sampler(sigma, sigma_next) * (self.s_noise * sigma_up)
)
return (
result_sample,
- sigma_from,
- sigma_from,
+ sigma,
+ sigma,
denoised,
)
@@ -392,14 +571,14 @@ class SonarEulerAncestral(SonarSampler):
extra_args=None,
callback=None,
disable=None,
- sonar_config=None,
+ sonar_config: SonarConfig | None = None,
+ sonar_params: dict | None = None,
eta=1.0,
s_noise=1.0,
noise_sampler: Callable | None = None,
):
- if sonar_config is None:
- sonar_config = SonarConfig()
- s_in = x.new_ones([x.shape[0]])
+ sonar_config = cls.get_config(sonar_config, sonar_params)
+ s_in = x.new_ones((x.shape[0],))
sonar = cls(
eta,
s_noise,
@@ -449,104 +628,134 @@ class SonarDPMPPSDE(SonarSampler):
self.s_noise = s_noise
@staticmethod
- def sigma_fn(t) -> float:
+ def sigma_fn(t: Tensor) -> float:
return t.neg().exp()
@staticmethod
- def t_fn(sigma) -> float:
- return sigma.log.neg()
+ def t_fn(sigma: Tensor) -> float:
+ return sigma.log().neg()
# DPM++ solver algorithm copied from ComfyUI source.
def momentum_step( # noqa: PLR0914
self,
- step_index,
+ step_index: int,
x: Tensor,
denoised: Tensor,
- sigma_from,
- sigma_to,
- sigma_down,
- ):
- if sigma_to == 0:
- derivative = sampling.to_d(x, sigma_from, denoised)
- dt = sigma_down - sigma_from
- return super().momentum_step(x, derivative, dt)
+ sigma: Tensor,
+ sigma_next: Tensor,
+ sigma_down: Tensor,
+ ) -> Tensor:
+ if sigma_next == 0:
+ return super().momentum_step(step_index, x, denoised, sigma, sigma_down)
- def sigma_fn(t):
- return t.neg().exp()
-
- def t_fn(sigma):
- return sigma.log().neg()
-
- hd = self.history_d
- p = (1.0 - self.cfg.momentum) * self.cfg.direction
+ cfg = self.cfg
+ # Halve the momentum proportion if there's history since we will use it twice.
+ adjusted_momentum = (
+ cfg.momentum + (1 - cfg.momentum) / 2
+ if self.history_d is not None
+ else cfg.momentum
+ )
r = 1 / 2
# DPM-Solver++
- t, t_next = t_fn(sigma_from), t_fn(sigma_to)
+ t, t_next = self.t_fn(sigma), self.t_fn(sigma_next)
h = t_next - t
s = t + h * r
fac = 1 / (2 * r)
# Step 1
- sd, su = sampling.get_ancestral_step(sigma_fn(t), sigma_fn(s), self.eta)
- s_ = t_fn(sd)
- diff_2 = (t - s_).expm1() * denoised
- momentum_d = (1.0 - p) * diff_2 + p * hd
- self.update_hist(momentum_d)
- hd = self.history_d
- x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - momentum_d
- 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
- sd, su = sampling.get_ancestral_step(
- sigma_fn(t),
- sigma_fn(t_next),
+ s_t, s_s = self.sigma_fn(t), self.sigma_fn(s)
+ sd, su = get_ancestral_step(
+ s_t,
+ s_s,
self.eta,
)
- t_next_ = t_fn(sd)
- denoised_d = (1 - fac) * denoised + fac * denoised_2
- diff_1 = (t - t_next_).expm1() * denoised_d
- momentum_d = (1.0 - p) * diff_1 + p * hd
- self.update_hist(momentum_d)
- x = (sigma_fn(t_next_) / sigma_fn(t)) * x - momentum_d
+ s_ = self.t_fn(sd)
+ momentum_denoised = self.get_momentum_denoised(
+ x,
+ denoised,
+ sigma,
+ step=step_index,
+ )
+ diff_2 = (t - s_).expm1() * momentum_denoised
+ momentum_d = self.get_momentum_d(
+ x,
+ momentum_denoised,
+ sigma,
+ step=step_index,
+ momentum=adjusted_momentum,
+ d=diff_2,
+ )
+ x_2 = ((self.sigma_fn(s_) / s_t) * x).sub_(momentum_d)
+ x_2 += self.noise_sampler(s_t, s_s).mul_(
+ self.s_noise * su,
+ )
+ sigma_2 = s_s
+ denoised_2 = self.call_model(x_2, sigma_2)
+ momentum_denoised_2 = self.get_momentum_denoised(
+ x,
+ denoised_2,
+ sigma_2,
+ step=step_index,
+ )
+
+ # Step 2
+ s_t_next = self.sigma_fn(t_next)
+ sd, su = get_ancestral_step(
+ s_t,
+ s_t_next,
+ self.eta,
+ )
+ t_down = self.t_fn(sd)
+ denoised_d = (1 - fac) * momentum_denoised + fac * momentum_denoised_2
+ diff_1 = (t - t_down).expm1() * denoised_d
+ momentum_d = self.get_momentum_d(
+ x,
+ momentum_denoised_2,
+ sigma_2,
+ step=step_index,
+ momentum=adjusted_momentum,
+ d=diff_1,
+ )
+ x = ((self.sigma_fn(t_down) / s_t) * x).sub_(momentum_d)
x = self.guidance_step(step_index, x, denoised_d)
- return x + self.noise_sampler(sigma_fn(t), sigma_fn(t_next)) * self.s_noise * su
+ x += self.noise_sampler(s_t, s_t_next).mul_(
+ self.s_noise * su,
+ )
+ return x
def step(
self,
step_index: int,
sample: torch.FloatTensor,
- ):
+ ) -> Tensor:
def sigma_fn(t):
return t.neg().exp()
def t_fn(sigma):
return sigma.log().neg()
- 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_from,
- sigma_to,
+ sigma, sigma_next = self.sigmas[step_index], self.sigmas[step_index + 1]
+ sigma_down, _sigma_up = get_ancestral_step(
+ sigma,
+ sigma_next,
eta=self.eta,
)
- denoised = self.model(sample, sigma_from * self.s_in, **self.extra_args)
+ denoised = self.call_model(sample, sigma)
result_sample = self.momentum_step(
step_index,
sample,
denoised,
- sigma_from,
- sigma_to,
+ sigma,
+ sigma_next,
sigma_down,
)
return (
result_sample,
- sigma_from,
- sigma_from,
+ sigma,
+ sigma,
denoised,
)
@@ -555,19 +764,19 @@ class SonarDPMPPSDE(SonarSampler):
def sampler(
cls,
model,
- x,
- sigmas,
- extra_args=None,
+ x: Tensor,
+ sigmas: Tensor,
+ extra_args: dict | None = None,
callback=None,
- disable=None,
- sonar_config=None,
+ disable: bool | None = None, # noqa: FBT001
+ sonar_config: SonarConfig | None = None,
+ sonar_params: dict | None = None,
eta=1.0,
s_noise=1.0,
noise_sampler=None,
- ):
- if sonar_config is None:
- sonar_config = SonarConfig()
- s_in = x.new_ones([x.shape[0]])
+ ) -> Tensor:
+ sonar_config = cls.get_config(sonar_config, sonar_params)
+ s_in = x.new_ones((x.shape[0],))
sonar = cls(
eta,
s_noise,
@@ -602,7 +811,7 @@ class SonarDPMPPSDE(SonarSampler):
return x
-def add_samplers():
+def add_samplers() -> None:
extra_samplers = {
"sonar_euler": SonarEuler.sampler,
"sonar_euler_ancestral": SonarEulerAncestral.sampler,
diff --git a/py/utils.py b/py/utils.py
new file mode 100644
index 0000000..c0cf222
--- /dev/null
+++ b/py/utils.py
@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from comfy.model_management import device_supports_non_blocking
+from comfy.utils import common_upscale
+
+from .external import MODULES as EXT
+
+BLENDING_MODES = {"lerp": torch.lerp}
+UPSCALE_METHODS = (
+ "bilinear",
+ "nearest-exact",
+ "nearest",
+ "area",
+ "bicubic",
+ "bislerp",
+)
+
+
+def scale_samples(
+ samples: torch.Tensor,
+ width: int,
+ height: int,
+ *,
+ mode: str = "bicubic",
+) -> torch.Tensor:
+ return common_upscale(samples, width, height, mode, None)
+
+
+def init_integrations(integrations) -> None:
+ global scale_samples, BLENDING_MODES, UPSCALE_METHODS # noqa: PLW0603
+
+ bleh = integrations.bleh
+ if bleh is None:
+ return
+ bleh_latentutils = bleh.py.latent_utils
+ BLENDING_MODES = bleh_latentutils.BLENDING_MODES
+ UPSCALE_METHODS = bleh_latentutils.UPSCALE_METHODS
+ scale_samples = bleh_latentutils.scale_samples
+
+
+EXT.register_init_handler(init_integrations)
+
+
+def scale_noise(
+ noise: torch.Tensor,
+ factor: float = 1.0,
+ *,
+ normalized: bool = True,
+ threshold_std_devs: float = 2.5,
+ normalize_dims: tuple | None = None,
+) -> torch.Tensor:
+ numel = noise.numel()
+ if not normalized or numel == 0:
+ return noise.mul_(factor) if factor != 1 else noise
+ if normalize_dims is not None:
+ std = noise.std(dim=normalize_dims, keepdim=True)
+ noise = noise / std # noqa: PLR6104
+ return noise.sub_(noise.mean(dim=normalize_dims, keepdim=True)).mul_(factor)
+ mean, std = noise.mean().item(), noise.std().item()
+ threshold = threshold_std_devs / math.sqrt(numel)
+ if abs(mean) > threshold:
+ noise -= mean
+ if abs(1.0 - std) > threshold:
+ noise /= std
+ return noise.mul_(factor) if factor != 1 else noise
+
+
+CAN_NONBLOCK = {}
+
+
+def tensor_to(
+ tensor: torch.Tensor,
+ dest: torch.Tensor | torch.Device | str,
+) -> torch.Tensor:
+ 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 quantile_normalize(
+ noise: torch.Tensor,
+ *,
+ quantile: float = 0.75,
+ dim: int | None = 1,
+ flatten: bool = True,
+ nq_fac: float = 1.0,
+ pow_fac: float = 0.5,
+) -> torch.Tensor:
+ if quantile is None or quantile <= 0 or quantile >= 1:
+ return noise
+ orig_shape = noise.shape
+ if isinstance(quantile, (tuple, list)):
+ quantile = torch.tensor(
+ quantile,
+ device=noise.device,
+ dtype=noise.dtype,
+ )
+ qdim = dim
+ if noise.ndim > 1 and flatten:
+ if qdim is not None and qdim >= noise.ndim:
+ qdim = 1 if noise.ndim > 2 else None
+ if qdim is None:
+ flatdim = 0
+ elif qdim in {0, 1}:
+ flatdim = qdim + 1
+ elif qdim in {2, 3}:
+ noise = noise.movedim(qdim, 1)
+ tempshape = noise.shape
+ flatdim = 2
+ else:
+ raise ValueError(
+ "Cannot handling quantile normalization flattening dims > 3",
+ )
+ else:
+ flatdim = None
+ nq = torch.quantile(
+ (noise if flatdim is None else noise.flatten(start_dim=flatdim)).abs(),
+ quantile,
+ dim=-1,
+ )
+ nq_shape = tuple(nq.shape) + (1,) * (noise.ndim - nq.ndim)
+ nq = nq.mul_(nq_fac).reshape(*nq_shape)
+ noise = noise.clamp(-nq, nq)
+ noise = torch.copysign(
+ torch.pow(torch.abs(noise), pow_fac),
+ noise,
+ )
+ if flatdim is not None and qdim in {2, 3}:
+ return (
+ noise.reshape(tempshape).movedim(1, qdim).reshape(orig_shape).contiguous()
+ )
+ return noise
+
+
+def adjust_slice(s: slice, size: int, offset: int) -> slice:
+ if offset == 0:
+ return s
+ # Input slice must have positive start/stop and be in bounds for the object that will be sliced here.
+ start = s.start if s.start is not None else 0
+ stop = s.stop if s.stop is not None else size
+ if offset < 0:
+ adj = min(start, abs(offset))
+ return slice(start - adj, stop - adj)
+ adj = min(size - stop, offset)
+ return slice(start + adj, stop + adj)
+
+
+def crop_samples(
+ tensor: torch.Tensor,
+ width: int,
+ height: int,
+ *,
+ mode="center",
+ offset_width: int = 0,
+ offset_height: int = 0,
+):
+ if tensor.ndim < 3:
+ raise ValueError("Can only handle >= 3 dimensional tensors")
+ th, tw = tensor.shape[-2:]
+ if (tw, th) == (width, height):
+ return tensor
+ if tw < width or th < height:
+ raise ValueError("Can't crop sample smaller than requested width or height")
+ if mode == "center":
+ hmode = wmode = "center"
+ else:
+ hmode, wmode, *splitextra = mode.split("_")
+ if splitextra:
+ raise ValueError("Bad composite mode")
+ if hmode == "top":
+ hslice = slice(0, height)
+ elif hmode == "center":
+ hoffs = (th - height) // 2
+ hslice = slice(hoffs, hoffs + height)
+ elif hmode == "bottom":
+ hslice = slice(th - height, th)
+ else:
+ raise ValueError("Bad height mode in composite mode")
+ if wmode == "left":
+ wslice = slice(0, width)
+ elif wmode == "center":
+ woffs = (tw - width) // 2
+ wslice = slice(woffs, woffs + width)
+ elif wmode == "right":
+ wslice = slice(tw - width, tw)
+ else:
+ raise ValueError("Bad width mode in composite mode")
+ wslice = adjust_slice(wslice, tw, offset_width)
+ hslice = adjust_slice(hslice, th, offset_height)
+ return tensor[..., hslice, wslice]
diff --git a/ruff.toml b/ruff.toml
index 3c32a79..e589ee0 100644
--- a/ruff.toml
+++ b/ruff.toml
@@ -15,6 +15,8 @@ ignore = [
"D102",
"D103",
"D104",
+ "D105",
+ "D106",
"D107",
"D211",
"D213",