Merge pull request #15 from sgmklp/dev

Since we now search automatically, the script also doesn't try to load anything if none are found

README.md

@@ -9,10 +9,14 @@ Since some Stable Diffusion models were trained using this information, for exam
 
 I created this script as a convenience tool since it reduces the need of switching back and forth between the web UI and a booru site to copy-paste tags.
 
-You can either download the files manually as described below, or use a pre-packaged version from [Releases](https://github.com/DominikDoom/a1111-sd-webui-tagcomplete/releases).
+You can either clone / download the files manually as described [below](#installation), or use a pre-packaged version from [Releases](https://github.com/DominikDoom/a1111-sd-webui-tagcomplete/releases).
 
-### Disclaimer:
-This script is definitely not optimized, and it's not very intelligent. The tags are simply recommended based on their natural order in the CSV, which is their respective image count for the default Danbooru tag list. Also, at least for now there's no way to turn the script off from the ui, but I plan to get around to that eventually.
+### Wildcard & Embedding support
+Autocompletion also works with wildcard files used by [this script](https://github.com/jtkelm2/stable-diffusion-webui-1/blob/master/scripts/wildcards.py) of the same name (demo video further down). This enables you to either insert categories to be replaced by the script, or even replace them with the actual wildcard file content in the same step.
+
+It also scans the embeddings folder and displays completion hints for the names of all .pt and .bin files inside if you start typing `<`. Note that some normal tags also use < in Kaomoji (like ">_<" for example), so the results will contain both.
+
+Both are now enabled by default and scan the `/embeddings` and `/scripts/wildcards` folders automatically.
 
 ### Known Issues:
 If `replaceUnderscores` is active, the script will currently only partly replace edited tags containing multiple words in brackets.
@@ -23,17 +27,22 @@ Demo video (with keyboard navigation):
 
 https://user-images.githubusercontent.com/34448969/195344430-2b5f9945-b98b-4943-9fbc-82cf633321b1.mp4
 
+Wildcard script support:
+
+https://user-images.githubusercontent.com/34448969/195632461-49d226ae-d393-453d-8f04-1e44b073234c.mp4
+
 Dark and Light mode supported, including tag colors:
 
 ![tagtypes](https://user-images.githubusercontent.com/34448969/195177127-f63949f8-271d-4767-bccd-f1b5e818a7f8.png)
 ![tagtypes_light](https://user-images.githubusercontent.com/34448969/195180061-ceebcc25-9e4c-424f-b0c9-ba8e8f4f17f4.png)
 
-
 ## Installation
-Simply put `tagAutocomplete.js` into the `javascript` folder of your web UI installation. It will run automatically the next time the web UI is started.
-For the script to work, you also need to download the `tags` folder from this repo and paste it and its contents into the web UI root, or create them there manually.
+Simply copy the `javascript`, `scripts` and `tags` folder into your web UI installation root. It will run automatically the next time the web UI is started.
 
 The tags folder contains `config.json` and the tag data the script uses for autocompletion. By default, Danbooru and e621 tags are included.
+After scanning for embeddings and wildcards, the script will also create a `temp` directory here which lists the found files so they can be accessed in the browser side of the script. You can delete the temp folder without consequences as it will be recreated on the next startup.
+### Important:
+The script needs **all three folders** to work properly.
 
 ### Config
 The config contains the following settings and defaults:
@@ -48,6 +57,8 @@ The config contains the following settings and defaults:
 	"maxResults": 5,
 	"replaceUnderscores": true,
 	"escapeParentheses": true,
+	"useWildcards": true,
+	"useEmbeddings": true,
 	"colors": {
 		"danbooru": {
 			"0": ["lightblue", "dodgerblue"],
@@ -74,9 +85,11 @@ The config contains the following settings and defaults:
 |---------|-------------|
 | tagFile | Specifies the tag file to use. You can provide a custom tag database of your liking, but since the script was developed with Danbooru tags in mind, it might not work properly with other configurations.|
 | activeIn | Allows to selectively (de)activate the script for txt2img, img2img, and the negative prompts for both. |
-| maxResults | How many results to show max. For the default tag set, the results are ordered by occurence count. |
+| maxResults | How many results to show max. For the default tag set, the results are ordered by occurence count. For embeddings and wildcards it will show all results in a scrollable list. |
 | replaceUnderscores | If true, undescores are replaced with spaces on clicking a tag. Might work better for some models. |
 | escapeParentheses | If true, escapes tags containing () so they don't contribute to the web UI's prompt weighting functionality. |
+| useWildcards | Used to toggle the wildcard completion functionality. |
+| useEmbeddings | Used to toggle the embedding completion functionality. |
 | colors | Contains customizable colors for the tag types, you can add new ones here for custom tag files (same name as filename, without the .csv). The first value is for dark, the second for light mode. Color names and hex codes should both work.|
 
 ### CSV tag data

javascript/tagAutocomplete.js

@@ -1,4 +1,5 @@
 var acConfig = null;
+var acActive = true;
 
 // Style for new elements. Gets appended to the Gradio root.
 let autocompleteCSS_dark = `
@@ -131,10 +132,10 @@ function difference(a, b) {
 }
 // Get the identifier for the text area to differentiate between positive and negative
 function getTextAreaIdentifier(textArea) {
-    let txt2img_n = gradioApp().querySelector('#negative_prompt > label > textarea');
+    let txt2img_n = gradioApp().querySelector('#txt2img_neg_prompt > label > textarea');
     let img2img = gradioApp().querySelector('#tab_img2img');
     let img2img_p = img2img.querySelector('#img2img_prompt > label > textarea');
-    let img2img_n = img2img.querySelector('#negative_prompt > label > textarea');
+    let img2img_n = img2img.querySelector('#img2img_neg_prompt > label > textarea');
 
     let modifier = "";
     if (textArea === img2img_p || textArea === img2img_n) {
@@ -164,8 +165,28 @@ function createResultsDiv(textArea) {
     return resultsDiv;
 }
 
+// Create the checkbox to enable/disable autocomplete
+function createCheckbox() {
+    let label = document.createElement("label");
+    let input = document.createElement("input");
+    let span = document.createElement("span");
+
+    label.setAttribute('id', 'acActiveCheckbox');
+    label.setAttribute('class', '"flex items-center text-gray-700 text-sm rounded-lg cursor-pointer dark:bg-transparent');
+    input.setAttribute('type', 'checkbox');
+    input.setAttribute('class', 'gr-check-radio gr-checkbox')
+    span.setAttribute('class', 'ml-2');
+
+    span.textContent = "Enable Autocomplete";
+
+    label.appendChild(input);
+    label.appendChild(span);
+    return label;
+}
+
 // The selected tag index. Needs to be up here so hide can access it.
 var selectedTag = null;
+var previousTags = [];
 
 // Show or hide the results div
 function isVisible(textArea) {
@@ -187,6 +208,10 @@ function hideResults(textArea) {
     selectedTag = null;
 }
 
+function escapeRegExp(string) {
+    return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string
+}
+
 let hideBlocked = false;
 // On click, insert the tag into the prompt textbox with respect to the cursor position
 function insertTextAtCursor(textArea, result, tagword) {
@@ -201,11 +226,19 @@ function insertTextAtCursor(textArea, result, tagword) {
         sanitizedText = "__" + text.replace("Wildcards: ", "") + "__";
     } else if (tagType === "wildcardTag") {
         sanitizedText = text.replace(/^.*?: /g, "");
+    } else if (tagType === "embedding") {
+        sanitizedText = `<${text.replace(/^.*?: /g, "")}>`;
     } else {
         sanitizedText = acConfig.replaceUnderscores ? text.replaceAll("_", " ") : text;
     }
 
-    sanitizedText = acConfig.escapeParentheses ? sanitizedText.replaceAll("(", "\\(").replaceAll(")", "\\)") : sanitizedText;
+    if (acConfig.escapeParentheses) {
+        sanitizedText = sanitizedText
+            .replaceAll("(", "\\(")
+            .replaceAll(")", "\\)")
+            .replaceAll("[", "\\[")
+            .replaceAll("]", "\\]");
+    }
 
     var prompt = textArea.value;
 
@@ -214,21 +247,22 @@ function insertTextAtCursor(textArea, result, tagword) {
     let editStart = Math.max(cursorPos - tagword.length, 0);
     let editEnd = Math.min(cursorPos + tagword.length, prompt.length);
     let surrounding = prompt.substring(editStart, editEnd);
-    let insert = surrounding.replace(tagword, sanitizedText);
+    let match = surrounding.match(new RegExp(escapeRegExp(`${tagword}`)));
+    let afterInsertCursorPos = editStart + match.index + sanitizedText.length;
 
-    // Add back start
-    var newPrompt = prompt.substring(0, editStart) + insert;
-    // Add comma if needed
     var optionalComma = "";
     if (tagType !== "wildcardFile") {
-        optionalComma = surrounding.match(`/${tagword},/g`) !== null ? "" : ", ";
+        optionalComma = surrounding.match(new RegExp(escapeRegExp(`${tagword},`))) !== null ? "" : ", ";
     }
-    // Set selection after insertion
-    textArea.selectionStart = editStart + newPrompt.length + optionalComma.length;
-    textArea.selectionEnd = textArea.selectionStart;
-    // Add back end
-    newPrompt += optionalComma + prompt.substring(editEnd);
+
+    // Replace partial tag word with new text, add comma if needed
+    let insert = surrounding.replace(tagword, sanitizedText + optionalComma);
+
+    // Add back start
+    var newPrompt = prompt.substring(0, editStart) + insert + prompt.substring(editEnd);
     textArea.value = newPrompt;
+    textArea.selectionStart = afterInsertCursorPos + optionalComma.length;
+    textArea.selectionEnd = textArea.selectionStart
 
     // Since we've modified a Gradio Textbox component manually, we need to simulate an `input` DOM event to ensure its
     // internal Svelte data binding remains in sync.
@@ -251,8 +285,13 @@ function insertTextAtCursor(textArea, result, tagword) {
 
 function addResultsToList(textArea, results, tagword) {
     let textAreaId = getTextAreaIdentifier(textArea);
-    let resultsList = gradioApp().querySelector('.autocompleteResults' + textAreaId + ' > ul');
+    let resultDiv = gradioApp().querySelector('.autocompleteResults' + textAreaId);
+    let resultsList = resultDiv.querySelector('ul');
+
+    // Reset list, selection and scrollTop since the list changed
     resultsList.innerHTML = "";
+    selectedTag = null;
+    resultDiv.scrollTop = 0;
 
     // Find right colors from config
     let tagFileName = acConfig.tagFile.split(".")[0];
@@ -263,10 +302,10 @@ function addResultsToList(textArea, results, tagword) {
     for (let i = 0; i < results.length; i++) {
         let result = results[i];
         let li = document.createElement("li");
-        li.innerHTML = result[0];
+        li.textContent = result[0];
 
-        // Wildcards have no tag type
-        if (!result[1].startsWith("wildcard")) {
+        // Wildcards & Embeds have no tag type
+        if (!result[1].startsWith("wildcard") && result[1] !== "embedding") {
             // Set the color of the tag
             let tagType = result[1];
             let colorGroup = tagColors[tagFileName];
@@ -304,12 +343,15 @@ function updateSelectionStyle(textArea, num) {
 
 wildcardFiles = [];
 wildcards = {};
+embeddings = [];
 allTags = [];
-previousTags = [];
 results = [];
 tagword = "";
 resultCount = 0;
 function autocomplete(textArea, prompt, fixedTag = null) {
+    // Return if the function is deactivated in the UI
+    if (!acActive) return;
+
     // Guard for empty prompt
     if (prompt.length === 0) {
         hideResults(textArea);
@@ -341,14 +383,14 @@ function autocomplete(textArea, prompt, fixedTag = null) {
 
     tagword = tagword.toLowerCase();
 
-    if ([...tagword.matchAll(/\b__([^,_ ]+)__([^, ]*)\b/g)].length > 0 && acConfig.useWildcards) {
+    if (acConfig.useWildcards && [...tagword.matchAll(/\b__([^,_ ]+)__([^, ]*)\b/g)].length > 0) {
         // Show wildcards from a file with that name
         wcMatch = [...tagword.matchAll(/\b__([^,_ ]+)__([^, ]*)\b/g)]
         let wcFile = wcMatch[0][1];
         let wcWord = wcMatch[0][2];
         results = wildcards[wcFile].filter(x => (wcWord !== null) ? x.toLowerCase().includes(wcWord) : x) // Filter by tagword
             .map(x => [wcFile + ": " + x.trim(), "wildcardTag"]); // Mark as wildcard
-    } else if ((tagword.startsWith("__") && !tagword.endsWith("__") || tagword === "__") && acConfig.useWildcards) {
+    } else if (acConfig.useWildcards && (tagword.startsWith("__") && !tagword.endsWith("__") || tagword === "__")) {
         // Show available wildcard files
         let tempResults = [];
         if (tagword !== "__") {
@@ -357,6 +399,17 @@ function autocomplete(textArea, prompt, fixedTag = null) {
             tempResults = wildcardFiles;
         }
         results = tempResults.map(x => ["Wildcards: " + x.trim(), "wildcardFile"]); // Mark as wildcard
+    } else if (acConfig.useEmbeddings && tagword.match(/<[^,> ]*>?/g)) {
+        // Show embeddings
+        let tempResults = [];
+        if (tagword !== "<") {
+            tempResults = embeddings.filter(x => x.toLowerCase().includes(tagword.replace("<", ""))) // Filter by tagword
+        } else {
+            tempResults = embeddings;
+        }
+        // Since some tags are kaomoji, we have to still get the normal results first.
+        genericResults = allTags.filter(x => x[0].toLowerCase().includes(tagword)).slice(0, acConfig.maxResults);
+        results = genericResults.concat(tempResults.map(x => ["Embeddings: " + x.trim(), "embedding"])); // Mark as embedding
     } else {
         results = allTags.filter(x => x[0].toLowerCase().includes(tagword)).slice(0, acConfig.maxResults);
     }
@@ -368,12 +421,14 @@ function autocomplete(textArea, prompt, fixedTag = null) {
         return;
     }
 
-    selectedTag = null; // Reset since the list changed
     showResults(textArea);
     addResultsToList(textArea, results, tagword);
 }
 
 function navigateInList(textArea, event) {
+    // Return if the function is deactivated in the UI
+    if (!acActive) return;
+
     validKeys = ["ArrowUp", "ArrowDown", "ArrowLeft", "ArrowRight", "Enter", "Escape"];
 
     if (!validKeys.includes(event.key)) return;
@@ -420,7 +475,7 @@ function navigateInList(textArea, event) {
 
 styleAdded = false;
 onUiUpdate(function () {
-    // One-time config, tags & wildcards loading
+    // Load config
     if (acConfig === null) {
         try {
             acConfig = JSON.parse(readFile("file/tags/config.json"));
@@ -429,6 +484,7 @@ onUiUpdate(function () {
             return;
         }
     }
+    // Load main tags
     if (allTags.length === 0) {
         try {
             allTags = loadCSV();
@@ -437,17 +493,17 @@ onUiUpdate(function () {
             return;
         }
     }
+    // Load wildcards
     if (wildcardFiles.length === 0 && acConfig.useWildcards) {
         try {
-            wildcardFiles = readFile("file/tags/wildcardNames.txt").split("\n")
-                .filter(x => !x.startsWith("//")) // Remove comments
-                .filter(x => x.toLowerCase().includes(tagword.substring(2))) // Filter by tagword
+            wildcardFiles = readFile("file/tags/temp/wc.txt").split("\n")
                 .filter(x => x.trim().length > 0) // Remove empty lines
+                .map(x => x.trim().replace(".txt", "")); // Remove file extension & newlines
 
             wildcardFiles.forEach(fName => {
                 try {
-                    wildcards[fName.trim()] = readFile(`file/scripts/wildcards/${fName}.txt`).split("\n")
-                        .filter(x => x.trim().length > 0) // Remove empty lines
+                    wildcards[fName] = readFile(`file/scripts/wildcards/${fName}.txt`).split("\n")
+                        .filter(x => x.trim().length > 0); // Remove empty lines
                 } catch (e) {
                     console.log(`Could not load wildcards for ${fName}`);
                 }
@@ -456,12 +512,25 @@ onUiUpdate(function () {
             console.error("Error loading wildcardNames.txt: " + e);
         }
     }
+    // Load embeddings
+    if (embeddings.length === 0 && acConfig.useEmbeddings) {
+        try {
+            embeddings = readFile("file/tags/temp/emb.txt").split("\n")
+                .filter(x => x.trim().length > 0) // Remove empty lines
+                .map(x => x.replace(".bin", "").replace(".pt", "")); // Remove file extensions
+        } catch (e) {
+            console.error("Error loading embeddings.txt: " + e);
+        }
+    }
 
     // Find all textareas
     let txt2imgTextArea = gradioApp().querySelector('#txt2img_prompt > label > textarea');
     let img2imgTextArea = gradioApp().querySelector('#img2img_prompt > label > textarea');
-    let negativeTextAreas = Array.from(gradioApp().querySelectorAll('#negative_prompt > label > textarea'));
-    let textAreas = [txt2imgTextArea, img2imgTextArea, negativeTextAreas[0], negativeTextAreas[1]];
+    let txt2imgTextArea_n = gradioApp().querySelector('#txt2img_neg_prompt > label > textarea');
+    let img2imgTextArea_n = gradioApp().querySelector('#img2img_neg_prompt > label > textarea');
+    let textAreas = [txt2imgTextArea, img2imgTextArea, txt2imgTextArea_n, img2imgTextArea_n];
+
+    let quicksettings = gradioApp().querySelector('#quicksettings');
 
     // Not found, we're on a page without prompt textareas
     if (textAreas.every(v => v === null || v === undefined)) return;
@@ -506,6 +575,16 @@ onUiUpdate(function () {
         }
     });
 
+    if (gradioApp().querySelector("#acActiveCheckbox") === null) {
+        // Add toggle switch
+        let cb = createCheckbox();
+        cb.querySelector("input").checked = acActive;
+        cb.querySelector("input").addEventListener("change", (e) => {
+            acActive = e.target.checked;
+        });
+        quicksettings.parentNode.insertBefore(cb, quicksettings.nextSibling);
+    }
+
     if (styleAdded) return;
 
     // Add style to dom

scripts/tag_autocomplete_helper.py

@@ -0,0 +1,48 @@
+# This helper script scans folders for wildcards and embeddings and writes them
+# to a temporary file to expose it to the javascript side
+
+import os
+
+# The path to the folder containing the wildcards and embeddings
+FILE_DIR = os.path.dirname(os.path.realpath("__file__"))
+WILDCARD_PATH = os.path.join(FILE_DIR, 'scripts/wildcards')
+EMB_PATH = os.path.join(FILE_DIR, 'embeddings')
+# The path to the temporary file
+TEMP_PATH = os.path.join(FILE_DIR, 'tags/temp')
+
+
+def get_wildcards():
+    """Returns a list of all wildcards"""
+    return filter(lambda f: f.endswith(".txt"), os.listdir(WILDCARD_PATH))
+
+
+def get_embeddings():
+    """Returns a list of all embeddings"""
+    return filter(lambda f: f.endswith(".bin") or f.endswith(".pt"), os.listdir(EMB_PATH))
+
+
+def write_to_temp_file(name, data):
+    """Writes the given data to a temporary file"""
+    with open(os.path.join(TEMP_PATH, name), 'w', encoding="utf-8") as f:
+        f.write(('\n'.join(data)))
+
+
+# Check if the temp path exists and create it if not
+if not os.path.exists(TEMP_PATH):
+    os.makedirs(TEMP_PATH)
+    # Set up files to ensure the script doesn't fail to load them
+    # even if no wildcards or embeddings are found
+    write_to_temp_file('wc.txt', [])
+    write_to_temp_file('emb.txt', [])
+
+# Write wildcards to wc.txt if found
+if os.path.exists(WILDCARD_PATH):
+    wildcards = get_wildcards()
+    if wildcards:
+        write_to_temp_file('wc.txt', wildcards)
+
+# Write embeddings to emb.txt if found
+if os.path.exists(EMB_PATH):
+    embeddings = get_embeddings()
+    if embeddings:
+        write_to_temp_file('emb.txt', embeddings)

tags/config.json

@@ -8,7 +8,8 @@
 	"maxResults": 5,
 	"replaceUnderscores": true,
 	"escapeParentheses": true,
-	"useWildcards": false,
+	"useWildcards": true,
+	"useEmbeddings": true,
 	"colors": {
 		"danbooru": {
 			"0": ["lightblue", "dodgerblue"],

tags/wildcardNames.txt

@@ -1,7 +0,0 @@
-// Put the file names of wildcard files you want to use here. Needed so that the script can access them.
-// The default ones are the following, you can uncomment them if you have them
-//adjective
-//artist
-//genre
-//site
-//style