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

Merge pull request #12 from newtextdoc1111/dev

Browse Source 
Merge Dev

- Fix #11
- Added a function to check if a new CSV exists
This commit is contained in:
newtextdoc1111
2025-05-23 13:11:39 +09:00
committed by GitHub
parent a04653856e 1b0de5148f
commit b2899b7643
4 changed files with 286 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+4 -1
View File
@@ -15,4 +15,7 @@ node_modules/
# csv files
data/*.csv
data/.download
data/.download
# csv meta information
csv_meta.json
__init__.py
+2 -1
View File
@@ -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 = {}
modules/downloader.py
+279 -143
View File
@@ -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()
web/js/autocomplete.js
+1 -1
View File
@@ -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;
}
}