Compare commits
16 Commits
2387a8d8ef
...
777e4890ae
| Author | SHA1 | Date | |
|---|---|---|---|
| 777e4890ae | |||
| 9477777124 | |||
| 831ed302e9 | |||
| c7003df0bb | |||
| bf5c8683a7 | |||
| 18d098a0da | |||
| ed39c5918a | |||
| 38e95c150a | |||
| 9891b3ec41 | |||
| ff42016009 | |||
| 800692d190 | |||
| 70edeb0631 | |||
| 3ed778a0d6 | |||
| b852c37786 | |||
| 901e3b722d | |||
| ab204056a0 |
@@ -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
@@ -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
@@ -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;">×</button>
|
<button id="nodes-stats-close" style="background:none;border:none;color:#888;font-size:20px;cursor:pointer;">×</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 <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 <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 <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>`;
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -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 == []
|
||||||
@@ -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
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user