import math import copy import torch 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 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 = torch.nn.functional.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 torch.nn.functional.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 = torch.nn.functional.interpolate(images.detach().clone().movedim(-1,1), scale_factor=0.1, mode='bilinear').movedim(1,-1).cpu().numpy() np_ref = torch.nn.functional.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 torch.nn.functional.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 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) class AlphaClean: def __init__(self): pass @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" 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 AlphaMatte: def __init__(self): pass @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" def alpha_matte(self, images, alpha_trimap, preblur, blackpoint, whitepoint, max_iterations, estimate_fg): d = preblur * 2 + 1 i_dup = copy.deepcopy(images.cpu().numpy().astype(np.float64)) a_dup = copy.deepcopy(alpha_trimap.cpu().numpy().astype(np.float64)) fg = copy.deepcopy(images.cpu().numpy().astype(np.float64)) bg = copy.deepcopy(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 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" 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 = torch.nn.functional.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: def __init__(self): pass @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" 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: def __init__(self): pass @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 = "mask/filters" 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" 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" 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: def __init__(self): pass @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 = "mask/filters" 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: def __init__(self): pass @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, "min": 0.0, "max": 100.0, "step": 0.1 }), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "enhance" CATEGORY = "image/filters" 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),) # DEPRECATED: use GuidedFilterImage instead class GuidedFilterAlpha: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "images": ("IMAGE",), "alpha": ("IMAGE",), "filter_radius": ("INT", { "default": 8, "min": 1, "max": 64, "step": 1 }), "sigma": ("FLOAT", { "default": 0.1, "min": 0.01, "max": 1.0, "step": 0.01 }), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "guided_filter_alpha" CATEGORY = "image/filters" def guided_filter_alpha(self, images: torch.Tensor, alpha: torch.Tensor, filter_radius: int, sigma: float): d = filter_radius * 2 + 1 s = sigma / 10 i_dup = copy.deepcopy(images.cpu().numpy()) a_dup = copy.deepcopy(alpha.cpu().numpy()) for index, image in enumerate(i_dup): alpha_work = a_dup[index] i_dup[index] = guidedFilter(image, alpha_work, d, s) return (torch.from_numpy(i_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" 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" 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" 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" 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" 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: def __init__(self): pass @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" 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),) 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" 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: def __init__(self): pass @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 = "latent/filters" 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: def __init__(self): pass @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 = "latent/filters" def batch_normalize(self, latents, reference, 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)): for i in range(t.size(1)): r_sd, r_mean = torch.std_mean(reference["samples"][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 latents_copy["samples"] = torch.lerp(latents["samples"], t.movedim(1,0), factor) # [B x C x H x W] return (latents_copy,) class AdainFilterLatent: def __init__(self): pass @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 = "latent/filters" def batch_normalize(self, latents, reference, filter_size, factor): latents_copy = copy.deepcopy(latents) t = latents_copy["samples"].movedim(1, -1) # [B x C x H x W] -> [B x H x W x C] ref = reference["samples"].movedim(1, -1) d = filter_size * 2 + 1 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) # [B x H x W x C] -> [B x C x H x W] latents_copy["samples"] = torch.lerp(latents["samples"], t, factor) return (latents_copy,) class SharpenFilterLatent: def __init__(self): pass @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 = "latent/filters" def filter_latent(self, latents, filter_size, factor): latents_copy = copy.deepcopy(latents) t = latents_copy["samples"].movedim(1, -1) # [B x C x H x W] -> [B x H x W x C] d = filter_size * 2 + 1 t_blurred = cv_blur_tensor(t, d, d) t = t - t_blurred t = t * factor t = t + t_blurred t = t.movedim(-1, 1) # [B x H x W x C] -> [B x C x H x W] latents_copy["samples"] = t return (latents_copy,) class AdainImage: def __init__(self): pass @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" 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: def __init__(self): pass @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 = "latent/filters" 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: def __init__(self): pass @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" 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: def __init__(self): pass @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}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "difference_checker" CATEGORY = "image/filters" def difference_checker(self, images1, images2, multiplier): t = copy.deepcopy(images1) t = torch.abs(images1 - images2) * multiplier return (torch.clamp(t, 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" 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" 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 = "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_dir": ("FLOAT", {"default": 0.0, "min": -1.5, "max": 1.5, "step": 0.01}), "y_dir": ("FLOAT", {"default": 0.0, "min": -1.5, "max": 1.5, "step": 0.01}), "brightness": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100, "step": 0.01}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "relight" CATEGORY = "image/filters" def relight(self, image, normals, x_dir, y_dir, 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 = torch.nn.functional.interpolate(norm.movedim(-1,1), size=(image.shape[1], image.shape[2]), mode='bilinear').movedim(1,-1) light = torch.tensor([x_dir, y_dir, abs(1 - math.sqrt(x_dir ** 2 + y_dir ** 2) * 0.7)]) light = torch.nn.functional.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 = "utils" def notify(self, latent): latents = latent["samples"] width, height = latents.size(3), latents.size(2) text = ["",] text[0] = f"batch size: {latents.size(0)}" text.append(f"width: {width} ({width * 8})") text.append(f"height: {height} ({height * 8})") cmean = [0,0,0,0] for i in range(4): minimum = torch.min(latents[:,i,:,:]).item() maximum = torch.max(latents[:,i,:,:]).item() std_dev, mean = torch.std_mean(latents[:,i,:,:], dim=None) 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" 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" 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"],), "use_tonemap": ("BOOLEAN", {"default": False}), "tonemap_scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 10, "step": 0.01}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "apply" CATEGORY = "image/filters" def apply(self, images, stops, input_mode, output_mode, use_tonemap, tonemap_scale): t = images.detach().clone().cpu().numpy().astype(np.float32) if input_mode == "sRGB": sRGBtoLinear(t[:,:,:,:3]) if use_tonemap: tonemapToLinear(t[:,:,:,:3], tonemap_scale) exposure(t[:,:,:,:3], stops) if use_tonemap: 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,) # 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" 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 = torch.nn.functional.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] = torch.nn.functional.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" 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" 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] = torch.nn.functional.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" 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 = torch.nn.functional.conv2d pad = torch.nn.functional.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 = torch.nn.functional.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" 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/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 = torch.nn.functional.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 = torch.nn.functional.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/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 = torch.nn.functional.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 = torch.nn.functional.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/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" 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 = "conditioning/instructpix2pix" 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 LatentNormalizeShuffle: def __init__(self): pass @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 = "latent/filters" 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 = torch.nn.functional.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 PrintSigmas: @classmethod def INPUT_TYPES(s): return {"required": { "sigmas": ("SIGMAS",) }} RETURN_TYPES = ("SIGMAS",) FUNCTION = "notify" OUTPUT_NODE = True CATEGORY = "utils" def notify(self, sigmas): print(sigmas) return (sigmas,) class VisualizeLatents: @classmethod def INPUT_TYPES(s): return {"required": {"latent": ("LATENT", ),}} RETURN_TYPES = ("IMAGE",) FUNCTION = "visualize" CATEGORY = "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" def game(self, width, height, cell_size, seed, threshold, steps, optional_start=None): F = torch.nn.functional 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) NODE_CLASS_MAPPINGS = { "AdainFilterLatent": AdainFilterLatent, "AdainImage": AdainImage, "AdainLatent": AdainLatent, "AlphaClean": AlphaClean, "AlphaMatte": AlphaMatte, "BatchAlign": BatchAlign, "BatchAverageImage": BatchAverageImage, "BatchAverageUnJittered": BatchAverageUnJittered, "BatchNormalizeImage": BatchNormalizeImage, "BatchNormalizeLatent": BatchNormalizeLatent, "BetterFilmGrain": BetterFilmGrain, "BilateralFilterImage": BilateralFilterImage, "BlurImageFast": BlurImageFast, "BlurMaskFast": BlurMaskFast, "ClampOutliers": ClampOutliers, "ColorMatchImage": ColorMatchImage, "ConvertNormals": ConvertNormals, "DepthToNormals": DepthToNormals, "DifferenceChecker": DifferenceChecker, "DilateErodeMask": DilateErodeMask, "EnhanceDetail": EnhanceDetail, "ExposureAdjust": ExposureAdjust, "FrequencyCombine": FrequencyCombine, "FrequencySeparate": FrequencySeparate, "GameOfLife": GameOfLife, "GuidedFilterAlpha": GuidedFilterAlpha, "GuidedFilterImage": GuidedFilterImage, "ImageConstant": ImageConstant, "ImageConstantHSV": ImageConstantHSV, "InstructPixToPixConditioningAdvanced": InstructPixToPixConditioningAdvanced, "JitterImage": JitterImage, "Keyer": Keyer, "LatentNormalizeShuffle": LatentNormalizeShuffle, "LatentStats": LatentStats, "MedianFilterImage": MedianFilterImage, "NormalMapSimple": NormalMapSimple, "OffsetLatentImage": OffsetLatentImage, "PrintSigmas": PrintSigmas, "RelightSimple": RelightSimple, "RemapRange": RemapRange, "RestoreDetail": RestoreDetail, "SharpenFilterLatent": SharpenFilterLatent, "ShuffleChannels": ShuffleChannels, "Tonemap": Tonemap, "UnJitterImage": UnJitterImage, "UnTonemap": UnTonemap, "VisualizeLatents": VisualizeLatents, } NODE_DISPLAY_NAME_MAPPINGS = { "AdainFilterLatent": "AdaIN Filter (Latent)", "AdainImage": "AdaIN (Image)", "AdainLatent": "AdaIN (Latent)", "AlphaClean": "Alpha Clean", "AlphaMatte": "Alpha Matte", "BatchAlign": "Batch Align (RAFT)", "BatchAverageImage": "Batch Average Image", "BatchAverageUnJittered": "Batch Average Un-Jittered", "BatchNormalizeImage": "Batch Normalize (Image)", "BatchNormalizeLatent": "Batch Normalize (Latent)", "BetterFilmGrain": "Better Film Grain", "BilateralFilterImage": "Bilateral Filter Image", "BlurImageFast": "Blur Image (Fast)", "BlurMaskFast": "Blur Mask (Fast)", "ClampOutliers": "Clamp Outliers", "ColorMatchImage": "Color Match Image", "ConvertNormals": "Convert Normals", "DepthToNormals": "Depth To Normals", "DifferenceChecker": "Difference Checker", "DilateErodeMask": "Dilate/Erode Mask", "EnhanceDetail": "Enhance Detail", "ExposureAdjust": "Exposure Adjust", "FrequencyCombine": "Frequency Combine", "FrequencySeparate": "Frequency Separate", "GameOfLife": "Game Of Life", "GuidedFilterAlpha": "(DEPRECATED) Guided Filter Alpha", "GuidedFilterImage": "Guided Filter Image", "ImageConstant": "Image Constant Color (RGB)", "ImageConstantHSV": "Image Constant Color (HSV)", "InstructPixToPixConditioningAdvanced": "InstructPixToPixConditioningAdvanced", "JitterImage": "Jitter Image", "Keyer": "Keyer", "LatentNormalizeShuffle": "LatentNormalizeShuffle", "LatentStats": "Latent Stats", "MedianFilterImage": "Median Filter Image", "NormalMapSimple": "Normal Map (Simple)", "OffsetLatentImage": "Offset Latent Image", "PrintSigmas": "PrintSigmas", "RelightSimple": "Relight (Simple)", "RemapRange": "Remap Range", "RestoreDetail": "Restore Detail", "SharpenFilterLatent": "Sharpen Filter (Latent)", "ShuffleChannels": "Shuffle Channels", "Tonemap": "Tonemap", "UnJitterImage": "Un-Jitter Image", "UnTonemap": "UnTonemap", "VisualizeLatents": "Visualize Latents", }