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
+46 -32
View File
@@ -4,14 +4,16 @@
<img src="docs/logo.svg" width="120" alt="Node Stats logo">
</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
- **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
- **Smart aging** — packages gradually move from "recently unused" to "safe to remove" over time
- **Uninstall detection** — removed packages are flagged separately, historical data preserved
- **Per-model classification** — models grouped by type (checkpoints, vae, …) with the same recency tiers
- **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
- **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
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
- **Sections** for each tier, sorted from most actionable to least
- **Expandable rows** — click any package to see per-node execution counts and timestamps
**Nodes tab**
- Summary bar with counts for each classification tier
- 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
@@ -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/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 |
```bash
@@ -114,42 +123,46 @@ curl http://localhost:8188/nodes-stats/packages | python3 -m json.tool
## How It Works
```
Queue Prompt ──> Prompt Handler ──> Extract class_types ──> Background Thread
─────────────────────────┘
SQLite DB
usage_stats.db
──────────┐
node_usage │ per-node counts & timestamps
prompt_log │ full node list per prompt
└──────────┘
GET /nodes-stats/packages ◄──┘
Mapper merges DB data
with NODE_CLASS_MAPPINGS
Classify by recency ──> JSON response ──> UI Dialog
Queue Prompt ──> Prompt Handler ──> Extract class_types + prompt ──> Background Thread
┌───────────────────────────────┘
SQLite DB
usage_stats.db
┌─────────────┐
node_usage │ per-node counts & timestamps
prompt_log │ full node list per prompt
│ model_usage │ per-model counts & timestamps
└─────────────┘
GET /nodes-stats/packages ◄─────────┤
GET /nodes-stats/models ◄─────────┘
Merge DB data with installed
nodes/models, classify by recency
JSON response ──> UI Dialog (Nodes tab / Models tab)
```
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)
4. Maps each class_type to its source package using `RELATIVE_PYTHON_MODULE`
5. Upserts per-node counts and timestamps into SQLite
6. On stats request, merges DB data with current node registry and classifies by recency
5. Detects model file selections by introspecting each node's `INPUT_TYPES()` for folder-dropdown inputs, then resolves filenames via `folder_paths`
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
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 |
|-------|----------|
| `node_usage` | Per-node: class_type, package, execution count, first/last seen |
| `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.
@@ -157,8 +170,9 @@ Use `POST /nodes-stats/reset` to clear all data and start fresh.
```
__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
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
tests/ Unit tests for tracker and mapper
```
+32 -4
View File
@@ -4,7 +4,7 @@ import threading
from aiohttp import web
from server import PromptServer
from .mapper import NodePackageMapper
from .mapper import NodePackageMapper, ModelMapper
from .tracker import UsageTracker
logger = logging.getLogger(__name__)
@@ -14,10 +14,11 @@ WEB_DIRECTORY = "./js"
mapper = NodePackageMapper()
tracker = UsageTracker()
model_mapper = ModelMapper()
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:
prompt = json_data.get("prompt", {})
class_types = set()
@@ -26,9 +27,11 @@ def on_prompt_handler(json_data):
if ct:
class_types.add(ct)
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(
target=tracker.record_usage,
args=(class_types, mapper),
target=_record_prompt,
args=(class_types, prompt),
daemon=True,
).start()
except Exception:
@@ -36,6 +39,19 @@ def on_prompt_handler(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)
@@ -62,11 +78,23 @@ async def get_node_stats(request):
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")
async def reset_stats(request):
try:
tracker.reset()
mapper.invalidate()
model_mapper.invalidate()
return web.json_response({"status": "ok"})
except Exception:
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() {
let data;
let data, modelData;
try {
const resp = await fetch("/nodes-stats/packages");
if (!resp.ok) {
alert("Failed to load node stats: HTTP " + resp.status);
return;
}
data = await resp.json();
if (!Array.isArray(data)) {
alert("Failed to load node stats: unexpected response format");
const [pkgResp, modelResp] = await Promise.all([
fetch("/nodes-stats/packages"),
fetch("/nodes-stats/models"),
]);
if (!pkgResp.ok) { alert("Failed to load node stats: HTTP " + pkgResp.status); return; }
if (!modelResp.ok) { alert("Failed to load model stats: HTTP " + modelResp.status); return; }
data = await pkgResp.json();
modelData = await modelResp.json();
if (!Array.isArray(data) || !Array.isArray(modelData)) {
alert("Failed to load stats: unexpected response format");
return;
}
} catch (e) {
alert("Failed to load node stats: " + e.message);
alert("Failed to load stats: " + e.message);
return;
}
const custom = data.filter((p) => p.package !== "__builtin__");
// Remove existing dialog if any
const existing = document.getElementById("nodes-stats-dialog");
if (existing) existing.remove();
@@ -67,69 +71,55 @@ async function showStatsDialog() {
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;";
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;">
<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>
</div>`;
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>
// Tab switcher — wired via addEventListener after insertion, no onclick globals
html += `
<div style="display:flex;gap:0;margin-bottom:20px;border-bottom:1px solid #333;">
<button id="ns-tab-nodes"
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;">
Nodes
</button>
<button id="ns-tab-models"
style="background:none;border:none;border-bottom:2px solid transparent;color:#888;padding:8px 18px;cursor:pointer;font-family:monospace;font-size:13px;">
Models
</button>
</div>`;
if (safeToRemove.length > 0) {
html += sectionHeader("Safe to Remove", "Unused for 2+ months", "#e44");
html += buildTable(safeToRemove, "safe_to_remove");
}
// Nodes tab content (existing content, wrapped)
html += `<div id="ns-content-nodes">`;
html += buildNodesTabContent(custom);
html += `</div>`;
if (considerRemoving.length > 0) {
html += sectionHeader("Consider Removing", "Unused for 1-2 months", "#e90");
html += buildTable(considerRemoving, "consider_removing");
}
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");
}
// Models tab content
html += `<div id="ns-content-models" style="display:none;">`;
html += buildModelsTabContent(modelData);
html += `</div>`;
dialog.innerHTML = html;
overlay.appendChild(dialog);
document.body.appendChild(overlay);
document
.getElementById("nodes-stats-close")
.addEventListener("click", () => overlay.remove());
// Tab switch — local function, no window pollution
function switchTab(tab) {
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
dialog.querySelectorAll(".pkg-row").forEach((row) => {
@@ -148,7 +138,7 @@ async function showStatsDialog() {
// Easter egg: click "used" badge 5 times to show podium
let eggClicks = 0;
let eggTimer = null;
const usedBadge = document.getElementById("nodes-stats-used-badge");
const usedBadge = dialog.querySelector("#nodes-stats-used-badge");
if (usedBadge) {
usedBadge.addEventListener("click", () => {
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) {
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>`;
+116
View File
@@ -66,3 +66,119 @@ class NodePackageMapper:
def invalidate(self):
"""Force rebuild on next access (e.g. after node reload)."""
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
);
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_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:
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):
"""Return raw per-node usage data."""
with self._lock:
@@ -216,6 +266,98 @@ class UsageTracker:
result.sort(key=lambda p: p["total_executions"])
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):
"""Return the timestamp of the earliest recorded prompt, or None."""
with self._lock:
@@ -237,6 +379,7 @@ class UsageTracker:
try:
conn.execute("DELETE FROM node_usage")
conn.execute("DELETE FROM prompt_log")
conn.execute("DELETE FROM model_usage")
conn.commit()
finally:
conn.close()