nodes/blepping-ComfyUI-sonar
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added sigma override options that can be set to allow Brownian initial noise. See changelog.

Browse Source 
Added experimental Collatz noise type.
Minor internal changes.
This commit is contained in:
blepping
2025-05-28 16:20:41 -06:00
parent 9dedbeb0b0
commit 83460f3b8f
5 changed files with 270 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
changelog.md
+5
View File
@@ -2,6 +2,11 @@
Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top.
## 20250528
* Added `override_sigma`, `override_sigma_next`, `override_sigma_min` and `override_sigma_max` options that can be set in the `SonarCustomNoiseAdv` node YAML options. This enables using noise generators that require a sigma in stuff like initial noise (for example, Brownian). You will need to manually find and set the correct values yourself.
* Added Collatz noise based on the Collatz conjecture. Very experimental, very slow, likely to change and quite possibly just plain bad. But you can try it.
## 20250505
* Added `SonarQuantileFilteredNoise` node.
py/nodes.py
+88
View File
@@ -969,6 +969,7 @@ class SonarResizedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMix
{
"default": 1152,
"min": 16,
"max": 1024 * 1024 * 1024,
"step": 8,
"tooltip": "Note: This should almost always be set to a higher value than the image you're actually sampling.",
},
@@ -978,6 +979,7 @@ class SonarResizedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMix
{
"default": 1152,
"min": 16,
"max": 1024 * 1024 * 1024,
"step": 8,
"tooltip": "Note: This should almost always be set to a higher value than the image you're actually sampling.",
},
@@ -1351,6 +1353,91 @@ class SonarAdvancedPowerLawNoiseNode(SonarCustomNoiseNodeBase):
)
class SonarAdvancedCollatzNoiseNode(SonarCustomNoiseNodeBase):
DESCRIPTION = "Custom noise type that allows specifying parameters for Collatz noise. Very experimental, also very slow."
@classmethod
def INPUT_TYPES(cls):
result = super().INPUT_TYPES()
result["required"] |= {
"adjust_scale": (
"BOOLEAN",
{
"default": True,
},
),
"use_initial": (
"BOOLEAN",
{
"default": True,
},
),
"iteration_sign_flipping": (
"BOOLEAN",
{
"default": False,
},
),
"chain_length": (
"STRING",
{
"default": "1, 2, 3, 4",
"tooltip": "Comma-separated list of chain lengths. Cannot be empty. Iterations will cycle through the list and wrap.",
},
),
"iterations": ("INT", {"default": 500, "min": 1, "max": 10000}),
"rmin": ("FLOAT", {"default": -100.0, "min": -100000.0, "max": 100000.0}),
"rmax": ("FLOAT", {"default": 100.0, "min": -100000.0, "max": 100000.0}),
"flatten": ("BOOLEAN", {"default": False}),
"dims": (
"STRING",
{
"default": "-1, -2",
"tooltip": "Comma-separated list of dimensions. Cannot be empty. May be negative to count from the end of the list. Iterations will cycle through the list and wrap.",
},
),
}
return result
@classmethod
def get_item_class(cls):
return noise.AdvancedCollatzNoise
def go(
self,
*,
factor,
rescale,
adjust_scale,
use_initial,
iteration_sign_flipping,
chain_length,
iterations,
rmin,
rmax,
flatten,
dims,
sonar_custom_noise_opt=None,
):
if rmin > rmax:
rmin, rmax = rmax, rmin
dims = tuple(int(i) for i in dims.split(","))
return super().go(
factor,
rescale=rescale,
sonar_custom_noise_opt=sonar_custom_noise_opt,
adjust_scale=adjust_scale,
use_initial=use_initial,
iteration_sign_flipping=iteration_sign_flipping,
chain_length=tuple(int(i) for i in chain_length.split(",")),
iterations=iterations,
rmin=rmin,
rmax=rmax,
flatten=flatten,
dims=dims,
)
class SonarQuantileFilteredNoiseNode(SonarCustomNoiseNodeBase):
DESCRIPTION = "Custom noise type that allows filtering noise based on the quantile"
@@ -2316,6 +2403,7 @@ NODE_CLASS_MAPPINGS = {
"SonarAdvancedPyramidNoise": SonarAdvancedPyramidNoiseNode,
"SonarAdvanced1fNoise": SonarAdvanced1fNoiseNode,
"SonarAdvancedPowerLawNoise": SonarAdvancedPowerLawNoiseNode,
"SonarAdvancedCollatzNoise": SonarAdvancedCollatzNoiseNode,
"SonarAdvancedDistroNoise": SonarAdvancedDistroNoiseNode,
"SonarCustomNoise": SonarCustomNoiseNode,
"SonarCustomNoiseAdv": SonarCustomNoiseAdvNode,
py/noise.py
+42 -6
View File
@@ -13,7 +13,7 @@ from torch import Tensor
from . import external, utils
from .noise_generation import *
from .sonar import SonarGuidanceMixin
from .utils import crop_samples, quantile_normalize, scale_noise
from .utils import crop_samples, fallback, quantile_normalize, scale_noise
# ruff: noqa: ANN002, ANN003, FBT001
@@ -81,13 +81,20 @@ class CustomNoiseItem(CustomNoiseItemBase):
**kwargs,
):
ns_kwargs = getattr(self, "ns_kwargs", {}).copy()
# print("NS KWARGS", ns_kwargs)
return get_noise_sampler(
override_sigma, override_sigma_next, override_sigma_min, override_sigma_max = (
ns_kwargs.pop(k, None)
for k in (
"override_sigma",
"override_sigma_next",
"override_sigma_min",
"override_sigma_max",
)
)
ns = get_noise_sampler(
self.noise_type,
x,
sigma_min,
sigma_max,
fallback(override_sigma_min, sigma_min),
fallback(override_sigma_max, sigma_max),
seed=ns_kwargs.pop("seed", seed),
cpu=ns_kwargs.pop("cpu", cpu),
factor=self.factor,
@@ -98,6 +105,16 @@ class CustomNoiseItem(CustomNoiseItemBase):
**ns_kwargs,
**kwargs,
)
if override_sigma is None and override_sigma_next is None:
return ns
def noise_sampler(sigma, sigma_next):
return ns(
fallback(override_sigma, sigma),
fallback(override_sigma_next, sigma_next),
)
return noise_sampler
class CustomNoiseChain:
@@ -296,6 +313,24 @@ class AdvancedDistroNoise(AdvancedNoiseBase):
return DistroNoiseGenerator
class AdvancedCollatzNoise(AdvancedNoiseBase):
ns_factory_arg_keys = (
"adjust_scale",
"use_initial",
"iteration_sign_flipping",
"chain_length",
"iterations",
"rmin",
"rmax",
"flatten",
"dims",
)
@property
def ns_factory(self):
return CollatzNoiseGenerator
class CompositeNoise(CustomNoiseItemBase):
def __init__(
self,
@@ -1616,6 +1651,7 @@ NOISE_SAMPLERS: dict[NoiseType, Callable] = {
),
),
),
NoiseType.COLLATZ: NoiseSampler.wrap(CollatzNoiseGenerator),
}
py/noise_generation.py
+117 -2
View File
@@ -18,13 +18,20 @@ except ImportError:
HAVE_WAVELETS = False
from . import utils
from .utils import quantile_normalize, scale_noise, tensor_to
from .utils import (
fallback,
normalize_to_scale,
quantile_normalize,
scale_noise,
tensor_to,
)
# ruff: noqa: D413, D417, D212, ANN002, ANN003
class NoiseType(Enum):
BROWNIAN = auto()
COLLATZ = auto()
DISTRO = auto()
GAUSSIAN = auto()
GREEN_TEST = auto()
@@ -1356,11 +1363,119 @@ class WaveletNoiseGenerator(FramesToChannelsNoiseGenerator):
for lidx in range(min(ht.shape[2], len(hscale))):
# print(">> SCALE IDX", lidx)
ht[:, :, lidx, :, :] *= hscale[lidx]
return self.fix_output_frames(self.wavelet_inverse((yl, yh)))
result = self.fix_output_frames(self.wavelet_inverse((yl, yh)))
if result.shape == noise.shape:
return result
return result[tuple(slice(0, dl) for dl in noise.shape)]
class CollatzNoiseGenerator(NoiseGenerator):
name = "collatz"
@classmethod
def ng_params(cls):
return super().ng_params() | {
"adjust_scale": True,
"use_initial": True,
"iteration_sign_flipping": False,
"chain_length": (1, 2, 3, 4),
"iterations": 500,
"rmin": -100.0,
"rmax": 100.0,
"flatten": False,
"dims": (-1, -2),
}
@staticmethod
def _get_iter_slices(n_dims, dim, offset, stride) -> tuple:
return tuple(
slice(None) if didx != dim else slice(offset, None, stride)
for didx in range(n_dims)
)
def _generate_iteration(
self,
*,
dim: int,
chain_length: int,
flatten: False,
shape=None,
):
dtype, device = self.dtype, self.device
out_shape = shape = fallback(shape, self.shape)
if dim >= len(shape):
raise ValueError("Requested dimension out of range")
rmin = self.rmin
rmaxsubmin = self.rmax - self.rmin
if flatten:
shape = torch.Size((*shape[:dim], math.prod(shape[dim:])))
size = shape[dim]
chain_length = min(size, chain_length)
n_chunks = math.ceil(size / chain_length)
result_shape = tuple(
(chain_length * n_chunks) if idx == dim else sz
for idx, sz in enumerate(shape)
)
chunk_shape = tuple(
n_chunks if idx == dim else sz for idx, sz in enumerate(shape)
)
result = torch.zeros(result_shape, dtype=dtype, device=device)
noise = torch.rand(
chunk_shape,
generator=self.generator,
dtype=dtype,
device=self.gen_device,
layout=self.layout,
)
if noise.device != self.device:
noise = tensor_to(noise, self.device)
for chainidx in range(chain_length):
if chainidx == 0 and self.use_initial:
result[self._get_iter_slices(result.ndim, dim, 0, chain_length)] = noise
continue
chunk = (
noise
if chainidx == 0
else result[
self._get_iter_slices(result.ndim, dim, chainidx - 1, chain_length)
]
)
result[self._get_iter_slices(result.ndim, dim, chainidx, chain_length)] = (
torch.where(
((chunk * rmaxsubmin + rmin).trunc() % 2) == 0,
chunk * 0.5,
chunk * 3.0 + chunk.sign(),
)
)
return result[
tuple(slice(None, sz) for sz in (shape if flatten else out_shape))
].reshape(out_shape)
def generate(self, *_args):
out_dims = len(self.shape)
dims = tuple(dim if dim >= 0 else out_dims + dim for dim in self.dims)
n_dims, n_chainlens = len(dims), len(self.chain_length)
if not all(0 <= d < out_dims for d in dims):
raise ValueError("Dimension out of range")
dtype, device = self.dtype, self.device
result = torch.zeros(self.shape, dtype=dtype, device=device)
for iteration in range(self.iterations):
temp = self._generate_iteration(
dim=dims[iteration % n_dims],
chain_length=self.chain_length[iteration % n_chainlens],
flatten=self.flatten,
)
if self.iteration_sign_flipping and (iteration & 1) == 1:
temp.neg_()
result += temp
if self.adjust_scale:
result = normalize_to_scale(result, -1.0, 1.0, dim=1)
return result
__all__ = (
"BrownianNoiseGenerator",
"CollatzNoiseGenerator",
"DistroNoiseGenerator",
"GaussianNoiseGenerator",
"GreenTestNoiseGenerator",
py/utils.py
+18
View File
@@ -83,6 +83,7 @@ def tensor_to(
return tensor.to(dest, non_blocking=non_blocking)
# Initial version based on Studentt distribution normalizatino from https://github.com/Clybius/ComfyUI-Extra-Samplers/
def quantile_normalize(
noise: torch.Tensor,
*,
@@ -138,6 +139,19 @@ def quantile_normalize(
return noise
def normalize_to_scale(latent, target_min, target_max, *, dim=(-3, -2, -1)):
min_val, max_val = (
latent.amin(dim=dim, keepdim=True),
latent.amax(dim=dim, keepdim=True),
)
normalized = (latent - min_val).div_(max_val - min_val)
return (
normalized.mul_(target_max - target_min)
.add_(target_min)
.clamp_(target_min, target_max)
)
def adjust_slice(s: slice, size: int, offset: int) -> slice:
if offset == 0:
return s
@@ -194,3 +208,7 @@ def crop_samples(
wslice = adjust_slice(wslice, tw, offset_width)
hslice = adjust_slice(hslice, th, offset_height)
return tensor[..., hslice, wslice]
def fallback(val, default=None):
return val if val is not None else default