130 lines
4.9 KiB
TypeScript
130 lines
4.9 KiB
TypeScript
import { updateProto } from "./util.js";
|
|
import type { LGraphNode, IWidget, LLink, INodeOutputSlot, IContextMenuItem, LGraphCanvas } from "@comfyorg/litegraph";
|
|
import type { ComfyExtension } from "@comfyorg/comfyui-frontend-types";
|
|
|
|
interface PromptNode extends LGraphNode {
|
|
updateOptions: () => void;
|
|
}
|
|
|
|
const promptProto: any = {
|
|
onWidgetChanged(this: PromptNode, name: string, value: any, old_Value: any, widget: IWidget) {
|
|
if (name === "insert_textual_inversion") {
|
|
const keyword = value?.value?.keyword;
|
|
if (!keyword) return;
|
|
const tag = `<${keyword}>`;
|
|
|
|
const textWidget = this.widgets?.find((w) => w.name === "prompt");
|
|
const text = textWidget?.value ?? "";
|
|
|
|
if (typeof text === 'string' && textWidget) {
|
|
if (text.includes(tag)) {
|
|
textWidget.value = text.split(tag).join("");
|
|
} else {
|
|
textWidget.value = `<${keyword}> ${text}`;
|
|
}
|
|
}
|
|
|
|
widget.value = "...";
|
|
}
|
|
},
|
|
|
|
onConnectionsChange(this: PromptNode, type: number, index: number, isConnected: boolean, link_info: LLink, inputOrOutput: INodeOutputSlot) {
|
|
if (app.extensionManager.setting.get("drawthings.node.color_prompts") === false) return;
|
|
|
|
let isPositive = false;
|
|
let isNegative = false;
|
|
|
|
for (const linkId of this.outputs?.[0]?.links ?? []) {
|
|
const link = this.graph?.getLink(linkId);
|
|
if (!link) continue;
|
|
const targetId = link.target_id;
|
|
const targetNode = this.graph?.getNodeById(targetId);
|
|
if ((targetNode as any)?.comfyClass === "DrawThingsSampler") {
|
|
const input = targetNode?.inputs?.[link.target_slot];
|
|
if (input?.name === "positive") isPositive = true;
|
|
if (input?.name === "negative") isNegative = true;
|
|
}
|
|
}
|
|
|
|
if (isPositive && isNegative) {
|
|
this.color = window.LGraphCanvas.node_colors.purple.color;
|
|
this.bgcolor = window.LGraphCanvas.node_colors.purple.bgcolor;
|
|
} else if (isPositive) {
|
|
this.color = window.LGraphCanvas.node_colors.green.color;
|
|
this.bgcolor = window.LGraphCanvas.node_colors.green.bgcolor;
|
|
} else if (isNegative) {
|
|
this.color = window.LGraphCanvas.node_colors.red.color;
|
|
this.bgcolor = window.LGraphCanvas.node_colors.red.bgcolor;
|
|
} else {
|
|
this.color = undefined;
|
|
this.bgcolor = undefined;
|
|
}
|
|
},
|
|
|
|
onNodeCreated(this: PromptNode) {
|
|
const output = this.outputs?.find((output) => output.name == "PROMPT");
|
|
if (output) {
|
|
(output as any).color_on = (output as any).color_off = app.canvas.default_connection_color_byType["CONDITIONING"];
|
|
}
|
|
|
|
const promptWidget = this.widgets?.find((w) => w.name === "prompt");
|
|
const promptNode = this;
|
|
if ((promptWidget as any)?.element) {
|
|
(promptWidget as any).element.addEventListener("change", () => {
|
|
promptNode.updateOptions();
|
|
});
|
|
}
|
|
},
|
|
|
|
getExtraMenuOptions(this: PromptNode, canvas: LGraphCanvas, options: IContextMenuItem[]) {
|
|
const promptColors = app.extensionManager.setting.get("drawthings.node.color_prompts")
|
|
options.push(
|
|
...[
|
|
null,
|
|
{
|
|
content: (promptColors ? "✓ " : "") + "Change colors when connections change",
|
|
callback: async () => {
|
|
try {
|
|
await app.extensionManager.setting.set("drawthings.node.color_prompts", !promptColors);
|
|
} catch (error) {
|
|
console.error(`Error changing setting: ${error}`);
|
|
}
|
|
},
|
|
},
|
|
null,
|
|
] as any
|
|
);
|
|
},
|
|
};
|
|
|
|
const extension: ComfyExtension = {
|
|
name: "promptNode",
|
|
|
|
beforeRegisterNodeDef(nodeType, nodeData, app) {
|
|
if ((nodeType as any).comfyClass === "DrawThingsPrompt") {
|
|
updateProto(nodeType, promptProto);
|
|
}
|
|
},
|
|
|
|
settings: [
|
|
{
|
|
id: "drawthings.node.color_prompts",
|
|
type: "boolean",
|
|
name: "Change prompt node colors when connections change",
|
|
defaultValue: true,
|
|
category: ["DrawThings", "Nodes", "Change prompt"],
|
|
onChange: (newVal: boolean, oldVal: boolean) => {
|
|
if (oldVal === false && newVal === true) {
|
|
app.graph.nodes
|
|
.filter((n) => n.type === "DrawThingsPrompt")
|
|
.forEach((n) => {
|
|
setTimeout(() => (n as any).onConnectionsChange(), 10)
|
|
})
|
|
}
|
|
},
|
|
},
|
|
],
|
|
};
|
|
|
|
export default extension;
|