allow for looping over the batch dim

README.md

@@ -1,4 +1,3 @@
 #  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 the nodes into your `custom_nodes/` folder

canny_edge_detect.py

@@ -1,7 +1,8 @@
-import numpy as np
 import cv2
+import numpy as np
 import torch
 
+
 class CannyEdgeDetection:
     def __init__(self):
         pass
@@ -32,11 +33,17 @@ class CannyEdgeDetection:
     CATEGORY = "postprocessing"
 
     def canny(self, image: torch.Tensor, lower_threshold: int, upper_threshold: int):
-        tensor_image = image.numpy()[0]
-        gray_image = (cv2.cvtColor(tensor_image, cv2.COLOR_BGR2GRAY) * 255).astype(np.uint8)
+        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).unsqueeze(0)
-        return (tensor,)
+            tensor = torch.from_numpy(canny)
+            result[b] = tensor
+
+        return (result,)
 
 NODE_CLASS_MAPPINGS = {
     "CannyEdgeDetection": CannyEdgeDetection

color_correct.py

@@ -1,8 +1,9 @@
-import numpy as np
 import cv2
+import numpy as np
 import torch
 from PIL import Image, ImageEnhance
 
+
 class ColorCorrect:
     def __init__(self):
         pass
@@ -57,7 +58,11 @@ class ColorCorrect:
     CATEGORY = "postprocessing"
 
     def color_correct(self, image: torch.Tensor, temperature: float, hue: float, brightness: float, contrast: float, saturation: float, gamma: float):
-        tensor_image = image.numpy()[0]
+        batch_size, height, width, _ = image.shape
+        result = torch.zeros_like(image)
+
+        for b in range(batch_size):
+            tensor_image = image[b].numpy()
 
             brightness /= 100
             contrast /= 100
@@ -101,8 +106,9 @@ class ColorCorrect:
             modified_image = modified_image.astype(np.uint8)
             modified_image = modified_image / 255
             modified_image = torch.from_numpy(modified_image).unsqueeze(0)
+            result[b] = modified_image
 
-        return (modified_image, )
+        return (result, )
 
 NODE_CLASS_MAPPINGS = {
     "ColorCorrect": ColorCorrect,

dither.py

@@ -1,5 +1,6 @@
 import torch
 
+
 class Dither:
     def __init__(self):
         pass
@@ -24,7 +25,11 @@ class Dither:
     CATEGORY = "postprocessing"
 
     def dither(self, image: torch.Tensor, bits: int):
-        tensor_image = image.numpy()[0]
+        batch_size, height, width, _ = image.shape
+        result = torch.zeros_like(image)
+
+        for b in range(batch_size):
+            tensor_image = image[b].numpy()
             img = (tensor_image * 255)
             height, width, _ = img.shape
 
@@ -49,7 +54,9 @@ class Dither:
 
             dithered = img / 255
             tensor = torch.from_numpy(dithered).unsqueeze(0)
-        return (tensor,)
+            result[b] = tensor
+
+        return (result,)
 
 NODE_CLASS_MAPPINGS = {
     "Dither": Dither

gaussian_blur.py

@@ -1,6 +1,7 @@
 import cv2
 import torch
 
+
 class GaussianBlur:
     def __init__(self):
         pass
@@ -31,10 +32,16 @@ class GaussianBlur:
     CATEGORY = "postprocessing"
 
     def blur(self, image: torch.Tensor, kernel_size: int, sigma: float):
-        tensor_image = image.numpy()[0]
+        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)
-        return (tensor,)
+            result[b] = tensor
+
+        return (result,)
 
 NODE_CLASS_MAPPINGS = {
     "GaussianBlur": GaussianBlur

kmeans_quantize.py

@@ -1,7 +1,8 @@
-import numpy as np
 import cv2
+import numpy as np
 import torch
 
+
 class KMeansQuantize:
     def __init__(self):
         pass
@@ -32,7 +33,11 @@ class KMeansQuantize:
     CATEGORY = "postprocessing"
 
     def kmeans_quantize(self, image: torch.Tensor, colors: int, precision: int):
-        tensor_image = image.numpy()[0].astype(np.float32)
+        batch_size, height, width, _ = image.shape
+        result = torch.zeros_like(image)
+
+        for b in range(batch_size):
+            tensor_image = image[b].numpy().astype(np.float32)
             img = tensor_image
 
             height, width, c = img.shape
@@ -48,6 +53,12 @@ class KMeansQuantize:
                 criteria, 1, cv2.KMEANS_PP_CENTERS
             )
 
-        result = center[label.flatten()].reshape(*img.shape)
-        tensor = torch.from_numpy(result).unsqueeze(0)
-        return (tensor,)
+            img = center[label.flatten()].reshape(*img.shape)
+            tensor = torch.from_numpy(img).unsqueeze(0)
+            result[b] = tensor
+
+        return (result,)
+
+NODE_CLASS_MAPPINGS = {
+    "KMeansQuantize": KMeansQuantize
+}

pixel_sort.py

@@ -29,17 +29,22 @@ class PixelSort:
 
     CATEGORY = "postprocessing"
 
-    def sort_pixels(self, image, mask, direction, span_limit, sort_by, order):
+    def sort_pixels(self, image: torch.Tensor, mask: torch.Tensor, direction: str, span_limit: int, sort_by: str, order: str):
         horizontal_sort = direction == "horizontal"
         reverse_sorting = order == "backward"
         sort_by = sort_by[0].upper()
         span_limit = span_limit if span_limit > 0 else None
 
-        tensor_img = image.numpy()[0]
-        tensor_mask = mask.numpy()[0]
+        batch_size = image.shape[0]
+        result = torch.zeros_like(image)
+
+        for b in range(batch_size):
+            tensor_img = image[b].numpy()
+            tensor_mask = mask[b].numpy()
             sorted_image = pixel_sort(tensor_img, tensor_mask, horizontal_sort, span_limit, sort_by, reverse_sorting)
-        tensor = torch.from_numpy(sorted_image).unsqueeze(0)
-        return (tensor,)
+            result[b] = torch.from_numpy(sorted_image)
+
+        return (result,)
 
 NODE_CLASS_MAPPINGS = {
     "PixelSort": PixelSort,

sharpen.py

@@ -1,7 +1,8 @@
-import numpy as np
 import cv2
+import numpy as np
 import torch
 
+
 class Sharpen:
     def __init__(self):
         pass
@@ -32,7 +33,11 @@ class Sharpen:
     CATEGORY = "postprocessing"
 
     def sharpen(self, image: torch.Tensor, kernel_size: int, alpha: float):
-        tensor_image = image.numpy()[0]
+        batch_size, height, width, _ = image.shape
+        result = torch.zeros_like(image)
+
+        for b in range(batch_size):
+            tensor_image = image[b].numpy()
 
             kernel = np.ones((kernel_size, kernel_size), dtype=np.float32) * -1
             center = kernel_size // 2
@@ -43,7 +48,9 @@ class Sharpen:
 
             tensor = torch.from_numpy(sharpened).unsqueeze(0)
             tensor = torch.clamp(tensor, 0, 1)
-        return (tensor,)
+            result[b] = tensor
+
+        return (result,)
 
 NODE_CLASS_MAPPINGS = {
     "Sharpen": Sharpen