diff --git a/.gitignore b/.gitignore index d5136f695..51692a3f2 100644 --- a/.gitignore +++ b/.gitignore @@ -5,7 +5,4 @@ input/ !input/example.png models/ temp/ -custom_nodes/ -!custom_nodes/example_node.py.example -!custom_nodes/mosaic_node.py extra_model_paths.yaml diff --git a/custom_nodes/mosaic_node.py b/custom_nodes/mosaic_node.py new file mode 100644 index 000000000..dc8394597 --- /dev/null +++ b/custom_nodes/mosaic_node.py @@ -0,0 +1,45 @@ +import torch + +class Mosaic: + interpolation_methods = ['nearest', 'bilinear', 'bicubic', 'area', 'nearest-exact'] + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "pixel_size": ("INT", { + "default": 4, + "min": 1, + "max": 512, + "step": 1 + }), + "interpolation_method": (s.interpolation_methods, {"default": s.interpolation_methods[0]}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "mosaic" + + CATEGORY = "image" + + def mosaic(self, image, interpolation_method, pixel_size): + samples = image.movedim(-1,1) + + starting_width = samples.shape[3] + starting_height = samples.shape[2] + + #downsample dimensions + dowsample_width = int(starting_width / pixel_size) | 1 + dowsample_height = int(starting_height / pixel_size) | 1 + + downsampled_image = torch.nn.functional.interpolate(samples, size=(dowsample_height, dowsample_width), mode=interpolation_method) + output_image = torch.nn.functional.interpolate(downsampled_image, size=(starting_height, starting_width), mode='nearest') + + output = output_image.movedim(1,-1) + + return (output,) + +NODE_CLASS_MAPPINGS = { + "Mosaic": Mosaic +} \ No newline at end of file