59 lines
1.6 KiB
Python
59 lines
1.6 KiB
Python
import torch
|
|
|
|
class Dither:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"bits": ("INT", {
|
|
"default": 4,
|
|
"min": 0,
|
|
"max": 8,
|
|
"step": 1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "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
|
|
|
|
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
|
|
|
|
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
|
|
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
|
|
|
|
out = torch.clamp(out, 0, 1).unsqueeze(0)
|
|
|
|
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
|
|
}
|