diff --git a/web/js/modelInfo.js b/web/js/modelInfo.js index d4c39ce..12f8c12 100644 --- a/web/js/modelInfo.js +++ b/web/js/modelInfo.js @@ -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); + }; }, });