Noise

README.md

@@ -1,2 +1,36 @@
-# cg-noisetools
- 
+# Noise Tools
+
+## Mix Noise
+
+This node takes two noise inputs and produces a weighted mix, optionally with a weight mask.
+
+Note that this is not the sort of noise masking you want for inpainting - the sampler will try to remove noise globally. 
+
+### Inputs 
+
+- *Required* - `noise1` - the original noise source.
+- *Optional* - `noise2` - the secondary noise source. If not connected, it is zero noise.
+- *Required* - `weight2` - the weight given to the second noise source. The first noise source has weight `1-weight2`.
+- *Optional* - `mask` - multiply `weight2` by the mask values. The mask will be rescaled to fit the latent.
+- *Required* - `renormalise` - should the noise be renormalised to mean of zero and stdev of one after mixing. Normally `yes`.
+
+### Outputs
+
+- `noise` A noise generator
+
+### Usage
+
+#### Generating small variations
+
+- Connect two noise sources to the node, and set `weight2` to `0`. Set the noise sources to have (different) fixed seeds.
+- Try different seeds on the first source until you get an image you like.
+- Increase `weight2` slowly (a weight of 0.2 is pretty big) to get variations on the image.
+- Try different second seeds as well
+
+#### Masked noise
+
+This is a bit more experimental!
+
+- Follow the first two steps above. 
+- Once you get an image you like, copy it into a `Load Image` node and edit a mask to pick the parts of the image you'd like to vary the noise for.
+- Connect the mask to the mask input, and then follow the third and fourth steps above.

__init__.py

@@ -0,0 +1,11 @@
+from noise_nodes import MixNoise, ShapeNoise
+
+VERSION = "1.0"
+
+NODE_CLASS_MAPPINGS = {
+    "Mix Noise"   : MixNoise,
+    "Shape Noise" : ShapeNoise,
+}
+__all__ = ['NODE_CLASS_MAPPINGS',]
+
+

noise.py

@@ -0,0 +1,23 @@
+from abc import ABC, abstractmethod
+import torch
+
+class Noise(ABC):
+    @abstractmethod
+    def generate_noise(self, input_latent:torch.Tensor) -> torch.Tensor: pass
+
+    @property
+    def seed(self): return None
+
+class NormalisableNoise(Noise):
+    def __init__(self, renormalise:bool):
+        self.renormalise = renormalise
+
+    def generate_noise(self, input_latent:torch.Tensor) -> torch.Tensor:
+        def normalise(noise:torch.Tensor, eps=1e-8):
+            std, mean = torch.std_mean(noise)
+            return (noise-mean)/(std+eps)
+        noise = self._generate_noise(input_latent)
+        return normalise(noise) if self.renormalise else noise
+    
+    @abstractmethod
+    def _generate_noise(self, input_latent:torch.Tensor) -> torch.Tensor: pass

noise_nodes.py

@@ -0,0 +1,90 @@
+import torch
+from typing import Optional
+from math import sin, pi
+from noise import Noise, NormalisableNoise
+
+class Noise_MixedNoise(NormalisableNoise):
+    def __init__(self, noise1:Noise, noise2:Optional[Noise], weight2:float, renormalise:bool, mask:Optional[torch.Tensor]):
+        super().__init__(renormalise)
+        self.noise1      = noise1
+        self.noise2      = noise2
+        self.weight2     = weight2
+        self.mask        = mask
+        
+    @property
+    def seed(self): return self.noise1.seed
+
+    def _generate_noise(self, input_latent:torch.Tensor) -> torch.Tensor:
+        noise1 = self.noise1.generate_noise(input_latent)
+        noise2 = self.noise2.generate_noise(input_latent) if self.noise2 is not None else torch.zeros_like(noise1)
+        mixed_noise = noise1 * (1.0-self.weight2) + noise2 * (self.weight2)
+        
+        if self.mask is not None:
+            while len(self.mask.shape)<4: self.mask.unsqueeze_(0)
+            mask:torch.Tensor = torch.nn.functional.interpolate(self.mask, size=input_latent['samples'].shape[-2:], mode='bilinear')
+            mask = mask.expand(-1,noise1.shape[1],-1,-1)
+            mixed_noise = mixed_noise * (mask) + noise1 * (1.0-mask)
+
+        return mixed_noise
+    
+class MixNoise:
+    CATEGORY = "quicknodes"
+    @classmethod    
+    def INPUT_TYPES(s):
+        return { 
+            "required":  { 
+                "noise1":      ("NOISE",), 
+                "weight2":     ("FLOAT", {"default":0.01, "step":0.001, "min":-1.0, "max":1.0}),
+                "renormalise": (["yes","no"],),
+                }, 
+            "optional" : {
+                "noise2": ("NOISE",), 
+                "mask":   ("MASK",),
+            }
+        }
+
+    RETURN_TYPES = ("NOISE",)
+    FUNCTION = "func"
+
+    def func(self, noise1, weight2, renormalise, noise2=None, mask=None):
+        return (Noise_MixedNoise(noise1, noise2, weight2, renormalise=='yes', mask),)
+    
+class Noise_ShapedNoise(NormalisableNoise):
+    def __init__(self, noise:NormalisableNoise, weight:float, renormalise:bool, x:bool, y:bool):
+        super().__init__(renormalise)
+        self.noise  = noise
+        self.weight = weight
+        self.x      = x
+        self.y      = y
+
+    @property
+    def seed(self): return self.noise.seed
+
+    def _generate_noise(self, input_latent:torch.Tensor) -> torch.Tensor:
+        def offset_sine(length:int) -> list[float]: return [ 2*sin(pi*x/length)-1 for x in range(length) ]
+        noise = self.noise.generate_noise(input_latent)
+        b,c,h,w = noise.shape
+        xscale = torch.ones((w,1)) + (self.weight * torch.Tensor([offset_sine(w),]) if self.x else 0)
+        yscale = torch.ones((h,1)) + (self.weight * torch.Tensor([offset_sine(h),]) if self.y else 0)
+        noise = noise * (torch.matmul(yscale.T,xscale))
+        return noise
+
+class ShapeNoise:
+    CATEGORY = "quicknodes"
+    @classmethod    
+    def INPUT_TYPES(s):
+        return { 
+            "required":  { 
+                "noise":       ("NOISE",), 
+                "weight":      ("FLOAT", {"default":0.01, "step":0.001, "min":-1.0, "max":1.0}),
+                "renormalise": (["yes","no"],),
+                "mode":        (["xy","x","y"],),
+            }, 
+        }
+    RETURN_TYPES = ("NOISE",)
+    FUNCTION = "func"   
+
+    def func(self, noise:NormalisableNoise, weight:float, renormalise:str, mode:str):
+        return (Noise_ShapedNoise(noise, weight, renormalise=="yes", 'x' in mode, 'y' in mode),)
+
+

pyproject.toml

@@ -0,0 +1,13 @@
+[project]
+name = "cg-noisetools"
+description = "A set of nodes that manipulate noise sources."
+version = "1.0"
+license = { file = "LICENSE" }
+
+[project.urls]
+Repository = "https://github.com/chrisgoringe/cg-noisetools"
+
+[tool.comfy]
+PublisherId = "chrisgoringe"
+DisplayName = "cg-noisetools"
+Icon = ""