nodes/blepping-ComfyUI-sonar
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow setting noise type for Euler A

Browse Source
This commit is contained in:
blepping
2024-02-02 13:58:45 -07:00
parent 3eb4e8292d
commit 347eb3f5cf
3 changed files with 368 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+22 -3
View File
@@ -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.
py/noise.py
+312
View File
@@ -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)
py/sonar.py
+34 -3
View File
@@ -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,
},