From 75892eb461b7e6968a98d092f791dca820e3da57 Mon Sep 17 00:00:00 2001
From: qrtal <90406181+qrtal@users.noreply.github.com>
Date: Mon, 8 Jan 2024 01:27:02 -0500
Subject: [PATCH] added autocomplete support for loras

---
 py/autocomplete.py      |  7 +++++++
 py/model_info.py        |  4 ++--
 web/js/autocompleter.js | 37 ++++++++++++++++++++++++++++++++++++-
 3 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/py/autocomplete.py b/py/autocomplete.py
index abf838d..d3a879e 100644
--- a/py/autocomplete.py
+++ b/py/autocomplete.py
@@ -1,6 +1,7 @@
 from server import PromptServer
 from aiohttp import web
 import os
+import folder_paths
 
 dir = os.path.abspath(os.path.join(__file__, "../../user"))
 if not os.path.exists(dir):
@@ -20,3 +21,9 @@ async def update_autocomplete(request):
     with open(file, "w", encoding="utf-8") as f:
         f.write(await request.text())
     return web.Response(status=200)
+
+
+@PromptServer.instance.routes.get("/pysssss/loras")
+async def get_loras(request):
+    loras = folder_paths.get_filename_list("loras")
+    return web.json_response(list(map(lambda a: os.path.splitext(a)[0], loras)))
\ No newline at end of file
diff --git a/py/model_info.py b/py/model_info.py
index b9b4daf..b360baf 100644
--- a/py/model_info.py
+++ b/py/model_info.py
@@ -31,7 +31,7 @@ async def save_notes(request):
     name = name[pos+1:]
 
     file_path = None
-    if type == "embeddings":
+    if type == "embeddings" or type == "loras":
         name = name.lower()
         files = folder_paths.get_filename_list(type)
         for f in files:
@@ -67,7 +67,7 @@ async def load_metadata(request):
     name = name[pos+1:]
 
     file_path = None
-    if type == "embeddings":
+    if type == "embeddings" or type == "loras":
         name = name.lower()
         files = folder_paths.get_filename_list(type)
         for f in files:
diff --git a/web/js/autocompleter.js b/web/js/autocompleter.js
index 99af50f..47ab948 100644
--- a/web/js/autocompleter.js
+++ b/web/js/autocompleter.js
@@ -143,6 +143,23 @@ class EmbeddingInfoDialog extends ModelInfoDialog {
 	}
 }
 
+class LoraInfoDialog extends ModelInfoDialog {
+	async addInfo() {
+		super.addInfo();
+		const info = await this.addCivitaiInfo();
+		if (info) {
+			$el("div", {
+				parent: this.content,
+				innerHTML: info.description,
+				style: {
+					maxHeight: "250px",
+					overflow: "auto",
+				},
+			});
+		}
+	}
+}
+
 class CustomWordsDialog extends ComfyDialog {
 	async show() {
 		const text = await getCustomWords();
@@ -268,7 +285,25 @@ app.registerExtension({
 			TextAreaAutoComplete.updateWords("pysssss.embeddings", words);
 		}
 
-		Promise.all([addEmbeddings(), addCustomWords()]);
+		async function addLoras() {
+			const loras = await api
+				.fetchApi("/pysssss/loras", { cache: "no-store" })
+				.then(res => res.json());
+			const words = {};
+			words["lora:"] = { text: "lora:" };
+
+			for (const lora of loras) {
+				const v = `<lora:${lora}:1.0>`;
+				words[v] = {
+					text: v,
+					info: () => new LoraInfoDialog(lora).show("loras", lora),
+				};
+			}
+
+			TextAreaAutoComplete.updateWords("pysssss.loras", words);
+		}
+
+		Promise.all([addEmbeddings(), addLoras(), addCustomWords()]);
 
 		const STRING = ComfyWidgets.STRING;
 		const SKIP_WIDGETS = new Set(["ttN xyPlot.x_values", "ttN xyPlot.y_values"]);