From 4ecef087742233476af415fa7a680f6433385c6b Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sun, 27 Aug 2023 21:24:36 +0100
Subject: [PATCH] Allow custom words

---
 py/autocomplete.py            | 19 +++++++++++++++++++
 user/autocomplete.txt         |  0
 web/js/autocompleter.js       | 19 ++++++++++++++++++-
 web/js/common/autocomplete.js |  4 ----
 4 files changed, 37 insertions(+), 5 deletions(-)
 create mode 100644 py/autocomplete.py
 create mode 100644 user/autocomplete.txt

diff --git a/py/autocomplete.py b/py/autocomplete.py
new file mode 100644
index 0000000..a456198
--- /dev/null
+++ b/py/autocomplete.py
@@ -0,0 +1,19 @@
+from server import PromptServer
+from aiohttp import web
+import os 
+
+dir = os.path.abspath(os.path.join(__file__, "../../user"))
+if not os.path.exists(dir):
+    os.mkdir(dir)
+file = os.path.join(dir, "autocomplete.txt")
+
+@PromptServer.instance.routes.get("/pysssss/autocomplete")
+async def get_autocomplete(request):
+    if os.path.isfile(file):
+        return web.FileResponse(file)
+    return web.Response(status=404)
+
+@PromptServer.instance.routes.post("/pysssss/autocomplete")
+async def update_autocomplete(request):
+    with open(file, "w") as f:
+        f.write(request.text())
\ No newline at end of file
diff --git a/user/autocomplete.txt b/user/autocomplete.txt
new file mode 100644
index 0000000..e69de29
diff --git a/web/js/autocompleter.js b/web/js/autocompleter.js
index 97265f5..166c0f5 100644
--- a/web/js/autocompleter.js
+++ b/web/js/autocompleter.js
@@ -41,7 +41,24 @@ app.registerExtension({
 			TextAreaAutoComplete.updateWords("pysssss.embeddings", words);
 		}
 
-		addEmbeddings();
+		async function addCustomWords() {
+			const resp = await api.fetchApi("/pysssss/autocomplete", { cache: "no-store" });
+			if (resp.status === 200) {
+				const text = await resp.text();
+				if (text) {
+					TextAreaAutoComplete.updateWords(
+						"pysssss.customwords",
+						text.split("\n").reduce((p, n) => {
+							n = n.trim();
+							p[n] = { text: n };
+							return p;
+						}, {})
+					);
+				}
+			}
+		}
+
+		Promise.all([addEmbeddings(), addCustomWords()]);
 
 		const STRING = ComfyWidgets.STRING;
 		ComfyWidgets.STRING = function (node, inputName, inputData) {
diff --git a/web/js/common/autocomplete.js b/web/js/common/autocomplete.js
index 95b4a06..585a6fd 100644
--- a/web/js/common/autocomplete.js
+++ b/web/js/common/autocomplete.js
@@ -341,23 +341,19 @@ export class TextAreaAutoComplete {
 			switch (e.key) {
 				case "ArrowUp":
 					e.preventDefault();
-					this.selected.el.classList.remove("pysssss-autocomplete-item--selected");
 					if (this.selected.index) {
 						this.#setSelected(this.currentWords[this.selected.index - 1]);
 					} else {
 						this.#setSelected(this.currentWords[this.currentWords.length - 1]);
 					}
-					this.selected.el.classList.add("pysssss-autocomplete-item--selected");
 					break;
 				case "ArrowDown":
 					e.preventDefault();
-					this.selected.el.classList.remove("pysssss-autocomplete-item--selected");
 					if (this.selected.index === this.currentWords.length - 1) {
 						this.#setSelected(this.currentWords[0]);
 					} else {
 						this.#setSelected(this.currentWords[this.selected.index + 1]);
 					}
-					this.selected.el.classList.add("pysssss-autocomplete-item--selected");
 					break;
 				case "Tab":
 					e.preventDefault();