File size: 6,150 Bytes
1e3b872 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import { app } from "../../../scripts/app.js";
import { addMenuHandler } from "./common/utils.js";
import { addNode } from "./common/utils.js";
const connectionMap = {
"KSampler (Efficient)": ["input", 5],
"KSampler Adv. (Efficient)": ["input", 5],
"KSampler SDXL (Eff.)": ["input", 3],
"XY Plot": ["output", 0],
"Noise Control Script": ["input & output", 0],
"HighRes-Fix Script": ["input & output", 0],
"Tiled Upscaler Script": ["input & output", 0],
"AnimateDiff Script": ["output", 0]
};
/**
* connect this node output to the input of another node
* @method connect
* @param {number_or_string} slot (could be the number of the slot or the string with the name of the slot)
* @param {LGraphNode} node the target node
* @param {number_or_string} target_slot the input slot of the target node (could be the number of the slot or the string with the name of the slot, or -1 to connect a trigger)
* @return {Object} the link_info is created, otherwise null
LGraphNode.prototype.connect = function(output_slot, target_node, input_slot)
**/
function addAndConnectScriptNode(scriptType, selectedNode) {
const selectedNodeType = connectionMap[selectedNode.type];
const newNodeType = connectionMap[scriptType];
// 1. Create the new node without position adjustments
const newNode = addNode(scriptType, selectedNode, { shiftX: 0, shiftY: 0 });
// 2. Adjust position of the new node based on conditions
if (newNodeType[0].includes("input") && selectedNodeType[0].includes("output")) {
newNode.pos[0] += selectedNode.size[0] + 50;
} else if (newNodeType[0].includes("output") && selectedNodeType[0].includes("input")) {
newNode.pos[0] -= (newNode.size[0] + 50);
}
// 3. Logic for connecting the nodes
switch (selectedNodeType[0]) {
case "output":
if (newNodeType[0] === "input & output") {
// For every node that was previously connected to the selectedNode's output
const connectedNodes = selectedNode.getOutputNodes(selectedNodeType[1]);
if (connectedNodes && connectedNodes.length) {
for (let connectedNode of connectedNodes) {
// Disconnect the node from selectedNode's output
selectedNode.disconnectOutput(selectedNodeType[1]);
// Connect the newNode's output to the previously connected node,
// using the appropriate slot based on the type of the connectedNode
const targetSlot = (connectedNode.type in connectionMap) ? connectionMap[connectedNode.type][1] : 0;
newNode.connect(0, connectedNode, targetSlot);
}
}
// Connect selectedNode's output to newNode's input
selectedNode.connect(selectedNodeType[1], newNode, newNodeType[1]);
}
break;
case "input":
if (newNodeType[0] === "output") {
newNode.connect(0, selectedNode, selectedNodeType[1]);
} else if (newNodeType[0] === "input & output") {
const ogInputNode = selectedNode.getInputNode(selectedNodeType[1]);
if (ogInputNode) {
ogInputNode.connect(0, newNode, 0);
}
newNode.connect(0, selectedNode, selectedNodeType[1]);
}
break;
case "input & output":
if (newNodeType[0] === "output") {
newNode.connect(0, selectedNode, 0);
} else if (newNodeType[0] === "input & output") {
const connectedNodes = selectedNode.getOutputNodes(0);
if (connectedNodes && connectedNodes.length) {
for (let connectedNode of connectedNodes) {
selectedNode.disconnectOutput(0);
newNode.connect(0, connectedNode, connectedNode.type in connectionMap ? connectionMap[connectedNode.type][1] : 0);
}
}
// Connect selectedNode's output to newNode's input
selectedNode.connect(selectedNodeType[1], newNode, newNodeType[1]);
}
break;
}
return newNode;
}
function createScriptEntry(node, scriptType) {
return {
content: scriptType,
callback: function() {
addAndConnectScriptNode(scriptType, node);
},
};
}
function getScriptOptions(nodeType, node) {
const allScriptTypes = [
"XY Plot",
"Noise Control Script",
"HighRes-Fix Script",
"Tiled Upscaler Script",
"AnimateDiff Script"
];
// Filter script types based on node type
const scriptTypes = allScriptTypes.filter(scriptType => {
const scriptBehavior = connectionMap[scriptType][0];
if (connectionMap[nodeType][0] === "output") {
return scriptBehavior.includes("input"); // Includes nodes that are "input" or "input & output"
} else {
return true;
}
});
return scriptTypes.map(script => createScriptEntry(node, script));
}
function showAddScriptMenu(_, options, e, menu, node) {
const scriptOptions = getScriptOptions(node.type, node);
new LiteGraph.ContextMenu(scriptOptions, {
event: e,
callback: null,
parentMenu: menu,
node: node
});
return false;
}
// Extension Definition
app.registerExtension({
name: "efficiency.addScripts",
async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (connectionMap[nodeData.name]) {
addMenuHandler(nodeType, function(insertOption) {
insertOption({
content: "📜 Add script...",
has_submenu: true,
callback: showAddScriptMenu
});
});
}
},
});
|