diff --git a/changelog.md b/changelog.md index fc0a7e2..8b4a83b 100644 --- a/changelog.md +++ b/changelog.md @@ -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. diff --git a/py/nodes.py b/py/nodes.py index 50b3a69..50a73cf 100644 --- a/py/nodes.py +++ b/py/nodes.py @@ -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, diff --git a/py/noise.py b/py/noise.py index 2acda40..10f292c 100644 --- a/py/noise.py +++ b/py/noise.py @@ -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), } diff --git a/py/noise_generation.py b/py/noise_generation.py index b802c36..136c080 100644 --- a/py/noise_generation.py +++ b/py/noise_generation.py @@ -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", diff --git a/py/utils.py b/py/utils.py index c0cf222..5ab6cb7 100644 --- a/py/utils.py +++ b/py/utils.py @@ -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