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

feat: Refactor autocomplete and data loader logic to support multiple tag sources (e.g. e621) (#8)

Browse Source
This commit is contained in:
newtextdoc1111
2025-05-18 05:46:59 +09:00
parent 49d9169ca0
commit 053ee9b7a6
6 changed files with 325 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/api.py
+74 -67
View File
@@ -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)