custom mosaic node

This commit is contained in:
xss 2023-03-21 13:29:57 -05:00
parent 9cfbe4d4c7
commit 708ee6ef56
2 changed files with 45 additions and 3 deletions

3
.gitignore vendored
View File

@ -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

View File

@ -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
}