nodes/spacepxl-ComfyUI-Image-Filters
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
main
spacepxl-ComfyUI-Image-Filters/nodes.py
T
spacepxl bbb3fb0045 Poisson noise
2025-12-14 21:00:47 -05:00

2406 lines
89 KiB
Python
Raw Permalink Blame History

import math
import copy
import torch
import torch.nn.functional as F
import numpy as np
import cv2
from pymatting import estimate_alpha_cf, estimate_foreground_ml, fix_trimap
from tqdm import trange
try:
from cv2.ximgproc import guidedFilter
except ImportError:
print("\033[33mUnable to import guidedFilter, make sure you have only opencv-contrib-python or run the import_error_install.bat script\033[m")
import comfy.model_management
import node_helpers
from server import PromptServer
from comfy.utils import ProgressBar
from comfy_extras.nodes_post_processing import gaussian_kernel
from .raft import *
MAX_RESOLUTION=8192
# gaussian blur a tensor image batch in format [B x H x W x C] on H/W (spatial, per-image, per-channel)
def cv_blur_tensor(images, dx, dy):
if min(dx, dy) > 100:
np_img = F.interpolate(images.detach().clone().movedim(-1,1), scale_factor=0.1, mode='bilinear').movedim(1,-1).cpu().numpy()
for index, image in enumerate(np_img):
np_img[index] = cv2.GaussianBlur(image, (dx // 20 * 2 + 1, dy // 20 * 2 + 1), 0)
return F.interpolate(torch.from_numpy(np_img).movedim(-1,1), size=(images.shape[1], images.shape[2]), mode='bilinear').movedim(1,-1)
else:
np_img = images.detach().clone().cpu().numpy()
for index, image in enumerate(np_img):
np_img[index] = cv2.GaussianBlur(image, (dx, dy), 0)
return torch.from_numpy(np_img)
# guided filter a tensor image batch in format [B x H x W x C] on H/W (spatial, per-image, per-channel)
def guided_filter_tensor(ref, images, d, s):
if d > 100:
np_img = F.interpolate(images.detach().clone().movedim(-1,1), scale_factor=0.1, mode='bilinear').movedim(1,-1).cpu().numpy()
np_ref = F.interpolate(ref.detach().clone().movedim(-1,1), scale_factor=0.1, mode='bilinear').movedim(1,-1).cpu().numpy()
for index, image in enumerate(np_img):
np_img[index] = guidedFilter(np_ref[index], image, d // 20 * 2 + 1, s)
return F.interpolate(torch.from_numpy(np_img).movedim(-1,1), size=(images.shape[1], images.shape[2]), mode='bilinear').movedim(1,-1)
else:
np_img = images.detach().clone().cpu().numpy()
np_ref = ref.cpu().numpy()
for index, image in enumerate(np_img):
np_img[index] = guidedFilter(np_ref[index], image, d, s)
return torch.from_numpy(np_img)
# std_dev and mean of tensor t within local spatial filter size d, per-image, per-channel [B x H x W x C]
def std_mean_filter(t, d):
t_mean = cv_blur_tensor(t, d, d)
t_diff_squared = (t - t_mean) ** 2
t_std = torch.sqrt(cv_blur_tensor(t_diff_squared, d, d))
return t_std, t_mean
def RGB2YCbCr(t):
YCbCr = t.detach().clone()
YCbCr[:,:,:,0] = 0.2123 * t[:,:,:,0] + 0.7152 * t[:,:,:,1] + 0.0722 * t[:,:,:,2]
YCbCr[:,:,:,1] = 0 - 0.1146 * t[:,:,:,0] - 0.3854 * t[:,:,:,1] + 0.5 * t[:,:,:,2]
YCbCr[:,:,:,2] = 0.5 * t[:,:,:,0] - 0.4542 * t[:,:,:,1] - 0.0458 * t[:,:,:,2]
return YCbCr
def YCbCr2RGB(t):
RGB = t.detach().clone()
RGB[:,:,:,0] = t[:,:,:,0] + 1.5748 * t[:,:,:,2]
RGB[:,:,:,1] = t[:,:,:,0] - 0.1873 * t[:,:,:,1] - 0.4681 * t[:,:,:,2]
RGB[:,:,:,2] = t[:,:,:,0] + 1.8556 * t[:,:,:,1]
return RGB
def hsv_to_rgb(h, s, v):
if s:
if h == 1.0: h = 0.0
i = int(h*6.0)
f = h*6.0 - i
w = v * (1.0 - s)
q = v * (1.0 - s * f)
t = v * (1.0 - s * (1.0 - f))
if i==0: return (v, t, w)
if i==1: return (q, v, w)
if i==2: return (w, v, t)
if i==3: return (w, q, v)
if i==4: return (t, w, v)
if i==5: return (v, w, q)
else: return (v, v, v)
def sRGBtoLinear(npArray):
less = npArray <= 0.0404482362771082
npArray[less] = npArray[less] / 12.92
npArray[~less] = np.power((npArray[~less] + 0.055) / 1.055, 2.4)
def linearToSRGB(npArray):
less = npArray <= 0.0031308
npArray[less] = npArray[less] * 12.92
npArray[~less] = np.power(npArray[~less], 1/2.4) * 1.055 - 0.055
def sRGBtoLinear_pt(t: torch.Tensor):
less = t <= 0.0404482362771082
t[less] = t[less] / 12.92
t[~less] = torch.pow((t[~less] + 0.055) / 1.055, 2.4)
return t
def linearToSRGB_pt(t: torch.Tensor):
less = t <= 0.0031308
t[less] = t[less] * 12.92
t[~less] = torch.pow(t[~less], 1 / 2.4) * 1.055 - 0.055
return t
def linearToTonemap(npArray, tonemap_scale):
npArray /= tonemap_scale
more = npArray > 0.06
SLog3 = np.clip((np.log10((npArray + 0.01)/0.19) * 261.5 + 420) / 1023, 0, 1)
npArray[more] = np.power(1 / (1 + (1 / np.power(SLog3[more] / (1 - SLog3[more]), 1.7))), 1.7)
npArray *= tonemap_scale
def tonemapToLinear(npArray, tonemap_scale):
npArray /= tonemap_scale
more = npArray > 0.06
x = np.power(np.clip(npArray, 0.000001, 1), 1/1.7)
ut = 1 / (1 + np.power((-1 / x) * (x - 1), 1/1.7))
npArray[more] = np.power(10, (ut[more] * 1023 - 420)/261.5) * 0.19 - 0.01
npArray *= tonemap_scale
def exposure(npArray, stops):
more = npArray > 0
npArray[more] *= pow(2, stops)
def randn_like_g(x, generator=None):
device = generator.device if generator is not None else x.device
r = torch.randn(x.size(), generator=generator, dtype=x.dtype, layout=x.layout, device=device)
return r.to(x.device)
class AlphaClean:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"radius": ("INT", {"default": 8, "min": 1, "max": 64, "step": 1}),
"fill_holes": ("INT", {"default": 1, "min": 0, "max": 16, "step": 1}),
"white_threshold": ("FLOAT", {"default": 0.9, "min": 0.01, "max": 1.0, "step": 0.01}),
"extra_clip": ("FLOAT", {"default": 0.98, "min": 0.01, "max": 1.0, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "alpha_clean"
CATEGORY = "Image-Filters/image"
DEPRECATED = True
def alpha_clean(self, images: torch.Tensor, radius: int, fill_holes: int, white_threshold: float, extra_clip: float):
d = radius * 2 + 1
i_dup = copy.deepcopy(images.cpu().numpy())
for index, image in enumerate(i_dup):
cleaned = cv2.bilateralFilter(image, 9, 0.05, 8)
alpha = np.clip((image - white_threshold) / (1 - white_threshold), 0, 1)
rgb = image * alpha
alpha = cv2.GaussianBlur(alpha, (d,d), 0) * 0.99 + np.average(alpha) * 0.01
rgb = cv2.GaussianBlur(rgb, (d,d), 0) * 0.99 + np.average(rgb) * 0.01
rgb = rgb / np.clip(alpha, 0.00001, 1)
rgb = rgb * extra_clip
cleaned = np.clip(cleaned / rgb, 0, 1)
if fill_holes > 0:
fD = fill_holes * 2 + 1
gamma = cleaned * cleaned
kD = np.ones((fD, fD), np.uint8)
kE = np.ones((fD + 2, fD + 2), np.uint8)
gamma = cv2.dilate(gamma, kD, iterations=1)
gamma = cv2.erode(gamma, kE, iterations=1)
gamma = cv2.GaussianBlur(gamma, (fD, fD), 0)
cleaned = np.maximum(cleaned, gamma)
i_dup[index] = cleaned
return (torch.from_numpy(i_dup),)
class MaskClean:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
"radius": ("INT", {"default": 8, "min": 1, "max": 64, "step": 1}),
"fill_holes": ("INT", {"default": 1, "min": 0, "max": 16, "step": 1}),
"white_threshold": ("FLOAT", {"default": 0.9, "min": 0.001, "max": 1.0, "step": 0.001}),
"extra_clip": ("FLOAT", {"default": 0.98, "min": 0.001, "max": 1.0, "step": 0.001}),
},
}
RETURN_TYPES = ("MASK",)
FUNCTION = "alpha_clean"
CATEGORY = "Image-Filters/mask"
def alpha_clean(self, mask, radius, fill_holes, white_threshold, extra_clip):
d = radius * 2 + 1
i_dup = mask.cpu().numpy()
for index, image in enumerate(i_dup):
cleaned = cv2.bilateralFilter(image, 9, 0.05, 8)
alpha = np.clip((image - white_threshold) / (1 - white_threshold), 0, 1)
rgb = image * alpha
alpha = cv2.GaussianBlur(alpha, (d,d), 0) * 0.99 + np.average(alpha) * 0.01
rgb = cv2.GaussianBlur(rgb, (d,d), 0) * 0.99 + np.average(rgb) * 0.01
rgb = rgb / np.clip(alpha, 0.00001, 1)
rgb = rgb * extra_clip
cleaned = np.clip(cleaned / rgb, 0, 1)
if fill_holes > 0:
fD = fill_holes * 2 + 1
gamma = cleaned * cleaned
kD = np.ones((fD, fD), np.uint8)
kE = np.ones((fD + 2, fD + 2), np.uint8)
gamma = cv2.dilate(gamma, kD, iterations=1)
gamma = cv2.erode(gamma, kE, iterations=1)
gamma = cv2.GaussianBlur(gamma, (fD, fD), 0)
cleaned = np.maximum(cleaned, gamma)
i_dup[index] = cleaned
return (torch.from_numpy(i_dup),)
class AlphaMatte:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"alpha_trimap": ("IMAGE",),
"preblur": ("INT", {"default": 8, "min": 0, "max": 256, "step": 1}),
"blackpoint": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 0.99, "step": 0.01}),
"whitepoint": ("FLOAT", {"default": 0.99, "min": 0.01, "max": 1.0, "step": 0.01}),
"max_iterations": ("INT", {"default": 1000, "min": 100, "max": 10000, "step": 100}),
"estimate_fg": (["true", "false"],),
},
}
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE",)
RETURN_NAMES = ("alpha", "fg", "bg",)
FUNCTION = "alpha_matte"
CATEGORY = "Image-Filters/image"
DEPRECATED = True
def alpha_matte(self, images, alpha_trimap, preblur, blackpoint, whitepoint, max_iterations, estimate_fg):
d = preblur * 2 + 1
i_dup = images.cpu().numpy().astype(np.float64)
a_dup = alpha_trimap.cpu().numpy().astype(np.float64)
fg = images.cpu().numpy().astype(np.float64)
bg = images.cpu().numpy().astype(np.float64)
for index, image in enumerate(i_dup):
trimap = a_dup[index][:,:,0] # convert to single channel
if preblur > 0:
trimap = cv2.GaussianBlur(trimap, (d, d), 0)
trimap = fix_trimap(trimap, blackpoint, whitepoint)
alpha = estimate_alpha_cf(image, trimap, laplacian_kwargs={"epsilon": 1e-6}, cg_kwargs={"maxiter":max_iterations})
if estimate_fg == "true":
fg[index], bg[index] = estimate_foreground_ml(image, alpha, return_background=True)
a_dup[index] = np.stack([alpha, alpha, alpha], axis = -1) # convert back to rgb
return (
torch.from_numpy(a_dup.astype(np.float32)), # alpha
torch.from_numpy(fg.astype(np.float32)), # fg
torch.from_numpy(bg.astype(np.float32)), # bg
)
class ImageMatting:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"trimap": ("MASK",),
"preblur": ("INT", {"default": 8, "min": 0, "max": 256, "step": 1}),
"blackpoint": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 0.99, "step": 0.01}),
"whitepoint": ("FLOAT", {"default": 0.99, "min": 0.01, "max": 1.0, "step": 0.01}),
"max_iterations": ("INT", {"default": 1000, "min": 10, "max": 10000, "step": 10}),
"estimate_fg": ("BOOLEAN", {"default": True}),
},
}
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE",)
RETURN_NAMES = ("matte", "fg", "bg",)
FUNCTION = "alpha_matte"
CATEGORY = "Image-Filters/image"
def alpha_matte(self, images, trimap, preblur, blackpoint, whitepoint, max_iterations, estimate_fg):
d = preblur * 2 + 1
i_dup = images.cpu().numpy().astype(np.float64)
a_dup = trimap.cpu().numpy().astype(np.float64)
fg = copy.deepcopy(i_dup)
bg = copy.deepcopy(i_dup)
for index, image in enumerate(i_dup):
trimap = a_dup[index]
if preblur > 0:
trimap = cv2.GaussianBlur(trimap, (d, d), 0)
trimap = fix_trimap(trimap, blackpoint, whitepoint)
alpha = estimate_alpha_cf(image, trimap, laplacian_kwargs={"epsilon": 1e-6}, cg_kwargs={"maxiter":max_iterations})
if estimate_fg:
fg[index], bg[index] = estimate_foreground_ml(image, alpha, return_background=True)
a_dup[index] = alpha
return (
torch.from_numpy(a_dup.astype(np.float32)), # matte
torch.from_numpy(fg.astype(np.float32)), # fg
torch.from_numpy(bg.astype(np.float32)), # bg
)
class BetterFilmGrain:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"scale": ("FLOAT", {"default": 0.5, "min": 0.25, "max": 2.0, "step": 0.05}),
"strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.01}),
"saturation": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01}),
"toe": ("FLOAT", {"default": 0.0, "min": -0.2, "max": 0.5, "step": 0.001}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "grain"
CATEGORY = "Image-Filters/image"
def grain(self, image, scale, strength, saturation, toe, seed):
t = image.detach().clone()
torch.manual_seed(seed)
grain = torch.rand(t.shape[0], int(t.shape[1] // scale), int(t.shape[2] // scale), 3)
YCbCr = RGB2YCbCr(grain)
YCbCr[:,:,:,0] = cv_blur_tensor(YCbCr[:,:,:,0], 3, 3)
YCbCr[:,:,:,1] = cv_blur_tensor(YCbCr[:,:,:,1], 15, 15)
YCbCr[:,:,:,2] = cv_blur_tensor(YCbCr[:,:,:,2], 11, 11)
grain = (YCbCr2RGB(YCbCr) - 0.5) * strength
grain[:,:,:,0] *= 2
grain[:,:,:,2] *= 3
grain += 1
grain = grain * saturation + grain[:,:,:,1].unsqueeze(3).repeat(1,1,1,3) * (1 - saturation)
grain = F.interpolate(grain.movedim(-1,1), size=(t.shape[1], t.shape[2]), mode='bilinear').movedim(1,-1)
t[:,:,:,:3] = torch.clip((1 - (1 - t[:,:,:,:3]) * grain) * (1 - toe) + toe, 0, 1)
return(t,)
class BlurImageFast:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"radius_x": ("INT", {"default": 1, "min": 0, "max": 1023, "step": 1}),
"radius_y": ("INT", {"default": 1, "min": 0, "max": 1023, "step": 1}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "blur_image"
CATEGORY = "Image-Filters/image"
def blur_image(self, images, radius_x, radius_y):
if radius_x + radius_y == 0:
return (images,)
dx = radius_x * 2 + 1
dy = radius_y * 2 + 1
dup = copy.deepcopy(images.cpu().numpy())
for index, image in enumerate(dup):
dup[index] = cv2.GaussianBlur(image, (dx, dy), 0)
return (torch.from_numpy(dup),)
class BlurMaskFast:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"masks": ("MASK",),
"radius_x": ("INT", {"default": 1, "min": 0, "max": 1023, "step": 1}),
"radius_y": ("INT", {"default": 1, "min": 0, "max": 1023, "step": 1}),
},
}
RETURN_TYPES = ("MASK",)
FUNCTION = "blur_mask"
CATEGORY = "Image-Filters/mask"
def blur_mask(self, masks, radius_x, radius_y):
if radius_x + radius_y == 0:
return (masks,)
dx = radius_x * 2 + 1
dy = radius_y * 2 + 1
dup = copy.deepcopy(masks.cpu().numpy())
for index, mask in enumerate(dup):
dup[index] = cv2.GaussianBlur(mask, (dx, dy), 0)
return (torch.from_numpy(dup),)
class ColorMatchImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", ),
"reference": ("IMAGE", ),
"blur_type": (["blur", "guidedFilter"],),
"blur_size": ("INT", {"default": 0, "min": 0, "max": 1023}),
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "batch_normalize"
CATEGORY = "Image-Filters/image"
def batch_normalize(self, images, reference, blur_type, blur_size, factor):
t = images.detach().clone() + 0.1
ref = reference.detach().clone() + 0.1
if ref.shape[0] < t.shape[0]:
ref = ref[0].unsqueeze(0).repeat(t.shape[0], 1, 1, 1)
if blur_size == 0:
mean = torch.mean(t, (1,2), keepdim=True)
mean_ref = torch.mean(ref, (1,2), keepdim=True)
for i in range(t.shape[0]):
for c in range(3):
t[i,:,:,c] /= mean[i,0,0,c]
t[i,:,:,c] *= mean_ref[i,0,0,c]
else:
d = blur_size * 2 + 1
if blur_type == "blur":
blurred = cv_blur_tensor(t, d, d)
blurred_ref = cv_blur_tensor(ref, d, d)
elif blur_type == "guidedFilter":
blurred = guided_filter_tensor(t, t, d, 0.01)
blurred_ref = guided_filter_tensor(ref, ref, d, 0.01)
for i in range(t.shape[0]):
for c in range(3):
t[i,:,:,c] /= blurred[i,:,:,c]
t[i,:,:,c] *= blurred_ref[i,:,:,c]
t = t - 0.1
torch.clamp(torch.lerp(images, t, factor), 0, 1)
return (t,)
class RestoreDetail:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", ),
"detail": ("IMAGE", ),
"mode": (["add", "multiply"],),
"blur_type": (["blur", "guidedFilter"],),
"blur_size": ("INT", {"default": 1, "min": 1, "max": 1023}),
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "batch_normalize"
CATEGORY = "Image-Filters/image"
def batch_normalize(self, images, detail, mode, blur_type, blur_size, factor):
t = images.detach().clone() + 0.1
ref = detail.detach().clone() + 0.1
if ref.shape[0] < t.shape[0]:
ref = ref[0].unsqueeze(0).repeat(t.shape[0], 1, 1, 1)
d = blur_size * 2 + 1
if blur_type == "blur":
blurred = cv_blur_tensor(t, d, d)
blurred_ref = cv_blur_tensor(ref, d, d)
elif blur_type == "guidedFilter":
blurred = guided_filter_tensor(t, t, d, 0.01)
blurred_ref = guided_filter_tensor(ref, ref, d, 0.01)
if mode == "multiply":
t = (ref / blurred_ref) * blurred
else:
t = (ref - blurred_ref) + blurred
t = t - 0.1
t = torch.clamp(torch.lerp(images, t, factor), 0, 1)
return (t,)
class DilateErodeMask:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"masks": ("MASK",),
"radius": ("INT", {"default": 0, "min": -1023, "max": 1023, "step": 1}),
"shape": (["box", "circle"],),
},
}
RETURN_TYPES = ("MASK",)
FUNCTION = "dilate_mask"
CATEGORY = "Image-Filters/mask"
def dilate_mask(self, masks, radius, shape):
if radius == 0:
return (masks,)
s = abs(radius)
d = s * 2 + 1
k = np.zeros((d, d), np.uint8)
if shape == "circle":
k = cv2.circle(k, (s,s), s, 1, -1)
else:
k += 1
dup = copy.deepcopy(masks.cpu().numpy())
for index, mask in enumerate(dup):
if radius > 0:
dup[index] = cv2.dilate(mask, k, iterations=1)
else:
dup[index] = cv2.erode(mask, k, iterations=1)
return (torch.from_numpy(dup),)
class EnhanceDetail:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"filter_radius": ("INT", {"default": 2, "min": 1, "max": 64, "step": 1}),
"sigma": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 100.0, "step": 0.01}),
"denoise": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 10.0, "step": 0.01}),
"detail_mult": ("FLOAT", {"default": 2.0, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "enhance"
CATEGORY = "Image-Filters/image"
def enhance(self, images: torch.Tensor, filter_radius: int, sigma: float, denoise: float, detail_mult: float):
if filter_radius == 0:
return (images,)
d = filter_radius * 2 + 1
s = sigma / 10
n = denoise / 10
dup = copy.deepcopy(images.cpu().numpy())
for index, image in enumerate(dup):
imgB = image
if denoise > 0.0:
imgB = cv2.bilateralFilter(image, d, n, d)
imgG = np.clip(guidedFilter(image, image, d, s), 0.001, 1)
details = (imgB/imgG - 1) * detail_mult + 1
dup[index] = np.clip(details*imgG - imgB + image, 0, 1)
return (torch.from_numpy(dup),)
class GuidedFilterImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", ),
"guide": ("IMAGE", ),
"size": ("INT", {"default": 4, "min": 0, "max": 1023}),
"sigma": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 100.0, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "filter_image"
CATEGORY = "Image-Filters/image"
def filter_image(self, images, guide, size, sigma):
d = size * 2 + 1
s = sigma / 10
filtered = guided_filter_tensor(guide, images, d, s)
return (filtered,)
class MedianFilterImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", ),
"size": ("INT", {"default": 1, "min": 1, "max": 1023}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "filter_image"
CATEGORY = "Image-Filters/image"
def filter_image(self, images, size):
np_images = images.detach().clone().cpu().numpy()
d = size * 2 + 1
for index, image in enumerate(np_images):
if d > 5:
work_image = image * 255
work_image = cv2.medianBlur(work_image.astype(np.uint8), d)
np_images[index] = work_image.astype(np.float32) / 255
else:
np_images[index] = cv2.medianBlur(image, d)
return (torch.from_numpy(np_images),)
class BilateralFilterImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", ),
"size": ("INT", {"default": 8, "min": 1, "max": 64}),
"sigma_color": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1000.0, "step": 0.01}),
"sigma_space": ("FLOAT", {"default": 100.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "filter_image"
CATEGORY = "Image-Filters/image"
def filter_image(self, images, size, sigma_color, sigma_space):
np_images = images.detach().clone().cpu().numpy()
d = size * 2 + 1
for index, image in enumerate(np_images):
np_images[index] = cv2.bilateralFilter(image, d, sigma_color, sigma_space)
return (torch.from_numpy(np_images),)
class FrequencyCombine:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"high_frequency": ("IMAGE", ),
"low_frequency": ("IMAGE", ),
"mode": (["subtract", "divide"],),
"eps": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 0.99, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "filter_image"
CATEGORY = "Image-Filters/image"
def filter_image(self, high_frequency, low_frequency, mode, eps):
t = low_frequency.detach().clone()
if mode == "subtract":
t = t + high_frequency - 0.5
else:
t = (high_frequency * 2) * (t + eps) - eps
return (torch.clamp(t, 0, 1),)
class FrequencySeparate:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"original": ("IMAGE", ),
"low_frequency": ("IMAGE", ),
"mode": (["subtract", "divide"],),
"eps": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 0.99, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("high_frequency",)
FUNCTION = "filter_image"
CATEGORY = "Image-Filters/image"
def filter_image(self, original, low_frequency, mode, eps):
t = original.detach().clone()
if mode == "subtract":
t = t - low_frequency + 0.5
else:
t = ((t + eps) / (low_frequency + eps)) * 0.5
return (t,)
class RemapRange:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"blackpoint": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"whitepoint": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 1.0, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "remap"
CATEGORY = "Image-Filters/image"
def remap(self, image: torch.Tensor, blackpoint: float, whitepoint: float):
bp = min(blackpoint, whitepoint - 0.001)
scale = 1 / (whitepoint - bp)
i_dup = copy.deepcopy(image.cpu().numpy())
i_dup = np.clip((i_dup - bp) * scale, 0.0, 1.0)
return (torch.from_numpy(i_dup),)
class ClampImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"blackpoint": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"whitepoint": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "clamp_image"
CATEGORY = "Image-Filters/image"
def clamp_image(self, image: torch.Tensor, blackpoint: float, whitepoint: float):
clamped_image = torch.clamp(torch.nan_to_num(image.detach().clone()), min=blackpoint, max=whitepoint)
return (clamped_image,)
Channel_List = ["red", "green", "blue", "alpha", "white", "black"]
Alpha_List = ["red", "green", "blue", "alpha", "white", "black", "none"]
class ShuffleChannels:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"red": (Channel_List, {"default": "red"}),
"green": (Channel_List, {"default": "green"}),
"blue": (Channel_List, {"default": "blue"}),
"alpha": (Alpha_List, {"default": "none"}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "shuffle"
CATEGORY = "Image-Filters/image"
def shuffle(self, image, red, green, blue, alpha):
ch = 3 if alpha == "none" else 4
t = torch.zeros((image.shape[0], image.shape[1], image.shape[2], ch), dtype=image.dtype, device=image.device)
image_copy = image.detach().clone()
ch_key = [red, green, blue, alpha]
for i in range(ch):
if ch_key[i] == "white":
t[:,:,:,i] = 1
elif ch_key[i] == "red":
t[:,:,:,i] = image_copy[:,:,:,0]
elif ch_key[i] == "green":
t[:,:,:,i] = image_copy[:,:,:,1]
elif ch_key[i] == "blue":
t[:,:,:,i] = image_copy[:,:,:,2]
elif ch_key[i] == "alpha":
if image.shape[3] > 3:
t[:,:,:,i] = image_copy[:,:,:,3]
else:
t[:,:,:,i] = 1
return(t,)
class ClampOutliers:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"latents": ("LATENT", ),
"std_dev": ("FLOAT", {"default": 3.0, "min": 0.1, "max": 100.0, "step": 0.1, "round": 0.1}),
},
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "clamp_outliers"
CATEGORY = "Image-Filters/latent"
def clamp_outliers(self, latents, std_dev):
latents_copy = copy.deepcopy(latents)
t = latents_copy["samples"]
for i, latent in enumerate(t):
for j, channel in enumerate(latent):
sd, mean = torch.std_mean(channel, dim=None)
t[i,j] = torch.clamp(channel, min = -sd * std_dev + mean, max = sd * std_dev + mean)
latents_copy["samples"] = t
return (latents_copy,)
class AdainLatent:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"latents": ("LATENT", ),
"reference": ("LATENT", ),
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
},
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "batch_normalize"
CATEGORY = "Image-Filters/latent"
def batch_normalize(self, latents, reference, factor):
latents_copy = copy.deepcopy(latents)
t = latents_copy["samples"]
t_std, t_mean = torch.std_mean(t, dim=(-2, -1), keepdim=True)
ref_std, ref_mean = torch.std_mean(reference["samples"], dim=(-2, -1), keepdim=True)
t = (t - t_mean) / t_std
t = t * ref_std + ref_mean
latents_copy["samples"] = torch.lerp(latents["samples"], t, factor)
return (latents_copy,)
class AdainFilterLatent:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"latents": ("LATENT", ),
"reference": ("LATENT", ),
"filter_size": ("INT", {"default": 1, "min": 1, "max": 128}),
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
},
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "batch_normalize"
CATEGORY = "Image-Filters/latent"
def batch_normalize(self, latents, reference, filter_size, factor):
latents_copy = copy.deepcopy(latents)
t = latents_copy["samples"].movedim(1, -1) # BCHW -> BHWC or BCFHW -> BFHWC
ref = reference["samples"].movedim(1, -1)
d = filter_size * 2 + 1
if t.dim() == 5:
t_std, t_mean, ref_std, ref_mean = [], [], [], []
for b in range(t.shape[0]):
tb_std, tb_mean = std_mean_filter(t[b], d)
rb_std, rb_mean = std_mean_filter(ref[b], d)
t_std.append(tb_std)
t_mean.append(tb_mean)
ref_std.append(rb_std)
ref_mean.append(rb_mean)
t_std = torch.stack(t_std, dim=0)
t_mean = torch.stack(t_mean, dim=0)
ref_std = torch.stack(ref_std, dim=0)
ref_mean = torch.stack(ref_mean, dim=0)
else:
t_std, t_mean = std_mean_filter(t, d)
ref_std, ref_mean = std_mean_filter(ref, d)
t = (t - t_mean) / t_std
t = t * ref_std + ref_mean
t = t.movedim(-1, 1) # BHWC -> BCHW or BFHWC -> BCFHW
latents_copy["samples"] = torch.lerp(latents["samples"], t, factor)
return (latents_copy,)
class SharpenFilterLatent:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"latents": ("LATENT", ),
"filter_size": ("INT", {"default": 1, "min": 1, "max": 128}),
"factor": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "round": 0.01}),
},
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "filter_latent"
CATEGORY = "Image-Filters/latent"
def filter_latent(self, latents, filter_size, factor):
latents_copy = copy.deepcopy(latents)
t = latents_copy["samples"].movedim(1, -1) # BCHW -> BHWC or BCFHW -> BFHWC
d = filter_size * 2 + 1
if t.dim() == 5:
t_blurred = []
for b in range(t.shape[0]):
t_blurred.append(cv_blur_tensor(t[b], d, d))
t_blurred = torch.stack(t_blurred, dim=0)
else:
t_blurred = cv_blur_tensor(t, d, d)
t = t - t_blurred
t = t * factor
t = t + t_blurred
t = t.movedim(-1, 1) # BHWC -> BCHW or BFHWC -> BCFHW
latents_copy["samples"] = t
return (latents_copy,)
class AdainImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", ),
"reference": ("IMAGE", ),
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "batch_normalize"
CATEGORY = "Image-Filters/image"
def batch_normalize(self, images, reference, factor):
t = copy.deepcopy(images) # [B x H x W x C]
t = t.movedim(-1,0) # [C x B x H x W]
for c in range(t.size(0)):
for i in range(t.size(1)):
r_sd, r_mean = torch.std_mean(reference[i, :, :, c], dim=None) # index by original dim order
i_sd, i_mean = torch.std_mean(t[c, i], dim=None)
t[c, i] = ((t[c, i] - i_mean) / i_sd) * r_sd + r_mean
t = torch.lerp(images, t.movedim(0,-1), factor) # [B x H x W x C]
return (t,)
class BatchNormalizeLatent:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"latents": ("LATENT", ),
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
},
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "batch_normalize"
CATEGORY = "Image-Filters/latent"
def batch_normalize(self, latents, factor):
latents_copy = copy.deepcopy(latents)
t = latents_copy["samples"] # [B x C x H x W]
t = t.movedim(0,1) # [C x B x H x W]
for c in range(t.size(0)):
c_sd, c_mean = torch.std_mean(t[c], dim=None)
for i in range(t.size(1)):
i_sd, i_mean = torch.std_mean(t[c, i], dim=None)
t[c, i] = (t[c, i] - i_mean) / i_sd
t[c] = t[c] * c_sd + c_mean
latents_copy["samples"] = torch.lerp(latents["samples"], t.movedim(1,0), factor) # [B x C x H x W]
return (latents_copy,)
class BatchNormalizeImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", ),
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "batch_normalize"
CATEGORY = "Image-Filters/image"
def batch_normalize(self, images, factor):
t = copy.deepcopy(images) # [B x H x W x C]
t = t.movedim(-1,0) # [C x B x H x W]
for c in range(t.size(0)):
c_sd, c_mean = torch.std_mean(t[c], dim=None)
for i in range(t.size(1)):
i_sd, i_mean = torch.std_mean(t[c, i], dim=None)
t[c, i] = (t[c, i] - i_mean) / i_sd
t[c] = t[c] * c_sd + c_mean
t = torch.lerp(images, t.movedim(0,-1), factor) # [B x H x W x C]
return (t,)
class DifferenceChecker:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images1": ("IMAGE", ),
"images2": ("IMAGE", ),
"multiplier": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 1000.0, "step": 0.01, "round": 0.01}),
"print_MAE": ("BOOLEAN", {"default": False}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "difference_checker"
OUTPUT_NODE = True
CATEGORY = "Image-Filters/image"
def difference_checker(self, images1, images2, multiplier, print_MAE):
t = copy.deepcopy(images1)
t = torch.abs(images1 - images2)
if print_MAE:
print(f"MAE = {torch.mean(t)}")
return (torch.clamp(t * multiplier, min=0, max=1),)
class ImageConstant:
def __init__(self, device="cpu"):
self.device = device
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
"height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
"red": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"green": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"blue": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "generate"
CATEGORY = "Image-Filters/image"
def generate(self, width, height, batch_size, red, green, blue):
r = torch.full([batch_size, height, width, 1], red)
g = torch.full([batch_size, height, width, 1], green)
b = torch.full([batch_size, height, width, 1], blue)
return (torch.cat((r, g, b), dim=-1), )
class ImageConstantHSV:
def __init__(self, device="cpu"):
self.device = device
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
"height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
"hue": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"saturation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"value": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "generate"
CATEGORY = "Image-Filters/image"
def generate(self, width, height, batch_size, hue, saturation, value):
red, green, blue = hsv_to_rgb(hue, saturation, value)
r = torch.full([batch_size, height, width, 1], red)
g = torch.full([batch_size, height, width, 1], green)
b = torch.full([batch_size, height, width, 1], blue)
return (torch.cat((r, g, b), dim=-1), )
class OffsetLatentImage:
def __init__(self):
self.device = comfy.model_management.intermediate_device()
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
"height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
"offset_0": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.1, "round": 0.1}),
"offset_1": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.1, "round": 0.1}),
"offset_2": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.1, "round": 0.1}),
"offset_3": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.1, "round": 0.1}),
},
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "generate"
CATEGORY = "Image-Filters/latent"
def generate(self, width, height, batch_size, offset_0, offset_1, offset_2, offset_3):
latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
latent[:,0,:,:] = offset_0
latent[:,1,:,:] = offset_1
latent[:,2,:,:] = offset_2
latent[:,3,:,:] = offset_3
return ({"samples":latent}, )
class RelightSimple:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"normals": ("IMAGE",),
"x": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001}),
"y": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001}),
"z": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 1.0, "step": 0.001}),
"brightness": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "relight"
CATEGORY = "Image-Filters/image"
def relight(self, image, normals, x, y, z, brightness):
if image.shape[0] != normals.shape[0]:
raise Exception("Batch size for image and normals must match")
norm = normals.detach().clone() * 2 - 1
norm = F.interpolate(norm.movedim(-1,1), size=(image.shape[1], image.shape[2]), mode='bilinear').movedim(1,-1)
light = torch.tensor([x, y, z])
light = F.normalize(light, dim=0)
diffuse = norm[:,:,:,0] * light[0] + norm[:,:,:,1] * light[1] + norm[:,:,:,2] * light[2]
diffuse = torch.clip(diffuse.unsqueeze(3).repeat(1,1,1,3), 0, 1)
relit = image.detach().clone()
relit[:,:,:,:3] = torch.clip(relit[:,:,:,:3] * diffuse * brightness, 0, 1)
return (relit,)
class LatentStats:
@classmethod
def INPUT_TYPES(s):
return {"required": {"latent": ("LATENT", ),}}
RETURN_TYPES = ("STRING", "FLOAT", "FLOAT", "FLOAT", "FLOAT")
RETURN_NAMES = ("stats", "c0_mean", "c1_mean", "c2_mean", "c3_mean")
FUNCTION = "notify"
OUTPUT_NODE = True
CATEGORY = "Image-Filters/utils"
def notify(self, latent):
latents = latent["samples"]
channels = latents.size(1)
width, height = latents.size(3), latents.size(2)
text = ["",]
text[0] = f"batch size: {latents.size(0)}"
text.append(f"channels: {channels}")
text.append(f"width: {width} ({width * 8})")
text.append(f"height: {height} ({height * 8})")
cmean = [0,0,0,0]
for i in range(channels):
minimum = torch.min(latents[:,i,:,:]).item()
maximum = torch.max(latents[:,i,:,:]).item()
std_dev, mean = torch.std_mean(latents[:,i,:,:], dim=None)
if i < 4:
cmean[i] = mean
text.append(f"c{i} mean: {mean:.1f} std_dev: {std_dev:.1f} min: {minimum:.1f} max: {maximum:.1f}")
printtext = "\033[36mLatent Stats:\033[m"
for t in text:
printtext += "\n " + t
returntext = ""
for i in range(len(text)):
if i > 0:
returntext += "\n"
returntext += text[i]
print(printtext)
return (returntext, cmean[0], cmean[1], cmean[2], cmean[3])
class Tonemap:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"input_mode": (["linear", "sRGB"],),
"output_mode": (["sRGB", "linear"],),
"tonemap_scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 10, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "apply"
CATEGORY = "Image-Filters/image"
def apply(self, images, input_mode, output_mode, tonemap_scale):
t = images.detach().clone().cpu().numpy().astype(np.float32)
if input_mode == "sRGB":
sRGBtoLinear(t[:,:,:,:3])
linearToTonemap(t[:,:,:,:3], tonemap_scale)
if output_mode == "sRGB":
linearToSRGB(t[:,:,:,:3])
t = np.clip(t, 0, 1)
t = torch.from_numpy(t)
return (t,)
class UnTonemap:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"input_mode": (["sRGB", "linear"],),
"output_mode": (["linear", "sRGB"],),
"tonemap_scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 10, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "apply"
CATEGORY = "Image-Filters/image"
def apply(self, images, input_mode, output_mode, tonemap_scale):
t = images.detach().clone().cpu().numpy().astype(np.float32)
if input_mode == "sRGB":
sRGBtoLinear(t[:,:,:,:3])
tonemapToLinear(t[:,:,:,:3], tonemap_scale)
if output_mode == "sRGB":
linearToSRGB(t[:,:,:,:3])
t = np.clip(t, 0, 1)
t = torch.from_numpy(t)
return (t,)
class ExposureAdjust:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"stops": ("FLOAT", {"default": 0.0, "min": -100, "max": 100, "step": 0.01}),
"input_mode": (["sRGB", "linear"],),
"output_mode": (["sRGB", "linear"],),
"tonemap": (["linear", "Reinhard", "linlog"], {"default": "Reinhard"}),
"tonemap_scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 10, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "adjust_exposure"
CATEGORY = "Image-Filters/image"
def adjust_exposure(self, images, stops, input_mode, output_mode, tonemap, tonemap_scale):
t = images.detach().clone().cpu().numpy().astype(np.float32)
if input_mode == "sRGB":
sRGBtoLinear(t[...,:3])
if tonemap == "linlog":
tonemapToLinear(t[...,:3], tonemap_scale)
elif tonemap == "Reinhard":
t = np.clip(t, 0, 0.999)
t[...,:3] = -t[...,:3] / (t[...,:3] - 1)
exposure(t[...,:3], stops)
if tonemap == "linlog":
linearToTonemap(t[...,:3], tonemap_scale)
elif tonemap == "Reinhard":
t[...,:3] = t[...,:3] / (t[...,:3] + 1)
if output_mode == "sRGB":
linearToSRGB(t[...,:3])
t = np.clip(t, 0, 1)
t = torch.from_numpy(t)
return (t,)
# Normal map standard coordinates: +r:+x:right, +g:+y:up, +b:+z:in
class ConvertNormals:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"normals": ("IMAGE",),
"input_mode": (["BAE", "MiDaS", "Standard"],),
"output_mode": (["BAE", "MiDaS", "Standard"],),
"scale_XY": ("FLOAT",{"default": 1, "min": 0, "max": 100, "step": 0.001}),
"normalize": ("BOOLEAN", {"default": True}),
"fix_black": ("BOOLEAN", {"default": True}),
},
"optional": {
"optional_fill": ("IMAGE",),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "convert_normals"
CATEGORY = "Image-Filters/image"
def convert_normals(self, normals, input_mode, output_mode, scale_XY, normalize, fix_black, optional_fill=None):
t = normals.detach().clone()
if input_mode == "BAE":
t[:,:,:,0] = 1 - t[:,:,:,0] # invert R
elif input_mode == "MiDaS":
t[:,:,:,:3] = torch.stack([1 - t[:,:,:,2], t[:,:,:,1], t[:,:,:,0]], dim=3) # BGR -> RGB and invert R
if fix_black:
key = torch.clamp(1 - t[:,:,:,2] * 2, min=0, max=1)
if optional_fill == None:
t[:,:,:,0] += key * 0.5
t[:,:,:,1] += key * 0.5
t[:,:,:,2] += key
else:
fill = optional_fill.detach().clone()
if fill.shape[1:3] != t.shape[1:3]:
fill = F.interpolate(fill.movedim(-1,1), size=(t.shape[1], t.shape[2]), mode='bilinear').movedim(1,-1)
if fill.shape[0] != t.shape[0]:
fill = fill[0].unsqueeze(0).expand(t.shape[0], -1, -1, -1)
t[:,:,:,:3] += fill[:,:,:,:3] * key.unsqueeze(3).expand(-1, -1, -1, 3)
t[:,:,:,:2] = (t[:,:,:,:2] - 0.5) * scale_XY + 0.5
if normalize:
t[:,:,:,:3] = F.normalize(t[:,:,:,:3] * 2 - 1, dim=3) / 2 + 0.5
if output_mode == "BAE":
t[:,:,:,0] = 1 - t[:,:,:,0] # invert R
elif output_mode == "MiDaS":
t[:,:,:,:3] = torch.stack([t[:,:,:,2], t[:,:,:,1], 1 - t[:,:,:,0]], dim=3) # invert R and BGR -> RGB
return (t,)
class BatchAverageImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"operation": (["mean", "median"],),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "apply"
CATEGORY = "Image-Filters/image"
def apply(self, images, operation):
t = images.detach().clone()
if operation == "mean":
return (torch.mean(t, dim=0, keepdim=True),)
elif operation == "median":
return (torch.median(t, dim=0, keepdim=True)[0],)
return(t,)
class NormalMapSimple:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"scale_XY": ("FLOAT",{"default": 1, "min": 0, "max": 100, "step": 0.001}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "normal_map"
CATEGORY = "Image-Filters/image"
def normal_map(self, images, scale_XY):
t = images.detach().clone().cpu().numpy().astype(np.float32)
L = np.mean(t[:,:,:,:3], axis=3)
for i in range(t.shape[0]):
t[i,:,:,0] = cv2.Scharr(L[i], -1, 1, 0, cv2.BORDER_REFLECT) * -1
t[i,:,:,1] = cv2.Scharr(L[i], -1, 0, 1, cv2.BORDER_REFLECT)
t[:,:,:,2] = 1
t = torch.from_numpy(t)
t[:,:,:,:2] *= scale_XY
t[:,:,:,:3] = F.normalize(t[:,:,:,:3], dim=3) / 2 + 0.5
return (t,)
class DepthToNormals:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"depth": ("IMAGE",),
"scale": ("FLOAT",{"default": 1, "min": 0.001, "max": 1000, "step": 0.001}),
"output_mode": (["Standard", "BAE", "MiDaS"],),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("normals",)
FUNCTION = "normal_map"
CATEGORY = "Image-Filters/image"
def normal_map(self, depth, scale, output_mode):
kernel_x = torch.Tensor([[0,0,0],[1,0,-1],[0,0,0]]).unsqueeze(0).unsqueeze(0).repeat(3, 1, 1, 1)
kernel_y = torch.Tensor([[0,1,0],[0,0,0],[0,-1,0]]).unsqueeze(0).unsqueeze(0).repeat(3, 1, 1, 1)
conv2d = F.conv2d
pad = F.pad
size_x = depth.size(2)
size_y = depth.size(1)
max_dim = max(size_x, size_y)
position_map = depth.detach().clone() * scale
xs = torch.linspace(-1 * size_x / max_dim, 1 * size_x / max_dim, steps=size_x)
ys = torch.linspace(-1 * size_y / max_dim, 1 * size_y / max_dim, steps=size_y)
grid_x, grid_y = torch.meshgrid(xs, ys, indexing='xy')
position_map[..., 0] = grid_x.unsqueeze(0)
position_map[..., 1] = grid_y.unsqueeze(0)
position_map = position_map.movedim(-1, 1) # BCHW
grad_x = conv2d(pad(position_map, (1,1,1,1), mode='replicate'), kernel_x, padding='valid', groups=3)
grad_y = conv2d(pad(position_map, (1,1,1,1), mode='replicate'), kernel_y, padding='valid', groups=3)
cross_product = torch.cross(grad_x, grad_y, dim=1)
normals = F.normalize(cross_product)
normals[:, 1] *= -1
if output_mode != "Standard":
normals[:, 0] *= -1
if output_mode == "MiDaS":
normals = torch.flip(normals, dims=[1,])
normals = normals.movedim(1, -1) * 0.5 + 0.5 # BHWC
return (normals,)
class Keyer:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"operation": (["luminance", "saturation", "max", "min", "red", "green", "blue", "redscreen", "greenscreen", "bluescreen"],),
"low": ("FLOAT",{"default": 0, "step": 0.001}),
"high": ("FLOAT",{"default": 1, "step": 0.001}),
"gamma": ("FLOAT",{"default": 1.0, "min": 0.001, "step": 0.001}),
"premult": ("BOOLEAN", {"default": True}),
},
}
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK")
RETURN_NAMES = ("image", "alpha", "mask")
FUNCTION = "keyer"
CATEGORY = "Image-Filters/image"
def keyer(self, images, operation, low, high, gamma, premult):
t = images[:,:,:,:3].detach().clone()
if operation == "luminance":
alpha = 0.2126 * t[:,:,:,0] + 0.7152 * t[:,:,:,1] + 0.0722 * t[:,:,:,2]
elif operation == "saturation":
minV = torch.min(t, 3)[0]
maxV = torch.max(t, 3)[0]
mask = maxV != 0
alpha = maxV
alpha[mask] = (maxV[mask] - minV[mask]) / maxV[mask]
elif operation == "max":
alpha = torch.max(t, 3)[0]
elif operation == "min":
alpha = torch.min(t, 3)[0]
elif operation == "red":
alpha = t[:,:,:,0]
elif operation == "green":
alpha = t[:,:,:,1]
elif operation == "blue":
alpha = t[:,:,:,2]
elif operation == "redscreen":
alpha = 0.7 * (t[:,:,:,1] + t[:,:,:,2]) - t[:,:,:,0] + 1
elif operation == "greenscreen":
alpha = 0.7 * (t[:,:,:,0] + t[:,:,:,2]) - t[:,:,:,1] + 1
elif operation == "bluescreen":
alpha = 0.7 * (t[:,:,:,0] + t[:,:,:,1]) - t[:,:,:,2] + 1
else: # should never be reached
alpha = t[:,:,:,0] * 0
if low == high:
alpha = (alpha > high).to(t.dtype)
else:
alpha = (alpha - low) / (high - low)
if gamma != 1.0:
alpha = torch.pow(alpha, 1/gamma)
alpha = torch.clamp(alpha, min=0, max=1).unsqueeze(3).repeat(1,1,1,3)
if premult:
t *= alpha
return (t, alpha, alpha[:,:,:,0])
jitter_matrix = torch.Tensor([[[1, 0, 0], [0, 1, 0]], [[1, 0, 1], [0, 1, 0]], [[1, 0, 1], [0, 1, 1]],
[[1, 0, 0], [0, 1, 1]], [[1, 0,-1], [0, 1, 1]], [[1, 0,-1], [0, 1, 0]],
[[1, 0,-1], [0, 1,-1]], [[1, 0, 0], [0, 1,-1]], [[1, 0, 1], [0, 1,-1]]])
class JitterImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"jitter_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "step": 0.1}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "jitter"
CATEGORY = "Image-Filters/image/jitter"
def jitter(self, images, jitter_scale):
t = images.detach().clone().movedim(-1,1) # [B x C x H x W]
theta = jitter_matrix.detach().clone().to(t.device)
theta[:,0,2] *= jitter_scale * 2 / t.shape[3]
theta[:,1,2] *= jitter_scale * 2 / t.shape[2]
affine = F.affine_grid(theta, torch.Size([9, t.shape[1], t.shape[2], t.shape[3]]))
batch = []
for i in range(t.shape[0]):
jb = t[i].repeat(9,1,1,1)
jb = F.grid_sample(jb, affine, mode='bilinear', padding_mode='border', align_corners=None)
batch.append(jb)
t = torch.cat(batch, dim=0).movedim(1,-1) # [B x H x W x C]
return (t,)
class UnJitterImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"jitter_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "step": 0.1}),
"oflow_align": ("BOOLEAN", {"default": False}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "jitter"
CATEGORY = "Image-Filters/image/jitter"
def jitter(self, images, jitter_scale, oflow_align):
t = images.detach().clone().movedim(-1,1) # [B x C x H x W]
if oflow_align:
pbar = ProgressBar(t.shape[0] // 9)
raft_model, raft_device = load_raft()
batch = []
for i in trange(t.shape[0] // 9):
batch1 = t[i*9].unsqueeze(0).repeat(9,1,1,1)
batch2 = t[i*9:i*9+9]
flows = raft_flow(raft_model, raft_device, batch1, batch2)
batch.append(flows)
pbar.update(1)
flows = torch.cat(batch, dim=0)
t = flow_warp(t, flows)
else:
theta = jitter_matrix.detach().clone().to(t.device)
theta[:,0,2] *= jitter_scale * -2 / t.shape[3]
theta[:,1,2] *= jitter_scale * -2 / t.shape[2]
affine = F.affine_grid(theta, torch.Size([9, t.shape[1], t.shape[2], t.shape[3]]))
batch = []
for i in range(t.shape[0] // 9):
jb = t[i*9:i*9+9]
jb = F.grid_sample(jb, affine, mode='bicubic', padding_mode='border', align_corners=None)
batch.append(jb)
t = torch.cat(batch, dim=0)
t = t.movedim(1,-1) # [B x H x W x C]
return (t,)
class BatchAverageUnJittered:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"operation": (["mean", "median"],),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "apply"
CATEGORY = "Image-Filters/image/jitter"
def apply(self, images, operation):
t = images.detach().clone()
batch = []
for i in range(t.shape[0] // 9):
if operation == "mean":
batch.append(torch.mean(t[i*9:i*9+9], dim=0, keepdim=True))
elif operation == "median":
batch.append(torch.median(t[i*9:i*9+9], dim=0, keepdim=True)[0])
return (torch.cat(batch, dim=0),)
class BatchAlign:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"ref_frame": ("INT", {"default": 0, "min": 0}),
"direction": (["forward", "backward"],),
"blur": ("INT", {"default": 0, "min": 0}),
},
}
RETURN_TYPES = ("IMAGE", "IMAGE")
RETURN_NAMES = ("aligned", "flow")
FUNCTION = "apply"
CATEGORY = "Image-Filters/image"
def apply(self, images, ref_frame, direction, blur):
t = images.detach().clone().movedim(-1,1) # [B x C x H x W]
rf = min(ref_frame, t.shape[0] - 1)
raft_model, raft_device = load_raft()
ref = t[rf].unsqueeze(0).repeat(t.shape[0],1,1,1)
if direction == "forward":
flows = raft_flow(raft_model, raft_device, ref, t)
else:
flows = raft_flow(raft_model, raft_device, t, ref) * -1
if blur > 0:
d = blur * 2 + 1
dup = flows.movedim(1,-1).detach().clone().cpu().numpy()
blurred = []
for img in dup:
blurred.append(torch.from_numpy(cv2.GaussianBlur(img, (d,d), 0)).unsqueeze(0).movedim(-1,1))
flows = torch.cat(blurred).to(flows.device)
t = flow_warp(t, flows)
t = t.movedim(1,-1) # [B x H x W x C]
f = images.detach().clone() * 0
f[:,:,:,:2] = flows.movedim(1,-1)
return (t,f)
class InstructPixToPixConditioningAdvanced:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"positive": ("CONDITIONING", ),
"negative": ("CONDITIONING", ),
"new": ("LATENT", ),
"new_scale": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}),
"original": ("LATENT", ),
"original_scale": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}),
},
}
RETURN_TYPES = ("CONDITIONING","CONDITIONING","CONDITIONING","LATENT")
RETURN_NAMES = ("cond1", "cond2", "negative", "latent")
FUNCTION = "encode"
CATEGORY = "Image-Filters/conditioning"
def encode(self, positive, negative, new, new_scale, original, original_scale):
new_shape, orig_shape = new["samples"].shape, original["samples"].shape
if new_shape != orig_shape:
raise Exception(f"Latent shape mismatch: {new_shape} and {orig_shape}")
out_latent = {}
out_latent["samples"] = new["samples"] * new_scale
out = []
for conditioning in [positive, negative]:
c = []
for t in conditioning:
d = t[1].copy()
d["concat_latent_image"] = original["samples"] * original_scale
n = [t[0], d]
c.append(n)
out.append(c)
return (out[0], out[1], negative, out_latent)
class InpaintConditionEncode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"vae": ("VAE", ),
"pixels": ("IMAGE", ),
"mask": ("MASK", ),
},}
RETURN_TYPES = ("INPAINT_CONDITION",)
RETURN_NAMES = ("inpaint_condition",)
FUNCTION = "encode"
CATEGORY = "Image-Filters/conditioning"
def encode(self, vae, pixels, mask):
x = (pixels.shape[1] // 8) * 8
y = (pixels.shape[2] // 8) * 8
mask = F.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")
pixels = pixels.clone()
if pixels.shape[1] != x or pixels.shape[2] != y:
x_offset = (pixels.shape[1] % 8) // 2
y_offset = (pixels.shape[2] % 8) // 2
pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]
m = (1.0 - mask.round()).squeeze(1)
for i in range(3):
pixels[:,:,:,i] -= 0.5
pixels[:,:,:,i] *= m
pixels[:,:,:,i] += 0.5
concat_latent = vae.encode(pixels)
return ({"concat_latent_image": concat_latent, "concat_mask": mask},)
class InpaintConditionApply:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"positive": ("CONDITIONING", ),
"negative": ("CONDITIONING", ),
"inpaint_condition": ("INPAINT_CONDITION", ),
"noise_mask": ("BOOLEAN", {"default": False, "tooltip": "Add a noise mask to the latent so sampling will only happen within the mask. Might improve results or completely break things depending on the model."}),
},
"optional": {
"latents_optional": ("LATENT",),
},}
RETURN_TYPES = ("CONDITIONING","CONDITIONING","LATENT")
RETURN_NAMES = ("positive", "negative", "latent")
FUNCTION = "encode"
CATEGORY = "Image-Filters/conditioning"
def encode(self, positive, negative, inpaint_condition, noise_mask=True, latents_optional=None):
concat_latent = inpaint_condition["concat_latent_image"]
concat_mask = inpaint_condition["concat_mask"]
if latents_optional is not None:
out_latent = latents_optional.copy()
else:
out_latent = {}
out_latent["samples"] = torch.zeros_like(concat_latent)
if noise_mask:
out_latent["noise_mask"] = concat_mask
out = []
for conditioning in [positive, negative]:
c = node_helpers.conditioning_set_values(conditioning, {"concat_latent_image": concat_latent,
"concat_mask": concat_mask})
out.append(c)
return (out[0], out[1], out_latent)
class LatentNormalizeShuffle:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"latents": ("LATENT", ),
"flatten": ("INT", {"default": 0, "min": 0, "max": 16}),
"normalize": ("BOOLEAN", {"default": True}),
"shuffle": ("BOOLEAN", {"default": True}),
},
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "batch_normalize"
CATEGORY = "Image-Filters/latent"
def batch_normalize(self, latents, flatten, normalize, shuffle):
latents_copy = copy.deepcopy(latents)
t = latents_copy["samples"] # [B x C x H x W]
if flatten > 0:
d = flatten * 2 + 1
channels = t.shape[1]
kernel = gaussian_kernel(d, 1, device=t.device).repeat(channels, 1, 1).unsqueeze(1)
t_blurred = F.conv2d(t, kernel, padding='same', groups=channels)
t = t - t_blurred
if normalize:
for b in range(t.shape[0]):
for c in range(4):
t_sd, t_mean = torch.std_mean(t[b,c])
t[b,c] = (t[b,c] - t_mean) / t_sd
if shuffle:
t_shuffle = []
for i in (1,2,3,0):
t_shuffle.append(t[:,i])
t = torch.stack(t_shuffle, dim=1)
latents_copy["samples"] = t
return (latents_copy,)
class RandnLikeLatent:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"latents": ("LATENT", ),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True, "tooltip": "The random seed used for creating the noise."}),
},
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "generate"
CATEGORY = "Image-Filters/latent"
def generate(self, latents, seed):
latents_copy = copy.deepcopy(latents)
gen_cpu = torch.Generator(device="cpu").manual_seed(seed)
latents_copy["samples"] = randn_like_g(latents_copy["samples"], generator=gen_cpu)
return (latents_copy,)
class PrintSigmas:
@classmethod
def INPUT_TYPES(s):
return {
"required": {"sigmas": ("SIGMAS",)}
}
RETURN_TYPES = ("SIGMAS",)
FUNCTION = "notify"
OUTPUT_NODE = True
CATEGORY = "Image-Filters/utils"
def notify(self, sigmas):
print(sigmas)
return (sigmas,)
class ShowSigmas:
@classmethod
def INPUT_TYPES(s):
return {
"required": {"sigmas": ("SIGMAS",)},
"hidden": {"unique_id": "UNIQUE_ID",},
}
RETURN_TYPES = ("SIGMAS",)
FUNCTION = "show_sigmas"
OUTPUT_NODE = True
CATEGORY = "Image-Filters/utils"
def show_sigmas(self, sigmas, unique_id=None):
if unique_id:
PromptServer.instance.send_progress_text(f"{sigmas}", unique_id)
return (sigmas,)
class VisualizeLatents:
@classmethod
def INPUT_TYPES(s):
return {
"required": {"latent": ("LATENT", ),}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "visualize"
CATEGORY = "Image-Filters/utils"
def visualize(self, latent):
latents = latent["samples"]
batch, channels, height, width = latents.size()
latents = latents - latents.mean()
latents = latents / latents.std()
latents = latents / 10 + 0.5
scale = int(channels ** 0.5)
vis = torch.zeros(batch, height * scale, width * scale)
for i in range(channels):
start_h = (i % scale) * height
end_h = start_h + height
start_w = (i // scale) * width
end_w = start_w + width
vis[:, start_h:end_h, start_w:end_w] = latents[:, i]
return (vis.unsqueeze(-1).repeat(1, 1, 1, 3),)
class GameOfLife:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"width": ("INT", { "default": 32, "min": 8, "max": 1024, "step": 1}),
"height": ("INT", { "default": 32, "min": 8, "max": 1024, "step": 1}),
"cell_size": ("INT", { "default": 16, "min": 8, "max": 1024, "step": 8}),
"seed": ("INT", { "default": 0, "min": 0, "max": 0xffffffffffffffff, "step": 1}),
"threshold": ("FLOAT", { "default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}),
"steps": ("INT", { "default": 64, "min": 1, "max": 999999, "step": 1}),
},
"optional": {
"optional_start": ("MASK", ),
},
}
RETURN_TYPES = ("IMAGE", "MASK", "MASK", "MASK")
RETURN_NAMES = ("image", "mask", "off", "on")
FUNCTION = "game"
CATEGORY = "Image-Filters/image"
def game(self, width, height, cell_size, seed, threshold, steps, optional_start=None):
if optional_start is None:
# base random initialization
torch.manual_seed(seed)
grid = torch.rand(1, 1, height, width)
else:
grid = optional_start[0].unsqueeze(0).unsqueeze(0)
grid = F.interpolate(grid, size=(height, width))
grid = (grid > threshold).type(torch.uint8)
empty = torch.zeros(1, 1, height, width, dtype=torch.uint8)
# neighbor convolution kernel
kernel = torch.ones(1, 1, 3, 3, dtype=torch.uint8)
kernel[0, 0, 1, 1] = 0
game_states = [[], [], []] # grid, turn_off, turn_on
game_states[0].append(grid.detach().clone())
game_states[1].append(empty.detach().clone())
game_states[2].append(empty.detach().clone())
for step in range(steps - 1):
new_state = grid.detach().clone()
neighbors = F.conv2d(F.pad(new_state, pad=(1, 1, 1, 1), mode="circular"), kernel) #, padding="same")
# If a cell is ON and has fewer than two neighbors that are ON, it turns OFF
new_state[(new_state == 1) == (neighbors < 2)] = 0
# If a cell is ON and has either two or three neighbors that are ON, it remains ON.
# If a cell is ON and has more than three neighbors that are ON, it turns OFF.
new_state[(new_state == 1) == (neighbors > 3)] = 0
# If a cell is OFF and has exactly three neighbors that are ON, it turns ON.
new_state[(new_state == 0) == (neighbors == 3)] = 1
turn_off = ((grid - new_state) == 1).type(torch.uint8)
turn_on = ((new_state - grid) == 1).type(torch.uint8)
game_states[0].append(new_state.detach().clone())
game_states[1].append(turn_off.detach().clone())
game_states[2].append(turn_on.detach().clone())
grid = new_state
def postprocess(tensorlist, to_image=False):
game_anim = torch.cat(tensorlist, dim=0).type(torch.float32)
game_anim = F.interpolate(game_anim, size=(height * cell_size, width * cell_size))
game_anim = torch.squeeze(game_anim, dim=1) # BCHW -> BHW
if to_image:
game_anim = game_anim.unsqueeze(-1).repeat(1,1,1,3) # BHWC
return game_anim
image = postprocess(game_states[0], to_image=True)
mask = postprocess(game_states[0])
off = postprocess(game_states[1])
on = postprocess(game_states[2])
return (image, mask, off, on)
modeltest_code_default = """d = model.model.model_config.unet_config
for k in d.keys():
print(k, d[k])"""
class ModelTest:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"code": ("STRING", {"multiline": True, "default": modeltest_code_default}),
},
}
RETURN_TYPES = ()
FUNCTION = "test"
OUTPUT_NODE = True
CATEGORY = "Image-Filters/utils"
def test(self, model, code):
exec(code)
return ()
class ConditioningSubtract:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"cond_orig": ("CONDITIONING", ),
"cond_subtract": ("CONDITIONING", ),
"subtract_strength": ("FLOAT", {"default": 1.0, "step": 0.01}),
},
}
RETURN_TYPES = ("CONDITIONING",)
FUNCTION = "addWeighted"
CATEGORY = "Image-Filters/conditioning"
def addWeighted(self, cond_orig, cond_subtract, subtract_strength):
out = []
if len(cond_subtract) > 1:
logging.warning("Warning: ConditioningSubtract cond_subtract contains more than 1 cond, only the first one will actually be applied to cond_orig.")
cond_from = cond_subtract[0][0]
pooled_output_from = cond_subtract[0][1].get("pooled_output", None)
for i in range(len(cond_orig)):
t1 = cond_orig[i][0]
pooled_output_to = cond_orig[i][1].get("pooled_output", pooled_output_from)
t0 = cond_from[:,:t1.shape[1]]
if t0.shape[1] < t1.shape[1]:
t0 = torch.cat([t0] + [torch.zeros((1, (t1.shape[1] - t0.shape[1]), t1.shape[2]))], dim=1)
tw = t1 - torch.mul(t0, subtract_strength)
t_to = cond_orig[i][1].copy()
if pooled_output_from is not None and pooled_output_to is not None:
t_to["pooled_output"] = pooled_output_to - torch.mul(pooled_output_from, subtract_strength)
elif pooled_output_from is not None:
t_to["pooled_output"] = pooled_output_from
n = [tw, t_to]
out.append(n)
return (out, )
class Noise_CustomNoise:
def __init__(self, noise_latent):
self.seed = 0
self.noise_latent = noise_latent
def generate_noise(self, input_latent):
return self.noise_latent.detach().clone().cpu()
class CustomNoise:
@classmethod
def INPUT_TYPES(s):
return {
"required":{"noise": ("LATENT",),}
}
RETURN_TYPES = ("NOISE",)
FUNCTION = "get_noise"
CATEGORY = "Image-Filters/sampling"
def get_noise(self, noise):
noise_latent = noise["samples"].detach().clone()
std, mean = torch.std_mean(noise_latent, dim=(-2, -1), keepdim=True)
noise_latent = (noise_latent - mean) / std
return (Noise_CustomNoise(noise_latent),)
class ExtractNFrames:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"frames": ("INT", {"default": 16, "min": 2}),
},
"optional": {
"images": ("IMAGE",),
"masks": ("MASK",),
},
}
RETURN_TYPES = ("LIST", "IMAGE", "MASK")
RETURN_NAMES = ("index_list", "images", "masks")
FUNCTION = "extract"
CATEGORY = "Image-Filters/image/frames"
def extract(self, frames, images=None, masks=None):
original_length = 2
if images is not None: original_length = max(original_length, len(images))
if masks is not None: original_length = max(original_length, len(masks))
n = min(original_length, frames)
step = step = (original_length - 1) / (n - 1)
ids = [round(i * step) for i in range(n)]
while len(ids) < frames:
ids.append(ids[-1])
new_images = []
new_masks = []
for i in ids:
if images is not None:
new_images.append(images[min(i, len(images) - 1)].detach().clone())
else:
new_images.append(torch.zeros(512, 512, 3))
if masks is not None:
new_masks.append(masks[min(i, len(masks) - 1)].detach().clone())
else:
new_masks.append(torch.zeros(512, 512))
return (ids, torch.stack(new_images, dim=0), torch.stack(new_masks, dim=0))
class MergeFramesByIndex:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"index_list": ("LIST",),
"orig_images": ("IMAGE",),
"images": ("IMAGE",),
},
"optional": {
"orig_masks": ("MASK",),
"masks": ("MASK",),
},
}
RETURN_TYPES = ("IMAGE", "MASK")
RETURN_NAMES = ("images", "masks")
FUNCTION = "merge"
CATEGORY = "Image-Filters/image/frames"
def merge(self, index_list, orig_images, images, orig_masks=None, masks=None):
new_images = orig_images.detach().clone()
new_masks = torch.ones_like(new_images[..., 0]) # BHW
if orig_masks is not None:
for i in range(len(new_masks)):
new_masks[i] = orig_masks[min(i, len(orig_masks) - 1)].detach().clone()
for i, frame in enumerate(index_list):
frame_mask = masks[i] if masks is not None else torch.ones_like(new_masks[i])
new_images[frame] *= (1 - frame_mask[..., None])
new_images[frame] += images[i].detach().clone() * frame_mask[..., None]
new_masks[frame] *= 0
return (new_images, new_masks)
class Hunyuan3Dv2LatentUpscaleBy:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"samples": ("LATENT",),
"scale_by": ("FLOAT", {"default": 2.0, "min": 0.01, "max": 8.0, "step": 0.01}),
},
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "upscale"
CATEGORY = "Image-Filters/latent"
def upscale(self, samples, scale_by):
s = samples.copy()
size = round(samples["samples"].shape[-1] * scale_by)
s["samples"] = F.interpolate(samples["samples"], size=(size,), mode="nearest-exact")
return (s,)
class PackVideoMask:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
"blend_mode": (["max", "min", "average"], {"default": "max"}),
"causal": ("BOOLEAN", {"default": True, "tooltip": "First latent frame is a single frame"}),
"stride": ("INT", {"default": 4, "min": 1, "tooltip": "downsampling factor to match VAE, ie 4 for Wan, 8 for LTXV"}),
},
}
RETURN_TYPES = ("MASK",)
FUNCTION = "pack_mask"
CATEGORY = "Image-Filters/mask"
def pack_mask(self, mask, blend_mode, causal, stride):
packed_mask = mask.detach().clone()
# repeat first frame to match stride
if causal:
dup_first_frame = packed_mask[0].unsqueeze(0).repeat(stride - 1, 1, 1)
packed_mask = torch.cat([dup_first_frame, packed_mask], dim=0)
# repeat last frame to match stride
remainder = packed_mask.shape[0] % stride
if remainder > 0:
dup_last_frame = packed_mask[-1].unsqueeze(0).repeat(stride - remainder, 1, 1)
packed_mask = torch.cat([packed_mask, dup_last_frame], dim=0)
# shuffle every n frame chunk to channels
B, H, W = packed_mask.shape
packed_mask = packed_mask.reshape(B // stride, stride, H, W)
# squash channels
if blend_mode == "max":
squashed_mask = packed_mask.max(dim=1).values
elif blend_mode == "min":
squashed_mask = packed_mask.min(dim=1).values
else: # average
squashed_mask = packed_mask.mean(dim=1)
return (squashed_mask,)
class PoissonNoise:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"gain": ("FLOAT", {"default": 1000, "min": 0.001, "max": 1_000_000, "step": 0.001}),
"gain_r": ("FLOAT", {"default": 1.0, "min": 0, "max": 1_000_000, "step": 0.001}),
"gain_g": ("FLOAT", {"default": 2.0, "min": 0, "max": 1_000_000, "step": 0.001}),
"gain_b": ("FLOAT", {"default": 0.5, "min": 0, "max": 1_000_000, "step": 0.001}),
"clamp": ("BOOLEAN", {"default": True}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "poissson_noise"
CATEGORY = "Image-Filters/image"
def poissson_noise(self, image, gain, gain_r, gain_g, gain_b, clamp, seed):
linear = sRGBtoLinear_pt(image.cpu().clone())
linear[..., 0] *= gain_r
linear[..., 1] *= gain_g
linear[..., 2] *= gain_b
generator = torch.Generator("cpu").manual_seed(seed)
noise = torch.poisson(linear * gain, generator) * (1 / gain)
noise[..., 0] *= 1 / gain_r
noise[..., 1] *= 1 / gain_g
noise[..., 2] *= 1 / gain_b
output = linearToSRGB_pt(noise)
if clamp: output = torch.clamp(output, min=0, max=1)
return(output,)
COMBINED_MAPPINGS = {
"AdainFilterLatent": (AdainFilterLatent, "AdaIN Filter (Latent)"),
"AdainImage": (AdainImage, "AdaIN (Image)"),
"AdainLatent": (AdainLatent, "AdaIN (Latent)"),
"AlphaClean": (AlphaClean, "Alpha Clean (DEPRECATED, use MaskClean)"),
"AlphaMatte": (AlphaMatte, "Alpha Matte (DEPRECATED, use ImageMatting)"),
"BatchAlign": (BatchAlign, "Batch Align (RAFT)"),
"BatchAverageImage": (BatchAverageImage, "Batch Average Image"),
"BatchAverageUnJittered": (BatchAverageUnJittered, "Batch Average Un-Jittered"),
"BatchNormalizeImage": (BatchNormalizeImage, "Batch Normalize (Image)"),
"BatchNormalizeLatent": (BatchNormalizeLatent, "Batch Normalize (Latent)"),
"BetterFilmGrain": (BetterFilmGrain, "Better Film Grain"),
"BilateralFilterImage": (BilateralFilterImage, "Bilateral Filter Image"),
"BlurImageFast": (BlurImageFast, "Blur Image (Fast)"),
"BlurMaskFast": (BlurMaskFast, "Blur Mask (Fast)"),
"ClampImage": (ClampImage, "Clamp Image"),
"ClampOutliers": (ClampOutliers, "Clamp Outliers"),
"ColorMatchImage": (ColorMatchImage, "Color Match Image"),
"ConditioningSubtract": (ConditioningSubtract, "ConditioningSubtract"),
"ConvertNormals": (ConvertNormals, "Convert Normals"),
"CustomNoise": (CustomNoise, "CustomNoise"),
"DepthToNormals": (DepthToNormals, "Depth To Normals"),
"DifferenceChecker": (DifferenceChecker, "Difference Checker"),
"DilateErodeMask": (DilateErodeMask, "Dilate/Erode Mask"),
"EnhanceDetail": (EnhanceDetail, "Enhance Detail"),
"ExposureAdjust": (ExposureAdjust, "Exposure Adjust"),
"ExtractNFrames": (ExtractNFrames, "Extract N Frames"),
"FrequencyCombine": (FrequencyCombine, "Frequency Combine"),
"FrequencySeparate": (FrequencySeparate, "Frequency Separate"),
"GameOfLife": (GameOfLife, "Game Of Life"),
"GuidedFilterImage": (GuidedFilterImage, "Guided Filter Image"),
"Hunyuan3Dv2LatentUpscaleBy": (Hunyuan3Dv2LatentUpscaleBy, "Upscale Hunyuan3Dv2 Latent By"),
"ImageConstant": (ImageConstant, "Image Constant Color (RGB)"),
"ImageConstantHSV": (ImageConstantHSV, "Image Constant Color (HSV)"),
"ImageMatting": (ImageMatting, "Image Matting"),
"InpaintConditionApply": (InpaintConditionApply, "Inpaint Condition Apply"),
"InpaintConditionEncode": (InpaintConditionEncode, "Inpaint Condition Encode"),
"InstructPixToPixConditioningAdvanced": (InstructPixToPixConditioningAdvanced, "IP2P Conditioning Advanced"),
"JitterImage": (JitterImage, "Jitter Image"),
"Keyer": (Keyer, "Keyer"),
"LatentNormalizeShuffle": (LatentNormalizeShuffle, "LatentNormalizeShuffle"),
"RandnLikeLatent": (RandnLikeLatent, "RandnLikeLatent"),
"LatentStats": (LatentStats, "Latent Stats"),
"MaskClean": (MaskClean, "Mask (Alpha) Clean"),
"MedianFilterImage": (MedianFilterImage, "Median Filter Image"),
"MergeFramesByIndex": (MergeFramesByIndex, "Merge Frames By Index"),
"ModelTest": (ModelTest, "Model Test"),
"NormalMapSimple": (NormalMapSimple, "Normal Map (Simple)"),
"OffsetLatentImage": (OffsetLatentImage, "Offset Latent Image"),
"PackVideoMask": (PackVideoMask, "Pack Video Mask"),
"PoissonNoise": (PoissonNoise, "Poisson Noise Image"),
"PrintSigmas": (PrintSigmas, "Print Sigmas"),
"RelightSimple": (RelightSimple, "Relight (Simple)"),
"RemapRange": (RemapRange, "Remap Range"),
"RestoreDetail": (RestoreDetail, "Restore Detail"),
"SharpenFilterLatent": (SharpenFilterLatent, "Sharpen Filter (Latent)"),
"ShowSigmas": (ShowSigmas, "Show Sigmas"),
"ShuffleChannels": (ShuffleChannels, "Shuffle Channels"),
"Tonemap": (Tonemap, "Tonemap"),
"UnJitterImage": (UnJitterImage, "Un-Jitter Image"),
"UnTonemap": (UnTonemap, "UnTonemap"),
"VisualizeLatents": (VisualizeLatents, "Visualize Latents"),
}
Reference in New Issue View Git Blame Copy Permalink