diff --git a/README.md b/README.md index ccbc452..2630f35 100644 --- a/README.md +++ b/README.md @@ -10,10 +10,12 @@ A ComfyUI extension that tracks model usage frequency, helping you identify aban - **Automatic Tracking** - Monitors usage of all model types: checkpoints, LoRAs, VAEs, ControlNets, CLIP, UNETs, upscalers, GGUF, and more - **Sidebar Panel** - View usage statistics directly in ComfyUI's sidebar +- **File Size Display** - Shows model file sizes (GB/MB) to help identify storage-heavy models - **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 +- **Stale Detection** - Models unused for 30+ days are highlighted (configurable threshold) - **Category Grouping** - Models organized by type with collapsible sections +- **Reset Tracking** - Clear all usage data from the settings panel - **Zero Dependencies** - Uses only Python standard library ## Installation diff --git a/js/modelpulse.js b/js/modelpulse.js index 74d35e4..7c4ec29 100644 --- a/js/modelpulse.js +++ b/js/modelpulse.js @@ -61,6 +61,21 @@ async function fetchUsageData(timeframe = "all", sortBy = "last_used") { } } +/** + * Format file size to human readable string (GB or MB) + */ +function formatFileSize(bytes) { + if (bytes === null || bytes === undefined) return null; + + const gb = bytes / (1024 * 1024 * 1024); + if (gb >= 1) { + return `${gb.toFixed(1)} GB`; + } + + const mb = bytes / (1024 * 1024); + return `${mb.toFixed(0)} MB`; +} + /** * Format ISO date to relative time string */ @@ -761,12 +776,14 @@ function renderModelList(container) { const countLabel = currentTimeframe === "all" ? "Uses" : "Uses (period)"; const count = currentTimeframe === "all" ? model.usage_count : model.timeframe_count; + const sizeStr = formatFileSize(model.file_size); modelDiv.innerHTML = `
${model.name}
${countLabel}: ${count} Last: ${formatLastUsed(model.last_used)} + ${sizeStr ? `${sizeStr}` : ""}
`; modelsDiv.appendChild(modelDiv); diff --git a/py/model_types.py b/py/model_types.py index 097451a..f96ff99 100644 --- a/py/model_types.py +++ b/py/model_types.py @@ -102,3 +102,17 @@ MODEL_CATEGORIES: dict[str, dict[str, str]] = { "icon": "pi-database", }, } + +# Map categories to ComfyUI model folder names +CATEGORY_FOLDERS: dict[str, str] = { + "checkpoint": "checkpoints", + "lora": "loras", + "vae": "vae", + "controlnet": "controlnet", + "clip": "clip", + "unet": "unet", + "upscaler": "upscale_models", + "style_model": "style_models", + "gligen": "gligen", + "gguf": "unet", # GGUF models are stored in unet folder +} diff --git a/py/tracking.py b/py/tracking.py index 262a7b8..811235a 100644 --- a/py/tracking.py +++ b/py/tracking.py @@ -5,11 +5,59 @@ Extracts models from prompts and records usage statistics. """ from datetime import datetime, timedelta +from pathlib import Path from typing import Any -from .model_types import MODEL_LOADERS, LOADER_PATTERNS +from .model_types import MODEL_LOADERS, LOADER_PATTERNS, CATEGORY_FOLDERS from .storage import load_data, save_data, cleanup_old_usage_logs +# Try to import ComfyUI's folder_paths for proper model path resolution +try: + import folder_paths + HAS_FOLDER_PATHS = True +except ImportError: + HAS_FOLDER_PATHS = False + + +def get_model_file_size(category: str, name: str) -> int | None: + """ + Get the file size of a model in bytes. + + Uses ComfyUI's folder_paths to find models in all configured locations. + + Args: + category: The model category (e.g., 'checkpoint', 'lora') + name: The model filename + + Returns: + File size in bytes, or None if file not found + """ + if not HAS_FOLDER_PATHS: + return None + + # Map our categories to ComfyUI's folder_paths keys + folder_key = CATEGORY_FOLDERS.get(category) + + # For GGUF, try multiple folders since models can be unet or clip + if category == "gguf": + folder_keys = ["unet", "diffusion_models", "clip"] + elif category == "unet": + folder_keys = ["unet", "diffusion_models"] + elif folder_key: + folder_keys = [folder_key] + else: + return None + + for key in folder_keys: + try: + full_path = folder_paths.get_full_path(key, name) + if full_path: + return Path(full_path).stat().st_size + except Exception: + continue + + return None + class ModelTracker: """Tracks model usage across ComfyUI workflow executions.""" @@ -206,6 +254,9 @@ class ModelTracker: else: timeframe_count = model_data["usage_count"] + # Get file size + file_size = get_model_file_size(model_data["category"], model_data["name"]) + models_list.append({ "model_id": model_id, "category": model_data["category"], @@ -214,6 +265,7 @@ class ModelTracker: "last_used": model_data["last_used"], "usage_count": model_data["usage_count"], "timeframe_count": timeframe_count, + "file_size": file_size, }) # Sort the results