feat: Add embeddings and loras autocompletion

modules/api.py

@@ -1,7 +1,10 @@
-import os
 import json
+import os
+
+import folder_paths
 import server
 from aiohttp import web
+
 from . import downloader as dl
 
 # Get the absolute path to the 'data' directory
@@ -227,3 +230,21 @@ async def get_last_check_time(_request):
     except (IOError, json.JSONDecodeError) as e:
         print(f"[Autocomplete-Plus] Error reading csv_meta.json: {e}")
         return web.json_response({"last_check_time": None, "error": str(e)}, status=500)
+
+
+@server.PromptServer.instance.routes.get("/autocomplete-plus/embeddings")
+async def get_embeddings(request):
+    """
+    Returns a list of embedding files.
+    """
+    embeddings = folder_paths.get_filename_list("embeddings")
+    return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings)))
+
+
+@server.PromptServer.instance.routes.get("/autocomplete-plus/loras")
+async def get_loras(request):
+    """
+    Returns a list of lora files.
+    """
+    loras = folder_paths.get_filename_list("loras")
+    return web.json_response(list(map(lambda a: os.path.splitext(a)[0], loras)))

tests/js/flexsearch.test.js

@@ -1,10 +1,10 @@
 
-import { 
+import {
     createFlexSearchDocument,
     __test__
- } from "../../web/js/searchengine.js";
+} from "../../web/js/searchengine.js";
 
-const { createTagEncoder, createCJKEncoder } = __test__;
+const { createTagEncoder, createCJKEncoder, createModelEncoder } = __test__;
 
 function parseCSVLine(line) {
     const result = [];
@@ -64,23 +64,23 @@ copyright_(series),2,1298,"copyright,コピーライト (シリーズ),コピー
 
     const ControlCSV = `
 __wildcard__,0,1000,
-<lora:my_lora1:1.0>,0,1000,
-Embedding: my_embedding,0,1000,
+<lora:my_lora1>,0,1000,
+embedding: my_embedding,0,1000,
 `;
 
     const mockCSV = [
         commonCSV, cjkAliasCSV, specialCharCSV, ControlCSV
     ].map(csv => csv.trim()).join('\n');
-    
+
     let mockTags;
 
-    let tagEncoder, cjkEncoder;
+    let tagEncoder, cjkEncoder, modelEncoder;
     let document;
 
     let performSearch = function (query, limit = 100) {
         const results = document.search(query, {
             field: ["tag", "alias"],
-            limit: limit, 
+            limit: limit,
             suggest: false,
             merge: true,
         });
@@ -98,6 +98,7 @@ Embedding: my_embedding,0,1000,
 
         tagEncoder = createTagEncoder();
         cjkEncoder = createCJKEncoder();
+        modelEncoder = createModelEncoder();
 
         document = createFlexSearchDocument();
 
@@ -134,6 +135,32 @@ Embedding: my_embedding,0,1000,
             const encoded = tagEncoder.encode('one_two_');
             expect(encoded).toEqual(['one', 'two']);
         });
+        test('should properly encode embedding notation', () => {
+            expect(
+                modelEncoder.encode('embedding:path/to/my_embed1')
+            ).toEqual(['embedding:', 'path', 'to', 'my', 'embed1']);
+
+            expect(
+                modelEncoder.encode('embedding:path\\to\\my-embed1')
+            ).toEqual(['embedding:', 'path', 'to', 'my', 'embed1']);
+
+            expect(
+                modelEncoder.encode('embedding:path\\to\\this is my embed. my-negative01 (v1)__by me')
+            ).toEqual(['embedding:', 'path', 'to', 'this', 'is', 'my', 'embed', 'my', 'negative01', 'v1', 'by', 'me']);
+        });
+        test('should properly encode lora notation', () => {
+            expect(
+                modelEncoder.encode('<lora:path/to/my_lora1>')
+            ).toEqual(['lora:', 'path', 'to', 'my', 'lora1']);
+
+            expect(
+                modelEncoder.encode('<lora:path\\to\\my-lora1>')
+            ).toEqual(['lora:', 'path', 'to', 'my', 'lora1']);
+
+            expect(
+                modelEncoder.encode('<lora:path\\to\\this is my lora. my-style01 (v1)__by me>')
+            ).toEqual(['lora:', 'path', 'to', 'this', 'is', 'my', 'lora', 'my', 'style01', 'v1', 'by', 'me']);
+        });
     });
 
     describe('Basic Search', () => {
@@ -300,7 +327,7 @@ Embedding: my_embedding,0,1000,
             const results = performSearch(tag);
             expect(results.length).toEqual(1);
 
-            expect(results).toContain("<lora:my_lora1:1.0>");
+            expect(results).toContain("<lora:my_lora1>");
         });
     });
 

web/js/autocomplete.js

@@ -1,6 +1,8 @@
 import {
+    ModelTagSource,
     TagCategory,
     TagData,
+    TagSource,
     autoCompleteData,
     getEnabledTagSourceInPriorityOrder
 } from './data.js';
@@ -9,6 +11,7 @@ import {
     hiraToKata,
     kataToHira,
     formatCountHumanReadable,
+    escapeHtml,
     isContainsLetterOrNumber,
     normalizeTagToInsert,
     normalizeTagToSearch,
@@ -89,7 +92,7 @@ function searchCompletionCandidates(textareaElement) {
     }
 
     // Generate Hiragana/Katakana variations if applicable
-    const queryVariations = new Set([partialTag, normalizeTagToSearch(partialTag)]);
+    const queryVariations = new Set([partialTag.toLocaleLowerCase(), normalizeTagToSearch(partialTag).toLocaleLowerCase()]);
     const kataQuery = hiraToKata(partialTag);
     if (kataQuery !== partialTag) {
         queryVariations.add(kataQuery);
@@ -128,15 +131,14 @@ function sequentialSearch(partialTag, queryVariations) {
             let matchedAlias = null;
 
             // Check primary tag against all variations for exact/partial match
-            const tagMatch = matchWord(tagData.tag, queryVariations);
+            const tagMatch = matchWord(tagData.tag.toLocaleLowerCase(), queryVariations);
             matched = tagMatch.matched;
             isExactMatch = tagMatch.isExactMatch;
 
             // If primary tag didn't match, check aliases against all variations
             if (!matched && tagData.alias && Array.isArray(tagData.alias) && tagData.alias.length > 0) {
                 for (const alias of tagData.alias) {
-                    const lowerAlias = alias.toLowerCase();
-                    const aliasMatch = matchWord(lowerAlias, queryVariations);
+                    const aliasMatch = matchWord(alias.toLocaleLowerCase(), queryVariations);
                     if (aliasMatch.matched) {
                         matched = true;
                         isExactMatch = aliasMatch.isExactMatch;
@@ -299,9 +301,9 @@ function getCurrentPartialTag(inputElement) {
  * Inserts the selected tag into the textarea, replacing the partial tag,
  * making the change undoable.
  * @param {HTMLTextAreaElement} inputElement
- * @param {string} tagToInsert The raw tag string to insert.
+ * @param {TagData} tagDataToInsert The raw tag string to insert.
  */
-function insertTagToTextArea(inputElement, tagToInsert) {
+function insertTagToTextArea(inputElement, tagDataToInsert) {
     const text = inputElement.value;
     const cursorPos = inputElement.selectionStart;
 
@@ -309,7 +311,13 @@ function insertTagToTextArea(inputElement, tagToInsert) {
     const replaceStart = Math.min(cursorPos, tagStart);
     let replaceEnd = cursorPos;
 
-    const normalizedTag = normalizeTagToInsert(tagToInsert);
+    let normalizedTag;
+    if(Object.values(ModelTagSource).includes(tagDataToInsert.source)){
+        // If the tag is from a model tag source, don't want to normalize it
+        normalizedTag = tagDataToInsert.tag;
+    }else{
+        normalizedTag = normalizeTagToInsert(tagDataToInsert.tag);
+    }
 
     const currentTagAfterCursor = text.substring(cursorPos, tagEnd).trimEnd();
     if (normalizedTag.lastIndexOf(currentTagAfterCursor) !== -1) {
@@ -378,7 +386,7 @@ class AutocompleteUI {
         this.tagsList.addEventListener('mousedown', (e) => {
             const row = e.target.closest('.autocomplete-plus-item');
             if (row && row.dataset.tag) {
-                this.#insertTag(row.dataset.tag);
+                this.#insertTag(row.dataset);
                 e.preventDefault(); // Prevent focus loss from input
                 e.stopPropagation();
             }
@@ -443,10 +451,12 @@ class AutocompleteUI {
         this.#highlightItem();
     }
 
-    /** Selects the currently highlighted item */
+    /** Selects the currently highlighted item
+     * @returns {TagData|null} The selected tag data.
+     */
     getSelectedTag() {
         if (this.selectedIndex >= 0 && this.selectedIndex < this.candidates.length) {
-            return this.candidates[this.selectedIndex].tag;
+            return this.candidates[this.selectedIndex];
         }
 
         return null; // No valid selection
@@ -492,9 +502,10 @@ class AutocompleteUI {
         if (settingValues.tagSourceIconPosition == 'hidden') {
             tagName.textContent = tagData.tag;
         } else {
+            const escapedTag = escapeHtml(tagData.tag);
             tagName.innerHTML = settingValues.tagSourceIconPosition == 'left'
-                ? `${tagSourceIconHtml} ${tagData.tag}`
-                : `${tagData.tag} ${tagSourceIconHtml}`;
+                ? `${tagSourceIconHtml} ${escapedTag}`
+                : `${escapedTag} ${tagSourceIconHtml}`;
         }
 
         // grayout tag name if it already exists
@@ -635,10 +646,10 @@ class AutocompleteUI {
 
     /**
      * Handles the selection of an item
-     * @param {string} selectedTag The tag to insert.
+     * @param {TagData} selectedTag The tag to insert.
      */
     #insertTag(selectedTag) {
-        if (!this.target || !selectedTag || selectedTag.length <= 0) {
+        if (!this.target || !selectedTag) {
             this.hide();
             return;
         }

web/js/data.js

@@ -1,14 +1,20 @@
 import { settingValues, updateMaxTagLength } from "./settings.js";
-import { createFlexSearchDocument } from "./searchengine.js";
+import { createFlexSearchDocument, createFlexSearchDocumentForModel } from "./searchengine.js";
 
 // --- Constants ---
 
-// Tag data sources
+// Tag sources for booru-like tag data.
 export const TagSource = {
     Danbooru: 'danbooru',
     E621: 'e621',
 }
 
+// Tag sources for model based tag data.
+export const ModelTagSource = {
+    Embeddings: 'embeddings',
+    Lora: 'lora'
+}
+
 export const TagCategory = {
     'danbooru': [
         'general',
@@ -28,6 +34,12 @@ export const TagCategory = {
         'invalid',
         'meta',
         'lore',
+    ],
+    'embeddings': [
+        'embeddings'
+    ],
+    'lora': [
+        'lora'
     ]
 }
 
@@ -40,19 +52,19 @@ export class TagData {
     /**
      * Create a tag data object
      * @param {string} tag - The tag name
-     * @param {string[]} [alias=[]] - Array of aliases for the tag
-     * @param {string} [category='general'] - Category of the tag
+     * @param {number} [category] - Category index of the tag
      * @param {number} [count=0] - Frequency count/popularity of the tag
-     * @param {string} [source=TagSources.Danbooru] - The source of the tag data
+     * @param {string[]} [alias=[]] - Array of aliases for the tag
+     * @param {string} [source=TagSource.Danbooru] - The source of the tag data
      */
-    constructor(tag, alias = [], category = 'general', count = 0, source = TagSource.Danbooru) {
+    constructor(tag, category, count = 0, alias = [], source = TagSource.Danbooru) {
         /** @type {string} */
         this.tag = tag;
 
         /** @type {string[]} */
         this.alias = alias;
 
-        /** @type {string} */
+        /** @type {number} */
         this.category = category;
 
         /** @type {number} */
@@ -98,24 +110,32 @@ export const autoCompleteData = {};
 const TAGS_CSV_HEADER = 'tag,category,count,alias';
 const TAGS_CSV_HEADER_COLUMNS = TAGS_CSV_HEADER.split(',');
 const TAG_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('tag');
-const ALIAS_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('alias');
 const CATEGORY_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('category');
 const COUNT_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('count');
+const ALIAS_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('alias');
 
 // --- Helder Functions ---
 
+
 /**
  * Get the available tag sources in priority order based on the current settings.
  * @returns {string[]} Array of available tag sources in priority order
  */
 export function getEnabledTagSourceInPriorityOrder() {
-    return Object.values(TagSource)
+    let enabledTagSources = Object.values(TagSource)
         .filter((s) => {
             return settingValues.tagSource === s || settingValues.tagSource === 'all';
         })
         .toSorted((a, b) => {
             return a === settingValues.primaryTagSource ? -1 : 1;
         });
+
+    // Append Loras and Embeddings if enabled
+    if (settingValues.enableModels) {
+        enabledTagSources = [...enabledTagSources, ...Object.values(ModelTagSource)];
+    }
+
+    return enabledTagSources;
 }
 
 // --- Data Loading Functions ---
@@ -159,7 +179,7 @@ async function loadTags(csvUrl, siteName) {
                 const aliases = aliasStr ? aliasStr.split(',').map(a => a.trim()).filter(a => a.length > 0) : [];
 
                 // Create a TagData instance instead of a plain object
-                const tagData = new TagData(tag, aliases, category, count, siteName);
+                const tagData = new TagData(tag, category, count, aliases, siteName);
 
                 updateMaxTagLength(tag.length);
 
@@ -202,7 +222,14 @@ async function buildFlexSearchIndex(siteName) {
             return;
         }
 
-        const document = createFlexSearchDocument();
+        let document = null;
+        if (Object.values(TagSource).includes(siteName)) {
+            document = createFlexSearchDocument();
+        } else if (Object.values(ModelTagSource).includes(siteName)) {
+            document = createFlexSearchDocumentForModel();
+        } else {
+            throw new Error(`[Autocomplete-Plus] Invalid site name: ${siteName}`);
+        }
 
         let startIdx = 0;
         const startTime = performance.now();
@@ -341,7 +368,11 @@ function parseCSVLine(line) {
     return result;
 }
 
-export async function fetchCsvList() {
+/**
+ * Fetch the list of CSV files from the API endpoint
+ * @returns {Promise<void>}
+ */
+async function fetchCsvList() {
     try {
         const response = await fetch('/autocomplete-plus/csv');
         if (!response.ok) {
@@ -357,9 +388,8 @@ export async function fetchCsvList() {
 
 /**
  * Initializes the autocomplete data by fetching the list of CSV files and loading them.
- * This function is called when the extension is initialized.
  */
-export async function initializeData(csvListData, source) {
+async function initializeDataFromCSV(csvListData, source) {
     if (autoCompleteData.hasOwnProperty(source) === false) {
         autoCompleteData[source] = new AutocompleteData();
     }
@@ -444,3 +474,88 @@ export async function initializeData(csvListData, source) {
         autoCompleteData[source].isInitializing = false;
     }
 }
+
+/**
+ * Load Embeddings data from the API endpoint
+ * @returns {Promise<void>}
+ */
+async function loadEmbeddings() {
+    try {
+        const response = await fetch('/autocomplete-plus/embeddings', { cache: "no-store" });
+        if (!response.ok) {
+            throw new Error(`HTTP error! status: ${response.status}`);
+        }
+        const embeddings = await response.json();
+        const source = ModelTagSource.Embeddings;
+
+        if (autoCompleteData.hasOwnProperty(source) === false) {
+            autoCompleteData[source] = new AutocompleteData();
+        }
+
+        embeddings.forEach(embedding => {
+            if (!autoCompleteData[source].tagMap.has(embedding)) {
+                const tagData = new TagData(`embedding:${embedding}`, 0, 0, [], source);
+                autoCompleteData[source].sortedTags.push(tagData);
+                autoCompleteData[source].tagMap.set(embedding, tagData);
+
+                updateMaxTagLength(embedding.length);
+            }
+        });
+
+        await buildFlexSearchIndex(ModelTagSource.Embeddings);
+
+        console.log(`[Autocomplete-Plus] Loaded ${embeddings.length} Embeddings`);
+    } catch (error) {
+        console.error(`[Autocomplete-Plus] Failed to fetch Embeddings data:`, error);
+    }
+}
+
+/**
+ * Load LoRA data from the API endpoint
+ * @returns {Promise<void>}
+ */
+async function loadLoras() {
+    try {
+        const response = await fetch('/autocomplete-plus/loras', { cache: "no-store" });
+        if (!response.ok) {
+            throw new Error(`HTTP error! status: ${response.status}`);
+        }
+        const loraNames = await response.json();
+        const source = ModelTagSource.Lora;
+
+        if (autoCompleteData.hasOwnProperty(source) === false) {
+            autoCompleteData[source] = new AutocompleteData();
+        }
+
+        loraNames.forEach(loraName => {
+            if (!autoCompleteData[source].tagMap.has(loraName)) {
+                const tagData = new TagData(`<lora:${loraName}>`, 0, 0, [], source);
+                autoCompleteData[source].sortedTags.push(tagData);
+                autoCompleteData[source].tagMap.set(loraName, tagData);
+
+                updateMaxTagLength(loraName.length);
+            }
+        });
+
+        await buildFlexSearchIndex(ModelTagSource.Lora);
+
+        console.log(`[Autocomplete-Plus] Loaded ${loraNames.length} LoRA models`);
+    } catch (error) {
+        console.error(`[Autocomplete-Plus] Failed to fetch LoRA data:`, error);
+    }
+}
+
+/**
+ * Load all data sources asynchronously.
+ */
+export async function loadDataAsync() {
+    return Promise.all([
+        fetchCsvList().then((csvList) => {
+            Object.values(TagSource).forEach((source) => {
+                initializeDataFromCSV(csvList, source);
+            });
+        }),
+        loadEmbeddings(),
+        loadLoras(),
+    ]);
+}

web/js/main.js

@@ -3,7 +3,7 @@ import { $el } from "/scripts/ui.js";
 import { ComfyWidgets } from "/scripts/widgets.js";
 import { settingValues } from "./settings.js";
 import { loadCSS } from "./utils.js";
-import { TagSource, fetchCsvList, initializeData } from "./data.js";
+import { TagSource, loadDataAsync } from "./data.js";
 import { AutocompleteEventHandler } from "./autocomplete.js";
 import { RelatedTagsEventHandler } from "./related-tags.js";
 
@@ -240,11 +240,7 @@ app.registerExtension({
         let rootPath = import.meta.url.replace("js/main.js", "");
         loadCSS(rootPath + "css/autocomplete-plus.css"); // Load CSS for autocomplete
 
-        fetchCsvList().then((csvList) => {
-            Object.values(TagSource).forEach((source) => {
-                initializeData(csvList, source);
-            });
-        });
+        await loadDataAsync();
     },
 
     // One the Settings Screen, displays reverse order in same category
@@ -297,6 +293,17 @@ app.registerExtension({
                 settingValues.useFastSearch = newVal;
             }
         },
+        {
+            id: id + ".enable_models",
+            name: "Enable Loras and Embeddings",
+            tooltip: "Enable Lora and Embedding suggestions",
+            type: "boolean",
+            defaultValue: true,
+            category: [name, "Autocompletion", "Enable Loras and Embeddings"],
+            onChange: (newVal, oldVal) => {
+                settingValues.enableModels = newVal;
+            }
+        },
         {
             id: id + ".max_suggestions",
             name: "Max suggestions",

web/js/searchengine.js

@@ -35,6 +35,23 @@ function createCJKEncoder() {
     });
 }
 
+/**
+ * Creates an encoder optimized for processing Embedding or Lora notation.
+ * @returns {Encoder} FlexSearch encoder
+ */
+function createModelEncoder() {
+    return new Encoder({
+        normalize: true,
+        dedupe: false,
+        numeric: true,
+        cache: true,
+        prepare: function (str) {
+            return str.replace(/(lora:|embedding:)/g, "$1 ");
+        },
+        split: /[<>_./\(\)\-\s\\]+/
+    });
+}
+
 /**
  * Creates a FlexSearch Document instance optimized for tag searching.
  * Configures separate encoders for English tags and CJK aliases with appropriate tokenization.
@@ -81,6 +98,52 @@ export function createFlexSearchDocument() {
     return document;
 }
 
+/**
+ * Creates a FlexSearch Document instance optimized for lora or embedding searching.
+ * @returns {Document} Configured FlexSearch document
+ */
+export function createFlexSearchDocumentForModel() {
+    const modelEncoder = createModelEncoder();
+    const cjkEncoder = createCJKEncoder();
+
+    // Custom encoding function for alias field that handles mixed language content
+    const encodeAlias = function (term) {
+        return term.split(",")
+            .flatMap(str => {
+                if (/[^\u0000-\u007f]/.test(str)) {
+                    // Contains non-ASCII characters (CJK text)
+                    return cjkEncoder.encode(str);
+                } else {
+                    // ASCII characters only (English text)
+                    return modelEncoder.encode(str);
+                }
+            })
+            .filter(Boolean);
+    }
+
+    // Configure the FlexSearch document with optimized indexing settings
+    // Note: alias field is not indexed for lora or embedding search
+    const document = new Document({
+        document: {
+            id: "id",
+            index: [
+                {
+                    field: "tag",
+                    tokenize: "bidirectional",  // Allow partial matching from both ends
+                    encoder: modelEncoder,
+                },
+                {
+                    field: "alias",             // Index the alias field for multi-language support
+                    tokenize: "full",           // Full tokenization for complete alias matching
+                    encode: encodeAlias,        // Use custom multi-language encoding function
+                }
+            ]
+        }
+    });
+
+    return document;
+}
+
 // Export functions for testing when in test environment
 const isTestEnvironment = typeof process !== 'undefined' && process.env.NODE_ENV === 'test';
-export const __test__ = isTestEnvironment ? { createTagEncoder, createCJKEncoder } : undefined;
+export const __test__ = isTestEnvironment ? { createTagEncoder, createCJKEncoder, createModelEncoder } : undefined;

web/js/settings.js

@@ -7,6 +7,7 @@ export const settingValues = {
     // Autocomplete feature settings
     enabled: true,
     maxSuggestions: 10,
+    enableModels: true, // Enable Lora and Embedding suggestions
     useFastSearch: false,
 
     // Related tags feature settings

web/js/utils.js

@@ -109,6 +109,30 @@ export function formatCountHumanReadable(num) {
     return (num / si[i].value).toFixed(1).replace(rx, "$1") + si[i].symbol;
 }
 
+/**
+ * Escapes HTML special characters in a string.
+ * @param {string} str The input string.
+ * @returns {string} The escaped string.
+ */
+export function escapeHtml(str) {
+    if (typeof str !== 'string') {
+        return str;
+    }
+
+    const escapeMap = {
+        '&': '&',
+        '<': '&lt;',
+        '>': '&gt;',
+        '"': '&quot;',
+        "'": '&#x27;',
+        '`': '&#x60;',
+        '/': '&#x2F;'
+    };
+
+    return str.replace(/[&<>"'`/]/g, match =>
+        escapeMap[match]);
+}
+
 /**
  * Escapes parentheses in a string for use in prompts.
  * Replaces '(' with '\(' and ')' with '\)'.