111 lines
4.6 KiB
JavaScript
111 lines
4.6 KiB
JavaScript
import { app } from "/scripts/app.js";
|
||
import { ComfyWidgets } from "/scripts/widgets.js";
|
||
|
||
/* ------------------------------------------------------------------
|
||
Extension: Comfy.GenerateTextNode
|
||
Adds a read‑only multiline STRING widget to the "Generate Text (🔗LLMToolkit)"
|
||
node that displays the LLM response directly on the node. The node still
|
||
outputs the normal "context" so it can be connected to an external
|
||
Display Text node as before.
|
||
-------------------------------------------------------------------*/
|
||
|
||
app.registerExtension({
|
||
name: "Comfy.GenerateTextNode",
|
||
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
||
// Only act on our LLMToolkitTextGenerator node
|
||
if (nodeData.name !== "LLMToolkitTextGenerator") return;
|
||
|
||
/* -----------------------------
|
||
1. Inject widget on creation
|
||
----------------------------- */
|
||
const origOnNodeCreated = nodeType.prototype.onNodeCreated;
|
||
nodeType.prototype.onNodeCreated = function () {
|
||
// Invoke original handler first
|
||
const ret = origOnNodeCreated ? origOnNodeCreated.apply(this, arguments) : undefined;
|
||
|
||
// Build a unique widget name so each node instance has its own key
|
||
const existingNodes = app.graph._nodes.filter(n => n.type === nodeData.name);
|
||
const widgetName = `${nodeData.name}_response_${existingNodes.length}`;
|
||
|
||
// Create read‑only STRING widget for displaying response
|
||
const widgetInfo = ComfyWidgets.STRING(
|
||
this,
|
||
widgetName,
|
||
["STRING", {
|
||
default: "",
|
||
placeholder: "LLM response will appear here …",
|
||
multiline: true,
|
||
}],
|
||
app
|
||
);
|
||
|
||
// Make read‑only
|
||
widgetInfo.widget.inputEl.readOnly = true;
|
||
|
||
// Store reference for quick updates
|
||
this.__llmResponseWidget = widgetInfo.widget || widgetInfo;
|
||
|
||
// Ensure node size grows to fit another widget
|
||
this.setSize(this.computeSize());
|
||
|
||
return ret;
|
||
};
|
||
|
||
/* -----------------------------
|
||
2. Helper to update widget
|
||
----------------------------- */
|
||
function updateResponseWidget(textContainer) {
|
||
if (!this.__llmResponseWidget) return;
|
||
|
||
// Extract final text string (textContainer may be array/string)
|
||
let valueToSet = "";
|
||
if (Array.isArray(textContainer)) {
|
||
valueToSet = textContainer.join("\n");
|
||
} else if (typeof textContainer === "string") {
|
||
valueToSet = textContainer;
|
||
}
|
||
|
||
// Trim and write
|
||
valueToSet = valueToSet.trim();
|
||
this.__llmResponseWidget.value = valueToSet;
|
||
if (this.__llmResponseWidget.inputEl) this.__llmResponseWidget.inputEl.value = valueToSet;
|
||
app.graph.setDirtyCanvas(true);
|
||
}
|
||
|
||
/* -----------------------------
|
||
3. Hook into onExecuted to receive message
|
||
----------------------------- */
|
||
const origOnExecuted = nodeType.prototype.onExecuted;
|
||
nodeType.prototype.onExecuted = function (message) {
|
||
if (origOnExecuted) origOnExecuted.apply(this, arguments);
|
||
|
||
// message is typically { ui: { string: [...] }, result: (…) }
|
||
let textData = undefined;
|
||
if (message?.ui?.string) {
|
||
textData = message.ui.string;
|
||
} else if (message?.string) {
|
||
textData = message.string;
|
||
} else if (Array.isArray(message)) {
|
||
textData = message;
|
||
}
|
||
updateResponseWidget.call(this, textData);
|
||
};
|
||
|
||
/* -----------------------------
|
||
4. Restore saved value on load
|
||
----------------------------- */
|
||
const origOnConfigure = nodeType.prototype.onConfigure;
|
||
nodeType.prototype.onConfigure = function (config) {
|
||
if (origOnConfigure) origOnConfigure.apply(this, arguments);
|
||
|
||
if (!this.__llmResponseWidget || !config?.widgets_values?.length) return;
|
||
|
||
const idx = this.widgets.findIndex(w => w === this.__llmResponseWidget);
|
||
if (idx !== -1 && config.widgets_values.length > idx) {
|
||
const savedVal = config.widgets_values[idx] || "";
|
||
this.__llmResponseWidget.value = savedVal;
|
||
if (this.__llmResponseWidget.inputEl) this.__llmResponseWidget.inputEl.value = savedVal;
|
||
}
|
||
};
|
||
},
|
||
});
|