diff --git a/modules/api.py b/modules/api.py index c200359..588de44 100644 --- a/modules/api.py +++ b/modules/api.py @@ -8,43 +8,64 @@ from aiohttp import web # os.path.join(..., '..', 'data') goes up one level and then into 'data' DATA_DIR = os.path.join(os.path.dirname(__file__), '..', 'data') -BASE_URL = '/autocomplete-plus/csv' +DANBOORU_PREFIX = 'danbooru' +E621_PREFIX = 'e621' -TAGS_BASE_FILE = 'danbooru_tags.csv' -COOCCURRENCE_BASE_FILE = 'danbooru_tags_cooccurrence.csv' +TAGS_SUFFIX = 'tags' +COOCCURRENCE_SUFFIX = 'tags_cooccurrence' def get_csv_file_status(): """ Returns a dictionary of csv file statuses. """ - tags_base_exists = os.path.exists(os.path.join(DATA_DIR, TAGS_BASE_FILE)) - cooccurrence_base_exists = os.path.exists(os.path.join(DATA_DIR, COOCCURRENCE_BASE_FILE)) - tags_extra_files = [] - cooccurrence_extra_files = [] + data = { + DANBOORU_PREFIX: { + 'base_tags': False, + 'extra_tags': [], + 'base_cooccurrence': False, + 'extra_cooccurrence': [], + }, + E621_PREFIX: { + 'base_tags': False, + 'extra_tags': [], + 'base_cooccurrence': False, + 'extra_cooccurrence': [], + } + } - all_csv_files = [f for f in os.listdir(DATA_DIR) if f.endswith('.csv')] - if len(all_csv_files) == 0: - print("[Autocomplete-Plus] No CSV files found in the data directory.") + for prefix in [DANBOORU_PREFIX, E621_PREFIX]: + base_tags_file = f"{prefix}_{TAGS_SUFFIX}.csv" + base_cooccurrence_file = f"{prefix}_{COOCCURRENCE_SUFFIX}.csv" - # Create extra CSV files list - for filename in all_csv_files: - if filename in [TAGS_BASE_FILE, COOCCURRENCE_BASE_FILE]: - continue # Skip base files - if 'cooccurrence' in filename.lower(): - cooccurrence_extra_files.append(filename) - elif 'tag' in filename.lower(): - tags_extra_files.append(filename) + tags_base_exists = os.path.exists(os.path.join(DATA_DIR, base_tags_file)) + cooccurrence_base_exists = os.path.exists(os.path.join(DATA_DIR, base_cooccurrence_file)) - # Return the lists of extra files - return { - 'danbooru':{ - 'base_tags': tags_base_exists, # exists + tags_extra_files = [] + cooccurrence_extra_files = [] + + all_csv_files = [f for f in os.listdir(DATA_DIR) if f.startswith(prefix) and f.endswith('.csv')] + if len(all_csv_files) == 0: + print("[Autocomplete-Plus] No CSV files found in the data directory.") + + # Create extra CSV files list + for filename in all_csv_files: + if filename in [base_tags_file, base_cooccurrence_file]: + continue # Skip base files + if COOCCURRENCE_SUFFIX in filename.lower(): + cooccurrence_extra_files.append(filename) + elif TAGS_SUFFIX in filename.lower(): + tags_extra_files.append(filename) + + data[prefix] = { + 'base_tags': tags_base_exists, 'extra_tags': tags_extra_files, 'base_cooccurrence': cooccurrence_base_exists, 'extra_cooccurrence': cooccurrence_extra_files, } - } + + # Return the lists of extra files + return data # --- API Endpoints --- @@ -55,41 +76,58 @@ async def get_csv_list(_request): base files: file exists boolean extra files: count of extra files """ - extra_csv_files = get_csv_file_status() + csv_file_status = get_csv_file_status() response = { - 'danbooru': { - 'base_tags': extra_csv_files['danbooru']['base_tags'], - 'extra_tags': len(extra_csv_files['danbooru']['extra_tags']), - 'base_cooccurrence': extra_csv_files['danbooru']['base_cooccurrence'], - 'extra_cooccurrence': len(extra_csv_files['danbooru']['extra_cooccurrence']), + DANBOORU_PREFIX: { + 'base_tags': csv_file_status[DANBOORU_PREFIX]['base_tags'], + 'extra_tags': len(csv_file_status[DANBOORU_PREFIX]['extra_tags']), + 'base_cooccurrence': csv_file_status[DANBOORU_PREFIX]['base_cooccurrence'], + 'extra_cooccurrence': len(csv_file_status[DANBOORU_PREFIX]['extra_cooccurrence']), + }, + E621_PREFIX: { + 'base_tags': csv_file_status[E621_PREFIX]['base_tags'], + 'extra_tags': len(csv_file_status[E621_PREFIX]['extra_tags']), + 'base_cooccurrence': csv_file_status[E621_PREFIX]['base_cooccurrence'], + 'extra_cooccurrence': len(csv_file_status[E621_PREFIX]['extra_cooccurrence']), } } return web.json_response(response) -@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/tags/base') -async def get_base_tags_file(_request): +@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/{source}/{suffix}/base') +async def get_base_tags_file(request): """ Returns the base tags CSV file. """ - file_path = os.path.join(DATA_DIR, TAGS_BASE_FILE) + source = str(request.match_info['source']) + suffix = str(request.match_info['suffix']) + if source not in [DANBOORU_PREFIX, E621_PREFIX] or suffix not in [TAGS_SUFFIX, COOCCURRENCE_SUFFIX]: + return web.json_response({"error": "Invalid tag source or suffix"}, status=400) + + file_path = os.path.join(DATA_DIR, f"{source}_{suffix}.csv") if not os.path.exists(file_path): return web.json_response({"error": "Base tags file not found"}, status=404) + return web.FileResponse(file_path) -@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/tags/extra/{index}') +@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/{source}/{suffix}/extra/{index}') async def get_extra_tags_file(request): """ Returns the extra tags CSV file at the specified index. """ try: - extra_csv_files = get_csv_file_status() + csv_file_status = get_csv_file_status() + source = str(request.match_info['source']) + suffix = str(request.match_info['suffix']) + if source not in [DANBOORU_PREFIX, E621_PREFIX] or suffix not in [TAGS_SUFFIX, COOCCURRENCE_SUFFIX]: + return web.json_response({"error": "Invalid tag source or suffix"}, status=400) + index = int(request.match_info['index']) - if index < 0 or index >= len(extra_csv_files['danbooru']['extra_tags']): + if index < 0 or index >= len(csv_file_status[source][f'extra_{suffix}']): return web.json_response({"error": "Invalid index"}, status=404) - file_path = os.path.join(DATA_DIR, extra_csv_files['danbooru']['extra_tags'][index]) + file_path = os.path.join(DATA_DIR, csv_file_status[source][f'extra_{suffix}'][index]) if not os.path.exists(file_path): return web.json_response({"error": "Extra tags file not found"}, status=404) @@ -97,34 +135,3 @@ async def get_extra_tags_file(request): except ValueError: return web.json_response({"error": "Invalid index format"}, status=400) - -@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/cooccurrence/base') -async def get_base_cooccurrence_file(_request): - """ - Returns the base cooccurrence CSV file. - """ - file_path = os.path.join(DATA_DIR, COOCCURRENCE_BASE_FILE) - if not os.path.exists(file_path): - return web.json_response({"error": "Base cooccurrence file not found"}, status=404) - return web.FileResponse(file_path) - -@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/cooccurrence/extra/{index}') -async def get_extra_cooccurrence_file(request): - """ - Returns the extra cooccurrence CSV file at the specified index. - """ - try: - extra_csv_files = get_csv_file_status() - - index = int(request.match_info['index']) - if index < 0 or index >= len(extra_csv_files['cooccurrence']): - return web.json_response({"error": "Invalid index"}, status=404) - - file_path = os.path.join(DATA_DIR, extra_csv_files['danbooru']['extra_cooccurrence'][index]) - if not os.path.exists(file_path): - return web.json_response({"error": "Extra cooccurrence file not found"}, status=404) - - return web.FileResponse(file_path) - - except ValueError: - return web.json_response({"error": "Invalid index format"}, status=400) diff --git a/web/js/autocomplete.js b/web/js/autocomplete.js index 50b739b..6750f33 100644 --- a/web/js/autocomplete.js +++ b/web/js/autocomplete.js @@ -1,7 +1,8 @@ import { - TagCategory, + DanbooruTagCategory, TagData, - autoCompleteData + autoCompleteData, + getTagSourceInPriorityOrder } from './data.js'; import { formatCountHumanReadable, @@ -99,54 +100,59 @@ function searchCompletionCandidates(textareaElement) { queryVariations.add(hiraQuery); } - // Search in sortedTags (already sorted by count) - for (const tagData of autoCompleteData.sortedTags) { - let matched = false; - let isExactMatch = false; - let matchedAlias = null; + const sources = getTagSourceInPriorityOrder(); + for (const source of sources) { + // Search in sortedTags (already sorted by count) + for (const tagData of autoCompleteData[source].sortedTags) { + let matched = false; + let isExactMatch = false; + let matchedAlias = null; - // Check primary tag against all variations for exact/partial match - const tagMatch = matchWord(tagData.tag, queryVariations); - matched = tagMatch.matched; - isExactMatch = tagMatch.isExactMatch; + // Check primary tag against all variations for exact/partial match + const tagMatch = matchWord(tagData.tag, 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); - if (aliasMatch.matched) { - matched = true; - isExactMatch = aliasMatch.isExactMatch; - matchedAlias = alias; - break; + // 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); + if (aliasMatch.matched) { + matched = true; + isExactMatch = aliasMatch.isExactMatch; + matchedAlias = alias; + break; + } } } - } - // Add candidate if matched and not already added - if (matched && !addedTags.has(tagData.tag)) { - // Add to exact matches or partial matches based on match type - if (isExactMatch) { - exactMatches.push(tagData); - } else { - partialMatches.push(tagData); - } + const tagSetKey = tagData.tag; - addedTags.add(tagData.tag); - - // Check if we've reached the maximum suggestions limit combining both arrays - if (exactMatches.length + partialMatches.length >= settingValues.maxSuggestions) { - // Return the combined results, prioritizing exact matches - const result = [...exactMatches, ...partialMatches].slice(0, settingValues.maxSuggestions); - - if (settingValues._logprocessingTime) { - const endTime = performance.now(); - const duration = endTime - startTime; - console.debug(`[Autocomplete-Plus] Search for "${partialTag}" took ${duration.toFixed(2)}ms. Found ${result.length} candidates (max reached).`); + // Add candidate if matched and not already added + if (matched && !addedTags.has(tagSetKey)) { + // Add to exact matches or partial matches based on match type + if (isExactMatch) { + exactMatches.push(tagData); + } else { + partialMatches.push(tagData); } - return result; // Early exit + addedTags.add(tagSetKey); + + // Check if we've reached the maximum suggestions limit combining both arrays + if (exactMatches.length + partialMatches.length >= settingValues.maxSuggestions) { + // Return the combined results, prioritizing exact matches + const result = [...exactMatches, ...partialMatches].slice(0, settingValues.maxSuggestions); + + if (settingValues._logprocessingTime) { + const endTime = performance.now(); + const duration = endTime - startTime; + console.debug(`[Autocomplete-Plus] Search for "${partialTag}" took ${duration.toFixed(2)}ms. Found ${result.length} candidates (max reached).`); + } + + return result; // Early exit + } } } } @@ -382,7 +388,7 @@ class AutocompleteUI { * @param {boolean} isExisting */ #createTagElement(tagData, isExisting) { - const categoryText = TagCategory[tagData.category] || "unknown"; + const categoryText = DanbooruTagCategory[tagData.category] || "unknown"; const tagRow = document.createElement('div'); tagRow.classList.add('autocomplete-plus-item'); @@ -392,7 +398,7 @@ class AutocompleteUI { // Tag name const tagName = document.createElement('span'); tagName.classList.add('autocomplete-plus-tag-name'); - tagName.textContent = tagData.tag; + tagName.textContent = `[${tagData.source[0]}] ${tagData.tag}`; // grayout tag name if it already exists if (isExisting) { diff --git a/web/js/data.js b/web/js/data.js index 3e337e4..49d749d 100644 --- a/web/js/data.js +++ b/web/js/data.js @@ -2,8 +2,14 @@ import { settingValues } from "./settings.js"; // --- Constants --- -// Tag categories for display -export const TagCategory = [ +// Tag data sources +export const TagSource = { + Danbooru: 'danbooru', + E621: 'e621', +} + +// Danbooru tag categories for display +export const DanbooruTagCategory = [ 'general', 'artist', 'unused', @@ -12,6 +18,20 @@ export const TagCategory = [ 'meta', ] +export const E621TagCategory = [ + 'general', + 'artist', + 'unused', + 'copyright', + 'character', + 'species', + 'invalid', + 'meta', + 'lore', +] + +// --- Data Structures --- + /** * Class representing a tag and its metadata */ @@ -22,45 +42,54 @@ export class TagData { * @param {string[]} [alias=[]] - Array of aliases for the tag * @param {string} [category='general'] - Category of the tag * @param {number} [count=0] - Frequency count/popularity of the tag + * @param {string} [source=TagSources.Danbooru] - The source of the tag data */ - constructor(tag, alias = [], category = 'general', count = 0) { + constructor(tag, alias = [], category = 'general', count = 0, source = TagSource.Danbooru) { /** @type {string} */ this.tag = tag; - + /** @type {string[]} */ this.alias = alias; - + /** @type {string} */ this.category = category; - + /** @type {number} */ this.count = count; + + this.source = source; } } -// Data storage -export const autoCompleteData = { - /** @type {TagData[]} */ - sortedTags: [], +class AutocompleteData { + constructor() { + /** @type {TagData[]} */ + this.sortedTags = []; - /** @type {Map} */ - tagMap: new Map(), // Stores tag data, mapping tag names to TagData objects - - /** @type {Map} */ - aliasMap: new Map(), // Maps aliases to their main tag names - - /** @type {Map>} */ - cooccurrenceMap: new Map(), // Stores co-occurrence data for related tags - - isInitializing: false, - initialized: false, + /** @type {Map} */ + this.tagMap = new Map(); - // Progress of "base" csv loading - baseLoadingProgress: { - // tags: 0, // Commented out because tags csv aren't loaded in chunks - cooccurrence: 0 + /** @type {Map} */ + this.aliasMap = new Map(); + + /** @type {Map>} */ + this.cooccurrenceMap = new Map(); + + this.isInitializing = false; + this.initialized = false; + + // Progress of "base" csv loading + this.baseLoadingProgress = { + // tags: 0, + cooccurrence: 0 + }; } -}; +} + +/** + * @type {Object} + */ +export const autoCompleteData = {}; // CSV Header for tags const TAGS_CSV_HEADER = 'tag,category,count,alias'; @@ -70,14 +99,24 @@ 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'); +// --- Helder Functions --- + +/** + * + */ +export function getTagSourceInPriorityOrder() { + return Object.values(TagSource).toSorted((a, b) => { return a === settingValues.priorityTagSource ? -1 : 1; }); +} + // --- Data Loading Functions --- /** * Loads tag data from a single CSV file. * @param {string} csvUrl - The URL of the CSV file to load. + * @param {string} siteName - The site name (e.g., 'danbooru', 'e621'). * @returns {Promise} */ -async function loadTags(csvUrl) { +async function loadTags(csvUrl, siteName) { try { const response = await fetch(csvUrl, { cache: "no-store" }); if (!response.ok) { @@ -100,9 +139,9 @@ async function loadTags(csvUrl) { const count = parseInt(columns[COUNT_INDEX].trim(), 10); if (!tag || isNaN(count)) continue; - + // Skip if tag already exists (priority to earlier loaded files - extra then base) - if (autoCompleteData.tagMap.has(tag)) { + if (autoCompleteData[siteName].tagMap.has(tag)) { continue; } @@ -110,26 +149,26 @@ async function loadTags(csvUrl) { 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); + const tagData = new TagData(tag, aliases, category, count, siteName); - autoCompleteData.sortedTags.push(tagData); + autoCompleteData[siteName].sortedTags.push(tagData); } else { console.warn(`[Autocomplete-Plus] Invalid CSV format in line ${i + 1} of ${csvUrl}: ${line}. Expected ${TAGS_CSV_HEADER_COLUMNS.length} columns, but got ${columns.length}.`); continue; } } - + // Sort by count in descending order - autoCompleteData.sortedTags.sort((a, b) => b.count - a.count); + autoCompleteData[siteName].sortedTags.sort((a, b) => b.count - a.count); // Build maps as before, but ensure not to overwrite if already processed from extra files - autoCompleteData.sortedTags.forEach(tagData => { - if (!autoCompleteData.tagMap.has(tagData.tag)) { - autoCompleteData.tagMap.set(tagData.tag, tagData); + autoCompleteData[siteName].sortedTags.forEach(tagData => { + if (!autoCompleteData[siteName].tagMap.has(tagData.tag)) { + autoCompleteData[siteName].tagMap.set(tagData.tag, tagData); if (tagData.alias && Array.isArray(tagData.alias)) { tagData.alias.forEach(alias => { - if (!autoCompleteData.aliasMap.has(alias)) { - autoCompleteData.aliasMap.set(alias, tagData.tag); // Map alias back to the main tag + if (!autoCompleteData[siteName].aliasMap.has(alias)) { + autoCompleteData[siteName].aliasMap.set(alias, tagData.tag); // Map alias back to the main tag } }); } @@ -144,9 +183,10 @@ async function loadTags(csvUrl) { /** * Loads co-occurrence data from a single CSV file. * @param {string} csvUrl - The URL of the CSV file to load. + * @param {string} siteName - The site name (e.g., 'danbooru', 'e621'). * @returns {Promise} */ -async function loadCooccurrence(csvUrl) { +async function loadCooccurrence(csvUrl, siteName) { try { const response = await fetch(csvUrl, { cache: "no-store" }); if (!response.ok) { @@ -158,7 +198,7 @@ async function loadCooccurrence(csvUrl) { const startIndex = lines[0].startsWith('tag_a,tag_b,count') ? 1 : 0; - await processInChunks(lines, startIndex, autoCompleteData.cooccurrenceMap, csvUrl); + await processInChunks(lines, startIndex, autoCompleteData[siteName].cooccurrenceMap, csvUrl, siteName); } catch (error) { console.error(`[Autocomplete-Plus] Failed to fetch or process cooccurrence data from ${csvUrl}:`, error); } @@ -168,7 +208,7 @@ async function loadCooccurrence(csvUrl) { * Process CSV data in chunks to avoid blocking the UI. * Modifies the targetMap directly. */ -function processInChunks(lines, startIndex, targetMap, sourceFileName = "CSV") { +function processInChunks(lines, startIndex, targetMap, csvUrl, siteName) { return new Promise((resolve) => { const CHUNK_SIZE = 10000; let i = startIndex; @@ -194,7 +234,7 @@ function processInChunks(lines, startIndex, targetMap, sourceFileName = "CSV") { } targetMap.get(tagA).set(tagB, count); - + // Add tagB -> tagA relationship (bidirectional) if (!targetMap.has(tagB)) { targetMap.set(tagB, new Map()); @@ -206,7 +246,7 @@ function processInChunks(lines, startIndex, targetMap, sourceFileName = "CSV") { } if (i < lines.length) { - autoCompleteData.baseLoadingProgress.cooccurrence = Math.round((i / lines.length) * 100); + autoCompleteData[siteName].baseLoadingProgress.cooccurrence = Math.round((i / lines.length) * 100); setTimeout(processChunk, 0); } else { resolve(); @@ -250,76 +290,101 @@ function parseCSVLine(line) { return result; } -/** - * 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() { - if (autoCompleteData.isInitializing || autoCompleteData.initialized) { - return; - } - - const startTime = performance.now(); - autoCompleteData.isInitializing = true; - // console.log("[Autocomplete-Plus] Initializing autocomplete data..."); - +export async function fetchCsvList() { try { const response = await fetch('/autocomplete-plus/csv'); if (!response.ok) { throw new Error(`[Autocomplete-Plus] Failed to fetch CSV list: ${response.status} ${response.statusText}`); } - const csvListData = await response.json(); + return await response.json(); + } catch (error) { + console.error("[Autocomplete-Plus] Error fetch csv data:", error); + } - const extraTagsCount = csvListData.danbooru.extra_tags || 0; - const extraCooccurrenceCount = csvListData.danbooru.extra_cooccurrence || 0; + return null; +} - const tagsUrl = '/autocomplete-plus/csv/tags'; - const cooccurrenceUrl = '/autocomplete-plus/csv/cooccurrence'; +/** + * 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) { + if (autoCompleteData.hasOwnProperty(source) === false) { + autoCompleteData[source] = new AutocompleteData(); + } - // Load extra tags first - let tagsLoadPromises = []; - let currentTagPromise = Promise.resolve(); - for (let i = 0; i < extraTagsCount; i++) { - currentTagPromise = currentTagPromise.then(() => loadTags(`${tagsUrl}/extra/${i}`)); - tagsLoadPromises.push(currentTagPromise); + if (autoCompleteData[source].isInitializing || autoCompleteData[source].initialized) { + return; + } + + const startTime = performance.now(); + autoCompleteData[source].isInitializing = true; + // console.log("[Autocomplete-Plus] Initializing autocomplete data..."); + + try { + // Store functions that return Promises (Promise Factories) + // These factories will be called later to start the actual loading. + const tagsLoadPromiseFactories = []; + const cooccurrenceLoadPromiseFactories = []; + + // Check if siteName exists in csvListData to prevent errors if a sourte is removed or misconfigured + if (!csvListData[source]) { + console.warn(`[Autocomplete-Plus] CSV list data not found for sourte: ${source}. Skipping.`); + return; } - // Then load base tags if it exists - if (csvListData.danbooru.base_tags) { - currentTagPromise = currentTagPromise.then(() => loadTags(`${tagsUrl}/base`)); - tagsLoadPromises.push(currentTagPromise); - } - - // Load extra cooccurrence first - let cooccurrenceLoadPromises = []; - let cooccurrencePromiseChain = Promise.resolve(); - for (let i = 0; i < extraCooccurrenceCount; i++) { - cooccurrencePromiseChain = cooccurrencePromiseChain.then(() => loadCooccurrence(`${cooccurrenceUrl}/extra/${i}`)); - cooccurrenceLoadPromises.push(cooccurrencePromiseChain); - } - - // Then load base cooccurrence if it exists - if (csvListData.danbooru.base_cooccurrence) { - cooccurrencePromiseChain = cooccurrencePromiseChain.then(() => loadCooccurrence(`${cooccurrenceUrl}/base`)); - cooccurrenceLoadPromises.push(cooccurrencePromiseChain); - } + const extraTagsCount = csvListData[source].extra_tags || 0; + const extraCooccurrenceCount = csvListData[source].extra_cooccurrence || 0; + const tagsUrl = `/autocomplete-plus/csv/${source}/tags`; + const cooccurrenceUrl = `/autocomplete-plus/csv/${source}/tags_cooccurrence`; + + // Factory for loading tags for the current sourte + const siteTagsLoaderFactory = async () => { + let promiseChain = Promise.resolve(); + for (let i = 0; i < extraTagsCount; i++) { + promiseChain = promiseChain.then(() => loadTags(`${tagsUrl}/extra/${i}`, source)); + } + if (csvListData[source].base_tags) { + promiseChain = promiseChain.then(() => loadTags(`${tagsUrl}/base`, source)); + } + return promiseChain; + }; + tagsLoadPromiseFactories.push(siteTagsLoaderFactory); + + // Factory for loading cooccurrence data for the current sourte + const siteCooccurrenceLoaderFactory = async () => { + let promiseChain = Promise.resolve(); + for (let i = 0; i < extraCooccurrenceCount; i++) { + promiseChain = promiseChain.then(() => loadCooccurrence(`${cooccurrenceUrl}/extra/${i}`, source)); + } + if (csvListData[source].base_cooccurrence) { + promiseChain = promiseChain.then(() => loadCooccurrence(`${cooccurrenceUrl}/base`, source)); + } + return promiseChain; + }; + cooccurrenceLoadPromiseFactories.push(siteCooccurrenceLoaderFactory); + + // Now, execute all promise factories and wait for their completion. + // The actual loading (fetch calls) will start when the factories are invoked here. await Promise.all([ - Promise.all(tagsLoadPromises).then(() => { + Promise.all(tagsLoadPromiseFactories.map(factory => factory())).then(() => { const endTime = performance.now(); - console.log(`[Autocomplete-Plus] Tags loading complete in ${(endTime - startTime).toFixed(2)}ms. Extra file count: ${extraTagsCount}`); - + console.log(`[Autocomplete-Plus] "${source}" Tags loading complete in ${(endTime - startTime).toFixed(2)}ms`); }), - Promise.all(cooccurrenceLoadPromises).then(() => { + Promise.all(cooccurrenceLoadPromiseFactories.map(factory => factory())).then(() => { const endTime = performance.now(); - console.log(`[Autocomplete-Plus] Co-occurrence loading complete in ${(endTime - startTime).toFixed(2)}ms. Extra file count: ${extraCooccurrenceCount}`); + // The original log for extraCooccurrenceCount was potentially misleading as it showed + // the count for the last processed sourte. Removed for clarity. + // If a total count is needed, it should be calculated across all sites. + console.log(`[Autocomplete-Plus] "${source}" Co-occurrence loading complete in ${(endTime - startTime).toFixed(2)}ms.`); }) ]); - autoCompleteData.initialized = true; + autoCompleteData[source].initialized = true; } catch (error) { console.error("[Autocomplete-Plus] Error initializing autocomplete data:", error); } finally { - autoCompleteData.isInitializing = false; + autoCompleteData[source].isInitializing = false; } } diff --git a/web/js/main.js b/web/js/main.js index 414d8b6..e572d00 100644 --- a/web/js/main.js +++ b/web/js/main.js @@ -2,7 +2,7 @@ import { app } from "/scripts/app.js"; import { ComfyWidgets } from "/scripts/widgets.js"; import { settingValues } from "./settings.js"; import { loadCSS } from "./utils.js"; -import { initializeData } from "./data.js"; +import { TagSource, fetchCsvList, getTagSourceInPriorityOrder, initializeData } from "./data.js"; import { AutocompleteEventHandler } from "./autocomplete.js"; import { RelatedTagsEventHandler } from "./related-tags.js"; @@ -126,17 +126,34 @@ const name = "Autocomplete Plus"; app.registerExtension({ id: id, name: name, - setup() { + async setup() { initializeEventHandlers(); let rootPath = import.meta.url.replace("js/main.js", ""); loadCSS(rootPath + "css/autocomplete-plus.css"); // Load CSS for autocomplete - initializeData(); + fetchCsvList().then((csvList) => { + getTagSourceInPriorityOrder().forEach((source) => { + initializeData(csvList, source); + }); + }); }, - + //One the Settings Screen, displays reverse order in same category - settings: [ + settings: [ + // --- General Settings --- + { + id: id + ".priority_tag_source", + name: "Priotized Tag Source", + tooltip: "If multiple tag source are available, Which tag source's tags should be displayed first.", + type: "combo", + options: Object.values(TagSource), + defaultValue: TagSource.Danbooru, + category: [name, "General", "Priority Tag Source"], + onChange: (newVal, oldVal) => { + settingValues.priorityTagSource = newVal; + } + }, // --- Autocomplete Settings --- { id: id + ".max_suggestions", @@ -164,7 +181,7 @@ app.registerExtension({ settingValues.enabled = newVal; } }, - + // --- Related Tags Settings --- { id: id + ".related_tags_trigger_mode", diff --git a/web/js/related-tags.js b/web/js/related-tags.js index 3fa0525..67e4d2f 100644 --- a/web/js/related-tags.js +++ b/web/js/related-tags.js @@ -1,4 +1,4 @@ -import { TagCategory, TagData, autoCompleteData } from './data.js'; +import { TagSource, DanbooruTagCategory, TagData, autoCompleteData, getTagSourceInPriorityOrder } from './data.js'; import { settingValues } from './settings.js'; import { extractTagsFromTextArea, @@ -15,19 +15,20 @@ import { /** * Calculates the Jaccard similarity between two tags. * Jaccard similarity = (A ∩ B) / (A ∪ B) = (A ∩ B) / (|A| + |B| - |A ∩ B|) + * @param {string} tagSource The name of the site (e.g., 'danbooru', 'e621') * @param {string} tagA The first tag * @param {string} tagB The second tag * @returns {number} Similarity score between 0 and 1 */ -function calculateJaccardSimilarity(tagA, tagB) { +function calculateJaccardSimilarity(tagSource, tagA, tagB) { // Get the count of tagA and tagB individually - const countA = autoCompleteData.tagMap.get(tagA)?.count || 0; - const countB = autoCompleteData.tagMap.get(tagB)?.count || 0; + const countA = autoCompleteData[tagSource].tagMap.get(tagA)?.count || 0; + const countB = autoCompleteData[tagSource].tagMap.get(tagB)?.count || 0; if (countA === 0 || countB === 0) return 0; // Get the cooccurrence count - const cooccurrenceAB = autoCompleteData.cooccurrenceMap.get(tagA)?.get(tagB) || 0; + const cooccurrenceAB = autoCompleteData[tagSource].cooccurrenceMap.get(tagA)?.get(tagB) || 0; // Calculate Jaccard similarity // (A ∩ B) / (A ∪ B) = (A ∩ B) / (|A| + |B| - |A ∩ B|) @@ -64,11 +65,13 @@ export function getTagFromCursorPosition(inputElement) { function searchRelatedTags(tag) { const startTime = performance.now(); // Record start time for performance measurement - if (!tag || !autoCompleteData.cooccurrenceMap.has(tag)) { + const tagSource = TagSource.Danbooru; // TODO: Leave the tag source as Danbooru until e621_tags_cooccurrence.csv is ready + + if (!tag || !autoCompleteData[tagSource].cooccurrenceMap.has(tag)) { return []; } - const cooccurrences = autoCompleteData.cooccurrenceMap.get(tag); + const cooccurrences = autoCompleteData[tagSource].cooccurrenceMap.get(tag); const relatedTags = []; // Convert to array for sorting @@ -77,11 +80,11 @@ function searchRelatedTags(tag) { if (coTag === tag) return; // Get tag data - const tagData = autoCompleteData.tagMap.get(coTag); + const tagData = autoCompleteData[tagSource].tagMap.get(coTag); if (!tagData) return; // Calculate similarity - const similarity = calculateJaccardSimilarity(tag, coTag); + const similarity = calculateJaccardSimilarity(tagSource, tag, coTag); relatedTags.push({ tag: coTag, @@ -320,7 +323,7 @@ class RelatedTagsUI { this.root.style.display = 'block'; // Update initialization status if not already done - if (!autoCompleteData.initialized) { + if (!autoCompleteData[TagSource.Danbooru].initialized) { if (this.autoRefreshTimerId) { clearTimeout(this.autoRefreshTimerId); } @@ -407,11 +410,11 @@ class RelatedTagsUI { #updateContent() { this.tagsContainer.innerHTML = ''; - if (!autoCompleteData.initialized) { + if (!autoCompleteData[TagSource.Danbooru].initialized) { // Show loading message const messageDiv = document.createElement('div'); messageDiv.className = 'related-tags-loading-message'; - messageDiv.textContent = `Initializing cooccurrence data... [${autoCompleteData.baseLoadingProgress.cooccurrence}%]`; + messageDiv.textContent = `Initializing cooccurrence data... [${autoCompleteData[TagSource.Danbooru].baseLoadingProgress.cooccurrence}%]`; this.tagsContainer.appendChild(messageDiv); return; } @@ -441,7 +444,7 @@ class RelatedTagsUI { * @returns {HTMLTableRowElement} The tag row element */ #createTagElement(tagData, isExisting) { - const categoryText = TagCategory[tagData.category] || "unknown"; + const categoryText = DanbooruTagCategory[tagData.category] || "unknown"; const tagRow = document.createElement('div'); tagRow.className = 'related-tag-item'; diff --git a/web/js/settings.js b/web/js/settings.js index 537fe53..cdb59b3 100644 --- a/web/js/settings.js +++ b/web/js/settings.js @@ -1,4 +1,8 @@ export const settingValues = { + // General settings + displayTagSource: 'danbooru', // 'danbooru', 'e621', 'all' + priorityTagSource: 'danbooru', // 'danbooru', 'e621', + // Autocomplete feature settings enabled: true, maxSuggestions: 10,