diff --git a/.gitignore b/.gitignore index 4a7efb4..da7b88a 100644 --- a/.gitignore +++ b/.gitignore @@ -15,4 +15,7 @@ node_modules/ # csv files data/*.csv -data/.download \ No newline at end of file +data/.download + +# csv meta information +csv_meta.json diff --git a/__init__.py b/__init__.py index c82cc73..70d30b1 100644 --- a/__init__.py +++ b/__init__.py @@ -2,7 +2,8 @@ from .modules.api import * from .modules import downloader # check and download necessary csv files -downloader.check_and_download_csv_files() +dl = downloader.Downloader() +dl.run_check_and_download() WEB_DIRECTORY = "./web" NODE_CLASS_MAPPINGS = {} diff --git a/modules/downloader.py b/modules/downloader.py index d675137..81f41c8 100644 --- a/modules/downloader.py +++ b/modules/downloader.py @@ -1,175 +1,311 @@ import os -import shutil # Required for moving files -import sys # Required for sys.stdout for progress bar -import urllib.error +import json import urllib.request - +import urllib.error +import shutil +import sys +from datetime import datetime, timezone, timedelta +from email.utils import parsedate_to_datetime from tqdm import tqdm -# --- Constants for Status --- -STATUS_NOT_STARTED = "not_started" -STATUS_DOWNLOADING = "downloading" # This status will be very brief for synchronous downloads -STATUS_COMPLETED = "completed" -STATUS_FAILED = "failed" - # --- File Definitions --- -TAGS_FILE_ID = "danbooru_tags" -TAGS_FILE = "danbooru_tags.csv" - -COOCCURRENCE_FILE_ID = "danbooru_tags_cooccurrence" -COOCCURRENCE_FILE = "danbooru_tags_cooccurrence.csv" # Get the absolute path to the "data" directory DATA_DIR = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "data")) TEMP_DOWNLOAD_DIR = os.path.join(DATA_DIR, ".download") -HUGGINGFACE_REPO_URL = "https://huggingface.co/datasets/newtextdoc1111/danbooru-tag-csv/resolve/main/" +# --- Metadata Constants --- +CSV_META_FILE_NAME = "csv_meta.json" +CSV_META_FILE = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", CSV_META_FILE_NAME)) -# --- State Management --- -# Stores the status of each file download -_download_states = { - TAGS_FILE_ID: { - "filename": TAGS_FILE, - "url": HUGGINGFACE_REPO_URL + TAGS_FILE, - "path": os.path.join(DATA_DIR, TAGS_FILE), - "status": STATUS_NOT_STARTED, - }, - COOCCURRENCE_FILE_ID: { - "filename": COOCCURRENCE_FILE, - "url": HUGGINGFACE_REPO_URL + COOCCURRENCE_FILE, - "path": os.path.join(DATA_DIR, COOCCURRENCE_FILE), - "status": STATUS_NOT_STARTED, - } +DEFAULT_CSV_METADATA = { + "version": 1, + "hf_datasets" : [ + { + "hf_dataset_id": "newtextdoc1111/danbooru-tag-csv", + "last_remote_check_timestamp": None, + "csv_files": [ + { + "file_name": "danbooru_tags.csv", + "last_download": None, + "last_modified_on_hf": None, + }, + { + "file_name": "danbooru_tags_cooccurrence.csv", + "last_download": None, + "last_modified_on_hf": None, + } + ] + } + ] } + # --- Helper Functions --- -def get_file_path(file_id: str) -> str: - """Returns the full local path for a given file_id.""" - return _download_states[file_id]["path"] +def get_file_path(file_name: str) -> str: + """Returns the full local path for a given file_name.""" + return os.path.join(DATA_DIR, file_name) -def get_temp_download_path(file_id: str) -> str: - """Returns the full temporary download path for a given file_id.""" - state = _download_states[file_id] - return os.path.join(TEMP_DOWNLOAD_DIR, state["filename"]) -def get_file_status(file_id: str) -> str: - """Gets the current download status of a file. Checks if file exists if status is not_started.""" - state = _download_states[file_id] - # If file exists and status is "not_started", it might have been placed manually or from a previous run. - # Update status to "completed" if the file is already there. - if os.path.exists(state["path"]) and state["status"] == STATUS_NOT_STARTED: - state["status"] = STATUS_COMPLETED - return state["status"] +def get_temp_download_path(file_name: str) -> str: + """Returns the full temporary download path for a given file_name.""" + return os.path.join(TEMP_DOWNLOAD_DIR, file_name) -def _download_file_with_progress_sync(file_id: str): - """Downloads a file synchronously with progress display to a temporary location.""" - state = _download_states[file_id] - final_path = state["path"] - temp_path = get_temp_download_path(file_id) - # This check is somewhat redundant if ensure_file_is_downloaded_sync calls it, - # but good for direct calls or safety. - if state["status"] == STATUS_COMPLETED and os.path.exists(final_path) and os.path.getsize(final_path) > 0: - return +def check_file_valid(file_path): + """Checks if a file is valid by verifying its existence and size.""" + return os.path.exists(file_path) and os.path.getsize(file_path) > 0 - state["status"] = STATUS_DOWNLOADING - print(f"[Autocomplete-Plus] Attempting to download {state['filename']} from {state['url']} to {temp_path}") - downloaded_size = 0 - total_size = 0 +class Downloader: + """ + Downloader class to manage downloading and checking of CSV files from HuggingFace. + """ - try: + def __init__(self): + """Initializes the Downloader.""" + self.csv_meta_file_exists_at_start = False + + self._ensure_directories_exist() + self.metadata = self._load_metadata() + + def get_default_csv_metadata(self): + """Returns the default CSV metadata.""" + return json.loads(json.dumps(DEFAULT_CSV_METADATA)) + + def _load_metadata(self) -> list: + """Loads metadata from CSV_META_FILE. Returns default if not found or error.""" + default_metadata = self.get_default_csv_metadata() + + if not os.path.exists(CSV_META_FILE): + print(f"[Autocomplete-Plus] Metadata file not found: {CSV_META_FILE}. Using default metadata.") + return default_metadata + + try: + 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.") + return default_metadata + else: + self.csv_meta_file_exists_at_start = True + return metadata + + except (IOError, json.JSONDecodeError) as e: + print(f"[Autocomplete-Plus] Error loading metadata from {CSV_META_FILE}: {e}. Using default metadata.") + return default_metadata + + def _save_metadata(self): + """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: + json.dump(self.metadata, f, indent=2) + except IOError as e: + print(f"[Autocomplete-Plus] Error saving metadata to {CSV_META_FILE}: {e}") + + def _get_hf_file_last_modified(self, dataset_repo_id: str, hf_filename: str) -> str | None: + """ + Retrieves the Last-Modified header for a file on HuggingFace and returns it as an ISO 8601 string. + hf_filename should be the full filename, e.g., "danbooru_tags.csv". + """ + 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)"}) + with urllib.request.urlopen(req, timeout=10) as response: + 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: + dt_object = dt_object.replace(tzinfo=timezone.utc) + else: + dt_object = dt_object.astimezone(timezone.utc) + return dt_object.isoformat() + else: + print(f"[Autocomplete-Plus] 'Last-Modified' header not found for {hf_filename} at {url}") + return None + except urllib.error.HTTPError as e: + print(f"[Autocomplete-Plus] HTTP error when checking last modified for {hf_filename}: {e.code} {e.reason}") + return None + except (urllib.error.URLError, TimeoutError) as e: + print(f"[Autocomplete-Plus] URL or network error when checking last modified for {hf_filename}: {e}") + return None + except Exception as e: + print(f"[Autocomplete-Plus] Unexpected error when checking last modified for {hf_filename}: {e}") + return None + + def _download_file_with_progress_sync(self, hf_dataset_id, file_name: str, metadata_entry: dict): + """ + Downloads a file synchronously with progress display to a temporary location. + """ + download_url = f"https://huggingface.co/datasets/{hf_dataset_id}/resolve/main/{file_name}" + final_path = get_file_path(file_name) + temp_path = get_temp_download_path(file_name) + now_utc = datetime.now(timezone.utc).isoformat() + + print(f"[Autocomplete-Plus] Attempting to download {file_name} from {download_url}") + + downloaded_size = 0 + total_size = 0 + + try: + 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)"}) + + 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: + while True: + chunk = response.read(chunk_size) + if not chunk: + break + + f_out.write(chunk) + downloaded_size += len(chunk) + pbar.update(len(chunk)) + + sys.stdout.write("\r" + " " * 100 + "\r") + sys.stdout.flush() + + shutil.move(temp_path, final_path) + print(f"[Autocomplete-Plus] Successfully downloaded and moved {file_name} to {final_path}.") + metadata_entry["last_download"] = now_utc + + except (urllib.error.URLError, OSError, TimeoutError) as e: + sys.stdout.write("\n") + sys.stdout.flush() + print(f"[Autocomplete-Plus] Error downloading {file_name}: {e}") + if os.path.exists(temp_path): + try: + os.remove(temp_path) + print(f"[Autocomplete-Plus] Removed partially downloaded file from temp: {temp_path}") + except OSError as rm_e: + print(f"[Autocomplete-Plus] Error removing temporary file {temp_path}: {rm_e}") + + 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): + os.remove(final_path) + print( + 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}") + + def _ensure_directories_exist(self): + """Ensures that DATA_DIR and TEMP_DOWNLOAD_DIR exist.""" os.makedirs(DATA_DIR, exist_ok=True) os.makedirs(TEMP_DOWNLOAD_DIR, exist_ok=True) - # Remove temporary file if it exists from a previous failed attempt - if os.path.exists(temp_path): - os.remove(temp_path) - - req = urllib.request.Request(state["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 # Use None for tqdm if unknown - 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 {state['filename']}", 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: - if total_size is None and downloaded_size > 0: # Content-Length was unknown - sys.stdout.write(f"\n[Autocomplete-Plus] Downloaded {state['filename']}: {downloaded_size / 1024:.2f} KiB (Content-Length unknown).\n") - sys.stdout.flush() - elif total_size is not None and downloaded_size < total_size: - # This case might indicate an incomplete download if loop breaks unexpectedly - print(f"\n[Autocomplete-Plus] Warning: Download of {state['filename']} may be incomplete. " - f"Expected {total_size} bytes, got {downloaded_size} bytes.") - break - - f_out.write(chunk) - downloaded_size += len(chunk) - pbar.update(len(chunk)) - - # Ensure the progress bar is cleared and a new line is printed if necessary - sys.stdout.write("\r" + " " * 100 + "\r") # Clear the tqdm line - sys.stdout.flush() - - shutil.move(temp_path, final_path) - print(f"[Autocomplete-Plus] Successfully downloaded and moved {state['filename']} to {final_path}.") - state["status"] = STATUS_COMPLETED - - except (urllib.error.URLError, OSError) as e: # Catch specific exceptions for network/file issues - sys.stdout.write("\n") # Ensure newline after any tqdm output - sys.stdout.flush() - print(f"[Autocomplete-Plus] Error downloading {state['filename']}: {e}") - state["status"] = STATUS_FAILED - if os.path.exists(temp_path): + def _check_new_csv_from_hf_dataset(self, dataset_meta: dict, now_utc: datetime): + """Checks HuggingFace for file updates and updates metadata.""" + perform_hf_check = True + if dataset_meta.get("last_remote_check_timestamp"): try: - os.remove(temp_path) - print(f"[Autocomplete-Plus] Removed partially downloaded file from temp: {temp_path}") - except OSError as rm_e: - print(f"[Autocomplete-Plus] Error removing temporary file {temp_path}: {rm_e}") + last_check_dt = datetime.fromisoformat(dataset_meta["last_remote_check_timestamp"]) + if now_utc - last_check_dt < timedelta(days=7): + perform_hf_check = False + except (ValueError, KeyError, TypeError): + print( + "[Autocomplete-Plus] Invalid or missing timestamp for last_remote_check_timestamp. Will perform remote check.") - # If the final file exists and might be corrupted (e.g. 0 bytes or smaller than expected) - if os.path.exists(final_path): - try: - file_size_at_final = os.path.getsize(final_path) - # Conditions for removal: 0 bytes, or smaller than total known size, - # or (if total size unknown but we downloaded something) smaller than what we downloaded. - if file_size_at_final == 0 or \ - (total_size > 0 and file_size_at_final < total_size) or \ - (total_size == 0 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}") - except OSError as rm_e: - print(f"[Autocomplete-Plus] Error removing potentially corrupted file {final_path}: {rm_e}") + if perform_hf_check: + huggingface_dataset_id = dataset_meta["hf_dataset_id"] -def ensure_file_is_downloaded_sync(file_id: str): - """Checks if a file needs downloading and initiates synchronous download if so.""" - state = _download_states[file_id] - file_path = state["path"] # Final destination path + print(f"[Autocomplete-Plus] Checking HuggingFace dataset {huggingface_dataset_id} for file updates...") + updated_all_hf_timestamps_successfully = True + for file_meta_info in dataset_meta["csv_files"]: + file_name_from_meta = file_meta_info.get("file_name") + if not file_name_from_meta: + print("[Autocomplete-Plus] Warning: file_name missing in csv_files metadata entry.") + updated_all_hf_timestamps_successfully = False + continue - # Check if the file already exists at the final destination and is valid (size > 0) - if os.path.exists(file_path) and os.path.getsize(file_path) > 0: - if state["status"] != STATUS_COMPLETED: - state["status"] = STATUS_COMPLETED - return # Skip download + last_modified = self._get_hf_file_last_modified(huggingface_dataset_id, file_name_from_meta) + if last_modified: + file_meta_info["last_modified_on_hf"] = last_modified + else: + updated_all_hf_timestamps_successfully = False + print(f"[Autocomplete-Plus] Failed to get remote last modified time for {file_name_from_meta}.") - # If file doesn't exist, is empty, or previous download failed / not started, then attempt download. - _download_file_with_progress_sync(file_id) + if updated_all_hf_timestamps_successfully: + 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.") -def check_and_download_csv_files(): - """Checks and downloads CSV files.""" - # print("[Autocomplete-Plus] --- Starting check and download for all necessary CSV files... ---") + def _download_csv_files_if_needed(self, dataset_meta: dict): + """Downloads CSV files if they are missing, outdated, or previously failed.""" - # Ensure data and temp directories exist - os.makedirs(DATA_DIR, exist_ok=True) - os.makedirs(TEMP_DOWNLOAD_DIR, exist_ok=True) + for file_meta_entry in dataset_meta["csv_files"]: + file_name = file_meta_entry.get("file_name") + if not file_name: + print("[Autocomplete-Plus] Warning: file_name missing in file_meta_entry during download check.") + continue - ensure_file_is_downloaded_sync(TAGS_FILE_ID) - ensure_file_is_downloaded_sync(COOCCURRENCE_FILE_ID) - # print("[Autocomplete-Plus] --- Finished check and download for all necessary CSV files. ---") + reason_for_download = self.check_csv_file_should_download(file_meta_entry, file_name) + + if reason_for_download: + print(f"[Autocomplete-Plus] Queuing download for {file_name}: {reason_for_download}") + self._download_file_with_progress_sync(dataset_meta["hf_dataset_id"], file_name, file_meta_entry) + + def check_csv_file_should_download(self, file_meta_entry, file_name): + """ + Determines if a file should be downloaded based on its metadata and local status. + Returns a reason string if a download is needed, or None if no action is required. + """ + local_file_path = get_file_path(file_name) + + # If the CSV_META_FILE was not found initially, we need to download the file. + if not self.csv_meta_file_exists_at_start: + return f"{CSV_META_FILE_NAME} was not found initially. Forcing download of {file_name}." + + # If the file is empty or missing, we need to download it. + if not check_file_valid(local_file_path): + return f"File {file_name} is missing or empty locally." + + # Check if the last modified date on HuggingFace is newer than the last download date + try: + last_download_dt = datetime.fromisoformat(file_meta_entry["last_download"]) + hf_modified_dt = datetime.fromisoformat(file_meta_entry["last_modified_on_hf"]) + if hf_modified_dt > last_download_dt: + return f"Remote file {file_name} is newer (HF: {last_download_dt}, Local Download: {hf_modified_dt})." + except (ValueError, TypeError): + file_meta_entry["last_download"] = None + file_meta_entry["last_modified_on_hf"] = None + return f"Invalid timestamp format for {file_name}. Forcing download to ensure integrity." + + # If the file is missing or empty, but the last download timestamp exists, we need to retry. + if not check_file_valid(local_file_path) and file_meta_entry.get("last_download") is not None: + return f"File {file_name} is missing or empty locally, but last download timestamp exists. Retrying." + + return None + + def run_check_and_download(self): + """ + Orchestrates the process of checking for updates and downloading CSV files. + This is the main entry point for the downloader logic. + """ + + now_utc = datetime.now(timezone.utc) + + datasets_meta = self.metadata.get("hf_datasets", []) + for dataset_meta in datasets_meta: + self._check_new_csv_from_hf_dataset(dataset_meta, now_utc) + self._download_csv_files_if_needed(dataset_meta) + + self._save_metadata() diff --git a/web/js/autocomplete.js b/web/js/autocomplete.js index d0456e1..d173081 100644 --- a/web/js/autocomplete.js +++ b/web/js/autocomplete.js @@ -889,7 +889,7 @@ export class AutocompleteEventHandler { // If UI is not visible, and the key is a non-character key (length > 1) // and not Delete or Backspace, then do nothing. // This prevents UI from appearing on ArrowUp, F1, Shift (alone), etc. - if (event.key.length > 1 && event.key !== "Delete" && event.key !== "Backspace") { + if (event.key.length > 1 && !["Delete", "Backspace", "Process"].includes(event.key)) { return; } }