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

refactor: Format code with ruff and update configuration

Browse Source
This commit is contained in:
newtextdoc1111
2025-07-05 13:40:17 +09:00
parent 57f80f0991
commit 98bf626ec1
4 changed files with 139 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+1 -1
View File
@@ -1,4 +1,4 @@
from .modules.api import *
from .modules.api import *
from .modules import downloader
# check and download necessary csv files
modules/api.py
+85 -83
View File
@@ -10,11 +10,12 @@ from . import downloader as dl
# os.path.join(..., '..', 'data') goes up one level and then into 'data'
DATA_DIR = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "data"))
DANBOORU_PREFIX = 'danbooru'
E621_PREFIX = 'e621'
DANBOORU_PREFIX = "danbooru"
E621_PREFIX = "e621"
TAGS_SUFFIX = "tags"
COOCCURRENCE_SUFFIX = "tags_cooccurrence"
TAGS_SUFFIX = 'tags'
COOCCURRENCE_SUFFIX = 'tags_cooccurrence'
def get_csv_file_status():
"""
@@ -23,17 +24,17 @@ def get_csv_file_status():
data = {
DANBOORU_PREFIX: {
'base_tags': False,
'extra_tags': [],
'base_cooccurrence': False,
'extra_cooccurrence': [],
"base_tags": False,
"extra_tags": [],
"base_cooccurrence": False,
"extra_cooccurrence": [],
},
E621_PREFIX: {
'base_tags': False,
'extra_tags': [],
'base_cooccurrence': False,
'extra_cooccurrence': [],
}
"base_tags": False,
"extra_tags": [],
"base_cooccurrence": False,
"extra_cooccurrence": [],
},
}
for prefix in [DANBOORU_PREFIX, E621_PREFIX]:
@@ -46,27 +47,28 @@ def get_csv_file_status():
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')]
all_csv_files = [f for f in os.listdir(DATA_DIR) if f.startswith(prefix) and f.endswith(".csv")]
# Create extra CSV files list
for filename in all_csv_files:
if filename in [base_tags_file, base_cooccurrence_file]:
continue # Skip base files
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,
"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
def get_last_check_time_from_metadata():
"""
Helper function to get the last remote check timestamp from csv_meta.json.
@@ -75,23 +77,25 @@ def get_last_check_time_from_metadata():
try:
if not os.path.exists(dl.CSV_META_FILE):
return None
with open(dl.CSV_META_FILE, 'r', encoding='utf-8') as f:
with open(dl.CSV_META_FILE, "r", encoding="utf-8") as f:
metadata = json.load(f)
datasets = metadata.get("hf_datasets", [])
if datasets and len(datasets) > 0:
return datasets[0].get("last_remote_check_timestamp")
return None
except (IOError, json.JSONDecodeError) as e:
print(f"[Autocomplete-Plus] Error reading csv_meta.json: {e}")
return None
# --- API Endpoints ---
@server.PromptServer.instance.routes.get('/autocomplete-plus/csv')
@server.PromptServer.instance.routes.get("/autocomplete-plus/csv")
async def get_csv_list(_request):
"""
Returns CSV file status.
@@ -102,43 +106,48 @@ async def get_csv_list(_request):
response = {
DANBOORU_PREFIX: {
'base_tags': csv_file_status[DANBOORU_PREFIX]['base_tags'],
'extra_tags': csv_file_status[DANBOORU_PREFIX]['extra_tags'],
'base_cooccurrence': csv_file_status[DANBOORU_PREFIX]['base_cooccurrence'],
'extra_cooccurrence': csv_file_status[DANBOORU_PREFIX]['extra_cooccurrence'],
"base_tags": csv_file_status[DANBOORU_PREFIX]["base_tags"],
"extra_tags": csv_file_status[DANBOORU_PREFIX]["extra_tags"],
"base_cooccurrence": csv_file_status[DANBOORU_PREFIX]["base_cooccurrence"],
"extra_cooccurrence": csv_file_status[DANBOORU_PREFIX]["extra_cooccurrence"],
},
E621_PREFIX: {
'base_tags': csv_file_status[E621_PREFIX]['base_tags'],
'extra_tags': csv_file_status[E621_PREFIX]['extra_tags'],
'base_cooccurrence': csv_file_status[E621_PREFIX]['base_cooccurrence'],
'extra_cooccurrence': csv_file_status[E621_PREFIX]['extra_cooccurrence'],
}
"base_tags": csv_file_status[E621_PREFIX]["base_tags"],
"extra_tags": csv_file_status[E621_PREFIX]["extra_tags"],
"base_cooccurrence": csv_file_status[E621_PREFIX]["base_cooccurrence"],
"extra_cooccurrence": csv_file_status[E621_PREFIX]["extra_cooccurrence"],
},
}
# Print csv file status to the console for debugging
print(f"""[Autocomplete-Plus] CSV file status:
* Danbooru -> base: {response[DANBOORU_PREFIX]['base_tags']}, extra: {", ".join(response[DANBOORU_PREFIX]['extra_tags'])}
* E621 -> base: {response[E621_PREFIX]['base_tags']}, extra: {", ".join(response[E621_PREFIX]['extra_tags'])}""")
* Danbooru -> base: {response[DANBOORU_PREFIX]["base_tags"]}, extra: {", ".join(response[DANBOORU_PREFIX]["extra_tags"])}
* E621 -> base: {response[E621_PREFIX]["base_tags"]}, extra: {", ".join(response[E621_PREFIX]["extra_tags"])}""")
return web.json_response(response)
@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/{source}/{suffix}/base')
@server.PromptServer.instance.routes.get("/autocomplete-plus/csv/{source}/{suffix}/base")
async def get_base_tags_file(request):
"""
Returns the base tags CSV 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]:
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/{source}/{suffix}/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.
@@ -146,16 +155,19 @@ async def get_extra_tags_file(request):
try:
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]:
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(csv_file_status[source][f'extra_{suffix}']):
index = int(request.match_info["index"])
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, csv_file_status[source][f'extra_{suffix}'][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)
@@ -164,7 +176,8 @@ async def get_extra_tags_file(request):
except ValueError:
return web.json_response({"error": "Invalid index format"}, status=400)
@server.PromptServer.instance.routes.post('/autocomplete-plus/csv/force-check-updates')
@server.PromptServer.instance.routes.post("/autocomplete-plus/csv/force-check-updates")
async def force_check_csv_updates(request):
"""
Forces a check for CSV file updates from HuggingFace, ignoring cooldown.
@@ -172,29 +185,29 @@ async def force_check_csv_updates(request):
"""
try:
print("[Autocomplete-Plus] Starting forced check for CSV updates from HuggingFace...")
downloader = dl.Downloader()
downloader.run_check_and_download(force_check=True)
print("[Autocomplete-Plus] Forced check completed successfully.")
# Get the updated last check time
last_check_time = get_last_check_time_from_metadata()
return web.json_response({
"success": True,
"message": "Force check completed successfully",
"last_check_time": last_check_time
})
return web.json_response(
{
"success": True,
"message": "Force check completed successfully",
"last_check_time": last_check_time,
}
)
except Exception as e:
print(f"[Autocomplete-Plus] Error during forced check: {e}")
return web.json_response({
"success": False,
"error": str(e)
}, status=500)
return web.json_response({"success": False, "error": str(e)}, status=500)
@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/last-check-time')
@server.PromptServer.instance.routes.get("/autocomplete-plus/csv/last-check-time")
async def get_last_check_time(_request):
"""
Returns the last remote check timestamp from csv_meta.json.
@@ -202,26 +215,15 @@ async def get_last_check_time(_request):
"""
try:
if not os.path.exists(dl.CSV_META_FILE):
return web.json_response({
"last_check_time": None,
"message": "csv_meta.json file not found"
})
return web.json_response({"last_check_time": None, "message": "csv_meta.json file not found"})
last_check_time = get_last_check_time_from_metadata()
if last_check_time is not None:
return web.json_response({
"last_check_time": last_check_time
})
return web.json_response({"last_check_time": last_check_time})
else:
return web.json_response({
"last_check_time": None,
"message": "No datasets found in metadata"
})
return web.json_response({"last_check_time": None, "message": "No datasets found in metadata"})
except (IOError, json.JSONDecodeError) as e:
print(f"[Autocomplete-Plus] Error reading csv_meta.json: {e}")
return web.json_response({
"last_check_time": None,
"error": str(e)
}, status=500)
return web.json_response({"last_check_time": None, "error": str(e)}, status=500)
modules/downloader.py
+44 -24
View File
@@ -20,7 +20,7 @@ CSV_META_FILE = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", C
DEFAULT_CSV_METADATA = {
"version": 1,
"hf_datasets" : [
"hf_datasets": [
{
"hf_dataset_id": "newtextdoc1111/danbooru-tag-csv",
"last_remote_check_timestamp": None,
@@ -34,10 +34,10 @@ DEFAULT_CSV_METADATA = {
"file_name": "danbooru_tags_cooccurrence.csv",
"last_download": None,
"last_modified_on_hf": None,
}
]
},
],
}
]
],
}
@@ -82,12 +82,14 @@ class Downloader:
return default_metadata
try:
with open(CSV_META_FILE, 'r', encoding='utf-8') as f:
with open(CSV_META_FILE, "r", encoding="utf-8") as f:
metadata = json.load(f)
if not isinstance(metadata, dict) or metadata.get("version") != DEFAULT_CSV_METADATA["version"]:
print(f"[Autocomplete-Plus] Metadata version mismatch. Expected {DEFAULT_CSV_METADATA['version']}, "
f"found {metadata.get('version')}. Using default metadata.")
print(
f"[Autocomplete-Plus] Metadata version mismatch. Expected {DEFAULT_CSV_METADATA['version']}, "
f"found {metadata.get('version')}. Using default metadata."
)
return default_metadata
else:
self.csv_meta_file_exists_at_start = True
@@ -101,7 +103,7 @@ class Downloader:
"""Saves metadata to CSV_META_FILE."""
try:
os.makedirs(os.path.dirname(CSV_META_FILE), exist_ok=True)
with open(CSV_META_FILE, 'w', encoding='utf-8') as f:
with open(CSV_META_FILE, "w", encoding="utf-8") as f:
json.dump(self.metadata, f, indent=2)
except IOError as e:
print(f"[Autocomplete-Plus] Error saving metadata to {CSV_META_FILE}: {e}")
@@ -113,10 +115,13 @@ class Downloader:
"""
url = f"https://huggingface.co/datasets/{dataset_repo_id}/resolve/main/{hf_filename}"
try:
req = urllib.request.Request(url, method='HEAD',
headers={"User-Agent": "ComfyUI-Autocomplete-Plus (Python urllib)"})
req = urllib.request.Request(
url,
method="HEAD",
headers={"User-Agent": "ComfyUI-Autocomplete-Plus (Python urllib)"},
)
with urllib.request.urlopen(req, timeout=10) as response:
last_modified_http = response.getheader('Last-Modified')
last_modified_http = response.getheader("Last-Modified")
if last_modified_http:
dt_object = parsedate_to_datetime(last_modified_http)
if dt_object.tzinfo is None or dt_object.tzinfo.utcoffset(dt_object) is None:
@@ -155,18 +160,28 @@ class Downloader:
if os.path.exists(temp_path):
os.remove(temp_path)
req = urllib.request.Request(download_url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 11.0; Win64)"})
req = urllib.request.Request(
download_url,
headers={"User-Agent": "Mozilla/5.0 (Windows NT 11.0; Win64)"},
)
with urllib.request.urlopen(req) as response:
total_size_str = response.getheader("Content-Length")
total_size = int(total_size_str) if total_size_str else None
chunk_size = 8192
with open(temp_path, "wb") as f_out, \
tqdm(total=total_size, unit='B', unit_scale=True, unit_divisor=1024,
desc=f"[Autocomplete-Plus] Downloading {file_name}", leave=False,
bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
) as pbar:
with (
open(temp_path, "wb") as f_out,
tqdm(
total=total_size,
unit="B",
unit_scale=True,
unit_divisor=1024,
desc=f"[Autocomplete-Plus] Downloading {file_name}",
leave=False,
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]",
) as pbar,
):
while True:
chunk = response.read(chunk_size)
if not chunk:
@@ -197,12 +212,15 @@ class Downloader:
if os.path.exists(final_path):
try:
file_size_at_final = os.path.getsize(final_path)
if file_size_at_final == 0 or \
(total_size and total_size > 0 and file_size_at_final < total_size) or \
(not total_size and downloaded_size > 0 and file_size_at_final < downloaded_size):
if (
file_size_at_final == 0
or (total_size and total_size > 0 and file_size_at_final < total_size)
or (not total_size and downloaded_size > 0 and file_size_at_final < downloaded_size)
):
os.remove(final_path)
print(
f"[Autocomplete-Plus] Removed potentially corrupted file at final destination: {final_path}")
f"[Autocomplete-Plus] Removed potentially corrupted file at final destination: {final_path}"
)
except OSError as rm_e:
print(f"[Autocomplete-Plus] Error removing potentially corrupted file {final_path}: {rm_e}")
@@ -221,7 +239,8 @@ class Downloader:
perform_hf_check = False
except (ValueError, KeyError, TypeError):
print(
"[Autocomplete-Plus] Invalid or missing timestamp for last_remote_check_timestamp. Will perform remote check.")
"[Autocomplete-Plus] Invalid or missing timestamp for last_remote_check_timestamp. Will perform remote check."
)
if perform_hf_check:
huggingface_dataset_id = dataset_meta["hf_dataset_id"]
@@ -246,7 +265,8 @@ class Downloader:
dataset_meta["last_remote_check_timestamp"] = now_utc.isoformat()
else:
print(
"[Autocomplete-Plus] Could not update all remote timestamps from HuggingFace. Will try again later.")
"[Autocomplete-Plus] Could not update all remote timestamps from HuggingFace. Will try again later."
)
def _download_csv_files_if_needed(self, dataset_meta: dict):
"""Downloads CSV files if they are missing, outdated, or previously failed."""
@@ -299,7 +319,7 @@ class Downloader:
"""
Orchestrates the process of checking for updates and downloading CSV files.
This is the main entry point for the downloader logic.
Args:
force_check: If True, forces a check of HuggingFace regardless of the last check timestamp.
"""
pyproject.toml
+9
View File
@@ -9,7 +9,16 @@ dependencies = ["",]
Repository = "https://github.com/newtextdoc1111/ComfyUI-Autocomplete-Plus"
# Used by Comfy Registry https://comfyregistry.org
[dependency-groups]
dev = [
"ruff>=0.12.2",
]
[tool.comfy]
PublisherId = "newtextdoc1111"
DisplayName = "ComfyUI-Autocomplete-Plus"
Icon = ""
[tool.ruff]
# Allow lines to be as long as 120.
line-length = 120