import warnings
warnings.filterwarnings('ignore', module="torchvision")
import ast
import math
import random
import os
import operator as op
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageColor, ImageFilter
import io
import torch
import torch.nn.functional as F
import torchvision.transforms.v2 as T
from nodes import MAX_RESOLUTION, SaveImage, common_ksampler
import folder_paths
import comfy.utils
import comfy.samplers
STOCHASTIC_SAMPLERS = ["euler_ancestral", "dpm_2_ancestral", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm"]
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 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.unsqueeze(1)
blurred = T.GaussianBlur(size, amount)(blurred)
blurred = blurred.squeeze(1)
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 KSamplerVariationsStochastic:
@classmethod
def INPUT_TYPES(s):
return {"required":{
"model": ("MODEL",),
"latent_image": ("LATENT", ),
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 25, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
"sampler": (comfy.samplers.KSampler.SAMPLERS, ),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
"positive": ("CONDITIONING", ),
"negative": ("CONDITIONING", ),
"variation_seed": ("INT:seed", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"variation_strength": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step":0.05, "round": 0.01}),
#"variation_sampler": (comfy.samplers.KSampler.SAMPLERS, ),
"cfg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.05, "round": 0.01}),
}}
RETURN_TYPES = ("LATENT", )
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, model, latent_image, noise_seed, steps, cfg, sampler, scheduler, positive, negative, variation_seed, variation_strength, cfg_scale, variation_sampler="dpmpp_2m_sde"):
# Stage 1: composition sampler
force_full_denoise = False # return with leftover noise = "enable"
disable_noise = False # add noise = "enable"
end_at_step = max(int(steps * (1-variation_strength)), 1)
start_at_step = 0
work_latent = latent_image.copy()
batch_size = work_latent["samples"].shape[0]
work_latent["samples"] = work_latent["samples"][0].unsqueeze(0)
stage1 = common_ksampler(model, noise_seed, steps, cfg, sampler, scheduler, positive, negative, work_latent, denoise=1.0, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)[0]
print(stage1)
if batch_size > 1:
stage1["samples"] = stage1["samples"].clone().repeat(batch_size, 1, 1, 1)
# Stage 2: variation sampler
force_full_denoise = True
disable_noise = True
cfg = max(cfg * cfg_scale, 1.0)
start_at_step = end_at_step
end_at_step = steps
return common_ksampler(model, variation_seed, steps, cfg, variation_sampler, scheduler, positive, negative, stage1, denoise=1.0, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
# From https://github.com/BlenderNeko/ComfyUI_Noise/
def slerp(val, low, high):
dims = low.shape
low = low.reshape(dims[0], -1)
high = high.reshape(dims[0], -1)
low_norm = low/torch.norm(low, dim=1, keepdim=True)
high_norm = high/torch.norm(high, dim=1, keepdim=True)
low_norm[low_norm != low_norm] = 0.0
high_norm[high_norm != high_norm] = 0.0
omega = torch.acos((low_norm*high_norm).sum(1))
so = torch.sin(omega)
res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high
return res.reshape(dims)
class KSamplerVariationsWithNoise:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"model": ("MODEL", ),
"latent_image": ("LATENT", ),
"main_seed": ("INT:seed", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
"positive": ("CONDITIONING", ),
"negative": ("CONDITIONING", ),
"variation_strength": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step":0.01, "round": 0.01}),
#"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
#"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
#"return_with_leftover_noise": (["disable", "enable"], ),
"variation_seed": ("INT:seed", {"default": random.randint(0, 0xffffffffffffffff), "min": 0, "max": 0xffffffffffffffff}),
}}
RETURN_TYPES = ("LATENT",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, model, latent_image, main_seed, steps, cfg, sampler_name, scheduler, positive, negative, variation_strength, variation_seed):
generator = torch.manual_seed(main_seed)
batch_size, _, height, width = latent_image["samples"].shape
base_noise = torch.randn((1, 4, height, width), dtype=torch.float32, device="cpu", generator=generator).repeat(batch_size, 1, 1, 1).cpu()
generator = torch.manual_seed(variation_seed)
variation_noise = torch.randn((batch_size, 4, height, width), dtype=torch.float32, device="cpu", generator=generator).cpu()
slerp_noise = slerp(variation_strength, base_noise, variation_noise)
device = comfy.model_management.get_torch_device()
end_at_step = steps #min(steps, end_at_step)
start_at_step = 0 #min(start_at_step, end_at_step)
real_model = None
comfy.model_management.load_model_gpu(model)
real_model = model.model
sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=1.0, model_options=model.model_options)
sigmas = sampler.sigmas
sigma = sigmas[start_at_step] - sigmas[end_at_step]
sigma /= model.model.latent_format.scale_factor
sigma = sigma.cpu().numpy()
work_latent = latent_image.copy()
work_latent["samples"] = latent_image["samples"].clone() + slerp_noise * sigma
force_full_denoise = True
#if return_with_leftover_noise == "enable":
# force_full_denoise = False
disable_noise = True
return common_ksampler(model, main_seed, steps, cfg, sampler_name, scheduler, positive, negative, work_latent, denoise=1.0, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
class SDXLEmptyLatentSizePicker:
def __init__(self):
self.device = comfy.model_management.intermediate_device()
@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)"}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
}}
RETURN_TYPES = ("LATENT","INT","INT",)
RETURN_NAMES = ("LATENT","width", "height",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, resolution, batch_size):
width, height = resolution.split(" ")[0].split("x")
width = int(width)
height = int(height)
latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
return ({"samples":latent}, 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 }),
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1 }),
}}
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, strength):
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 resource efficient? should we use a batch instead?
lut_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]
lut_img = lut_img * dom_scale + lut.domain[0]
if log_colorspace:
lut_img = lut_img ** (1/2.2)
lut_img = lut.apply(lut_img)
if log_colorspace:
lut_img = lut_img ** (2.2)
if is_non_default_domain:
lut_img = (lut_img - lut.domain[0]) / dom_scale
lut_img = torch.from_numpy(lut_img)
if strength < 1.0:
lut_img = strength * lut_img + (1 - strength) * img
out.append(lut_img)
out = torch.stack(out)
return (out, )
FONTS_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "fonts")
class DrawText:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"text": ("STRING", { "multiline": True, "default": "Hello, World!" }),
"font": ([f for f in os.listdir(FONTS_DIR) if f.endswith('.ttf') or f.endswith('.otf')], ),
"size": ("INT", { "default": 56, "min": 1, "max": 9999, "step": 1 }),
"color": ("STRING", { "multiline": False, "default": "#FFFFFF" }),
"background_color": ("STRING", { "multiline": False, "default": "#00000000" }),
"shadow_distance": ("INT", { "default": 0, "min": 0, "max": 100, "step": 1 }),
"shadow_blur": ("INT", { "default": 0, "min": 0, "max": 100, "step": 1 }),
"shadow_color": ("STRING", { "multiline": False, "default": "#000000" }),
"alignment": (["left", "center", "right"],),
"width": ("INT", { "default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1 }),
"height": ("INT", { "default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1 }),
},
}
RETURN_TYPES = ("IMAGE", "MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, text, font, size, color, background_color, shadow_distance, shadow_blur, shadow_color, alignment, width, height):
font = ImageFont.truetype(os.path.join(FONTS_DIR, font), size)
lines = text.split("\n")
# Calculate the width and height of the text
text_width = max(font.getbbox(line)[2] for line in lines)
line_height = font.getmask(text).getbbox()[3] + font.getmetrics()[1] # add descent to height
text_height = line_height * len(lines)
width = width if width > 0 else text_width
height = height if height > 0 else text_height
background_color = ImageColor.getrgb(background_color)
image = Image.new('RGBA', (width + shadow_distance, height + shadow_distance), color=background_color)
image_shadow = None
if shadow_distance > 0:
image_shadow = Image.new('RGBA', (width + shadow_distance, height + shadow_distance), color=background_color)
for i, line in enumerate(lines):
line_width = font.getbbox(line)[2]
#text_height =font.getbbox(line)[3]
if alignment == "left":
x = 0
elif alignment == "center":
x = (width - line_width) / 2
elif alignment == "right":
x = width - line_width
y = i * line_height
draw = ImageDraw.Draw(image)
draw.text((x, y), line, font=font, fill=color)
if image_shadow is not None:
draw = ImageDraw.Draw(image_shadow)
draw.text((x + shadow_distance, y + shadow_distance), line, font=font, fill=shadow_color)
if image_shadow is not None:
image_shadow = image_shadow.filter(ImageFilter.GaussianBlur(shadow_blur))
image = Image.alpha_composite(image_shadow, image)
image = pb(T.ToTensor()(image).unsqueeze(0))
mask = image[:, :, :, 3] if image.shape[3] == 4 else torch.ones_like(image[:, :, :, 0])
return (image[:, :, :, :3], mask,)
class RemBGSession:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": (["u2net: general purpose", "u2netp: lightweight general purpose", "u2net_human_seg: human segmentation", "u2net_cloth_seg: cloths Parsing", "silueta: very small u2net", "isnet-general-use: general purpose", "isnet-anime: anime illustrations", "sam: general purpose"],),
"providers": (['CPU', 'CUDA', 'ROCM', 'DirectML', 'OpenVINO', 'CoreML', 'Tensorrt', 'Azure'],),
},
}
RETURN_TYPES = ("REMBG_SESSION",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, model, providers):
from rembg import new_session as rembg_new_session
model = model.split(":")[0]
return (rembg_new_session(model, providers=[providers+"ExecutionProvider"]),)
class ImageRemoveBackground:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"rembg_session": ("REMBG_SESSION",),
"image": ("IMAGE",),
},
}
RETURN_TYPES = ("IMAGE", "MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, rembg_session, image):
from rembg import remove as rembg
image = p(image)
output = []
for img in image:
img = T.ToPILImage()(img)
img = rembg(img, session=rembg_session)
output.append(T.ToTensor()(img))
output = torch.stack(output, dim=0)
output = pb(output)
mask = output[:, :, :, 3] if output.shape[3] == 4 else torch.ones_like(output[:, :, :, 0])
return(output[:, :, :, :3], mask,)
class NoiseFromImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"noise_size": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01 }),
"color_noise": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01 }),
"mask_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01 }),
"mask_scale_diff": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01 }),
"noise_strenght": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01 }),
"saturation": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 100.0, "step": 0.1 }),
"contrast": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1 }),
"blur": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1 }),
},
"optional": {
"noise_mask": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE","IMAGE",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, image, noise_size, color_noise, mask_strength, mask_scale_diff, noise_strenght, saturation, contrast, blur, noise_mask=None):
torch.manual_seed(0)
elastic_alpha = max(image.shape[1], image.shape[2])# * noise_size
elastic_sigma = elastic_alpha / 400 * noise_size
blur_size = int(6 * blur+1)
if blur_size % 2 == 0:
blur_size+= 1
if noise_mask is None:
noise_mask = image
# Ensure noise mask is the same size as the image
if noise_mask.shape[1:] != image.shape[1:]:
noise_mask = F.interpolate(p(noise_mask), size=(image.shape[1], image.shape[2]), mode='bicubic', align_corners=False)
noise_mask = pb(noise_mask)
# Ensure we have the same number of masks and images
if noise_mask.shape[0] > image.shape[0]:
noise_mask = noise_mask[:image.shape[0]]
else:
noise_mask = torch.cat((noise_mask, noise_mask[-1:].repeat((image.shape[0]-noise_mask.shape[0], 1, 1, 1))), dim=0)
# Convert image to grayscale mask
noise_mask = noise_mask.mean(dim=3).unsqueeze(-1)
# add color noise
imgs = p(image.clone())
if color_noise > 0:
color_noise = torch.normal(torch.zeros_like(imgs), std=color_noise)
#color_noise = torch.rand_like(imgs) * (color_noise * 2) - color_noise
color_noise *= (imgs - imgs.min()) / (imgs.max() - imgs.min())
imgs = imgs + color_noise
imgs = imgs.clamp(0, 1)
# create fine noise
fine_noise = []
for n in imgs:
avg_color = n.mean(dim=[1,2])
tmp_noise = T.ElasticTransform(alpha=elastic_alpha, sigma=elastic_sigma, fill=avg_color.tolist())(n)
#tmp_noise = T.functional.adjust_saturation(tmp_noise, 2.0)
tmp_noise = T.GaussianBlur(blur_size, blur)(tmp_noise)
tmp_noise = T.ColorJitter(contrast=(contrast,contrast), saturation=(saturation,saturation))(tmp_noise)
fine_noise.append(tmp_noise)
#tmp_noise = F.interpolate(tmp_noise, scale_factor=.1, mode='bilinear', align_corners=False)
#tmp_noise = F.interpolate(tmp_noise, size=(tmp_noise.shape[1], tmp_noise.shape[2]), mode='bilinear', align_corners=False)
#tmp_noise = T.ElasticTransform(alpha=elastic_alpha, sigma=elastic_sigma/3, fill=avg_color.tolist())(n)
#tmp_noise = T.GaussianBlur(blur_size, blur)(tmp_noise)
#tmp_noise = T.functional.adjust_saturation(tmp_noise, saturation)
#tmp_noise = T.ColorJitter(contrast=(contrast,contrast), saturation=(saturation,saturation))(tmp_noise)
#fine_noise.append(tmp_noise)
imgs = None
del imgs
fine_noise = torch.stack(fine_noise, dim=0)
fine_noise = pb(fine_noise)
#fine_noise = torch.stack(fine_noise, dim=0)
#fine_noise = pb(fine_noise)
mask_scale_diff = min(mask_scale_diff, 0.99)
if mask_scale_diff > 0:
coarse_noise = F.interpolate(p(fine_noise), scale_factor=1-mask_scale_diff, mode='area')
coarse_noise = F.interpolate(coarse_noise, size=(fine_noise.shape[1], fine_noise.shape[2]), mode='bilinear', align_corners=False)
coarse_noise = pb(coarse_noise)
else:
coarse_noise = fine_noise
#noise_mask = noise_mask * mask_strength + (1 - mask_strength)
# merge fine and coarse noise
output = (1 - noise_mask) * coarse_noise + noise_mask * fine_noise
#noise_mask = noise_mask * mask_strength
if mask_strength < 1:
noise_mask = noise_mask.pow(mask_strength)
noise_mask = torch.nan_to_num(noise_mask).clamp(0, 1)
output = noise_mask * output + (1 - noise_mask) * image
# apply noise to image
output = output * noise_strenght + image * (1 - noise_strenght)
output = output.clamp(0, 1)
return (output,noise_mask.repeat(1,1,1,3),)
class RemoveLatentMask:
@classmethod
def INPUT_TYPES(s):
return {"required": { "samples": ("LATENT",),}}
RETURN_TYPES = ("LATENT",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, samples):
s = samples.copy()
if "noise_mask" in s:
del s["noise_mask"]
return (s,)
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,
"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,
"KSamplerVariationsStochastic+": KSamplerVariationsStochastic,
"KSamplerVariationsWithNoise+": KSamplerVariationsWithNoise,
"CLIPTextEncodeSDXL+": CLIPTextEncodeSDXLSimplified,
"SDXLEmptyLatentSizePicker+": SDXLEmptyLatentSizePicker,
"DrawText+": DrawText,
"RemBGSession+": RemBGSession,
"ImageRemoveBackground+": ImageRemoveBackground,
"RemoveLatentMask+": RemoveLatentMask,
#"NoiseFromImage~": NoiseFromImage,
}
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",
"MaskBlur+": "🔧 Mask Blur",
"MaskFlip+": "🔧 Mask Flip",
"MaskPreview+": "🔧 Mask Preview",
"MaskBatch+": "🔧 Mask Batch",
"MaskExpandBatch+": "🔧 Mask Expand Batch",
"TransitionMask+": "🔧 Transition Mask",
"MaskFromColor+": "🔧 Mask From Color",
"MaskFromBatch+": "🔧 Mask From Batch",
"SimpleMath+": "🔧 Simple Math",
"ConsoleDebug+": "🔧 Console Debug",
"DebugTensorShape+": "🔧 Tensor Shape Debug",
"ModelCompile+": "🔧 Compile Model",
"BatchCount+": "🔧 Batch Count",
"KSamplerVariationsStochastic+": "🔧 KSampler Stochastic Variations",
"KSamplerVariationsWithNoise+": "🔧 KSampler Variations with Noise Injection",
"CLIPTextEncodeSDXL+": "🔧 SDXLCLIPTextEncode",
"SDXLEmptyLatentSizePicker+": "🔧 SDXL Empty Latent Size Picker",
"DrawText+": "🔧 Draw Text",
"RemBGSession+": "🔧 RemBG Session",
"ImageRemoveBackground+": "🔧 Image Remove Background",
"RemoveLatentMask+": "🔧 Remove Latent Mask",
#"NoiseFromImage~": "🔧 Noise From Image",
}