From e9e3ade3f01508da8242e2943c159f3f1f392b72 Mon Sep 17 00:00:00 2001 From: Setmaster Date: Wed, 14 Jan 2026 22:56:28 -0500 Subject: [PATCH] Implement Phase 1: core tracking and sidebar UI - Add model usage tracking via PromptExecutor hook - Support 12+ loader node types (checkpoints, LoRAs, VAEs, etc.) - Pattern matching for custom node loaders - Cross-platform file locking (Windows/Unix) - JSON storage with schema migrations - REST API endpoints for usage data - Sidebar panel with timeframe filtering and sorting - Stale model highlighting (>30 days unused) --- .gitignore | 3 + README.md | 90 ++++++++ __init__.py | 80 ++++++++ js/modelpulse.js | 511 ++++++++++++++++++++++++++++++++++++++++++++++ py/__init__.py | 1 + py/model_types.py | 92 +++++++++ py/routes.py | 138 +++++++++++++ py/storage.py | 258 +++++++++++++++++++++++ py/tracking.py | 263 ++++++++++++++++++++++++ 9 files changed, 1436 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 js/modelpulse.js create mode 100644 py/__init__.py create mode 100644 py/model_types.py create mode 100644 py/routes.py create mode 100644 py/storage.py create mode 100644 py/tracking.py diff --git a/.gitignore b/.gitignore index 4c5f206..173cbc6 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,4 @@ .claude/ +__pycache__/ +*.pyc +*.pyo diff --git a/README.md b/README.md new file mode 100644 index 0000000..4aabb55 --- /dev/null +++ b/README.md @@ -0,0 +1,90 @@ +# ModelPulse + +A ComfyUI extension that tracks model usage frequency, helping you identify abandoned or underutilized models to manage storage efficiently. + +![ComfyUI](https://img.shields.io/badge/ComfyUI-0.3.0+-blue) +![Python](https://img.shields.io/badge/Python-3.10+-green) +![Platform](https://img.shields.io/badge/Platform-Windows%20%7C%20Linux%20%7C%20macOS-lightgrey) + +## Features + +- **Automatic Tracking** - Monitors usage of all model types: checkpoints, LoRAs, VAEs, ControlNets, CLIP, UNETs, upscalers, and more +- **Sidebar Panel** - View usage statistics directly in ComfyUI's sidebar +- **Time Filtering** - View usage for all time, this month, or this week +- **Smart Sorting** - Sort by last used, usage count, or name +- **Stale Detection** - Models unused for 30+ days are highlighted +- **Category Grouping** - Models organized by type with collapsible sections +- **Zero Dependencies** - Uses only Python standard library + +## Installation + +### Option 1: ComfyUI Manager (Recommended) + +Search for "ModelPulse" in ComfyUI Manager and click Install. + +### Option 2: Manual Installation + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/your-username/ComfyUI-ModelPulse.git +``` + +Restart ComfyUI after installation. + +## Usage + +1. Click the **chart icon** in the ComfyUI sidebar to open ModelPulse +2. Run workflows as usual - model usage is tracked automatically +3. Use the tabs to filter by time period (All Time / Month / Week) +4. Use the dropdown to sort by Last Used, Usage Count, or Name +5. Click category headers to collapse/expand sections + +## Tracked Model Types + +| Category | Loader Nodes | +|----------|-------------| +| Checkpoints | CheckpointLoaderSimple, CheckpointLoader | +| LoRAs | LoraLoader, LoraLoaderModelOnly | +| VAEs | VAELoader | +| ControlNets | ControlNetLoader | +| CLIP | CLIPLoader, DualCLIPLoader | +| UNETs | UNETLoader | +| Upscalers | UpscaleModelLoader | +| Style Models | StyleModelLoader | +| GLIGEN | GLIGENLoader | + +Custom node loaders (Impact Pack, Efficiency Nodes, etc.) are also detected via pattern matching. + +## API Endpoints + +ModelPulse exposes REST endpoints for programmatic access: + +| Endpoint | Method | Description | +|----------|--------|-------------| +| `/modelpulse/usage` | GET | Get all usage data (supports `?timeframe=` and `?sort=` params) | +| `/modelpulse/model/{id}` | GET | Get detailed data for a specific model | +| `/modelpulse/categories` | GET | Get category list with counts | +| `/modelpulse/reset` | POST | Reset all tracking data (requires `{"confirm": true}`) | + +## Data Storage + +Usage data is stored in: +``` +ComfyUI/user/default/modelpulse/usage_data.json +``` + +The data persists across ComfyUI restarts and can be backed up or edited manually. + +## Requirements + +- ComfyUI 0.3.0 or later (Vue 3 frontend) +- Python 3.10+ +- No external Python dependencies + +## License + +MIT License - see [LICENSE](LICENSE) for details. + +## Contributing + +Contributions are welcome! Please feel free to submit issues and pull requests. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d85e6aa --- /dev/null +++ b/__init__.py @@ -0,0 +1,80 @@ +""" +ModelPulse - ComfyUI Model Usage Tracker + +Tracks model usage frequency to help identify abandoned or underutilized models. +""" + +from .py.tracking import ModelTracker +from .py.routes import setup_routes + +# Global tracker instance +tracker = ModelTracker() + +# Store original execute for restoration if needed +_original_execute = None + + +def init_tracking(): + """Initialize the tracking hook on PromptExecutor.""" + global _original_execute + + try: + from execution import PromptExecutor + except ImportError: + print("[ModelPulse] Warning: Could not import PromptExecutor. Tracking disabled.") + return + + _original_execute = PromptExecutor.execute + + def tracked_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]): + """Wrapper that tracks model usage before execution.""" + try: + # Extract and record models from the prompt + tracker.extract_and_record_models(prompt) + except Exception as e: + # Don't let tracking errors break execution + print(f"[ModelPulse] Error tracking models: {e}") + + # Call original execute (synchronous) + return _original_execute(self, prompt, prompt_id, extra_data, execute_outputs) + + PromptExecutor.execute = tracked_execute + print("[ModelPulse] Tracking initialized") + + +def init_routes(): + """Initialize API routes.""" + try: + from server import PromptServer + + if PromptServer.instance is None: + print("[ModelPulse] Warning: PromptServer not initialized. Routes disabled.") + return + + setup_routes(PromptServer.instance.routes, tracker) + print("[ModelPulse] API routes registered") + + except ImportError: + print("[ModelPulse] Warning: Could not import PromptServer. Routes disabled.") + except Exception as e: + print(f"[ModelPulse] Error registering routes: {e}") + + +# Initialize on module load +init_tracking() +init_routes() + + +# ComfyUI node registration (empty - we don't add any nodes) +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +# Tell ComfyUI where to find our JavaScript files +WEB_DIRECTORY = "./js" + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "WEB_DIRECTORY", + "tracker", +] diff --git a/js/modelpulse.js b/js/modelpulse.js new file mode 100644 index 0000000..e372e3d --- /dev/null +++ b/js/modelpulse.js @@ -0,0 +1,511 @@ +import { app } from "../../scripts/app.js"; + +/** + * ModelPulse - Usage Tracker Extension + * + * Displays model usage statistics in a sidebar panel. + */ + +// State management +let currentTimeframe = "all"; +let currentSort = "last_used"; +let usageData = null; +let isLoading = false; + +/** + * Fetch usage data from the API + */ +async function fetchUsageData(timeframe = "all", sortBy = "last_used") { + try { + const response = await fetch( + `/modelpulse/usage?timeframe=${timeframe}&sort=${sortBy}` + ); + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + return await response.json(); + } catch (error) { + console.error("[ModelPulse] Error fetching usage data:", error); + return null; + } +} + +/** + * Format ISO date to relative time string + */ +function formatLastUsed(isoDate) { + if (!isoDate) return "Never"; + + const date = new Date(isoDate); + const now = new Date(); + const diffMs = now - date; + const diffMins = Math.floor(diffMs / (1000 * 60)); + const diffHours = Math.floor(diffMs / (1000 * 60 * 60)); + const diffDays = Math.floor(diffHours / 24); + + if (diffMins < 1) return "just now"; + if (diffMins < 60) return `${diffMins}m ago`; + if (diffHours < 24) return `${diffHours}h ago`; + if (diffDays < 30) return `${diffDays}d ago`; + + const diffMonths = Math.floor(diffDays / 30); + if (diffMonths < 12) return `${diffMonths}mo ago`; + + const diffYears = Math.floor(diffDays / 365); + return `${diffYears}y ago`; +} + +/** + * Get category display info + */ +function getCategoryInfo(category) { + const categories = { + checkpoint: { name: "Checkpoints", icon: "📦" }, + lora: { name: "LoRAs", icon: "✨" }, + vae: { name: "VAEs", icon: "⚙️" }, + controlnet: { name: "ControlNets", icon: "🎛️" }, + clip: { name: "CLIP", icon: "💬" }, + unet: { name: "UNETs", icon: "🖥️" }, + upscaler: { name: "Upscalers", icon: "🔍" }, + style_model: { name: "Style Models", icon: "🎨" }, + gligen: { name: "GLIGEN", icon: "📐" }, + }; + return categories[category] || { name: category, icon: "📁" }; +} + +/** + * Check if a model is stale (unused for > 30 days) + */ +function isStale(lastUsed) { + if (!lastUsed) return true; + const date = new Date(lastUsed); + const now = new Date(); + const diffDays = Math.floor((now - date) / (1000 * 60 * 60 * 24)); + return diffDays > 30; +} + +/** + * Group models by category + */ +function groupByCategory(models) { + const groups = {}; + for (const model of models) { + const cat = model.category; + if (!groups[cat]) { + groups[cat] = []; + } + groups[cat].push(model); + } + return groups; +} + +/** + * Create the styles for the UI + */ +function createStyles() { + const style = document.createElement("style"); + style.textContent = ` + .modelpulse-container { + display: flex; + flex-direction: column; + height: 100%; + font-family: var(--comfy-font, sans-serif); + font-size: 13px; + color: var(--input-text, #ddd); + } + + .modelpulse-header { + display: flex; + justify-content: space-between; + align-items: center; + padding: 12px; + border-bottom: 1px solid var(--border-color, #444); + } + + .modelpulse-title { + font-size: 15px; + font-weight: 600; + } + + .modelpulse-refresh { + background: transparent; + border: none; + color: var(--input-text, #ddd); + cursor: pointer; + padding: 4px 8px; + border-radius: 4px; + font-size: 14px; + } + + .modelpulse-refresh:hover { + background: var(--comfy-menu-bg-hover, #333); + } + + .modelpulse-refresh.loading { + opacity: 0.5; + cursor: wait; + } + + .modelpulse-tabs { + display: flex; + border-bottom: 1px solid var(--border-color, #444); + } + + .modelpulse-tab { + flex: 1; + padding: 10px; + text-align: center; + cursor: pointer; + background: transparent; + border: none; + color: var(--input-text, #ddd); + border-bottom: 2px solid transparent; + transition: all 0.15s ease; + } + + .modelpulse-tab:hover { + background: var(--comfy-menu-bg-hover, #333); + } + + .modelpulse-tab.active { + border-bottom-color: var(--primary-color, #4a9eff); + color: var(--primary-color, #4a9eff); + } + + .modelpulse-controls { + display: flex; + padding: 8px 12px; + gap: 8px; + align-items: center; + border-bottom: 1px solid var(--border-color, #444); + } + + .modelpulse-sort-label { + font-size: 12px; + color: var(--descrip-text, #999); + } + + .modelpulse-sort { + flex: 1; + padding: 6px 8px; + background: var(--comfy-input-bg, #222); + border: 1px solid var(--border-color, #444); + border-radius: 4px; + color: var(--input-text, #ddd); + font-size: 12px; + cursor: pointer; + } + + .modelpulse-list { + flex: 1; + overflow-y: auto; + padding: 8px 0; + } + + .modelpulse-category { + margin-bottom: 4px; + } + + .modelpulse-category-header { + display: flex; + align-items: center; + gap: 8px; + padding: 8px 12px; + background: var(--comfy-menu-bg, #2a2a2a); + cursor: pointer; + user-select: none; + } + + .modelpulse-category-header:hover { + background: var(--comfy-menu-bg-hover, #333); + } + + .modelpulse-category-icon { + font-size: 14px; + } + + .modelpulse-category-name { + flex: 1; + font-weight: 500; + } + + .modelpulse-category-count { + font-size: 11px; + color: var(--descrip-text, #999); + background: var(--comfy-input-bg, #222); + padding: 2px 6px; + border-radius: 10px; + } + + .modelpulse-category-arrow { + font-size: 10px; + transition: transform 0.15s ease; + } + + .modelpulse-category.collapsed .modelpulse-category-arrow { + transform: rotate(-90deg); + } + + .modelpulse-category-models { + display: block; + } + + .modelpulse-category.collapsed .modelpulse-category-models { + display: none; + } + + .modelpulse-model { + display: flex; + flex-direction: column; + padding: 10px 12px 10px 32px; + border-bottom: 1px solid var(--border-color, #333); + cursor: default; + } + + .modelpulse-model:hover { + background: var(--comfy-menu-bg-hover, #333); + } + + .modelpulse-model-name { + font-weight: 500; + margin-bottom: 4px; + word-break: break-all; + } + + .modelpulse-model-stats { + display: flex; + gap: 16px; + font-size: 11px; + color: var(--descrip-text, #999); + } + + .modelpulse-model.stale .modelpulse-model-name { + color: var(--error-text, #ff6b6b); + } + + .modelpulse-model.stale .modelpulse-last-used { + color: var(--error-text, #ff6b6b); + } + + .modelpulse-empty { + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + height: 200px; + color: var(--descrip-text, #999); + text-align: center; + padding: 20px; + } + + .modelpulse-empty-icon { + font-size: 48px; + margin-bottom: 12px; + opacity: 0.5; + } + + .modelpulse-loading { + display: flex; + align-items: center; + justify-content: center; + height: 100px; + color: var(--descrip-text, #999); + } + `; + return style; +} + +/** + * Render the ModelPulse UI into the container element + */ +async function renderModelPulseUI(container) { + // Add styles + container.appendChild(createStyles()); + + // Create main container + const main = document.createElement("div"); + main.className = "modelpulse-container"; + container.appendChild(main); + + // Header + const header = document.createElement("div"); + header.className = "modelpulse-header"; + header.innerHTML = ` + Usage + + `; + main.appendChild(header); + + // Timeframe tabs + const tabs = document.createElement("div"); + tabs.className = "modelpulse-tabs"; + tabs.innerHTML = ` + + + + `; + main.appendChild(tabs); + + // Controls (sort dropdown) + const controls = document.createElement("div"); + controls.className = "modelpulse-controls"; + controls.innerHTML = ` + Sort: + + `; + main.appendChild(controls); + + // Model list container + const list = document.createElement("div"); + list.className = "modelpulse-list"; + main.appendChild(list); + + // Event handlers + const refreshBtn = header.querySelector(".modelpulse-refresh"); + refreshBtn.addEventListener("click", () => refreshData(list, refreshBtn)); + + tabs.addEventListener("click", (e) => { + const tab = e.target.closest(".modelpulse-tab"); + if (tab) { + currentTimeframe = tab.dataset.timeframe; + tabs.querySelectorAll(".modelpulse-tab").forEach(t => t.classList.remove("active")); + tab.classList.add("active"); + refreshData(list, refreshBtn); + } + }); + + const sortSelect = controls.querySelector(".modelpulse-sort"); + sortSelect.addEventListener("change", () => { + currentSort = sortSelect.value; + refreshData(list, refreshBtn); + }); + + // Initial load + await refreshData(list, refreshBtn); +} + +/** + * Refresh the data and re-render the list + */ +async function refreshData(listContainer, refreshBtn) { + if (isLoading) return; + + isLoading = true; + refreshBtn.classList.add("loading"); + listContainer.innerHTML = '
Loading...
'; + + try { + usageData = await fetchUsageData(currentTimeframe, currentSort); + renderModelList(listContainer); + } catch (error) { + listContainer.innerHTML = ` +
+
⚠️
+
Error loading data
+
+ `; + } finally { + isLoading = false; + refreshBtn.classList.remove("loading"); + } +} + +/** + * Render the model list grouped by category + */ +function renderModelList(container) { + container.innerHTML = ""; + + if (!usageData || !usageData.models || usageData.models.length === 0) { + container.innerHTML = ` +
+
📊
+
No model usage recorded yet
+
Run some workflows to start tracking
+
+ `; + return; + } + + const grouped = groupByCategory(usageData.models); + + // Sort categories by total usage + const sortedCategories = Object.keys(grouped).sort((a, b) => { + const aTotal = grouped[a].reduce((sum, m) => sum + m.timeframe_count, 0); + const bTotal = grouped[b].reduce((sum, m) => sum + m.timeframe_count, 0); + return bTotal - aTotal; + }); + + for (const category of sortedCategories) { + const models = grouped[category]; + const catInfo = getCategoryInfo(category); + + const catDiv = document.createElement("div"); + catDiv.className = "modelpulse-category"; + + // Category header + const catHeader = document.createElement("div"); + catHeader.className = "modelpulse-category-header"; + catHeader.innerHTML = ` + ${catInfo.icon} + ${catInfo.name} + ${models.length} + ▼ + `; + catHeader.addEventListener("click", () => { + catDiv.classList.toggle("collapsed"); + }); + catDiv.appendChild(catHeader); + + // Models list + const modelsDiv = document.createElement("div"); + modelsDiv.className = "modelpulse-category-models"; + + for (const model of models) { + const modelDiv = document.createElement("div"); + modelDiv.className = "modelpulse-model"; + if (isStale(model.last_used)) { + modelDiv.classList.add("stale"); + } + + const countLabel = currentTimeframe === "all" ? "Uses" : "Uses (period)"; + const count = currentTimeframe === "all" ? model.usage_count : model.timeframe_count; + + modelDiv.innerHTML = ` +
${model.name}
+
+ ${countLabel}: ${count} + Last: ${formatLastUsed(model.last_used)} +
+ `; + modelsDiv.appendChild(modelDiv); + } + + catDiv.appendChild(modelsDiv); + container.appendChild(catDiv); + } +} + +// Register the extension +app.registerExtension({ + name: "ModelPulse", + async setup() { + // Register sidebar tab + app.extensionManager.registerSidebarTab({ + id: "modelpulse", + icon: "pi pi-chart-bar", + title: "Usage", + tooltip: "Track model usage frequency", + type: "custom", + render: (el) => { + renderModelPulseUI(el); + } + }); + + console.log("[ModelPulse] Extension loaded"); + } +}); diff --git a/py/__init__.py b/py/__init__.py new file mode 100644 index 0000000..dcad39e --- /dev/null +++ b/py/__init__.py @@ -0,0 +1 @@ +"""ModelPulse Python package.""" diff --git a/py/model_types.py b/py/model_types.py new file mode 100644 index 0000000..cd105a3 --- /dev/null +++ b/py/model_types.py @@ -0,0 +1,92 @@ +""" +ModelPulse - Model Type Definitions + +Maps ComfyUI loader node types to their model categories and input keys. +""" + +# Standard ComfyUI model loaders +# Format: "NodeClassName": ("category", "input_key" or ["input_key1", "input_key2"]) +MODEL_LOADERS: dict[str, tuple[str, str | list[str]]] = { + # Checkpoint loaders + "CheckpointLoaderSimple": ("checkpoint", "ckpt_name"), + "CheckpointLoader": ("checkpoint", "ckpt_name"), + + # LoRA loaders + "LoraLoader": ("lora", "lora_name"), + "LoraLoaderModelOnly": ("lora", "lora_name"), + + # VAE loaders + "VAELoader": ("vae", "vae_name"), + + # ControlNet loaders + "ControlNetLoader": ("controlnet", "control_net_name"), + + # CLIP loaders + "CLIPLoader": ("clip", "clip_name"), + "DualCLIPLoader": ("clip", ["clip_name1", "clip_name2"]), + + # UNET loaders + "UNETLoader": ("unet", "unet_name"), + + # Upscaler loaders + "UpscaleModelLoader": ("upscaler", "model_name"), + + # Style model loaders + "StyleModelLoader": ("style_model", "style_model_name"), + + # GLIGEN loaders + "GLIGENLoader": ("gligen", "gligen_name"), +} + +# Patterns for detecting custom node loaders (Impact Pack, Efficiency Nodes, etc.) +# These are checked if the node class name isn't in MODEL_LOADERS +LOADER_PATTERNS: list[tuple[str, str, list[str]]] = [ + # (pattern_in_class_name, category, possible_input_keys) + ("Checkpoint", "checkpoint", ["ckpt_name", "checkpoint", "model_name"]), + ("Lora", "lora", ["lora_name", "lora"]), + ("VAE", "vae", ["vae_name", "vae"]), + ("ControlNet", "controlnet", ["control_net_name", "controlnet", "control_net"]), + ("CLIP", "clip", ["clip_name", "clip"]), + ("UNET", "unet", ["unet_name", "unet"]), + ("Upscale", "upscaler", ["model_name", "upscale_model"]), +] + +# Model categories for display purposes +MODEL_CATEGORIES: dict[str, dict[str, str]] = { + "checkpoint": { + "name": "Checkpoints", + "icon": "pi-box", + }, + "lora": { + "name": "LoRAs", + "icon": "pi-sparkles", + }, + "vae": { + "name": "VAEs", + "icon": "pi-cog", + }, + "controlnet": { + "name": "ControlNets", + "icon": "pi-sitemap", + }, + "clip": { + "name": "CLIP Models", + "icon": "pi-comments", + }, + "unet": { + "name": "UNETs", + "icon": "pi-server", + }, + "upscaler": { + "name": "Upscalers", + "icon": "pi-expand", + }, + "style_model": { + "name": "Style Models", + "icon": "pi-palette", + }, + "gligen": { + "name": "GLIGEN", + "icon": "pi-th-large", + }, +} diff --git a/py/routes.py b/py/routes.py new file mode 100644 index 0000000..2cfa679 --- /dev/null +++ b/py/routes.py @@ -0,0 +1,138 @@ +""" +ModelPulse - API Routes + +Custom HTTP routes for the ModelPulse extension. +""" + +from aiohttp import web +from urllib.parse import unquote + +from .tracking import ModelTracker + + +def setup_routes(routes, tracker: ModelTracker) -> None: + """ + Register ModelPulse API routes. + + Args: + routes: PromptServer.instance.routes + tracker: The ModelTracker instance + """ + + @routes.get("/modelpulse/usage") + async def get_usage_data(request: web.Request) -> web.Response: + """ + Get all model usage data. + + Query parameters: + timeframe: "all" (default), "month", or "week" + sort: "last_used" (default), "usage_count", or "name" + category: Optional category filter + """ + timeframe = request.query.get("timeframe", "all") + sort_by = request.query.get("sort", "last_used") + category = request.query.get("category") + + # Validate parameters + if timeframe not in ("all", "month", "week"): + timeframe = "all" + if sort_by not in ("last_used", "usage_count", "name"): + sort_by = "last_used" + + data = tracker.get_usage_data( + timeframe=timeframe, + sort_by=sort_by, + category=category, + ) + return web.json_response(data) + + @routes.get("/modelpulse/model/{model_id:.*}") + async def get_model_detail(request: web.Request) -> web.Response: + """ + Get detailed usage data for a specific model. + + Path parameters: + model_id: The model identifier (category/filename) + """ + model_id = unquote(request.match_info["model_id"]) + data = tracker.get_model_detail(model_id) + + if data is None: + return web.json_response( + {"error": "Model not found"}, + status=404, + ) + + return web.json_response(data) + + @routes.get("/modelpulse/categories") + async def get_categories(request: web.Request) -> web.Response: + """Get list of model categories with counts.""" + from .model_types import MODEL_CATEGORIES + + usage_data = tracker.get_usage_data() + category_counts: dict[str, int] = {} + + for model in usage_data["models"]: + cat = model["category"] + category_counts[cat] = category_counts.get(cat, 0) + 1 + + categories = [] + for cat_id, cat_info in MODEL_CATEGORIES.items(): + categories.append({ + "id": cat_id, + "name": cat_info["name"], + "icon": cat_info["icon"], + "count": category_counts.get(cat_id, 0), + }) + + # Sort by count descending + categories.sort(key=lambda x: x["count"], reverse=True) + + return web.json_response({"categories": categories}) + + @routes.post("/modelpulse/reset") + async def reset_tracking(request: web.Request) -> web.Response: + """ + Reset all tracking data. + + Requires confirmation in request body: {"confirm": true} + """ + try: + body = await request.json() + except Exception: + return web.json_response( + {"error": "Invalid JSON body"}, + status=400, + ) + + if body.get("confirm") is not True: + return web.json_response( + {"error": "Confirmation required. Send {\"confirm\": true}"}, + status=400, + ) + + tracker.reset() + return web.json_response({"status": "ok", "message": "Tracking data reset"}) + + @routes.post("/modelpulse/cleanup") + async def cleanup_old_data(request: web.Request) -> web.Response: + """ + Clean up old usage log entries. + + Optional body parameter: {"max_days": 365} + """ + try: + body = await request.json() + max_days = body.get("max_days", 365) + except Exception: + max_days = 365 + + if not isinstance(max_days, int) or max_days < 1: + max_days = 365 + + tracker.cleanup(max_days=max_days) + return web.json_response({ + "status": "ok", + "message": f"Cleaned up entries older than {max_days} days", + }) diff --git a/py/storage.py b/py/storage.py new file mode 100644 index 0000000..bfc2604 --- /dev/null +++ b/py/storage.py @@ -0,0 +1,258 @@ +""" +ModelPulse - Storage Layer + +JSON-based persistence for model usage data with cross-platform file locking. +""" + +import json +import sys +import time +from datetime import datetime +from pathlib import Path +from typing import Any + +# Cross-platform file locking +if sys.platform == "win32": + import msvcrt + + def _lock_file(f, exclusive: bool = False) -> None: + """Lock file on Windows using msvcrt.""" + msvcrt.locking(f.fileno(), msvcrt.LK_NBLCK if exclusive else msvcrt.LK_LOCK, 1) + + def _unlock_file(f) -> None: + """Unlock file on Windows using msvcrt.""" + try: + f.seek(0) + msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK, 1) + except OSError: + pass # Already unlocked or file closed +else: + import fcntl + + def _lock_file(f, exclusive: bool = False) -> None: + """Lock file on Unix using fcntl.""" + fcntl.flock(f.fileno(), fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH) + + def _unlock_file(f) -> None: + """Unlock file on Unix using fcntl.""" + fcntl.flock(f.fileno(), fcntl.LOCK_UN) + + +class FileLock: + """Cross-platform file lock context manager using a .lock file.""" + + def __init__(self, path: Path, timeout: float = 10.0): + self.lock_path = path.with_suffix(path.suffix + ".lock") + self.timeout = timeout + self._lock_file = None + + def __enter__(self): + start = time.time() + while True: + try: + # Create lock file exclusively + self._lock_file = open(self.lock_path, "x") + return self + except FileExistsError: + # Check if lock is stale (older than timeout) + try: + lock_age = time.time() - self.lock_path.stat().st_mtime + if lock_age > self.timeout: + # Stale lock, remove it + self.lock_path.unlink() + continue + except FileNotFoundError: + continue # Lock was released, retry + + if time.time() - start > self.timeout: + raise TimeoutError(f"Could not acquire lock on {self.lock_path}") + time.sleep(0.1) + + def __exit__(self, exc_type, exc_val, exc_tb): + if self._lock_file: + self._lock_file.close() + try: + self.lock_path.unlink() + except FileNotFoundError: + pass + + +SCHEMA_VERSION = 1 + + +def get_storage_path() -> Path: + """Get the path to the usage data file.""" + # ComfyUI stores user data in user/default/ + # We need to find the ComfyUI root directory + # This file is in custom_nodes/ComfyUI-ModelPulse/py/ + current_file = Path(__file__).resolve() + custom_nodes_dir = current_file.parent.parent.parent # Go up to custom_nodes + comfyui_root = custom_nodes_dir.parent # Go up to ComfyUI root + + storage_dir = comfyui_root / "user" / "default" / "modelpulse" + storage_dir.mkdir(parents=True, exist_ok=True) + + return storage_dir / "usage_data.json" + + +def get_backup_path() -> Path: + """Get the path for backup files.""" + storage_path = get_storage_path() + return storage_path.with_suffix(".backup.json") + + +def create_empty_data() -> dict[str, Any]: + """Create an empty data structure.""" + now = datetime.utcnow().isoformat() + "Z" + return { + "version": SCHEMA_VERSION, + "models": {}, + "metadata": { + "tracking_started": now, + "last_updated": now, + }, + } + + +def load_data() -> dict[str, Any]: + """ + Load usage data from JSON file. + + Creates the file with empty data if it doesn't exist. + Handles schema migrations if needed. + """ + storage_path = get_storage_path() + + if not storage_path.exists(): + data = create_empty_data() + save_data(data) + return data + + try: + with FileLock(storage_path): + with open(storage_path, "r", encoding="utf-8") as f: + data = json.load(f) + + # Handle schema migrations + data = migrate_schema(data) + return data + + except (json.JSONDecodeError, KeyError) as e: + # Corrupted file - backup and create fresh + print(f"[ModelPulse] Warning: Corrupted data file, creating backup. Error: {e}") + backup_corrupted_file(storage_path) + data = create_empty_data() + save_data(data) + return data + + +def save_data(data: dict[str, Any]) -> None: + """ + Save usage data to JSON file with file locking. + + Uses exclusive lock to prevent concurrent writes. + """ + storage_path = get_storage_path() + + # Update last_updated timestamp + data["metadata"]["last_updated"] = datetime.utcnow().isoformat() + "Z" + + # Write to temp file first, then rename (atomic operation) + temp_path = storage_path.with_suffix(".tmp") + + try: + with FileLock(storage_path): + with open(temp_path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + + # Atomic rename + temp_path.replace(storage_path) + + except Exception as e: + # Clean up temp file if it exists + if temp_path.exists(): + temp_path.unlink() + raise e + + +def migrate_schema(data: dict[str, Any]) -> dict[str, Any]: + """ + Migrate data to the current schema version. + + Creates a backup before migration. + """ + current_version = data.get("version", 0) + + if current_version == SCHEMA_VERSION: + return data + + # Backup before migration + backup_path = get_backup_path() + with open(backup_path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2) + print(f"[ModelPulse] Created backup at {backup_path} before schema migration") + + # Apply migrations sequentially + if current_version < 1: + data = migrate_to_v1(data) + + # Future migrations would go here: + # if current_version < 2: + # data = migrate_to_v2(data) + + data["version"] = SCHEMA_VERSION + return data + + +def migrate_to_v1(data: dict[str, Any]) -> dict[str, Any]: + """Migrate from pre-v1 (or no version) to v1 schema.""" + # Ensure required structure exists + if "models" not in data: + data["models"] = {} + + if "metadata" not in data: + now = datetime.utcnow().isoformat() + "Z" + data["metadata"] = { + "tracking_started": now, + "last_updated": now, + } + + # Ensure all models have required fields + for model_id, model_data in data.get("models", {}).items(): + if "usage_log" not in model_data: + model_data["usage_log"] = [] + if "first_used" not in model_data: + model_data["first_used"] = model_data.get( + "last_used", datetime.utcnow().isoformat() + "Z" + ) + + return data + + +def backup_corrupted_file(storage_path: Path) -> None: + """Backup a corrupted file for manual recovery.""" + if storage_path.exists(): + timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S") + backup_path = storage_path.with_suffix(f".corrupted_{timestamp}.json") + storage_path.rename(backup_path) + print(f"[ModelPulse] Backed up corrupted file to {backup_path}") + + +def cleanup_old_usage_logs(data: dict[str, Any], max_days: int = 365) -> dict[str, Any]: + """ + Clean up old daily usage logs to prevent unbounded growth. + + Keeps the last `max_days` days of daily data. + """ + cutoff_date = datetime.utcnow().date() + from datetime import timedelta + + cutoff_str = (cutoff_date - timedelta(days=max_days)).isoformat() + + for model_data in data.get("models", {}).values(): + usage_log = model_data.get("usage_log", []) + model_data["usage_log"] = [ + entry for entry in usage_log if entry.get("date", "") >= cutoff_str + ] + + return data diff --git a/py/tracking.py b/py/tracking.py new file mode 100644 index 0000000..262a7b8 --- /dev/null +++ b/py/tracking.py @@ -0,0 +1,263 @@ +""" +ModelPulse - Core Tracking Logic + +Extracts models from prompts and records usage statistics. +""" + +from datetime import datetime, timedelta +from typing import Any + +from .model_types import MODEL_LOADERS, LOADER_PATTERNS +from .storage import load_data, save_data, cleanup_old_usage_logs + + +class ModelTracker: + """Tracks model usage across ComfyUI workflow executions.""" + + def __init__(self): + self._data: dict[str, Any] | None = None + + @property + def data(self) -> dict[str, Any]: + """Lazy-load data from storage.""" + if self._data is None: + self._data = load_data() + return self._data + + def _save(self) -> None: + """Save current data to storage.""" + if self._data is not None: + save_data(self._data) + + def extract_models_from_prompt(self, prompt: dict[str, Any]) -> list[dict[str, str]]: + """ + Extract all model references from a ComfyUI prompt. + + Args: + prompt: The prompt dictionary from ComfyUI execution + + Returns: + List of dicts with 'category', 'name', and 'model_id' keys + """ + models = [] + seen = set() # Avoid duplicates within same prompt + + for node_id, node_data in prompt.items(): + class_type = node_data.get("class_type", "") + inputs = node_data.get("inputs", {}) + + # Check standard loaders first + if class_type in MODEL_LOADERS: + category, input_keys = MODEL_LOADERS[class_type] + extracted = self._extract_from_inputs(category, input_keys, inputs) + for model in extracted: + if model["model_id"] not in seen: + seen.add(model["model_id"]) + models.append(model) + + # Check pattern-based loaders for custom nodes + elif "Loader" in class_type: + extracted = self._extract_from_patterns(class_type, inputs) + for model in extracted: + if model["model_id"] not in seen: + seen.add(model["model_id"]) + models.append(model) + + return models + + def _extract_from_inputs( + self, category: str, input_keys: str | list[str], inputs: dict[str, Any] + ) -> list[dict[str, str]]: + """Extract model names from node inputs.""" + models = [] + + if isinstance(input_keys, str): + input_keys = [input_keys] + + for key in input_keys: + value = inputs.get(key) + if value and isinstance(value, str): + model_id = f"{category}/{value}" + models.append({ + "category": category, + "name": value, + "model_id": model_id, + }) + + return models + + def _extract_from_patterns( + self, class_type: str, inputs: dict[str, Any] + ) -> list[dict[str, str]]: + """Extract models using pattern matching for custom node loaders.""" + models = [] + + for pattern, category, possible_keys in LOADER_PATTERNS: + if pattern.lower() in class_type.lower(): + for key in possible_keys: + value = inputs.get(key) + if value and isinstance(value, str): + model_id = f"{category}/{value}" + models.append({ + "category": category, + "name": value, + "model_id": model_id, + }) + break # Found a match for this pattern + break # Only match first pattern + + return models + + def record_usage(self, models: list[dict[str, str]]) -> None: + """ + Record usage for a list of models. + + Args: + models: List of model dicts from extract_models_from_prompt + """ + if not models: + return + + now = datetime.utcnow() + now_iso = now.isoformat() + "Z" + today = now.date().isoformat() + + for model in models: + model_id = model["model_id"] + + if model_id not in self.data["models"]: + # New model - create entry + self.data["models"][model_id] = { + "category": model["category"], + "name": model["name"], + "path": model_id, # We don't have full path, use model_id + "first_used": now_iso, + "last_used": now_iso, + "usage_count": 0, + "usage_log": [], + } + + model_data = self.data["models"][model_id] + + # Update timestamps and count + model_data["last_used"] = now_iso + model_data["usage_count"] += 1 + + # Update daily usage log + usage_log = model_data["usage_log"] + if usage_log and usage_log[-1]["date"] == today: + usage_log[-1]["count"] += 1 + else: + usage_log.append({"date": today, "count": 1}) + + self._save() + + def extract_and_record_models(self, prompt: dict[str, Any]) -> None: + """ + Extract models from prompt and record their usage. + + This is the main entry point called from the PromptExecutor hook. + """ + models = self.extract_models_from_prompt(prompt) + self.record_usage(models) + + def get_usage_data( + self, + timeframe: str = "all", + sort_by: str = "last_used", + category: str | None = None, + ) -> dict[str, Any]: + """ + Get filtered and sorted usage data. + + Args: + timeframe: "all", "month", or "week" + sort_by: "last_used", "usage_count", or "name" + category: Optional category filter + + Returns: + Dict with 'models' list and 'metadata' + """ + models_list = [] + now = datetime.utcnow() + + # Calculate timeframe boundaries + if timeframe == "week": + cutoff = now - timedelta(days=7) + elif timeframe == "month": + cutoff = now - timedelta(days=30) + else: + cutoff = None + + cutoff_date = cutoff.date().isoformat() if cutoff else None + + for model_id, model_data in self.data["models"].items(): + # Filter by category if specified + if category and model_data["category"] != category: + continue + + # Calculate timeframe-specific usage count + if cutoff_date: + timeframe_count = sum( + entry["count"] + for entry in model_data.get("usage_log", []) + if entry["date"] >= cutoff_date + ) + else: + timeframe_count = model_data["usage_count"] + + models_list.append({ + "model_id": model_id, + "category": model_data["category"], + "name": model_data["name"], + "first_used": model_data["first_used"], + "last_used": model_data["last_used"], + "usage_count": model_data["usage_count"], + "timeframe_count": timeframe_count, + }) + + # Sort the results + if sort_by == "usage_count": + models_list.sort(key=lambda x: x["timeframe_count"], reverse=True) + elif sort_by == "name": + models_list.sort(key=lambda x: x["name"].lower()) + else: # last_used (default) + models_list.sort(key=lambda x: x["last_used"], reverse=True) + + return { + "models": models_list, + "metadata": self.data["metadata"], + "timeframe": timeframe, + "sort_by": sort_by, + } + + def get_model_detail(self, model_id: str) -> dict[str, Any] | None: + """ + Get detailed usage data for a specific model. + + Args: + model_id: The model identifier (category/filename) + + Returns: + Model data dict or None if not found + """ + model_data = self.data["models"].get(model_id) + if not model_data: + return None + + return { + "model_id": model_id, + **model_data, + } + + def reset(self) -> None: + """Reset all tracking data.""" + from .storage import create_empty_data + + self._data = create_empty_data() + self._save() + + def cleanup(self, max_days: int = 365) -> None: + """Clean up old usage logs.""" + self._data = cleanup_old_usage_logs(self.data, max_days) + self._save()