diff --git a/README.md b/README.md index 37705dd..79cfeda 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,9 @@ # ComfyUI-post-processing-nodes -A collection of post processing nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI), simply download this repo and drag the nodes into your `custom_nodes/` folder \ No newline at end of file +A collection of post processing nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI), simply download this repo and drag `combined_nodes.py` into your `custom_nodes/` folder + +## Combine Nodes + +By default `combined_nodes.py` should have all of the combined nodes. If you want a subset of nodes, you can run + + python combine_files.py [--files FILES [FILES ...]] [--output OUTPUT] diff --git a/combine_files.py b/combine_files.py new file mode 100644 index 0000000..7b3727d --- /dev/null +++ b/combine_files.py @@ -0,0 +1,111 @@ +from pathlib import Path +import sys +import os +import glob +import ast +import argparse + + +def get_python_files(path): + for file in Path(path).glob("*.py"): + if file.is_file() and not file.name.startswith("combine"): + yield str(file) + +def parse_files(files): + imports = set() + class_definitions = set() + node_class_mappings = set() + functions = set() + + for file in files: + # read file as lines + with open(file, "r") as f: + lines = f.readlines() + + # remove comments + lines = [line for line in lines if not line.startswith("#")] + + num_lines = len(lines) + i = 0 + while i < num_lines: + line = lines[i] + if line.startswith("import") or line.startswith("from"): + imports.add(line.strip()) + + elif line.startswith("class"): + class_info = line + j = i + 1 + while not lines[j].startswith("NODE_CLASS_MAPPINGS"): + class_info += lines[j] + j += 1 + class_definitions.add(class_info) + i = j - 1 + + elif line.startswith("NODE_CLASS_MAPPINGS"): + node_class_mappings.add(lines[i+1]) + + elif line.startswith("def"): + function_info = line + j = i + 1 + while j < num_lines and not lines[j].startswith("NODE_CLASS_MAPPINGS") and not lines[j].startswith("def"): + function_info += lines[j] + j += 1 + functions.add(function_info) + i = j - 1 + + i += 1 + + return imports, class_definitions, node_class_mappings, functions + +def write_combined(imports: list[str], class_definitions: list[str], node_class_mappings: list[str], functions: list[str], output_file: str): + with open(output_file, "w") as f: + # write imports + for line in imports: + f.write(line + "\n") + + # write 2 blank lines + f.write("\n\n") + + # write class definitions + for line in class_definitions: + f.write(line) + + # write functions + for line in functions: + f.write(line + "\n") + + # write node class mappings + f.write("NODE_CLASS_MAPPINGS = {\n") + for line in node_class_mappings: + if not line.endswith(",\n"): + if line.endswith("\n"): + line = line[:-1] + ",\n" + else: + line += ",\n" + + f.write(line) + f.write("}\n") + + +def main(): + parser = argparse.ArgumentParser(description="Collect unique imports from Python files") + parser.add_argument("--all", action="store_true", help="Include all Python files in the specified directory") + parser.add_argument("--files", nargs="+", help="Specify Python files to parse") + parser.add_argument("--output", default="combined_nodes.py", help="Specify the output file name") + args = parser.parse_args() + + if args.all: + args.path = "." + files = get_python_files(args.path) + imports, class_definitions, node_class_mappings, functions = parse_files(files) + write_combined(imports, class_definitions, node_class_mappings, functions, args.output) + else: + if args.files is None: + print("No files specified") + return + + imports, class_definitions, node_class_mappings, functions = parse_files(args.files) + write_combined(imports, class_definitions, node_class_mappings, functions, args.output) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/combined_nodes.py b/combined_nodes.py new file mode 100644 index 0000000..b9788cb --- /dev/null +++ b/combined_nodes.py @@ -0,0 +1,92 @@ +import cv2 +import numpy as np +import torch + + +class CannyEdgeDetection: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "lower_threshold": ("INT", { + "default": 100, + "min": 0, + "max": 500, + "step": 10 + }), + "upper_threshold": ("INT", { + "default": 200, + "min": 0, + "max": 500, + "step": 10 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "canny" + + CATEGORY = "postprocessing" + + def canny(self, image: torch.Tensor, lower_threshold: int, upper_threshold: int): + batch_size, height, width, _ = image.shape + result = torch.zeros(batch_size, height, width) + + for b in range(batch_size): + tensor_image = image[b].numpy().copy() + gray_image = (cv2.cvtColor(tensor_image, cv2.COLOR_RGB2GRAY) * 255).astype(np.uint8) + canny = cv2.Canny(gray_image, lower_threshold, upper_threshold) + tensor = torch.from_numpy(canny) + result[b] = tensor + + return (result,) + +class GaussianBlur: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "kernel_size": ("INT", { + "default": 5, + "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 blur(self, image: torch.Tensor, kernel_size: int, sigma: float): + batch_size, height, width, _ = image.shape + result = torch.zeros_like(image) + + for b in range(batch_size): + tensor_image = image[b].numpy() + blurred = cv2.GaussianBlur(tensor_image, (kernel_size, kernel_size), sigma) + tensor = torch.from_numpy(blurred).unsqueeze(0) + result[b] = tensor + + return (result,) + +NODE_CLASS_MAPPINGS = { + "CannyEdgeDetection": CannyEdgeDetection, + "GaussianBlur": GaussianBlur, +}