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

@@ -1,5 +1,7 @@
 # Booru tag autocompletion for A1111
 
+[![GitHub release (latest SemVer)](https://img.shields.io/github/v/release/DominikDoom/a1111-sd-webui-tagcomplete)](https://github.com/DominikDoom/a1111-sd-webui-tagcomplete/releases)
+
 This custom script serves as a drop-in extension for the popular [AUTOMATIC1111 web UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) for Stable Diffusion.
 
 It displays autocompletion hints for recognized tags from "image booru" boards such as Danbooru, which are primarily used for browsing Anime-style illustrations.
@@ -7,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, completion for negative or img2img prompt textboxes isn't supported, and there's no way to turn the script off from the ui, but I plan to get around to those features 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.
@@ -21,25 +27,38 @@ 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:
 ```json
 {
 	"tagFile": "danbooru.csv",
+	"activeIn": {
+		"txt2img": true,
+		"img2img": true,
+		"negativePrompts": true
+	},
 	"maxResults": 5,
 	"replaceUnderscores": true,
+	"escapeParentheses": true,
+	"useWildcards": true,
+	"useEmbeddings": true,
 	"colors": {
 		"danbooru": {
 			"0": ["lightblue", "dodgerblue"],
@@ -65,8 +84,12 @@ The config contains the following settings and defaults:
 | Setting	| Description |
 |---------|-------------|
 | 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.|
-| maxResults | How many results to show max. For the default tag set, the results are ordered by occurence count. |
+| 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. 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

@@ -0,0 +1,600 @@
+var acConfig = null;
+var acActive = true;
+
+// Style for new elements. Gets appended to the Gradio root.
+let autocompleteCSS_dark = `
+    .autocompleteResults {
+        position: absolute;
+        z-index: 999;
+        margin: 5px 0 0 0;
+        background-color: #0b0f19 !important;
+        border: 1px solid #4b5563 !important;
+        border-radius: 12px !important;
+        overflow-y: auto;
+    }
+    .autocompleteResultsList > li:nth-child(odd) {
+        background-color: #111827;
+    }
+    .autocompleteResultsList > li {
+        list-style-type: none;
+        padding: 10px;
+        cursor: pointer;
+    }
+    .autocompleteResultsList > li:hover {
+        background-color: #1f2937;
+    }
+    .autocompleteResultsList > li.selected {
+        background-color: #374151;
+    }
+`;
+let autocompleteCSS_light = `
+    .autocompleteResults {
+        position: absolute;
+        z-index: 999;
+        margin: 5px 0 0 0;
+        background-color: #ffffff !important;
+        border: 1.5px solid #e5e7eb !important;
+        border-radius: 12px !important;
+        overflow-y: auto;
+    }
+    .autocompleteResultsList > li:nth-child(odd) {
+        background-color: #f9fafb;
+    }
+    .autocompleteResultsList > li {
+        list-style-type: none;
+        padding: 10px;
+        cursor: pointer;
+    }
+    .autocompleteResultsList > li:hover {
+        background-color: #f5f6f8;
+    }
+    .autocompleteResultsList > li.selected {
+        background-color: #e5e7eb;
+    }
+`;
+
+// Parse the CSV file into a 2D array. Doesn't use regex, so it is very lightweight.
+function parseCSV(str) {
+    var arr = [];
+    var quote = false;  // 'true' means we're inside a quoted field
+
+    // Iterate over each character, keep track of current row and column (of the returned array)
+    for (var row = 0, col = 0, c = 0; c < str.length; c++) {
+        var cc = str[c], nc = str[c + 1];        // Current character, next character
+        arr[row] = arr[row] || [];             // Create a new row if necessary
+        arr[row][col] = arr[row][col] || '';   // Create a new column (start with empty string) if necessary
+
+        // If the current character is a quotation mark, and we're inside a
+        // quoted field, and the next character is also a quotation mark,
+        // add a quotation mark to the current column and skip the next character
+        if (cc == '"' && quote && nc == '"') { arr[row][col] += cc; ++c; continue; }
+
+        // If it's just one quotation mark, begin/end quoted field
+        if (cc == '"') { quote = !quote; continue; }
+
+        // If it's a comma and we're not in a quoted field, move on to the next column
+        if (cc == ',' && !quote) { ++col; continue; }
+
+        // If it's a newline (CRLF) and we're not in a quoted field, skip the next character
+        // and move on to the next row and move to column 0 of that new row
+        if (cc == '\r' && nc == '\n' && !quote) { ++row; col = 0; ++c; continue; }
+
+        // If it's a newline (LF or CR) and we're not in a quoted field,
+        // move on to the next row and move to column 0 of that new row
+        if (cc == '\n' && !quote) { ++row; col = 0; continue; }
+        if (cc == '\r' && !quote) { ++row; col = 0; continue; }
+
+        // Otherwise, append the current character to the current column
+        arr[row][col] += cc;
+    }
+    return arr;
+}
+
+// Load file
+function readFile(filePath) {
+    let request = new XMLHttpRequest();
+    request.open("GET", filePath, false);
+    request.send(null);
+    return request.responseText;
+}
+
+function loadCSV() {
+    let text = readFile(`file/tags/${acConfig.tagFile}`);
+    return parseCSV(text);
+}
+
+// Debounce function to prevent spamming the autocomplete function
+var dbTimeOut;
+const debounce = (func, wait = 300) => {
+    return function (...args) {
+        if (dbTimeOut) {
+            clearTimeout(dbTimeOut);
+        }
+
+        dbTimeOut = setTimeout(() => {
+            func.apply(this, args);
+        }, wait);
+    }
+}
+
+// Difference function to fix duplicates not being seen as changes in normal filter
+function difference(a, b) {
+    if (a.length == 0) {
+        return b;
+    }
+    if (b.length == 0) {
+        return a;
+    }
+
+    return [...b.reduce((acc, v) => acc.set(v, (acc.get(v) || 0) - 1),
+        a.reduce((acc, v) => acc.set(v, (acc.get(v) || 0) + 1), new Map())
+    )].reduce((acc, [v, count]) => acc.concat(Array(Math.abs(count)).fill(v)), []);
+}
+// Get the identifier for the text area to differentiate between positive and negative
+function getTextAreaIdentifier(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('#img2img_neg_prompt > label > textarea');
+
+    let modifier = "";
+    if (textArea === img2img_p || textArea === img2img_n) {
+        modifier += ".img2img";
+    }
+    if (textArea === txt2img_n || textArea === img2img_n) {
+        modifier += ".n";
+    } else {
+        modifier += ".p";
+    }
+    return modifier;
+}
+
+// Create the result list div and necessary styling
+function createResultsDiv(textArea) {
+    let resultsDiv = document.createElement("div");
+    let resultsList = document.createElement('ul');
+
+    let textAreaId = getTextAreaIdentifier(textArea);
+    let typeClass = textAreaId.replaceAll(".", " ");
+
+    resultsDiv.style.setProperty("max-height", acConfig.maxResults * 50 + "px");
+    resultsDiv.setAttribute('class', `autocompleteResults ${typeClass}`);
+    resultsList.setAttribute('class', 'autocompleteResultsList');
+    resultsDiv.appendChild(resultsList);
+
+    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) {
+    let textAreaId = getTextAreaIdentifier(textArea);
+    let resultsDiv = gradioApp().querySelector('.autocompleteResults' + textAreaId);
+    return resultsDiv.style.display === "block";
+}
+function showResults(textArea) {
+    let textAreaId = getTextAreaIdentifier(textArea);
+
+    let resultsDiv = gradioApp().querySelector('.autocompleteResults' + textAreaId);
+    resultsDiv.style.display = "block";
+}
+function hideResults(textArea) {
+    let textAreaId = getTextAreaIdentifier(textArea);
+
+    let resultsDiv = gradioApp().querySelector('.autocompleteResults' + textAreaId);
+    resultsDiv.style.display = "none";
+    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) {
+    let text = result[0];
+    let tagType = result[1];
+
+    let cursorPos = textArea.selectionStart;
+    var sanitizedText = text
+
+    // Replace differently depending on if it's a tag or wildcard
+    if (tagType === "wildcardFile") {
+        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;
+    }
+
+    if (acConfig.escapeParentheses) {
+        sanitizedText = sanitizedText
+            .replaceAll("(", "\\(")
+            .replaceAll(")", "\\)")
+            .replaceAll("[", "\\[")
+            .replaceAll("]", "\\]");
+    }
+
+    var prompt = textArea.value;
+
+
+    // Edit prompt text
+    let editStart = Math.max(cursorPos - tagword.length, 0);
+    let editEnd = Math.min(cursorPos + tagword.length, prompt.length);
+    let surrounding = prompt.substring(editStart, editEnd);
+    let match = surrounding.match(new RegExp(escapeRegExp(`${tagword}`)));
+    let afterInsertCursorPos = editStart + match.index + sanitizedText.length;
+
+    var optionalComma = "";
+    if (tagType !== "wildcardFile") {
+        optionalComma = surrounding.match(new RegExp(escapeRegExp(`${tagword},`))) !== null ? "" : ", ";
+    }
+
+    // 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.
+    textArea.dispatchEvent(new Event("input", { bubbles: true }));
+
+    // Update previous tags with the edited prompt to prevent re-searching the same term
+    let tags = newPrompt.match(/[^, ]+/g);
+    previousTags = tags;
+
+    // Hide results after inserting
+    if (tagType === "wildcardFile") {
+        // If it's a wildcard, we want to keep the results open so the user can select another wildcard
+        hideBlocked = true;
+        autocomplete(textArea, prompt, sanitizedText);
+        setTimeout(() => { hideBlocked = false; }, 100);
+    } else {
+        hideResults(textArea);
+    }
+}
+
+function addResultsToList(textArea, results, tagword) {
+    let textAreaId = getTextAreaIdentifier(textArea);
+    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];
+    let tagColors = acConfig.colors;
+
+    let mode = gradioApp().querySelector('.dark') ? 0 : 1;
+
+    for (let i = 0; i < results.length; i++) {
+        let result = results[i];
+        let li = document.createElement("li");
+        li.textContent = result[0];
+
+        // 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];
+            // Default to danbooru scheme if no matching one is found
+            if (colorGroup === undefined) colorGroup = tagColors["danbooru"];
+
+            li.style = `color: ${colorGroup[tagType][mode]};`;
+        }
+
+        // Add listener
+        li.addEventListener("click", function () { insertTextAtCursor(textArea, result, tagword); });
+        // Add element to list
+        resultsList.appendChild(li);
+    }
+}
+
+function updateSelectionStyle(textArea, num) {
+    let textAreaId = getTextAreaIdentifier(textArea);
+    let resultDiv = gradioApp().querySelector('.autocompleteResults' + textAreaId);
+    let resultsList = resultDiv.querySelector('ul');
+    let items = resultsList.getElementsByTagName('li');
+
+    for (let i = 0; i < items.length; i++) {
+        items[i].classList.remove('selected');
+    }
+
+    items[num].classList.add('selected');
+
+    // Set scrolltop to selected item if we are showing more than max results
+    if (items.length > acConfig.maxResults) {
+        let selected = items[num];
+        resultDiv.scrollTop = selected.offsetTop - resultDiv.offsetTop;
+    }
+}
+
+wildcardFiles = [];
+wildcards = {};
+embeddings = [];
+allTags = [];
+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);
+        return;
+    }
+
+    if (fixedTag === null) {
+        // Match tags with RegEx to get the last edited one
+        let tags = prompt.match(/[^, ]+/g);
+        let diff = difference(tags, previousTags)
+        previousTags = tags;
+
+        // Guard for no difference / only whitespace remaining
+        if (diff === null || diff.length === 0) {
+            if (!hideBlocked) hideResults(textArea);
+            return;
+        }
+
+        tagword = diff[0]
+
+        // Guard for empty tagword
+        if (tagword === null || tagword.length === 0) {
+            hideResults(textArea);
+            return;
+        }
+    } else {
+        tagword = fixedTag;
+    }
+
+    tagword = tagword.toLowerCase();
+
+    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 (acConfig.useWildcards && (tagword.startsWith("__") && !tagword.endsWith("__") || tagword === "__")) {
+        // Show available wildcard files
+        let tempResults = [];
+        if (tagword !== "__") {
+            tempResults = wildcardFiles.filter(x => x.toLowerCase().includes(tagword.replace("__", ""))) // Filter by tagword
+        } else {
+            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);
+    }
+    resultCount = results.length;
+
+    // Guard for empty results
+    if (resultCount === 0) {
+        hideResults(textArea);
+        return;
+    }
+
+    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;
+    if (!isVisible(textArea)) return
+
+    switch (event.key) {
+        case "ArrowUp":
+            if (selectedTag === null) {
+                selectedTag = resultCount - 1;
+            } else {
+                selectedTag = (selectedTag - 1 + resultCount) % resultCount;
+            }
+            break;
+        case "ArrowDown":
+            if (selectedTag === null) {
+                selectedTag = 0;
+            } else {
+                selectedTag = (selectedTag + 1) % resultCount;
+            }
+            break;
+        case "ArrowLeft":
+            selectedTag = 0;
+            break;
+        case "ArrowRight":
+            selectedTag = resultCount - 1;
+            break;
+        case "Enter":
+            if (selectedTag !== null) {
+                insertTextAtCursor(textArea, results[selectedTag], tagword);
+            }
+            break;
+        case "Escape":
+            hideResults(textArea);
+            break;
+    }
+    // Update highlighting
+    if (selectedTag !== null)
+        updateSelectionStyle(textArea, selectedTag);
+
+    // Prevent default behavior
+    event.preventDefault();
+    event.stopPropagation();
+}
+
+styleAdded = false;
+onUiUpdate(function () {
+    // Load config
+    if (acConfig === null) {
+        try {
+            acConfig = JSON.parse(readFile("file/tags/config.json"));
+        } catch (e) {
+            console.error("Error loading config.json: " + e);
+            return;
+        }
+    }
+    // Load main tags
+    if (allTags.length === 0) {
+        try {
+            allTags = loadCSV();
+        } catch (e) {
+            console.error("Error loading tags file: " + e);
+            return;
+        }
+    }
+    // Load wildcards
+    if (wildcardFiles.length === 0 && acConfig.useWildcards) {
+        try {
+            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] = 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}`);
+                }
+            });
+        } catch (e) {
+            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 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;
+    // Already added?
+    if (gradioApp().querySelector('.autocompleteResults.p') !== null
+        && (gradioApp().querySelector('.autocompleteResults.n') === null
+            && !acConfig.activeIn.negativePrompts)) {
+        return;
+    }
+
+    textAreas.forEach(area => {
+        // Skip directly if not found on the page
+        if (area === null || area === undefined) return;
+
+        // Return if autocomplete is disabled for the current area type in config
+        let textAreaId = getTextAreaIdentifier(area);
+        if (textAreaId.includes("p") || (textAreaId.includes("n") && acConfig.activeIn.negativePrompts)) {
+            if (textAreaId.includes("img2img")) {
+                if (!acConfig.activeIn.img2img) return;
+            } else {
+                if (!acConfig.activeIn.txt2img) return;
+            }
+        }
+
+        // Only add listeners once
+        if (!area.classList.contains('autocomplete')) {
+            // Add our new element
+            var resultsDiv = createResultsDiv(area);
+            area.parentNode.insertBefore(resultsDiv, area.nextSibling);
+            // Hide by default so it doesn't show up on page load
+            hideResults(area);
+
+            // Add autocomplete event listener
+            area.addEventListener('input', debounce(() => autocomplete(area, area.value), 100));
+            // Add focusout event listener
+            area.addEventListener('focusout', debounce(() => hideResults(area), 400));
+            // Add up and down arrow event listener
+            area.addEventListener('keydown', (e) => navigateInList(area, e));
+
+            // Add class so we know we've already added the listeners
+            area.classList.add('autocomplete');
+        }
+    });
+
+    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
+    let acStyle = document.createElement('style');
+    let css = gradioApp().querySelector('.dark') ? autocompleteCSS_dark : autocompleteCSS_light;
+    if (acStyle.styleSheet) {
+        acStyle.styleSheet.cssText = css;
+    } else {
+        acStyle.appendChild(document.createTextNode(css));
+    }
+    gradioApp().appendChild(acStyle);
+    styleAdded = true;
+});

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)

tagAutocomplete.js

@@ -1,359 +0,0 @@
-// Style for new elements. Gets appended to the Gradio root.
-const autocompleteCSS_dark = `
-    #autocompleteResults {
-        position: absolute;
-        z-index: 999;
-        margin: 5px 0 0 0;
-        background-color: #0b0f19 !important;
-        border: 1px solid #4b5563 !important;
-        border-radius: 12px !important;
-        overflow: hidden;
-    }
-    #autocompleteResultsList > li:nth-child(odd) {
-        background-color: #111827;
-    }
-    #autocompleteResultsList > li {
-        list-style-type: none;
-        padding: 10px;
-        cursor: pointer;
-    }
-    #autocompleteResultsList > li:hover {
-        background-color: #1f2937;
-    }
-    #autocompleteResultsList > li.selected {
-        background-color: #374151;
-    }
-`;
-const autocompleteCSS_light = `
-    #autocompleteResults {
-        position: absolute;
-        z-index: 999;
-        margin: 5px 0 0 0;
-        background-color: #ffffff !important;
-        border: 1.5px solid #e5e7eb !important;
-        border-radius: 12px !important;
-        overflow: hidden;
-    }
-    #autocompleteResultsList > li:nth-child(odd) {
-        background-color: #f9fafb;
-    }
-    #autocompleteResultsList > li {
-        list-style-type: none;
-        padding: 10px;
-        cursor: pointer;
-    }
-    #autocompleteResultsList > li:hover {
-        background-color: #f5f6f8;
-    }
-    #autocompleteResultsList > li.selected {
-        background-color: #e5e7eb;
-    }
-`;
-
-var acConfig = null;
-
-// Parse the CSV file into a 2D array. Doesn't use regex, so it is very lightweight.
-function parseCSV(str) {
-    var arr = [];
-    var quote = false;  // 'true' means we're inside a quoted field
-
-    // Iterate over each character, keep track of current row and column (of the returned array)
-    for (var row = 0, col = 0, c = 0; c < str.length; c++) {
-        var cc = str[c], nc = str[c+1];        // Current character, next character
-        arr[row] = arr[row] || [];             // Create a new row if necessary
-        arr[row][col] = arr[row][col] || '';   // Create a new column (start with empty string) if necessary
-
-        // If the current character is a quotation mark, and we're inside a
-        // quoted field, and the next character is also a quotation mark,
-        // add a quotation mark to the current column and skip the next character
-        if (cc == '"' && quote && nc == '"') { arr[row][col] += cc; ++c; continue; }
-
-        // If it's just one quotation mark, begin/end quoted field
-        if (cc == '"') { quote = !quote; continue; }
-
-        // If it's a comma and we're not in a quoted field, move on to the next column
-        if (cc == ',' && !quote) { ++col; continue; }
-
-        // If it's a newline (CRLF) and we're not in a quoted field, skip the next character
-        // and move on to the next row and move to column 0 of that new row
-        if (cc == '\r' && nc == '\n' && !quote) { ++row; col = 0; ++c; continue; }
-
-        // If it's a newline (LF or CR) and we're not in a quoted field,
-        // move on to the next row and move to column 0 of that new row
-        if (cc == '\n' && !quote) { ++row; col = 0; continue; }
-        if (cc == '\r' && !quote) { ++row; col = 0; continue; }
-
-        // Otherwise, append the current character to the current column
-        arr[row][col] += cc;
-    }
-    return arr;
-}
-
-// Load file
-function readFile(filePath) {
-    let request = new XMLHttpRequest();
-    request.open("GET", filePath, false);
-    request.send(null);
-    return request.responseText;
-}
-
-function loadCSV() {
-    let text = readFile(`file/tags/${acConfig.tagFile}`);
-    return parseCSV(text);
-}
-
-// Debounce function to prevent spamming the autocomplete function
-var dbTimeOut;
-const debounce = (func, wait = 300) => {
-    return function(...args) {
-        if (dbTimeOut) {
-            clearTimeout(dbTimeOut);
-        }
-
-        dbTimeOut = setTimeout(() => {
-            func.apply(this, args);
-        }, wait);
-    }
-}
-
-// Difference function to fix duplicates not being seen as changes in normal filter
-function difference(a, b) {
-    if (a.length == 0) {
-        return b;
-    }
-    if (b.length == 0) {
-        return a;
-    }
-
-    return [...b.reduce( (acc, v) => acc.set(v, (acc.get(v) || 0) - 1),
-            a.reduce( (acc, v) => acc.set(v, (acc.get(v) || 0) + 1), new Map() ) 
-    )].reduce( (acc, [v, count]) => acc.concat(Array(Math.abs(count)).fill(v)), [] );
-}
-
-// Create the result list div and necessary styling
-function createResultsDiv() {
-    let resultsDiv = document.createElement("div");
-    let resultsList = document.createElement('ul');
-    
-    resultsDiv.setAttribute('id', 'autocompleteResults');
-    resultsList.setAttribute('id', 'autocompleteResultsList');
-    resultsDiv.appendChild(resultsList);
-
-    return resultsDiv;
-}
-
-// The selected tag index. Needs to be up here so hide can access it.
-var selectedTag = null;
-
-// Show or hide the results div
-var isVisible = false;
-function showResults() {
-    let resultsDiv = gradioApp().querySelector('#autocompleteResults');
-    resultsDiv.style.display = "block";
-    isVisible = true;
-}
-function hideResults() {
-    let resultsDiv = gradioApp().querySelector('#autocompleteResults');
-    resultsDiv.style.display = "none";
-    isVisible = false;
-    selectedTag = null;
-}
-
-// On click, insert the tag into the prompt textbox with respect to the cursor position
-function insertTextAtCursor(text, tagword) {
-    let promptTextbox = gradioApp().querySelector('#txt2img_prompt > label > textarea');
-    let cursorPos = promptTextbox.selectionStart;
-    let sanitizedText = acConfig.replaceUnderscores ? text.replaceAll("_", " ") : text;
-    sanitizedText = acConfig.escapeParentheses ? sanitizedText.replaceAll("(", "\\(").replaceAll(")", "\\)") : sanitizedText;
-    
-    var prompt = promptTextbox.value;
-    let optionalComma = (prompt[cursorPos] === "," || prompt[cursorPos + tagword.length] === ",") ? "" : ", ";
-
-    // Edit prompt text
-    let direction = prompt.substring(cursorPos, cursorPos + tagword.length) === tagword ? 1 : -1;
-    if (direction === 1) {
-        promptTextbox.value = prompt.substring(0, cursorPos) + sanitizedText + optionalComma + prompt.substring(cursorPos + tagword.length)
-        // Update cursor position to after the inserted text
-        promptTextbox.selectionStart = cursorPos + sanitizedText.length + optionalComma.length;
-    } else {  
-        promptTextbox.value = prompt.substring(0, cursorPos - tagword.length) + sanitizedText + optionalComma + prompt.substring(cursorPos)
-        promptTextbox.selectionStart = cursorPos - tagword.length + sanitizedText.length + optionalComma.length;
-    }
-    prompt = promptTextbox.value;
-    promptTextbox.selectionEnd = promptTextbox.selectionStart;
-
-    // Hide results after inserting
-    hideResults();
-
-    // Update previous tags with the edited prompt to prevent re-searching the same term
-    let tags = prompt.match(/[^, ]+/g);
-    previousTags = tags;
-}
-
-function addResultsToList(results, tagword) {
-    let resultsList = gradioApp().querySelector('#autocompleteResultsList');
-    resultsList.innerHTML = "";
-
-    // Find right colors from config
-    let tagFileName = acConfig.tagFile.split(".")[0];
-    let tagColors = acConfig.colors;
-    //let colorIndex = Object.keys(tagColors).findIndex(key => key === tagFileName);
-    //let colorValues = Object.values(tagColors)[colorIndex];
-
-    let mode = gradioApp().querySelector('.dark') ? 0 : 1;
-
-    for (let i = 0; i < results.length; i++) {
-        let result = results[i];
-        let li = document.createElement("li");
-        li.innerHTML = result[0];
-
-        // Set the color of the tag
-        let tagType = result[1];
-        li.style = `color: ${tagColors[tagFileName][tagType][mode]};`;
-        // Add listener
-        li.addEventListener("click", function() { insertTextAtCursor(result[0], tagword); });
-        // Add element to list
-        resultsList.appendChild(li);
-    }
-}
-
-function updateSelectionStyle(num) {
-    let resultsList = gradioApp().querySelector('#autocompleteResultsList');
-    let items = resultsList.getElementsByTagName('li');
-
-    for (let i = 0; i < items.length; i++) {
-        items[i].classList.remove('selected');
-    }
-
-    items[num].classList.add('selected');
-}
-
-allTags = [];
-previousTags = [];
-results = [];
-tagword = "";
-resultCount = 0;
-function autocomplete(prompt) {
-    // Guard for empty prompt
-    if (prompt.length === 0) {
-        hideResults();
-        return;
-    }
-
-    // Match tags with RegEx to get the last edited one
-    let tags = prompt.match(/[^, ]+/g);
-    let diff = difference(tags, previousTags)
-    previousTags = tags;
-
-    // Guard for no difference / only whitespace remaining
-    if (diff === undefined || diff.length === 0) {
-        hideResults();
-        return;
-    }
-
-    tagword = diff[0]
-
-    // Guard for empty tagword
-    if (tagword === undefined || tagword.length === 0) {
-        hideResults();
-        return;
-    }
-    
-    results = allTags.filter(x => x[0].includes(tagword)).slice(0, acConfig.maxResults);
-    resultCount = results.length;
-
-    // Guard for empty results
-    if (resultCount === 0) {
-        hideResults();
-        return;
-    }
-
-    showResults();
-    addResultsToList(results, tagword);
-}
-
-function navigateInList(event) {
-    validKeys = ["ArrowUp", "ArrowDown", "ArrowLeft", "ArrowRight", "Enter", "Escape"];
-
-    if (!validKeys.includes(event.key)) return;
-    if (!isVisible) return
-
-    switch (event.key) {
-        case "ArrowUp":
-            if (selectedTag === null) {
-                selectedTag = resultCount - 1;
-            } else {
-                selectedTag = (selectedTag - 1 + resultCount) % resultCount;
-            }
-            break;
-        case "ArrowDown":
-            if (selectedTag === null) {
-                selectedTag = 0;
-            } else {
-                selectedTag = (selectedTag + 1) % resultCount;
-            }
-            break;
-        case "ArrowLeft":
-            selectedTag = 0;
-            break;
-        case "ArrowRight":
-            selectedTag = resultCount - 1;
-            break;
-        case "Enter":
-            if (selectedTag !== null) {
-                insertTextAtCursor(results[selectedTag][0], tagword);
-            }
-            break;
-        case "Escape":
-            hideResults();
-            break;
-    }
-    // Update highlighting
-    if (selectedTag !== null)
-        updateSelectionStyle(selectedTag);
-
-    // Prevent default behavior
-    event.preventDefault();
-    event.stopPropagation();
-}
-
-onUiUpdate(function(){
-    // One-time CSV setup
-    if (acConfig === null) acConfig = JSON.parse(readFile("file/tags/config.json"));
-    if (allTags.length === 0) allTags = loadCSV();
-
-	let promptTextbox = gradioApp().querySelector('#txt2img_prompt > label > textarea');
-	
-    if (promptTextbox === null) return;
-    if (gradioApp().querySelector('#autocompleteResults') != null) return;
-
-    // Only add listeners once
-    if (!promptTextbox.classList.contains('autocomplete')) {
-        // Add our new element
-        var resultsDiv = gradioApp().querySelector('#autocompleteResults') ?? createResultsDiv();
-        promptTextbox.parentNode.insertBefore(resultsDiv, promptTextbox.nextSibling);
-        // Hide by default so it doesn't show up on page load
-        hideResults();
-        
-        // Add autocomplete event listener
-        promptTextbox.addEventListener('input', debounce(() => autocomplete(promptTextbox.value), 100));
-        // Add focusout event listener
-        promptTextbox.addEventListener('focusout', debounce(() => hideResults(), 400));
-        // Add up and down arrow event listener
-        promptTextbox.addEventListener('keydown', (e) => navigateInList(e));
-
-        // Add class so we know we've already added the listeners
-        promptTextbox.classList.add('autocomplete');
-
-        // Add style to dom
-        let acStyle = document.createElement('style');
-
-        let css = gradioApp().querySelector('.dark') ? autocompleteCSS_dark : autocompleteCSS_light;
-        if (acStyle.styleSheet) {
-            acStyle.styleSheet.cssText = css;
-        } else {
-            acStyle.appendChild(document.createTextNode(css));
-        }
-        gradioApp().appendChild(acStyle);
-    }
-});

tags/config.json

@@ -1,8 +1,15 @@
 {
 	"tagFile": "danbooru.csv",
+	"activeIn": {
+		"txt2img": true,
+		"img2img": true,
+		"negativePrompts": true
+	},
 	"maxResults": 5,
 	"replaceUnderscores": true,
 	"escapeParentheses": true,
+	"useWildcards": true,
+	"useEmbeddings": true,
 	"colors": {
 		"danbooru": {
 			"0": ["lightblue", "dodgerblue"],