Fix widget inputs

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();
 

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({

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

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 {

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)