diff --git a/js/node_frequency_tracker.js b/js/node_frequency_tracker.js index e92551a..fac8077 100644 --- a/js/node_frequency_tracker.js +++ b/js/node_frequency_tracker.js @@ -8,6 +8,23 @@ const originalAdd = LGraph.prototype.add; LGraph.prototype.add = function (node, skip_compute_order) { originalAdd.apply(this, arguments); + + // Everytime created a node, update the frequency + function generateNodeKey(node) { + let type = node.type; + let inputs = node.inputs ? node.inputs.map(input => input.type).join(",") : ""; + let outputs = node.outputs ? node.outputs.map(output => output.type).join(",") : ""; + return `${type}|${inputs}|${outputs}`; + } + + var selectionFrequency = JSON.parse(localStorage.getItem("nodeSelectionFrequency") || "{}"); + let key = generateNodeKey(node); + + if (!selectionFrequency[key]) { + selectionFrequency[key] = 0; + } + selectionFrequency[key] += 1; + localStorage.setItem("nodeSelectionFrequency", JSON.stringify(selectionFrequency)); }; //From: https://github.com/jagenjo/litegraph.js/blob/0555a2f2a3df5d4657593c6d45eb192359888195/src/litegraph.js#L11182 @@ -66,6 +83,15 @@ LGraphCanvas.prototype.showConnectionMenu = function (optPass) { // addNodeMenu // get defaults nodes for this slottype var fromSlotType = slotX.type == LiteGraph.EVENT ? "_event_" : slotX.type; + + var topFrequentNodes = findTopFrequentNodes(isFrom, fromSlotType, 3); // Get top 3 frequent nodes + if (topFrequentNodes.length) { + topFrequentNodes.forEach(nodeType => { + options.push(nodeType); + }); + options.push(null); + } + var slotTypesDefault = isFrom ? LiteGraph.slot_types_default_out : LiteGraph.slot_types_default_in; if (slotTypesDefault && slotTypesDefault[fromSlotType]) { if (typeof slotTypesDefault[fromSlotType] == "object" || typeof slotTypesDefault[fromSlotType] == "array") { @@ -121,4 +147,24 @@ LGraphCanvas.prototype.showConnectionMenu = function (optPass) { // addNodeMenu } return false; -}; \ No newline at end of file +}; + +function findTopFrequentNodes(isFrom, fromSlotType, topN = 3) { + var selectionFrequency = JSON.parse(localStorage.getItem("nodeSelectionFrequency") || "{}"); + let candidates = []; + + for (let key in selectionFrequency) { + let [type, inputs, outputs] = key.split("|"); + if (isFrom ? inputs.includes(fromSlotType) : outputs.includes(fromSlotType)) { + candidates.push({ type: type, frequency: selectionFrequency[key] }); + } + } + + // Sort candidates by frequency in descending order + candidates.sort((a, b) => b.frequency - a.frequency); + + // Return the types of the top N candidates + return candidates.slice(0, topN).map(c => c.type); +} + +// To reset the tracker, use localStorage.removeItem('nodeSelectionFrequency'); \ No newline at end of file