diff --git a/__init__.py b/__init__.py index 70d30b1..3f5ee6f 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,4 @@ -from .modules.api import * +from .modules.api import * from .modules import downloader # check and download necessary csv files diff --git a/modules/api.py b/modules/api.py index 313cfbf..4564179 100644 --- a/modules/api.py +++ b/modules/api.py @@ -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) \ No newline at end of file + return web.json_response({"last_check_time": None, "error": str(e)}, status=500) diff --git a/modules/downloader.py b/modules/downloader.py index 1704679..60a8969 100644 --- a/modules/downloader.py +++ b/modules/downloader.py @@ -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. """ diff --git a/pyproject.toml b/pyproject.toml index a3f2352..31f2bfb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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