Show prompts in UI

README.md

@@ -1,6 +1,5 @@
 # ExLlama nodes for ComfyUI
-A simple prompt generator for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) utilizing [ExLlama](https://github.com/turboderp/exllama).  
-Outputs are printed in the console, if you'd like to see them in the UI you can use the `Show Text` node from [here](https://github.com/pythongosssss/ComfyUI-Custom-Scripts).
+A simple prompt generator for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) utilizing [ExLlama](https://github.com/turboderp/exllama).
 
 ## Installation
 Clone the repository to `custom_nodes` in your ComfyUI directory and install dependencies:

__init__.py

@@ -9,3 +9,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ZuellniExLlamaGenerator": "ExLlama Generator",
     "ZuellniExLlamaLoader": "ExLlama Loader",
 }
+
+WEB_DIRECTORY = "."

nodes.py

@@ -32,6 +32,11 @@ class Generator:
     RETURN_TYPES = ("STRING",)
 
     def generate(self, model, max_tokens, temperature, top_k, top_p, typical_p, penalty, seed, prompt):
+        prompt = prompt.strip()
+
+        if not prompt:
+            return ("",)
+
         settings = ExLlamaAltGenerator.Settings()
         settings.temperature = temperature
         settings.top_k = top_k
@@ -41,7 +46,7 @@ class Generator:
 
         torch.manual_seed(seed)
         stop_conditions = [model.tokenizer.eos_token_id, model.tokenizer.newline_token_id]
-        model.begin_stream(prompt.strip(), stop_conditions, max_tokens, settings)
+        model.begin_stream(prompt, stop_conditions, max_tokens, settings)
 
         progress = ProgressBar(max_tokens)
         eos = False
@@ -49,12 +54,11 @@ class Generator:
 
         while not eos:
             chunk, eos = model.stream()
-            text += chunk
             progress.update(1)
+            text += chunk
 
         text = text.strip()
-        print("\n[\033[94mExLlama\033[0m]: " + text, end="\n\n")
-        return (text,)
+        return {"ui": {"text": text}, "result": (text,)}
 
 
 class Loader:

scripts.js

@@ -0,0 +1,49 @@
+import { app } from "../../../scripts/app.js";
+import { ComfyWidgets } from "../../../scripts/widgets.js";
+
+app.registerExtension({
+	name: "Zuellni.ExLlama.Generator",
+	async beforeRegisterNodeDef(nodeType, nodeData, app) {
+		if (nodeData.name === "ZuellniExLlamaGenerator") {
+			const onExecuted = nodeType.prototype.onExecuted;
+
+			nodeType.prototype.onExecuted = function (message) {
+				onExecuted?.apply(this, arguments);
+
+				if (this.widgets) {
+					const pos = this.widgets.findIndex((w) => w.name === "text");
+
+					if (pos !== -1) {
+						for (let i = pos; i < this.widgets.length; i++) {
+							this.widgets[i].onRemove?.();
+						}
+
+						this.widgets.length = pos;
+					}
+
+					const string = ["STRING", { multiline: true }]
+					const widget = ComfyWidgets["STRING"](this, "text", string, app).widget;
+
+					widget.inputEl.readOnly = true;
+					widget.inputEl.style.opacity = 0.7;
+					widget.value = message.text.join("");
+
+					requestAnimationFrame(() => {
+						const size = this.computeSize();
+
+						if (size[0] < this.size[0]) {
+							size[0] = this.size[0];
+						}
+
+						if (size[1] < this.size[1]) {
+							size[1] = this.size[1];
+						}
+
+						this.onResize?.(size);
+						app.graph.setDirtyCanvas(true, false);
+					});
+				}
+			};
+		}
+	},
+});