From 4c7a9dbcb66d3a53764d4725f92f7c116bcb4821 Mon Sep 17 00:00:00 2001 From: EllangoK Date: Sun, 2 Apr 2023 18:44:27 -0400 Subject: [PATCH 1/9] adds Blend, Blur, Dither, Sharpen nodes --- comfy_extras/nodes_post_processing.py | 215 ++++++++++++++++++++++++++ nodes.py | 3 +- 2 files changed, 217 insertions(+), 1 deletion(-) create mode 100644 comfy_extras/nodes_post_processing.py diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py new file mode 100644 index 00000000000..3f3bddd755b --- /dev/null +++ b/comfy_extras/nodes_post_processing.py @@ -0,0 +1,215 @@ +import torch +import torch.nn.functional as F + + +class Blend: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "blend_factor": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "blend_mode": (["normal", "multiply", "screen", "overlay", "soft_light"],), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "blend_images" + + CATEGORY = "postprocessing" + + def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str): + blended_image = self.blend_mode(image1, image2, blend_mode) + blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor + blended_image = torch.clamp(blended_image, 0, 1) + return (blended_image,) + + def blend_mode(self, img1, img2, mode): + if mode == "normal": + return img2 + elif mode == "multiply": + return img1 * img2 + elif mode == "screen": + return 1 - (1 - img1) * (1 - img2) + elif mode == "overlay": + return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2)) + elif mode == "soft_light": + return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (self.g(img1) - img1)) + else: + raise ValueError(f"Unsupported blend mode: {mode}") + + def g(self, x): + return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x)) + +class Blur: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "blur_radius": ("INT", { + "default": 1, + "min": 1, + "max": 31, + "step": 1 + }), + "sigma": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 10.0, + "step": 0.1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "blur" + + CATEGORY = "postprocessing" + + def gaussian_kernel(self, kernel_size: int, sigma: float): + x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij") + d = torch.sqrt(x * x + y * y) + g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) + return g / g.sum() + + def blur(self, image: torch.Tensor, blur_radius: int, sigma: float): + if blur_radius == 0: + return (image,) + + batch_size, height, width, channels = image.shape + + kernel_size = blur_radius * 2 + 1 + kernel = self.gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1) + + image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) + blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels) + blurred = blurred.permute(0, 2, 3, 1) + + return (blurred,) + +class Dither: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "bits": ("INT", { + "default": 4, + "min": 1, + "max": 8, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "dither" + + CATEGORY = "postprocessing" + + def dither(self, image: torch.Tensor, bits: int): + batch_size, height, width, _ = image.shape + result = torch.zeros_like(image) + + for b in range(batch_size): + tensor_image = image[b] + img = (tensor_image * 255) + height, width, _ = img.shape + + scale = 255 / (2**bits - 1) + + for y in range(height): + for x in range(width): + old_pixel = img[y, x].clone() + new_pixel = torch.round(old_pixel / scale) * scale + img[y, x] = new_pixel + + quant_error = old_pixel - new_pixel + + if x + 1 < width: + img[y, x + 1] += quant_error * 7 / 16 + if y + 1 < height: + if x - 1 >= 0: + img[y + 1, x - 1] += quant_error * 3 / 16 + img[y + 1, x] += quant_error * 5 / 16 + if x + 1 < width: + img[y + 1, x + 1] += quant_error * 1 / 16 + + dithered = img / 255 + tensor = dithered.unsqueeze(0) + result[b] = tensor + + return (result,) + +class Sharpen: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "sharpen_radius": ("INT", { + "default": 1, + "min": 1, + "max": 31, + "step": 1 + }), + "alpha": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 5.0, + "step": 0.1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "sharpen" + + CATEGORY = "postprocessing" + + def sharpen(self, image: torch.Tensor, sharpen_radius: int, alpha: float): + if sharpen_radius == 0: + return (image,) + + batch_size, height, width, channels = image.shape + + kernel_size = sharpen_radius * 2 + 1 + kernel = torch.ones((kernel_size, kernel_size), dtype=torch.float32) * -1 + center = kernel_size // 2 + kernel[center, center] = kernel_size**2 + kernel *= alpha + kernel = kernel.repeat(channels, 1, 1).unsqueeze(1) + + tensor_image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) + sharpened = F.conv2d(tensor_image, kernel, padding=center, groups=channels) + sharpened = sharpened.permute(0, 2, 3, 1) + + result = torch.clamp(sharpened, 0, 1) + + return (result,) + +NODE_CLASS_MAPPINGS = { + "Blend": Blend, + "Blur": Blur, + "Dither": Dither, + "Sharpen": Sharpen, +} diff --git a/nodes.py b/nodes.py index 963ff32a018..a93f0410867 100644 --- a/nodes.py +++ b/nodes.py @@ -1112,4 +1112,5 @@ def load_custom_nodes(): def init_custom_nodes(): load_custom_nodes() - load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py")) \ No newline at end of file + load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py")) + load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_post_processing.py")) From fa2febc0624678362cc758d316bb59afce9c8f06 Mon Sep 17 00:00:00 2001 From: EllangoK Date: Mon, 3 Apr 2023 09:52:04 -0400 Subject: [PATCH 2/9] blend supports any size, dither -> quantize --- comfy_extras/nodes_post_processing.py | 74 ++++++++++++++++----------- 1 file changed, 44 insertions(+), 30 deletions(-) diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py index 3f3bddd755b..322f3ca89f5 100644 --- a/comfy_extras/nodes_post_processing.py +++ b/comfy_extras/nodes_post_processing.py @@ -1,5 +1,7 @@ +import numpy as np import torch import torch.nn.functional as F +from PIL import Image class Blend: @@ -28,6 +30,9 @@ def INPUT_TYPES(s): CATEGORY = "postprocessing" def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str): + if image1.shape != image2.shape: + image2 = self.crop_and_resize(image2, image1.shape) + blended_image = self.blend_mode(image1, image2, blend_mode) blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor blended_image = torch.clamp(blended_image, 0, 1) @@ -50,6 +55,29 @@ def blend_mode(self, img1, img2, mode): def g(self, x): return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x)) + def crop_and_resize(self, img: torch.Tensor, target_shape: tuple): + batch_size, img_h, img_w, img_c = img.shape + _, target_h, target_w, _ = target_shape + img_aspect_ratio = img_w / img_h + target_aspect_ratio = target_w / target_h + + # Crop center of the image to the target aspect ratio + if img_aspect_ratio > target_aspect_ratio: + new_width = int(img_h * target_aspect_ratio) + left = (img_w - new_width) // 2 + img = img[:, :, left:left + new_width, :] + else: + new_height = int(img_w / target_aspect_ratio) + top = (img_h - new_height) // 2 + img = img[:, top:top + new_height, :, :] + + # Resize to target size + img = img.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) + img = F.interpolate(img, size=(target_h, target_w), mode='bilinear', align_corners=False) + img = img.permute(0, 2, 3, 1) + + return img + class Blur: def __init__(self): pass @@ -100,7 +128,7 @@ def blur(self, image: torch.Tensor, blur_radius: int, sigma: float): return (blurred,) -class Dither: +class Quantize: def __init__(self): pass @@ -109,51 +137,37 @@ def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), - "bits": ("INT", { - "default": 4, + "colors": ("INT", { + "default": 256, "min": 1, - "max": 8, + "max": 256, "step": 1 }), + "dither": (["none", "floyd-steinberg"],), }, } RETURN_TYPES = ("IMAGE",) - FUNCTION = "dither" + FUNCTION = "quantize" CATEGORY = "postprocessing" - def dither(self, image: torch.Tensor, bits: int): + def quantize(self, image: torch.Tensor, colors: int = 256, dither: str = "FLOYDSTEINBERG"): batch_size, height, width, _ = image.shape result = torch.zeros_like(image) + dither_option = Image.Dither.FLOYDSTEINBERG if dither == "floyd-steinberg" else Image.Dither.NONE + for b in range(batch_size): tensor_image = image[b] - img = (tensor_image * 255) - height, width, _ = img.shape - - scale = 255 / (2**bits - 1) - - for y in range(height): - for x in range(width): - old_pixel = img[y, x].clone() - new_pixel = torch.round(old_pixel / scale) * scale - img[y, x] = new_pixel - - quant_error = old_pixel - new_pixel + img = (tensor_image * 255).to(torch.uint8).numpy() + pil_image = Image.fromarray(img, mode='RGB') - if x + 1 < width: - img[y, x + 1] += quant_error * 7 / 16 - if y + 1 < height: - if x - 1 >= 0: - img[y + 1, x - 1] += quant_error * 3 / 16 - img[y + 1, x] += quant_error * 5 / 16 - if x + 1 < width: - img[y + 1, x + 1] += quant_error * 1 / 16 + palette = pil_image.quantize(colors=colors) # Required as described in https://github.com/python-pillow/Pillow/issues/5836 + quantized_image = pil_image.quantize(colors=colors, palette=palette, dither=dither_option) - dithered = img / 255 - tensor = dithered.unsqueeze(0) - result[b] = tensor + quantized_array = torch.tensor(np.array(quantized_image.convert("RGB"))).float() / 255 + result[b] = quantized_array return (result,) @@ -210,6 +224,6 @@ def sharpen(self, image: torch.Tensor, sharpen_radius: int, alpha: float): NODE_CLASS_MAPPINGS = { "Blend": Blend, "Blur": Blur, - "Dither": Dither, + "Quantize": Quantize, "Sharpen": Sharpen, } From 56196ab0f72c8f671bd85b425744f80f02c823ea Mon Sep 17 00:00:00 2001 From: EllangoK Date: Tue, 4 Apr 2023 10:57:34 -0400 Subject: [PATCH 3/9] use common_upcale in blend --- comfy_extras/nodes_post_processing.py | 29 +++++---------------------- 1 file changed, 5 insertions(+), 24 deletions(-) diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py index 322f3ca89f5..703deaabf63 100644 --- a/comfy_extras/nodes_post_processing.py +++ b/comfy_extras/nodes_post_processing.py @@ -3,6 +3,8 @@ import torch.nn.functional as F from PIL import Image +import comfy.utils + class Blend: def __init__(self): @@ -31,7 +33,9 @@ def INPUT_TYPES(s): def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str): if image1.shape != image2.shape: - image2 = self.crop_and_resize(image2, image1.shape) + image2 = image2.permute(0, 3, 1, 2) + image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center') + image2 = image2.permute(0, 2, 3, 1) blended_image = self.blend_mode(image1, image2, blend_mode) blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor @@ -55,29 +59,6 @@ def blend_mode(self, img1, img2, mode): def g(self, x): return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x)) - def crop_and_resize(self, img: torch.Tensor, target_shape: tuple): - batch_size, img_h, img_w, img_c = img.shape - _, target_h, target_w, _ = target_shape - img_aspect_ratio = img_w / img_h - target_aspect_ratio = target_w / target_h - - # Crop center of the image to the target aspect ratio - if img_aspect_ratio > target_aspect_ratio: - new_width = int(img_h * target_aspect_ratio) - left = (img_w - new_width) // 2 - img = img[:, :, left:left + new_width, :] - else: - new_height = int(img_w / target_aspect_ratio) - top = (img_h - new_height) // 2 - img = img[:, top:top + new_height, :, :] - - # Resize to target size - img = img.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) - img = F.interpolate(img, size=(target_h, target_w), mode='bilinear', align_corners=False) - img = img.permute(0, 2, 3, 1) - - return img - class Blur: def __init__(self): pass From 080c758cda19288039de6941876dbdf6f3f9d357 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Tue, 4 Apr 2023 18:16:23 +0200 Subject: [PATCH 4/9] Ask for confirmation before clearing nodes --- web/scripts/ui.js | 1 + 1 file changed, 1 insertion(+) diff --git a/web/scripts/ui.js b/web/scripts/ui.js index 679f10b20d2..68bfc792aa7 100644 --- a/web/scripts/ui.js +++ b/web/scripts/ui.js @@ -510,6 +510,7 @@ export class ComfyUI { $el("button", { textContent: "Load", onclick: () => fileInput.click() }), $el("button", { textContent: "Refresh", onclick: () => app.refreshComboInNodes() }), $el("button", { textContent: "Clear", onclick: () => { + if (!confirm("Are you sure you want to remove all nodes?")) return; app.clean(); app.graph.clear(); }}), From af291e6f69a66bce6460de58e6e9328f48640dd5 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 4 Apr 2023 13:56:13 -0400 Subject: [PATCH 5/9] Convert line endings to unix. --- comfy_extras/nodes_post_processing.py | 420 +++++++++++++------------- 1 file changed, 210 insertions(+), 210 deletions(-) diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py index 703deaabf63..de9ef083803 100644 --- a/comfy_extras/nodes_post_processing.py +++ b/comfy_extras/nodes_post_processing.py @@ -1,210 +1,210 @@ -import numpy as np -import torch -import torch.nn.functional as F -from PIL import Image - -import comfy.utils - - -class Blend: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image1": ("IMAGE",), - "image2": ("IMAGE",), - "blend_factor": ("FLOAT", { - "default": 0.5, - "min": 0.0, - "max": 1.0, - "step": 0.01 - }), - "blend_mode": (["normal", "multiply", "screen", "overlay", "soft_light"],), - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "blend_images" - - CATEGORY = "postprocessing" - - def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str): - if image1.shape != image2.shape: - image2 = image2.permute(0, 3, 1, 2) - image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center') - image2 = image2.permute(0, 2, 3, 1) - - blended_image = self.blend_mode(image1, image2, blend_mode) - blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor - blended_image = torch.clamp(blended_image, 0, 1) - return (blended_image,) - - def blend_mode(self, img1, img2, mode): - if mode == "normal": - return img2 - elif mode == "multiply": - return img1 * img2 - elif mode == "screen": - return 1 - (1 - img1) * (1 - img2) - elif mode == "overlay": - return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2)) - elif mode == "soft_light": - return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (self.g(img1) - img1)) - else: - raise ValueError(f"Unsupported blend mode: {mode}") - - def g(self, x): - return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x)) - -class Blur: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "blur_radius": ("INT", { - "default": 1, - "min": 1, - "max": 31, - "step": 1 - }), - "sigma": ("FLOAT", { - "default": 1.0, - "min": 0.1, - "max": 10.0, - "step": 0.1 - }), - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "blur" - - CATEGORY = "postprocessing" - - def gaussian_kernel(self, kernel_size: int, sigma: float): - x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij") - d = torch.sqrt(x * x + y * y) - g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) - return g / g.sum() - - def blur(self, image: torch.Tensor, blur_radius: int, sigma: float): - if blur_radius == 0: - return (image,) - - batch_size, height, width, channels = image.shape - - kernel_size = blur_radius * 2 + 1 - kernel = self.gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1) - - image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) - blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels) - blurred = blurred.permute(0, 2, 3, 1) - - return (blurred,) - -class Quantize: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "colors": ("INT", { - "default": 256, - "min": 1, - "max": 256, - "step": 1 - }), - "dither": (["none", "floyd-steinberg"],), - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "quantize" - - CATEGORY = "postprocessing" - - def quantize(self, image: torch.Tensor, colors: int = 256, dither: str = "FLOYDSTEINBERG"): - batch_size, height, width, _ = image.shape - result = torch.zeros_like(image) - - dither_option = Image.Dither.FLOYDSTEINBERG if dither == "floyd-steinberg" else Image.Dither.NONE - - for b in range(batch_size): - tensor_image = image[b] - img = (tensor_image * 255).to(torch.uint8).numpy() - pil_image = Image.fromarray(img, mode='RGB') - - palette = pil_image.quantize(colors=colors) # Required as described in https://github.com/python-pillow/Pillow/issues/5836 - quantized_image = pil_image.quantize(colors=colors, palette=palette, dither=dither_option) - - quantized_array = torch.tensor(np.array(quantized_image.convert("RGB"))).float() / 255 - result[b] = quantized_array - - return (result,) - -class Sharpen: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "sharpen_radius": ("INT", { - "default": 1, - "min": 1, - "max": 31, - "step": 1 - }), - "alpha": ("FLOAT", { - "default": 1.0, - "min": 0.1, - "max": 5.0, - "step": 0.1 - }), - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "sharpen" - - CATEGORY = "postprocessing" - - def sharpen(self, image: torch.Tensor, sharpen_radius: int, alpha: float): - if sharpen_radius == 0: - return (image,) - - batch_size, height, width, channels = image.shape - - kernel_size = sharpen_radius * 2 + 1 - kernel = torch.ones((kernel_size, kernel_size), dtype=torch.float32) * -1 - center = kernel_size // 2 - kernel[center, center] = kernel_size**2 - kernel *= alpha - kernel = kernel.repeat(channels, 1, 1).unsqueeze(1) - - tensor_image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) - sharpened = F.conv2d(tensor_image, kernel, padding=center, groups=channels) - sharpened = sharpened.permute(0, 2, 3, 1) - - result = torch.clamp(sharpened, 0, 1) - - return (result,) - -NODE_CLASS_MAPPINGS = { - "Blend": Blend, - "Blur": Blur, - "Quantize": Quantize, - "Sharpen": Sharpen, -} +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +import comfy.utils + + +class Blend: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "blend_factor": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "blend_mode": (["normal", "multiply", "screen", "overlay", "soft_light"],), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "blend_images" + + CATEGORY = "postprocessing" + + def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str): + if image1.shape != image2.shape: + image2 = image2.permute(0, 3, 1, 2) + image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center') + image2 = image2.permute(0, 2, 3, 1) + + blended_image = self.blend_mode(image1, image2, blend_mode) + blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor + blended_image = torch.clamp(blended_image, 0, 1) + return (blended_image,) + + def blend_mode(self, img1, img2, mode): + if mode == "normal": + return img2 + elif mode == "multiply": + return img1 * img2 + elif mode == "screen": + return 1 - (1 - img1) * (1 - img2) + elif mode == "overlay": + return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2)) + elif mode == "soft_light": + return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (self.g(img1) - img1)) + else: + raise ValueError(f"Unsupported blend mode: {mode}") + + def g(self, x): + return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x)) + +class Blur: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "blur_radius": ("INT", { + "default": 1, + "min": 1, + "max": 31, + "step": 1 + }), + "sigma": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 10.0, + "step": 0.1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "blur" + + CATEGORY = "postprocessing" + + def gaussian_kernel(self, kernel_size: int, sigma: float): + x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij") + d = torch.sqrt(x * x + y * y) + g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) + return g / g.sum() + + def blur(self, image: torch.Tensor, blur_radius: int, sigma: float): + if blur_radius == 0: + return (image,) + + batch_size, height, width, channels = image.shape + + kernel_size = blur_radius * 2 + 1 + kernel = self.gaussian_kernel(kernel_size, sigma).repeat(channels, 1, 1).unsqueeze(1) + + image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) + blurred = F.conv2d(image, kernel, padding=kernel_size // 2, groups=channels) + blurred = blurred.permute(0, 2, 3, 1) + + return (blurred,) + +class Quantize: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "colors": ("INT", { + "default": 256, + "min": 1, + "max": 256, + "step": 1 + }), + "dither": (["none", "floyd-steinberg"],), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "quantize" + + CATEGORY = "postprocessing" + + def quantize(self, image: torch.Tensor, colors: int = 256, dither: str = "FLOYDSTEINBERG"): + batch_size, height, width, _ = image.shape + result = torch.zeros_like(image) + + dither_option = Image.Dither.FLOYDSTEINBERG if dither == "floyd-steinberg" else Image.Dither.NONE + + for b in range(batch_size): + tensor_image = image[b] + img = (tensor_image * 255).to(torch.uint8).numpy() + pil_image = Image.fromarray(img, mode='RGB') + + palette = pil_image.quantize(colors=colors) # Required as described in https://github.com/python-pillow/Pillow/issues/5836 + quantized_image = pil_image.quantize(colors=colors, palette=palette, dither=dither_option) + + quantized_array = torch.tensor(np.array(quantized_image.convert("RGB"))).float() / 255 + result[b] = quantized_array + + return (result,) + +class Sharpen: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "sharpen_radius": ("INT", { + "default": 1, + "min": 1, + "max": 31, + "step": 1 + }), + "alpha": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 5.0, + "step": 0.1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "sharpen" + + CATEGORY = "postprocessing" + + def sharpen(self, image: torch.Tensor, sharpen_radius: int, alpha: float): + if sharpen_radius == 0: + return (image,) + + batch_size, height, width, channels = image.shape + + kernel_size = sharpen_radius * 2 + 1 + kernel = torch.ones((kernel_size, kernel_size), dtype=torch.float32) * -1 + center = kernel_size // 2 + kernel[center, center] = kernel_size**2 + kernel *= alpha + kernel = kernel.repeat(channels, 1, 1).unsqueeze(1) + + tensor_image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) + sharpened = F.conv2d(tensor_image, kernel, padding=center, groups=channels) + sharpened = sharpened.permute(0, 2, 3, 1) + + result = torch.clamp(sharpened, 0, 1) + + return (result,) + +NODE_CLASS_MAPPINGS = { + "Blend": Blend, + "Blur": Blur, + "Quantize": Quantize, + "Sharpen": Sharpen, +} From de3d5f46ce0544339884fe454a59b342fcf28cf3 Mon Sep 17 00:00:00 2001 From: missionfloyd Date: Tue, 4 Apr 2023 19:32:49 -0600 Subject: [PATCH 6/9] Fix .graphdialog style --- web/style.css | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/web/style.css b/web/style.css index 9162bbba99f..393d1667e54 100644 --- a/web/style.css +++ b/web/style.css @@ -237,3 +237,28 @@ button.comfy-queue-btn { visibility:hidden } } + +.graphdialog { + min-height: 1em; +} + +.graphdialog .name { + font-size: 14px; + font-family: sans-serif; + color: #999999; +} + +.graphdialog button { + margin-top: unset; + vertical-align: unset; + height: 1.6em; + padding-right: 8px; +} + +.graphdialog input, .graphdialog textarea, .graphdialog select { + background-color: #222; + border: 2px solid; + border-color: #444444; + color: #ddd; + border-radius: 12px 0 0 12px; +} From e46b1c3034a23eeb048e279d0d285737d39a4b1a Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 4 Apr 2023 22:22:02 -0400 Subject: [PATCH 7/9] Disable xformers in VAE when xformers == 0.0.18 --- comfy/ldm/modules/diffusionmodules/model.py | 4 ++-- comfy/model_management.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/comfy/ldm/modules/diffusionmodules/model.py b/comfy/ldm/modules/diffusionmodules/model.py index 94f5510b9f7..788a6fc4f0a 100644 --- a/comfy/ldm/modules/diffusionmodules/model.py +++ b/comfy/ldm/modules/diffusionmodules/model.py @@ -9,7 +9,7 @@ from ldm.modules.attention import MemoryEfficientCrossAttention import model_management -if model_management.xformers_enabled(): +if model_management.xformers_enabled_vae(): import xformers import xformers.ops @@ -364,7 +364,7 @@ def forward(self, x, context=None, mask=None): def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None): assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown' - if model_management.xformers_enabled() and attn_type == "vanilla": + if model_management.xformers_enabled_vae() and attn_type == "vanilla": attn_type = "vanilla-xformers" if model_management.pytorch_attention_enabled() and attn_type == "vanilla": attn_type = "vanilla-pytorch" diff --git a/comfy/model_management.py b/comfy/model_management.py index 4aa47ff1637..052dfb77521 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -199,11 +199,25 @@ def get_autocast_device(dev): return dev.type return "cuda" + def xformers_enabled(): if vram_state == CPU: return False return XFORMERS_IS_AVAILBLE + +def xformers_enabled_vae(): + enabled = xformers_enabled() + if not enabled: + return False + try: + #0.0.18 has a bug where Nan is returned when inputs are too big (1152x1920 res images and above) + if xformers.version.__version__ == "0.0.18": + return False + except: + pass + return enabled + def pytorch_attention_enabled(): return ENABLE_PYTORCH_ATTENTION From 10ad4c1d17d8ea469565d904a8f47f1d2eeee459 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 4 Apr 2023 22:48:11 -0400 Subject: [PATCH 8/9] Move unclip stuff out of _for_testing --- nodes.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/nodes.py b/nodes.py index 28beb37b14e..935e28b84a2 100644 --- a/nodes.py +++ b/nodes.py @@ -197,7 +197,7 @@ def INPUT_TYPES(s): RETURN_TYPES = ("MODEL", "CLIP", "VAE") FUNCTION = "load_checkpoint" - CATEGORY = "loaders" + CATEGORY = "advanced/loaders" def load_checkpoint(self, config_name, ckpt_name, output_vae=True, output_clip=True): config_path = folder_paths.get_full_path("configs", config_name) @@ -227,7 +227,7 @@ def INPUT_TYPES(s): RETURN_TYPES = ("MODEL", "CLIP", "VAE", "CLIP_VISION") FUNCTION = "load_checkpoint" - CATEGORY = "_for_testing/unclip" + CATEGORY = "loaders" def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True): ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) @@ -450,7 +450,7 @@ def INPUT_TYPES(s): RETURN_TYPES = ("CONDITIONING",) FUNCTION = "apply_adm" - CATEGORY = "_for_testing/unclip" + CATEGORY = "conditioning" def apply_adm(self, conditioning, clip_vision_output, strength, noise_augmentation): c = [] @@ -1038,7 +1038,6 @@ def expand_image(self, image, left, top, right, bottom, feathering): NODE_CLASS_MAPPINGS = { "KSampler": KSampler, - "CheckpointLoader": CheckpointLoader, "CheckpointLoaderSimple": CheckpointLoaderSimple, "CLIPTextEncode": CLIPTextEncode, "CLIPSetLastLayer": CLIPSetLastLayer, @@ -1077,6 +1076,7 @@ def expand_image(self, image, left, top, right, bottom, feathering): "VAEEncodeTiled": VAEEncodeTiled, "TomePatchModel": TomePatchModel, "unCLIPCheckpointLoader": unCLIPCheckpointLoader, + "CheckpointLoader": CheckpointLoader, } def load_custom_node(module_path): From 871a76b77b9cafa8615da1cedeaafc1b10cf85e3 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 4 Apr 2023 22:54:33 -0400 Subject: [PATCH 9/9] Rename and reorganize post processing nodes. --- comfy_extras/nodes_post_processing.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py index de9ef083803..ba699e2b8a0 100644 --- a/comfy_extras/nodes_post_processing.py +++ b/comfy_extras/nodes_post_processing.py @@ -29,7 +29,7 @@ def INPUT_TYPES(s): RETURN_TYPES = ("IMAGE",) FUNCTION = "blend_images" - CATEGORY = "postprocessing" + CATEGORY = "image/postprocessing" def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str): if image1.shape != image2.shape: @@ -86,7 +86,7 @@ def INPUT_TYPES(s): RETURN_TYPES = ("IMAGE",) FUNCTION = "blur" - CATEGORY = "postprocessing" + CATEGORY = "image/postprocessing" def gaussian_kernel(self, kernel_size: int, sigma: float): x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size), torch.linspace(-1, 1, kernel_size), indexing="ij") @@ -131,7 +131,7 @@ def INPUT_TYPES(s): RETURN_TYPES = ("IMAGE",) FUNCTION = "quantize" - CATEGORY = "postprocessing" + CATEGORY = "image/postprocessing" def quantize(self, image: torch.Tensor, colors: int = 256, dither: str = "FLOYDSTEINBERG"): batch_size, height, width, _ = image.shape @@ -179,7 +179,7 @@ def INPUT_TYPES(s): RETURN_TYPES = ("IMAGE",) FUNCTION = "sharpen" - CATEGORY = "postprocessing" + CATEGORY = "image/postprocessing" def sharpen(self, image: torch.Tensor, sharpen_radius: int, alpha: float): if sharpen_radius == 0: @@ -203,8 +203,8 @@ def sharpen(self, image: torch.Tensor, sharpen_radius: int, alpha: float): return (result,) NODE_CLASS_MAPPINGS = { - "Blend": Blend, - "Blur": Blur, - "Quantize": Quantize, - "Sharpen": Sharpen, + "ImageBlend": Blend, + "ImageBlur": Blur, + "ImageQuantize": Quantize, + "ImageSharpen": Sharpen, }