Setmaster
fb28135b82
- Show model file sizes (GB/MB) in the sidebar next to last used time - Use ComfyUI's folder_paths for reliable model path resolution - Support all configured model locations including extra_model_paths - Handle GGUF models that can be in unet or clip folders - Update README with new features
316 lines
9.9 KiB
Python
316 lines
9.9 KiB
Python
"""
|
|
ModelPulse - Core Tracking Logic
|
|
|
|
Extracts models from prompts and records usage statistics.
|
|
"""
|
|
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .model_types import MODEL_LOADERS, LOADER_PATTERNS, CATEGORY_FOLDERS
|
|
from .storage import load_data, save_data, cleanup_old_usage_logs
|
|
|
|
# Try to import ComfyUI's folder_paths for proper model path resolution
|
|
try:
|
|
import folder_paths
|
|
HAS_FOLDER_PATHS = True
|
|
except ImportError:
|
|
HAS_FOLDER_PATHS = False
|
|
|
|
|
|
def get_model_file_size(category: str, name: str) -> int | None:
|
|
"""
|
|
Get the file size of a model in bytes.
|
|
|
|
Uses ComfyUI's folder_paths to find models in all configured locations.
|
|
|
|
Args:
|
|
category: The model category (e.g., 'checkpoint', 'lora')
|
|
name: The model filename
|
|
|
|
Returns:
|
|
File size in bytes, or None if file not found
|
|
"""
|
|
if not HAS_FOLDER_PATHS:
|
|
return None
|
|
|
|
# Map our categories to ComfyUI's folder_paths keys
|
|
folder_key = CATEGORY_FOLDERS.get(category)
|
|
|
|
# For GGUF, try multiple folders since models can be unet or clip
|
|
if category == "gguf":
|
|
folder_keys = ["unet", "diffusion_models", "clip"]
|
|
elif category == "unet":
|
|
folder_keys = ["unet", "diffusion_models"]
|
|
elif folder_key:
|
|
folder_keys = [folder_key]
|
|
else:
|
|
return None
|
|
|
|
for key in folder_keys:
|
|
try:
|
|
full_path = folder_paths.get_full_path(key, name)
|
|
if full_path:
|
|
return Path(full_path).stat().st_size
|
|
except Exception:
|
|
continue
|
|
|
|
return None
|
|
|
|
|
|
class ModelTracker:
|
|
"""Tracks model usage across ComfyUI workflow executions."""
|
|
|
|
def __init__(self):
|
|
self._data: dict[str, Any] | None = None
|
|
|
|
@property
|
|
def data(self) -> dict[str, Any]:
|
|
"""Lazy-load data from storage."""
|
|
if self._data is None:
|
|
self._data = load_data()
|
|
return self._data
|
|
|
|
def _save(self) -> None:
|
|
"""Save current data to storage."""
|
|
if self._data is not None:
|
|
save_data(self._data)
|
|
|
|
def extract_models_from_prompt(self, prompt: dict[str, Any]) -> list[dict[str, str]]:
|
|
"""
|
|
Extract all model references from a ComfyUI prompt.
|
|
|
|
Args:
|
|
prompt: The prompt dictionary from ComfyUI execution
|
|
|
|
Returns:
|
|
List of dicts with 'category', 'name', and 'model_id' keys
|
|
"""
|
|
models = []
|
|
seen = set() # Avoid duplicates within same prompt
|
|
|
|
for node_id, node_data in prompt.items():
|
|
class_type = node_data.get("class_type", "")
|
|
inputs = node_data.get("inputs", {})
|
|
|
|
# Check standard loaders first
|
|
if class_type in MODEL_LOADERS:
|
|
category, input_keys = MODEL_LOADERS[class_type]
|
|
extracted = self._extract_from_inputs(category, input_keys, inputs)
|
|
for model in extracted:
|
|
if model["model_id"] not in seen:
|
|
seen.add(model["model_id"])
|
|
models.append(model)
|
|
|
|
# Check pattern-based loaders for custom nodes
|
|
elif "Loader" in class_type:
|
|
extracted = self._extract_from_patterns(class_type, inputs)
|
|
for model in extracted:
|
|
if model["model_id"] not in seen:
|
|
seen.add(model["model_id"])
|
|
models.append(model)
|
|
|
|
return models
|
|
|
|
def _extract_from_inputs(
|
|
self, category: str, input_keys: str | list[str], inputs: dict[str, Any]
|
|
) -> list[dict[str, str]]:
|
|
"""Extract model names from node inputs."""
|
|
models = []
|
|
|
|
if isinstance(input_keys, str):
|
|
input_keys = [input_keys]
|
|
|
|
for key in input_keys:
|
|
value = inputs.get(key)
|
|
if value and isinstance(value, str):
|
|
model_id = f"{category}/{value}"
|
|
models.append({
|
|
"category": category,
|
|
"name": value,
|
|
"model_id": model_id,
|
|
})
|
|
|
|
return models
|
|
|
|
def _extract_from_patterns(
|
|
self, class_type: str, inputs: dict[str, Any]
|
|
) -> list[dict[str, str]]:
|
|
"""Extract models using pattern matching for custom node loaders."""
|
|
models = []
|
|
|
|
for pattern, category, possible_keys in LOADER_PATTERNS:
|
|
if pattern.lower() in class_type.lower():
|
|
for key in possible_keys:
|
|
value = inputs.get(key)
|
|
if value and isinstance(value, str):
|
|
model_id = f"{category}/{value}"
|
|
models.append({
|
|
"category": category,
|
|
"name": value,
|
|
"model_id": model_id,
|
|
})
|
|
break # Found a match for this pattern
|
|
break # Only match first pattern
|
|
|
|
return models
|
|
|
|
def record_usage(self, models: list[dict[str, str]]) -> None:
|
|
"""
|
|
Record usage for a list of models.
|
|
|
|
Args:
|
|
models: List of model dicts from extract_models_from_prompt
|
|
"""
|
|
if not models:
|
|
return
|
|
|
|
now = datetime.utcnow()
|
|
now_iso = now.isoformat() + "Z"
|
|
today = now.date().isoformat()
|
|
|
|
for model in models:
|
|
model_id = model["model_id"]
|
|
|
|
if model_id not in self.data["models"]:
|
|
# New model - create entry
|
|
self.data["models"][model_id] = {
|
|
"category": model["category"],
|
|
"name": model["name"],
|
|
"path": model_id, # We don't have full path, use model_id
|
|
"first_used": now_iso,
|
|
"last_used": now_iso,
|
|
"usage_count": 0,
|
|
"usage_log": [],
|
|
}
|
|
|
|
model_data = self.data["models"][model_id]
|
|
|
|
# Update timestamps and count
|
|
model_data["last_used"] = now_iso
|
|
model_data["usage_count"] += 1
|
|
|
|
# Update daily usage log
|
|
usage_log = model_data["usage_log"]
|
|
if usage_log and usage_log[-1]["date"] == today:
|
|
usage_log[-1]["count"] += 1
|
|
else:
|
|
usage_log.append({"date": today, "count": 1})
|
|
|
|
self._save()
|
|
|
|
def extract_and_record_models(self, prompt: dict[str, Any]) -> None:
|
|
"""
|
|
Extract models from prompt and record their usage.
|
|
|
|
This is the main entry point called from the PromptExecutor hook.
|
|
"""
|
|
models = self.extract_models_from_prompt(prompt)
|
|
self.record_usage(models)
|
|
|
|
def get_usage_data(
|
|
self,
|
|
timeframe: str = "all",
|
|
sort_by: str = "last_used",
|
|
category: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Get filtered and sorted usage data.
|
|
|
|
Args:
|
|
timeframe: "all", "month", or "week"
|
|
sort_by: "last_used", "usage_count", or "name"
|
|
category: Optional category filter
|
|
|
|
Returns:
|
|
Dict with 'models' list and 'metadata'
|
|
"""
|
|
models_list = []
|
|
now = datetime.utcnow()
|
|
|
|
# Calculate timeframe boundaries
|
|
if timeframe == "week":
|
|
cutoff = now - timedelta(days=7)
|
|
elif timeframe == "month":
|
|
cutoff = now - timedelta(days=30)
|
|
else:
|
|
cutoff = None
|
|
|
|
cutoff_date = cutoff.date().isoformat() if cutoff else None
|
|
|
|
for model_id, model_data in self.data["models"].items():
|
|
# Filter by category if specified
|
|
if category and model_data["category"] != category:
|
|
continue
|
|
|
|
# Calculate timeframe-specific usage count
|
|
if cutoff_date:
|
|
timeframe_count = sum(
|
|
entry["count"]
|
|
for entry in model_data.get("usage_log", [])
|
|
if entry["date"] >= cutoff_date
|
|
)
|
|
else:
|
|
timeframe_count = model_data["usage_count"]
|
|
|
|
# Get file size
|
|
file_size = get_model_file_size(model_data["category"], model_data["name"])
|
|
|
|
models_list.append({
|
|
"model_id": model_id,
|
|
"category": model_data["category"],
|
|
"name": model_data["name"],
|
|
"first_used": model_data["first_used"],
|
|
"last_used": model_data["last_used"],
|
|
"usage_count": model_data["usage_count"],
|
|
"timeframe_count": timeframe_count,
|
|
"file_size": file_size,
|
|
})
|
|
|
|
# Sort the results
|
|
if sort_by == "usage_count":
|
|
models_list.sort(key=lambda x: x["timeframe_count"], reverse=True)
|
|
elif sort_by == "name":
|
|
models_list.sort(key=lambda x: x["name"].lower())
|
|
else: # last_used (default)
|
|
models_list.sort(key=lambda x: x["last_used"], reverse=True)
|
|
|
|
return {
|
|
"models": models_list,
|
|
"metadata": self.data["metadata"],
|
|
"timeframe": timeframe,
|
|
"sort_by": sort_by,
|
|
}
|
|
|
|
def get_model_detail(self, model_id: str) -> dict[str, Any] | None:
|
|
"""
|
|
Get detailed usage data for a specific model.
|
|
|
|
Args:
|
|
model_id: The model identifier (category/filename)
|
|
|
|
Returns:
|
|
Model data dict or None if not found
|
|
"""
|
|
model_data = self.data["models"].get(model_id)
|
|
if not model_data:
|
|
return None
|
|
|
|
return {
|
|
"model_id": model_id,
|
|
**model_data,
|
|
}
|
|
|
|
def reset(self) -> None:
|
|
"""Reset all tracking data."""
|
|
from .storage import create_empty_data
|
|
|
|
self._data = create_empty_data()
|
|
self._save()
|
|
|
|
def cleanup(self, max_days: int = 365) -> None:
|
|
"""Clean up old usage logs."""
|
|
self._data = cleanup_old_usage_logs(self.data, max_days)
|
|
self._save()
|