diff --git a/changelog.md b/changelog.md index b578480..38d12d4 100644 --- a/changelog.md +++ b/changelog.md @@ -2,6 +2,12 @@ Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top. +## 20240823 + +* Added descriptions and tooltips for most nodes. +* Added `repeat_batch` parameter to `NoisyLatentLike` node. +* Added a `SONAR_CUSTOM_NOISE to NOISE` node to allow converting from Sonar's custom noise type to the built in ComfyUI `NOISE` (used by `SamplerCustomAdvanced` and possibly other nodes). + ## 20240521 Mega update! Many new features, documentation reorganized. diff --git a/docs/advanced_noise_nodes.md b/docs/advanced_noise_nodes.md index a39abce..7e5f28a 100644 --- a/docs/advanced_noise_nodes.md +++ b/docs/advanced_noise_nodes.md @@ -74,6 +74,12 @@ This node can be used to override configuration settings for other samplers, inc *** +### `SONAR_CUSTOM_NOISE to NOISE` + +This node can be used to convert Sonar custom noise to the `NOISE` type used by the builtin `SamplerCustomAdvanced` (and any other nodes that take a `NOISE` input). + +*** + ### `SonarModulatedNoise` Experimental noise modulation based on code stolen from diff --git a/py/freeu_extreme.py b/py/freeu_extreme.py index 1b643c3..d9b401d 100644 --- a/py/freeu_extreme.py +++ b/py/freeu_extreme.py @@ -36,6 +36,7 @@ BLEND_OPS = ( class FreeUExtremeConfigNode: + DESCRIPTION = "Allows setting configuration for FreeU Extreme." RETURN_TYPES = ("FRUX_CONFIG",) FUNCTION = "go" CATEGORY = "model_patches" @@ -44,10 +45,33 @@ class FreeUExtremeConfigNode: def INPUT_TYPES(cls): return { "required": { - "stage_1": ("BOOLEAN", {"default": True}), - "stage_2": ("BOOLEAN", {"default": False}), - "stage_3": ("BOOLEAN", {"default": False}), - "target": (("backbone", "skip", "both"),), + "stage_1": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Controls whether this configuration applies to stage 1.", + }, + ), + "stage_2": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Controls whether this configuration applies to stage 2.", + }, + ), + "stage_3": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Controls whether this configuration applies to stage 3.", + }, + ), + "target": ( + ("backbone", "skip", "both"), + { + "tooltip": "Controls whether this filter applies to backbone or skip layers (or both).", + }, + ), "start": ( "FLOAT", { @@ -56,6 +80,7 @@ class FreeUExtremeConfigNode: "max": 1.0, "step": 0.1, "round": False, + "tooltip": "Start time as percentage of sampling this configuration applies to. Inclusive.", }, ), "end": ( @@ -66,6 +91,7 @@ class FreeUExtremeConfigNode: "max": 1.0, "step": 0.1, "round": False, + "tooltip": "End time as percentage of sampling this configuration applies to. Inclusive.", }, ), "slice": ( @@ -76,6 +102,7 @@ class FreeUExtremeConfigNode: "max": 1.0, "step": 0.1, "round": False, + "tooltip": "Percentage of the layer the FreeU effect is applied to.", }, ), "slice_offset": ( @@ -86,6 +113,7 @@ class FreeUExtremeConfigNode: "max": 1.0, "step": 0.1, "round": False, + "tooltip": "Offset as a percentage the layer is applied to. For example if slice is 0.25 and slice_offset is 0.25 then the filter will apply to the range 25% through 50%.", }, ), "filter_norm": ( @@ -96,6 +124,7 @@ class FreeUExtremeConfigNode: "max": 10.0, "step": 0.1, "round": False, + "tooltip": "Normalization factor applied to the filter. 1.0 means 100% normalized.", }, ), "scale": ( @@ -106,6 +135,7 @@ class FreeUExtremeConfigNode: "max": 100.0, "step": 0.1, "round": False, + "tooltip": "Strength of the effects applied by this configuration.", }, ), "blend": ( @@ -116,19 +146,48 @@ class FreeUExtremeConfigNode: "max": 10.0, "step": 0.1, "round": False, + "tooltip": "Blends the filtered result based on the specified strength where 1.0 means 100% filtered.", + }, + ), + "blend_mode": ( + tuple(BLEND_OPS.keys()), + { + "tooltip": "Mode used when blending. Generally only has an effect when blend is set to values other than 0 or 1", + }, + ), + "hidden_mean": ( + "BOOLEAN", + { + "default": True, + "tooltip": "You can think of this as FreeU V2 mode.", + }, + ), + "final": ( + "BOOLEAN", + { + "default": True, + "tooltip": "When enabled, other configurations won't be considered if this one matched. Otherwise, multiple configurations/filter effects can be stacked.", }, ), - "blend_mode": (tuple(BLEND_OPS.keys()),), - "hidden_mean": ("BOOLEAN", {"default": True}), - "final": ("BOOLEAN", {"default": True}), }, "optional": { - "sonar_power_filter_opt": ("SONAR_POWER_FILTER",), - "frux_config_opt": ("FRUX_CONFIG",), + "sonar_power_filter_opt": ( + "SONAR_POWER_FILTER", + { + "tooltip": "Optionally attach a Power Filter here to set filtering parameters.", + }, + ), + "frux_config_opt": ( + "FRUX_CONFIG", + { + "tooltip": "Optionally attach another configuration node here.", + }, + ), }, } - def go(self, **kwargs: dict): + @classmethod + def go(cls, **kwargs: dict): return (FreeUExtremeConfig(**kwargs),) @@ -223,12 +282,10 @@ class FreeUExtremeConfig: return False if not getattr(self, f"stage_{stage}"): return False - if self.target not in ("skip" if is_skip else "backbone", "both"): - return False - return True + return not self.target not in {"skip" if is_skip else "backbone", "both"} def apply(self, idx, x, filter_cache, cpu_fft=False): - batch, features, height, width = x.shape + _batch, features, _height, _width = x.shape scale = self.get_scale(x) slice_size = int(features * self.slice) slice_offs = int(features * self.slice_offset) @@ -280,6 +337,7 @@ class FreeUExtremeConfig: class FreeUExtremeNode: + DESCRIPTION = "Main FreeU Extreme node. Allows patching a model with the FreeU (V2) effect with more control." RETURN_TYPES = ("MODEL",) FUNCTION = "go" CATEGORY = "model_patches" @@ -288,18 +346,45 @@ class FreeUExtremeNode: def INPUT_TYPES(cls): return { "required": { - "model": ("MODEL",), - "cpu_fft": ("BOOLEAN", {"default": False}), + "model": ( + "MODEL", + { + "tooltip": "Model to patch.", + }, + ), + "cpu_fft": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Controls whether to perform FFT calculations on the CPU. May be necessary for some GPUs that don't have native support for FFT operations at the cost of performance.", + }, + ), }, "optional": { - "input_config": ("FRUX_CONFIG",), - "middle_config": ("FRUX_CONFIG",), - "output_config": ("FRUX_CONFIG",), + "input_config": ( + "FRUX_CONFIG", + { + "tooltip": "Allows specifying configuration for input blocks.", + }, + ), + "middle_config": ( + "FRUX_CONFIG", + { + "tooltip": "Allows specifying configuration for middle blocks.", + }, + ), + "output_config": ( + "FRUX_CONFIG", + { + "tooltip": "Allows specifying configuration for output blocks.", + }, + ), }, } + @classmethod def go( - self, + cls, model, cpu_fft, input_config=None, diff --git a/py/nodes.py b/py/nodes.py index d7c20a5..49c0f40 100644 --- a/py/nodes.py +++ b/py/nodes.py @@ -2,9 +2,11 @@ from __future__ import annotations import abc import inspect +import random from types import SimpleNamespace from typing import Any, Callable +import numpy as np import torch from comfy import samplers @@ -19,39 +21,98 @@ from .sonar import ( SonarDPMPPSDE, SonarEuler, SonarEulerAncestral, + SonarGuidanceMixin, ) class NoisyLatentLikeNode: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "noise_type": (tuple(NoiseType.get_names(skip=(NoiseType.BROWNIAN,))),), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), - "latent": ("LATENT",), - "multiplier": ("FLOAT", {"default": 1.0}), - "add_to_latent": ("BOOLEAN", {"default": False}), - }, - "optional": { - "custom_noise_opt": ("SONAR_CUSTOM_NOISE",), - "mul_by_sigmas_opt": ("SIGMAS",), - "model_opt": ("MODEL",), - }, - } - + 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.",) CATEGORY = "latent/noise" FUNCTION = "go" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "noise_type": ( + tuple(NoiseType.get_names()), + { + "default": "gaussian", + "tooltip": "Sets the type of noise to generate. Has no effect when the custom_noise_opt input is connected.", + }, + ), + "seed": ( + "INT", + { + "default": 0, + "min": 0, + "max": 0xFFFFFFFFFFFFFFFF, + "tooltip": "Seed to use for generated noise.", + }, + ), + "latent": ( + "LATENT", + { + "tooltip": "Latent used as a reference for generating noise.", + }, + ), + "multiplier": ( + "FLOAT", + { + "default": 1.0, + "tooltip": "Multiplier for the strength of the generated noise. Performed after mul_by_sigmas_opt.", + }, + ), + "add_to_latent": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Add the generated noise to the reference latent rather than adding it to an empty latent. Generally should be enabled for img2img workflows.", + }, + ), + "repeat_batch": ( + "INT", + { + "default": 1, + "tooltip": "Repeats the noise generation the specified number of times. For example, if set to two and your reference latent is also batch two you will get a batch of four as output.", + }, + ), + }, + "optional": { + "custom_noise_opt": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Allows connecting a custom noise chain. When connected, noise_type has no effect.", + }, + ), + "mul_by_sigmas_opt": ( + "SIGMAS", + { + "tooltip": "When connected, will scale the generated noise by the first sigma. Must also connect model_opt to enable.", + }, + ), + "model_opt": ( + "MODEL", + { + "tooltip": "Used when mul_by_sigmas_opt is connected, no effect otherwise.", + }, + ), + }, + } + + @classmethod def go( - self, + cls, + *, noise_type: str, seed: None | int, latent: dict, multiplier: float = 1.0, add_to_latent=False, + repeat_batch=1, custom_noise_opt: object | None = None, mul_by_sigmas_opt: None | torch.Tensor = None, model_opt: object | None = None, @@ -99,17 +160,24 @@ class NoisyLatentLikeNode: randst = torch.random.get_rng_state() try: torch.random.manual_seed(seed) - result = ns(sigma, sigma_next) + result = torch.cat( + tuple(ns(sigma, sigma_next) for _ in range(repeat_batch)), + dim=0, + ) finally: torch.random.set_rng_state(randst) result = scale_noise(result, multiplier, normalized=True) if add_to_latent: - result += latent_samples.to(result) + result += latent_samples.repeat( + *(repeat_batch if i == 0 else 1 for i in range(latent_samples.ndim)), + ).to(result) return ({"samples": result},) class SonarCustomNoiseNodeBase(abc.ABC): + DESCRIPTION = "A custom noise item." RETURN_TYPES = ("SONAR_CUSTOM_NOISE",) + OUTPUT_TOOLTIPS = ("A custom noise chain.",) CATEGORY = "advanced/noise" FUNCTION = "go" @@ -129,6 +197,7 @@ class SonarCustomNoiseNodeBase(abc.ABC): "max": 100.0, "step": 0.001, "round": False, + "tooltip": "Scaling factor for the generated noise of this type.", }, ), }, @@ -144,12 +213,18 @@ class SonarCustomNoiseNodeBase(abc.ABC): "max": 100.0, "step": 0.001, "round": False, + "tooltip": "When non-zero, this custom noise item and other custom noise items items connected to it will have their factor scaled to add up to the specified rescale value. When set to 0, rescaling is disabled.", }, ), } if include_chain: result["optional"] |= { - "sonar_custom_noise_opt": ("SONAR_CUSTOM_NOISE",), + "sonar_custom_noise_opt": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Optional input for more custom noise items.", + }, + ), } return result @@ -175,11 +250,17 @@ class SonarCustomNoiseNode(SonarCustomNoiseNodeBase): def INPUT_TYPES(cls): result = super().INPUT_TYPES() result["required"] |= { - "noise_type": (tuple(NoiseType.get_names()),), + "noise_type": ( + tuple(NoiseType.get_names()), + { + "tooltip": "Sets the type of noise to generate.", + }, + ), } return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return noise.CustomNoiseItem @@ -190,11 +271,18 @@ class SonarNormalizeNoiseNodeMixin: class SonarModulatedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin): + DESCRIPTION = "Custom noise type that allows modulating the output of another custom noise generator." + @classmethod def INPUT_TYPES(cls): result = super().INPUT_TYPES(include_rescale=False, include_chain=False) result["required"] |= { - "sonar_custom_noise": ("SONAR_CUSTOM_NOISE",), + "sonar_custom_noise": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Input custom noise to modulate.", + }, + ), "modulation_type": ( ( "intensity", @@ -202,24 +290,58 @@ class SonarModulatedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeM "spectral_signum", "none", ), + { + "tooltip": "Type of modulation to use.", + }, + ), + "dims": ( + "INT", + { + "default": 3, + "min": 1, + "max": 3, + "tooltip": "Dimensions to modulate over. 1 - channels only, 2 - height and width, 3 - both", + }, + ), + "strength": ( + "FLOAT", + { + "default": 2.0, + "min": -100.0, + "max": 100.0, + "tooltip": "Controls the strength of the modulation effect.", + }, + ), + "normalize_result": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the final result is normalized to 1.0 strength.", + }, + ), + "normalize_noise": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the generated noise is normalized to 1.0 strength.", + }, ), - "dims": ("INT", {"default": 3, "min": 1, "max": 3}), - "strength": ("FLOAT", {"default": 2.0, "min": -100.0, "max": 100.0}), - "normalize_result": (("default", "forced", "disabled"),), - "normalize_noise": (("default", "forced", "disabled"),), "normalize_ref": ( "BOOLEAN", - {"default": True}, + { + "default": True, + "tooltip": "Controls whether the reference latent (when present) is normalized to 1.0 strength.", + }, ), } result["optional"] |= {"ref_latent_opt": ("LATENT",)} return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return noise.ModulatedNoise def go( self, + *, factor, sonar_custom_noise, modulation_type, @@ -246,23 +368,58 @@ class SonarModulatedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeM class SonarRepeatedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin): + DESCRIPTION = "Custom noise type that allows caching the output of other custom noise generators." + @classmethod def INPUT_TYPES(cls): result = super().INPUT_TYPES(include_rescale=False, include_chain=False) result["required"] |= { - "sonar_custom_noise": ("SONAR_CUSTOM_NOISE",), - "repeat_length": ("INT", {"default": 8, "min": 1, "max": 100}), - "max_recycle": ("INT", {"default": 1000, "min": 1, "max": 1000}), - "normalize": (("default", "forced", "disabled"),), - "permute": (("enabled", "disabled", "always"),), + "sonar_custom_noise": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Custom noise input for items to repeat. Note: Unlike most other custom noise nodes, this is treated like a list.", + }, + ), + "repeat_length": ( + "INT", + { + "default": 8, + "min": 1, + "max": 100, + "tooltip": "Number of items to cache.", + }, + ), + "max_recycle": ( + "INT", + { + "default": 1000, + "min": 1, + "max": 1000, + "tooltip": "Number of times an individual item will be used before it is replaced with fresh noise.", + }, + ), + "normalize": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the generated noise is normalized to 1.0 strength.", + }, + ), + "permute": ( + ("enabled", "disabled", "always"), + { + "tooltip": "When enabled, recycled noise will be permuted by randomly flipping it, rolling the channels, etc. If set to always, the noise will be permuted the first time it's used as well.", + }, + ), } return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return noise.RepeatedNoise def go( self, + *, factor, sonar_custom_noise, repeat_length, @@ -281,24 +438,66 @@ 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." + @classmethod def INPUT_TYPES(cls): result = super().INPUT_TYPES(include_rescale=False, include_chain=False) result["required"] |= { - "model": ("MODEL",), - "sonar_custom_noise": ("SONAR_CUSTOM_NOISE",), - "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}), - "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}), - "normalize": (("default", "forced", "disabled"),), + "model": ( + "MODEL", + { + "tooltip": "The model input is required to calculate sampling percentages.", + }, + ), + "sonar_custom_noise": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Custom noise to use when start_percent and end_percent matches.", + }, + ), + "start_percent": ( + "FLOAT", + { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "tooltip": "Time the custom noise becomes active. Note: Sampling percentage where 1.0 indicates 100%, not based on steps.", + }, + ), + "end_percent": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "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.", + }, + ), + "normalize": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the generated noise is normalized to 1.0 strength.", + }, + ), + } + result["optional"] |= { + "fallback_sonar_custom_noise": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Optional input for noise to use when outside of the start_percent, end_percent range. NOTE: When not connected, defaults to NO NOISE which is probably not what you want.", + }, + ), } - result["optional"] |= {"fallback_sonar_custom_noise": ("SONAR_CUSTOM_NOISE",)} return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return noise.ScheduledNoise def go( self, + *, model, factor, sonar_custom_noise, @@ -321,24 +520,58 @@ class SonarScheduledNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeM class SonarCompositeNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin): + DESCRIPTION = "Custom noise type that allows compositing two other custom noise generators based on a mask." + @classmethod def INPUT_TYPES(cls): result = super().INPUT_TYPES(include_rescale=False, include_chain=False) result["required"] |= { - "sonar_custom_noise_dst": ("SONAR_CUSTOM_NOISE",), - "sonar_custom_noise_src": ("SONAR_CUSTOM_NOISE",), - "normalize_dst": (("default", "forced", "disabled"),), - "normalize_src": (("default", "forced", "disabled"),), - "normalize_result": (("default", "forced", "disabled"),), - "mask": ("MASK",), + "sonar_custom_noise_dst": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Custom noise input for noise where the mask is not set.", + }, + ), + "sonar_custom_noise_src": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Custom noise input for noise where the mask is set..", + }, + ), + "normalize_dst": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether noise generated for dst is normalized to 1.0 strength.", + }, + ), + "normalize_src": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether noise generated for src is normalized to 1.0 strength.", + }, + ), + "normalize_result": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the final result after composition is normalized to 1.0 strength.", + }, + ), + "mask": ( + "MASK", + { + "tooltip": "Mask to use when compositing noise. Where the mask is 1.0, you will get 100% src, where it is 0.75 you will get 75% src and 25% dst. The mask will be rescaled to match the latent size if necessary.", + }, + ), } return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return noise.CompositeNoise def go( self, + *, factor, sonar_custom_noise_dst, sonar_custom_noise_src, @@ -359,13 +592,30 @@ class SonarCompositeNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeM class SonarGuidedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin): + DESCRIPTION = "Custom noise type that mixes a references with another custom noise generator to guide the generation." + @classmethod def INPUT_TYPES(cls): result = super().INPUT_TYPES(include_rescale=False, include_chain=False) result["required"] |= { - "latent": ("LATENT",), - "sonar_custom_noise": ("SONAR_CUSTOM_NOISE",), - "method": (("euler", "linear"),), + "latent": ( + "LATENT", + { + "tooltip": "Latent to use for guidance.", + }, + ), + "sonar_custom_noise": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Custom noise input to combine with the guidance.", + }, + ), + "method": ( + ("euler", "linear"), + { + "tooltip": "Method to use when calculating guidance. When set to linear, will simply LERP the guidance at the specified strength. When set to Euler, will do a Euler step toward the guidance instead.", + }, + ), "guidance_factor": ( "FLOAT", { @@ -374,22 +624,38 @@ class SonarGuidedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixi "max": 100.0, "step": 0.001, "round": False, + "tooltip": "Strength of the guidance to apply. Generally should be a relatively slow value to avoid overpowering the generation.", + }, + ), + "normalize_noise": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the generated noise is normalized to 1.0 strength.", + }, + ), + "normalize_result": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the final result is normalized to 1.0 strength.", }, ), - "normalize_noise": (("default", "forced", "disabled"),), - "normalize_result": (("default", "forced", "disabled"),), "normalize_ref": ( "BOOLEAN", - {"default": True}, + { + "default": True, + "tooltip": "Controls whether the reference latent (when present) is normalized to 1.0 strength.", + }, ), } return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return noise.GuidedNoise def go( self, + *, factor, latent, sonar_custom_noise, @@ -399,8 +665,6 @@ class SonarGuidedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixi method="euler", guidance_factor=0.5, ): - from .sonar import SonarGuidanceMixin - return super().go( factor, ref_latent=scale_noise( @@ -416,18 +680,39 @@ class SonarGuidedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixi class SonarRandomNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin): + DESCRIPTION = "Custom noise type that randomly selects between other custom noise items connected to it." + @classmethod def INPUT_TYPES(cls): result = super().INPUT_TYPES(include_rescale=False, include_chain=False) result["required"] |= { - "sonar_custom_noise": ("SONAR_CUSTOM_NOISE",), - "mix_count": ("INT", {"default": 1, "min": 1, "max": 100}), - "normalize": (("default", "forced", "disabled"),), + "sonar_custom_noise": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Custom noise input for noise items to randomize. Note: Unlike most other custom noise nodes, this is treated like a list.", + }, + ), + "mix_count": ( + "INT", + { + "default": 1, + "min": 1, + "max": 100, + "tooltip": "Number of items to select each time noise is generated.", + }, + ), + "normalize": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the generated noise is normalized to 1.0 strength.", + }, + ), } return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return noise.RandomNoise def go( @@ -445,7 +730,135 @@ class SonarRandomNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixi ) +class CustomNOISE: + def __init__( + self, + custom_noise, + seed, + *, + cpu_noise=True, + normalize=True, + multiplier=1.0, + ): + self.custom_noise = custom_noise + self.seed = seed + self.cpu_noise = cpu_noise + self.normalize = normalize + self.multiplier = multiplier + + def _sample_noise(self, latent_image, seed): + result = self.custom_noise.make_noise_sampler( + latent_image, + None, + None, + seed=seed, + cpu=self.cpu_noise, + normalized=self.normalize, + )(None, None).to( + device="cpu", + dtype=latent_image.dtype, + ) + if result.layout != latent_image.layout: + if latent_image.layout == torch.sparse_coo: + return result.to_sparse() + errstr = f"Cannot handle latent layout {type(latent_image.layout).__name__}" + raise NotImplementedError(errstr) + return result if self.multiplier == 1.0 else result.mul_(self.multiplier) + + def generate_noise(self, input_latent): + latent_image = input_latent["samples"] + batch_inds = input_latent.get("batch_index") + torch.manual_seed(self.seed) + random.seed(self.seed) + if self.multiplier == 0.0: + return torch.zeros( + latent_image.shape, + dtype=latent_image.dtype, + layout=latent_image.layout, + device="cpu", + ) + if batch_inds is None: + return self._sample_noise(latent_image, self.seed) + unique_inds, inverse_inds = np.unique(batch_inds, return_inverse=True) + result = [] + batch_size = latent_image.shape[0] + for idx in range(unique_inds[-1] + 1): + noise = self._sample_noise( + latent_image[idx % batch_size].unsqueeze(0), + self.seed + idx, + ) + if idx in unique_inds: + result.append(noise) + return torch.cat(tuple(result[i] for i in inverse_inds), axis=0) + + +class SonarToComfyNOISENode: + 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" + FUNCTION = "go" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "custom_noise": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Custom noise type to convert.", + }, + ), + "seed": ( + "INT", + { + "default": 0, + "min": 0, + "max": 0xFFFFFFFFFFFFFFFF, + "tooltip": "Seed to use for generated noise.", + }, + ), + "cpu_noise": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Controls whether noise is generated on CPU or GPU.", + }, + ), + "normalize": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Controls whether generated noise is normalized to 1.0 strength.", + }, + ), + "multiplier": ( + "FLOAT", + { + "default": 1.0, + "step": 0.001, + "round": False, + "tooltip": "Simple multiplier applied to noise after all other scaling and normalization effects. If set to 0, no noise will be generated (same as disabling noise).", + }, + ), + }, + } + + @classmethod + def go(cls, *, custom_noise, seed, cpu_noise=True, normalize=True, multiplier=1.0): + return ( + CustomNOISE( + custom_noise, + seed, + cpu_noise=cpu_noise, + normalize=normalize, + multiplier=multiplier, + ), + ) + + class GuidanceConfigNode: + DESCRIPTION = "Allows specifying extended guidance parameters for Sonar samplers." + @classmethod def INPUT_TYPES(cls): return { @@ -458,12 +871,35 @@ class GuidanceConfigNode: "max": 2.0, "step": 0.001, "round": False, + "tooltip": "Controls the strength of the guidance. You'll generally want to use fairly low values here.", }, ), - "guidance_type": (tuple(t.name.lower() for t in GuidanceType),), - "start_step": ("INT", {"default": 1, "min": 1}), - "end_step": ("INT", {"default": 9999, "min": 1}), - "latent": ("LATENT",), + "guidance_type": ( + tuple(t.name.lower() for t in GuidanceType), + { + "tooltip": "Method to use when calculating guidance. When set to linear, will simply LERP the guidance at the specified strength. When set to Euler, will do a Euler step toward the guidance instead.", + }, + ), + "start_step": ( + "INT", + { + "default": 0, + "min": 0, + "tooltip": "First zero-based step the guidance is active.", + }, + ), + "end_step": ( + "INT", + { + "default": 9999, + "min": 0, + "tooltip": "Last zero-based step the guidance is active.", + }, + ), + "latent": ( + "LATENT", + {"tooltip": "Latent to use as a reference for guidance."}, + ), }, } @@ -472,8 +908,9 @@ class GuidanceConfigNode: FUNCTION = "make_guidance_cfg" + @classmethod def make_guidance_cfg( - self, + cls, guidance_type, factor, start_step, @@ -492,6 +929,8 @@ class GuidanceConfigNode: class SamplerNodeSonarBase: + DESCRIPTION = "Sonar - momentum based sampler node." + @classmethod def INPUT_TYPES(cls): return { @@ -504,6 +943,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.", }, ), "momentum_hist": ( @@ -514,9 +954,15 @@ class SamplerNodeSonarBase: "max": 1.5, "step": 0.01, "round": False, + "tooltip": "Strength of momentum history", + }, + ), + "momentum_init": ( + tuple(t.name for t in HistoryType), + { + "tooltip": "Initial value used for momentum history. ZERO - history starts zeroed out. RAND - History is initialized with a random value. SAMPLE - History is initialized from the latent at the start of sampling.", }, ), - "momentum_init": (tuple(t.name for t in HistoryType),), "direction": ( "FLOAT", { @@ -525,14 +971,23 @@ class SamplerNodeSonarBase: "max": 15.0, "step": 0.01, "round": False, + "tooltip": "Multiplier applied to the result of normal sampling.", }, ), "rand_init_noise_type": ( tuple(NoiseType.get_names(skip=(NoiseType.BROWNIAN,))), + { + "tooltip": "Noise type to use when momentum_init is set to RANDOM.", + }, ), }, "optional": { - "guidance_cfg_opt": ("SONAR_GUIDANCE_CFG",), + "guidance_cfg_opt": ( + "SONAR_GUIDANCE_CFG", + { + "tooltip": "Optional input for extended guidance parameters.", + }, + ), }, } @@ -554,6 +1009,7 @@ class SamplerNodeSonarEuler(SamplerNodeSonarBase): "max": 100.0, "step": 0.01, "round": False, + "tooltip": "Multiplier for noise added during ancestral or SDE sampling.", }, ), }, @@ -565,8 +1021,10 @@ class SamplerNodeSonarEuler(SamplerNodeSonarBase): FUNCTION = "get_sampler" + @classmethod def get_sampler( - self, + cls, + *, momentum, momentum_hist, momentum_init, @@ -608,20 +1066,33 @@ class SamplerNodeSonarEulerAncestral(SamplerNodeSonarEuler): "max": 100.0, "step": 0.01, "round": False, + "tooltip": "Basically controls the ancestralness of the sampler. When set to 0, you will get a non-ancestral (or SDE) sampler.", + }, + ), + "noise_type": ( + tuple(NoiseType.get_names()), + { + "tooltip": "Noise type used during ancestral or SDE sampling. Only used when the custom noise input is not connected.", }, ), - "noise_type": (tuple(NoiseType.get_names()),), }, ) result["optional"].update( { - "custom_noise_opt": ("SONAR_CUSTOM_NOISE",), + "custom_noise_opt": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Optional input for custom noise used during ancestral or SDE sampling. When connected, the built-in noise_type selector is ignored.", + }, + ), }, ) return result + @classmethod def get_sampler( - self, + cls, + *, momentum, momentum_hist, momentum_init, @@ -669,20 +1140,33 @@ class SamplerNodeSonarDPMPPSDE(SamplerNodeSonarEuler): "max": 100.0, "step": 0.01, "round": False, + "tooltip": "Basically controls the ancestralness of the sampler. When set to 0, you will get a non-ancestral (or SDE) sampler.", + }, + ), + "noise_type": ( + tuple(NoiseType.get_names(default=NoiseType.BROWNIAN)), + { + "tooltip": "Noise type used during ancestral or SDE sampling. Only used when the custom noise input is not connected.", }, ), - "noise_type": (tuple(NoiseType.get_names(default=NoiseType.BROWNIAN)),), }, ) result["optional"].update( { - "custom_noise_opt": ("SONAR_CUSTOM_NOISE",), + "custom_noise_opt": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Optional input for custom noise used during ancestral or SDE sampling. When connected, the built-in noise_type selector is ignored.", + }, + ), }, ) return result + @classmethod def get_sampler( - self, + cls, + *, momentum, momentum_hist, momentum_init, @@ -717,6 +1201,7 @@ class SamplerNodeSonarDPMPPSDE(SamplerNodeSonarEuler): class SamplerNodeConfigOverride: + 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 @@ -730,6 +1215,7 @@ class SamplerNodeConfigOverride: "default": 1.0, "step": 0.01, "round": False, + "tooltip": "Basically controls the ancestralness of the sampler. When set to 0, you will get a non-ancestral (or SDE) sampler.", }, ), "s_noise": ( @@ -738,6 +1224,7 @@ class SamplerNodeConfigOverride: "default": 1.0, "step": 0.01, "round": False, + "tooltip": "Multiplier for noise added during ancestral or SDE sampling.", }, ), "s_churn": ( @@ -747,6 +1234,7 @@ class SamplerNodeConfigOverride: "min": 0.0, "step": 0.01, "round": False, + "tooltip": "Churn was the predececessor of ETA. Only used by a few types of samplers (notably Euler non-ancestral). Not used by any ancestral or SDE samplers.", }, ), "r": ( @@ -755,15 +1243,43 @@ class SamplerNodeConfigOverride: "default": 0.5, "step": 0.01, "round": False, + "tooltip": "Used by dpmpp_sde.", + }, + ), + "sde_solver": ( + ("midpoint", "heun"), + { + "tooltip": "Solver used by dpmpp_2m_sde.", + }, + ), + "cpu_noise": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Controls whether noise is generated on CPU or GPU.", + }, + ), + "normalize": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Controls whether generated noise is normalized to 1.0 strength.", }, ), - "sde_solver": (("midpoint", "heun"),), - "cpu_noise": ("BOOLEAN", {"default": True}), - "normalize": ("BOOLEAN", {"default": True}), }, "optional": { - "noise_type": (tuple(NoiseType.get_names()),), - "custom_noise_opt": ("SONAR_CUSTOM_NOISE",), + "noise_type": ( + tuple(NoiseType.get_names()), + { + "tooltip": "Noise type used during ancestral or SDE sampling. Only used when the custom noise input is not connected.", + }, + ), + "custom_noise_opt": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Optional input for custom noise used during ancestral or SDE sampling. When connected, the built-in noise_type selector is ignored.", + }, + ), }, } @@ -774,6 +1290,7 @@ class SamplerNodeConfigOverride: def get_sampler( self, + *, sampler, eta, s_noise, @@ -857,7 +1374,7 @@ class SamplerNodeConfigOverride: ) sig = inspect.signature(sampler.sampler_function) params = sig.parameters - kwargs = kwargs | {} + kwargs = kwargs.copy() if "noise_sampler" in params: kwargs["noise_sampler"] = noise_sampler for k in cls.KWARG_OVERRIDES: @@ -888,12 +1405,15 @@ NODE_CLASS_MAPPINGS = { "SonarScheduledNoise": SonarScheduledNoiseNode, "SonarGuidedNoise": SonarGuidedNoiseNode, "SonarRandomNoise": SonarRandomNoiseNode, + "SONAR_CUSTOM_NOISE to NOISE": SonarToComfyNOISENode, } NODE_DISPLAY_NAME_MAPPINGS = {} if "bleh" in external.MODULES: + import ast + bleh = external.MODULES["bleh"] bleh_latentutils = bleh.py.latent_utils @@ -901,6 +1421,8 @@ if "bleh" in external.MODULES: SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin, ): + DESCRIPTION = "Custom noise type that allows blending and filtering the output of another noise generator." + @classmethod def INPUT_TYPES(cls): result = super().INPUT_TYPES(include_rescale=False, include_chain=False) @@ -934,11 +1456,13 @@ if "bleh" in external.MODULES: } return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return noise.BlendFilterNoise def go( self, + *, factor, sonar_custom_noise, blend_mode, @@ -953,8 +1477,6 @@ if "bleh" in external.MODULES: normalize_result, normalize_noise, ): - import ast - ffilter_custom = ffilter_custom.strip() normalize_result = ( None if normalize_result == "default" else normalize_result == "forced" @@ -987,6 +1509,8 @@ 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( @@ -1032,8 +1556,10 @@ if "restart" in external.MODULES: FUNCTION = "sample" CATEGORY = "sampling" + @classmethod def sample( - self, + cls, + *, model, add_noise, noise_seed, @@ -1079,6 +1605,8 @@ if "restart" in external.MODULES: if hasattr(rs.restart_sampling, "RestartSampler"): 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 { @@ -1095,7 +1623,8 @@ if "restart" in external.MODULES: FUNCTION = "go" CATEGORY = "sampling/custom_sampling/samplers" - def go(self, sampler, chunked_mode, custom_noise_opt=None): + @classmethod + def go(cls, sampler, chunked_mode, custom_noise_opt=None): restart_options = { "restart_chunked": chunked_mode, "restart_wrapped_sampler": sampler, diff --git a/py/noise.py b/py/noise.py index 0b026b6..36ffb6b 100644 --- a/py/noise.py +++ b/py/noise.py @@ -10,6 +10,7 @@ from torch import Tensor from . import external from .noise_generation import * +from .sonar import SonarGuidanceMixin # ruff: noqa: D412, D413, D417, D212, D407, ANN002, ANN003, FBT001, FBT002, S311 @@ -220,7 +221,7 @@ class CompositeNoise(CustomNoiseItemBase): ) def clone_key(self, k): - if k in ("mask", "src_noise", "dst_noise"): + if k in {"mask", "src_noise", "dst_noise"}: return getattr(self, k).clone() return super().clone_key(k) @@ -286,13 +287,11 @@ class GuidedNoise(CustomNoiseItemBase): ) def clone_key(self, k): - if k in ("noise", "ref_latent"): + if k in {"noise", "ref_latent"}: return getattr(self, k).clone() return super().clone_key(k) def make_noise_sampler(self, x, *args, normalized=True, **kwargs): - from .sonar import SonarGuidanceMixin - factor, guidance_factor = self.factor, self.guidance_factor normalize_noise, normalize_result = ( self.get_normalize(f"normalize_{k}", normalized) @@ -325,6 +324,7 @@ class GuidedNoise(CustomNoiseItemBase): factor, normalized=normalize_result, ) + case "euler": def noise_sampler(s, sn): @@ -884,8 +884,8 @@ if "bleh" in external.MODULES: normalized or num_samplers > 1, ) normalize_result = self.get_normalize("normalize_result", normalized) - noise_effects = self.affect in ("noise", "both") - result_effects = self.affect in ("result", "both") + noise_effects = self.affect in {"noise", "both"} + result_effects = self.affect in {"result", "both"} noise_init = torch.zeros_like(x) def noise_sampler(s, sn): diff --git a/py/noise_generation.py b/py/noise_generation.py index 1b3d580..d41c1f3 100644 --- a/py/noise_generation.py +++ b/py/noise_generation.py @@ -133,7 +133,7 @@ def perlin_noise_tensor( step -- smooth step function [0, 1] -> [0, 1] (default: `smooth_step`) Raises: - Exception: if position and vector shapes do not match + NoiseError: if position and vector shapes do not match Returns: (batch_size, block_height * grid_height, block_width * grid_width) @@ -148,11 +148,11 @@ def perlin_noise_tensor( bh, bw = positions.shape[1:3] for i in range(2): - if positions.shape[i + 3] not in (1, vectors.shape[i + 2]): + if positions.shape[i + 3] not in {1, vectors.shape[i + 2]}: msg = f"Blocks shapes do not match: vectors ({vectors.shape[1]}, {vectors.shape[2]}), positions {gh}, {gw})" raise NoiseError(msg) - if positions.shape[0] not in (1, batch_size): + if positions.shape[0] not in {1, batch_size}: msg = f"Batch sizes do not match: vectors ({vectors.shape[0]}), positions ({positions.shape[0]})" raise NoiseError(msg) @@ -206,7 +206,7 @@ def perlin_noise( generator -- random generator used for grid vectors (default: {None}) Raises: - Exception: if grid and out shapes do not match + NoiseError: if grid and out shapes do not match Returns: Noise image shaped (batch_size, height, width) @@ -410,9 +410,8 @@ def power_noise_like(tensor, alpha=2, k=1): # This doesn't work properly right __all__ = ( - "NoiseType", "NoiseError", - "scale_noise", + "NoiseType", "green_noise_like", "highres_pyramid_noise_like", "laplacian_noise_like", @@ -421,6 +420,7 @@ __all__ = ( "pyramid_noise_like", "pyramid_old_noise_like", "rand_perlin_like", + "scale_noise", "studentt_noise_like", "uniform_noise_like", ) diff --git a/py/powernoise.py b/py/powernoise.py index a5eedeb..acb92bc 100644 --- a/py/powernoise.py +++ b/py/powernoise.py @@ -516,7 +516,7 @@ class PowerFilterNoiseItem(PowerNoiseItem): x, ns, self.make_filter(x.shape), - self.normalize_result in (True, None), + self.normalize_result in {True, None}, ) filtered_noise = filtered_ns( torch.scalar_tensor(14.0), @@ -533,11 +533,19 @@ class PowerFilterNoiseItem(PowerNoiseItem): class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): + DESCRIPTION = "Custom noise type that applies a filter to generated noise." + @classmethod def INPUT_TYPES(cls, *args: list, **kwargs: dict): result = super().INPUT_TYPES(*args, **kwargs) result["required"] |= { - "time_brownian": ("BOOLEAN", {"default": False}), + "time_brownian": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Controls whether brownian noise is used when mix isn't 1.0.", + }, + ), "alpha": ( "FLOAT", { @@ -546,6 +554,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): "max": 5.0, "step": 0.001, "round": False, + "tooltip": "Values above 0 will amplify low frequencies, negative values will amplify high frequencies.", }, ), "max_freq": ( @@ -556,6 +565,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): "max": 0.7071, "step": 0.001, "round": False, + "tooltip": "Maximum frequency to pass through the filter.", }, ), "min_freq": ( @@ -566,6 +576,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): "max": 0.7071, "step": 0.001, "round": False, + "tooltip": "Minimum frequency to pass through the filter.", }, ), "stretch": ( @@ -576,6 +587,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): "max": 100, "step": 0.1, "round": False, + "tooltip": "Stretches the filter's shape by the specified factor.", }, ), "rotate": ( @@ -586,6 +598,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): "max": 90, "step": 5, "round": False, + "tooltip": "Rotates the filter.", }, ), "pnorm": ( @@ -596,6 +609,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): "max": 100, "step": 0.1, "round": False, + "tooltip": "Factor used for cushioning the band-pass region.", }, ), "mix": ( @@ -606,6 +620,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): "max": 1.0, "step": 0.001, "round": False, + "tooltip": "Controls the ratio of filtered noise. For example, 0.75 means 75% noise with the filter effects applied, 25% raw noise.", }, ), "common_mode": ( @@ -616,6 +631,7 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): "max": 100.0, "step": 0.001, "round": False, + "tooltip": "Attempts to desaturate thelatent by injecting the average across channels (controlled by channel_correction). Applied after mix.", }, ), "channel_correlation": ( @@ -624,13 +640,20 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): "default": "1, 1, 1, 1, 1, 1", "multiline": False, "dynamicPrompts": False, + "tooltip": "Comma-separated list of channel correlation strengths.", + }, + ), + "preview": ( + ("none", "no_mix", "mix"), + { + "tooltip": "When enabled, displays a preview of the filter shape and a sample of noise. Mix - previews noise after mix is applied. no_mix - only previews the filtered noise.", }, ), - "preview": (("none", "no_mix", "mix"),), } return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return PowerNoiseItem def go( @@ -648,6 +671,8 @@ class SonarPowerNoiseNode(SonarCustomNoiseNodeBase): class SonarPowerFilterNoiseNode(SonarPowerNoiseNode, SonarNormalizeNoiseNodeMixin): + DESCRIPTION = "Custom noise type that allows applying a Power Filter to another custom noise generator." + @classmethod def INPUT_TYPES(cls): result = super().INPUT_TYPES(include_rescale=False, include_chain=False) @@ -662,8 +687,18 @@ class SonarPowerFilterNoiseNode(SonarPowerNoiseNode, SonarNormalizeNoiseNodeMixi ): del result["required"][k] result["required"] |= { - "sonar_custom_noise": ("SONAR_CUSTOM_NOISE",), - "sonar_power_filter": ("SONAR_POWER_FILTER",), + "sonar_custom_noise": ( + "SONAR_CUSTOM_NOISE", + { + "tooltip": "Custom noise type to filter.", + }, + ), + "sonar_power_filter": ( + "SONAR_POWER_FILTER", + { + "tooltip": "Filter to use.", + }, + ), "filter_norm_factor": ( "FLOAT", { @@ -672,15 +707,32 @@ class SonarPowerFilterNoiseNode(SonarPowerNoiseNode, SonarNormalizeNoiseNodeMixi "max": 1.0, "step": 0.1, "round": False, + "tooltip": "Normalization factor applied to the specified filter. 1.0 means 100% normalized.", + }, + ), + "normalize_result": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the final result is normalized to 1.0 strength.", + }, + ), + "normalize_noise": ( + ("default", "forced", "disabled"), + { + "tooltip": "Controls whether the generated noise is normalized to 1.0 strength.", }, ), - "normalize_result": (("default", "forced", "disabled"),), - "normalize_noise": (("default", "forced", "disabled"),), } - result["required"]["preview"] = ((*result["required"]["preview"][0], "custom"),) + result["required"]["preview"] = ( + (*result["required"]["preview"][0], "custom"), + { + "tooltip": "When enabled, displays a preview of the filter shape and a sample of noise. Mix - previews noise after mix is applied. no_mix - only previews the filtered noise. custom - Like no_mix, but will use a latent previewer to display a color preview of the generated noise. Works best when previewer is set to TAESD.", + }, + ) return result - def get_item_class(self): + @classmethod + def get_item_class(cls): return PowerFilterNoiseItem def go( @@ -714,69 +766,23 @@ class SonarPowerFilterNode: @classmethod def INPUT_TYPES(cls): + include_keys = {"alpha", "max_freq", "min_freq", "stretch", "rotate", "pnorm"} return { "required": { - "alpha": ( - "FLOAT", + k: v + for k, v in SonarPowerNoiseNode.INPUT_TYPES()["required"].items() + if k in include_keys + } + | { + "oversample": ( + "INT", { - "default": 0.0, - "min": -5.0, - "max": 5.0, - "step": 0.001, - "round": False, + "default": 4, + "min": 1, + "max": 128, + "tooltip": "Oversampling factor used for the filter size.", }, ), - "max_freq": ( - "FLOAT", - { - "default": 0.7071, - "min": 0.0, - "max": 0.7071, - "step": 0.001, - "round": False, - }, - ), - "min_freq": ( - "FLOAT", - { - "default": 0.0, - "min": 0.0, - "max": 0.7071, - "step": 0.001, - "round": False, - }, - ), - "stretch": ( - "FLOAT", - { - "default": 1.0, - "min": 0.01, - "max": 100, - "step": 0.1, - "round": False, - }, - ), - "rotate": ( - "FLOAT", - { - "default": 0, - "min": -90, - "max": 90, - "step": 5, - "round": False, - }, - ), - "pnorm": ( - "FLOAT", - { - "default": 2, - "min": 0.125, - "max": 100, - "step": 0.1, - "round": False, - }, - ), - "oversample": ("INT", {"default": 4, "min": 1, "max": 128}), "blur": ( "FLOAT", { @@ -785,6 +791,7 @@ class SonarPowerFilterNode: "max": 10.0, "step": 0.01, "round": False, + "tooltip": "Slightly blurs the filter to reduce artifacts.", }, ), "scale": ( @@ -795,17 +802,24 @@ class SonarPowerFilterNode: "max": 100.0, "step": 0.1, "round": False, + "tooltip": "Scales the filter to the specified strength. May be negative.", + }, + ), + "compose_mode": ( + ("max", "min", "add", "sub", "mul"), + { + "tooltip": "Controls composition of the option attached filter. For example, when set to MUL the result will be this filter multiplied by the attached filter. No effect if the optional filter input is not attached.", }, ), - "compose_mode": (("max", "min", "add", "sub", "mul"),), }, "optional": { "power_filter_opt": ("SONAR_POWER_FILTER",), }, } + @classmethod def go( - self, + cls, min_freq=0.0, max_freq=0.7071, stretch=1.0, @@ -836,6 +850,7 @@ class SonarPowerFilterNode: class SonarPreviewFilterNode: + DESCRIPTION = "Allows previewing a Power Filter." RETURN_TYPES = ("SONAR_POWER_FILTER",) CATEGORY = "advanced/noise" FUNCTION = "go" @@ -845,7 +860,12 @@ class SonarPreviewFilterNode: def INPUT_TYPES(cls): return { "required": { - "sonar_power_filter": ("SONAR_POWER_FILTER",), + "sonar_power_filter": ( + "SONAR_POWER_FILTER", + { + "tooltip": "Power Filter to preview.", + }, + ), "filter_gain": ( "FLOAT", { @@ -854,6 +874,7 @@ class SonarPreviewFilterNode: "max": 1000000.0, "step": 0.1, "round": False, + "tooltip": "Gain factor applied to the filter part of the preview.", }, ), "kernel_gain": ( @@ -864,6 +885,7 @@ class SonarPreviewFilterNode: "max": 1000000.0, "step": 0.1, "round": False, + "tooltip": "Gain factor applied to the kernel part of the preview.", }, ), "norm_factor": ( @@ -874,6 +896,7 @@ class SonarPreviewFilterNode: "max": 1.0, "step": 0.1, "round": False, + "tooltip": "Normalization factor applied to the filter before previewing. 1.0 means 100% normalized.", }, ), "preview_size": ( @@ -888,12 +911,16 @@ class SonarPreviewFilterNode: "128x127", "127x128", ), + { + "tooltip": "Controls the size of the generated preview. Note: Sizes are in latent pixels. For most models, one latent pixel equals eight pixels", + }, ), }, } + @classmethod def go( - self, + cls, sonar_power_filter, filter_gain=1 / 3, kernel_gain=1 / 3, diff --git a/py/sonar.py b/py/sonar.py index ac6cc85..c1a1631 100644 --- a/py/sonar.py +++ b/py/sonar.py @@ -2,12 +2,14 @@ from __future__ import annotations +import importlib from enum import Enum, auto from sys import stderr from typing import Any, Callable, NamedTuple import torch from comfy.k_diffusion import sampling +from comfy.samplers import KSampler, k_diffusion_sampling from torch import Tensor from tqdm.auto import trange @@ -60,10 +62,10 @@ class SonarBase: seed: int | None = None, ): sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() - if noise_sampler is not None and self.cfg.noise_type not in ( + if noise_sampler is not None and self.cfg.noise_type not in { None, self.DEFAULT_NOISE_TYPE, - ): + }: print( "Sonar: Warning: Noise sampler supplied, overriding noise type from settings", file=stderr, @@ -127,7 +129,7 @@ class SonarBase: momentum_d = (1.0 - p) * d + p * hd # Euler method with momentum - x = x + momentum_d * dt + x = x + momentum_d * dt # noqa: PLR6104 self.update_hist(momentum_d) @@ -155,9 +157,8 @@ class SonarGuidanceMixin: return ((latent - avg_s) / std_s).to(latent.dtype) def guidance_step(self, step_index: int, x: Tensor, denoised: Tensor): - if (self.guidance is None or self.guidance.factor == 0.0) or not ( - self.guidance.start_step <= (step_index + 1) <= self.guidance.end_step - ): + step_matched = self.guidance.start_step <= step_index <= self.guidance.end_step + if self.guidance is None or self.guidance.factor == 0.0 or not step_matched: return x if self.ref_latent.device != x.device: self.ref_latent = self.ref_latent.to(device=x.device) @@ -263,7 +264,7 @@ class SonarEuler(SonarSampler): else torch.randn_like(sample) ) eps = noise * self.s_noise - sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5 + sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5 # noqa: PLR6104 denoised = self.model(sample, sigma_hat * self.s_in, **self.extra_args) derivative = sampling.to_d(sample, sigma, denoised) @@ -320,7 +321,7 @@ class SonarEuler(SonarSampler): ) for i in trange(len(sigmas) - 1, disable=disable): - x, sigma, sigma_hat, denoised = sonar.step( + x, _sigma, sigma_hat, denoised = sonar.step( i, x, ) @@ -370,7 +371,7 @@ class SonarEulerAncestral(SonarSampler): result_sample = self.momentum_step(sample, derivative, dt) if sigma_to > 0: result_sample = self.guidance_step(step_index, result_sample, denoised) - result_sample = ( + result_sample = ( # noqa: PLR6104 result_sample + self.noise_sampler(sigma_from, sigma_to) * self.s_noise * sigma_up ) @@ -417,7 +418,7 @@ class SonarEulerAncestral(SonarSampler): ) for i in trange(len(sigmas) - 1, disable=disable): - x, sigma, sigma_hat, denoised = sonar.step( + x, _sigma, sigma_hat, denoised = sonar.step( i, x, ) @@ -457,7 +458,7 @@ class SonarDPMPPSDE(SonarSampler): return sigma.log.neg() # DPM++ solver algorithm copied from ComfyUI source. - def momentum_step( + def momentum_step( # noqa: PLR0914 self, step_index, x: Tensor, @@ -495,7 +496,7 @@ class SonarDPMPPSDE(SonarSampler): self.update_hist(momentum_d) hd = self.history_d x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - momentum_d - x_2 = x_2 + self.noise_sampler(sigma_fn(t), sigma_fn(s)) * self.s_noise * su + x_2 += self.noise_sampler(sigma_fn(t), sigma_fn(s)) * self.s_noise * su denoised_2 = self.model(x_2, sigma_fn(s) * self.s_in, **self.extra_args) # Step 2 @@ -527,7 +528,7 @@ class SonarDPMPPSDE(SonarSampler): self.init_hist_d(sample) sigma_from, sigma_to = self.sigmas[step_index], self.sigmas[step_index + 1] - sigma_down, sigma_up = sampling.get_ancestral_step( + sigma_down, _sigma_up = sampling.get_ancestral_step( sigma_from, sigma_to, eta=self.eta, @@ -585,7 +586,7 @@ class SonarDPMPPSDE(SonarSampler): ) for i in trange(len(sigmas) - 1, disable=disable): - x, sigma, sigma_hat, denoised = sonar.step( + x, _sigma, sigma_hat, denoised = sonar.step( i, x, ) @@ -603,10 +604,6 @@ class SonarDPMPPSDE(SonarSampler): def add_samplers(): - import importlib - - from comfy.samplers import KSampler, k_diffusion_sampling - extra_samplers = { "sonar_euler": SonarEuler.sampler, "sonar_euler_ancestral": SonarEulerAncestral.sampler, diff --git a/ruff.toml b/ruff.toml index ef82bca..3c32a79 100644 --- a/ruff.toml +++ b/ruff.toml @@ -8,6 +8,8 @@ ignore = [ "ANN204", "ANN206", "C901", + "CPY001", + "DOC201", "D100", "D101", "D102", @@ -26,6 +28,7 @@ ignore = [ "PLR0912", "PLR0913", "PLR0915", + "PLR0917", "PLR2004", "T201", "TRY003",