nodes/newtextdoc1111-ComfyUI-Autocomplete-Plus
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: get csv from python api and allow extra csv loading

Browse Source
This commit is contained in:
newtextdoc1111
2025-05-08 22:09:45 +09:00
parent 6536ccaef1
commit 27174964b3
6 changed files with 238 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+2
View File
@@ -1,3 +1,5 @@
from .modules.api import *
# --- Original Node Mappings (Keep or modify as needed) ---
WEB_DIRECTORY = "./web"
NODE_CLASS_MAPPINGS = {}
modules/api.py
+116
View File
@@ -0,0 +1,116 @@
import os
import server
from aiohttp import web
# Get the absolute path to the 'data' directory
# __file__ is the path to the current script (api.py)
# os.path.dirname(__file__) is the directory of the current script (modules)
# 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'
TAGS_BASE_FILE = 'danbooru_tags.csv'
COOCCURRENCE_BASE_FILE = 'danbooru_tags_cooccurrence.csv'
def get_extra_csv_files():
"""
Returns a dictionary with extra CSV files for tags and cooccurrence.
"""
tags_extra_files = []
cooccurrence_extra_files = []
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.")
# 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)
return {
'tags': tags_extra_files,
'cooccurrence': cooccurrence_extra_files,
}
# --- API Endpoints ---
@server.PromptServer.instance.routes.get('/autocomplete-plus/csv')
async def get_csv_list(_request):
"""
Returns number of extra CSV files available.
"""
extra_csv_files = get_extra_csv_files()
response = {
'tags': len(extra_csv_files['tags']),
'cooccurrence': len(extra_csv_files['cooccurrence']),
}
return web.json_response(response)
@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/tags/base')
async def get_base_tags_file(_request):
"""
Returns the base tags CSV file.
"""
file_path = os.path.join(DATA_DIR, TAGS_BASE_FILE)
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}')
async def get_extra_tags_file(request):
"""
Returns the extra tags CSV file at the specified index.
"""
try:
extra_csv_files = get_extra_csv_files()
index = int(request.match_info['index'])
if index < 0 or index >= len(extra_csv_files['tags']):
return web.json_response({"error": "Invalid index"}, status=404)
file_path = os.path.join(DATA_DIR, extra_csv_files['tags'][index])
if not os.path.exists(file_path):
return web.json_response({"error": "Extra tags file not found"}, status=404)
return web.FileResponse(file_path)
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_extra_csv_files()
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['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)
web/js/autocomplete.js
+1 -1
View File
@@ -273,7 +273,7 @@ class AutocompleteUI {
// 3rd cell: Category
const catCell = document.createElement('td');
catCell.textContent = TagCategory[candidate.category].substring(0, 2);
catCell.textContent = categoryText.substring(0, 2);
catCell.style.padding = '4px 8px';
catCell.style.minWidth = '50px'; // Ensure some minimum space for count alignment
web/js/data.js
+114 -93
View File
@@ -14,70 +14,73 @@ export const TagCategory = [
// Data storage
class AutoCompleteData {
tagsLoaded = false;
tagMap = new Map();
aliasMap = new Map();
sortedTags = [];
export const autoCompleteData = {
/** @type {Map<string, TagData>} */
tagMap: new Map(), // Stores tag data, mapping tag names to TagData objects
cooccurrenceLoaded = false;
cooccurrenceMap = new Map();
cooccurrenceInitProgress = 0;
}
sortedTags: [],
/** @type {Map<string, Map<string, number>>} */
cooccurrenceMap: new Map(), // Stores co-occurrence data for related tags
isInitializing: false,
initialized: false,
export const autoCompleteData = new AutoCompleteData();
// Progress of "base" csv loading
baseLoadingProgress: {
// tags: 0, // Commented out because tags csv aren't loaded in chunks
cooccurrence: 0
}
};
// --- Data Loading Functions ---
/**
* Loads and processes tag data from the CSV file.
* Loads tag data from a single CSV file.
* @param {string} csvUrl - The URL of the CSV file to load.
* @returns {Promise<void>}
*/
async function loadTags(rootPath) {
const startTime = performance.now();
const url = rootPath + 'data/danbooru_tags.csv';
async function loadTags(csvUrl) {
try {
const response = await fetch(url); //TODO: ignore browser cache
const response = await fetch(csvUrl, { cache: "no-store" });
if (!response.ok) {
throw new Error(`HTTP error! status: ${response.status}`);
}
const csvText = await response.text(); // Get raw CSV text
const csvText = await response.text();
const lines = csvText.split('\n').filter(line => line.trim().length > 0);
const totalLines = lines.length;
// Skip header row if present (tag,alias,count)
const startIndex = lines[0].startsWith('tag,alias,category,count') ? 1 : 0;
const parsedData = [];
for (let i = startIndex; i < lines.length; i++) {
const line = lines[i];
// Handle CSV parsing properly (consider quotes and commas in values)
const columns = parseCSVLine(line);
if (columns.length >= 3) {
if (columns.length == 4) {
const tag = columns[0].trim();
const aliasStr = columns[1].trim();
const category = columns[2].trim();
const count = parseInt(columns[3].trim(), 10);
// Skip invalid entries
if (!tag || isNaN(count)) continue;
// Parse aliases - might be comma-separated list inside quotes
const aliases = aliasStr ? aliasStr.split(',').map(a => a.trim()).filter(a => a.length > 0) : [];
parsedData.push({
autoCompleteData.sortedTags.push({
tag,
alias: aliases,
category,
count
});
}else{
console.warn(`[Autocomplete-Plus] Invalid CSV format in line ${i + 1} of ${csvUrl}: ${line}`);
continue;
}
}
// Sort by count in descending order
parsedData.sort((a, b) => b.count - a.count);
autoCompleteData.sortedTags = parsedData;
autoCompleteData.sortedTags.sort((a, b) => b.count - a.count);
// Build maps as before
autoCompleteData.sortedTags.forEach(tagData => {
@@ -91,26 +94,19 @@ async function loadTags(rootPath) {
}
});
autoCompleteData.tagsLoaded = true;
if(settingValues._logprocessingTime) {
const endTime = performance.now();
const duration = endTime - startTime;
console.debug(`[Autocomplete-Plus] Processed ${autoCompleteData.sortedTags.length} tags from CSV in ${duration.toFixed(2)}ms.`);
}
} catch (error) {
console.error(`[Autocomplete-Plus] Failed to fetch or process tags from ${url}:`, error);
autoCompleteData.tagsLoaded = false;
console.error(`[Autocomplete-Plus] Failed to fetch or process tags from ${csvUrl}:`, error);
}
}
/**
* Loads and processes cooccurrence data from the CSV file using chunked processing.
* Loads co-occurrence data from a single CSV file.
* @param {string} csvUrl - The URL of the CSV file to load.
* @returns {Promise<void>}
*/
async function loadCooccurrence(rootPath) {
const url = rootPath + 'data/danbooru_tags_cooccurrence.csv';
async function loadCooccurrence(csvUrl) {
try {
const response = await fetch(url);
const response = await fetch(csvUrl, { cache: "no-store" });
if (!response.ok) {
throw new Error(`HTTP error! status: ${response.status}`);
}
@@ -118,43 +114,21 @@ async function loadCooccurrence(rootPath) {
const csvText = await response.text();
const lines = csvText.split('\n').filter(line => line.trim().length > 0);
// Skip header row if present (tag_a,tag_b,count)
const startIndex = lines[0].startsWith('tag_a,tag_b,count') ? 1 : 0;
// Create a new Map to store the bidirectional relationships
const bidirectionalMap = new Map();
// Process CSV data in chunks
await processInChunks(lines, startIndex, bidirectionalMap);
// Assign the bidirectional map to the cooccurrenceMap
let actualPairCount = 0;
let primaryTagCount = 0;
autoCompleteData.cooccurrenceMap.clear();
for (const [tag, relatedTags] of bidirectionalMap) {
autoCompleteData.cooccurrenceMap.set(tag, relatedTags);
actualPairCount += relatedTags.size;
primaryTagCount++;
}
autoCompleteData.cooccurrenceLoaded = true;
if(settingValues._logprocessingTime) {
console.log(`[Autocomplete-Plus] Processed bidirectional relationships for ${primaryTagCount} tags from CSV.`);
}
await processInChunks(lines, startIndex, autoCompleteData.cooccurrenceMap, csvUrl);
} catch (error) {
console.error(`[Autocomplete-Plus] Failed to fetch or process cooccurrence data from ${url}:`, error);
autoCompleteData.cooccurrenceLoaded = false;
console.error(`[Autocomplete-Plus] Failed to fetch or process cooccurrence data from ${csvUrl}:`, error);
}
}
/**
* Process CSV data in chunks to avoid blocking the UI
* Process CSV data in chunks to avoid blocking the UI.
* Modifies the targetMap directly.
*/
function processInChunks(lines, startIndex, bidirectionalMap) {
function processInChunks(lines, startIndex, targetMap, sourceFileName = "CSV") {
return new Promise((resolve) => {
const CHUNK_SIZE = 10000; // Process 10,000 lines at a time
const CHUNK_SIZE = 10000;
let i = startIndex;
let pairCount = 0;
@@ -170,38 +144,33 @@ function processInChunks(lines, startIndex, bidirectionalMap) {
const tagB = columns[1].trim();
const count = parseInt(columns[2].trim(), 10);
// Skip invalid entries
if (!tagA || !tagB || isNaN(count)) continue;
// Add tagA -> tagB relationship
if (!bidirectionalMap.has(tagA)) {
bidirectionalMap.set(tagA, new Map());
if (!targetMap.has(tagA)) {
targetMap.set(tagA, new Map());
}
bidirectionalMap.get(tagA).set(tagB, count);
targetMap.get(tagA).set(tagB, count);
// Add tagB -> tagA relationship (bidirectional)
if (!bidirectionalMap.has(tagB)) {
bidirectionalMap.set(tagB, new Map());
if (!targetMap.has(tagB)) {
targetMap.set(tagB, new Map());
}
bidirectionalMap.get(tagB).set(tagA, count);
targetMap.get(tagB).set(tagA, count);
pairCount++;
}
}
if (i < lines.length) {
// Update progress in the data object
autoCompleteData.cooccurrenceInitProgress = Math.round((i / lines.length) * 100);
// Schedule next chunk with setTimeout to allow UI updates
autoCompleteData.baseLoadingProgress.cooccurrence = Math.round((i / lines.length) * 100);
setTimeout(processChunk, 0);
} else {
console.log(`[Autocomplete-Plus] Finished processing ${pairCount} one-way cooccurrence pairs`);
resolve();
}
}
// Start processing the first chunk
processChunk();
});
}
@@ -221,15 +190,12 @@ function parseCSVLine(line) {
if (char === '"') {
if (inQuotes && i + 1 < line.length && line[i + 1] === '"') {
// Escaped quote (double quote inside quotes)
current += '"';
i++; // Skip the next quote
i++;
} else {
// Toggle quote mode
inQuotes = !inQuotes;
}
} else if (char === ',' && !inQuotes) {
// End of column
result.push(current);
current = '';
} else {
@@ -237,17 +203,72 @@ function parseCSVLine(line) {
}
}
// Don't forget to add the last column
result.push(current);
return result;
}
export function loadAllData(rootPath) {
return Promise.all([
loadTags(rootPath),
loadCooccurrence(rootPath)
]).catch(error => {
console.error("[Autocomplete-Plus] Error loading data:", error);
});
/**
* 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...");
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();
const extraTagsCount = csvListData.tags || 0;
const extraCooccurrenceCount = csvListData.cooccurrence || 0;
const tagsUrl = '/autocomplete-plus/csv/tags';
const tagsLoadPromises = [
loadTags(`${tagsUrl}/base`)
];
let currentTagPromise = tagsLoadPromises[0];
for (let i = 0; i < extraTagsCount; i++) {
currentTagPromise = currentTagPromise.then(loadTags(`${tagsUrl}/extra/${i}`));
tagsLoadPromises.push(currentTagPromise);
}
const cooccurrenceUrl = '/autocomplete-plus/csv/cooccurrence';
const cooccurrenceLoadPromises = [
loadCooccurrence(`${cooccurrenceUrl}/base`)
];
let cooccurrencePromiseChain = cooccurrenceLoadPromises[0];
for (let i = 0; i < extraCooccurrenceCount; i++) {
cooccurrencePromiseChain = cooccurrencePromiseChain.then(loadCooccurrence(`${cooccurrenceUrl}/extra/${i}`));
cooccurrenceLoadPromises.push(cooccurrencePromiseChain);
}
await Promise.all([
Promise.all(tagsLoadPromises).then(() => {
const endTime = performance.now();
console.log(`[Autocomplete-Plus] Tags loading complete in ${(endTime - startTime).toFixed(2)}ms. Extra file count: ${extraTagsCount}`);
}),
Promise.all(cooccurrenceLoadPromises).then(() => {
const endTime = performance.now();
console.log(`[Autocomplete-Plus] Co-occurrence loading complete in ${(endTime - startTime).toFixed(2)}ms. Extra file count: ${extraCooccurrenceCount}`);
})
]);
autoCompleteData.initialized = true;
} catch (error) {
console.error("[Autocomplete-Plus] Error initializing autocomplete data:", error);
} finally {
autoCompleteData.isInitializing = false;
}
}
web/js/main.js
+2 -2
View File
@@ -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 { loadAllData } from "./data.js";
import { initializeData } from "./data.js";
import { AutocompleteEventHandler } from "./autocomplete.js";
import { RelatedTagsEventHandler } from "./related-tags.js";
@@ -132,7 +132,7 @@ app.registerExtension({
let rootPath = import.meta.url.replace("js/main.js", "");
loadCSS(rootPath + "css/autocomplete-plus.css"); // Load CSS for autocomplete
loadAllData(rootPath);
initializeData();
},
settings: [
{
web/js/related-tags.js
+3 -3
View File
@@ -77,7 +77,7 @@ class RelatedTagsUI {
// Make visible
this.root.style.display = 'block';
if (!autoCompleteData.cooccurrenceLoaded) {
if (!autoCompleteData.initialized) {
if (this.autoRefreshTimerId) {
clearTimeout(this.autoRefreshTimerId);
}
@@ -119,10 +119,10 @@ class RelatedTagsUI {
tagNameSpan.textContent = this.currentTag;
this.header.appendChild(tagNameSpan);
if (!autoCompleteData.cooccurrenceLoaded) {
if (!autoCompleteData.initialized) {
const messageDiv = document.createElement('div');
messageDiv.className = 'related-tags-loading-message';
messageDiv.textContent = `Initializing cooccurrence data... [${autoCompleteData.cooccurrenceInitProgress}%]`;
messageDiv.textContent = `Initializing cooccurrence data... [${autoCompleteData.baseLoadingProgress.cooccurrence}%]`;
this.tagsContainer.appendChild(messageDiv);
return;
}