diff --git a/py/autocomplete.py b/py/autocomplete.py index a456198..26810e0 100644 --- a/py/autocomplete.py +++ b/py/autocomplete.py @@ -1,19 +1,22 @@ from server import PromptServer from aiohttp import web -import os +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 + f.write(await request.text()) + return web.Response(status=200) diff --git a/web/js/autocompleter.js b/web/js/autocompleter.js index 686e034..4720e51 100644 --- a/web/js/autocompleter.js +++ b/web/js/autocompleter.js @@ -1,10 +1,43 @@ import { app } from "../../../scripts/app.js"; import { ComfyWidgets } from "../../../scripts/widgets.js"; import { api } from "../../../../scripts/api.js"; -import { $el } from "../../../../scripts/ui.js"; +import { $el, ComfyDialog } from "../../../../scripts/ui.js"; import { TextAreaAutoComplete } from "./common/autocomplete.js"; import { ModelInfoDialog } from "./common/modelInfoDialog.js"; +async function getCustomWords() { + const resp = await api.fetchApi("/pysssss/autocomplete", { cache: "no-store" }); + if (resp.status === 200) { + return await resp.text(); + } + return undefined; +} + +async function addCustomWords(text) { + if (!text) { + text = await getCustomWords(); + } + if (text) { + TextAreaAutoComplete.updateWords( + "pysssss.customwords", + text.split("\n").reduce((p, n) => { + n = n.trim(); + const pos = n.lastIndexOf(","); + let priority = undefined; + if (pos > -1) { + const v = parseInt(n.substr(pos + 1).trim()); + if (!isNaN(v)) { + priority = v; + n = n.substr(0, pos).trim(); + } + } + p[n] = { text: n, priority }; + return p; + }, {}) + ); + } +} + class EmbeddingInfoDialog extends ModelInfoDialog { async addInfo() { super.addInfo(); @@ -22,8 +55,64 @@ class EmbeddingInfoDialog extends ModelInfoDialog { } } +class CustomWordsDialog extends ComfyDialog { + async show() { + const text = await getCustomWords(); + this.words = $el("textarea", { + textContent: text, + style: { + width: "70vw", + height: "70vh", + }, + }); + super.show( + $el("div", [ + $el("h2", { + textContent: "Custom Autocomplete Words", + style: { + color: "#fff", + marginTop: 0, + textAlign: "center", + fontFamily: "sans-serif", + }, + }), + this.words, + ]) + ); + } + + createButtons() { + const btns = super.createButtons(); + const save = $el("button", { + type: "button", + textContent: "Save", + onclick: async (e) => { + try { + const res = await api.fetchApi("/pysssss/autocomplete", { method: "POST", body: this.words.value }); + if (res.status !== 200) { + throw new Error("Error saving: " + res.status + " " + res.statusText); + } + save.textContent = "Saved!"; + addCustomWords(this.words.value); + setTimeout(() => { + save.textContent = "Save"; + }, 500); + } catch (error) { + alert("Error saving word list!"); + console.error(error); + } + }, + }); + + btns.unshift(save); + return btns; + } +} + +const id = "pysssss.AutoCompleter"; + app.registerExtension({ - name: "pysssss.AutoCompleter", + name: id, init() { async function addEmbeddings() { const embeddings = await api.getEmbeddings(); @@ -41,32 +130,6 @@ app.registerExtension({ TextAreaAutoComplete.updateWords("pysssss.embeddings", words); } - 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(); - const pos = n.lastIndexOf(","); - let priority = undefined; - if (pos > -1) { - const v = parseInt(n.substr(pos + 1).trim()); - if (!isNaN(v)) { - priority = v; - n = n.substr(0, pos).trim(); - } - } - p[n] = { text: n, priority }; - return p; - }, {}) - ); - } - } - } - Promise.all([addEmbeddings(), addCustomWords()]); const STRING = ComfyWidgets.STRING; @@ -79,5 +142,45 @@ app.registerExtension({ return r; }; + + app.ui.settings.addSetting({ + id, + name: "🐍 Text Autocomplete", + defaultValue: true, + type: (name, setter, value) => { + return $el("tr", [ + $el("td", [ + $el("label", { + for: id.replaceAll(".", "-"), + textContent: name, + }), + ]), + $el("td", [ + $el("input", { + id: id.replaceAll(".", "-"), + type: "checkbox", + checked: value, + onchange: (event) => { + const checked = !!event.target.checked; + TextAreaAutoComplete.enabled = checked; + setter(checked); + }, + }), + $el("button", { + textContent: "Manage Custom Words", + onclick: () => { + app.ui.settings.element.close(); + new CustomWordsDialog().show(); + }, + style: { + fontSize: "14px", + display: "block", + marginTop: "5px", + }, + }), + ]), + ]); + }, + }); }, }); diff --git a/web/js/common/autocomplete.js b/web/js/common/autocomplete.js index b0f4daa..96c2268 100644 --- a/web/js/common/autocomplete.js +++ b/web/js/common/autocomplete.js @@ -302,6 +302,8 @@ class TextAreaCaretHelper { /*********************/ export class TextAreaAutoComplete { + static enabled = true; + static groups = {}; /** @type {Record} */ @@ -333,6 +335,8 @@ export class TextAreaAutoComplete { * @param {KeyboardEvent} e */ #keyDown(e) { + if (!TextAreaAutoComplete.enabled) return; + if (this.dropdown.parentElement) { // We are visible switch (e.key) { @@ -364,6 +368,7 @@ export class TextAreaAutoComplete { * @param {KeyboardEvent} e */ #keyPress(e) { + if (!TextAreaAutoComplete.enabled) return; if (this.dropdown.parentElement) { // We are visible switch (e.key) { @@ -380,6 +385,7 @@ export class TextAreaAutoComplete { } #keyUp(e) { + if (!TextAreaAutoComplete.enabled) return; if (this.dropdown.parentElement) { // We are visible switch (e.key) {