From 053ee9b7a67137ab46c2afcbf6605db8f7eae181 Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Sun, 18 May 2025 05:46:59 +0900 Subject: [PATCH 1/6] feat: Refactor autocomplete and data loader logic to support multiple tag sources (e.g. e621) (#8) --- modules/api.py | 141 ++++++++++++----------- web/js/autocomplete.js | 94 +++++++-------- web/js/data.js | 251 ++++++++++++++++++++++++++--------------- web/js/main.js | 29 ++++- web/js/related-tags.js | 29 ++--- web/js/settings.js | 4 + 6 files changed, 325 insertions(+), 223 deletions(-) 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, From e89a048815f1645d8ab1180945d1bce53304354f Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Sun, 18 May 2025 06:28:42 +0900 Subject: [PATCH 2/6] feat: Added "Display tags source" setting to filter displayed tags (UI only) --- web/js/main.js | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/web/js/main.js b/web/js/main.js index e572d00..abbbd21 100644 --- a/web/js/main.js +++ b/web/js/main.js @@ -139,13 +139,13 @@ app.registerExtension({ }); }, - //One the Settings Screen, displays reverse order in same category + // One the Settings Screen, displays reverse order in same category 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.", + tooltip: "If all tag source are enabled, which tag source's tags should be displayed first.", type: "combo", options: Object.values(TagSource), defaultValue: TagSource.Danbooru, @@ -153,6 +153,18 @@ app.registerExtension({ onChange: (newVal, oldVal) => { settingValues.priorityTagSource = newVal; } + }, + { + id: id + ".display_tag_source", + name: "Display Tag Source", + tooltip: "If multiple tag source are avirable, which tag source's tags should be displayed.", + type: "combo", + options: [...Object.values(TagSource), "all"], + defaultValue: "all", + category: [name, "General", "Display Tag Source"], + onChange: (newVal, oldVal) => { + settingValues.displayTagSource = newVal; + } }, // --- Autocomplete Settings --- { From 421112eb41af517463fdfccc5f91bca8921e3f86 Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Mon, 19 May 2025 12:21:33 +0900 Subject: [PATCH 3/6] feat: Refactor tag source handling and enhance UI with icon support for autocomplete suggestions --- web/css/autocomplete-plus.css | 191 +++++++++++++++++++++++----------- web/js/autocomplete.js | 39 ++++--- web/js/data.js | 55 +++++----- web/js/main.js | 39 ++++--- web/js/related-tags.js | 39 +++++-- web/js/settings.js | 7 +- web/js/utils.js | 33 +++++- 7 files changed, 276 insertions(+), 127 deletions(-) diff --git a/web/css/autocomplete-plus.css b/web/css/autocomplete-plus.css index 0432f28..29b52b8 100644 --- a/web/css/autocomplete-plus.css +++ b/web/css/autocomplete-plus.css @@ -1,39 +1,29 @@ /* Light Theme */ body { - --autocomplete-plus-text-color-cat-general: var(--p-blue-600); - --autocomplete-plus-text-color-cat-artist: var(--p-red-600); - --autocomplete-plus-text-color-cat-copyright: var(--p-purple-600); - --autocomplete-plus-text-color-cat-character: var(--p-green-600); - --autocomplete-plus-text-color-cat-meta: var(--p-yellow-600); - --autocomplete-plus-text-color-cat-unknown: var(--p-gray-600); + --autocomplete-plus-text-color-cat-blue: var(--p-blue-600); + --autocomplete-plus-text-color-cat-red: var(--p-red-600); + --autocomplete-plus-text-color-cat-purple: var(--p-purple-600); + --autocomplete-plus-text-color-cat-green: var(--p-green-600); + --autocomplete-plus-text-color-cat-yellow: var(--p-yellow-600); + --autocomplete-plus-text-color-cat-gray: var(--p-gray-600); + --autocomplete-plus-text-color-cat-sky: var(--p-sky-500); + --autocomplete-plus-text-color-cat-orange: var(--p-orange-600); + --autocomplete-plus-text-color-cat-white: var(--p-neutral-700); --autocomplete-plus-text-color-disabled: var(--p-gray-400); - - /* Selected color */ - --autocomplete-plus-text-color-cat-general-selected: var(--p-blue-200); - --autocomplete-plus-text-color-cat-artist-selected: var(--p-red-200); - --autocomplete-plus-text-color-cat-copyright-selected: var(--p-purple-200); - --autocomplete-plus-text-color-cat-character-selected: var(--p-green-200); - --autocomplete-plus-text-color-cat-meta-selected: var(--p-yellow-200); - --autocomplete-plus-text-color-cat-unknown-selected: var(--p-gray-200); } /* Dark Theme */ body.dark-theme { - --autocomplete-plus-text-color-cat-general: var(--p-blue-400); - --autocomplete-plus-text-color-cat-artist: var(--p-red-400); - --autocomplete-plus-text-color-cat-copyright: var(--p-purple-400); - --autocomplete-plus-text-color-cat-character: var(--p-green-400); - --autocomplete-plus-text-color-cat-meta: var(--p-yellow-400); - --autocomplete-plus-text-color-cat-unknown: var(--p-gray-500); + --autocomplete-plus-text-color-cat-blue: var(--p-blue-400); + --autocomplete-plus-text-color-cat-red: var(--p-red-400); + --autocomplete-plus-text-color-cat-purple: var(--p-purple-400); + --autocomplete-plus-text-color-cat-green: var(--p-green-400); + --autocomplete-plus-text-color-cat-yellow: var(--p-yellow-400); + --autocomplete-plus-text-color-cat-gray: var(--p-gray-400); + --autocomplete-plus-text-color-cat-sky: var(--p-sky-300); + --autocomplete-plus-text-color-cat-orange: var(--p-orange-400); + --autocomplete-plus-text-color-cat-white: var(--p-neutral-200); --autocomplete-plus-text-color-disabled: var(--p-gray-500); - - /* Selected color */ - --autocomplete-plus-text-color-cat-general-selected: var(--p-blue-200); - --autocomplete-plus-text-color-cat-artist-selected: var(--p-red-200); - --autocomplete-plus-text-color-cat-copyright-selected: var(--p-purple-200); - --autocomplete-plus-text-color-cat-character-selected: var(--p-green-200); - --autocomplete-plus-text-color-cat-meta-selected: var(--p-yellow-200); - --autocomplete-plus-text-color-cat-unknown-selected: var(--p-gray-200); } #autocomplete-plus-root { @@ -62,28 +52,60 @@ body.dark-theme { grid-template-columns: subgrid; } -.autocomplete-plus-item[data-tag-category="general"] { - color: var(--autocomplete-plus-text-color-cat-general); +.autocomplete-plus-item.danbooru[data-tag-category="general"] { + color: var(--autocomplete-plus-text-color-cat-blue); } -.autocomplete-plus-item[data-tag-category="artist"] { - color: var(--autocomplete-plus-text-color-cat-artist); +.autocomplete-plus-item.danbooru[data-tag-category="artist"] { + color: var(--autocomplete-plus-text-color-cat-red); } -.autocomplete-plus-item[data-tag-category="copyright"] { - color: var(--autocomplete-plus-text-color-cat-copyright); +.autocomplete-plus-item.danbooru[data-tag-category="copyright"] { + color: var(--autocomplete-plus-text-color-cat-purple); } -.autocomplete-plus-item[data-tag-category="character"] { - color: var(--autocomplete-plus-text-color-cat-character); +.autocomplete-plus-item.danbooru[data-tag-category="character"] { + color: var(--autocomplete-plus-text-color-cat-green); } -.autocomplete-plus-item[data-tag-category="meta"] { - color: var(--autocomplete-plus-text-color-cat-meta); +.autocomplete-plus-item.danbooru[data-tag-category="meta"] { + color: var(--autocomplete-plus-text-color-cat-yellow); } -.autocomplete-plus-item[data-tag-category="unknown"] { - color: var(--autocomplete-plus-text-color-cat-unknown); +.autocomplete-plus-item.danbooru[data-tag-category="unknown"] { + color: var(--autocomplete-plus-text-color-cat-gray); +} + +.autocomplete-plus-item.e621[data-tag-category="general"] { + color: var(--autocomplete-plus-text-color-cat-sky); +} + +.autocomplete-plus-item.e621[data-tag-category="artist"] { + color: var(--autocomplete-plus-text-color-cat-orange); +} + +.autocomplete-plus-item.e621[data-tag-category="copyright"] { + color: var(--autocomplete-plus-text-color-cat-purple); +} + +.autocomplete-plus-item.e621[data-tag-category="character"] { + color: var(--autocomplete-plus-text-color-cat-green); +} + +.autocomplete-plus-item.e621[data-tag-category="species"] { + color: var(--autocomplete-plus-text-color-cat-red); +} + +.autocomplete-plus-item.e621[data-tag-category="meta"] { + color: var(--autocomplete-plus-text-color-cat-white); +} + +.autocomplete-plus-item.e621[data-tag-category="lore"] { + color: var(--autocomplete-plus-text-color-cat-red); +} + +.autocomplete-plus-item.e621[data-tag-category="unknown"] { + color: var(--autocomplete-plus-text-color-cat-purple); } .autocomplete-plus-item span { @@ -113,22 +135,31 @@ body.dark-theme { background-color: var(--comfy-hover-bg); } -.autocomplete-plus-tag-name { - padding: 8px; +.autocomplete-plus-item .autocomplete-plus-tag-name { white-space: nowrap; } +.autocomplete-plus-item .autocomplete-plus-tag-icon-svg { + width: 1em; + height: 1em; + vertical-align: middle; +} + .autocomplete-plus-tag-name.autocomplete-plus-already-exists { color: var(--autocomplete-plus-text-color-disabled); } -.autocomplete-plus-alias { +.autocomplete-plus-item .autocomplete-plus-alias { overflow: hidden; color: var(--descrip-text); text-overflow: ellipsis; white-space: nowrap; } +.autocomplete-plus-item .autocomplete-plus-tag-count { + text-align: right; +} + /* Related Tags UI Styles */ #related-tags-root { position: absolute; @@ -163,9 +194,16 @@ body.dark-theme { font-weight: bold; } +.related-tags-header-tag-name .autocomplete-plus-tag-icon-svg { + width: 1em; + height: 1em; + vertical-align: middle; +} + .related-tags-header-controls { display: flex; align-items: center; + margin-left: 16px; gap: 4px; } @@ -210,28 +248,60 @@ body.dark-theme { grid-template-columns: subgrid; } -.related-tag-item[data-tag-category="general"] { - color: var(--autocomplete-plus-text-color-cat-general); +.related-tag-item.danbooru[data-tag-category="general"] { + color: var(--autocomplete-plus-text-color-cat-blue); } -.related-tag-item[data-tag-category="artist"] { - color: var(--autocomplete-plus-text-color-cat-artist); +.related-tag-item.danbooru[data-tag-category="artist"] { + color: var(--autocomplete-plus-text-color-cat-red); } -.related-tag-item[data-tag-category="copyright"] { - color: var(--autocomplete-plus-text-color-cat-copyright); +.related-tag-item.danbooru[data-tag-category="copyright"] { + color: var(--autocomplete-plus-text-color-cat-purple); } -.related-tag-item[data-tag-category="character"] { - color: var(--autocomplete-plus-text-color-cat-character); +.related-tag-item.danbooru[data-tag-category="character"] { + color: var(--autocomplete-plus-text-color-cat-green); } -.related-tag-item[data-tag-category="meta"] { - color: var(--autocomplete-plus-text-color-cat-meta); +.related-tag-item.danbooru[data-tag-category="meta"] { + color: var(--autocomplete-plus-text-color-cat-yellow); } -.related-tag-item[data-tag-category="unknown"] { - color: var(--autocomplete-plus-text-color-cat-unknown); +.related-tag-item.danbooru[data-tag-category="unknown"] { + color: var(--autocomplete-plus-text-color-cat-gray); +} + +.related-tag-item.e621[data-tag-category="general"] { + color: var(--autocomplete-plus-text-color-cat-sky); +} + +.related-tag-item.e621[data-tag-category="artist"] { + color: var(--autocomplete-plus-text-color-cat-orange); +} + +.related-tag-item.e621[data-tag-category="copyright"] { + color: var(--autocomplete-plus-text-color-cat-purple); +} + +.related-tag-item.e621[data-tag-category="character"] { + color: var(--autocomplete-plus-text-color-cat-green); +} + +.related-tag-item.e621[data-tag-category="species"] { + color: var(--autocomplete-plus-text-color-cat-red); +} + +.related-tag-item.e621[data-tag-category="meta"] { + color: var(--autocomplete-plus-text-color-cat-white); +} + +.related-tag-item.e621[data-tag-category="lore"] { + color: var(--autocomplete-plus-text-color-cat-red); +} + +.related-tag-item.e621[data-tag-category="unknown"] { + color: var(--autocomplete-plus-text-color-cat-gray); } /* stylelint-disable-next-line no-descending-specificity */ @@ -261,23 +331,22 @@ body.dark-theme { font-weight: bold; } -.related-tag-name { - padding: 8px; +.related-tag-item .related-tag-name { white-space: nowrap; } -.related-tag-name.related-tag-already-exists { +.related-tag-item .related-tag-name.related-tag-already-exists { color: var(--autocomplete-plus-text-color-disabled); } -.related-tag-alias { +.related-tag-item .related-tag-alias { overflow: hidden; color: var(--descrip-text); text-overflow: ellipsis; white-space: nowrap; } -.related-tag-similarity { +.related-tag-item .related-tag-similarity { color: var(--descrip-text); text-align: right; white-space: nowrap; @@ -294,4 +363,4 @@ body.dark-theme { color: var(--error-text); font-style: italic; text-align: center; -} +} \ No newline at end of file diff --git a/web/js/autocomplete.js b/web/js/autocomplete.js index 6750f33..d0456e1 100644 --- a/web/js/autocomplete.js +++ b/web/js/autocomplete.js @@ -1,8 +1,8 @@ import { - DanbooruTagCategory, + TagCategory, TagData, autoCompleteData, - getTagSourceInPriorityOrder + getEnabledTagSourceInPriorityOrder } from './data.js'; import { formatCountHumanReadable, @@ -13,7 +13,8 @@ import { normalizeTagToSearch, extractTagsFromTextArea, getCurrentTagRange, - getViewportMargin + getViewportMargin, + IconSvgHtmlString } from './utils.js'; import { settingValues } from './settings.js'; @@ -100,7 +101,7 @@ function searchCompletionCandidates(textareaElement) { queryVariations.add(hiraQuery); } - const sources = getTagSourceInPriorityOrder(); + const sources = getEnabledTagSourceInPriorityOrder(); for (const source of sources) { // Search in sortedTags (already sorted by count) for (const tagData of autoCompleteData[source].sortedTags) { @@ -274,6 +275,13 @@ class AutocompleteUI { this.root = document.createElement('div'); // Use table instead of div this.root.id = 'autocomplete-plus-root'; + // Create svg icon element as definition + this.iconSvgDef = document.createElement('div'); + this.iconSvgDef.style.position = 'absolute'; + this.iconSvgDef.style.display = 'none'; + this.iconSvgDef.innerHTML = IconSvgHtmlString; + this.root.appendChild(this.iconSvgDef); + this.tagsList = document.createElement('div'); this.tagsList.id = 'autocomplete-plus-list'; this.root.appendChild(this.tagsList); @@ -388,17 +396,24 @@ class AutocompleteUI { * @param {boolean} isExisting */ #createTagElement(tagData, isExisting) { - const categoryText = DanbooruTagCategory[tagData.category] || "unknown"; + const categoryText = TagCategory[tagData.source][tagData.category] || "unknown"; const tagRow = document.createElement('div'); - tagRow.classList.add('autocomplete-plus-item'); + tagRow.classList.add('autocomplete-plus-item', tagData.source); tagRow.dataset.tag = tagData.tag; tagRow.dataset.tagCategory = categoryText; - // Tag name + // Tag icon and name + const tagSourceIconHtml = ``; const tagName = document.createElement('span'); - tagName.classList.add('autocomplete-plus-tag-name'); - tagName.textContent = `[${tagData.source[0]}] ${tagData.tag}`; + tagName.className = 'autocomplete-plus-tag-name'; + if (settingValues.tagSourceIconPosition == 'hidden') { + tagName.textContent = tagData.tag; + } else { + tagName.innerHTML = settingValues.tagSourceIconPosition == 'left' + ? `${tagSourceIconHtml} ${tagData.tag}` + : `${tagData.tag} ${tagSourceIconHtml}`; + } // grayout tag name if it already exists if (isExisting) { @@ -407,7 +422,7 @@ class AutocompleteUI { // Alias const alias = document.createElement('span'); - alias.classList.add('autocomplete-plus-alias'); + alias.className = 'autocomplete-plus-alias'; // Display alias if available if (tagData.alias && tagData.alias.length > 0) { @@ -420,10 +435,11 @@ class AutocompleteUI { const category = document.createElement('span'); category.className = `autocomplete-plus-category`; category.textContent = `${categoryText.substring(0, 2)}`; + category.title = categoryText; // Full category on hover // Count const tagCount = document.createElement('span'); - category.className = `autocomplete-plus-tag-count`; + tagCount.className = `autocomplete-plus-tag-count`; tagCount.textContent = formatCountHumanReadable(tagData.count); tagRow.appendChild(tagName); @@ -458,7 +474,6 @@ class AutocompleteUI { const viewportHeight = window.innerHeight; const margin = getViewportMargin(); - const targetRect = this.target.getBoundingClientRect(); const targetElmOffset = this.#calculateElementOffset(this.target); const { top: caretTop, left: caretLeft, lineHeight: caretLineHeight } = this.#getCaretCoordinates(this.target); diff --git a/web/js/data.js b/web/js/data.js index 49d749d..6e85c8f 100644 --- a/web/js/data.js +++ b/web/js/data.js @@ -8,27 +8,27 @@ export const TagSource = { E621: 'e621', } -// Danbooru tag categories for display -export const DanbooruTagCategory = [ - 'general', - 'artist', - 'unused', - 'copyright', - 'character', - 'meta', -] - -export const E621TagCategory = [ - 'general', - 'artist', - 'unused', - 'copyright', - 'character', - 'species', - 'invalid', - 'meta', - 'lore', -] +export const TagCategory = { + 'danbooru': [ + 'general', + 'artist', + 'unused', + 'copyright', + 'character', + 'meta', + ], + 'e621': [ + 'general', + 'artist', + 'unused', + 'copyright', + 'character', + 'species', + 'invalid', + 'meta', + 'lore', + ] +} // --- Data Structures --- @@ -102,10 +102,17 @@ const COUNT_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('count'); // --- 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 getTagSourceInPriorityOrder() { - return Object.values(TagSource).toSorted((a, b) => { return a === settingValues.priorityTagSource ? -1 : 1; }); +export function getEnabledTagSourceInPriorityOrder() { + return Object.values(TagSource) + .filter((s) => { + return settingValues.tagSource === s || settingValues.tagSource === 'all'; + }) + .toSorted((a, b) => { + return a === settingValues.primaryTagSource ? -1 : 1; + }); } // --- Data Loading Functions --- diff --git a/web/js/main.js b/web/js/main.js index abbbd21..ead0f80 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 { TagSource, fetchCsvList, getTagSourceInPriorityOrder, initializeData } from "./data.js"; +import { TagSource, fetchCsvList, initializeData } from "./data.js"; import { AutocompleteEventHandler } from "./autocomplete.js"; import { RelatedTagsEventHandler } from "./related-tags.js"; @@ -133,7 +133,7 @@ app.registerExtension({ loadCSS(rootPath + "css/autocomplete-plus.css"); // Load CSS for autocomplete fetchCsvList().then((csvList) => { - getTagSourceInPriorityOrder().forEach((source) => { + Object.values(TagSource).forEach((source) => { initializeData(csvList, source); }); }); @@ -141,29 +141,40 @@ app.registerExtension({ // One the Settings Screen, displays reverse order in same category settings: [ - // --- General Settings --- + // --- Tag source Settings --- { - id: id + ".priority_tag_source", - name: "Priotized Tag Source", - tooltip: "If all tag source are enabled, which tag source's tags should be displayed first.", + id: id + ".tag_source_icon_position", + name: "Tag Source Icon Position", + type: "combo", + options: ["left", "right", "hidden"], + defaultValue: "left", + category: [name, "Tag Source", "Tag Source Icon Position"], + onChange: (newVal, oldVal) => { + settingValues.tagSourceIconPosition = newVal; + } + }, + { + id: id + ".primary_tag_source", + name: "Primary source for 'all' Source", + tooltip: "When 'Autocomplete Tag Source' is 'all', this determines which source's tags appear first in suggestions.", type: "combo", options: Object.values(TagSource), defaultValue: TagSource.Danbooru, - category: [name, "General", "Priority Tag Source"], + category: [name, "Tag Source", "Prioritize Tag Source"], onChange: (newVal, oldVal) => { - settingValues.priorityTagSource = newVal; + settingValues.primaryTagSource = newVal; } }, - { - id: id + ".display_tag_source", - name: "Display Tag Source", - tooltip: "If multiple tag source are avirable, which tag source's tags should be displayed.", + { + id: id + ".tag_source", + name: "Autocomplete Tag Source", + tooltip: "Select the tag source for autocomplete suggestions. 'all' includes tags from all loaded sources.", type: "combo", options: [...Object.values(TagSource), "all"], defaultValue: "all", - category: [name, "General", "Display Tag Source"], + category: [name, "Tag Source", "Tag Source"], onChange: (newVal, oldVal) => { - settingValues.displayTagSource = newVal; + settingValues.tagSource = newVal; } }, // --- Autocomplete Settings --- diff --git a/web/js/related-tags.js b/web/js/related-tags.js index 67e4d2f..22a4d0c 100644 --- a/web/js/related-tags.js +++ b/web/js/related-tags.js @@ -1,4 +1,4 @@ -import { TagSource, DanbooruTagCategory, TagData, autoCompleteData, getTagSourceInPriorityOrder } from './data.js'; +import { TagSource, TagCategory, TagData, autoCompleteData } from './data.js'; import { settingValues } from './settings.js'; import { extractTagsFromTextArea, @@ -7,7 +7,7 @@ import { isValidTag, normalizeTagToInsert, normalizeTagToSearch, - getCurrentTagRange + getCurrentTagRange, } from './utils.js'; // --- RelatedTags Logic --- @@ -91,6 +91,7 @@ function searchRelatedTags(tag) { similarity: similarity, alias: tagData.alias, category: tagData.category, + source: tagData.source, count: tagData.count, }); }); @@ -387,13 +388,29 @@ class RelatedTagsUI { * Updates header content */ #updateHeader() { + // Find the tag data for the current tag + const tagData = Object.values(TagSource) + .map((source) => { + if (source in autoCompleteData && autoCompleteData[source].tagMap.has(this.currentTag)) { + return autoCompleteData[source].tagMap.get(this.currentTag); + } + }) + .find((tagData) => tagData !== undefined); + // Update header text with current tag this.headerText.innerHTML = ''; // Clear previous content this.headerText.textContent = 'Tags related to: '; - const tagNameSpan = document.createElement('span'); - tagNameSpan.className = 'related-tags-header-tag-name'; - tagNameSpan.textContent = this.currentTag; - this.headerText.appendChild(tagNameSpan); + const tagName = document.createElement('span'); + tagName.className = 'related-tags-header-tag-name'; + tagName.textContent = this.currentTag; + if (tagData && ['left', 'right'].includes(settingValues.tagSourceIconPosition)) { + const tagSourceIconHtml = ``; + tagName.innerHTML = settingValues.tagSourceIconPosition == 'left' + ? `${tagSourceIconHtml} ${tagData.tag}` + : `${tagData.tag} ${tagSourceIconHtml}`; + } + + this.headerText.appendChild(tagName); // Update pin button this.pinBtn.textContent = this.isPinned ? '🎯' : '📌'; @@ -421,9 +438,9 @@ class RelatedTagsUI { if (!this.relatedTags || this.relatedTags.length === 0) { // Show no related tags message - const messageCell = document.createElement('div'); - messageCell.textContent = 'No related tags found'; - this.tagsContainer.appendChild(messageCell); + const messageDiv = document.createElement('div'); + messageDiv.textContent = 'No related tags found'; + this.tagsContainer.appendChild(messageDiv); return; } @@ -444,10 +461,10 @@ class RelatedTagsUI { * @returns {HTMLTableRowElement} The tag row element */ #createTagElement(tagData, isExisting) { - const categoryText = DanbooruTagCategory[tagData.category] || "unknown"; + const categoryText = TagCategory[tagData.source][tagData.category] || "unknown"; const tagRow = document.createElement('div'); - tagRow.className = 'related-tag-item'; + tagRow.classList.add('related-tag-item', tagData.source); tagRow.dataset.tag = tagData.tag; tagRow.dataset.tagCategory = categoryText; diff --git a/web/js/settings.js b/web/js/settings.js index cdb59b3..6b93e98 100644 --- a/web/js/settings.js +++ b/web/js/settings.js @@ -1,7 +1,8 @@ export const settingValues = { - // General settings - displayTagSource: 'danbooru', // 'danbooru', 'e621', 'all' - priorityTagSource: 'danbooru', // 'danbooru', 'e621', + // Tag source settings + tagSource: 'all', // 'danbooru', 'e621', 'all' + primaryTagSource: 'danbooru', // 'danbooru', 'e621' + tagSourceIconPosition: 'left', // 'left', 'right', 'hidden' // Autocomplete feature settings enabled: true, diff --git a/web/js/utils.js b/web/js/utils.js index 056ff81..ac0246d 100644 --- a/web/js/utils.js +++ b/web/js/utils.js @@ -1,3 +1,29 @@ +// --- Html String constants --- + +/** + * HTML string for the tag source icon. + */ +export const IconSvgHtmlString = ` + + + + + + + + + + + + e + + + + `; + // --- String Helper Functions --- const MAX_PROMPT_WEIGHT_VALUE = 9.9; @@ -491,7 +517,10 @@ export function getCurrentTagRange(text, cursorPos) { // --- End String Helper Functions --- -// Function to load a CSS file +/** + * Load a CSS file dynamically. + * @param {string} href + */ export function loadCSS(href) { const link = document.createElement('link'); link.rel = 'stylesheet'; @@ -500,7 +529,7 @@ export function loadCSS(href) { // Ensure the CSS is loaded before other scripts might rely on its styles // by adding it to the head. document.head.appendChild(link); - console.debug(`Loaded CSS: ${href}`); // Optional: Log loading + // console.debug(`[Autocomplete-Plus] Loaded CSS: ${href}`); // Optional: Log loading } /** From 4458fef0806a11e5e3d16dd5f2b509b18a0a2bc1 Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Mon, 19 May 2025 13:27:20 +0900 Subject: [PATCH 4/6] feat: Update README to support for Danbooru and e621 tags (#8) --- README.md | 20 +++++++++++++++++--- docs/README_jp.md | 21 ++++++++++++++++++--- 2 files changed, 35 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index d5ee45b..b11ff7f 100644 --- a/README.md +++ b/README.md @@ -9,8 +9,7 @@ ## Overview -**ComfyUI-Autocomplete-Plus** is a custom node that provides multiple input assistance features for any text area in [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Currently, it mainly supports Danbooru-style tag input. - +**ComfyUI-Autocomplete-Plus** is a custom node that provides multiple input assistance features for any text area in [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Currently, it mainly supports Currently, it supports Danbooru and e621 tags (e621 does not support some functions). ## Features - **:zap:No setup required**: Automatically downloads CSV data optimized for Danbooru tags. @@ -38,6 +37,7 @@ When you type in a text input area, tags that partially match the text are displ - Tag aliases are also included in the search. Japanese hiragana and katakana are searched without distinction. - Tags are color-coded by category. The color-coding rules are the same as Danbooru. - Tags that have already been entered are displayed grayed out. +- You can display Danbooru and e621 tags at the same time. You can also change the priority from the settings. ## Related Tags @@ -76,6 +76,11 @@ This is a CSV file for related tag calculation, recording tag pairs and their co Tag pairs are further filtered from the tag information CSV under the following conditions: - Co-occurrence count of 100 or more +### e621 CSV + +Currently, automatic download of CSV for e621 is not supported, so please manually place a CSV with the same structure as `danbooru_tags.csv` in the data folder with the name `e621_tags.csv`. +Also, displaying related tags is not supported. + ### User CSV Users can also use their own CSV files. For example, frequently used meta tags can be added to the autocomplete suggestions by placing a file named `danbooru_tags_meta.csv` in the `data` folder. @@ -97,6 +102,15 @@ worst_quality,5,9999999, ## Settings +### Tag Source + +> [!TIP] +> The source of tag data such as Danbooru or e621 is called the "tag source". + +- **Autocomplete Tag Source**: The tag source to display in the autocomplete suggestions. Select "all" to display all loaded tag sources. +- **Primary source for 'all' Source**: When `Autocomplete Tag Source` is set to "all", the tag source specified here will be displayed with priority. +- **Tag Source Icon Position**: Where to display the icon of the tag source. Select "hidden" to hide it. + ### Autocomplete - **Enable Autocomplete**: Enable/disable the autocomplete feature. @@ -107,7 +121,7 @@ worst_quality,5,9999999, - **Enable Related Tags**: Enable/disable the related tags feature. - **Max related tags**: Maximum number of related tags to display. - **Default Display Position**: Default display position when ComfyUI starts. -- **Related Tags Trigger Mode**: Trigger for displaying related tags (click, Ctrl+click). +- **Related Tags Trigger Mode**: Which action will trigger displaying related tags (click only, Ctrl+click) ## Known Issues diff --git a/docs/README_jp.md b/docs/README_jp.md index 457dec0..dc205c7 100644 --- a/docs/README_jp.md +++ b/docs/README_jp.md @@ -7,7 +7,7 @@ ## 概要 -**ComfyUI-Autocomplete-Plus** は、[ComfyUI](https://github.com/comfyanonymous/ComfyUI) の任意のテキストエリアに複数の入力支援機能を提供するカスタムノードです。現在は主にDanbooruスタイルのタグ入力に対応しています。 +**ComfyUI-Autocomplete-Plus** は、[ComfyUI](https://github.com/comfyanonymous/ComfyUI) の任意のテキストエリアに複数の入力支援機能を提供するカスタムノードです。現在はDanbooruとe621のタグに対応しています(e621は一部の機能が未対応です)。 ## 特徴 @@ -36,6 +36,7 @@ - タグのエイリアスも検索対象に含まれます。日本語のひらがな、カタカナは区別せず検索されます - タグのカテゴリ毎に色分けされます。色分けのルールは Danbooru と同じです - 入力済みのタグはグレーアウトで表示されます +- Danbooruとe621のタグを同時に表示出来ます。設定から優先順位を変更できます ## 関連タグ @@ -74,9 +75,14 @@ タグペアはタグ情報 CSV から、さらに以下の条件でフィルタリングされています。 - 共起回数が100件以上 +### e621 CSV + +現在、e621用 CSV の自動ダウンロードは未対応なため `danbooru_tags.csv` と同じ構造の CSV を `e621_tags.csv` という名前でデータフォルダーに手動配置してください。 +また、関連タグ表示も同様に未対応です。 + ### ユーザーCSV -ユーザーが自身で用意したCSVを使用することも可能です。例として、よく使うメタタグを `danbooru_tags_meta.csv` の名前で `data` フォルダーに配置することでオートコンプリート候補に追加できます。 +ユーザーが自身で用意した CSV を使用することも可能です。例として、よく使うメタタグを `danbooru_tags_meta.csv` の名前で `data` フォルダーに配置することでオートコンプリート候補に追加できます。 ヘッダー行はなくても構いません。反映にはブラウザのリロードが必要です。 **メタタグの例** @@ -95,6 +101,15 @@ worst_quality,5,9999999, ## 設定 +### タグソース + +> [!TIP] +> Danbooruやe621等のタグデータの提供元を「タグソース」と呼びます + +- **Autocomplete Tag Source**: オートコンプリート候補に表示するタグソース。「all」を選択するとロード済みの全てのタグソースを表示します +- **Primary source for 'all' Source**: `Autocomplete Tag Source` が「all」のとき、ここで指定したタグソースが優先して表示されます +- **Tag Source Icon Position**: タグの情報源のアイコンをどの位置に表示するか。「hidden」を選択すると非表示になります + ### オートコンプリート - **Enable Autocomplete**: オートコンプリート機能の有効化/無効化 @@ -105,7 +120,7 @@ worst_quality,5,9999999, - **Enable Related Tags**: 関連タグ機能の有効化/無効化 - **Max related tags**: 関連タグの最大表示件数 - **Default Display Position**: ComfyUI起動時のデフォルト表示位置 -- **Related Tags Trigger Mode** : 関連タグの表示トリガー(クリック、Ctrl+クリック) +- **Related Tags Trigger Mode** : 関連タグの表示する際、どの操作をトリガーとするか(クリックのみ、Ctrl+クリック) ## 既知の問題 From 920c1a1dadf52224478e11863d722e2073a33a1f Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Mon, 19 May 2025 23:31:01 +0900 Subject: [PATCH 5/6] fix: Correct wording in the overview section of README --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index b11ff7f..54fccaf 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ ## Overview -**ComfyUI-Autocomplete-Plus** is a custom node that provides multiple input assistance features for any text area in [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Currently, it mainly supports Currently, it supports Danbooru and e621 tags (e621 does not support some functions). +**ComfyUI-Autocomplete-Plus** is a custom node that provides multiple input assistance features for any text area in [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Currently, it supports Danbooru and e621 tags (e621 does not support some functions). ## Features - **:zap:No setup required**: Automatically downloads CSV data optimized for Danbooru tags. From 7fc8d79792d773dfd86bba92201939a052b4e14a Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Mon, 19 May 2025 23:34:16 +0900 Subject: [PATCH 6/6] fix: Correct wording of README_jp --- docs/README_jp.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/README_jp.md b/docs/README_jp.md index dc205c7..820ba6c 100644 --- a/docs/README_jp.md +++ b/docs/README_jp.md @@ -120,7 +120,7 @@ worst_quality,5,9999999, - **Enable Related Tags**: 関連タグ機能の有効化/無効化 - **Max related tags**: 関連タグの最大表示件数 - **Default Display Position**: ComfyUI起動時のデフォルト表示位置 -- **Related Tags Trigger Mode** : 関連タグの表示する際、どの操作をトリガーとするか(クリックのみ、Ctrl+クリック) +- **Related Tags Trigger Mode** : 関連タグを表示する際、どの操作をトリガーとするか(クリックのみ、Ctrl+クリック) ## 既知の問題