From ca3ee58750f8627a006e863bbc40830580d73042 Mon Sep 17 00:00:00 2001 From: blepping Date: Thu, 30 Jan 2025 06:10:36 -0700 Subject: [PATCH] Internal cleanups and refactoring. Some integration improvements. Bump date in changelog --- changelog.md | 2 +- py/nodes.py | 168 +++++++++++++++++++++++++++------------------------ py/noise.py | 4 ++ 3 files changed, 93 insertions(+), 81 deletions(-) diff --git a/changelog.md b/changelog.md index 0621ed5..94a68bf 100644 --- a/changelog.md +++ b/changelog.md @@ -2,7 +2,7 @@ Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top. -## 20241219 +## 20250130 *Note*: May change seeds. diff --git a/py/nodes.py b/py/nodes.py index e6637ce..b306d0b 100644 --- a/py/nodes.py +++ b/py/nodes.py @@ -2,8 +2,8 @@ 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 @@ -158,7 +158,7 @@ class NoisyLatentLikeNode(metaclass=IntegratedNode): } @classmethod - def go( + def go( # noqa: PLR0914 cls, *, noise_type: str, @@ -182,9 +182,11 @@ class NoisyLatentLikeNode(metaclass=IntegratedNode): 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( @@ -1404,68 +1406,6 @@ class SonarAdvancedDistroNoiseNode(SonarCustomNoiseNodeBase): ) -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 SonarWaveletFilteredNoiseNode( SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin, @@ -1535,6 +1475,68 @@ class SonarWaveletFilteredNoiseNode( ) +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(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",) @@ -2244,6 +2246,8 @@ class SonarBlendFilterNoiseNode( normalize_result, normalize_noise, ): + if bleh is None: + raise RuntimeError("bleh not available") import ast # noqa: PLC0415 ffilter_custom = ffilter_custom.strip() @@ -2321,6 +2325,8 @@ class SonarBlehOpsNoiseNode( rules, normalize, ): + if bleh is None: + raise RuntimeError("bleh not available") return super().go( factor, noise=sonar_custom_noise.clone(), @@ -2384,11 +2390,11 @@ class KRestartSamplerCustomNoise(metaclass=IntegratedNode): RETURN_TYPES = ("LATENT", "LATENT") RETURN_NAMES = ("output", "denoised_output") - FUNCTION = "sample" + FUNCTION = "go" CATEGORY = "sampling" @classmethod - def sample( + def go( cls, *, model, @@ -2409,6 +2415,8 @@ class KRestartSamplerCustomNoise(metaclass=IntegratedNode): chunked_mode=False, custom_noise_opt=None, ): + if restart is None: + raise RuntimeError("Restart not available") return restart.restart_sampling.restart_sampling( model, noise_seed, @@ -2458,6 +2466,8 @@ class RestartSamplerCustomNoise(metaclass=IntegratedNode): @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, @@ -2473,20 +2483,18 @@ class RestartSamplerCustomNoise(metaclass=IntegratedNode): return (restart_sampler,) +NODE_CLASS_MAPPINGS |= { + "KRestartSamplerCustomNoise": KRestartSamplerCustomNoise, + "RestartSamplerCustomNoise": RestartSamplerCustomNoise, + "SonarBlendFilterNoise": SonarBlendFilterNoiseNode, + "SonarBlehOpsNoise": SonarBlehOpsNoiseNode, +} + + def init_integrations(integrations): - global NODE_CLASS_MAPPINGS, restart, bleh # noqa: PLW0603 + global restart, bleh # noqa: PLW0603 restart = integrations.restart - if restart is not None: - NODE_CLASS_MAPPINGS["KRestartSamplerCustomNoise"] = KRestartSamplerCustomNoise - if hasattr(restart.restart_sampling, "RestartSampler"): - NODE_CLASS_MAPPINGS["RestartSamplerCustomNoise"] = RestartSamplerCustomNoise bleh = integrations.bleh - if bleh is None: - return - NODE_CLASS_MAPPINGS |= { - "SonarBlendFilterNoise": SonarBlendFilterNoiseNode, - "SonarBlehOpsNoise": SonarBlehOpsNoiseNode, - } external.MODULES.register_init_handler(init_integrations) diff --git a/py/noise.py b/py/noise.py index 80bb949..54a2180 100644 --- a/py/noise.py +++ b/py/noise.py @@ -489,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)