diff --git a/py/utils/usage_stats.py b/py/utils/usage_stats.py index d886ab9f..ba4ebfbb 100644 --- a/py/utils/usage_stats.py +++ b/py/utils/usage_stats.py @@ -57,6 +57,9 @@ class UsageStats: "last_save_time": 0 } + # Track if stats have been modified since last save + self._is_dirty = False + # Queue for prompt_ids to process self.pending_prompt_ids = set() @@ -180,27 +183,39 @@ class UsageStats: async def save_stats(self, force=False): """Save statistics to file""" try: - # Only save if it's been at least save_interval since last save or force is True + # Only save if: + # 1. force is True, OR + # 2. stats have been modified (is_dirty) AND save_interval has passed current_time = time.time() - if not force and (current_time - self.stats.get("last_save_time", 0)) < self.save_interval: - return False - + time_since_last_save = current_time - self.stats.get("last_save_time", 0) + + if not force: + if not self._is_dirty: + # No changes to save + return False + if time_since_last_save < self.save_interval: + # Too soon since last save + return False + # Use a lock to prevent concurrent writes async with self._lock: # Update last save time self.stats["last_save_time"] = current_time - + # Create directory if it doesn't exist os.makedirs(os.path.dirname(self._stats_file_path), exist_ok=True) - + # Write to a temporary file first, then move it to avoid corruption temp_path = f"{self._stats_file_path}.tmp" with open(temp_path, 'w', encoding='utf-8') as f: json.dump(self.stats, f, indent=2, ensure_ascii=False) - + # Replace the old file with the new one os.replace(temp_path, self._stats_file_path) - + + # Clear dirty flag since we've saved + self._is_dirty = False + logger.debug(f"Saved usage statistics to {self._stats_file_path}") return True except Exception as e: @@ -218,25 +233,32 @@ class UsageStats: while True: # Wait a short interval before checking for new prompt_ids await asyncio.sleep(5) # Check every 5 seconds - + # Process any pending prompt_ids if self.pending_prompt_ids: async with self._lock: # Get a copy of the set and clear original prompt_ids = self.pending_prompt_ids.copy() self.pending_prompt_ids.clear() - + # Process each prompt_id - registry = MetadataRegistry() - for prompt_id in prompt_ids: - try: - metadata = registry.get_metadata(prompt_id) - await self._process_metadata(metadata) - except Exception as e: - logger.error(f"Error processing prompt_id {prompt_id}: {e}") - - # Periodically save stats - await self.save_stats() + try: + registry = MetadataRegistry() + except NameError: + # MetadataRegistry not available (standalone mode) + registry = None + + if registry: + for prompt_id in prompt_ids: + try: + metadata = registry.get_metadata(prompt_id) + await self._process_metadata(metadata) + except Exception as e: + logger.error(f"Error processing prompt_id {prompt_id}: {e}") + + # Periodically save stats (only if there are changes) + if self._is_dirty: + await self.save_stats() except asyncio.CancelledError: # Task was cancelled, clean up await self.save_stats(force=True) @@ -254,9 +276,10 @@ class UsageStats: """Process metadata from an execution""" if not metadata or not isinstance(metadata, dict): return - + # Increment total executions count self.stats["total_executions"] += 1 + self._is_dirty = True # Get today's date in YYYY-MM-DD format today = datetime.datetime.now().strftime("%Y-%m-%d") @@ -373,7 +396,11 @@ class UsageStats: """Process a prompt execution immediately (synchronous approach)""" if not prompt_id: return - + + if standalone_mode: + # Usage statistics are not available in standalone mode + return + try: # Process metadata for this prompt_id registry = MetadataRegistry() diff --git a/web/comfyui/settings.js b/web/comfyui/settings.js index e1bdaa63..35b64486 100644 --- a/web/comfyui/settings.js +++ b/web/comfyui/settings.js @@ -16,6 +16,9 @@ const PROMPT_TAG_AUTOCOMPLETE_DEFAULT = true; const TAG_SPACE_REPLACEMENT_SETTING_ID = "loramanager.tag_space_replacement"; const TAG_SPACE_REPLACEMENT_DEFAULT = false; +const USAGE_STATISTICS_SETTING_ID = "loramanager.usage_statistics"; +const USAGE_STATISTICS_DEFAULT = true; + // ============================================================================ // Helper Functions // ============================================================================ @@ -124,6 +127,32 @@ const getTagSpaceReplacementPreference = (() => { }; })(); +const getUsageStatisticsPreference = (() => { + let settingsUnavailableLogged = false; + + return () => { + const settingManager = app?.extensionManager?.setting; + if (!settingManager || typeof settingManager.get !== "function") { + if (!settingsUnavailableLogged) { + console.warn("LoRA Manager: settings API unavailable, using default usage statistics setting."); + settingsUnavailableLogged = true; + } + return USAGE_STATISTICS_DEFAULT; + } + + try { + const value = settingManager.get(USAGE_STATISTICS_SETTING_ID); + return value ?? USAGE_STATISTICS_DEFAULT; + } catch (error) { + if (!settingsUnavailableLogged) { + console.warn("LoRA Manager: unable to read usage statistics setting, using default.", error); + settingsUnavailableLogged = true; + } + return USAGE_STATISTICS_DEFAULT; + } + }; +})(); + // ============================================================================ // Register Extension with All Settings // ============================================================================ @@ -168,6 +197,14 @@ app.registerExtension({ tooltip: "When enabled, tag names with underscores will have them replaced with spaces when inserted (e.g., 'blonde_hair' becomes 'blonde hair').", category: ["LoRA Manager", "Autocomplete", "Tag Formatting"], }, + { + id: USAGE_STATISTICS_SETTING_ID, + name: "Enable usage statistics tracking", + type: "boolean", + defaultValue: USAGE_STATISTICS_DEFAULT, + tooltip: "When enabled, LoRA Manager will track model usage statistics during workflow execution. Disabling this will prevent unnecessary disk writes.", + category: ["LoRA Manager", "Statistics", "Usage Tracking"], + }, ], }); @@ -175,4 +212,4 @@ app.registerExtension({ // Exports // ============================================================================ -export { getWheelSensitivity, getAutoPathCorrectionPreference, getPromptTagAutocompletePreference, getTagSpaceReplacementPreference }; +export { getWheelSensitivity, getAutoPathCorrectionPreference, getPromptTagAutocompletePreference, getTagSpaceReplacementPreference, getUsageStatisticsPreference }; diff --git a/web/comfyui/usage_stats.js b/web/comfyui/usage_stats.js index e4b8f28c..2099c879 100644 --- a/web/comfyui/usage_stats.js +++ b/web/comfyui/usage_stats.js @@ -2,7 +2,7 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; import { showToast } from "./utils.js"; -import { getAutoPathCorrectionPreference } from "./settings.js"; +import { getAutoPathCorrectionPreference, getUsageStatisticsPreference } from "./settings.js"; // Define target nodes and their widget configurations const PATH_CORRECTION_TARGETS = [ @@ -25,6 +25,11 @@ app.registerExtension({ setup() { // Listen for successful executions api.addEventListener("execution_success", ({ detail }) => { + // Skip if usage statistics is disabled + if (!getUsageStatisticsPreference()) { + return; + } + if (detail && detail.prompt_id) { this.updateUsageStats(detail.prompt_id); }