Files
cubiq-ComfyUI_essentials/essentials.py
T
2024-01-28 20:30:41 +01:00

1226 lines
40 KiB
Python

import warnings
warnings.filterwarnings('ignore', module="torchvision")
import ast
import math
import random
import os
import operator as op
import numpy as np
import torch
import torch.nn.functional as F
import torchvision.transforms.v2 as T
from nodes import MAX_RESOLUTION, SaveImage
import folder_paths
import comfy.utils
def p(image):
return image.permute([0,3,1,2])
def pb(image):
return image.permute([0,2,3,1])
# from https://github.com/pythongosssss/ComfyUI-Custom-Scripts
class AnyType(str):
def __ne__(self, __value: object) -> bool:
return False
any = AnyType("*")
EPSILON = 1e-5
class GetImageSize:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
}
}
RETURN_TYPES = ("INT", "INT")
RETURN_NAMES = ("width", "height")
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image):
return (image.shape[2], image.shape[1],)
class ImageResize:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }),
"height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }),
"interpolation": (["nearest", "bilinear", "bicubic", "area", "nearest-exact", "lanczos"],),
"keep_proportion": ("BOOLEAN", { "default": False }),
"condition": (["always", "only if bigger", "only if smaller"],),
}
}
RETURN_TYPES = ("IMAGE", "INT", "INT",)
RETURN_NAMES = ("IMAGE", "width", "height",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, width, height, keep_proportion, interpolation="nearest", condition="always"):
if keep_proportion is True:
_, oh, ow, _ = image.shape
if width == 0 and oh < height:
width = MAX_RESOLUTION
elif width == 0 and oh >= height:
width = ow
if height == 0 and ow < width:
height = MAX_RESOLUTION
elif height == 0 and ow >= width:
height = ow
#width = ow if width == 0 else width
#height = oh if height == 0 else height
ratio = min(width / ow, height / oh)
width = round(ow*ratio)
height = round(oh*ratio)
outputs = p(image)
if "always" in condition or ("bigger" in condition and (oh > height or ow > width)) or ("smaller" in condition and (oh < height or ow < width)):
if interpolation == "lanczos":
outputs = comfy.utils.lanczos(outputs, width, height)
else:
outputs = F.interpolate(outputs, size=(height, width), mode=interpolation)
outputs = pb(outputs)
return(outputs, outputs.shape[2], outputs.shape[1],)
class ImageFlip:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"axis": (["x", "y", "xy"],),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, axis):
dim = ()
if "y" in axis:
dim += (1,)
if "x" in axis:
dim += (2,)
image = torch.flip(image, dim)
return(image,)
class ImageCrop:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"width": ("INT", { "default": 256, "min": 0, "max": MAX_RESOLUTION, "step": 8, }),
"height": ("INT", { "default": 256, "min": 0, "max": MAX_RESOLUTION, "step": 8, }),
"position": (["top-left", "top-center", "top-right", "right-center", "bottom-right", "bottom-center", "bottom-left", "left-center", "center"],),
"x_offset": ("INT", { "default": 0, "min": -99999, "step": 1, }),
"y_offset": ("INT", { "default": 0, "min": -99999, "step": 1, }),
}
}
RETURN_TYPES = ("IMAGE","INT","INT",)
RETURN_NAMES = ("IMAGE","x","y",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, width, height, position, x_offset, y_offset):
_, oh, ow, _ = image.shape
width = min(ow, width)
height = min(oh, height)
if "center" in position:
x = round((ow-width) / 2)
y = round((oh-height) / 2)
if "top" in position:
y = 0
if "bottom" in position:
y = oh-height
if "left" in position:
x = 0
if "right" in position:
x = ow-width
x += x_offset
y += y_offset
x2 = x+width
y2 = y+height
if x2 > ow:
x2 = ow
if x < 0:
x = 0
if y2 > oh:
y2 = oh
if y < 0:
y = 0
image = image[:, y:y2, x:x2, :]
return(image, x, y, )
class ImageDesaturate:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"factor": ("FLOAT", { "default": 1.00, "min": 0.00, "max": 1.00, "step": 0.05, }),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, factor):
grayscale = 0.299 * image[..., 0] + 0.587 * image[..., 1] + 0.114 * image[..., 2]
grayscale = (1.0 - factor) * image + factor * grayscale.unsqueeze(-1).repeat(1, 1, 1, 3)
return(grayscale,)
class ImagePosterize:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"threshold": ("FLOAT", { "default": 0.50, "min": 0.00, "max": 1.00, "step": 0.05, }),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, threshold):
image = 0.299 * image[..., 0] + 0.587 * image[..., 1] + 0.114 * image[..., 2]
#image = image.mean(dim=3, keepdim=True)
image = (image > threshold).float()
image = image.unsqueeze(-1).repeat(1, 1, 1, 3)
return(image,)
class ImageEnhanceDifference:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image1": ("IMAGE",),
"image2": ("IMAGE",),
"exponent": ("FLOAT", { "default": 0.75, "min": 0.00, "max": 1.00, "step": 0.05, }),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image1, image2, exponent):
if image1.shape != image2.shape:
image2 = p(image2)
image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center')
image2 = pb(image2)
diff_image = image1 - image2
diff_image = torch.pow(diff_image, exponent)
diff_image = torch.clamp(diff_image, 0, 1)
return(diff_image,)
class ImageExpandBatch:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"size": ("INT", { "default": 16, "min": 1, "step": 1, }),
"method": (["expand", "repeat all", "repeat first", "repeat last"],)
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, size, method):
orig_size = image.shape[0]
if orig_size == size:
return (image,)
if size <= 1:
return (image[:size],)
if 'expand' in method:
out = torch.empty([size] + list(image.shape)[1:], dtype=image.dtype, device=image.device)
if size < orig_size:
scale = (orig_size - 1) / (size - 1)
for i in range(size):
out[i] = image[min(round(i * scale), orig_size - 1)]
else:
scale = orig_size / size
for i in range(size):
out[i] = image[min(math.floor((i + 0.5) * scale), orig_size - 1)]
elif 'all' in method:
out = image.repeat([math.ceil(size / image.shape[0])] + [1] * (len(image.shape) - 1))[:size]
elif 'first' in method:
if size < image.shape[0]:
out = image[:size]
else:
out = torch.cat([image[:1].repeat(size-image.shape[0], 1, 1, 1), image], dim=0)
elif 'last' in method:
if size < image.shape[0]:
out = image[:size]
else:
out = torch.cat((image, image[-1:].repeat((size-image.shape[0], 1, 1, 1))), dim=0)
return (out,)
class ExtractKeyframes:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"threshold": ("FLOAT", { "default": 0.85, "min": 0.00, "max": 1.00, "step": 0.01, }),
}
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("KEYFRAMES", "indexes")
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, threshold):
window_size = 2
variations = torch.sum(torch.abs(image[1:] - image[:-1]), dim=[1, 2, 3])
#variations = torch.sum((image[1:] - image[:-1]) ** 2, dim=[1, 2, 3])
threshold = torch.quantile(variations.float(), threshold).item()
keyframes = []
for i in range(image.shape[0] - window_size + 1):
window = image[i:i + window_size]
variation = torch.sum(torch.abs(window[-1] - window[0])).item()
if variation > threshold:
keyframes.append(i + window_size - 1)
return (image[keyframes], ','.join(map(str, keyframes)),)
"""
class NoiseFromImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"adjust_levels": ("FLOAT", { "default": 1.00, "min": 0.00, "max": 20.00, "step": 0.05, }),
#"noise_intensity": ("FLOAT", { "default": 1.00, "min": 0.00, "max": 1.00, "step": 0.05, }),
"noise_density": ("FLOAT", { "default": 0.05, "min": 0.00, "max": 1.00, "step": 0.05, }),
"noise_scale": ("FLOAT", { "default": 0.2, "min": 0.00, "max": 1.00, "step": 0.05, }),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, noise_seed, adjust_levels, noise_density, noise_scale):
generator = torch.manual_seed(noise_seed)
image = image.mean(dim=3).unsqueeze(-1).repeat(1, 1, 1, 3)
# Adjust image levels
image = (1 - adjust_levels) * torch.mean(image) + adjust_levels * image
image = torch.clamp(image, 0, 1)
# Create noise
fine_noise = torch.rand([image.shape[0], image.shape[1], image.shape[2], image.shape[3]], dtype=image.dtype, layout=image.layout, generator=generator, device="cpu")
fine_noise = fine_noise * (fine_noise > 1-noise_density).float() # Lower density
fine_noise = (fine_noise * 16).round() / 16
coarse_noise = F.interpolate(p(fine_noise), scale_factor=noise_scale, mode='bilinear', align_corners=False)
coarse_noise = F.interpolate(coarse_noise, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False)
coarse_noise = pb(coarse_noise)
# Merge noises
noise = ((1 - image) * coarse_noise + image * fine_noise)
noise = torch.clamp(noise, 0, 1)
noise = image * noise
# Change noise intensity
#noise = noise * noise_intensity
#print(noise.min(), noise.max())
#noise = torch.clamp(noise, 0, 1)
# Apply noise to image
#noise = torch.clamp((1-noise_intensity) * image + noise, 0, 1)
#out = image + fine_noise * mask * noise_intensity
return (noise,)
"""
class MaskFlip:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
"axis": (["x", "y", "xy"],),
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, mask, axis):
dim = ()
if "y" in axis:
dim += (1,)
if "x" in axis:
dim += (2,)
mask = torch.flip(mask, dims=dim)
return(mask,)
class MaskBlur:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
"amount": ("FLOAT", { "default": 6.0, "min": 0, "step": 0.5, }),
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, mask, amount):
size = int(6 * amount +1)
if size % 2 == 0:
size+= 1
blurred = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 1)
blurred = p(blurred)
blurred = T.GaussianBlur(size, amount)(blurred)
blurred = pb(blurred)
blurred = blurred[:, :, :, 0]
return(blurred,)
class MaskPreview(SaveImage):
def __init__(self):
self.output_dir = folder_paths.get_temp_directory()
self.type = "temp"
self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
self.compress_level = 4
@classmethod
def INPUT_TYPES(s):
return {
"required": {"mask": ("MASK",), },
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
}
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, mask, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
return self.save_images(preview, filename_prefix, prompt, extra_pnginfo)
class MaskBatch:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask1": ("MASK",),
"mask2": ("MASK",),
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, mask1, mask2):
if mask1.shape[1:] != mask2.shape[1:]:
mask2 = F.interpolate(mask2.unsqueeze(1), size=(mask1.shape[1], mask1.shape[2]), mode="bicubic").squeeze(1)
out = torch.cat((mask1, mask2), dim=0)
return (out,)
class MaskExpandBatch:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
"size": ("INT", { "default": 16, "min": 1, "step": 1, }),
"method": (["expand", "repeat all", "repeat first", "repeat last"],)
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, mask, size, method):
orig_size = mask.shape[0]
if orig_size == size:
return (mask,)
if size <= 1:
return (mask[:size],)
if 'expand' in method:
out = torch.empty([size] + list(mask.shape)[1:], dtype=mask.dtype, device=mask.device)
if size < orig_size:
scale = (orig_size - 1) / (size - 1)
for i in range(size):
out[i] = mask[min(round(i * scale), orig_size - 1)]
else:
scale = orig_size / size
for i in range(size):
out[i] = mask[min(math.floor((i + 0.5) * scale), orig_size - 1)]
elif 'all' in method:
out = mask.repeat([math.ceil(size / mask.shape[0])] + [1] * (len(mask.shape) - 1))[:size]
elif 'first' in method:
if size < mask.shape[0]:
out = mask[:size]
else:
out = torch.cat([mask[:1].repeat(size-mask.shape[0], 1, 1), mask], dim=0)
elif 'last' in method:
if size < mask.shape[0]:
out = mask[:size]
else:
out = torch.cat((mask, mask[-1:].repeat((size-mask.shape[0], 1, 1))), dim=0)
return (out,)
def cubic_bezier(t, p):
p0, p1, p2, p3 = p
return (1 - t)**3 * p0 + 3 * (1 - t)**2 * t * p1 + 3 * (1 - t) * t**2 * p2 + t**3 * p3
class MaskFromColor:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"red": ("INT", { "default": 255, "min": 0, "max": 255, "step": 1, }),
"green": ("INT", { "default": 255, "min": 0, "max": 255, "step": 1, }),
"blue": ("INT", { "default": 255, "min": 0, "max": 255, "step": 1, }),
"threshold": ("INT", { "default": 0, "min": 0, "max": 127, "step": 1, }),
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, red, green, blue, threshold):
temp = (torch.clamp(image, 0, 1.0) * 255.0).round().to(torch.int)
color = torch.tensor([red, green, blue])
lower_bound = (color - threshold).clamp(min=0)
upper_bound = (color + threshold).clamp(max=255)
lower_bound = lower_bound.view(1, 1, 1, 3)
upper_bound = upper_bound.view(1, 1, 1, 3)
mask = (temp >= lower_bound) & (temp <= upper_bound)
mask = mask.all(dim=-1)
mask = mask.float()
return (mask, )
class MaskFromBatch:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK", ),
"start": ("INT", { "default": 0, "min": 0, "step": 1, }),
"length": ("INT", { "default": -1, "min": -1, "step": 1, }),
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, mask, start, length):
if length<0:
length = mask.shape[0]
start = min(start, mask.shape[0]-1)
length = min(mask.shape[0]-start, length)
return (mask[start:start + length], )
class ImageFromBatch:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"start": ("INT", { "default": 0, "min": 0, "step": 1, }),
"length": ("INT", { "default": -1, "min": -1, "step": 1, }),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, start, length):
if length<0:
length = image.shape[0]
start = min(start, image.shape[0]-1)
length = min(image.shape[0]-start, length)
return (image[start:start + length], )
class ImageCompositeFromMaskBatch:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image_from": ("IMAGE", ),
"image_to": ("IMAGE", ),
"mask": ("MASK", )
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image_from, image_to, mask):
frames = mask.shape[0]
if image_from.shape[1] != image_to.shape[1] or image_from.shape[2] != image_to.shape[2]:
image_to = p(image_to)
image_to = comfy.utils.common_upscale(image_to, image_from.shape[2], image_from.shape[1], upscale_method='bicubic', crop='center')
image_to = pb(image_to)
if frames < image_from.shape[0]:
image_from = image_from[:frames]
elif frames > image_from.shape[0]:
image_from = torch.cat((image_from, image_from[-1].unsqueeze(0).repeat(frames-image_from.shape[0], 1, 1, 1)), dim=0)
mask = mask.unsqueeze(3).repeat(1, 1, 1, 3)
if image_from.shape[1] != mask.shape[1] or image_from.shape[2] != mask.shape[2]:
mask = p(mask)
mask = comfy.utils.common_upscale(mask, image_from.shape[2], image_from.shape[1], upscale_method='bicubic', crop='center')
mask = pb(mask)
out = mask * image_to + (1 - mask) * image_from
return (out, )
class TransitionMask:
@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, }),
"frames": ("INT", { "default": 16, "min": 1, "max": 9999, "step": 1, }),
"start_frame": ("INT", { "default": 0, "min": 0, "step": 1, }),
"end_frame": ("INT", { "default": 9999, "min": 0, "step": 1, }),
"transition_type": (["horizontal slide", "vertical slide", "horizontal bar", "vertical bar", "center box", "horizontal door", "vertical door", "circle", "fade"],),
"timing_function": (["linear", "in", "out", "in-out"],)
}
}
RETURN_TYPES = ("MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, width, height, frames, start_frame, end_frame, transition_type, timing_function):
if timing_function == 'in':
tf = [0.0, 0.0, 0.5, 1.0]
elif timing_function == 'out':
tf = [0.0, 0.5, 1.0, 1.0]
elif timing_function == 'in-out':
tf = [0, 1, 0, 1]
#elif timing_function == 'back':
# tf = [0, 1.334, 1.334, 0]
else:
tf = [0, 0, 1, 1]
out = []
end_frame = min(frames, end_frame)
transition = end_frame - start_frame
if start_frame > 0:
out = out + [torch.full((height, width), 0.0, dtype=torch.float32, device="cpu")] * start_frame
for i in range(transition):
frame = torch.full((height, width), 0.0, dtype=torch.float32, device="cpu")
progress = i/(transition-1)
if timing_function != 'linear':
progress = cubic_bezier(progress, tf)
if "horizontal slide" in transition_type:
pos = round(width*progress)
frame[:, :pos] = 1.0
elif "vertical slide" in transition_type:
pos = round(height*progress)
frame[:pos, :] = 1.0
elif "box" in transition_type:
box_w = round(width*progress)
box_h = round(height*progress)
x1 = (width - box_w) // 2
y1 = (height - box_h) // 2
x2 = x1 + box_w
y2 = y1 + box_h
frame[y1:y2, x1:x2] = 1.0
elif "circle" in transition_type:
radius = math.ceil(math.sqrt(pow(width,2)+pow(height,2))*progress/2)
c_x = width // 2
c_y = height // 2
# is this real life? Am I hallucinating?
x = torch.arange(0, width, dtype=torch.float32, device="cpu")
y = torch.arange(0, height, dtype=torch.float32, device="cpu")
y, x = torch.meshgrid((y, x), indexing="ij")
circle = ((x - c_x) ** 2 + (y - c_y) ** 2) <= (radius ** 2)
frame[circle] = 1.0
elif "horizontal bar" in transition_type:
bar = round(height*progress)
y1 = (height - bar) // 2
y2 = y1 + bar
frame[y1:y2, :] = 1.0
elif "vertical bar" in transition_type:
bar = round(width*progress)
x1 = (width - bar) // 2
x2 = x1 + bar
frame[:, x1:x2] = 1.0
elif "horizontal door" in transition_type:
bar = math.ceil(height*progress/2)
if bar > 0:
frame[:bar, :] = 1.0
frame[-bar:, :] = 1.0
elif "vertical door" in transition_type:
bar = math.ceil(width*progress/2)
if bar > 0:
frame[:, :bar] = 1.0
frame[:, -bar:] = 1.0
elif "fade" in transition_type:
frame[:,:] = progress
out.append(frame)
if end_frame < frames:
out = out + [torch.full((height, width), 1.0, dtype=torch.float32, device="cpu")] * (frames - end_frame)
out = torch.stack(out, dim=0)
return (out, )
def min_(tensor_list):
# return the element-wise min of the tensor list.
x = torch.stack(tensor_list)
mn = x.min(axis=0)[0]
return torch.clamp(mn, min=0)
def max_(tensor_list):
# return the element-wise max of the tensor list.
x = torch.stack(tensor_list)
mx = x.max(axis=0)[0]
return torch.clamp(mx, max=1)
# From https://github.com/Jamy-L/Pytorch-Contrast-Adaptive-Sharpening/
class ImageCAS:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"amount": ("FLOAT", {"default": 0.8, "min": 0, "max": 1, "step": 0.05}),
},
}
RETURN_TYPES = ("IMAGE",)
CATEGORY = "essentials"
FUNCTION = "execute"
def execute(self, image, amount):
img = F.pad(p(image), pad=(1, 1, 1, 1)).cpu()
a = img[..., :-2, :-2]
b = img[..., :-2, 1:-1]
c = img[..., :-2, 2:]
d = img[..., 1:-1, :-2]
e = img[..., 1:-1, 1:-1]
f = img[..., 1:-1, 2:]
g = img[..., 2:, :-2]
h = img[..., 2:, 1:-1]
i = img[..., 2:, 2:]
# Computing contrast
cross = (b, d, e, f, h)
mn = min_(cross)
mx = max_(cross)
diag = (a, c, g, i)
mn2 = min_(diag)
mx2 = max_(diag)
mx = mx + mx2
mn = mn + mn2
# Computing local weight
inv_mx = torch.reciprocal(mx + EPSILON)
amp = inv_mx * torch.minimum(mn, (2 - mx))
# scaling
amp = torch.sqrt(amp)
w = - amp * (amount * (1/5 - 1/8) + 1/8)
div = torch.reciprocal(1 + 4*w)
output = ((b + d + f + h)*w + e) * div
output = output.clamp(0, 1)
#output = torch.nan_to_num(output) # this seems the only way to ensure there are no NaNs
output = pb(output)
return (output,)
operators = {
ast.Add: op.add,
ast.Sub: op.sub,
ast.Mult: op.mul,
ast.Div: op.truediv,
ast.FloorDiv: op.floordiv,
ast.Pow: op.pow,
ast.BitXor: op.xor,
ast.USub: op.neg,
ast.Mod: op.mod,
}
op_functions = {
'min': min,
'max': max
}
class SimpleMath:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"optional": {
"a": ("INT,FLOAT", { "default": 0.0, "step": 0.1 }),
"b": ("INT,FLOAT", { "default": 0.0, "step": 0.1 }),
},
"required": {
"value": ("STRING", { "multiline": False, "default": "" }),
},
}
RETURN_TYPES = ("INT", "FLOAT", )
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, value, a = 0.0, b = 0.0):
def eval_(node):
if isinstance(node, ast.Num): # number
return node.n
elif isinstance(node, ast.Name): # variable
if node.id == "a":
return a
if node.id == "b":
return b
elif isinstance(node, ast.BinOp): # <left> <operator> <right>
return operators[type(node.op)](eval_(node.left), eval_(node.right))
elif isinstance(node, ast.UnaryOp): # <operator> <operand> e.g., -1
return operators[type(node.op)](eval_(node.operand))
elif isinstance(node, ast.Call): # custom function
if node.func.id in op_functions:
args =[eval_(arg) for arg in node.args]
return op_functions[node.func.id](*args)
else:
return 0
result = eval_(ast.parse(value, mode='eval').body)
if math.isnan(result):
result = 0.0
return (round(result), result, )
class ModelCompile():
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"fullgraph": ("BOOLEAN", { "default": False }),
"dynamic": ("BOOLEAN", { "default": False }),
"mode": (["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"],),
},
}
RETURN_TYPES = ("MODEL", )
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, model, fullgraph, dynamic, mode):
work_model = model.clone()
torch._dynamo.config.suppress_errors = True
work_model.model.diffusion_model = torch.compile(work_model.model.diffusion_model, dynamic=dynamic, fullgraph=fullgraph, mode=mode)
return( work_model, )
class ConsoleDebug:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"value": (any, {}),
},
"optional": {
"prefix": ("STRING", { "multiline": False, "default": "Value:" })
}
}
RETURN_TYPES = ()
FUNCTION = "execute"
CATEGORY = "essentials"
OUTPUT_NODE = True
def execute(self, value, prefix):
print(f"\033[96m{prefix} {value}\033[0m")
return (None,)
class DebugTensorShape:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"tensor": (any, {}),
},
}
RETURN_TYPES = ()
FUNCTION = "execute"
CATEGORY = "essentials"
OUTPUT_NODE = True
def execute(self, tensor):
shapes = []
def tensorShape(tensor):
if isinstance(tensor, dict):
for k in tensor:
tensorShape(tensor[k])
elif isinstance(tensor, list):
for i in range(len(tensor)):
tensorShape(tensor[i])
elif hasattr(tensor, 'shape'):
shapes.append(list(tensor.shape))
tensorShape(tensor)
print(f"\033[96mShapes found: {shapes}\033[0m")
return (None,)
class BatchCount:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"batch": (any, {}),
},
}
RETURN_TYPES = ("INT",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, batch):
count = 0
if hasattr(batch, 'shape'):
count = batch.shape[0]
elif isinstance(batch, dict) and 'samples' in batch:
count = batch['samples'].shape[0]
elif isinstance(batch, list) or isinstance(batch, dict):
count = len(batch)
return (count, )
class ImageSeamCarving:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"width": ("INT", { "default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1, }),
"height": ("INT", { "default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1, }),
"energy": (["backward", "forward"],),
"order": (["width-first", "height-first"],),
},
"optional": {
"keep_mask": ("MASK",),
"drop_mask": ("MASK",),
}
}
RETURN_TYPES = ("IMAGE",)
CATEGORY = "essentials"
FUNCTION = "execute"
def execute(self, image, width, height, energy, order, keep_mask=None, drop_mask=None):
try:
from .carve import seam_carving
except ImportError as e:
raise Exception(e)
img = p(image)
if keep_mask is not None:
#keep_mask = keep_mask.reshape((-1, 1, keep_mask.shape[-2], keep_mask.shape[-1])).movedim(1, -1)
keep_mask = p(keep_mask.unsqueeze(-1))
if keep_mask.shape[2] != img.shape[2] or keep_mask.shape[3] != img.shape[3]:
keep_mask = F.interpolate(keep_mask, size=(img.shape[2], img.shape[3]), mode="bilinear")
if drop_mask is not None:
drop_mask = p(drop_mask.unsqueeze(-1))
if drop_mask.shape[2] != img.shape[2] or drop_mask.shape[3] != img.shape[3]:
drop_mask = F.interpolate(drop_mask, size=(img.shape[2], img.shape[3]), mode="bilinear")
out = []
for i in range(img.shape[0]):
resized = seam_carving(
T.ToPILImage()(img[i]),
size=(width, height),
energy_mode=energy,
order=order,
keep_mask=T.ToPILImage()(keep_mask[i]) if keep_mask is not None else None,
drop_mask=T.ToPILImage()(drop_mask[i]) if drop_mask is not None else None,
)
out.append(T.ToTensor()(resized))
out = torch.stack(out)
out = pb(out)
return(out, )
class CLIPTextEncodeSDXLSimplified:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
"height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
"text": ("STRING", {"multiline": True, "default": ""}),
"clip": ("CLIP", ),
}}
RETURN_TYPES = ("CONDITIONING",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, clip, width, height, text):
crop_w = 0
crop_h = 0
width = width*4
height = height*4
target_width = width
target_height = height
text_g = text_l = text
tokens = clip.tokenize(text_g)
tokens["l"] = clip.tokenize(text_l)["l"]
if len(tokens["l"]) != len(tokens["g"]):
empty = clip.tokenize("")
while len(tokens["l"]) < len(tokens["g"]):
tokens["l"] += empty["l"]
while len(tokens["l"]) > len(tokens["g"]):
tokens["g"] += empty["g"]
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
return ([[cond, {"pooled_output": pooled, "width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}]], )
class SDXLResolutionPicker:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"resolution": (["704x1408 (0.5)","704x1344 (0.52)","768x1344 (0.57)","768x1280 (0.6)","832x1216 (0.68)","832x1152 (0.72)","896x1152 (0.78)","896x1088 (0.82)","960x1088 (0.88)","960x1024 (0.94)","1024x1024 (1.0)","1024x960 (1.07)","1088x960 (1.13)","1088x896 (1.21)","1152x896 (1.29)","1152x832 (1.38)","1216x832 (1.46)","1280x768 (1.67)","1344x768 (1.75)","1344x704 (1.91)","1408x704 (2.0)","1472x704 (2.09)","1536x640 (2.4)","1600x640 (2.5)","1664x576 (2.89)","1728x576 (3.0)",], {"default": "1024x1024 (1.0)"}),
}}
RETURN_TYPES = ("INT","INT",)
RETURN_NAMES = ("width", "height",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, resolution):
width, height = resolution.split(" ")[0].split("x")
return (width, height,)
LUTS_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "luts")
# From https://github.com/yoonsikp/pycubelut/blob/master/pycubelut.py (MIT license)
class ImageApplyLUT:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"lut_file": ([f for f in os.listdir(LUTS_DIR) if f.endswith('.cube')], ),
"log_colorspace": ("BOOLEAN", { "default": False }),
"clip_values": ("BOOLEAN", { "default": False }),
}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
# TODO: check if we can do without numpy
def execute(self, image, lut_file, log_colorspace, clip_values):
from colour.io.luts.iridas_cube import read_LUT_IridasCube
lut = read_LUT_IridasCube(os.path.join(LUTS_DIR, lut_file))
lut.name = lut_file
if clip_values:
if lut.domain[0].max() == lut.domain[0].min() and lut.domain[1].max() == lut.domain[1].min():
lut.table = np.clip(lut.table, lut.domain[0, 0], lut.domain[1, 0])
else:
if len(lut.table.shape) == 2: # 3x1D
for dim in range(3):
lut.table[:, dim] = np.clip(lut.table[:, dim], lut.domain[0, dim], lut.domain[1, dim])
else: # 3D
for dim in range(3):
lut.table[:, :, :, dim] = np.clip(lut.table[:, :, :, dim], lut.domain[0, dim], lut.domain[1, dim])
out = []
for img in image: # TODO: is this more resrouce efficient?
img = img.numpy().copy()
is_non_default_domain = not np.array_equal(lut.domain, np.array([[0., 0., 0.], [1., 1., 1.]]))
dom_scale = None
if is_non_default_domain:
dom_scale = lut.domain[1] - lut.domain[0]
img = img * dom_scale + lut.domain[0]
if log_colorspace:
img = img ** (1/2.2)
img = lut.apply(img)
if log_colorspace:
img = img ** (2.2)
if is_non_default_domain:
img = (img - lut.domain[0]) / dom_scale
img = torch.from_numpy(img)
out.append(img)
out = torch.stack(out)
return (out, )
NODE_CLASS_MAPPINGS = {
"GetImageSize+": GetImageSize,
"ImageResize+": ImageResize,
"ImageCrop+": ImageCrop,
"ImageFlip+": ImageFlip,
"ImageDesaturate+": ImageDesaturate,
"ImagePosterize+": ImagePosterize,
"ImageCASharpening+": ImageCAS,
"ImageSeamCarving+": ImageSeamCarving,
"ImageEnhanceDifference+": ImageEnhanceDifference,
"ImageExpandBatch+": ImageExpandBatch,
"ImageFromBatch+": ImageFromBatch,
"ImageCompositeFromMaskBatch+": ImageCompositeFromMaskBatch,
"ExtractKeyframes+": ExtractKeyframes,
"ImageApplyLUT+": ImageApplyLUT,
#"NoiseFromImage+": NoiseFromImage,
"MaskBlur+": MaskBlur,
"MaskFlip+": MaskFlip,
"MaskPreview+": MaskPreview,
"MaskBatch+": MaskBatch,
"MaskExpandBatch+": MaskExpandBatch,
"TransitionMask+": TransitionMask,
"MaskFromColor+": MaskFromColor,
"MaskFromBatch+": MaskFromBatch,
"SimpleMath+": SimpleMath,
"ConsoleDebug+": ConsoleDebug,
"DebugTensorShape+": DebugTensorShape,
"ModelCompile+": ModelCompile,
"BatchCount+": BatchCount,
"CLIPTextEncodeSDXL+": CLIPTextEncodeSDXLSimplified,
"SDXLResolutionPicker+": SDXLResolutionPicker,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"GetImageSize+": "🔧 Get Image Size",
"ImageResize+": "🔧 Image Resize",
"ImageCrop+": "🔧 Image Crop",
"ImageFlip+": "🔧 Image Flip",
"ImageDesaturate+": "🔧 Image Desaturate",
"ImagePosterize+": "🔧 Image Posterize",
"ImageCASharpening+": "🔧 Image Contrast Adaptive Sharpening",
"ImageSeamCarving+": "🔧 Image Seam Carving",
"ImageEnhanceDifference+": "🔧 Image Enhance Difference",
"ImageExpandBatch+": "🔧 Image Expand Batch",
"ImageFromBatch+": "🔧 Image From Batch",
"ImageCompositeFromMaskBatch+": "🔧 Image Composite From Mask Batch",
"ExtractKeyframes+": "🔧 Extract Keyframes (experimental)",
"ImageApplyLUT+": "🔧 Image Apply LUT",
#"NoiseFromImage+": "🔧 Noise From Image",
"MaskBlur+": "🔧 Mask Blur",
"MaskFlip+": "🔧 Mask Flip",
"MaskPreview+": "🔧 Mask Preview",
"MaskBatch+": "🔧 Mask Batch",
"MaskExpandBatch+": "🔧 Mask Expand Batch",
"TransitionMask+": "🔧 Transition Mask",
"MaskFromColor+": "🔧 Mask From Color",
"MaskFromBatch+": "🔧 MaskFromBatch",
"SimpleMath+": "🔧 Simple Math",
"ConsoleDebug+": "🔧 Console Debug",
"DebugTensorShape+": "🔧 Tensor Shape Debug",
"ModelCompile+": "🔧 Compile Model",
"BatchCount+": "🔧 Batch Count",
"CLIPTextEncodeSDXL+": "🔧 SDXLCLIPTextEncode",
"SDXLResolutionPicker+": "🔧 SDXL Resolutions",
}