From 9f9d25dd497389940bc7fb3be352e077c29ada47 Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Sun, 4 Aug 2024 12:57:22 +0200
Subject: [PATCH] Rename the convert node, some fixes and cleanup

---
 README.md  |  6 +++---
 exllama.py |  8 ++++----
 text.js    |  8 ++++----
 text.py    | 28 ++++++++++++++++------------
 4 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/README.md b/README.md
index 2ebf834..872b613 100644
--- a/README.md
+++ b/README.md
@@ -28,7 +28,7 @@ git clone https://huggingface.co/turboderp/Llama-3.1-8B-Instruct-exl2 -b 4.0bpw
 > You can add your own `llm` path to the [extra_model_paths.yaml](https://github.com/comfyanonymous/ComfyUI/blob/master/extra_model_paths.yaml.example) file and put the models there instead.
 
 ## Nodes
-<table>
+<table width="100%">
   <tr>
     <td colspan="3" align="center"><b>ExLlama Nodes</b></td>
   </tr>
@@ -106,8 +106,8 @@ git clone https://huggingface.co/turboderp/Llama-3.1-8B-Instruct-exl2 -b 4.0bpw
     <td colspan="3" align="center"><b>Text Nodes</b></td>
   </tr>
   <tr>
-    <td><b>Convert</b></td>
-    <td colspan="2">Strips punctuation, whitespace, and changes case for input.</td>
+    <td><b>Clean</b></td>
+    <td colspan="2">Strips punctuation, fixes whitespace, and changes case for input text.</td>
   </tr>
   <tr>
     <td><b>Message</b></td>
diff --git a/exllama.py b/exllama.py
index 6163cf8..b5d713f 100644
--- a/exllama.py
+++ b/exllama.py
@@ -50,7 +50,7 @@ class Loader:
                 "fast_tensors": ("BOOLEAN", {"default": True}),
                 "flash_attention": ("BOOLEAN", {"default": True}),
                 "max_seq_len": ("INT", {"default": 2048, "max": 2**20, "step": 256}),
-            },
+            }
         }
 
     _CACHES = {
@@ -127,7 +127,7 @@ class Formatter:
                 "model": ("EXL_MODEL",),
                 "messages": ("EXL_MESSAGES",),
                 "add_assistant_role": ("BOOLEAN", {"default": True}),
-            },
+            }
         }
 
     CATEGORY = _CATEGORY
@@ -182,7 +182,7 @@ class Tokenizer:
                 "text": ("STRING", {"forceInput": True}),
                 "add_bos_token": ("BOOLEAN", {"default": True}),
                 "encode_special_tokens": ("BOOLEAN", {"default": True}),
-            },
+            }
         }
 
     CATEGORY = _CATEGORY
@@ -214,7 +214,7 @@ class Settings:
                 "tfs": ("FLOAT", {"max": 1, "step": 0.01}),
                 "typical": ("FLOAT", {"max": 1, "step": 0.01}),
                 "temperature_last": ("BOOLEAN", {"default": True}),
-            },
+            }
         }
 
     CATEGORY = _CATEGORY
diff --git a/text.js b/text.js
index 75919a5..879e6de 100644
--- a/text.js
+++ b/text.js
@@ -10,7 +10,7 @@ app.registerExtension({
 		const onExecuted = nodeType.prototype.onExecuted
 
 		if (nodeData.name == "ZuellniTextPreview") {
-			nodeType.prototype.onNodeCreated = function() {
+			nodeType.prototype.onNodeCreated = function () {
 				const output = this.widgets.find(w => w.name == "output")
 
 				if (output) {
@@ -24,13 +24,13 @@ app.registerExtension({
 				return onNodeCreated?.apply(this, arguments)
 			}
 
-			nodeType.prototype.onExecuted = function(message) {
+			nodeType.prototype.onExecuted = function (message) {
 				const output = this.widgets.find(w => w.name == "output")
 				output && (output.value = message.text)
 				return onExecuted?.apply(this, arguments)
 			}
 		} else if (nodeData.name == "ZuellniTextReplace") {
-			nodeType.prototype.onNodeCreated = function() {
+			nodeType.prototype.onNodeCreated = function () {
 				const count = this.widgets.find(w => w.name == "count")
 
 				if (count) {
@@ -41,7 +41,7 @@ app.registerExtension({
 				return onNodeCreated?.apply(this, arguments)
 			}
 
-			nodeType.prototype.onChanged = function(count) {
+			nodeType.prototype.onChanged = function (count) {
 				!this.inputs && (this.inputs = [])
 				const current = this.inputs.length
 
diff --git a/text.py b/text.py
index e47400f..3d9e2d7 100644
--- a/text.py
+++ b/text.py
@@ -4,26 +4,26 @@ _CATEGORY = "Zuellni/Text"
 _MAPPING = "ZuellniText"
 
 
-class Convert:
+class Clean:
     @classmethod
     def INPUT_TYPES(cls):
         return {
             "required": {
                 "text": ("STRING", {"forceInput": True}),
-                "strip": (("punctuation", "whitespace", "both", "none"),),
+                "strip": (("both", "punctuation", "whitespace", "none"),),
                 "case": (("lower", "upper", "capitalize", "title", "none"),),
-            },
+                "fix": ("BOOLEAN", {"default": True}),
+            }
         }
 
     CATEGORY = _CATEGORY
-    FUNCTION = "convert"
+    FUNCTION = "clean"
     RETURN_NAMES = ("TEXT",)
     RETURN_TYPES = ("STRING",)
 
-    def convert(self, text, strip, case):
+    def clean(self, text, strip, case, fix):
         if strip == "both":
             text = text.strip(string.punctuation + string.whitespace)
-            text = " ".join(text.split()).strip()
         elif strip != "none":
             text = text.strip(getattr(string, strip))
 
@@ -32,6 +32,10 @@ class Convert:
         elif case != "none":
             text = getattr(text, case)()
 
+        if fix:
+            text = "\n".join([t for t in text.splitlines() if t])
+            text = " ".join(text.split())
+
         return (text,)
 
 
@@ -47,11 +51,11 @@ class Message:
         }
 
     CATEGORY = _CATEGORY
-    FUNCTION = "append"
+    FUNCTION = "add"
     RETURN_NAMES = ("MESSAGES",)
     RETURN_TYPES = ("EXL_MESSAGES",)
 
-    def append(self, role, content, messages=[]):
+    def add(self, role, content, messages=[]):
         return (messages + [{"role": role, "content": content}],)
 
 
@@ -62,7 +66,7 @@ class Preview:
             "required": {
                 "text": ("STRING", {"forceInput": True}),
                 "output": ("STRING", {"multiline": True}),
-            },
+            }
         }
 
     CATEGORY = _CATEGORY
@@ -81,7 +85,7 @@ class Replace:
             "required": {
                 "count": ("INT", {"default": 1, "min": 1, "max": 26}),
                 "text": ("STRING", {"multiline": True}),
-            },
+            }
         }
 
     CATEGORY = _CATEGORY
@@ -114,7 +118,7 @@ class String:
 
 
 NODE_CLASS_MAPPINGS = {
-    f"{_MAPPING}Convert": Convert,
+    f"{_MAPPING}Clean": Clean,
     f"{_MAPPING}Message": Message,
     f"{_MAPPING}Preview": Preview,
     f"{_MAPPING}Replace": Replace,
@@ -122,7 +126,7 @@ NODE_CLASS_MAPPINGS = {
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    f"{_MAPPING}Convert": "Convert",
+    f"{_MAPPING}Clean": "Clean",
     f"{_MAPPING}Message": "Message",
     f"{_MAPPING}Preview": "Preview",
     f"{_MAPPING}Replace": "Replace",