Support custom + multiple models per node

2024-04-07 17:03:31 +01:00
parent b8f21cfb4a
commit 3f2c021e50
1 changed files with 83 additions and 30 deletions
@@ -228,41 +228,94 @@ class CheckpointInfoDialog extends ModelInfoDialog {
 	}
 }

-const infoHandler = {
-	LoraLoader: "loras",
-	"LoraLoader|pysssss": "loras",
-	CheckpointLoader: "checkpoints",
-	CheckpointLoaderSimple: "checkpoints",
-	"CheckpointLoader|pysssss": "checkpoints",
-	"Efficient Loader": "checkpoints",
-	"Eff. Loader SDXL": "checkpoints",
-};
+const lookups = {};
+
+function addInfoOption(node, typeName, type, infoClass, widgetName, opts, useValue) {
+	let value = node.widgets.find((w) => w.name === widgetName)?.value;
+	if (value?.content) {
+		value = value.content;
+	}
+	if (!value) {
+		return;
+	}
+	let optName;
+	if (useValue) {
+		const split = value.split(/\.|\\|\//);
+		optName = split[split.length - 2];
+	} else {
+		optName = `View ${typeName} info...`;
+	}
+	opts.unshift({
+		content: optName,
+		callback: async () => {
+			new infoClass(value).show(type, value);
+		},
+	});
+}
+
+function addTypeOptions(node, typeName, options) {
+	const type = typeName.toLowerCase() + "s";
+	const values = lookups[typeName][node.type];
+	if (!values) return;
+
+	const widgets = Object.keys(values);
+	const cls = type === "loras" ? LoraInfoDialog : CheckpointInfoDialog;
+	if (widgets.length === 1) {
+		addInfoOption(node, typeName, type, cls, widgets[0] || node.widgets[0].name, options);
+	} else {
+		const opts = [];
+		for (const w of widgets) {
+			addInfoOption(node, typeName, type, cls, w, opts, true);
+		}
+		if (opts.length) {
+			options.unshift({
+				title: `View ${typeName} info...`,
+				has_submenu: true,
+				submenu: {
+					options: opts,
+				},
+			});
+		}
+	}
+}

 app.registerExtension({
 	name: "pysssss.ModelInfo",
+	setup() {
+		const addSetting = (type, defaultValue) => {
+			app.ui.settings.addSetting({
+				id: `pysssss.ModelInfo.${type}Nodes`,
+				name: `🐍 Model Info - ${type} Nodes/Widgets`,
+				type: "text",
+				defaultValue,
+				tooltip: `Comma separated list of NodeTypeName or NodeTypeName.WidgetName that contain ${type} node names that should have the View Info option available.\nIf no widget name is specifed the first widget will be used.`,
+				onChange(value) {
+					lookups[type] = value.split(",").reduce((p, n) => {
+						const split = n.trim().split(".");
+						p[split[0]] ??= {};
+						p[split[0]][split[1] ?? ""] = true;
+						return p;
+					}, {});
+				},
+			});
+		};
+		addSetting("Lora", ["LoraLoader.lora_name", "LoraLoader|pysssss", "LoraLoaderModelOnly.lora_name"].join(","));
+		addSetting(
+			"Checkpoint",
+			["CheckpointLoader.ckpt_name", "CheckpointLoaderSimple", "CheckpointLoader|pysssss", "Efficient Loader", "Eff. Loader SDXL"].join(",")
+		);
+		console.log(lookups);
+	},
 	beforeRegisterNodeDef(nodeType) {
-		const type = infoHandler[nodeType.comfyClass];
-
-		if (type) {
-			const cls = type === "loras" ? LoraInfoDialog : CheckpointInfoDialog;
-			const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions;
-			nodeType.prototype.getExtraMenuOptions = function (_, options) {
-				let value = this.widgets[0].value;
-				if (!value) {
-					return;
+		const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions;
+		nodeType.prototype.getExtraMenuOptions = function (_, options) {
+			if (this.widgets) {
+				for (const type in lookups) {
+					addTypeOptions(this, type, options);
 				}
-				if (value.content) {
-					value = value.content;
-				}
-				options.unshift({
-					content: "View info...",
-					callback: async () => {
-						new cls(value).show(type, value);
-					},
-				});
+			}

-				return getExtraMenuOptions?.apply(this, arguments);
-			};
-		}
+			return getExtraMenuOptions?.apply(this, arguments);
+		};
 	},
 });