diff --git a/README.md b/README.md index e9e1cee..5cefead 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # ComfyUI-sonar -Extremely WIP and untested implementation of Sonar sampling. Currently it may not be even close to working properly. +Extremely WIP and untested implementation of Sonar sampling for [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Currently it may not be even close to working _properly_ but it does produce pretty reasonable results. Only supports Euler and Euler Ancestral sampling. @@ -12,8 +12,27 @@ The `direction` parameter should (unless I screwed it up) work like setting sign Like the original documentation says, you normally would not want to set `momentum` to a value below `0.85`. The default values are considered reasonable, doing stuff like using a negative direction may not produce good results. +## Parameters + +Very abbreviated section. The init type can make a big difference. If you use `RANDOM` you can get away with setting `direction` to high values (like up to `2.25` or so) and absurdly low values (like `-30.0`). It's also possible to set `momentum` and `momentum_hist` to negative values, although whether it's a good idea... + +## Noise + +I basically just copied a bunch of noise functions without really knowing what they do. The main thing I can say is they produce a semi-reasonable result and it's different from the other noise samplers. See credits below. + +1. `gaussian`: This is the default noise type. +2. `uniform`: Might enhance background details? +3. `brownian`: This is the noise type SDE samplers use. +4. `perlin` +5. `studentt`: There's a comment that says it may enhance subject details. It seemed to produce a fairly dark result. +6. `studentt_test`: An experiment that may be removed, it doesn't seem to be adding enough noise. You can possibly compensate by increasing `s_noise`. +7. `pink` +8. `highres_pyramid`: Not extensively tested, but it is slower than the other noise types. I would guess it does something like enhance details. + ## Credits -Original implementation: https://github.com/Kahsolt/stable-diffusion-webui-sonar +Original Sonar Sampler implementation (for A1111): https://github.com/Kahsolt/stable-diffusion-webui-sonar -My version basically just rips off this implementation for Diffusers: https://github.com/alexblattner/modified-euler-samplers-for-sonar-diffusers/ +My version basically just rips off this Sonar sampler implementation for Diffusers: https://github.com/alexblattner/modified-euler-samplers-for-sonar-diffusers/ + +Noise generation functions copied from https://github.com/Clybius/ComfyUI-Extra-Samplers with only minor modifications. I may have broken some of them in the process _or_ they may not have been suitable for use and I took them anyway. If they don't work it is not a reflection on the original source. diff --git a/py/noise.py b/py/noise.py new file mode 100644 index 0000000..d9615f4 --- /dev/null +++ b/py/noise.py @@ -0,0 +1,312 @@ +# Noise generation functions shamelessly yoinked from https://github.com/Clybius/ComfyUI-Extra-Samplers +from __future__ import annotations + +import math +import random +from typing import Callable + +import torch +from comfy.k_diffusion import sampling +from torch import FloatTensor, Generator, Tensor + +# ruff: noqa: D417,D212, D407, ANN002, ANN003, FBT002, S311 + + +class NoiseError(Exception): + pass + + +def get_positions(block_shape: tuple[int, int]) -> Tensor: + """ + Generate position tensor. + + Arguments: + block_shape -- (height, width) of position tensor + + Returns: + position vector shaped (1, height, width, 1, 1, 2) + """ + bh, bw = block_shape + return torch.stack( + torch.meshgrid( + [(torch.arange(b) + 0.5) / b for b in (bw, bh)], + indexing="xy", + ), + -1, + ).view(1, bh, bw, 1, 1, 2) + + +def unfold_grid(vectors: Tensor) -> Tensor: + """ + Unfold vector grid to batched vectors. + + Arguments: + vectors -- grid vectors + + Returns: + batched grid vectors + """ + batch_size, _, gpy, gpx = vectors.shape + return ( + torch.nn.functional.unfold(vectors, (2, 2)) + .view(batch_size, 2, 4, -1) + .permute(0, 2, 3, 1) + .view(batch_size, 4, gpy - 1, gpx - 1, 2) + ) + + +def smooth_step(t: Tensor) -> Tensor: + """ + Smooth step function [0, 1] -> [0, 1]. + + Arguments: + t -- input values (any shape) + + Returns: + output values (same shape as input values) + """ + return t * t * (3.0 - 2.0 * t) + + +def perlin_noise_tensor( + vectors: Tensor, + positions: Tensor, + step: Callable | None = None, +) -> Tensor: + """ + Generate perlin noise from batched vectors and positions. + + Arguments: + vectors -- batched grid vectors shaped (batch_size, 4, grid_height, grid_width, 2) + positions -- batched grid positions shaped (batch_size or 1, block_height, block_width, grid_height or 1, grid_width or 1, 2) + + Keyword Arguments: + step -- smooth step function [0, 1] -> [0, 1] (default: `smooth_step`) + + Raises: + Exception: if position and vector shapes do not match + + Returns: + (batch_size, block_height * grid_height, block_width * grid_width) + """ + if step is None: + step = smooth_step + + batch_size = vectors.shape[0] + # grid height, grid width + gh, gw = vectors.shape[2:4] + # block height, block width + bh, bw = positions.shape[1:3] + + for i in range(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): + msg = f"Batch sizes do not match: vectors ({vectors.shape[0]}), positions ({positions.shape[0]})" + raise NoiseError(msg) + + vectors = vectors.view(batch_size, 4, 1, gh * gw, 2) + positions = positions.view(positions.shape[0], bh * bw, -1, 2) + + step_x = step(positions[..., 0]) + step_y = step(positions[..., 1]) + + row0 = torch.lerp( + (vectors[:, 0] * positions).sum(dim=-1), + (vectors[:, 1] * (positions - positions.new_tensor((1, 0)))).sum(dim=-1), + step_x, + ) + row1 = torch.lerp( + (vectors[:, 2] * (positions - positions.new_tensor((0, 1)))).sum(dim=-1), + (vectors[:, 3] * (positions - positions.new_tensor((1, 1)))).sum(dim=-1), + step_x, + ) + noise = torch.lerp(row0, row1, step_y) + return ( + noise.view( + batch_size, + bh, + bw, + gh, + gw, + ) + .permute(0, 3, 1, 4, 2) + .reshape(batch_size, gh * bh, gw * bw) + ) + + +def perlin_noise( + grid_shape: tuple[int, int], + out_shape: tuple[int, int], + batch_size: int = 1, + generator: Generator | None = None, + *args, + **kwargs, +) -> Tensor: + """ + Generate perlin noise with given shape. `*args` and `**kwargs` are forwarded to `Tensor` creation. + + Arguments: + grid_shape -- Shape of grid (height, width). + out_shape -- Shape of output noise image (height, width). + + Keyword Arguments: + batch_size -- (default: {1}) + generator -- random generator used for grid vectors (default: {None}) + + Raises: + Exception: if grid and out shapes do not match + + Returns: + Noise image shaped (batch_size, height, width) + """ + # grid height and width + gh, gw = grid_shape + # output height and width + oh, ow = out_shape + # block height and width + bh, bw = oh // gh, ow // gw + + if oh != bh * gh: + msg = f"Output height {oh} must be divisible by grid height {gh}" + raise NoiseError(msg) + if ow != bw * gw != 0: + msg = f"Output width {ow} must be divisible by grid width {gw}" + raise NoiseError(msg) + + angle = torch.empty( + [batch_size] + [s + 1 for s in grid_shape], + *args, + **kwargs, + ).uniform_(to=2.0 * math.pi, generator=generator) + # random vectors on grid points + vectors = unfold_grid(torch.stack((torch.cos(angle), torch.sin(angle)), dim=1)) + # positions inside grid cells [0, 1) + positions = get_positions((bh, bw)).to(vectors) + return perlin_noise_tensor(vectors, positions).squeeze(0) + + +def rand_perlin_like(x): + noise = torch.randn_like(x) / 2.0 + noise_height = noise.size(dim=2) + noise_width = noise.size(dim=3) + for _ in range(2): + noise += perlin_noise( + (noise_height, noise_width), + (noise_height, noise_width), + batch_size=4, + ).to(x.device) + return noise / noise.std() + + +def uniform_noise_like(x): + return (torch.rand_like(x) - 0.5) * 3.46 + + +def highres_pyramid_noise_like(x, discount=0.7): + ( + b, + c, + h, + w, + ) = x.shape # EDIT: w and h get over-written, rename for a different variant! + orig_h = h + orig_w = w + u = torch.nn.Upsample(size=(orig_h, orig_w), mode="bilinear") + noise = uniform_noise_like(x) + for i in range(4): + r = random.random() * 2 + 2 # Rather than always going 2x, + h, w = min(orig_h * 15, int(h * (r**i))), min(orig_w * 15, int(w * (r**i))) + noise += u(torch.randn(b, c, h, w).to(x)) * discount**i + if h >= orig_h * 15 or w >= orig_w * 15: + break # Lowest resolution is 1x1 + return noise / noise.std() # Scaled back to roughly unit variance + + +def studentt_noise_like(x): + from torch.distributions import StudentT + + noise = StudentT(loc=0, scale=0.2, df=1).rsample(x.size()) + s: FloatTensor = torch.quantile(noise.flatten(start_dim=1).abs(), 0.75, dim=-1) + s = s.reshape(*s.shape, 1, 1, 1) + noise = noise.clamp(-s, s) + return torch.copysign(torch.pow(torch.abs(noise), 0.5), noise) + + +def studentt_noise_sampler( + x, +): # Produces more subject-focused outputs due to distribution, unsure if this works + noise = studentt_noise_like(x) + return lambda _sigma, _sigma_next: noise.to(x.device) / (7 / 3) + + +def green_noise_like(x): + # The comments said this didn't work and I had to learn the hard way. Turns out it's true! + (_, _, height, width) = x.shape + noise = torch.randn_like(x) + scale = 1.0 / (width * height) + fy = torch.fft.fftfreq(width, device=x.device)[:, None] ** 2 + fx = torch.fft.fftfreq(height, device=x.device) ** 2 + f = fy + fx + power = torch.sqrt(f) + power[0, 0] = 1 + noise = torch.fft.ifft2(torch.fft.fft2(noise) / torch.sqrt(power)) + noise *= scale / noise.std() + noise = torch.real(noise).to(x.device) + return noise / noise.std() + + +def generate_1f_noise(tensor, alpha, k, generator=None): + """Generate 1/f noise for a given tensor. + + Args: + tensor: The tensor to add noise to. + alpha: The parameter that determines the slope of the spectrum. + k: A constant. + + Returns: + A tensor with the same shape as `tensor` containing 1/f noise. + """ + fft = torch.fft.fft2(tensor) + freq = torch.arange(1, len(fft) + 1, dtype=torch.float) + spectral_density = k / freq**alpha + return torch.randn(tensor.shape, generator=generator) * spectral_density + + +def pink_noise_like(x): + noise = generate_1f_noise(x, 2.0, 1.0) + noise_mean = torch.mean(noise) + noise_std = torch.std(noise) + return noise.sub_(noise_mean).div_(noise_std).to(x.device) + + +NOISE_SAMPLERS = { + # No brownian as it is a special case that requires extra stuff like seed. + "gaussian": sampling.default_noise_sampler, + "uniform": lambda x: lambda _s, _sn: (torch.rand_like(x) - 0.5) * 3.46, + "perlin": lambda x: lambda _s, _sn: rand_perlin_like(x), + "studentt": studentt_noise_sampler, + "studentt_test": lambda x: lambda _s, _sn: studentt_noise_like(x).to(x.device), + "pink": lambda x: lambda _s, _sn: pink_noise_like(x), + "green_ish": lambda x: lambda _s, _sn: green_noise_like(x), + "highres_pyramid": lambda x: lambda _s, _sn: highres_pyramid_noise_like(x), +} + + +def get_noise_sampler(name, x, sigma_min, sigma_max, seed=None, use_cpu=True): + if name == "brownian": + return sampling.BrownianTreeNoiseSampler( + x, + sigma_min, + sigma_max, + seed=seed, + cpu=use_cpu, + ) + if name == "default" or not name: + name = "gaussian" + ns = NOISE_SAMPLERS.get(name, None) + if ns is None: + raise ValueError("Unknown noise sampler") + return ns(x) diff --git a/py/sonar.py b/py/sonar.py index c64f707..38d6636 100644 --- a/py/sonar.py +++ b/py/sonar.py @@ -1,4 +1,4 @@ -# Adapted from https://github.com/alexblattner/modified-euler-samplers-for-sonar-diffusers and https://github.com/Kahsolt/stable-diffusion-webui-sonar +# Sonar sampler part adapted from https://github.com/alexblattner/modified-euler-samplers-for-sonar-diffusers and https://github.com/Kahsolt/stable-diffusion-webui-sonar from __future__ import annotations @@ -8,6 +8,8 @@ from comfy.k_diffusion import sampling from torch import Tensor from tqdm.auto import trange +from . import noise + class SonarEuler: def __init__( @@ -197,6 +199,7 @@ def sample_sonar_euler_ancestral( momentum=0.95, momentum_hist=0.75, momentum_init="ZERO", + noise_type="gaussian", direction=1.0, eta=1.0, s_noise=1.0, @@ -210,9 +213,22 @@ def sample_sonar_euler_ancestral( ) s.sigmas = sigmas extra_args = {} if extra_args is None else extra_args - noise_sampler = ( - sampling.default_noise_sampler(x) if noise_sampler is None else noise_sampler + if noise_type != "gaussian" and noise_sampler is not None: + raise ValueError( + "Unexpected noise_sampler presence with non-default noise type requested", + ) + sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() + noise_sampler = noise.get_noise_sampler( + noise_type, + x, + sigma_min, + sigma_max, + seed=None, + use_cpu=True, ) + # noise_sampler = ( + # sampling.default_noise_sampler(x) if noise_sampler is None else noise_sampler + # ) s_in = x.new_ones([x.shape[0]]) for i in trange(len(sigmas) - 1, disable=disable): @@ -322,6 +338,19 @@ class SamplerSonarEulerAncestral(SamplerSonarEuler): "round": False, }, ) + result["required"]["noise_type"] = ( + ( + "gaussian", + "uniform", + "brownian", + "perlin", + "studentt", + "studentt_test", + "highres_pyramid", + "pink", + # "green_ish", + ), + ) return result def get_sampler( @@ -330,6 +359,7 @@ class SamplerSonarEulerAncestral(SamplerSonarEuler): momentum_hist, momentum_init, direction, + noise_type, eta, s_noise, ): @@ -341,6 +371,7 @@ class SamplerSonarEulerAncestral(SamplerSonarEuler): "momentum": momentum, "momentum_hist": momentum_hist, "direction": direction, + "noise_type": noise_type, "eta": eta, "s_noise": s_noise, },