adds pixelize effect

post_processing/color_correct.py

@@ -61,18 +61,18 @@ class ColorCorrect:
         batch_size, height, width, _ = image.shape
         result = torch.zeros_like(image)
 
+        brightness /= 100
+        contrast /= 100
+        saturation /= 100
+        temperature /= 100
+
+        brightness = 1 + brightness
+        contrast = 1 + contrast
+        saturation = 1 + saturation
+
         for b in range(batch_size):
             tensor_image = image[b].numpy()
 
-            brightness /= 100
-            contrast /= 100
-            saturation /= 100
-            temperature /= 100
-
-            brightness = 1 + brightness
-            contrast = 1 + contrast
-            saturation = 1 + saturation
-
             modified_image = Image.fromarray((tensor_image * 255).astype(np.uint8))
 
             # brightness

post_processing/pixelize.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn.functional as F
+
+class Pixelize:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "pixel_size": ("INT", {
+                    "default": 8,
+                    "min": 2,
+                    "max": 128,
+                    "step": 1
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "apply_pixelize"
+
+    CATEGORY = "postprocessing"
+
+    def apply_pixelize(self, image: torch.Tensor, pixel_size: int):
+        pixelized_image = self.pixelize_image(image, pixel_size)
+        pixelized_image = torch.clamp(pixelized_image, 0, 1)
+        return (pixelized_image,)
+
+    def pixelize_image(self, image: torch.Tensor, pixel_size: int):
+        batch_size, height, width, channels = image.shape
+        new_height = height // pixel_size
+        new_width = width // pixel_size
+
+        image = image.permute(0, 3, 1, 2)
+        image = F.avg_pool2d(image, kernel_size=pixel_size, stride=pixel_size)
+        image = F.interpolate(image, size=(height, width), mode='nearest')
+        image = image.permute(0, 2, 3, 1)
+
+        return image
+
+NODE_CLASS_MAPPINGS = {
+    "Pixelize": Pixelize,
+}