From 7a7e3288ee413a600bb945d6b2df235991904b99 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Thu, 23 Mar 2023 21:37:19 +0000 Subject: [PATCH] Added support for converting widgets to inputs (and back) --- web/extensions/core/widgetInputs.js | 315 ++++++++++++++++++++++++++++ web/jsconfig.json | 9 + web/scripts/app.js | 31 +-- web/scripts/widgets.js | 50 ++++- 4 files changed, 385 insertions(+), 20 deletions(-) create mode 100644 web/extensions/core/widgetInputs.js create mode 100644 web/jsconfig.json diff --git a/web/extensions/core/widgetInputs.js b/web/extensions/core/widgetInputs.js new file mode 100644 index 00000000..4ee77802 --- /dev/null +++ b/web/extensions/core/widgetInputs.js @@ -0,0 +1,315 @@ +import { ComfyWidgets, addRandomizeWidget } from "/scripts/widgets.js"; +import { app } from "/scripts/app.js"; + +const CONVERTED_TYPE = "converted-widget"; +const VALID_TYPES = ["STRING", "combo", "number"]; + +function isConvertableWidget(widget, config) { + return VALID_TYPES.includes(widget.type) || VALID_TYPES.includes(config[0]); +} + +function hideWidget(node, widget, suffix = "") { + widget.origType = widget.type; + widget.origComputeSize = widget.computeSize; + widget.origSerializeValue = widget.serializeValue; + widget.computeSize = () => [0, -4]; // -4 is due to the gap litegraph adds between widgets automatically + widget.type = CONVERTED_TYPE + suffix; + widget.serializeValue = () => { + // Prevent serializing the widget if we have no input linked + const { link } = node.inputs.find((i) => i.widget?.name === widget.name); + if (link == null) { + return undefined; + } + return widget.value; + }; + + // Hide any linked widgets, e.g. seed+randomize + if (widget.linkedWidgets) { + for (const w of widget.linkedWidgets) { + hideWidget(node, w, ":" + widget.name); + } + } +} + +function showWidget(widget) { + widget.type = widget.origType; + widget.computeSize = widget.origComputeSize; + widget.serializeValue = widget.origSerializeValue; + + delete widget.origType; + delete widget.origComputeSize; + delete widget.origSerializeValue; + + // Hide any linked widgets, e.g. seed+randomize + if (widget.linkedWidgets) { + for (const w of widget.linkedWidgets) { + showWidget(w); + } + } +} + +function convertToInput(node, widget, config) { + hideWidget(node, widget); + + const { linkType } = getWidgetType(config); + + // Add input and store widget config for creating on primitive node + node.addInput(widget.name, linkType, { + widget: { name: widget.name, config }, + }); +} + +function convertToWidget(node, widget) { + showWidget(widget); + node.removeInput(node.inputs.findIndex((i) => i.widget?.name === widget.name)); +} + +function getWidgetType(config) { + // Special handling for COMBO so we restrict links based on the entries + let type = config[0]; + let linkType = type; + if (type instanceof Array) { + type = "COMBO"; + linkType = linkType.join(","); + } + return { type, linkType }; +} + +app.registerExtension({ + name: "Comfy.WidgetInputs", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + // Add menu options to conver to/from widgets + const origGetExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; + nodeType.prototype.getExtraMenuOptions = function (_, options) { + const r = origGetExtraMenuOptions ? origGetExtraMenuOptions.apply(this, arguments) : undefined; + + if (this.widgets) { + let toInput = []; + let toWidget = []; + for (const w of this.widgets) { + if (w.type === CONVERTED_TYPE) { + toWidget.push({ + content: `Convert ${w.name} to widget`, + callback: () => convertToWidget(this, w), + }); + } else { + const config = nodeData?.input?.required[w.name] || [w.type, w.options || {}]; + if (isConvertableWidget(w, config)) { + toInput.push({ + content: `Convert ${w.name} to input`, + callback: () => convertToInput(this, w, config), + }); + } + } + } + if (toInput.length) { + options.push(...toInput, null); + } + + if (toWidget.length) { + options.push(...toWidget, null); + } + } + + return r; + }; + + // On initial configure of nodes hide all converted widgets + const origOnConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function () { + const r = origOnConfigure ? origOnConfigure.apply(this, arguments) : undefined; + + if (this.inputs) { + for (const input of this.inputs) { + if (input.widget) { + const w = this.widgets.find((w) => w.name === input.widget.name); + hideWidget(this, w); + } + } + } + + return r; + }; + + // Double click a widget input to automatically attach a primitive + const origOnInputDblClick = nodeType.prototype.onInputDblClick; + nodeType.prototype.onInputDblClick = function (slot) { + const r = origOnInputDblClick ? origOnInputDblClick.apply(this, arguments) : undefined; + + if (this.inputs[slot].widget) { + const node = LiteGraph.createNode("PrimitiveNode"); + app.graph.add(node); + node.pos = [this.pos[0] - node.size[0] - 30, this.pos[1]]; + node.connect(0, this, slot); + } + + return r; + }; + }, + registerCustomNodes() { + class PrimitiveNode { + constructor() { + this.addOutput("connect to widget input", "*"); + this.serialize_widgets = true; + this.isVirtualNode = true; + } + + applyToGraph() { + if (!this.outputs[0].links?.length) return; + + // For each output link copy our value over the original widget value + for (const l of this.outputs[0].links) { + const linkInfo = app.graph.links[l]; + const node = this.graph.getNodeById(linkInfo.target_id); + const input = node.inputs[linkInfo.target_slot]; + const widgetName = input.widget.name; + if (widgetName) { + const widget = node.widgets.find((w) => w.name === widgetName); + if (widget) { + widget.value = this.widgets[0].value; + if (widget.callback) { + widget.callback(widget.value, app.canvas, node, app.canvas.graph_mouse, {}); + } + } + } + } + } + + onConnectionsChange(_, index, connected) { + if (connected) { + if (this.outputs[0].links?.length) { + if (!this.widgets?.length) { + this.#onFirstConnection(); + } + if (!this.widgets?.length && this.outputs[0].widget) { + // On first load it often cant recreate the widget as the other node doesnt exist yet + // Manually recreate it from the output info + this.#createWidget(this.outputs[0].widget.config); + } + } + } else if (!this.outputs[0].links?.length) { + this.#onLastDisconnect(); + } + } + + onConnectOutput(slot, type, input, target_node, target_slot) { + // Fires before the link is made allowing us to reject it if it isn't valid + + // No widget, we cant connect + if (!input.widget) return false; + + if (this.outputs[slot].links?.length) { + return this.#isValidConnection(input); + } + } + + #onFirstConnection() { + // First connection can fire before the graph is ready on initial load so random things can be missing + const linkId = this.outputs[0].links[0]; + const link = this.graph.links[linkId]; + if (!link) return; + + const theirNode = this.graph.getNodeById(link.target_id); + if (!theirNode || !theirNode.inputs) return; + + const input = theirNode.inputs[link.target_slot]; + if (!input) return; + + const widget = input.widget; + const { type, linkType } = getWidgetType(widget.config); + + // Update our output to restrict to the widget type + this.outputs[0].type = linkType; + this.outputs[0].name = type; + this.outputs[0].widget = widget; + + this.#createWidget(widget.config, theirNode, widget.name); + } + + #createWidget(inputData, node, widgetName) { + let type = inputData[0]; + + if (type instanceof Array) { + type = "COMBO"; + } + + let widget; + if (type in ComfyWidgets) { + widget = (ComfyWidgets[type](this, "value", inputData, app) || {}).widget; + } else { + widget = this.addWidget(type, "value", null, () => {}, {}); + } + + if (node?.widgets && widget) { + const theirWidget = node.widgets.find((w) => w.name === widgetName); + if (theirWidget) { + widget.value = theirWidget.value; + } + } + + if (widget.type === "number") { + addRandomizeWidget(this, widget, "Random after every gen"); + } + + // Grow our node if required + const sz = this.computeSize(); + if (this.size[0] < sz[0]) { + this.size[0] = sz[0]; + } + if (this.size[1] < sz[1]) { + this.size[1] = sz[1]; + } + + requestAnimationFrame(() => { + if (this.onResize) { + this.onResize(this.size); + } + }); + } + + #isValidConnection(input) { + // Only allow connections where the configs match + const config1 = this.outputs[0].widget.config; + const config2 = input.widget.config; + + if (config1[0] !== config2[0]) return false; + + for (const k in config1[1]) { + if (k !== "default") { + if (config1[1][k] !== config2[1][k]) { + return false; + } + } + } + + return true; + } + + #onLastDisconnect() { + // We cant remove + re-add the output here as if you drag a link over the same link + // it removes, then re-adds, causing it to break + this.outputs[0].type = "*"; + this.outputs[0].name = "connect to widget input"; + delete this.outputs[0].widget; + + if (this.widgets) { + // Allow widgets to cleanup + for (const w of this.widgets) { + if (w.onRemove) { + w.onRemove(); + } + } + this.widgets.length = 0; + } + } + } + + LiteGraph.registerNodeType( + "PrimitiveNode", + Object.assign(PrimitiveNode, { + title: "Primitive", + }) + ); + PrimitiveNode.category = "utils"; + }, +}); diff --git a/web/jsconfig.json b/web/jsconfig.json new file mode 100644 index 00000000..57403d8c --- /dev/null +++ b/web/jsconfig.json @@ -0,0 +1,9 @@ +{ + "compilerOptions": { + "baseUrl": ".", + "paths": { + "/*": ["./*"] + } + }, + "include": ["."] +} diff --git a/web/scripts/app.js b/web/scripts/app.js index fd410cd3..a33917b3 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -494,7 +494,7 @@ class ComfyApp { // Create and mount the LiteGraph in the DOM const canvasEl = (this.canvasEl = Object.assign(document.createElement("canvas"), { id: "graph-canvas" })); - canvasEl.tabIndex = "1" + canvasEl.tabIndex = "1"; document.body.prepend(canvasEl); this.graph = new LGraph(); @@ -525,7 +525,10 @@ class ComfyApp { this.loadGraphData(workflow); restored = true; } - } catch (err) {} + } catch (err) { + console.error("Error loading previous workflow", err); + debugger; + } // We failed to restore a workflow so load the default if (!restored) { @@ -572,12 +575,8 @@ class ComfyApp { const type = inputData[0]; if (Array.isArray(type)) { - // Enums e.g. latent rotation - let defaultValue = type[0]; - if (inputData[1] && inputData[1].default) { - defaultValue = inputData[1].default; - } - this.addWidget("combo", inputName, defaultValue, () => {}, { values: type }); + // Enums + Object.assign(config, widgets.COMBO(this, inputName, inputData, app) || {}); } else if (`${type}:${inputName}` in widgets) { // Support custom widgets by Type:Name Object.assign(config, widgets[`${type}:${inputName}`](this, inputName, inputData, app) || {}); @@ -667,11 +666,15 @@ class ComfyApp { async graphToPrompt() { const workflow = this.graph.serialize(); const output = {}; - for (const n of workflow.nodes) { - const node = this.graph.getNodeById(n.id); + // Process nodes in order of execution + for (const node of this.graph.computeExecutionOrder(false)) { + const n = workflow.nodes.find((n) => n.id === node.id); if (node.isVirtualNode) { - // Don't serialize frontend only nodes + // Don't serialize frontend only nodes but let them make changes + if (node.applyToGraph) { + node.applyToGraph(workflow); + } continue; } @@ -695,7 +698,11 @@ class ComfyApp { let link = node.getInputLink(i); while (parent && parent.isVirtualNode) { link = parent.getInputLink(link.origin_slot); - parent = parent.getInputNode(link.origin_slot); + if (link) { + parent = parent.getInputNode(link.origin_slot); + } else { + parent = null; + } } if (link) { diff --git a/web/scripts/widgets.js b/web/scripts/widgets.js index 30a02e72..8c78d955 100644 --- a/web/scripts/widgets.js +++ b/web/scripts/widgets.js @@ -10,9 +10,8 @@ function getNumberDefaults(inputData, defaultStep) { return { val: defaultVal, config: { min, max, step: 10.0 * step } }; } -function seedWidget(node, inputName, inputData) { - const seed = ComfyWidgets.INT(node, inputName, inputData); - const randomize = node.addWidget("toggle", "Random seed after every gen", true, function (v) {}, { +export function addRandomizeWidget(node, targetWidget, name, defaultValue = false) { + const randomize = node.addWidget("toggle", name, defaultValue, function (v) {}, { on: "enabled", off: "disabled", serialize: false, // Don't include this in prompt. @@ -20,14 +19,28 @@ function seedWidget(node, inputName, inputData) { randomize.afterQueued = () => { if (randomize.value) { - seed.widget.value = Math.floor(Math.random() * 1125899906842624); + const min = targetWidget.options?.min; + const max = targetWidget.options?.max; + if (min != null || max != null) { + targetWidget.value = Math.floor(Math.random() * ((max ?? 9999999999) - (min ?? 0) + 1) + (min ?? 0)); + } else { + targetWidget.value = Math.floor(Math.random() * 1125899906842624); + } } }; + return randomize; +} +function seedWidget(node, inputName, inputData) { + const seed = ComfyWidgets.INT(node, inputName, inputData); + const randomize = addRandomizeWidget(node, seed.widget, "Random seed after every gen", true); + + seed.widget.linkedWidgets = [randomize]; return { widget: seed, randomize }; } const MultilineSymbol = Symbol(); +const MultilineResizeSymbol = Symbol(); function addMultilineWidget(node, name, opts, app) { const MIN_SIZE = 50; @@ -95,7 +108,7 @@ function addMultilineWidget(node, name, opts, app) { // Calculate it here instead computeSize(node.size); } - const visible = app.canvas.ds.scale > 0.5; + const visible = app.canvas.ds.scale > 0.5 && this.type === "customtext"; const t = ctx.getTransform(); const margin = 10; Object.assign(this.inputEl.style, { @@ -149,9 +162,22 @@ function addMultilineWidget(node, name, opts, app) { } }; - if (!(MultilineSymbol in node)) { - node[MultilineSymbol] = true; - const onResize = node.onResize; + widget.onRemove = () => { + widget.inputEl?.remove(); + + // Restore original size handler if we are the last + if (!--node[MultilineSymbol]) { + node.onResize = node[MultilineResizeSymbol]; + delete node[MultilineSymbol]; + delete node[MultilineResizeSymbol]; + } + }; + + if (node[MultilineSymbol]) { + node[MultilineSymbol]++; + } else { + node[MultilineSymbol] = 1; + const onResize = (node[MultilineResizeSymbol] = node.onResize); node.onResize = function (size) { computeSize(size); @@ -199,6 +225,14 @@ export const ComfyWidgets = { return { widget: node.addWidget("text", inputName, defaultVal, () => {}, {}) }; } }, + COMBO(node, inputName, inputData) { + const type = inputData[0]; + let defaultValue = type[0]; + if (inputData[1] && inputData[1].default) { + defaultValue = inputData[1].default; + } + return { widget: node.addWidget("combo", inputName, defaultValue, () => {}, { values: type }) }; + }, IMAGEUPLOAD(node, inputName, inputData, app) { const imageWidget = node.widgets.find((w) => w.name === "image"); let uploadWidget;