diff --git a/.gitignore b/.gitignore index 9a11d94..a35878e 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ __pycache__ .vscode *.ckpt *.safetensors -*.pth \ No newline at end of file +*.pth +types \ No newline at end of file diff --git a/nodes.py b/nodes.py index 37059dd..62df59b 100644 --- a/nodes.py +++ b/nodes.py @@ -1,4 +1,10 @@ import nodes +import torch +import torch.nn.functional as F +import scipy.ndimage +import numpy as np + +from nodes import MAX_RESOLUTION class INTConstant: @classmethod @@ -16,7 +22,93 @@ class INTConstant: def get_value(self, value): return (value,) + +def gaussian_kernel(kernel_size: int, sigma: float, device=None): + x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size, device=device), torch.linspace(-1, 1, kernel_size, device=device), indexing="ij") + d = torch.sqrt(x * x + y * y) + g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) + return g / g.sum() +class GrowMaskWithBlur: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK",), + "expand": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}), + "tapered_corners": ("BOOLEAN", {"default": True}), + "flip_input": ("BOOLEAN", {"default": False}), + "blur_radius": ("INT", { + "default": 1, + "min": 1, + "max": 31, + "step": 1 + }), + "sigma": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 10.0, + "step": 0.1 + }), + }, + } + + CATEGORY = "KJNodes" + + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "expand_mask" + + def expand_mask(self, mask, expand, tapered_corners, flip_input, blur_radius, sigma): + if( flip_input ): + mask = 1.0 - mask + c = 0 if tapered_corners else 1 + kernel = np.array([[c, 1, c], + [1, 1, 1], + [c, 1, c]]) + growmask = mask.reshape((-1, mask.shape[-2], mask.shape[-1])) + out = [] + for m in growmask: + output = m.numpy() + for _ in range(abs(expand)): + if expand < 0: + output = scipy.ndimage.grey_erosion(output, footprint=kernel) + else: + output = scipy.ndimage.grey_dilation(output, footprint=kernel) + output = torch.from_numpy(output) + out.append(output) + + blurred = torch.stack(out, dim=0).reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + batch_size, height, width, channels = blurred.shape + if blur_radius != 0: + blurkernel_size = blur_radius * 2 + 1 + blurkernel = gaussian_kernel(blurkernel_size, sigma, device=blurred.device).repeat(channels, 1, 1).unsqueeze(1) + blurred = blurred.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) + padded_image = F.pad(blurred, (blur_radius,blur_radius,blur_radius,blur_radius), 'reflect') + blurred = F.conv2d(padded_image, blurkernel, padding=blurkernel_size // 2, groups=channels)[:,:,blur_radius:-blur_radius, blur_radius:-blur_radius] + blurred = blurred.permute(0, 2, 3, 1) + blurred = blurred[:, :, :, 0] + + return (blurred, 1.0 - blurred,) + + + +class PlotNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "start": ("FLOAT", {"default": 0.5, "min": 0.5, "max": 1.0}), + "max_frames": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }} + + RETURN_TYPES = ("FLOAT", "INT",) + FUNCTION = "plot" + CATEGORY = "KJNodes" + + def plot(self, start, max_frames): + result = start + max_frames + return (result,) + class ConditioningMultiCombine: @classmethod def INPUT_TYPES(s): @@ -97,14 +189,18 @@ class ConditioningSetMaskAndCombine: n[1]['mask_strength'] = strength c2.append(n) return (c, c2) + + NODE_CLASS_MAPPINGS = { "INTConstant": INTConstant, "ConditioningMultiCombine": ConditioningMultiCombine, "ConditioningSetMaskAndCombine": ConditioningSetMaskAndCombine, + "GrowMaskWithBlur": GrowMaskWithBlur, } NODE_DISPLAY_NAME_MAPPINGS = { "INTConstant": "INT Constant", "ConditioningMultiCombine": "Conditioning Multi Combine", "ConditioningSetMaskAndCombine": "ConditioningSetMaskAndCombine", + "GrowMaskWithBlur": "GrowMaskWithBlur", } \ No newline at end of file diff --git a/web/js/jsnodes.js b/web/js/jsnodes.js index bd8998b..579a462 100644 --- a/web/js/jsnodes.js +++ b/web/js/jsnodes.js @@ -5,6 +5,9 @@ app.registerExtension({ async beforeRegisterNodeDef(nodeType, nodeData, app) { switch (nodeData.name) { case "ConditioningMultiCombine": + nodeType.prototype.onNodeMoved = function () { + console.log(this.pos[0]) + } nodeType.prototype.onNodeCreated = function () { this.inputs_offset = nodeData.name.includes("selective")?1:0 this.cond_type = "CONDITIONING" diff --git a/web/js/plotnode.js b/web/js/plotnode.js new file mode 100644 index 0000000..67c519a --- /dev/null +++ b/web/js/plotnode.js @@ -0,0 +1,30 @@ +import { app } from "../../../scripts/app.js"; +//WIP doesn't do anything +app.registerExtension({ + name: "KJNodes.PlotNode", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + switch (nodeData.name) { + case "PlotNode": + + nodeType.prototype.onNodeCreated = function () { + + this.addWidget("button", "Update", null, () => { + + console.log("start x:" + this.pos[0]) + console.log("start y:" +this.pos[1]) + console.log(this.graph.links); + const toNode = this.graph._nodes.find((otherNode) => otherNode.id == this.graph.links[1].target_id); + console.log("target x:" + toNode.pos[0]) + const a = this.pos[0] + const b = toNode.pos[0] + const distance = Math.abs(a - b); + const maxDistance = 1000 + const finalDistance = (distance - 0) / (maxDistance - 0); + + this.widgets[0].value = finalDistance; + }); + } + break; + } + }, +}); \ No newline at end of file diff --git a/web/js/setgetnodes.js b/web/js/setgetnodes.js index ef75e37..151b7ae 100644 --- a/web/js/setgetnodes.js +++ b/web/js/setgetnodes.js @@ -26,7 +26,9 @@ app.registerExtension({ '', (s, t, u, v, x) => { node.validateName(node.graph); - this.title = "Set_" + this.widgets[0].value; + if(this.widgets[0].value !== ''){ + this.title = "Set_" + this.widgets[0].value; + } this.update(); this.properties.previousName = this.widgets[0].value; }, @@ -47,9 +49,11 @@ app.registerExtension({ ) { //On Disconnect if (slotType == 1 && !isChangeConnect) { - this.inputs[slot].type = '*'; - this.inputs[slot].name = '*'; - this.title = "Set" + if(this.inputs[slot].name === ''){ + this.inputs[slot].type = '*'; + this.inputs[slot].name = '*'; + this.title = "Set" + } } if (slotType == 2 && !isChangeConnect) { this.outputs[slot].type = '*'; @@ -61,10 +65,12 @@ app.registerExtension({ const fromNode = node.graph._nodes.find((otherNode) => otherNode.id == link_info.origin_id); const type = fromNode.outputs[link_info.origin_slot].type; - if (this.title == "Set"){ + if (this.title === "Set"){ + console.log("setting title to Set_" + type); this.title = "Set_" + type; } - if (this.widgets[0].value == ''){ + if (this.widgets[0].value === '*'){ + console.log("setting default value to " + type); this.widgets[0].value = type } @@ -162,6 +168,7 @@ app.registerExtension({ this.update = function() { if (node.graph) { this.findGetters(node.graph).forEach((getter) => { + getter.setType(this.inputs[0].type); }); if (this.widgets[0].value) {