diff --git a/README.md b/README.md new file mode 100644 index 0000000..bf4ea46 --- /dev/null +++ b/README.md @@ -0,0 +1,7 @@ +## Install + +From within the python environment you already use for ComfyUI install the requirements. +```bash +pip install -r comfy_mtb/requirements.txt +``` + diff --git a/__init__.py b/__init__.py index 6a0c303..ea748da 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,6 @@ +from .nodes.deep_bump import DeepBump +from .nodes.latent_processing import LatentLerp +from .nodes.fun import QRNode from .nodes.image_processing import ( ImageCompare, Denoise, @@ -13,6 +16,9 @@ from .nodes.graph_utils import IntToNumber, Modulo # NODE MAPPING NODE_CLASS_MAPPINGS = { + "Latent Lerp (mtb) [DEPRECATED]": LatentLerp, + "Deep Bump (mtb)": DeepBump, + "Int to Number (mtb)": IntToNumber, "Bounding Box (mtb)": BoundingBox, "Crop (mtb)": Crop, @@ -24,4 +30,5 @@ NODE_CLASS_MAPPINGS = { "HSV to RGB (mtb)": HSVtoRGB, "Color Correct (mtb)": ColorCorrect, "Modulo (mtb)": Modulo, + "QR Code (mtb)": QRNode, } diff --git a/nodes/deep_bump.py b/nodes/deep_bump.py new file mode 100644 index 0000000..652b2f2 --- /dev/null +++ b/nodes/deep_bump.py @@ -0,0 +1,306 @@ +import onnxruntime as ort +import numpy as np +import pathlib +import onnxruntime as ort +import torch +from .. import utils as utils_inference + +# Disable MS telemetry +ort.disable_telemetry_events() + +# - COLOR to NORMALS +def color_to_normals(color_img, overlap, progress_callback): + """Computes a normal map from the given color map. 'color_img' must be a numpy array + in C,H,W format (with C as RGB). 'overlap' must be one of 'SMALL', 'MEDIUM', 'LARGE'.""" + + # Remove alpha & convert to grayscale + img = np.mean(color_img[:3], axis=0, keepdimss=True) + + # Split image in tiles + print("DeepBump Color → Normals : tilling") + tile_size = 256 + overlaps = { + "SMALL": tile_size // 6, + "MEDIUM": tile_size // 4, + "LARGE": tile_size // 2, + } + stride_size = tile_size - overlaps[overlap] + tiles, paddings = utils_inference.tiles_split( + img, (tile_size, tile_size), (stride_size, stride_size) + ) + + # Load model + print("DeepBump Color → Normals : loading model") + addon_path = str(pathlib.Path(__file__).parent.absolute()) + ort_session = ort.InferenceSession(f"{addon_path}/models/deepbump256.onnx") + + # Predict normal map for each tile + print("DeepBump Color → Normals : generating") + pred_tiles = utils_inference.tiles_infer( + tiles, ort_session, progress_callback=progress_callback + ) + + # Merge tiles + print("DeepBump Color → Normals : merging") + pred_img = utils_inference.tiles_merge( + pred_tiles, + (stride_size, stride_size), + (3, img.shape[1], img.shape[2]), + paddings, + ) + + # Normalize each pixel to unit vector + pred_img = utils_inference.normalize(pred_img) + + return pred_img + + +# - NORMALS to CURVATURE +def conv_1d(array, kernel_1d): + """Performs row by row 1D convolutions of the given 2D image with the given 1D kernel.""" + + # Input kernel length must be odd + k_l = len(kernel_1d) + assert k_l % 2 != 0 + # Convolution is repeat-padded + extended = np.pad(array, k_l // 2, mode="wrap") + # Output has same size as input (padded, valid-mode convolution) + output = np.empty(array.shape) + for i in range(array.shape[0]): + output[i] = np.convolve(extended[i + (k_l // 2)], kernel_1d, mode="valid") + + return output * -1 + + +def gaussian_kernel(length, sigma): + """Returns a 1D gaussian kernel of size 'length'.""" + + space = np.linspace(-(length - 1) / 2, (length - 1) / 2, length) + kernel = np.exp(-0.5 * np.square(space) / np.square(sigma)) + return kernel / np.sum(kernel) + + +def normalize(np_array): + """Normalize all elements of the given numpy array to [0,1]""" + + return (np_array - np.min(np_array)) / (np.max(np_array) - np.min(np_array)) + + +def normals_to_curvature(normals_img, blur_radius, progress_callback): + """Computes a curvature map from the given normal map. 'normals_img' must be a numpy array + in C,H,W format (with C as RGB). 'blur_radius' must be one of 'SMALLEST', 'SMALLER', 'SMALL', + 'MEDIUM', 'LARGE', 'LARGER', 'LARGEST'.""" + + # Convolutions on normal map red & green channels + if progress_callback is not None: + progress_callback(0, 4) + diff_kernel = np.array([-1, 0, 1]) + h_conv = conv_1d(normals_img[0, :, :], diff_kernel) + if progress_callback is not None: + progress_callback(1, 4) + v_conv = conv_1d(-1 * normals_img[1, :, :].T, diff_kernel).T + if progress_callback is not None: + progress_callback(2, 4) + + # Sum detected edges + edges_conv = h_conv + v_conv + + # Blur radius size is proportional to img sizes + blur_factors = { + "SMALLEST": 1 / 256, + "SMALLER": 1 / 128, + "SMALL": 1 / 64, + "MEDIUM": 1 / 32, + "LARGE": 1 / 16, + "LARGER": 1 / 8, + "LARGEST": 1 / 4, + } + assert blur_radius in blur_factors + blur_radius_px = int(np.mean(normals_img.shape[1:3]) * blur_factors[blur_radius]) + + # If blur radius too small, do not blur + if blur_radius_px < 2: + edges_conv = normalize(edges_conv) + return np.stack([edges_conv, edges_conv, edges_conv]) + + # Make sure blur kernel length is odd + if blur_radius_px % 2 == 0: + blur_radius_px += 1 + + # Blur curvature with separated convolutions + sigma = blur_radius_px // 8 + if sigma == 0: + sigma = 1 + g_kernel = gaussian_kernel(blur_radius_px, sigma) + h_blur = conv_1d(edges_conv, g_kernel) + if progress_callback is not None: + progress_callback(3, 4) + v_blur = conv_1d(h_blur.T, g_kernel).T + if progress_callback is not None: + progress_callback(4, 4) + + # Normalize to [0,1] + curvature = normalize(v_blur) + + # Expand single channel the three channels (RGB) + return np.stack([curvature, curvature, curvature]) + + +# - NORMALS to HEIGHT +def normals_to_grad(normals_img): + return (normals_img[0] - 0.5) * 2, (normals_img[1] - 0.5) * 2 + + +def copy_flip(grad_x, grad_y): + """Concat 4 flipped copies of input gradients (makes them wrap). + Output is twice bigger in both dimensions.""" + + grad_x_top = np.hstack([grad_x, -np.flip(grad_x, axis=1)]) + grad_x_bottom = np.hstack([np.flip(grad_x, axis=0), -np.flip(grad_x)]) + new_grad_x = np.vstack([grad_x_top, grad_x_bottom]) + + grad_y_top = np.hstack([grad_y, np.flip(grad_y, axis=1)]) + grad_y_bottom = np.hstack([-np.flip(grad_y, axis=0), -np.flip(grad_y)]) + new_grad_y = np.vstack([grad_y_top, grad_y_bottom]) + + return new_grad_x, new_grad_y + + +def frankot_chellappa(grad_x, grad_y, progress_callback=None): + """Frankot-Chellappa depth-from-gradient algorithm.""" + + if progress_callback is not None: + progress_callback(0, 3) + + rows, cols = grad_x.shape + + rows_scale = (np.arange(rows) - (rows // 2 + 1)) / (rows - rows % 2) + cols_scale = (np.arange(cols) - (cols // 2 + 1)) / (cols - cols % 2) + + u_grid, v_grid = np.meshgrid(cols_scale, rows_scale) + + u_grid = np.fft.ifftshift(u_grid) + v_grid = np.fft.ifftshift(v_grid) + + if progress_callback is not None: + progress_callback(1, 3) + + grad_x_F = np.fft.fft2(grad_x) + grad_y_F = np.fft.fft2(grad_y) + + if progress_callback is not None: + progress_callback(2, 3) + + nominator = (-1j * u_grid * grad_x_F) + (-1j * v_grid * grad_y_F) + denominator = (u_grid**2) + (v_grid**2) + 1e-16 + + Z_F = nominator / denominator + Z_F[0, 0] = 0.0 + + Z = np.real(np.fft.ifft2(Z_F)) + + if progress_callback is not None: + progress_callback(3, 3) + + return (Z - np.min(Z)) / (np.max(Z) - np.min(Z)) + + +def normals_to_height(normals_img, seamless, progress_callback): + """Computes a height map from the given normal map. 'normals_img' must be a numpy array + in C,H,W format (with C as RGB). 'seamless' is a bool that should indicates if 'normals_img' + is seamless.""" + + # Flip height axis + flip_img = np.flip(normals_img, axis=1) + + # Get gradients from normal map + grad_x, grad_y = normals_to_grad(flip_img) + grad_x = np.flip(grad_x, axis=0) + grad_y = np.flip(grad_y, axis=0) + + # If non-seamless chosen, expand gradients + if not seamless: + grad_x, grad_y = copy_flip(grad_x, grad_y) + + # Compute height + pred_img = frankot_chellappa(-grad_x, grad_y, progress_callback=progress_callback) + + # Cut to valid part if gradients were expanded + if not seamless: + height, width = normals_img.shape[1], normals_img.shape[2] + pred_img = pred_img[:height, :width] + + # Expand single channel the three channels (RGB) + return np.stack([pred_img, pred_img, pred_img]) + + +# - ADDON +import numpy as np + +# import imageio.v3 as iio + + +class DeepBump: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "mode": ( + ["Color to Normals", "Normals to Curvature", "Normals to Height"], + ), + "color_to_normals_overlap": (["SMALL", "MEDIUM", "LARGE"],), + "normals_to_curvature_blur_radius": ( + [ + "SMALLEST", + "SMALLER", + "SMALL", + "MEDIUM", + "LARGE", + "LARGER", + "LARGEST", + ], + ), + "normals_to_height_seamless": (["TRUE", "FALSE"],), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply" + + CATEGORY = "image processing" + + def apply( + self, + image, + mode="Color to Normals", + color_to_normals_overlap="SMALL", + normals_to_curvature_blur_radius="SMALL", + normals_to_height_seamless="TRUE", + ): + + + image = utils_inference.tensor2pil(image) + + in_img = np.transpose(image, (2, 0, 1)) / 255 + + print(f"Input image shape: {in_img.shape}") + + # Apply processing + if mode == "Color to Normals": + out_img = color_to_normals(in_img, color_to_normals_overlap, None) + if mode == "Normals to Curvature": + out_img = normals_to_curvature( + in_img, normals_to_curvature_blur_radius, None + ) + if mode == "Normals to Height": + out_img = normals_to_height( + in_img, normals_to_height_seamless == "TRUE", None + ) + + out_img = (np.transpose(out_img, (1, 2, 0)) * 255).astype(np.uint8) + + return (utils_inference.pil2tensor(out_img),) diff --git a/nodes/fun.py b/nodes/fun.py new file mode 100644 index 0000000..591650c --- /dev/null +++ b/nodes/fun.py @@ -0,0 +1,43 @@ +import qrcode +from ..utils import pil2tensor, tensor2pil +from PIL import Image +class QRNode: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ("STRING", {"default": "https://www.github.com"}), + "width": ( + "INT", + {"default": 256, "max": 8096, "min": 0, "step": 1}, + ), + "height": ( + "INT", + {"default": 256, "max": 8096, "min": 0, "step": 1}, + ), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "do_qr" + CATEGORY = "fun" + + def do_qr(self, url, width, height): + qr = qrcode.QRCode( + version=1, + error_correction=qrcode.constants.ERROR_CORRECT_L, + box_size=10, + border=4, + ) + qr.add_data(url) + qr.make(fit=True) + # make the pil image + code = img = qr.make_image(back_color=(0, 0, 0), fill_color=(255, 255, 255)) + + # that we now resize without filtering + code = code.resize((width, height), Image.NEAREST) + + return (pil2tensor(code),) \ No newline at end of file diff --git a/nodes/latent_processing.py b/nodes/latent_processing.py new file mode 100644 index 0000000..b63e603 --- /dev/null +++ b/nodes/latent_processing.py @@ -0,0 +1,28 @@ +import torch + +class LatentLerp: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "A": ("LATENT",), + "B": ("LATENT",), + "t": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "lerp_latent" + + CATEGORY = "latent" + + def lerp_latent(self, A, B, t): + a = A.copy() + b = B.copy() + + torch.lerp(a["samples"], b["samples"], t, out=a["samples"]) + + return (a,) diff --git a/requirements.txt b/requirements.txt index da6cc83..d1946c9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ onnxruntime -imageio \ No newline at end of file +imageio +qrcode[pil] \ No newline at end of file diff --git a/utils.py b/utils.py index 0c42c8e..482dff4 100644 --- a/utils.py +++ b/utils.py @@ -1,6 +1,11 @@ from PIL import Image import numpy as np import torch +from pathlib import Path +import sys + +# Get the absolute path of the parent directory of the current script +here = Path(__file__).parent.resolve() # Tensor to PIL (grabbed from WAS Suite) def tensor2pil(image: torch.Tensor) -> Image.Image: return Image.fromarray(