Compare commits

..

16 Commits

Author SHA1 Message Date
Ethanfel 777e4890ae docs: update README for model tracking feature
Document new Models tab, /nodes-stats/models endpoint, model_usage
table, updated How It Works diagram, and file structure.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-09 18:43:21 +02:00
Ethanfel 9477777124 fix: scope input-scanning loop inside per-node try/except in extract_models_from_prompt
If INPUT_TYPES() succeeded but returned a non-dict, the inner loop
calling input_types.get() would raise AttributeError outside the guard,
aborting model tracking for the entire prompt. Moving the loop inside
the try/except means only the offending node is skipped via continue.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-09 18:34:10 +02:00
Ethanfel 831ed302e9 fix: update model_type on upsert conflict and scope close button to dialog
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-09 18:20:18 +02:00
Ethanfel c7003df0bb fix: scope badge querySelector to dialog, update dialog title 2026-04-09 18:14:01 +02:00
Ethanfel bf5c8683a7 feat: add Models tab with per-type usage stats 2026-04-09 18:11:52 +02:00
Ethanfel 18d098a0da fix: add independent exception handling to _record_prompt 2026-04-09 18:08:53 +02:00
Ethanfel ed39c5918a feat: extend prompt handler and add /nodes-stats/models endpoint 2026-04-09 18:04:16 +02:00
Ethanfel 38e95c150a fix: atomic ModelMapper build and isolate test NODE_CLASS_MAPPINGS mutations
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-09 18:03:16 +02:00
Ethanfel 9891b3ec41 feat: add ModelMapper with folder_paths introspection
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-09 18:00:06 +02:00
Ethanfel ff42016009 test: add tier coverage for safe_to_remove and consider_removing 2026-04-09 17:58:26 +02:00
Ethanfel 800692d190 feat: add model_usage table and tracker methods
Adds model_usage schema, record_model_usage(), get_raw_model_stats(), and get_model_stats() to UsageTracker, with reset() support and 8 passing tests.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-09 17:55:27 +02:00
Ethanfel 70edeb0631 fix: configure folder_paths mock to prevent disk pollution in tests 2026-04-09 17:53:18 +02:00
Ethanfel 3ed778a0d6 test: add conftest with ComfyUI module stubs 2026-04-09 17:47:16 +02:00
Ethanfel b852c37786 Fix 3 bugs in model tracking plan
- Add conftest.py task to stub ComfyUI modules for tests
- Move extract_models_from_prompt into background thread
- Replace window.nsShowTab with local switchTab + addEventListener

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-09 17:43:08 +02:00
Ethanfel 901e3b722d Add model tracking implementation plan
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-09 17:39:54 +02:00
Ethanfel ab204056a0 Add model tracking feature design doc
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-09 17:36:36 +02:00
10 changed files with 1832 additions and 98 deletions
+37 -23
View File
@@ -4,14 +4,16 @@
<img src="docs/logo.svg" width="120" alt="Node Stats logo"> <img src="docs/logo.svg" width="120" alt="Node Stats logo">
</p> </p>
A ComfyUI custom node package that silently tracks which nodes and packages you actually use. Helps identify unused packages that are safe to remove — keeping your ComfyUI install lean. A ComfyUI custom node package that silently tracks which nodes, packages, and model files you actually use. Helps identify unused packages and models that are safe to remove — keeping your ComfyUI install lean.
## Features ## Features
- **Silent tracking** — hooks into every prompt submission, zero config needed - **Silent tracking** — hooks into every prompt submission, zero config needed
- **Node & model tracking** — tracks custom node packages and model files (checkpoints, VAEs, ControlNets, etc.) separately
- **Per-package classification** — packages are sorted into tiers based on usage recency - **Per-package classification** — packages are sorted into tiers based on usage recency
- **Smart aging** — packages gradually move from "recently unused" to "safe to remove" over time - **Per-model classification** — models grouped by type (checkpoints, vae, …) with the same recency tiers
- **Uninstall detection** — removed packages are flagged separately, historical data preserved - **Smart aging** — items gradually move from "recently unused" to "safe to remove" over time
- **Uninstall detection** — removed packages/models are flagged separately, historical data preserved
- **Expandable detail** — click any package to see individual node-level stats - **Expandable detail** — click any package to see individual node-level stats
- **Non-blocking** — DB writes happen in a background thread, no impact on workflow execution - **Non-blocking** — DB writes happen in a background thread, no impact on workflow execution
@@ -55,11 +57,17 @@ Restart ComfyUI. Tracking starts immediately and silently.
### UI ### UI
Click the **Node Stats** button (bar chart icon) in the ComfyUI top menu bar. A dialog shows: Click the **Node Stats** button (bar chart icon) in the ComfyUI top menu bar. A dialog opens with two tabs:
- **Summary bar** with counts for each classification tier **Nodes tab**
- **Sections** for each tier, sorted from most actionable to least - Summary bar with counts for each classification tier
- **Expandable rows** — click any package to see per-node execution counts and timestamps - Sections for each tier, sorted from most actionable to least
- Expandable rows — click any package to see per-node execution counts and timestamps
**Models tab**
- Summary bar with counts for each tier across all model types
- Sections per model type (checkpoints, vae, controlnet, …)
- Per-model table showing execution count, last used date, and status
### API ### API
@@ -67,6 +75,7 @@ Click the **Node Stats** button (bar chart icon) in the ComfyUI top menu bar. A
|----------|--------|-------------| |----------|--------|-------------|
| `/nodes-stats/packages` | GET | Per-package aggregated stats with classification | | `/nodes-stats/packages` | GET | Per-package aggregated stats with classification |
| `/nodes-stats/usage` | GET | Raw per-node usage data | | `/nodes-stats/usage` | GET | Raw per-node usage data |
| `/nodes-stats/models` | GET | Per-type model stats with classification |
| `/nodes-stats/reset` | POST | Clear all tracked data | | `/nodes-stats/reset` | POST | Clear all tracked data |
```bash ```bash
@@ -114,42 +123,46 @@ curl http://localhost:8188/nodes-stats/packages | python3 -m json.tool
## How It Works ## How It Works
``` ```
Queue Prompt ──> Prompt Handler ──> Extract class_types ──> Background Thread Queue Prompt ──> Prompt Handler ──> Extract class_types + prompt ──> Background Thread
─────────────────────────┘ ┌───────────────────────────────┘
SQLite DB SQLite DB
usage_stats.db usage_stats.db
──────────┐ ┌─────────────┐
node_usage │ per-node counts & timestamps node_usage │ per-node counts & timestamps
prompt_log │ full node list per prompt prompt_log │ full node list per prompt
└──────────┘ │ model_usage │ per-model counts & timestamps
└─────────────┘
GET /nodes-stats/packages ◄── GET /nodes-stats/packages ◄─────────┤
GET /nodes-stats/models ◄─────────┘
Mapper merges DB data Merge DB data with installed
with NODE_CLASS_MAPPINGS nodes/models, classify by recency
Classify by recency ──> JSON response ──> UI Dialog JSON response ──> UI Dialog (Nodes tab / Models tab)
``` ```
1. Registers a prompt handler via `PromptServer.instance.add_on_prompt_handler()` 1. Registers a prompt handler via `PromptServer.instance.add_on_prompt_handler()`
2. On every prompt submission, extracts `class_type` from each node in the workflow 2. On every prompt submission, extracts `class_type` from each node and the full prompt dict
3. Offloads recording to a background thread (non-blocking) 3. Offloads recording to a background thread (non-blocking)
4. Maps each class_type to its source package using `RELATIVE_PYTHON_MODULE` 4. Maps each class_type to its source package using `RELATIVE_PYTHON_MODULE`
5. Upserts per-node counts and timestamps into SQLite 5. Detects model file selections by introspecting each node's `INPUT_TYPES()` for folder-dropdown inputs, then resolves filenames via `folder_paths`
6. On stats request, merges DB data with current node registry and classifies by recency 6. Upserts per-node and per-model counts and timestamps into SQLite
7. On stats request, merges DB data with current node registry / installed models and classifies by recency
## Data Storage ## Data Storage
All data is stored in `usage_stats.db` in the package directory. All data is stored in `<ComfyUI user dir>/nodes_stats/usage_stats.db` (survives extension reinstalls).
| Table | Contents | | Table | Contents |
|-------|----------| |-------|----------|
| `node_usage` | Per-node: class_type, package, execution count, first/last seen | | `node_usage` | Per-node: class_type, package, execution count, first/last seen |
| `prompt_log` | Per-prompt: timestamp, JSON array of all class_types used | | `prompt_log` | Per-prompt: timestamp, JSON array of all class_types used |
| `model_usage` | Per-model: filename, type, execution count, first/last seen |
Use `POST /nodes-stats/reset` to clear all data and start fresh. Use `POST /nodes-stats/reset` to clear all data and start fresh.
@@ -157,8 +170,9 @@ Use `POST /nodes-stats/reset` to clear all data and start fresh.
``` ```
__init__.py Entry point: prompt handler, API routes __init__.py Entry point: prompt handler, API routes
mapper.py class_type → package name mapping mapper.py class_type → package mapping; model filename → type mapping
tracker.py SQLite persistence and stats aggregation tracker.py SQLite persistence and stats aggregation
js/nodes_stats.js Frontend: menu button + stats dialog js/nodes_stats.js Frontend: menu button + stats dialog (Nodes/Models tabs)
pyproject.toml Package metadata pyproject.toml Package metadata
tests/ Unit tests for tracker and mapper
``` ```
+32 -4
View File
@@ -4,7 +4,7 @@ import threading
from aiohttp import web from aiohttp import web
from server import PromptServer from server import PromptServer
from .mapper import NodePackageMapper from .mapper import NodePackageMapper, ModelMapper
from .tracker import UsageTracker from .tracker import UsageTracker
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -14,10 +14,11 @@ WEB_DIRECTORY = "./js"
mapper = NodePackageMapper() mapper = NodePackageMapper()
tracker = UsageTracker() tracker = UsageTracker()
model_mapper = ModelMapper()
def on_prompt_handler(json_data): def on_prompt_handler(json_data):
"""Called on every prompt submission. Extracts class_types and records usage.""" """Called on every prompt submission. Extracts class_types and queues recording."""
try: try:
prompt = json_data.get("prompt", {}) prompt = json_data.get("prompt", {})
class_types = set() class_types = set()
@@ -26,9 +27,11 @@ def on_prompt_handler(json_data):
if ct: if ct:
class_types.add(ct) class_types.add(ct)
if class_types: if class_types:
# Pass the full prompt to the thread — model extraction (which calls
# INPUT_TYPES() on every node) happens off the main request thread.
threading.Thread( threading.Thread(
target=tracker.record_usage, target=_record_prompt,
args=(class_types, mapper), args=(class_types, prompt),
daemon=True, daemon=True,
).start() ).start()
except Exception: except Exception:
@@ -36,6 +39,19 @@ def on_prompt_handler(json_data):
return json_data return json_data
def _record_prompt(class_types, prompt):
try:
tracker.record_usage(class_types, mapper)
except Exception:
logger.warning("nodes-stats: error recording node usage", exc_info=True)
try:
models = model_mapper.extract_models_from_prompt(prompt)
if models:
tracker.record_model_usage(models)
except Exception:
logger.warning("nodes-stats: error recording model usage", exc_info=True)
PromptServer.instance.add_on_prompt_handler(on_prompt_handler) PromptServer.instance.add_on_prompt_handler(on_prompt_handler)
@@ -62,11 +78,23 @@ async def get_node_stats(request):
return web.json_response({"error": "internal error"}, status=500) return web.json_response({"error": "internal error"}, status=500)
@routes.get("/nodes-stats/models")
async def get_model_stats(request):
try:
installed_by_type = model_mapper.get_all_models()
stats = tracker.get_model_stats(installed_by_type)
return web.json_response(stats)
except Exception:
logger.error("nodes-stats: error getting model stats", exc_info=True)
return web.json_response({"error": "internal error"}, status=500)
@routes.post("/nodes-stats/reset") @routes.post("/nodes-stats/reset")
async def reset_stats(request): async def reset_stats(request):
try: try:
tracker.reset() tracker.reset()
mapper.invalidate() mapper.invalidate()
model_mapper.invalidate()
return web.json_response({"status": "ok"}) return web.json_response({"status": "ok"})
except Exception: except Exception:
logger.error("nodes-stats: error resetting stats", exc_info=True) logger.error("nodes-stats: error resetting stats", exc_info=True)
@@ -0,0 +1,117 @@
# Model Usage Tracking — Design
**Date:** 2026-04-08
**Status:** Approved
## Goal
Extend ComfyUI Node Stats to also track which model files are used across prompts, and surface models that are never used — making it easy to identify files that can be safely deleted.
**Scope:** All model folder types registered in `folder_paths` except LoRAs (and non-model folders: `configs`, `custom_nodes`, `temp`, `output`, `input`). Discovered dynamically so custom-node-added types are picked up automatically.
## Design Decisions
- **Tracked per filename globally** — `dreamshaper.safetensors` is one entry regardless of which node loads it. `model_type` is stored for display grouping only.
- **Detection via node introspection (Approach A)** — At prompt time, for each node look up its `INPUT_TYPES()`, find inputs whose type is a list (ComfyUI folder dropdown pattern), cross-reference with `folder_paths` to identify the folder type, extract the selected value.
- **UI: new tab** in the existing dialog alongside the current "Nodes" tab.
- **Same tier classification** as packages: `used`, `unused_new`, `consider_removing`, `safe_to_remove`, `uninstalled`.
- **Existing node stats are unaffected** — additive only.
## Data Collection
In `on_prompt_handler` (`__init__.py`), after extracting `class_types`, also extract model references:
1. For each node in the prompt, look up its class in `nodes.NODE_CLASS_MAPPINGS`
2. Call `INPUT_TYPES()` and find inputs whose declared type is a `list`
3. Map that list back to its `folder_paths` folder type (via reverse lookup built at startup)
4. Extract the actual selected value from the prompt's `inputs` dict
5. Record `(model_name, model_type)` via background thread
A `ModelMapper` class (in `mapper.py` or new file) handles the reverse lookup: `folder_type → set of filenames`, cached and invalidated on reset.
Excluded folder types (not model files):
```python
EXCLUDED_FOLDER_TYPES = {"loras", "configs", "custom_nodes", "temp", "output", "input", "upscale_models"}
```
> Note: `upscale_models` may be included/excluded per preference; default include.
## Storage
New table in existing `usage_stats.db`:
```sql
CREATE TABLE IF NOT EXISTS model_usage (
model_name TEXT PRIMARY KEY,
model_type TEXT NOT NULL,
count INTEGER NOT NULL DEFAULT 0,
first_seen TEXT NOT NULL,
last_seen TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_model_usage_type ON model_usage(model_type);
```
Same upsert pattern as `node_usage`. `reset` clears this table too.
## API
New endpoint:
```
GET /nodes-stats/models
```
Returns a list of model types, each containing classified models:
```json
[
{
"model_type": "checkpoints",
"models": [
{
"model_name": "dreamshaper.safetensors",
"count": 42,
"first_seen": "2026-01-01T00:00:00+00:00",
"last_seen": "2026-03-01T00:00:00+00:00",
"installed": true,
"status": "used"
},
{
"model_name": "old_model.ckpt",
"count": 0,
"first_seen": null,
"last_seen": null,
"installed": true,
"status": "safe_to_remove"
}
]
}
]
```
"All installed models" sourced from `folder_paths.get_filename_list(type)` for each tracked type. Models on disk but never seen appear with `count: 0`.
Existing endpoints unchanged:
- `GET /nodes-stats/packages`
- `GET /nodes-stats/usage`
- `POST /nodes-stats/reset` — extended to also clear `model_usage`
## UI
Two tabs at top of dialog: **Nodes** (existing content, untouched) | **Models** (new).
Models tab layout:
- Summary badge bar: counts per status tier across all model types
- One section per model type (only types with ≥1 model shown), titled e.g. "Checkpoints", "VAE", "ControlNet"
- Within each section: models sorted by status tier (safe_to_remove → consider_removing → unused_new → used), then alphabetically
- Each row: model name | execution count | last used date | status color
- No expandable rows (models are leaves)
- Uninstalled models: collapsed section at the bottom, same pattern as node packages
## File Changes
| File | Change |
|------|--------|
| `tracker.py` | Add `model_usage` table to schema; add `record_model_usage()`, `get_model_stats()` methods; extend `reset()` |
| `mapper.py` | Add `ModelMapper` class with folder-type reverse lookup and model filename introspection |
| `__init__.py` | Extend `on_prompt_handler` to extract and record model usage; add `GET /nodes-stats/models` endpoint |
| `js/nodes_stats.js` | Add tab switcher UI; add Models tab rendering (summary badges + per-type sections) |
File diff suppressed because it is too large Load Diff
+162 -62
View File
@@ -34,23 +34,27 @@ app.registerExtension({
}); });
async function showStatsDialog() { async function showStatsDialog() {
let data; let data, modelData;
try { try {
const resp = await fetch("/nodes-stats/packages"); const [pkgResp, modelResp] = await Promise.all([
if (!resp.ok) { fetch("/nodes-stats/packages"),
alert("Failed to load node stats: HTTP " + resp.status); fetch("/nodes-stats/models"),
return; ]);
} if (!pkgResp.ok) { alert("Failed to load node stats: HTTP " + pkgResp.status); return; }
data = await resp.json(); if (!modelResp.ok) { alert("Failed to load model stats: HTTP " + modelResp.status); return; }
if (!Array.isArray(data)) { data = await pkgResp.json();
alert("Failed to load node stats: unexpected response format"); modelData = await modelResp.json();
if (!Array.isArray(data) || !Array.isArray(modelData)) {
alert("Failed to load stats: unexpected response format");
return; return;
} }
} catch (e) { } catch (e) {
alert("Failed to load node stats: " + e.message); alert("Failed to load stats: " + e.message);
return; return;
} }
const custom = data.filter((p) => p.package !== "__builtin__");
// Remove existing dialog if any // Remove existing dialog if any
const existing = document.getElementById("nodes-stats-dialog"); const existing = document.getElementById("nodes-stats-dialog");
if (existing) existing.remove(); if (existing) existing.remove();
@@ -67,69 +71,55 @@ async function showStatsDialog() {
dialog.style.cssText = dialog.style.cssText =
"background:#1e1e1e;color:#ddd;border-radius:8px;padding:24px;max-width:800px;width:90%;max-height:85vh;overflow-y:auto;font-family:monospace;font-size:13px;"; "background:#1e1e1e;color:#ddd;border-radius:8px;padding:24px;max-width:800px;width:90%;max-height:85vh;overflow-y:auto;font-family:monospace;font-size:13px;";
const custom = data.filter((p) => p.package !== "__builtin__");
const safeToRemove = custom.filter((p) => p.status === "safe_to_remove");
const considerRemoving = custom.filter((p) => p.status === "consider_removing");
const unusedNew = custom.filter((p) => p.status === "unused_new");
const used = custom.filter((p) => p.status === "used");
const uninstalled = custom.filter((p) => p.status === "uninstalled");
let html = `<div style="display:flex;justify-content:space-between;align-items:center;margin-bottom:16px;"> let html = `<div style="display:flex;justify-content:space-between;align-items:center;margin-bottom:16px;">
<h2 style="margin:0;color:#fff;font-size:18px;">Node Package Stats</h2> <h2 style="margin:0;color:#fff;font-size:18px;">Usage Stats</h2>
<button id="nodes-stats-close" style="background:none;border:none;color:#888;font-size:20px;cursor:pointer;">&times;</button> <button id="nodes-stats-close" style="background:none;border:none;color:#888;font-size:20px;cursor:pointer;">&times;</button>
</div>`; </div>`;
html += `<div style="display:flex;gap:10px;margin-bottom:20px;flex-wrap:wrap;"> // Tab switcher — wired via addEventListener after insertion, no onclick globals
<div style="background:#3a1a1a;padding:8px 14px;border-radius:4px;border-left:3px solid #e44;"> html += `
<span style="font-size:20px;font-weight:bold;color:#e44;">${safeToRemove.length}</span> <div style="display:flex;gap:0;margin-bottom:20px;border-bottom:1px solid #333;">
<span style="color:#c99;margin-left:6px;">safe to remove</span> <button id="ns-tab-nodes"
</div> style="background:none;border:none;border-bottom:2px solid #4a4;color:#4a4;padding:8px 18px;cursor:pointer;font-family:monospace;font-size:13px;font-weight:bold;">
<div style="background:#2a2215;padding:8px 14px;border-radius:4px;border-left:3px solid #e90;"> Nodes
<span style="font-size:20px;font-weight:bold;color:#e90;">${considerRemoving.length}</span> </button>
<span style="color:#ca8;margin-left:6px;">consider removing</span> <button id="ns-tab-models"
</div> style="background:none;border:none;border-bottom:2px solid transparent;color:#888;padding:8px 18px;cursor:pointer;font-family:monospace;font-size:13px;">
<div style="background:#1a1a2a;padding:8px 14px;border-radius:4px;border-left:3px solid #68f;"> Models
<span style="font-size:20px;font-weight:bold;color:#68f;">${unusedNew.length}</span> </button>
<span style="color:#99b;margin-left:6px;">unused &lt;1 month</span>
</div>
<div id="nodes-stats-used-badge" style="background:#1a2a1a;padding:8px 14px;border-radius:4px;border-left:3px solid #4a4;cursor:default;user-select:none;">
<span style="font-size:20px;font-weight:bold;color:#4a4;">${used.length}</span>
<span style="color:#9c9;margin-left:6px;">used</span>
</div>
</div>`; </div>`;
if (safeToRemove.length > 0) { // Nodes tab content (existing content, wrapped)
html += sectionHeader("Safe to Remove", "Unused for 2+ months", "#e44"); html += `<div id="ns-content-nodes">`;
html += buildTable(safeToRemove, "safe_to_remove"); html += buildNodesTabContent(custom);
} html += `</div>`;
if (considerRemoving.length > 0) { // Models tab content
html += sectionHeader("Consider Removing", "Unused for 1-2 months", "#e90"); html += `<div id="ns-content-models" style="display:none;">`;
html += buildTable(considerRemoving, "consider_removing"); html += buildModelsTabContent(modelData);
} html += `</div>`;
if (unusedNew.length > 0) {
html += sectionHeader("Recently Unused", "Unused for less than 1 month", "#68f");
html += buildTable(unusedNew, "unused_new");
}
if (used.length > 0) {
html += sectionHeader("Used", "", "#4a4");
html += buildTable(used, "used");
}
if (uninstalled.length > 0) {
html += sectionHeader("Uninstalled", "Previously tracked, no longer installed", "#555");
html += buildTable(uninstalled, "uninstalled");
}
dialog.innerHTML = html; dialog.innerHTML = html;
overlay.appendChild(dialog); overlay.appendChild(dialog);
document.body.appendChild(overlay); document.body.appendChild(overlay);
document // Tab switch — local function, no window pollution
.getElementById("nodes-stats-close") function switchTab(tab) {
.addEventListener("click", () => overlay.remove()); dialog.querySelector("#ns-content-nodes").style.display = tab === "nodes" ? "" : "none";
dialog.querySelector("#ns-content-models").style.display = tab === "models" ? "" : "none";
const nodeBtn = dialog.querySelector("#ns-tab-nodes");
const modelBtn = dialog.querySelector("#ns-tab-models");
nodeBtn.style.borderBottomColor = tab === "nodes" ? "#4a4" : "transparent";
nodeBtn.style.color = tab === "nodes" ? "#4a4" : "#888";
nodeBtn.style.fontWeight = tab === "nodes" ? "bold" : "normal";
modelBtn.style.borderBottomColor = tab === "models" ? "#4a4" : "transparent";
modelBtn.style.color = tab === "models" ? "#4a4" : "#888";
modelBtn.style.fontWeight = tab === "models" ? "bold" : "normal";
}
dialog.querySelector("#ns-tab-nodes").addEventListener("click", () => switchTab("nodes"));
dialog.querySelector("#ns-tab-models").addEventListener("click", () => switchTab("models"));
dialog.querySelector("#nodes-stats-close").addEventListener("click", () => overlay.remove());
// Toggle expandable rows // Toggle expandable rows
dialog.querySelectorAll(".pkg-row").forEach((row) => { dialog.querySelectorAll(".pkg-row").forEach((row) => {
@@ -148,7 +138,7 @@ async function showStatsDialog() {
// Easter egg: click "used" badge 5 times to show podium // Easter egg: click "used" badge 5 times to show podium
let eggClicks = 0; let eggClicks = 0;
let eggTimer = null; let eggTimer = null;
const usedBadge = document.getElementById("nodes-stats-used-badge"); const usedBadge = dialog.querySelector("#nodes-stats-used-badge");
if (usedBadge) { if (usedBadge) {
usedBadge.addEventListener("click", () => { usedBadge.addEventListener("click", () => {
eggClicks++; eggClicks++;
@@ -165,6 +155,116 @@ async function showStatsDialog() {
} }
} }
function buildNodesTabContent(custom) {
const safeToRemove = custom.filter((p) => p.status === "safe_to_remove");
const considerRemoving = custom.filter((p) => p.status === "consider_removing");
const unusedNew = custom.filter((p) => p.status === "unused_new");
const used = custom.filter((p) => p.status === "used");
const uninstalled = custom.filter((p) => p.status === "uninstalled");
let html = `<div style="display:flex;gap:10px;margin-bottom:20px;flex-wrap:wrap;">
<div style="background:#3a1a1a;padding:8px 14px;border-radius:4px;border-left:3px solid #e44;">
<span style="font-size:20px;font-weight:bold;color:#e44;">${safeToRemove.length}</span>
<span style="color:#c99;margin-left:6px;">safe to remove</span>
</div>
<div style="background:#2a2215;padding:8px 14px;border-radius:4px;border-left:3px solid #e90;">
<span style="font-size:20px;font-weight:bold;color:#e90;">${considerRemoving.length}</span>
<span style="color:#ca8;margin-left:6px;">consider removing</span>
</div>
<div style="background:#1a1a2a;padding:8px 14px;border-radius:4px;border-left:3px solid #68f;">
<span style="font-size:20px;font-weight:bold;color:#68f;">${unusedNew.length}</span>
<span style="color:#99b;margin-left:6px;">unused &lt;1 month</span>
</div>
<div id="nodes-stats-used-badge" style="background:#1a2a1a;padding:8px 14px;border-radius:4px;border-left:3px solid #4a4;cursor:default;user-select:none;">
<span style="font-size:20px;font-weight:bold;color:#4a4;">${used.length}</span>
<span style="color:#9c9;margin-left:6px;">used</span>
</div>
</div>`;
if (safeToRemove.length > 0) html += sectionHeader("Safe to Remove", "Unused for 2+ months", "#e44") + buildTable(safeToRemove, "safe_to_remove");
if (considerRemoving.length > 0) html += sectionHeader("Consider Removing", "Unused for 1-2 months", "#e90") + buildTable(considerRemoving, "consider_removing");
if (unusedNew.length > 0) html += sectionHeader("Recently Unused", "Unused for less than 1 month", "#68f") + buildTable(unusedNew, "unused_new");
if (used.length > 0) html += sectionHeader("Used", "", "#4a4") + buildTable(used, "used");
if (uninstalled.length > 0) html += sectionHeader("Uninstalled", "Previously tracked, no longer installed", "#555") + buildTable(uninstalled, "uninstalled");
return html;
}
function buildModelsTabContent(modelData) {
// Flatten for summary counts
const allModels = modelData.flatMap((g) => g.models);
const safeCount = allModels.filter((m) => m.status === "safe_to_remove").length;
const considerCount = allModels.filter((m) => m.status === "consider_removing").length;
const unusedNewCount = allModels.filter((m) => m.status === "unused_new").length;
const usedCount = allModels.filter((m) => m.status === "used").length;
let html = `<div style="display:flex;gap:10px;margin-bottom:20px;flex-wrap:wrap;">
<div style="background:#3a1a1a;padding:8px 14px;border-radius:4px;border-left:3px solid #e44;">
<span style="font-size:20px;font-weight:bold;color:#e44;">${safeCount}</span>
<span style="color:#c99;margin-left:6px;">safe to remove</span>
</div>
<div style="background:#2a2215;padding:8px 14px;border-radius:4px;border-left:3px solid #e90;">
<span style="font-size:20px;font-weight:bold;color:#e90;">${considerCount}</span>
<span style="color:#ca8;margin-left:6px;">consider removing</span>
</div>
<div style="background:#1a1a2a;padding:8px 14px;border-radius:4px;border-left:3px solid #68f;">
<span style="font-size:20px;font-weight:bold;color:#68f;">${unusedNewCount}</span>
<span style="color:#99b;margin-left:6px;">unused &lt;1 month</span>
</div>
<div style="background:#1a2a1a;padding:8px 14px;border-radius:4px;border-left:3px solid #4a4;">
<span style="font-size:20px;font-weight:bold;color:#4a4;">${usedCount}</span>
<span style="color:#9c9;margin-left:6px;">used</span>
</div>
</div>`;
if (allModels.length === 0) {
html += `<p style="color:#666;">No models tracked yet. Run a workflow to start.</p>`;
return html;
}
for (const group of modelData) {
if (group.models.length === 0) continue;
const title = group.model_type.charAt(0).toUpperCase() + group.model_type.slice(1).replace(/_/g, " ");
html += sectionHeader(title, `${group.models.length} model${group.models.length !== 1 ? "s" : ""}`, "#4a4");
html += buildModelTable(group.models);
}
return html;
}
function buildModelTable(models) {
let html = `<table style="width:100%;border-collapse:collapse;margin-bottom:12px;">
<thead><tr style="color:#888;text-align:left;border-bottom:1px solid #333;">
<th style="padding:6px 8px;">Model</th>
<th style="padding:6px 8px;text-align:right;">Executions</th>
<th style="padding:6px 8px;">Last Used</th>
<th style="padding:6px 8px;">Status</th>
</tr></thead><tbody>`;
for (const m of models) {
const { bg, hover } = STATUS_COLORS[m.status] || STATUS_COLORS.used;
const lastSeen = m.last_seen ? new Date(m.last_seen).toLocaleDateString() : "—";
const statusLabel = {
safe_to_remove: { text: "safe to remove", color: "#e44" },
consider_removing: { text: "consider removing", color: "#e90" },
unused_new: { text: "unused <1mo", color: "#68f" },
used: { text: "used", color: "#4a4" },
uninstalled: { text: "uninstalled", color: "#555" },
}[m.status] || { text: m.status, color: "#888" };
html += `<tr style="background:${bg};border-bottom:1px solid #222;"
onmouseover="this.style.background='${hover}'" onmouseout="this.style.background='${bg}'">
<td style="padding:6px 8px;color:#fff;">${escapeHtml(m.model_name)}</td>
<td style="padding:6px 8px;text-align:right;">${m.count}</td>
<td style="padding:6px 8px;color:#888;">${lastSeen}</td>
<td style="padding:6px 8px;"><span style="color:${statusLabel.color};font-size:11px;">${statusLabel.text}</span></td>
</tr>`;
}
html += `</tbody></table>`;
return html;
}
function sectionHeader(title, subtitle, color) { function sectionHeader(title, subtitle, color) {
let html = `<h3 style="color:${color};margin:16px 0 8px;font-size:14px;">${title}`; let html = `<h3 style="color:${color};margin:16px 0 8px;font-size:14px;">${title}`;
if (subtitle) html += ` <span style="color:#666;font-size:12px;font-weight:normal;">— ${subtitle}</span>`; if (subtitle) html += ` <span style="color:#666;font-size:12px;font-weight:normal;">— ${subtitle}</span>`;
+116
View File
@@ -66,3 +66,119 @@ class NodePackageMapper:
def invalidate(self): def invalidate(self):
"""Force rebuild on next access (e.g. after node reload).""" """Force rebuild on next access (e.g. after node reload)."""
self._map = None self._map = None
# Folder types that are not model files and should not be tracked
EXCLUDED_FOLDER_TYPES = {
"loras",
"configs",
"custom_nodes",
"temp",
"output",
"input",
"annotators",
"assets",
}
class ModelMapper:
"""Tracks which folder_paths model types exist and resolves filenames to types."""
def __init__(self):
self._folder_files = None # {folder_type: frozenset(filenames)}
self._reverse = None # {filename: folder_type}
def _build(self):
try:
import folder_paths
folder_files = {}
for folder_type in folder_paths.folder_names_and_paths:
if folder_type in EXCLUDED_FOLDER_TYPES:
continue
try:
files = folder_paths.get_filename_list(folder_type)
except Exception:
files = []
if files:
folder_files[folder_type] = frozenset(files)
# Build reverse map: filename -> folder_type (last write wins on collision)
reverse = {}
for folder_type, files in folder_files.items():
for f in files:
reverse[f] = folder_type
# Assign atomically only on success
self._folder_files = folder_files
self._reverse = reverse
except Exception:
logger.warning("ModelMapper: failed to build model map", exc_info=True)
self._folder_files = {}
self._reverse = {}
def _ensure(self):
if self._folder_files is None:
self._build()
def get_model_type(self, filename):
"""Return the folder type for a filename, or None if not tracked."""
self._ensure()
return self._reverse.get(filename)
def get_all_models(self):
"""Return {folder_type: [filename, ...]} for all tracked types."""
self._ensure()
return {k: sorted(v) for k, v in self._folder_files.items()}
def extract_models_from_prompt(self, prompt):
"""Scan a prompt dict and return (model_name, model_type) pairs.
For each node, inspects INPUT_TYPES() to find list-type (folder dropdown)
inputs, then resolves the selected value against the folder_paths reverse map.
"""
self._ensure()
try:
import nodes as comfy_nodes
except ImportError:
return []
seen = set()
results = []
for node_data in prompt.values():
class_type = node_data.get("class_type")
node_inputs = node_data.get("inputs", {})
if not class_type or not node_inputs:
continue
node_cls = comfy_nodes.NODE_CLASS_MAPPINGS.get(class_type)
if node_cls is None:
continue
try:
input_types = node_cls.INPUT_TYPES()
for category in ("required", "optional"):
for input_name, input_def in input_types.get(category, {}).items():
if not isinstance(input_def, (list, tuple)) or not input_def:
continue
# ComfyUI folder dropdowns have a list as their type
if not isinstance(input_def[0], list):
continue
value = node_inputs.get(input_name)
if not isinstance(value, str) or value in seen:
continue
model_type = self.get_model_type(value)
if model_type:
seen.add(value)
results.append((value, model_type))
except Exception:
continue
return results
def invalidate(self):
"""Force rebuild on next access."""
self._folder_files = None
self._reverse = None
+15
View File
@@ -0,0 +1,15 @@
# tests/conftest.py
import sys
import os
from unittest.mock import MagicMock
# Put the project root on sys.path so tests can import tracker, mapper directly
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
# Stub ComfyUI-only modules before any test file imports project code
_folder_paths_mock = MagicMock()
_folder_paths_mock.get_user_directory.return_value = "/tmp"
sys.modules["folder_paths"] = _folder_paths_mock
for mod in ("nodes", "server"):
if mod not in sys.modules:
sys.modules[mod] = MagicMock()
+96
View File
@@ -0,0 +1,96 @@
import pytest
from unittest.mock import patch, MagicMock
from mapper import ModelMapper
FAKE_FOLDER_NAMES = {
"checkpoints": ([], {}),
"vae": ([], {}),
"loras": ([], {}),
"configs": ([], {}),
}
FAKE_FILES = {
"checkpoints": ["dream.safetensors", "v15.ckpt"],
"vae": ["vae.safetensors"],
"loras": ["style.safetensors"],
}
def _make_mapper():
# conftest.py already put a MagicMock in sys.modules["folder_paths"],
# so we can configure it directly here.
import folder_paths as fp
fp.folder_names_and_paths = FAKE_FOLDER_NAMES
fp.get_filename_list.side_effect = lambda t: FAKE_FILES.get(t, [])
m = ModelMapper()
m._build()
return m
def test_get_model_type_known():
m = _make_mapper()
assert m.get_model_type("dream.safetensors") == "checkpoints"
assert m.get_model_type("vae.safetensors") == "vae"
def test_loras_excluded():
m = _make_mapper()
assert m.get_model_type("style.safetensors") is None
def test_get_all_models():
m = _make_mapper()
all_models = m.get_all_models()
assert "checkpoints" in all_models
assert "vae" in all_models
assert "loras" not in all_models
assert "dream.safetensors" in all_models["checkpoints"]
def test_unknown_filename_returns_none():
m = _make_mapper()
assert m.get_model_type("nonexistent.ckpt") is None
def test_extract_models_from_prompt(monkeypatch):
m = _make_mapper()
fake_node_cls = MagicMock()
fake_node_cls.INPUT_TYPES.return_value = {
"required": {
"ckpt_name": (["dream.safetensors", "v15.ckpt"],),
"steps": ("INT", {"default": 20}),
}
}
fake_prompt = {
"1": {
"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "dream.safetensors", "steps": 20},
}
}
import nodes as comfy_nodes
monkeypatch.setattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {"CheckpointLoaderSimple": fake_node_cls})
results = m.extract_models_from_prompt(fake_prompt)
assert ("dream.safetensors", "checkpoints") in results
def test_extract_models_skips_non_list_inputs(monkeypatch):
m = _make_mapper()
fake_node_cls = MagicMock()
fake_node_cls.INPUT_TYPES.return_value = {
"required": {
"text": ("STRING", {}),
}
}
fake_prompt = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "hello"}}}
import nodes as comfy_nodes
monkeypatch.setattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {"CLIPTextEncode": fake_node_cls})
results = m.extract_models_from_prompt(fake_prompt)
assert results == []
+103
View File
@@ -0,0 +1,103 @@
import pytest
import tempfile
import os
from unittest.mock import patch
from datetime import datetime, timezone, timedelta
from tracker import UsageTracker
@pytest.fixture
def tracker(tmp_path):
return UsageTracker(db_path=str(tmp_path / "test.db"))
def test_record_and_retrieve_model_usage(tracker):
tracker.record_model_usage([("dreamshaper.safetensors", "checkpoints")])
tracker.record_model_usage([("dreamshaper.safetensors", "checkpoints")])
raw = tracker.get_raw_model_stats()
assert len(raw) == 1
assert raw[0]["model_name"] == "dreamshaper.safetensors"
assert raw[0]["model_type"] == "checkpoints"
assert raw[0]["count"] == 2
def test_record_multiple_models(tracker):
tracker.record_model_usage([
("dreamshaper.safetensors", "checkpoints"),
("vae.safetensors", "vae"),
])
raw = tracker.get_raw_model_stats()
assert len(raw) == 2
def test_reset_clears_model_usage(tracker):
tracker.record_model_usage([("model.safetensors", "checkpoints")])
tracker.reset()
assert tracker.get_raw_model_stats() == []
def test_empty_models_returns_empty(tracker):
assert tracker.get_raw_model_stats() == []
def test_get_model_stats_used(tracker):
tracker.record_model_usage([("model.safetensors", "checkpoints")])
installed = {"checkpoints": ["model.safetensors"]}
result = tracker.get_model_stats(installed)
assert len(result) == 1
assert result[0]["model_type"] == "checkpoints"
assert result[0]["models"][0]["status"] == "used"
assert result[0]["models"][0]["count"] == 1
def test_get_model_stats_never_used_new(tracker):
installed = {"checkpoints": ["unused.safetensors"]}
result = tracker.get_model_stats(installed)
assert result[0]["models"][0]["status"] == "unused_new"
assert result[0]["models"][0]["count"] == 0
def test_get_model_stats_uninstalled(tracker):
tracker.record_model_usage([("gone.safetensors", "checkpoints")])
installed = {} # no longer on disk
result = tracker.get_model_stats(installed)
assert result[0]["models"][0]["status"] == "uninstalled"
assert result[0]["models"][0]["installed"] is False
def test_get_model_stats_sorted_by_status(tracker):
tracker.record_model_usage([("active.safetensors", "checkpoints")])
installed = {"checkpoints": ["active.safetensors", "unused.safetensors"]}
result = tracker.get_model_stats(installed)
models = result[0]["models"]
statuses = [m["status"] for m in models]
# unused_new (2) comes before used (3) in STATUS_ORDER
assert statuses.index("unused_new") < statuses.index("used")
def test_get_model_stats_safe_to_remove(tracker):
"""A model last used 70 days ago should be classified safe_to_remove."""
tracker.record_model_usage([("old.safetensors", "checkpoints")])
installed = {"checkpoints": ["old.safetensors"]}
# Patch datetime in tracker so "now" is 70 days after last_seen
future_now = datetime.now(timezone.utc) + timedelta(days=70)
with patch("tracker.datetime") as mock_dt:
mock_dt.now.return_value = future_now
result = tracker.get_model_stats(installed)
assert result[0]["models"][0]["status"] == "safe_to_remove"
def test_get_model_stats_consider_removing(tracker):
"""A model last used 40 days ago should be classified consider_removing."""
tracker.record_model_usage([("medium.safetensors", "checkpoints")])
installed = {"checkpoints": ["medium.safetensors"]}
future_now = datetime.now(timezone.utc) + timedelta(days=40)
with patch("tracker.datetime") as mock_dt:
mock_dt.now.return_value = future_now
result = tracker.get_model_stats(installed)
assert result[0]["models"][0]["status"] == "consider_removing"
+143
View File
@@ -40,8 +40,17 @@ CREATE TABLE IF NOT EXISTS prompt_log (
class_types TEXT NOT NULL class_types TEXT NOT NULL
); );
CREATE TABLE IF NOT EXISTS model_usage (
model_name TEXT PRIMARY KEY,
model_type TEXT NOT NULL,
count INTEGER NOT NULL DEFAULT 0,
first_seen TEXT NOT NULL,
last_seen TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_node_usage_package ON node_usage(package); CREATE INDEX IF NOT EXISTS idx_node_usage_package ON node_usage(package);
CREATE INDEX IF NOT EXISTS idx_prompt_log_timestamp ON prompt_log(timestamp); CREATE INDEX IF NOT EXISTS idx_prompt_log_timestamp ON prompt_log(timestamp);
CREATE INDEX IF NOT EXISTS idx_model_usage_type ON model_usage(model_type);
""" """
@@ -116,6 +125,47 @@ class UsageTracker:
finally: finally:
conn.close() conn.close()
def record_model_usage(self, models):
"""Record usage of model files from a single prompt.
models: list of (model_name, model_type) tuples
"""
if not models:
return
now = datetime.now(timezone.utc).isoformat()
with self._lock:
self._ensure_db()
conn = self._connect()
try:
for model_name, model_type in models:
conn.execute(
"""INSERT INTO model_usage (model_name, model_type, count, first_seen, last_seen)
VALUES (?, ?, 1, ?, ?)
ON CONFLICT(model_name) DO UPDATE SET
count = count + 1,
last_seen = excluded.last_seen,
model_type = excluded.model_type""",
(model_name, model_type, now, now),
)
conn.commit()
finally:
conn.close()
def get_raw_model_stats(self):
"""Return raw per-model usage rows from DB."""
with self._lock:
self._ensure_db()
conn = self._connect()
try:
conn.row_factory = sqlite3.Row
rows = conn.execute(
"SELECT model_name, model_type, count, first_seen, last_seen "
"FROM model_usage ORDER BY count DESC"
).fetchall()
return [dict(r) for r in rows]
finally:
conn.close()
def get_node_stats(self): def get_node_stats(self):
"""Return raw per-node usage data.""" """Return raw per-node usage data."""
with self._lock: with self._lock:
@@ -216,6 +266,98 @@ class UsageTracker:
result.sort(key=lambda p: p["total_executions"]) result.sort(key=lambda p: p["total_executions"])
return result return result
def get_model_stats(self, installed_by_type):
"""Return per-type grouped model stats with tier classification.
installed_by_type: {model_type: [model_name, ...]} from ModelMapper
"""
db_rows = self.get_raw_model_stats()
db_models = {r["model_name"]: r for r in db_rows}
now = datetime.now(timezone.utc)
one_month_ago = (now - timedelta(days=30)).isoformat()
two_months_ago = (now - timedelta(days=60)).isoformat()
tracking_start = self._get_first_prompt_time()
STATUS_ORDER = {
"safe_to_remove": 0,
"consider_removing": 1,
"unused_new": 2,
"used": 3,
"uninstalled": 4,
}
result_by_type = {}
# Process installed models
for model_type, filenames in installed_by_type.items():
entries = []
for model_name in filenames:
if model_name in db_models:
row = db_models[model_name]
last_seen = row["last_seen"]
if last_seen < two_months_ago:
status = "safe_to_remove"
elif last_seen < one_month_ago:
status = "consider_removing"
else:
status = "used"
entry = {
"model_name": model_name,
"model_type": model_type,
"count": row["count"],
"first_seen": row["first_seen"],
"last_seen": last_seen,
"installed": True,
"status": status,
}
else:
if tracking_start is None:
status = "unused_new"
elif tracking_start < two_months_ago:
status = "safe_to_remove"
elif tracking_start < one_month_ago:
status = "consider_removing"
else:
status = "unused_new"
entry = {
"model_name": model_name,
"model_type": model_type,
"count": 0,
"first_seen": None,
"last_seen": None,
"installed": True,
"status": status,
}
entries.append(entry)
result_by_type[model_type] = entries
# Add uninstalled (in DB but not on disk)
installed_names = {
name for names in installed_by_type.values() for name in names
}
for model_name, row in db_models.items():
if model_name not in installed_names:
model_type = row["model_type"]
result_by_type.setdefault(model_type, []).append({
"model_name": model_name,
"model_type": model_type,
"count": row["count"],
"first_seen": row["first_seen"],
"last_seen": row["last_seen"],
"installed": False,
"status": "uninstalled",
})
# Sort each type's models by status tier then name
result = []
for model_type in sorted(result_by_type):
models = result_by_type[model_type]
models.sort(key=lambda m: (STATUS_ORDER.get(m["status"], 5), m["model_name"]))
result.append({"model_type": model_type, "models": models})
return result
def _get_first_prompt_time(self): def _get_first_prompt_time(self):
"""Return the timestamp of the earliest recorded prompt, or None.""" """Return the timestamp of the earliest recorded prompt, or None."""
with self._lock: with self._lock:
@@ -237,6 +379,7 @@ class UsageTracker:
try: try:
conn.execute("DELETE FROM node_usage") conn.execute("DELETE FROM node_usage")
conn.execute("DELETE FROM prompt_log") conn.execute("DELETE FROM prompt_log")
conn.execute("DELETE FROM model_usage")
conn.commit() conn.commit()
finally: finally:
conn.close() conn.close()