Merge pull request #3 from HalfMAI/fix_input_event

Note: This will break coloring at the moment since e621 uses more tag types than Danbooru. I'll try to implement custom color support soon.

README.md

@@ -7,15 +7,19 @@ 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.
 
-[[Setup instructions]](#installation)
+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).
 
 ### 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, neither keyboard selection for tags nor completion for negative or img2img prompt textboxes is supported, and there's no way to turn the feature off from the ui, but I plan to get around to those features eventually.
+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.
+
+### Known Issues:
+If `replaceUnderscores` is active, the script will currently only partly replace edited tags containing multiple words in brackets.
+For example, editing `atago (azur lane)`, it would be replaced with e.g. `taihou (azur lane), lane)`, since the script currently doesn't see the second part of the bracket as the same tag. So in those cases you should delete the old tag beforehand.
 
 ## Screenshots
-Demo video
+Demo video (with keyboard navigation):
 
-https://user-images.githubusercontent.com/34448969/195185810-547d8d0a-bf87-465d-91f1-7fb5c3259c3f.mp4
+https://user-images.githubusercontent.com/34448969/195344430-2b5f9945-b98b-4943-9fbc-82cf633321b1.mp4
 
 Dark and Light mode supported, including tag colors:
 
@@ -27,7 +31,7 @@ Dark and Light mode supported, including tag colors:
 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.
 
-The tags folder contains two files: `config.json` and `danbooru.csv`. This is the data the script uses for autocompletion.
+The tags folder contains `config.json` and the tag data the script uses for autocompletion. By default, Danbooru and e621 tags are included.
 
 ### Config
 The config contains the following settings and defaults:
@@ -35,7 +39,27 @@ The config contains the following settings and defaults:
 {
 	"tagFile": "danbooru.csv",
 	"maxResults": 5,
-	"replaceUnderscores": true
+	"replaceUnderscores": true,
+	"colors": {
+		"danbooru": {
+			"0": ["lightblue", "dodgerblue"],
+			"1": ["indianred", "firebrick"],
+			"3": ["violet", "darkorchid"],
+			"4": ["lightgreen", "darkgreen"],
+			"5": ["orange", "darkorange"]
+		},
+		"e621": {
+			"-1": ["red", "maroon"],
+			"0": ["lightblue", "dodgerblue"],
+			"1": ["gold", "goldenrod"],
+			"3": ["violet", "darkorchid"],
+			"4": ["lightgreen", "darkgreen"],
+			"5": ["tomato", "darksalmon"],
+			"6": ["red", "maroon"],
+			"7": ["whitesmoke", "black"],
+			"8": ["seagreen", "darkseagreen"]
+		}
+	}
 }
 ```
 | Setting	| Description |
@@ -43,6 +67,7 @@ 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.|
 | maxResults | How many results to show max. For the default tag set, the results are ordered by occurence count. |
 | replaceUnderscores | If true, undescores are replaced with spaces on clicking a tag. Might work better for some models. |
+| 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
 The script expects a CSV file with tags saved in the following way:
@@ -63,4 +88,17 @@ The numbering system follows the [tag API docs](https://danbooru.donmai.us/wiki_
 |4	    | Character   |
 |5	    | Meta        |
 
+or of e621:
+| Value	| Description |
+|-------|-------------|
+|-1	    | Invalid     |
+|0	    | General     |
+|1	    | Artist      |
+|3	    | Copyright   |
+|4	    | Character   |
+|5	    | Species     |
+|6	    | Invalid     |
+|7	    | Meta        |
+|8	    | Lore        |
+
 The tag type is used for coloring entries in the result list.

tagAutocomplete.js

@@ -20,6 +20,9 @@ const autocompleteCSS_dark = `
     #autocompleteResultsList > li:hover {
         background-color: #1f2937;
     }
+    #autocompleteResultsList > li.selected {
+        background-color: #374151;
+    }
 `;
 const autocompleteCSS_light = `
     #autocompleteResults {
@@ -42,6 +45,9 @@ const autocompleteCSS_light = `
     #autocompleteResultsList > li:hover {
         background-color: #f5f6f8;
     }
+    #autocompleteResultsList > li.selected {
+        background-color: #e5e7eb;
+    }
 `;
 
 var acConfig = null;
@@ -110,6 +116,20 @@ const debounce = (func, wait = 300) => {
     }
 }
 
+// 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");
@@ -122,14 +142,21 @@ function createResultsDiv() {
     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
@@ -137,17 +164,28 @@ function insertTextAtCursor(text, tagword) {
     let promptTextbox = gradioApp().querySelector('#txt2img_prompt > label > textarea');
     let cursorPos = promptTextbox.selectionStart;
     let sanitizedText = acConfig.replaceUnderscores ? text.replaceAll("_", " ") : text;
-    let optionalComma = (promptTextbox.value[cursorPos] == ",") ? "" : ", ";
+    sanitizedText = acConfig.escapeParentheses ? sanitizedText.replaceAll("(", "\\(").replaceAll(")", "\\)") : sanitizedText;
+    
+    var prompt = promptTextbox.value;
+    let optionalComma = (prompt[cursorPos] === "," || prompt[cursorPos + tagword.length] === ",") ? "" : ", ";
 
     // Edit prompt text
-    var prompt = promptTextbox.value;
-    promptTextbox.value = prompt.substring(0, cursorPos - tagword.length) + sanitizedText + optionalComma + prompt.substring(cursorPos);
+    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;
-
-    // Update cursor position to after the inserted text
-    promptTextbox.selectionStart = cursorPos + sanitizedText.length;
     promptTextbox.selectionEnd = promptTextbox.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.
+    promptTextbox.dispatchEvent(new Event("input", { bubbles: true }));
+
     // Hide results after inserting
     hideResults();
 
@@ -156,57 +194,80 @@ function insertTextAtCursor(text, tagword) {
     previousTags = tags;
 }
 
-const colors_dark = ["lightblue", "indianred", "unused", "violet", "lightgreen", "orange"];
-const colors_light = ["dodgerblue", "firebrick", "unused", "darkorchid", "darkgreen", "darkorange" ]
 function addResultsToList(results, tagword) {
     let resultsList = gradioApp().querySelector('#autocompleteResultsList');
     resultsList.innerHTML = "";
 
-    let colors = gradioApp().querySelector('.dark') ? colors_dark : colors_light;
+    // 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];
-        li.style = `color: ${colors[result[1]]};`;
+
+        // 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) {
+    if (prompt.length === 0) {
         hideResults();
         return;
     }
 
     // Match tags with RegEx to get the last edited one
     let tags = prompt.match(/[^, ]+/g);
-    let difference = tags.filter(x => !previousTags.includes(x));
+    let diff = difference(tags, previousTags)
     previousTags = tags;
 
     // Guard for no difference / only whitespace remaining
-    if (difference == undefined || difference.length == 0) {
+    if (diff === undefined || diff.length === 0) {
         hideResults();
         return;
     }
 
-    let tagword = difference[0]
+    tagword = diff[0]
 
     // Guard for empty tagword
-    if (tagword == undefined || tagword.length == 0) {
+    if (tagword === undefined || tagword.length === 0) {
         hideResults();
         return;
     }
     
-    let results = allTags.filter(x => x[0].includes(tagword)).slice(0, acConfig.maxResults);
+    results = allTags.filter(x => x[0].includes(tagword)).slice(0, acConfig.maxResults);
+    resultCount = results.length;
 
     // Guard for empty results
-    if (results.length == 0) {
+    if (resultCount === 0) {
         hideResults();
         return;
     }
@@ -215,14 +276,59 @@ function autocomplete(prompt) {
     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();
+    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 (promptTextbox === null) return;
     if (gradioApp().querySelector('#autocompleteResults') != null) return;
 
     // Only add listeners once
@@ -237,6 +343,8 @@ onUiUpdate(function(){
         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');

tags/config.json

@@ -1,5 +1,26 @@
 {
 	"tagFile": "danbooru.csv",
-	"maxResults": 10,
-	"replaceUnderscores": false
+	"maxResults": 5,
+	"replaceUnderscores": true,
+	"escapeParentheses": true,
+	"colors": {
+		"danbooru": {
+			"0": ["lightblue", "dodgerblue"],
+			"1": ["indianred", "firebrick"],
+			"3": ["violet", "darkorchid"],
+			"4": ["lightgreen", "darkgreen"],
+			"5": ["orange", "darkorange"]
+		},
+		"e621": {
+			"-1": ["red", "maroon"],
+			"0": ["lightblue", "dodgerblue"],
+			"1": ["gold", "goldenrod"],
+			"3": ["violet", "darkorchid"],
+			"4": ["lightgreen", "darkgreen"],
+			"5": ["tomato", "darksalmon"],
+			"6": ["red", "maroon"],
+			"7": ["whitesmoke", "black"],
+			"8": ["seagreen", "darkseagreen"]
+		}
+	}
 }