add combined_nodes file and generator

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
+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]

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()

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,
+}