* 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.
This commit is contained in:
Lt.Dr.Data 2023-06-05 18:12:52 +09:00
parent 2ec980bb9f
commit 82af034827
4 changed files with 71 additions and 17 deletions

View File

@ -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;

View File

@ -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);
}
}
}
}
}

View File

@ -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,

35
web/scripts/util.js Normal file
View File

@ -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;
}
};