slight refactoring of existing nodes

dither.py

@@ -3,15 +3,15 @@ import torch
 class Dither:
     def __init__(self):
         pass
-    
+
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
                 "image": ("IMAGE",),
                 "bits": ("INT", {
-                    "default": 4, 
-                    "min": 0,
+                    "default": 4,
+                    "min": 1,
                     "max": 8,
                     "step": 1
                 }),
@@ -23,36 +23,34 @@ class Dither:
 
     CATEGORY = "postprocessing"
 
-    def dither(self, image, bits):
-        tensor_img = image[0]
-        height, width, _ = tensor_img.shape
-        out = tensor_img.clone()
-        levels = 2 ** bits - 1
+    def dither(self, image: torch.Tensor, bits: int):
+        tensor_image = image.numpy()[0]
+        img = (tensor_image * 255)
+        height, width, _ = img.shape
+
+        scale = 255 / (2**bits - 1)
 
         for y in range(height):
             for x in range(width):
-                old_pixel = out[y, x].clone()
-                new_pixel = torch.round(old_pixel * levels) / levels
-                out[y, x] = new_pixel
+                old_pixel = img[y, x].copy()
+                new_pixel = np.round(old_pixel / scale) * scale
+                img[y, x] = new_pixel
 
-                error = old_pixel - new_pixel
+                quant_error = old_pixel - new_pixel
 
                 if x + 1 < width:
-                    out[y, x + 1] += error * (7 / 16)
-                if x - 1 >= 0 and y + 1 < height:
-                    out[y + 1, x - 1] += error * 3/16
+                    img[y, x + 1] += quant_error * 7 / 16
                 if y + 1 < height:
-                    out[y + 1, x] += error * 5/16
-                if x + 1 < width and y + 1 < height:
-                    out[y + 1, x + 1] += error * 1/16
+                    if x - 1 >= 0:
+                        img[y + 1, x - 1] += quant_error * 3 / 16
+                    img[y + 1, x] += quant_error * 5 / 16
+                    if x + 1 < width:
+                        img[y + 1, x + 1] += quant_error * 1 / 16
 
-        out = torch.clamp(out, 0, 1).unsqueeze(0)
+        dithered = img / 255
+        tensor = torch.from_numpy(dithered).unsqueeze(0)
+        return (tensor,)
 
-        return (out,)
-
-
-# A dictionary that contains all nodes you want to export with their names
-# NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {
     "Dither": Dither
 }