diff --git a/web/extensions/core/contextMenuFilter.js b/web/extensions/core/contextMenuFilter.js index a7f9e7099..c2480c81c 100644 --- a/web/extensions/core/contextMenuFilter.js +++ b/web/extensions/core/contextMenuFilter.js @@ -7,8 +7,9 @@ import { hook } from "../../scripts/utils.js"; const ext = { name: "Comfy.ContextMenuFilter", init() { - hook(LiteGraph, "onContextMenuCreated", (orig, contextMenu) => { - orig?.(contextMenu); + hook(LiteGraph, "onContextMenuCreated", (orig, args) => { + orig?.(...args); + const contextMenu = args[0]; // If we are a dark menu (only used for combo boxes) then add a filter input if (contextMenu.options?.className === "dark" && contextMenu.values?.length > 10) { @@ -32,7 +33,7 @@ const ext = { let selectedIndex = clickedComboValue ? contextMenu.values.findIndex(v => v === clickedComboValue) : 0; if (selectedIndex < 0) { selectedIndex = 0; - } + } let selectedItem = displayedItems[selectedIndex]; updateSelected(); diff --git a/web/extensions/core/invertMenuScrolling.js b/web/extensions/core/invertMenuScrolling.js index fa5c44c57..7d555d724 100644 --- a/web/extensions/core/invertMenuScrolling.js +++ b/web/extensions/core/invertMenuScrolling.js @@ -8,8 +8,9 @@ const id = "Comfy.InvertMenuScrolling"; app.registerExtension({ name: id, init() { - hook(LiteGraph, "onContextMenuCreated", (orig, contextMenu) => { - orig?.(contextMenu); + hook(LiteGraph, "onContextMenuCreated", (orig, args) => { + orig?.(...args) + const contextMenu = args[0]; contextMenu.options.invert_scrolling = localStorage[`Comfy.Settings.${id}`] === "true"; }) app.ui.settings.addSetting({ diff --git a/web/extensions/core/nodeTemplates.js b/web/extensions/core/nodeTemplates.js index f6ba71cf8..0c5753003 100644 --- a/web/extensions/core/nodeTemplates.js +++ b/web/extensions/core/nodeTemplates.js @@ -1,4 +1,5 @@ import { app } from "../../scripts/app.js"; +import { hook } from "../../scripts/utils.js"; import { ComfyDialog, $el } from "../../scripts/ui.js"; import { LGraphCanvas } from "../../lib/litegraph.core.js" @@ -131,9 +132,8 @@ app.registerExtension({ localStorage.setItem("litegrapheditor_clipboard", old); }; - const orig = LGraphCanvas.prototype.getCanvasMenuOptions; - LGraphCanvas.prototype.getCanvasMenuOptions = function () { - const options = orig.apply(this, arguments); + hook(LGraphCanvas, "getCanvasMenuOptions", function (orig, args) { + const options = orig.apply(this, args); options.push(null); options.push({ @@ -180,6 +180,6 @@ app.registerExtension({ } return options; - }; + }); }, }); diff --git a/web/extensions/core/widgetInputs.js b/web/extensions/core/widgetInputs.js index 7e00771cd..33cd4530e 100644 --- a/web/extensions/core/widgetInputs.js +++ b/web/extensions/core/widgetInputs.js @@ -1,5 +1,6 @@ import { ComfyWidgets, addValueControlWidget } from "../../scripts/widgets.js"; import { app } from "../../scripts/app.js"; +import { hook } from "../../scripts/utils.js"; import { ComfyGraphNode } from "../../scripts/graphNode.js"; import { LiteGraph } from "../../lib/litegraph.core.js" @@ -97,9 +98,9 @@ app.registerExtension({ name: "Comfy.WidgetInputs", async beforeRegisterNodeDef(nodeType, nodeData, app) { // Add menu options to conver to/from widgets - const origGetExtraMenuOptions = nodeType.class.prototype.getExtraMenuOptions; - nodeType.class.prototype.getExtraMenuOptions = function (_, options) { - const r = origGetExtraMenuOptions ? origGetExtraMenuOptions.apply(this, arguments) : undefined; + hook(nodeType.class, "getExtraMenuOptions", function(origGetExtraMenuOptions, args) { + const r = origGetExtraMenuOptions ? origGetExtraMenuOptions.apply(this, args) : undefined; + const options = args[1]; if (this.widgets) { let toInput = []; @@ -130,12 +131,11 @@ app.registerExtension({ } return r; - }; + }) // On initial configure of nodes hide all converted widgets - const origOnConfigure = nodeType.class.prototype.onConfigure; - nodeType.class.prototype.onConfigure = function () { - const r = origOnConfigure ? origOnConfigure.apply(this, arguments) : undefined; + hook(nodeType.class, "onConfigure", function(origOnConfigure, args) { + const r = origOnConfigure ? origOnConfigure.apply(this, args) : undefined; if (this.inputs) { for (const input of this.inputs) { @@ -151,7 +151,7 @@ app.registerExtension({ } return r; - }; + }); function isNodeAtPos(pos) { for (const n of app.graph._nodes) { @@ -163,10 +163,10 @@ app.registerExtension({ } // Double click a widget input to automatically attach a primitive - const origOnInputDblClick = nodeType.class.prototype.onInputDblClick; const ignoreDblClick = Symbol(); - nodeType.class.prototype.onInputDblClick = function (slot) { - const r = origOnInputDblClick ? origOnInputDblClick.apply(this, arguments) : undefined; + hook(nodeType.class, "onInputDblClick", function (origOnInputDblClick, args) { + const r = origOnInputDblClick ? origOnInputDblClick.apply(this, args) : undefined; + const slot = args[0]; const input = this.inputs[slot]; if (!input.widget || !input[ignoreDblClick]) { @@ -197,7 +197,7 @@ app.registerExtension({ }, 300); return r; - }; + }); }, registerCustomNodes() { class PrimitiveNode extends ComfyGraphNode { diff --git a/web/scripts/utils.js b/web/scripts/utils.js index 686f570af..b21c6705e 100644 --- a/web/scripts/utils.js +++ b/web/scripts/utils.js @@ -2,22 +2,25 @@ export function range(size, startAt = 0) { return [...Array(size).keys()].map(i => i + startAt); } +function isClass(obj) { + const isCtorClass = obj.constructor + && obj.constructor.toString().substring(0, 5) === 'class' + if(obj.prototype === undefined) { + return isCtorClass + } + const isPrototypeCtorClass = obj.prototype.constructor + && obj.prototype.constructor.toString + && obj.prototype.constructor.toString().substring(0, 5) === 'class' + return isCtorClass || isPrototypeCtorClass +} + export function hook(klass, fnName, cb) { - const fnLocation = klass; + let fnLocation = klass; + if (isClass(klass)) { + fnLocation = klass.prototype; + } const orig = fnLocation[fnName]; - fnLocation[fnName] = (...args) => { - cb(orig, ...args); + fnLocation[fnName] = function(...args) { + return cb.bind(this)(orig, args); } } - -const a = { - b: (c, d, e) => { - console.log(c, d, e); - } -} - -a.b(1,2,3) -hook(a, "b", (orig, c, d, e) => { - orig(c, d, e); -}) -a.b(1,2,3)