From 82af034827cf4e002fbae8e2516c54d9110cc6d4 Mon Sep 17 00:00:00 2001 From: "Lt.Dr.Data" Date: Mon, 5 Jun 2023 18:12:52 +0900 Subject: [PATCH] * Support of Combo refresh in PrimitiveNode * Security filter of workflow - if there are image files or model files in the config, they will be converted into an empty array. --- web/extensions/core/widgetInputs.js | 10 ++++---- web/scripts/app.js | 36 +++++++++++++++++++++-------- web/scripts/ui.js | 7 ++++-- web/scripts/util.js | 35 ++++++++++++++++++++++++++++ 4 files changed, 71 insertions(+), 17 deletions(-) create mode 100644 web/scripts/util.js diff --git a/web/extensions/core/widgetInputs.js b/web/extensions/core/widgetInputs.js index 4fe0a6013..c4e71ac6b 100644 --- a/web/extensions/core/widgetInputs.js +++ b/web/extensions/core/widgetInputs.js @@ -51,7 +51,7 @@ function showWidget(widget) { function convertToInput(node, widget, config) { hideWidget(node, widget); - const { linkType } = getWidgetType(config); + const { linkType } = getWidgetType(`${node.type}\n${widget.name}`, config); // Add input and store widget config for creating on primitive node const sz = node.size; @@ -72,13 +72,13 @@ function convertToWidget(node, widget) { node.setSize([Math.max(sz[0], node.size[0]), Math.max(sz[1], node.size[1])]); } -function getWidgetType(config) { +function getWidgetType(name, 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(","); + linkType = `${name}`; } return { type, linkType }; } @@ -86,7 +86,7 @@ function getWidgetType(config) { app.registerExtension({ name: "Comfy.WidgetInputs", async beforeRegisterNodeDef(nodeType, nodeData, app) { - // Add menu options to conver to/from widgets + // Add menu options to convert to/from widgets const origGetExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; nodeType.prototype.getExtraMenuOptions = function (_, options) { const r = origGetExtraMenuOptions ? origGetExtraMenuOptions.apply(this, arguments) : undefined; @@ -270,7 +270,7 @@ app.registerExtension({ } const widget = _widget; - const { type, linkType } = getWidgetType(widget.config); + const { type, linkType } = getWidgetType(`${theirNode.type}\n${widget.name}`, widget.config); // Update our output to restrict to the widget type this.outputs[0].type = linkType; this.outputs[0].name = type; diff --git a/web/scripts/app.js b/web/scripts/app.js index 95447ffa0..7f281f709 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -3,6 +3,7 @@ import { ComfyUI, $el } from "./ui.js"; import { api } from "./api.js"; import { defaultGraph } from "./defaultGraph.js"; import { getPngMetadata, importA1111, getLatentMetadata } from "./pnginfo.js"; +import { Util } from "./util.js"; /** * @typedef {import("types/comfy").ComfyExtension} ComfyExtension @@ -1355,6 +1356,7 @@ export class ComfyApp { const p = await this.graphToPrompt(); try { + p.workflow = Util.workflow_security_filter(p.workflow); await api.queuePrompt(number, p); } catch (error) { const formattedError = this.#formatPromptError(error) @@ -1441,20 +1443,34 @@ export class ComfyApp { const def = defs[node.type]; - // HOTFIX: The current patch is designed to prevent the rest of the code from breaking due to primitive nodes, - // and additional work is needed to consider the primitive logic in the refresh logic. - if(!def) - continue; - for(const widgetNum in node.widgets) { const widget = node.widgets[widgetNum] - if(widget.type == "combo" && def["input"]["required"][widget.name] !== undefined) { - widget.options.values = def["input"]["required"][widget.name][0]; + if(widget.type == "combo") { + if(def) { + // normal node + if(def["input"]["required"][widget.name]) { + widget.options.values = def["input"]["required"][widget.name][0]; - if(widget.name != 'image' && !widget.options.values.includes(widget.value)) { - widget.value = widget.options.values[0]; - widget.callback(widget.value); + if(widget.name != 'image' && !widget.options.values.includes(widget.value)) { + widget.value = widget.options.values[0]; + widget.callback(widget.value); + } + } } + else if(widget.name != "control_after_generate") { + // primitive node + let info = node.outputs[0].type.split("\n"); + let node_type = info[0]; + let slot_name = info[1]; + + widget.options.values = defs[node_type].input.required[slot_name][0]; + + if(widget.name != 'image' && !widget.options.values.includes(widget.value)) { + widget.value = widget.options.values[0]; + widget.callback(widget.value); + } + } + } } } diff --git a/web/scripts/ui.js b/web/scripts/ui.js index a26eedec3..1004357ff 100644 --- a/web/scripts/ui.js +++ b/web/scripts/ui.js @@ -1,4 +1,5 @@ import { api } from "./api.js"; +import { Util } from "./util.js"; export function $el(tag, propsOrChildren, children) { const split = tag.split("."); @@ -581,8 +582,10 @@ export class ComfyUI { filename += ".json"; } } - const json = JSON.stringify(app.graph.serialize(), null, 2); // convert the data to a JSON string - const blob = new Blob([json], { type: "application/json" }); + var json = app.graph.serialize(); + json = Util.workflow_security_filter(json); + const json_str = JSON.stringify(json, null, 2); // convert the data to a JSON string + const blob = new Blob([json_str], { type: "application/json" }); const url = URL.createObjectURL(blob); const a = $el("a", { href: url, diff --git a/web/scripts/util.js b/web/scripts/util.js new file mode 100644 index 000000000..ff4963da5 --- /dev/null +++ b/web/scripts/util.js @@ -0,0 +1,35 @@ +// If there are files with specific extensions in the array inside the config of nodes/inputs/widget in the workflow, empty that array. +const excludeExtensions = new Set(["png", "jpg", "webp", "jpeg", "safetensors", "ckpt", "pt", "pth"]); + +function getFileExtension(filename) { + return filename.slice((filename.lastIndexOf('.') - 1 >>> 0) + 2); +} + +export class Util { + static workflow_security_filter(workflow) { + workflow.nodes.forEach((node) => { + if (node.inputs) { + node.inputs.forEach((input) => { + if (input.widget && input.widget.config) { + const configArray = input.widget.config[0]; + if (Array.isArray(configArray) && configArray.every((filename) => excludeExtensions.has(getFileExtension(filename)))) { + input.widget.config[0] = []; + } + } + }); + } + if (node.outputs) { + node.outputs.forEach((output) => { + if (output.widget && output.widget.config) { + const configArray = output.widget.config[0]; + if (Array.isArray(configArray) && configArray.every((filename) => excludeExtensions.has(getFileExtension(filename)))) { + output.widget.config[0] = []; + } + } + }); + } + }); + + return workflow; + } +};