From 0ec46fb285e4826b82372e7a469771673db12575 Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Sat, 16 Sep 2023 22:30:29 +0200
Subject: [PATCH] Add a separate preview node to prevent generator from running
 twice, unique_id is buggy

---
 README.md   | 11 +++++++----
 __init__.py |  8 +++++---
 nodes.py    | 17 +++++++++++++++--
 scripts.js  |  4 ++--
 4 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index cea5462..9d0f3ad 100644
--- a/README.md
+++ b/README.md
@@ -19,18 +19,21 @@ pip install https://github.com/jllllll/exllama/releases/download/0.0.17/exllama-
 ## Nodes
 Comes with the following nodes:
 
-### ExLlama Loader
+### Loader
 Used to load 4-bit GPTQ Llama/2 models. You can find a lot of them over at [Hugging Face](https://huggingface.co/TheBloke).  
 ExLlama allocates [memory](https://github.com/turboderp/exllama/issues/259) according to `max_seq_len`. Lowering it is a good way to save on GPU RAM.  
 It's currently not possible to [offload](https://github.com/turboderp/exllama/issues/177) the model to CPU RAM.
 
-### ExLlama Generator
+### Generator
 Generates a `string` based on the given `prompt` for use with other nodes.  
 Default values correspond to the `simple-1` preset from [text-generation-webui](https://github.com/oobabooga/text-generation-webui).  
 ExLlama isn't [deterministic](https://github.com/turboderp/exllama/issues/201), so the outputs may differ even with the same seed.
 
-## Example
-The workflow can be opened directly in ComfyUI.  
+### Previewer
+Displays generated prompts in the UI.
+
+## Workflow
+Can be opened directly in ComfyUI.  
 Model used: [MythoLogic-Mini-7B](https://huggingface.co/TheBloke/MythoLogic-Mini-7B-GPTQ).
 
 ![workflow](https://github.com/Zuellni/ComfyUI-ExLlama-Nodes/assets/123005779/2a8c71b2-3014-49b8-9252-29a30eba01c9)
diff --git a/__init__.py b/__init__.py
index 6fd7ec2..eb7005d 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,13 +1,15 @@
-from .nodes import Generator, Loader
+from .nodes import Generator, Loader, Previewer
 
 NODE_CLASS_MAPPINGS = {
-    "ZuellniExLlamaGenerator": Generator,
     "ZuellniExLlamaLoader": Loader,
+    "ZuellniExLlamaGenerator": Generator,
+    "ZuellniExLlamaPreviewer": Previewer,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "ZuellniExLlamaGenerator": "ExLlama Generator",
     "ZuellniExLlamaLoader": "ExLlama Loader",
+    "ZuellniExLlamaGenerator": "ExLlama Generator",
+    "ZuellniExLlamaPreviewer": "ExLlama Previewer",
 }
 
 WEB_DIRECTORY = "."
diff --git a/nodes.py b/nodes.py
index 7fa49d2..6c4a682 100644
--- a/nodes.py
+++ b/nodes.py
@@ -27,7 +27,6 @@ class Generator:
 
     CATEGORY = "Zuellni/ExLlama"
     FUNCTION = "generate"
-    OUTPUT_NODE = True
     RETURN_NAMES = ("TEXT",)
     RETURN_TYPES = ("STRING",)
 
@@ -58,7 +57,7 @@ class Generator:
             text += chunk
 
         text = text.strip()
-        return {"ui": {"text": [text]}, "result": (text,)}
+        return (text,)
 
 
 class Loader:
@@ -91,3 +90,17 @@ class Loader:
         tokenizer = ExLlamaTokenizer(str(model_dir / "tokenizer.model"))
         generator = ExLlamaAltGenerator(model, tokenizer, cache)
         return (generator,)
+
+
+class Previewer:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"text": ("STRING", {"forceInput": True})}}
+
+    CATEGORY = "Zuellni/ExLlama"
+    FUNCTION = "preview"
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+
+    def preview(self, text):
+        return {"ui": {"text": [text]}}
diff --git a/scripts.js b/scripts.js
index 59fd4e9..44c9f3c 100644
--- a/scripts.js
+++ b/scripts.js
@@ -4,9 +4,9 @@ import { app } from "../../../scripts/app.js";
 import { ComfyWidgets } from "../../../scripts/widgets.js";
 
 app.registerExtension({
-	name: "Zuellni.ExLlama.Generator",
+	name: "ZuellniExLlamaPreviewer",
 	async beforeRegisterNodeDef(nodeType, nodeData, app) {
-		if (nodeData.name === "ZuellniExLlamaGenerator") {
+		if (nodeData.name === "ZuellniExLlamaPreviewer") {
 			const onExecuted = nodeType.prototype.onExecuted;
 
 			nodeType.prototype.onExecuted = function (message) {