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
                });
            });
        }
    },
});