diff --git a/modules/api.py b/modules/api.py index 588de44..d1cf78f 100644 --- a/modules/api.py +++ b/modules/api.py @@ -1,12 +1,14 @@ import os +import json import server from aiohttp import web +from . import downloader as dl # Get the absolute path to the 'data' directory # __file__ is the path to the current script (api.py) # os.path.dirname(__file__) is the directory of the current script (modules) # os.path.join(..., '..', 'data') goes up one level and then into 'data' -DATA_DIR = os.path.join(os.path.dirname(__file__), '..', 'data') +DATA_DIR = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "data")) DANBOORU_PREFIX = 'danbooru' E621_PREFIX = 'e621' @@ -45,8 +47,6 @@ def get_csv_file_status(): 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: @@ -67,6 +67,28 @@ def get_csv_file_status(): # 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. + Returns the timestamp string if found, None otherwise. + """ + try: + if not os.path.exists(dl.CSV_META_FILE): + return None + + 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') @@ -135,3 +157,65 @@ 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') +async def force_check_csv_updates(request): + """ + Forces a check for CSV file updates from HuggingFace, ignoring cooldown. + This allows users to manually trigger an update check at any time. + """ + 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 + }) + + except Exception as e: + print(f"[Autocomplete-Plus] Error during forced check: {e}") + return web.json_response({ + "success": False, + "error": str(e) + }, status=500) + +@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. + Returns null if the file doesn't exist or if there's an error reading it. + """ + 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" + }) + + 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 + }) + else: + 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 diff --git a/modules/downloader.py b/modules/downloader.py index 81f41c8..1704679 100644 --- a/modules/downloader.py +++ b/modules/downloader.py @@ -211,10 +211,10 @@ class Downloader: os.makedirs(DATA_DIR, exist_ok=True) os.makedirs(TEMP_DOWNLOAD_DIR, exist_ok=True) - def _check_new_csv_from_hf_dataset(self, dataset_meta: dict, now_utc: datetime): + def _check_new_csv_from_hf_dataset(self, dataset_meta: dict, now_utc: datetime, force_check: bool = False): """Checks HuggingFace for file updates and updates metadata.""" perform_hf_check = True - if dataset_meta.get("last_remote_check_timestamp"): + if not force_check and dataset_meta.get("last_remote_check_timestamp"): try: last_check_dt = datetime.fromisoformat(dataset_meta["last_remote_check_timestamp"]) if now_utc - last_check_dt < timedelta(days=7): @@ -295,17 +295,20 @@ class Downloader: return None - def run_check_and_download(self): + def run_check_and_download(self, force_check: bool = False): """ 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. """ 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._check_new_csv_from_hf_dataset(dataset_meta, now_utc, force_check) self._download_csv_files_if_needed(dataset_meta) self._save_metadata() diff --git a/web/js/main.js b/web/js/main.js index ead0f80..376a2cc 100644 --- a/web/js/main.js +++ b/web/js/main.js @@ -1,4 +1,5 @@ import { app } from "/scripts/app.js"; +import { $el } from "/scripts/ui.js"; import { ComfyWidgets } from "/scripts/widgets.js"; import { settingValues } from "./settings.js"; import { loadCSS } from "./utils.js"; @@ -6,6 +7,14 @@ import { TagSource, fetchCsvList, initializeData } from "./data.js"; import { AutocompleteEventHandler } from "./autocomplete.js"; import { RelatedTagsEventHandler } from "./related-tags.js"; +// --- Constants --- +const id = "AutocompletePlus"; +const name = "Autocomplete Plus"; + +// --- Functions --- +/** + * Initialize event handlers for the autocomplete and related tags features. + */ function initializeEventHandlers() { const autocompleteEventHandler = new AutocompleteEventHandler(); const relatedTagsEventHandler = new RelatedTagsEventHandler(); @@ -121,14 +130,113 @@ function initializeEventHandlers() { } } -const id = "AutocompletePlus"; -const name = "Autocomplete Plus"; +/** + * Add Miscellaneous settings to the settings screen + */ +async function addExtraSettings() { + // Function to perform the update check + async function performUpdateCheck(checkButton, lastCheckSpan) { + checkButton.textContent = "Checking..."; + checkButton.disabled = true; + + app.extensionManager.toast.add({ + severity: "info", + summary: "Check new CSV", + detail: "Checking the CSV updates, see console for more details.", + life: 5000 + }); + + try { + const response = await fetch('/autocomplete-plus/csv/force-check-updates', { + method: 'POST', + headers: { + 'Content-Type': 'application/json' + } + }); + + const result = await response.json(); + if (result.success) { + // Update last check time display using the response data + if (result.last_check_time) { + const newLastCheckDate = new Date(result.last_check_time); + lastCheckSpan.textContent = "Last checked: " + newLastCheckDate.toLocaleString(); + } else { + lastCheckSpan.textContent = "Last checked: Never"; + } + } + } catch (error) { + console.error("[Autocomplete-Plus] Error during force check:", error); + } finally { + checkButton.textContent = "Check now"; + checkButton.disabled = false; + } + } + + // Fetch last check time from API + let lastCheckTimeText = "Loading..."; + try { + const response = await fetch('/autocomplete-plus/csv/last-check-time'); + const data = await response.json(); + + if (data.last_check_time) { + const lastCheckDate = new Date(data.last_check_time); + lastCheckTimeText = "Last checked: " + lastCheckDate.toLocaleString(); + } else { + lastCheckTimeText = "Last checked: Never"; + } + } catch (error) { + console.error("[Autocomplete-Plus] Error fetching last check time:", error); + lastCheckTimeText = "Last checked: Error loading"; + } + + // Add extra setting for checking new CSV updates + app.ui.settings.addSetting({ + id: id + ".check_new_csv", + defaultValue: null, + name: "Check CSV updates", + category: [name, "Misc", "Check new CSV"], + type: () => { + const lastCheckSpan = $el("span", { + textContent: lastCheckTimeText, + className: "text-sm text-gray-500", + style: { + marginRight: "16px" + } + }); + + const checkButton = $el("button", { + textContent: "Check now", + className: "p-button p-component p-button-primary", + onclick: async () => { + await performUpdateCheck(checkButton, lastCheckSpan); + } + }); + + return $el("div", { + className: "flex-row items-center gap-2", + }, [ + $el("div", { + className: "p-component", + }, [ + lastCheckSpan, + checkButton, + ]), + ]); + } + }); +} + +/** + * Registration of the extension + */ app.registerExtension({ id: id, name: name, async setup() { initializeEventHandlers(); + addExtraSettings(); + let rootPath = import.meta.url.replace("js/main.js", ""); loadCSS(rootPath + "css/autocomplete-plus.css"); // Load CSS for autocomplete