Compare commits

..

15 Commits

Author SHA1 Message Date
Dominik Reh
bec222f2b3 Fix for 1-letter completion
Completion would sometimes not show if the prompt was only one letter long and identical to the previous completion
2023-01-12 15:54:57 +01:00
Dominik Reh
d4db6a7907 Option to show ? wiki links for danbooru/e621 tags
Disabled by default since the wiki pages likely contain NSFW images.
Closes #109
2023-01-12 15:49:53 +01:00
Dominik Reh
52593e6ac8 Update setting descriptions for black/whitelist 2023-01-12 14:45:16 +01:00
Dominik Reh
849e346924 Black/whitelisting options for models
Enables selective (de)activation based on model hash.
Closes #14
2023-01-12 14:35:54 +01:00
Dominik Reh
25b285bea3 Styling adjustments 2023-01-10 15:10:13 +01:00
Dominik Reh
984a7e772a File comments 2023-01-10 15:01:22 +01:00
Dominik Reh
964b4fcff3 Rework results system
Now uses object properties instead of array indices, much less confusing
2023-01-10 14:59:09 +01:00
Dominik Reh
54641ddbfc Move utility functions to their own file 2023-01-10 14:58:25 +01:00
Dominik Reh
c048684909 Load embeds recursively in fallback
Webui now supports recursive embedding loading, so we also use it here.
This shouldn't happen since the newer version uses the non-fallback, but it doesn't hurt
2023-01-06 15:52:40 +01:00
Dominik Reh
da9acfea2a Rework embedding load, now uses callback.
Should hopefully fix #100
2023-01-03 17:30:30 +01:00
Dominik Reh
552c6517b8 Make new settings id the default behavior instead of fallback 2023-01-03 11:14:16 +01:00
DominikDoom
f626eb3467 Merge pull request #101 from stysmmaker/fix/apply-settings-button-fallback
Add fallback to applySettingsButton variable
2023-01-03 10:57:53 +01:00
MMaker
2ba513bedc fix: Add fallback to applySettingsButton var
Need due to layout change in recent webui update
269f6e8676
2023-01-03 00:07:15 -05:00
Dominik Reh
89d36da47e Add fallback for embedding loading
Fixes error on outdated webuis, as mentioned in #98 and #99
2023-01-02 16:09:36 +01:00
Dominik Reh
5f2f746310 Skipped embeddings now also hold shape info
so we don't need to guess the type anymore if the model didn't load any.
2023-01-02 12:45:56 +01:00
4 changed files with 366 additions and 198 deletions

29
javascript/_result.js Normal file
View File

@@ -0,0 +1,29 @@
// Result data type for cleaner use of optional completion result properties
// Type enum
const ResultType = Object.freeze({
"tag": 1,
"embedding": 2,
"wildcardTag": 3,
"wildcardFile": 4,
"yamlWildcard": 5
});
// Class to hold result data and annotations to make it clearer to use
class AutocompleteResult {
// Main properties
text = "";
type = ResultType.tag;
// Additional info, only used in some cases
category = null;
count = null;
aliases = null;
meta = null;
// Constructor
constructor(text, type) {
this.text = text;
this.type = type;
}
}

96
javascript/_utils.js Normal file
View File

@@ -0,0 +1,96 @@
// Utility functions for tag autocomplete
// 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
async function readFile(filePath, json = false) {
let response = await fetch(`file=${filePath}`);
if (response.status != 200) {
console.error(`Error loading file "${filePath}": ` + response.status, response.statusText);
return null;
}
if (json)
return await response.json();
else
return await response.text();
}
// Load CSV
async function loadCSV(path) {
let text = await readFile(path);
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)), []);
}
function escapeRegExp(string) {
return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string
}
function escapeHTML(unsafeText) {
let div = document.createElement('div');
div.textContent = unsafeText;
return div.innerHTML;
}

View File

@@ -7,7 +7,7 @@ const styleColors = {
"--results-bg-odd": ["#111827", "#f9fafb"],
"--results-hover": ["#1f2937", "#f5f6f8"],
"--results-selected": ["#374151", "#e5e7eb"],
"--post-count-color": ["#6b6f7b", "#a2a9b4"],
"--meta-text-color": ["#6b6f7b", "#a2a9b4"],
"--embedding-v1-color": ["lightsteelblue", "#2b5797"],
"--embedding-v2-color": ["skyblue", "#2d89ef"],
}
@@ -23,7 +23,9 @@ const autocompleteCSS = `
background-color: transparent;
min-width: fit-content;
align-self: center;
margin: 0 5px;
}
#quicksettings [id^=setting_tac] > label > span {
margin-bottom: 0px;
}
[id^=refresh_tac] {
max-width: 2.5em;
@@ -61,12 +63,19 @@ const autocompleteCSS = `
overflow: hidden;
white-space: nowrap;
}
.acPostCount {
.acMetaText {
position: relative;
text-align: end;
padding: 0 0 0 15px;
flex-grow: 1;
color: var(--post-count-color);
color: var(--meta-text-color);
}
.acWikiLink {
padding: 0.5rem;
margin: -0.5rem 0 -0.5rem -0.5rem;
}
.acWikiLink:hover {
text-decoration: underline;
}
.acListItem.acEmbeddingV1 {
color: var(--embedding-v1-color);
@@ -76,64 +85,6 @@ const autocompleteCSS = `
}
`;
// 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
async function readFile(filePath, json = false) {
let response = await fetch(`file=${filePath}`);
if (response.status != 200) {
console.error(`Error loading file "${filePath}": ` + response.status, response.statusText);
return null;
}
if (json)
return await response.json();
else
return await response.text();
}
// Load CSV
async function loadCSV(path) {
let text = await readFile(path);
return parseCSV(text);
}
var tagBasePath = "";
var allTags = [];
var translations = new Map();
@@ -214,7 +165,9 @@ async function syncOptions() {
txt2img: opts["tac_activeIn.txt2img"],
img2img: opts["tac_activeIn.img2img"],
negativePrompts: opts["tac_activeIn.negativePrompts"],
thirdParty: opts["tac_activeIn.thirdParty"]
thirdParty: opts["tac_activeIn.thirdParty"],
modelList: opts["tac_activeIn.modelList"],
modelListMode: opts["tac_activeIn.modelListMode"]
},
// Results related settings
maxResults: opts["tac_maxResults"],
@@ -223,6 +176,7 @@ async function syncOptions() {
delayTime: opts["tac_delayTime"],
useWildcards: opts["tac_useWildcards"],
useEmbeddings: opts["tac_useEmbeddings"],
showWikiLinks: opts["tac_showWikiLinks"],
// Insertion related settings
replaceUnderscores: opts["tac_replaceUnderscores"],
escapeParentheses: opts["tac_escapeParentheses"],
@@ -274,34 +228,6 @@ async function syncOptions() {
CFG = newCFG;
}
// 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(textArea) {
let resultsDiv = document.createElement("div");
@@ -340,13 +266,26 @@ function hideResults(textArea) {
selectedTag = null;
}
function escapeRegExp(string) {
return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string
}
function escapeHTML(unsafeText) {
let div = document.createElement('div');
div.textContent = unsafeText;
return div.innerHTML;
var currentModelHash = "";
// Function to check activation criteria
function isEnabled() {
if (CFG.activeIn.global) {
let modelList = CFG.activeIn.modelList
.split(",")
.map(x => x.trim())
.filter(x => x.length > 0);
if (CFG.activeIn.modelListMode.toLowerCase() === "blacklist") {
// If the current model is in the blacklist, disable
return !modelList.includes(currentModelHash);
} else {
// If the current model is in the whitelist, enable.
// An empty whitelist is ignored.
return modelList.length === 0 || modelList.includes(currentModelHash);
}
} else {
return false;
}
}
const WEIGHT_REGEX = /[([]([^,()[\]:| ]+)(?::(?:\d+(?:\.\d+)?|\.\d+))?[)\]]/g;
@@ -354,24 +293,25 @@ const TAG_REGEX = /(<[^\t\n\r,>]+>?|[^\s,|<>]+|<)/g
const WC_REGEX = /\b__([^, ]+)__([^, ]*)\b/g;
const UMI_PROMPT_REGEX = /<[^\s]*?\[[^,<>]*[\]|]?>?/gi;
const UMI_TAG_REGEX = /(?:\[|\||--)([^<>\[\]\-|]+)/gi;
const MODEL_HASH_REGEX = /\[(.+)\]/g;
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 text = result.text;
let tagType = result.type;
let cursorPos = textArea.selectionStart;
var sanitizedText = text
// Replace differently depending on if it's a tag or wildcard
if (tagType === "wildcardFile") {
if (tagType === ResultType.wildcardFile) {
sanitizedText = "__" + text.replace("Wildcards: ", "") + "__";
} else if (tagType === "wildcardTag") {
} else if (tagType === ResultType.wildcardTag) {
sanitizedText = text.replace(/^.*?: /g, "");
} else if (tagType === "yamlWildcard" && !yamlWildcards.includes(text)) {
} else if (tagType === ResultType.yamlWildcard && !yamlWildcards.includes(text)) {
sanitizedText = text.replaceAll("_", " "); // Replace underscores only if the yaml tag is not using them
} else if (tagType === "embedding") {
} else if (tagType === ResultType.embedding) {
sanitizedText = `${text.replace(/^.*?: /g, "")}`;
} else {
sanitizedText = CFG.replaceUnderscores ? text.replaceAll("_", " ") : text;
@@ -395,7 +335,7 @@ function insertTextAtCursor(textArea, result, tagword) {
let afterInsertCursorPos = editStart + match.index + sanitizedText.length;
var optionalComma = "";
if (CFG.appendComma && tagType !== "wildcardFile" && tagType !== "yamlWildcard") {
if (CFG.appendComma && ![ResultType.wildcardFile, ResultType.yamlWildcard].includes(tagType)) {
optionalComma = surrounding.match(new RegExp(`${escapeRegExp(tagword)}[,:]`, "i")) !== null ? "" : ", ";
}
@@ -423,13 +363,7 @@ function insertTextAtCursor(textArea, result, tagword) {
previousTags = tags;
// If it was a yaml wildcard, also update the umiPreviousTags
if (tagType === "yamlWildcard" && originalTagword.length > 0) {
let editStart = Math.max(cursorPos - tagword.length, 0);
let editEnd = Math.min(cursorPos + tagword.length, originalTagword.length);
let surrounding = originalTagword.substring(editStart, editEnd);
let match = surrounding.match(new RegExp(escapeRegExp(`${tagword}`), "i"));
let insert = surrounding.replace(match, sanitizedText);
if (tagType === ResultType.yamlWildcard && originalTagword.length > 0) {
let umiSubPrompts = [...newPrompt.matchAll(UMI_PROMPT_REGEX)];
let umiTags = [];
@@ -443,7 +377,7 @@ function insertTextAtCursor(textArea, result, tagword) {
}
// Hide results after inserting
if (tagType === "wildcardFile") {
if (tagType === ResultType.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);
@@ -482,17 +416,16 @@ function addResultsToList(textArea, results, tagword, resetList) {
let itemText = document.createElement("div");
itemText.classList.add("acListItem");
flexDiv.appendChild(itemText);
let displayText = "";
// If the tag matches the tagword, we don't need to display the alias
if (result[3] && !result[0].includes(tagword)) { // Alias
let splitAliases = result[3].split(",");
if (result.aliases && !result.text.includes(tagword)) { // Alias
let splitAliases = result.aliases.split(",");
let bestAlias = splitAliases.find(a => a.toLowerCase().includes(tagword));
// search in translations if no alias matches
if (!bestAlias) {
let tagOrAlias = pair => pair[0] === result[0] || result[3].split(",").includes(pair[0]);
let tagOrAlias = pair => pair[0] === result.text || splitAliases.includes(pair[0]);
var tArray = [...translations];
if (tArray) {
var translationKey = [...translations].find(pair => tagOrAlias(pair) && pair[1].includes(tagword));
@@ -504,43 +437,68 @@ function addResultsToList(textArea, results, tagword, resetList) {
displayText = escapeHTML(bestAlias);
// Append translation for alias if it exists and is not what the user typed
if (translations.has(bestAlias) && translations.get(bestAlias) !== bestAlias && bestAlias !== result[0])
if (translations.has(bestAlias) && translations.get(bestAlias) !== bestAlias && bestAlias !== result.text)
displayText += `[${translations.get(bestAlias)}]`;
if (!CFG.alias.onlyShowAlias && result[0] !== bestAlias)
displayText += " ➝ " + result[0];
if (!CFG.alias.onlyShowAlias && result.text !== bestAlias)
displayText += " ➝ " + result.text;
} else { // No alias
displayText = escapeHTML(result[0]);
displayText = escapeHTML(result.text);
}
// Append translation for result if it exists
if (translations.has(result[0]))
displayText += `[${translations.get(result[0])}]`;
if (translations.has(result.text))
displayText += `[${translations.get(result.text)}]`;
// Print search term bolded in result
itemText.innerHTML = displayText.replace(tagword, `<b>${tagword}</b>`);
// Add wiki link if the setting is enabled and a supported tag set loaded
if (CFG.showWikiLinks && (tagFileName.toLowerCase().startsWith("danbooru") || tagFileName.toLowerCase().startsWith("e621"))) {
let wikiLink = document.createElement("a");
wikiLink.classList.add("acWikiLink");
wikiLink.innerText = "?";
let linkPart = displayText;
// Only use alias result if it is one
if (displayText.includes("➝"))
linkPart = displayText.split(" ➝ ")[1];
// Set link based on selected file
let tagFileNameLower = tagFileName.toLowerCase();
if (tagFileNameLower.startsWith("danbooru")) {
wikiLink.href = `https://danbooru.donmai.us/wiki_pages/${linkPart}`;
} else if (tagFileNameLower.startsWith("e621")) {
wikiLink.href = `https://e621.net/wiki_pages/${linkPart}`;
}
wikiLink.target = "_blank";
flexDiv.appendChild(wikiLink);
}
flexDiv.appendChild(itemText);
// Add post count & color if it's a tag
// Wildcards & Embeds have no tag type
if (!result[1].startsWith("wildcard") && result[1] !== "embedding") {
if (!result[1].startsWith("yaml")) {
// Wildcards & Embeds have no tag category
if (![ResultType.wildcardFile, ResultType.wildcardTag, ResultType.embedding].includes(result.type)) {
if (result.category) {
// Set the color of the tag
let tagType = result[1];
let cat = result.category;
let colorGroup = tagColors[tagFileName];
// Default to danbooru scheme if no matching one is found
if (!colorGroup)
colorGroup = tagColors["danbooru"];
// Set tag type to invalid if not found
if (!colorGroup[tagType])
tagType = "-1";
if (!colorGroup[cat])
cat = "-1";
itemText.style = `color: ${colorGroup[tagType][mode]};`;
flexDiv.style = `color: ${colorGroup[cat][mode]};`;
}
// Post count
if (result[2] && !isNaN(result[2])) {
let postCount = result[2];
if (result.count && !isNaN(result.count)) {
let postCount = result.count;
let formatter;
// Danbooru formats numbers with a padded fraction for 1M or 1k, but not for 10/100k
@@ -553,20 +511,23 @@ function addResultsToList(textArea, results, tagword, resetList) {
let countDiv = document.createElement("div");
countDiv.textContent = formattedCount;
countDiv.classList.add("acPostCount");
countDiv.classList.add("acMetaText");
flexDiv.appendChild(countDiv);
}
} else if (result[1] === "embedding" && result[2]) { // Check if it is an embedding we have version info for
let versionDiv = document.createElement("div");
versionDiv.textContent = result[2];
versionDiv.classList.add("acPostCount");
} else if (result.meta) { // Check if it is an embedding we have version info for
let metaDiv = document.createElement("div");
metaDiv.textContent = result.meta;
metaDiv.classList.add("acMetaText");
if (result[2].startsWith("v1"))
itemText.classList.add("acEmbeddingV1");
else if (result[2].startsWith("v2"))
itemText.classList.add("acEmbeddingV2");
// Add version info classes if it is an embedding
if (result.type === ResultType.embedding) {
if (result.meta.startsWith("v1"))
itemText.classList.add("acEmbeddingV1");
else if (result.meta.startsWith("v2"))
itemText.classList.add("acEmbeddingV2");
}
flexDiv.appendChild(versionDiv);
flexDiv.appendChild(metaDiv);
}
// Add listener
@@ -610,11 +571,13 @@ var originalTagword = "";
var resultCount = 0;
async function autocomplete(textArea, prompt, fixedTag = null) {
// Return if the function is deactivated in the UI
if (!CFG.activeIn.global) return;
if (!isEnabled()) return;
// Guard for empty prompt
if (prompt.length === 0) {
hideResults(textArea);
previousTags = [];
tagword = "";
return;
}
@@ -629,6 +592,14 @@ async function autocomplete(textArea, prompt, fixedTag = null) {
.concat(weightedTags);
}
// Guard for no tags
if (!tags || tags.length === 0) {
previousTags = [];
tagword = "";
hideResults(textArea);
return;
}
let tagCountChange = tags.length - previousTags.length;
let diff = difference(tags, previousTags);
previousTags = tags;
@@ -650,6 +621,7 @@ async function autocomplete(textArea, prompt, fixedTag = null) {
tagword = fixedTag;
}
results = [];
tagword = tagword.toLowerCase().replace(/[\n\r]/g, "");
if (CFG.useWildcards && [...tagword.matchAll(WC_REGEX)].length > 0) {
@@ -669,8 +641,13 @@ async function autocomplete(textArea, prompt, fixedTag = null) {
let wildcards = (await readFile(`${wcPair[0]}/${wcPair[1]}.txt?${new Date().getTime()}`)).split("\n")
.filter(x => x.trim().length > 0 && !x.startsWith('#')); // Remove empty lines and comments
results = wildcards.filter(x => (wcWord !== null && wcWord.length > 0) ? x.toLowerCase().includes(wcWord) : x) // Filter by tagword
.map(x => [wcFile + ": " + x.trim(), "wildcardTag"]); // Mark as wildcard
let tempResults = wildcards.filter(x => (wcWord !== null && wcWord.length > 0) ? x.toLowerCase().includes(wcWord) : x) // Filter by tagword
tempResults.forEach(t => {
let result = new AutocompleteResult(t.trim(), ResultType.wildcardTag);
result.meta = wcFile;
results.push(result);
});
} else if (CFG.useWildcards && (tagword.startsWith("__") && !tagword.endsWith("__") || tagword === "__")) {
// Show available wildcard files
let tempResults = [];
@@ -680,7 +657,13 @@ async function autocomplete(textArea, prompt, fixedTag = null) {
} else {
tempResults = wildcardFiles.concat(wildcardExtFiles);
}
results = tempResults.map(x => ["Wildcards: " + x[1].trim(), "wildcardFile"]); // Mark as wildcard
// Add final results
tempResults.forEach(wcFile => {
let result = new AutocompleteResult(wcFile[1].trim(), ResultType.wildcardFile);
result.meta = "Wildcard file";
results.push(result);
})
} else if (CFG.useWildcards && [...tagword.matchAll(UMI_PROMPT_REGEX)].length > 0) {
// We are in a UMI yaml tag definition, parse further
let umiSubPrompts = [...prompt.matchAll(UMI_PROMPT_REGEX)];
@@ -812,16 +795,36 @@ async function autocomplete(textArea, prompt, fixedTag = null) {
let baseFilter = x => x[0].toLowerCase().search(searchRegex) > -1;
let spaceIncludeFilter = x => x[0].toLowerCase().replaceAll(" ", "_").search(searchRegex) > -1;
tempResults = filteredWildcardsSorted.filter(x => baseFilter(x) || spaceIncludeFilter(x)) // Filter by tagword
results = tempResults.map(x => [x[0].trim(), "yamlWildcard", x[1]]); // Mark as yaml wildcard
// Add final results
tempResults.forEach(t => {
let result = new AutocompleteResult(t[0].trim(), ResultType.yamlWildcard)
result.count = t[1];
results.push(result);
});
} else if (showAll) {
let filteredWildcardsSorted = filteredWildcards("");
results = filteredWildcardsSorted.map(x => [x[0].trim(), "yamlWildcard", x[1]]); // Mark as yaml wildcard
// Add final results
filteredWildcardsSorted.forEach(t => {
let result = new AutocompleteResult(t[0].trim(), ResultType.yamlWildcard)
result.count = t[1];
results.push(result);
});
originalTagword = tagword;
tagword = "";
}
} else {
let filteredWildcardsSorted = filteredWildcards("");
results = filteredWildcardsSorted.map(x => [x[0].trim(), "yamlWildcard", x[1]]); // Mark as yaml wildcard
// Add final results
filteredWildcardsSorted.forEach(t => {
let result = new AutocompleteResult(t[0].trim(), ResultType.yamlWildcard)
result.count = t[1];
results.push(result);
});
originalTagword = tagword;
tagword = "";
}
@@ -851,8 +854,21 @@ async function autocomplete(textArea, prompt, fixedTag = null) {
} else {
searchRegex = new RegExp(`(^|[^a-zA-Z])${escapeRegExp(tagword)}`, 'i');
}
genericResults = allTags.filter(x => x[0].toLowerCase().search(searchRegex) > -1).slice(0, CFG.maxResults);
results = tempResults.map(x => [x[0].trim(), "embedding", x[1] + " Embedding"]).concat(genericResults); // Mark as embedding
let genericResults = allTags.filter(x => x[0].toLowerCase().search(searchRegex) > -1).slice(0, CFG.maxResults);
// Add final results
tempResults.forEach(t => {
let result = new AutocompleteResult(t[0].trim(), ResultType.embedding)
result.meta = t[1] + " Embedding";
results.push(result);
});
genericResults.forEach(g => {
let result = new AutocompleteResult(g[0].trim(), ResultType.tag)
result.category = g[1];
result.count = g[2];
result.aliases = g[3];
results.push(result);
});
} else {
// Create escaped search regex with support for * as a start placeholder
let searchRegex;
@@ -882,7 +898,14 @@ async function autocomplete(textArea, prompt, fixedTag = null) {
else
fil = (x) => baseFilter(x);
results = allTags.filter(fil);
// Add final results
allTags.filter(fil).forEach(t => {
let result = new AutocompleteResult(t[0].trim(), ResultType.tag)
result.category = t[1];
result.count = t[2];
result.aliases = t[3];
results.push(result);
});
}
// Slice if the user has set a max result count
if (!CFG.showAllResults) {
@@ -903,8 +926,8 @@ async function autocomplete(textArea, prompt, fixedTag = null) {
var oldSelectedTag = null;
function navigateInList(textArea, event) {
// Return if the function is deactivated in the UI
if (!CFG.activeIn.global) return;
// Return if the function is deactivated in the UI or the current model is excluded due to white/blacklist settings
if (!isEnabled()) return;
validKeys = ["ArrowUp", "ArrowDown", "PageUp", "PageDown", "Home", "End", "Enter", "Tab", "Escape"];
@@ -1059,7 +1082,7 @@ async function setup() {
let textAreas = getTextAreas();
// Add event listener to apply settings button so we can mirror the changes to our internal config
let applySettingsButton = gradioApp().querySelector("#tab_settings > div > .gr-button-primary");
let applySettingsButton = gradioApp().querySelector("#tab_settings #settings_submit") || gradioApp().querySelector("#tab_settings > div > .gr-button-primary");
applySettingsButton.addEventListener("click", () => {
// Wait 500ms to make sure the settings have been applied to the webui opts object
setTimeout(async () => {
@@ -1076,6 +1099,14 @@ async function setup() {
}, 500);
});
});
// Add change listener to model dropdown to react to model changes
let modelDropdown = gradioApp().querySelector("#setting_sd_model_checkpoint select");
currentModelHash = [...modelDropdown.value.matchAll(MODEL_HASH_REGEX)][0][1]; // Set initial model hash
modelDropdown.addEventListener("change", () => {
setTimeout(() => {
currentModelHash = [...modelDropdown.value.matchAll(MODEL_HASH_REGEX)][0][1];
}, 100);
});
// Not found, we're on a page without prompt textareas
if (textAreas.every(v => v === null || v === undefined)) return;
@@ -1158,4 +1189,4 @@ onUiUpdate(async () => {
await syncOptions();
// Rest of setup
setup();
});
});

View File

@@ -81,51 +81,59 @@ def get_ext_wildcard_tags():
return output
def get_embeddings():
def get_embeddings(sd_model):
"""Write a list of all embeddings with their version"""
# Get a list of all embeddings in the folder
embs_in_dir = [str(e.relative_to(EMB_PATH)) for e in EMB_PATH.glob("**/*") if e.suffix in {".bin", ".pt", ".png",'.webp', '.jxl', '.avif'}]
# Remove file extensions
embs_in_dir = [e[:e.rfind('.')] for e in embs_in_dir]
# Wait for all embeddings to be loaded
while len(sd_hijack.model_hijack.embedding_db.word_embeddings) + len(sd_hijack.model_hijack.embedding_db.skipped_embeddings) < len(embs_in_dir):
time.sleep(2) # Sleep for 2 seconds
# Get embedding dict from sd_hijack to separate v1/v2 embeddings
emb_type_a = sd_hijack.model_hijack.embedding_db.word_embeddings
emb_type_b = sd_hijack.model_hijack.embedding_db.skipped_embeddings
# Get the shape of the first item in the dict
emb_a_shape = -1
if (len(emb_type_a) > 0):
emb_a_shape = next(iter(emb_type_a.items()))[1].shape
# Add embeddings to the correct list
# Version constants
V1_SHAPE = 768
V2_SHAPE = 1024
emb_v1 = []
emb_v2 = []
if (emb_a_shape == V1_SHAPE):
emb_v1 = list(emb_type_a.keys())
emb_v2 = list(emb_type_b)
elif (emb_a_shape == V2_SHAPE):
emb_v1 = list(emb_type_b)
emb_v2 = list(emb_type_a.keys())
# Create a new list to store the modified strings
results = []
# Iterate through each string in the big list
for string in embs_in_dir:
if string in emb_v1:
results.append(string + ",v1")
elif string in emb_v2:
results.append(string + ",v2")
# If the string is not in either, default to v1
# (we can't know what it is since the startup model loaded none of them, but it's probably v1 since v2 is newer)
else:
results.append(string + ",v1")
try:
# Get embedding dict from sd_hijack to separate v1/v2 embeddings
emb_type_a = sd_hijack.model_hijack.embedding_db.word_embeddings
emb_type_b = sd_hijack.model_hijack.embedding_db.skipped_embeddings
# Get the shape of the first item in the dict
emb_a_shape = -1
emb_b_shape = -1
if (len(emb_type_a) > 0):
emb_a_shape = next(iter(emb_type_a.items()))[1].shape
if (len(emb_type_b) > 0):
emb_b_shape = next(iter(emb_type_b.items()))[1].shape
# Add embeddings to the correct list
if (emb_a_shape == V1_SHAPE):
emb_v1 = list(emb_type_a.keys())
elif (emb_a_shape == V2_SHAPE):
emb_v2 = list(emb_type_a.keys())
if (emb_b_shape == V1_SHAPE):
emb_v1 = list(emb_type_b.keys())
elif (emb_b_shape == V2_SHAPE):
emb_v2 = list(emb_type_b.keys())
# Get shape of current model
#vec = sd_model.cond_stage_model.encode_embedding_init_text(",", 1)
#model_shape = vec.shape[1]
# Show relevant entries at the top
#if (model_shape == V1_SHAPE):
# results = [e + ",v1" for e in emb_v1] + [e + ",v2" for e in emb_v2]
#elif (model_shape == V2_SHAPE):
# results = [e + ",v2" for e in emb_v2] + [e + ",v1" for e in emb_v1]
#else:
# raise AttributeError # Fallback to old method
results = sorted([e + ",v1" for e in emb_v1] + [e + ",v2" for e in emb_v2], key=lambda x: x.lower())
except AttributeError:
print("tag_autocomplete_helper: Old webui version or unrecognized model shape, using fallback for embedding completion.")
# Get a list of all embeddings in the folder
all_embeds = [str(e.relative_to(EMB_PATH)) for e in EMB_PATH.rglob("*") if e.suffix in {".bin", ".pt", ".png",'.webp', '.jxl', '.avif'}]
# Remove files with a size of 0
all_embeds = [e for e in all_embeds if EMB_PATH.joinpath(e).stat().st_size > 0]
# Remove file extensions
all_embeds = [e[:e.rfind('.')] for e in all_embeds]
results = [e + "," for e in all_embeds]
write_to_temp_file('emb.txt', results)
@@ -170,7 +178,9 @@ if not TEMP_PATH.exists():
write_to_temp_file('wc.txt', [])
write_to_temp_file('wce.txt', [])
write_to_temp_file('wcet.txt', [])
write_to_temp_file('emb.txt', [])
# Only reload embeddings if the file doesn't exist, since they are already re-written on model load
if not TEMP_PATH.joinpath("emb.txt").exists():
write_to_temp_file('emb.txt', [])
# Write wildcards to wc.txt if found
if WILDCARD_PATH.exists():
@@ -190,9 +200,8 @@ if WILDCARD_EXT_PATHS is not None:
# Write embeddings to emb.txt if found
if EMB_PATH.exists():
# We need to load the embeddings in a separate thread since we wait for them to be checked (after the model loads)
thread = threading.Thread(target=get_embeddings)
thread.start()
# Get embeddings after the model loaded callback
script_callbacks.on_model_loaded(get_embeddings)
# Register autocomplete options
@@ -206,6 +215,8 @@ def on_ui_settings():
shared.opts.add_option("tac_activeIn.img2img", shared.OptionInfo(True, "Active in img2img (Requires restart)", section=TAC_SECTION))
shared.opts.add_option("tac_activeIn.negativePrompts", shared.OptionInfo(True, "Active in negative prompts (Requires restart)", section=TAC_SECTION))
shared.opts.add_option("tac_activeIn.thirdParty", shared.OptionInfo(True, "Active in third party textboxes [Dataset Tag Editor] (Requires restart)", section=TAC_SECTION))
shared.opts.add_option("tac_activeIn.modelList", shared.OptionInfo("", "List of model hashes to use as black/whitelist, separated by commas.", section=TAC_SECTION))
shared.opts.add_option("tac_activeIn.modelListMode", shared.OptionInfo("Blacklist", "Mode to use for model hash list", gr.Dropdown, lambda: {"choices": ["Blacklist","Whitelist"]}, section=TAC_SECTION))
# Results related settings
shared.opts.add_option("tac_maxResults", shared.OptionInfo(5, "Maximum results", section=TAC_SECTION))
shared.opts.add_option("tac_showAllResults", shared.OptionInfo(False, "Show all results", section=TAC_SECTION))
@@ -213,6 +224,7 @@ def on_ui_settings():
shared.opts.add_option("tac_delayTime", shared.OptionInfo(100, "Time in ms to wait before triggering completion again (Requires restart)", section=TAC_SECTION))
shared.opts.add_option("tac_useWildcards", shared.OptionInfo(True, "Search for wildcards", section=TAC_SECTION))
shared.opts.add_option("tac_useEmbeddings", shared.OptionInfo(True, "Search for embeddings", section=TAC_SECTION))
shared.opts.add_option("tac_showWikiLinks", shared.OptionInfo(False, "Show '?' next to tags, linking to its Danbooru or e621 wiki page (Warning: This is an external site and very likely contains NSFW examples!)", section=TAC_SECTION))
# Insertion related settings
shared.opts.add_option("tac_replaceUnderscores", shared.OptionInfo(True, "Replace underscores with spaces on insertion", section=TAC_SECTION))
shared.opts.add_option("tac_escapeParentheses", shared.OptionInfo(True, "Escape parentheses on insertion", section=TAC_SECTION))