# CivitAI Extended Stats Shim — Implementation Plan
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
**Goal:** Add CivitAI stats (downloads, rating, thumbs up) to LoRA Manager cards with sorting, without modifying the original manager codebase.
**Architecture:** A shim inside ComfyUI-LM-Remote with its own SQLite DB (`civitai_stats.db`). The proxy middleware enriches `/api/lm/loras/list` responses with stats. New API endpoints handle bulk fetching. Frontend JS patches cards and sort dropdown.
**Tech Stack:** Python/aiohttp (backend), SQLite (storage), vanilla JS (frontend), CivitAI REST API v1
---
### Task 1: Stats SQLite Database Layer
**Files:**
- Create: `custom_nodes/ComfyUI-LM-Remote/stats_db.py`
- Create: `custom_nodes/ComfyUI-LM-Remote/tests/test_stats_db.py`
**Step 1: Write failing tests**
```python
# tests/test_stats_db.py
"""Tests for the CivitAI stats database layer."""
import asyncio
import pytest
from pathlib import Path
@pytest.fixture
def db_path(tmp_path):
return tmp_path / "test_stats.db"
@pytest.fixture
def stats_db(db_path):
from stats_db import StatsDB
db = StatsDB(db_path)
db.init()
return db
def test_init_creates_table(stats_db, db_path):
"""DB file is created and table exists."""
assert db_path.exists()
def test_upsert_and_get_single(stats_db):
"""Insert a stat row and retrieve it by sha256."""
stats_db.upsert("abc123", {
"civitai_model_id": 1,
"civitai_version_id": 10,
"download_count": 5000,
"rating": 4.8,
"rating_count": 120,
"thumbs_up_count": 89,
})
result = stats_db.get_by_hashes(["abc123"])
assert "abc123" in result
assert result["abc123"]["download_count"] == 5000
assert result["abc123"]["rating"] == 4.8
assert result["abc123"]["thumbs_up_count"] == 89
def test_upsert_updates_existing(stats_db):
"""Upserting same sha256 updates values."""
stats_db.upsert("abc123", {"download_count": 100, "rating": 3.0,
"rating_count": 10, "thumbs_up_count": 5})
stats_db.upsert("abc123", {"download_count": 200, "rating": 4.0,
"rating_count": 20, "thumbs_up_count": 15})
result = stats_db.get_by_hashes(["abc123"])
assert result["abc123"]["download_count"] == 200
def test_get_by_hashes_batch(stats_db):
"""Batch retrieval returns only matching hashes."""
for i in range(5):
stats_db.upsert(f"hash_{i}", {"download_count": i * 100,
"rating": 0, "rating_count": 0,
"thumbs_up_count": 0})
result = stats_db.get_by_hashes(["hash_1", "hash_3", "nonexistent"])
assert len(result) == 2
assert "hash_1" in result
assert "hash_3" in result
assert "nonexistent" not in result
def test_get_by_hashes_empty_input(stats_db):
"""Empty input returns empty dict."""
assert stats_db.get_by_hashes([]) == {}
def test_upsert_batch(stats_db):
"""Batch upsert inserts multiple rows."""
rows = [
("h1", {"civitai_model_id": 1, "download_count": 100,
"rating": 4.0, "rating_count": 10, "thumbs_up_count": 5}),
("h2", {"civitai_model_id": 2, "download_count": 200,
"rating": 3.5, "rating_count": 20, "thumbs_up_count": 15}),
]
stats_db.upsert_batch(rows)
result = stats_db.get_by_hashes(["h1", "h2"])
assert len(result) == 2
assert result["h1"]["download_count"] == 100
assert result["h2"]["download_count"] == 200
def test_get_all_stats(stats_db):
"""get_all returns every row keyed by sha256."""
stats_db.upsert("a", {"download_count": 1, "rating": 0,
"rating_count": 0, "thumbs_up_count": 0})
stats_db.upsert("b", {"download_count": 2, "rating": 0,
"rating_count": 0, "thumbs_up_count": 0})
result = stats_db.get_all()
assert len(result) == 2
def test_stats_count(stats_db):
"""count returns number of rows."""
assert stats_db.count() == 0
stats_db.upsert("a", {"download_count": 1, "rating": 0,
"rating_count": 0, "thumbs_up_count": 0})
assert stats_db.count() == 1
```
**Step 2: Run tests to verify they fail**
Run: `cd custom_nodes/ComfyUI-LM-Remote && python -m pytest tests/test_stats_db.py -v`
Expected: FAIL — `ModuleNotFoundError: No module named 'stats_db'`
**Step 3: Write implementation**
```python
# stats_db.py
"""SQLite storage for CivitAI extended stats."""
from __future__ import annotations
import sqlite3
import time
from pathlib import Path
_PACKAGE_DIR = Path(__file__).resolve().parent
_DEFAULT_DB_PATH = _PACKAGE_DIR / "civitai_stats.db"
_SCHEMA = """
CREATE TABLE IF NOT EXISTS model_stats (
sha256 TEXT PRIMARY KEY,
civitai_model_id INTEGER,
civitai_version_id INTEGER,
download_count INTEGER DEFAULT 0,
rating REAL DEFAULT 0,
rating_count INTEGER DEFAULT 0,
thumbs_up_count INTEGER DEFAULT 0,
fetched_at REAL
);
"""
class StatsDB:
"""Thin wrapper around a SQLite database for CivitAI stats."""
def __init__(self, db_path: Path | None = None):
self._db_path = db_path or _DEFAULT_DB_PATH
self._conn: sqlite3.Connection | None = None
def init(self) -> None:
"""Create the database and table if they don't exist."""
self._conn = sqlite3.connect(str(self._db_path))
self._conn.row_factory = sqlite3.Row
self._conn.execute("PRAGMA journal_mode=WAL")
self._conn.executescript(_SCHEMA)
def _ensure_conn(self) -> sqlite3.Connection:
if self._conn is None:
self.init()
return self._conn # type: ignore[return-value]
def upsert(self, sha256: str, data: dict) -> None:
"""Insert or update stats for a single model."""
conn = self._ensure_conn()
conn.execute(
"""INSERT INTO model_stats
(sha256, civitai_model_id, civitai_version_id,
download_count, rating, rating_count, thumbs_up_count, fetched_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(sha256) DO UPDATE SET
civitai_model_id = COALESCE(excluded.civitai_model_id, civitai_model_id),
civitai_version_id = COALESCE(excluded.civitai_version_id, civitai_version_id),
download_count = excluded.download_count,
rating = excluded.rating,
rating_count = excluded.rating_count,
thumbs_up_count = excluded.thumbs_up_count,
fetched_at = excluded.fetched_at
""",
(
sha256,
data.get("civitai_model_id"),
data.get("civitai_version_id"),
data.get("download_count", 0),
data.get("rating", 0),
data.get("rating_count", 0),
data.get("thumbs_up_count", 0),
time.time(),
),
)
conn.commit()
def upsert_batch(self, rows: list[tuple[str, dict]]) -> None:
"""Insert or update stats for multiple models."""
conn = self._ensure_conn()
now = time.time()
conn.executemany(
"""INSERT INTO model_stats
(sha256, civitai_model_id, civitai_version_id,
download_count, rating, rating_count, thumbs_up_count, fetched_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(sha256) DO UPDATE SET
civitai_model_id = COALESCE(excluded.civitai_model_id, civitai_model_id),
civitai_version_id = COALESCE(excluded.civitai_version_id, civitai_version_id),
download_count = excluded.download_count,
rating = excluded.rating,
rating_count = excluded.rating_count,
thumbs_up_count = excluded.thumbs_up_count,
fetched_at = excluded.fetched_at
""",
[
(
sha256,
d.get("civitai_model_id"),
d.get("civitai_version_id"),
d.get("download_count", 0),
d.get("rating", 0),
d.get("rating_count", 0),
d.get("thumbs_up_count", 0),
now,
)
for sha256, d in rows
],
)
conn.commit()
def get_by_hashes(self, hashes: list[str]) -> dict[str, dict]:
"""Return stats keyed by sha256 for the given hashes."""
if not hashes:
return {}
conn = self._ensure_conn()
placeholders = ",".join("?" for _ in hashes)
rows = conn.execute(
f"SELECT * FROM model_stats WHERE sha256 IN ({placeholders})",
hashes,
).fetchall()
return {row["sha256"]: dict(row) for row in rows}
def get_all(self) -> dict[str, dict]:
"""Return all stats keyed by sha256."""
conn = self._ensure_conn()
rows = conn.execute("SELECT * FROM model_stats").fetchall()
return {row["sha256"]: dict(row) for row in rows}
def count(self) -> int:
"""Return number of rows in model_stats."""
conn = self._ensure_conn()
return conn.execute("SELECT COUNT(*) FROM model_stats").fetchone()[0]
def close(self) -> None:
if self._conn:
self._conn.close()
self._conn = None
```
**Step 4: Run tests to verify they pass**
Run: `cd custom_nodes/ComfyUI-LM-Remote && python -m pytest tests/test_stats_db.py -v`
Expected: All 8 tests PASS
**Step 5: Commit**
```bash
cd custom_nodes/ComfyUI-LM-Remote
git add stats_db.py tests/test_stats_db.py
git commit -m "feat: add SQLite stats database layer for CivitAI stats"
```
---
### Task 2: Stats Fetch Service (CivitAI API Client)
**Files:**
- Create: `custom_nodes/ComfyUI-LM-Remote/stats_service.py`
- Create: `custom_nodes/ComfyUI-LM-Remote/tests/test_stats_service.py`
**Step 1: Write failing tests**
```python
# tests/test_stats_service.py
"""Tests for the CivitAI stats fetch service."""
import asyncio
import json
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from pathlib import Path
@pytest.fixture
def db(tmp_path):
from stats_db import StatsDB
db = StatsDB(tmp_path / "test.db")
db.init()
return db
@pytest.fixture
def service(db):
from stats_service import StatsFetchService
return StatsFetchService(db)
def test_extract_version_stats():
"""Extract stats from a CivitAI model API response."""
from stats_service import extract_version_stats
model_response = {
"id": 123,
"modelVersions": [
{
"id": 456,
"files": [{"hashes": {"SHA256": "AABB"}}],
"stats": {
"downloadCount": 5000,
"ratingCount": 120,
"rating": 4.8,
"thumbsUpCount": 89,
},
},
{
"id": 789,
"files": [{"hashes": {"SHA256": "CCDD"}}],
"stats": {
"downloadCount": 1000,
"ratingCount": 30,
"rating": 3.5,
"thumbsUpCount": 20,
},
},
],
}
result = extract_version_stats(model_response)
assert len(result) == 2
assert result[0][0] == "aabb" # sha256 lowercased
assert result[0][1]["download_count"] == 5000
assert result[0][1]["civitai_model_id"] == 123
assert result[0][1]["civitai_version_id"] == 456
assert result[1][0] == "ccdd"
def test_extract_version_stats_no_hash():
"""Versions without a SHA256 hash are skipped."""
from stats_service import extract_version_stats
model_response = {
"id": 1,
"modelVersions": [
{"id": 10, "files": [{"hashes": {}}], "stats": {"downloadCount": 1}},
],
}
result = extract_version_stats(model_response)
assert result == []
def test_extract_version_stats_no_stats():
"""Versions without stats get zeros."""
from stats_service import extract_version_stats
model_response = {
"id": 1,
"modelVersions": [
{"id": 10, "files": [{"hashes": {"SHA256": "AABB"}}]},
],
}
result = extract_version_stats(model_response)
assert len(result) == 1
assert result[0][1]["download_count"] == 0
@pytest.mark.asyncio
async def test_fetch_stats_for_models(service):
"""fetch_stats_for_models calls CivitAI API and stores results."""
models = [
{"sha256": "aabb", "civitai_model_id": 123},
{"sha256": "ccdd", "civitai_model_id": 456},
]
mock_response_123 = {
"id": 123,
"modelVersions": [{
"id": 10,
"files": [{"hashes": {"SHA256": "AABB"}}],
"stats": {"downloadCount": 500, "ratingCount": 10,
"rating": 4.0, "thumbsUpCount": 8},
}],
}
mock_response_456 = {
"id": 456,
"modelVersions": [{
"id": 20,
"files": [{"hashes": {"SHA256": "CCDD"}}],
"stats": {"downloadCount": 100, "ratingCount": 5,
"rating": 3.0, "thumbsUpCount": 2},
}],
}
async def mock_fetch(model_id):
return {123: mock_response_123, 456: mock_response_456}.get(model_id)
with patch.object(service, "_fetch_model", side_effect=mock_fetch):
count = await service.fetch_stats_for_models(models)
assert count == 2
result = service.db.get_by_hashes(["aabb", "ccdd"])
assert result["aabb"]["download_count"] == 500
assert result["ccdd"]["download_count"] == 100
```
**Step 2: Run tests to verify they fail**
Run: `cd custom_nodes/ComfyUI-LM-Remote && python -m pytest tests/test_stats_service.py -v`
Expected: FAIL — `ModuleNotFoundError: No module named 'stats_service'`
**Step 3: Write implementation**
```python
# stats_service.py
"""Service for fetching CivitAI stats and storing them in the local DB."""
from __future__ import annotations
import asyncio
import logging
from typing import Any
import aiohttp
from .stats_db import StatsDB
logger = logging.getLogger(__name__)
_CIVITAI_API = "https://civitai.com/api/v1"
_BATCH_SIZE = 20 # models per API request
_RATE_LIMIT_DELAY = 1.5 # seconds between batches
def extract_version_stats(model_data: dict) -> list[tuple[str, dict]]:
"""Extract (sha256, stats_dict) pairs from a CivitAI model response.
Each model has multiple versions; each version has files with hashes.
Returns one entry per version that has a SHA256 hash.
"""
model_id = model_data.get("id")
results = []
for version in model_data.get("modelVersions", []):
version_id = version.get("id")
sha256 = None
for f in version.get("files", []):
sha256 = (f.get("hashes") or {}).get("SHA256")
if sha256:
break
if not sha256:
continue
stats = version.get("stats") or {}
results.append((
sha256.lower(),
{
"civitai_model_id": model_id,
"civitai_version_id": version_id,
"download_count": stats.get("downloadCount", 0),
"rating": stats.get("rating", 0),
"rating_count": stats.get("ratingCount", 0),
"thumbs_up_count": stats.get("thumbsUpCount", 0),
},
))
return results
class StatsFetchService:
"""Fetches CivitAI stats for models and stores them in StatsDB."""
def __init__(self, db: StatsDB, api_key: str | None = None):
self.db = db
self._api_key = api_key
self._session: aiohttp.ClientSession | None = None
async def _get_session(self) -> aiohttp.ClientSession:
if self._session is None or self._session.closed:
headers = {}
if self._api_key:
headers["Authorization"] = f"Bearer {self._api_key}"
timeout = aiohttp.ClientTimeout(total=30)
self._session = aiohttp.ClientSession(
headers=headers, timeout=timeout
)
return self._session
async def _fetch_model(self, model_id: int) -> dict | None:
"""Fetch a single model's data from CivitAI API."""
url = f"{_CIVITAI_API}/models/{model_id}"
session = await self._get_session()
try:
async with session.get(url) as resp:
if resp.status == 429:
logger.warning("[LM-Remote] CivitAI rate limited, backing off")
await asyncio.sleep(5)
return None
if resp.status != 200:
logger.debug("[LM-Remote] CivitAI returned %d for model %d",
resp.status, model_id)
return None
return await resp.json()
except Exception as exc:
logger.warning("[LM-Remote] CivitAI fetch failed for model %d: %s",
model_id, exc)
return None
async def fetch_stats_for_models(
self,
models: list[dict],
progress_callback: Any = None,
) -> int:
"""Fetch stats from CivitAI for a list of models.
Args:
models: list of dicts with 'sha256' and 'civitai_model_id' keys.
progress_callback: optional async callable(current, total).
Returns:
Number of models successfully updated.
"""
# Deduplicate by model_id (multiple versions share the same model)
seen_model_ids: set[int] = set()
unique_models: list[dict] = []
for m in models:
mid = m.get("civitai_model_id")
if mid and mid not in seen_model_ids:
seen_model_ids.add(mid)
unique_models.append(m)
total = len(unique_models)
updated = 0
for i, model in enumerate(unique_models):
model_id = model["civitai_model_id"]
data = await self._fetch_model(model_id)
if data:
rows = extract_version_stats(data)
if rows:
self.db.upsert_batch(rows)
updated += len(rows)
if progress_callback:
await progress_callback(i + 1, total)
# Rate limiting between requests
if i < total - 1:
await asyncio.sleep(_RATE_LIMIT_DELAY)
return updated
async def close(self) -> None:
if self._session and not self._session.closed:
await self._session.close()
self._session = None
```
**Step 4: Run tests to verify they pass**
Run: `cd custom_nodes/ComfyUI-LM-Remote && python -m pytest tests/test_stats_service.py -v`
Expected: All 5 tests PASS
**Step 5: Commit**
```bash
cd custom_nodes/ComfyUI-LM-Remote
git add stats_service.py tests/test_stats_service.py
git commit -m "feat: add CivitAI stats fetch service"
```
---
### Task 3: Proxy Response Enrichment
**Files:**
- Modify: `custom_nodes/ComfyUI-LM-Remote/proxy.py`
- Create: `custom_nodes/ComfyUI-LM-Remote/tests/test_proxy_enrichment.py`
**Step 1: Write failing tests**
```python
# tests/test_proxy_enrichment.py
"""Tests for proxy response enrichment with stats."""
import json
import pytest
from pathlib import Path
@pytest.fixture
def db(tmp_path):
from stats_db import StatsDB
db = StatsDB(tmp_path / "test.db")
db.init()
return db
def test_enrich_list_response(db):
"""Stats are merged into list response items."""
from proxy import _enrich_list_response
db.upsert("hash_a", {"download_count": 5000, "rating": 4.8,
"rating_count": 120, "thumbs_up_count": 89})
response_data = {
"items": [
{"sha256": "hash_a", "file_name": "lora_a.safetensors"},
{"sha256": "hash_b", "file_name": "lora_b.safetensors"},
],
"total": 2,
}
enriched = _enrich_list_response(response_data, db)
assert enriched["items"][0]["download_count"] == 5000
assert enriched["items"][0]["rating"] == 4.8
assert enriched["items"][0]["thumbs_up_count"] == 89
# Item without stats gets defaults
assert enriched["items"][1].get("download_count") is None
def test_enrich_empty_items(db):
"""Empty items list returns unchanged."""
from proxy import _enrich_list_response
response_data = {"items": [], "total": 0}
enriched = _enrich_list_response(response_data, db)
assert enriched["items"] == []
```
**Step 2: Run tests to verify they fail**
Run: `cd custom_nodes/ComfyUI-LM-Remote && python -m pytest tests/test_proxy_enrichment.py -v`
Expected: FAIL — `ImportError: cannot import name '_enrich_list_response'`
**Step 3: Add enrichment function and integrate into proxy**
Add to `proxy.py` (after the existing imports):
```python
from .stats_db import StatsDB
# Lazy-initialized stats DB
_stats_db: StatsDB | None = None
def _get_stats_db() -> StatsDB:
global _stats_db
if _stats_db is None:
_stats_db = StatsDB()
_stats_db.init()
return _stats_db
def _enrich_list_response(response_data: dict, db: StatsDB | None = None) -> dict:
"""Merge stats from the local DB into a list API response."""
items = response_data.get("items", [])
if not items:
return response_data
if db is None:
db = _get_stats_db()
hashes = [item["sha256"] for item in items if item.get("sha256")]
if not hashes:
return response_data
stats_map = db.get_by_hashes(hashes)
for item in items:
sha = item.get("sha256")
if sha and sha in stats_map:
s = stats_map[sha]
item["download_count"] = s["download_count"]
item["rating"] = s["rating"]
item["rating_count"] = s["rating_count"]
item["thumbs_up_count"] = s["thumbs_up_count"]
return response_data
```
Then modify `_proxy_http` to intercept list responses. In the middleware function `lm_remote_proxy_middleware`, add after the `_should_proxy` check, a special case for list endpoints:
```python
# In lm_remote_proxy_middleware, replace the simple _should_proxy block:
if _should_proxy(path):
response = await _proxy_http(request)
# Enrich lora/checkpoint/embedding list responses with stats
if (path.endswith("/list") and path.startswith("/api/lm/")
and response.status == 200):
try:
body = json.loads(response.body)
enriched = _enrich_list_response(body)
response = web.json_response(
enriched,
status=response.status,
headers=dict(response.headers),
)
except Exception:
pass # Return original response on any error
return response
```
Add `import json` to the top of proxy.py.
**Step 4: Run tests to verify they pass**
Run: `cd custom_nodes/ComfyUI-LM-Remote && python -m pytest tests/test_proxy_enrichment.py -v`
Expected: All 2 tests PASS
**Step 5: Commit**
```bash
cd custom_nodes/ComfyUI-LM-Remote
git add proxy.py tests/test_proxy_enrichment.py
git commit -m "feat: enrich list API responses with CivitAI stats from local DB"
```
---
### Task 4: Fetch Stats API Endpoint + Toolbar Button Handler
**Files:**
- Modify: `custom_nodes/ComfyUI-LM-Remote/proxy.py` (add new routes)
- Modify: `custom_nodes/ComfyUI-LM-Remote/__init__.py` (register routes)
**Step 1: Add the fetch-stats and stats-status endpoints to proxy.py**
Add two new handler functions and register them in the middleware:
```python
async def _handle_fetch_stats(request: web.Request) -> web.Response:
"""Bulk fetch CivitAI stats for all models."""
from .stats_service import StatsFetchService
from .remote_client import RemoteLoraClient
db = _get_stats_db()
client = RemoteLoraClient.get_instance()
# Get all loras from remote
try:
items = await client._get_lora_list_cached()
except Exception as exc:
return web.json_response(
{"success": False, "error": f"Failed to fetch lora list: {exc}"},
status=500,
)
# Build list of models with civitai IDs
models = []
for item in items:
civitai = item.get("civitai") or {}
model_id = civitai.get("modelId")
sha256 = item.get("sha256")
if model_id and sha256:
models.append({"sha256": sha256, "civitai_model_id": model_id})
if not models:
return web.json_response({"success": True, "updated": 0,
"message": "No models with CivitAI IDs found"})
# Fetch stats (with progress via send_sync if available)
service = StatsFetchService(db)
try:
server = _get_prompt_server()
async def progress(current, total):
server.send_sync("lm_stats_progress", {
"current": current, "total": total,
})
updated = await service.fetch_stats_for_models(models, progress_callback=progress)
except Exception:
updated = await service.fetch_stats_for_models(models)
finally:
await service.close()
return web.json_response({"success": True, "updated": updated, "total": len(models)})
async def _handle_stats_status(request: web.Request) -> web.Response:
"""Check if stats DB is populated."""
db = _get_stats_db()
count = db.count()
return web.json_response({"success": True, "count": count})
```
Add these to a new `_EXTRA_API_HANDLERS` dict:
```python
_EXTRA_API_HANDLERS = {
"/api/lm-extra/fetch-stats": _handle_fetch_stats,
"/api/lm-extra/stats-status": _handle_stats_status,
}
```
In `lm_remote_proxy_middleware`, add a check before the proxy check:
```python
# Extra API endpoints (stats, etc.)
extra_handler = _EXTRA_API_HANDLERS.get(path)
if extra_handler is not None:
return await extra_handler(request)
```
**Step 2: Test manually (no automated test — depends on remote + CivitAI API)**
Start ComfyUI, call `POST /api/lm-extra/fetch-stats` via curl or browser console.
**Step 3: Commit**
```bash
cd custom_nodes/ComfyUI-LM-Remote
git add proxy.py
git commit -m "feat: add fetch-stats and stats-status API endpoints"
```
---
### Task 5: Sort Handling for Stats-Based Sort Keys
**Files:**
- Modify: `custom_nodes/ComfyUI-LM-Remote/proxy.py`
- Create: `custom_nodes/ComfyUI-LM-Remote/tests/test_stats_sort.py`
**Step 1: Write failing test**
```python
# tests/test_stats_sort.py
"""Tests for stats-based sorting."""
import pytest
from pathlib import Path
@pytest.fixture
def db(tmp_path):
from stats_db import StatsDB
db = StatsDB(tmp_path / "test.db")
db.init()
return db
def test_sort_by_downloads_desc(db):
"""Items sorted by download_count descending."""
from proxy import _sort_items_by_stats
db.upsert("h1", {"download_count": 100, "rating": 0, "rating_count": 0, "thumbs_up_count": 0})
db.upsert("h2", {"download_count": 5000, "rating": 0, "rating_count": 0, "thumbs_up_count": 0})
db.upsert("h3", {"download_count": 500, "rating": 0, "rating_count": 0, "thumbs_up_count": 0})
items = [
{"sha256": "h1", "file_name": "a"},
{"sha256": "h2", "file_name": "b"},
{"sha256": "h3", "file_name": "c"},
]
sorted_items = _sort_items_by_stats(items, "downloads", "desc", db)
assert [i["sha256"] for i in sorted_items] == ["h2", "h3", "h1"]
def test_sort_by_rating_desc(db):
"""Items sorted by rating descending."""
from proxy import _sort_items_by_stats
db.upsert("h1", {"download_count": 0, "rating": 3.0, "rating_count": 0, "thumbs_up_count": 0})
db.upsert("h2", {"download_count": 0, "rating": 4.8, "rating_count": 0, "thumbs_up_count": 0})
items = [
{"sha256": "h1", "file_name": "a"},
{"sha256": "h2", "file_name": "b"},
]
sorted_items = _sort_items_by_stats(items, "rating", "desc", db)
assert sorted_items[0]["sha256"] == "h2"
def test_sort_items_without_stats_go_last(db):
"""Items without stats sort to the end."""
from proxy import _sort_items_by_stats
db.upsert("h1", {"download_count": 100, "rating": 0, "rating_count": 0, "thumbs_up_count": 0})
items = [
{"sha256": "no_stats", "file_name": "x"},
{"sha256": "h1", "file_name": "a"},
]
sorted_items = _sort_items_by_stats(items, "downloads", "desc", db)
assert sorted_items[0]["sha256"] == "h1"
assert sorted_items[1]["sha256"] == "no_stats"
def test_is_stats_sort():
"""Identify stats-based sort keys."""
from proxy import _is_stats_sort
assert _is_stats_sort("downloads") is True
assert _is_stats_sort("rating") is True
assert _is_stats_sort("thumbsup") is True
assert _is_stats_sort("name") is False
assert _is_stats_sort("date") is False
```
**Step 2: Run tests to verify they fail**
Run: `cd custom_nodes/ComfyUI-LM-Remote && python -m pytest tests/test_stats_sort.py -v`
Expected: FAIL
**Step 3: Implement sort functions**
Add to `proxy.py`:
```python
_STATS_SORT_KEYS = {"downloads", "rating", "thumbsup"}
_STATS_SORT_COLUMN = {
"downloads": "download_count",
"rating": "rating",
"thumbsup": "thumbs_up_count",
}
def _is_stats_sort(sort_key: str) -> bool:
return sort_key in _STATS_SORT_KEYS
def _sort_items_by_stats(
items: list[dict],
sort_key: str,
order: str,
db: StatsDB | None = None,
) -> list[dict]:
"""Sort items by a stats column, items without stats go last."""
if db is None:
db = _get_stats_db()
column = _STATS_SORT_COLUMN.get(sort_key, "download_count")
hashes = [item["sha256"] for item in items if item.get("sha256")]
stats_map = db.get_by_hashes(hashes)
reverse = order == "desc"
def sort_key_fn(item):
sha = item.get("sha256")
if sha and sha in stats_map:
return (0, stats_map[sha].get(column, 0))
return (1, 0) # no stats → sort last
return sorted(items, key=sort_key_fn, reverse=reverse)
```
Then update the list response enrichment in the middleware to handle stats sorting. When `sort_by` is a stats key, the proxy needs to:
1. Request ALL items from remote (override page_size)
2. Sort locally
3. Paginate locally
Update the middleware's list interception logic:
```python
if _should_proxy(path):
# Stats-based sorting: intercept list requests with stats sort keys
if (path.endswith("/list") and path.startswith("/api/lm/")
and request.method == "GET"):
sort_by = request.query.get("sort_by", "name")
sort_key = sort_by.split(":")[0] if ":" in sort_by else sort_by
if _is_stats_sort(sort_key):
return await _proxy_list_with_stats_sort(request, sort_by)
response = await _proxy_http(request)
# Enrich list responses with stats
if (path.endswith("/list") and path.startswith("/api/lm/")
and response.status == 200):
try:
body = json.loads(response.body)
enriched = _enrich_list_response(body)
return web.json_response(enriched)
except Exception:
pass
return response
```
Add the full-list-sort handler:
```python
async def _proxy_list_with_stats_sort(
request: web.Request, sort_by: str
) -> web.Response:
"""Fetch all items from remote, sort by stats locally, paginate."""
import copy
from urllib.parse import urlencode, parse_qs, urlparse
page = int(request.query.get("page", "1"))
page_size = int(request.query.get("page_size", "20"))
sort_key, _, order = sort_by.partition(":")
order = order or "desc"
# Build a modified request to get all items
query = dict(request.query)
query["page"] = "1"
query["page_size"] = "9999"
query["sort_by"] = "name:asc" # use default sort on remote side
remote_url = f"{remote_config.remote_url}{request.path}?{urlencode(query)}"
session = await _get_proxy_session()
try:
async with session.get(remote_url) as resp:
if resp.status != 200:
return await _proxy_http(request)
data = await resp.json()
except Exception:
return await _proxy_http(request)
items = data.get("items", [])
db = _get_stats_db()
# Enrich all items with stats
_enrich_list_response(data, db)
# Sort by stats
sorted_items = _sort_items_by_stats(items, sort_key, order, db)
# Paginate
total = len(sorted_items)
start = (page - 1) * page_size
end = start + page_size
page_items = sorted_items[start:end]
return web.json_response({
"items": page_items,
"total": total,
"page": page,
"page_size": page_size,
"total_pages": (total + page_size - 1) // page_size,
})
```
**Step 4: Run tests to verify they pass**
Run: `cd custom_nodes/ComfyUI-LM-Remote && python -m pytest tests/test_stats_sort.py -v`
Expected: All 4 tests PASS
**Step 5: Commit**
```bash
cd custom_nodes/ComfyUI-LM-Remote
git add proxy.py tests/test_stats_sort.py
git commit -m "feat: add stats-based sorting (downloads, rating, thumbsup)"
```
---
### Task 6: Frontend — Card Badges + Sort Dropdown + Toolbar Button
**Files:**
- Create: `custom_nodes/ComfyUI-LM-Remote/web/comfyui/lm_stats_ui.js`
**Step 1: Write the JS extension**
```javascript
// web/comfyui/lm_stats_ui.js
/**
* CivitAI Stats UI — card badges, sort dropdown options, fetch button.
*
* Injected by ComfyUI-LM-Remote. Patches the LoRA Manager web UI
* (served from the remote) to show download counts, ratings, and
* thumbs-up counts on model cards, and adds stats-based sort options.
*/
(function () {
"use strict";
// ── Compact number formatting ──────────────────────────────────
function formatCompact(n) {
if (n == null) return null;
if (n >= 1_000_000) return (n / 1_000_000).toFixed(1).replace(/\.0$/, "") + "M";
if (n >= 1_000) return (n / 1_000).toFixed(1).replace(/\.0$/, "") + "k";
return String(n);
}
// ── Badge creation ─────────────────────────────────────────────
function createStatBadge(icon, value, title) {
if (value == null) return null;
const badge = document.createElement("span");
badge.className = "lm-stat-badge";
badge.title = title;
badge.innerHTML = ` ${formatCompact(value)}`;
return badge;
}
// ── Inject CSS ─────────────────────────────────────────────────
function injectStyles() {
if (document.getElementById("lm-stats-styles")) return;
const style = document.createElement("style");
style.id = "lm-stats-styles";
style.textContent = `
.lm-stat-badges {
display: flex;
gap: 6px;
flex-wrap: wrap;
padding: 2px 6px;
}
.lm-stat-badge {
display: inline-flex;
align-items: center;
gap: 3px;
font-size: 11px;
padding: 1px 5px;
border-radius: 4px;
background: rgba(255,255,255,0.1);
color: rgba(255,255,255,0.8);
white-space: nowrap;
}
.lm-stat-badge i {
font-size: 10px;
opacity: 0.7;
}
`;
document.head.appendChild(style);
}
// ── Patch model cards ──────────────────────────────────────────
function patchCards() {
const cards = document.querySelectorAll(".model-card:not([data-stats-patched])");
cards.forEach((card) => {
card.setAttribute("data-stats-patched", "1");
const downloads = card.dataset.downloadCount;
const rating = card.dataset.rating;
const thumbsUp = card.dataset.thumbsUpCount;
if (!downloads && !rating && !thumbsUp) return;
const container = document.createElement("div");
container.className = "lm-stat-badges";
const dlBadge = createStatBadge("download", downloads ? Number(downloads) : null, "Downloads");
const ratingBadge = createStatBadge("star", rating ? Number(Number(rating).toFixed(1)) : null, "Rating");
const thumbsBadge = createStatBadge("thumbs-up", thumbsUp ? Number(thumbsUp) : null, "Likes");
[dlBadge, ratingBadge, thumbsBadge].forEach((b) => {
if (b) container.appendChild(b);
});
if (container.children.length > 0) {
const footer = card.querySelector(".card-footer") || card.querySelector(".card-name");
if (footer) {
footer.parentNode.insertBefore(container, footer);
}
}
});
}
// ── Patch sort dropdown ────────────────────────────────────────
function patchSortDropdown() {
const select = document.getElementById("sortSelect");
if (!select || select.querySelector('[value="downloads:desc"]')) return;
const group = document.createElement("optgroup");
group.label = "CivitAI Stats";
const options = [
["downloads:desc", "Most downloaded"],
["downloads:asc", "Least downloaded"],
["rating:desc", "Highest rated"],
["rating:asc", "Lowest rated"],
["thumbsup:desc", "Most liked"],
["thumbsup:asc", "Least liked"],
];
options.forEach(([value, label]) => {
const opt = document.createElement("option");
opt.value = value;
opt.textContent = label;
group.appendChild(opt);
});
select.appendChild(group);
}
// ── Toolbar "Fetch Stats" button ───────────────────────────────
function addFetchStatsButton() {
const toolbar = document.querySelector(".action-buttons");
if (!toolbar || document.getElementById("fetchStatsBtn")) return;
const group = document.createElement("div");
group.className = "control-group";
group.innerHTML = `
`;
// Insert before the bulk operations button
const bulkBtn = document.getElementById("bulkOperationsBtn");
if (bulkBtn && bulkBtn.parentNode) {
toolbar.insertBefore(group, bulkBtn.parentNode);
} else {
toolbar.appendChild(group);
}
group.querySelector("button").addEventListener("click", async () => {
const btn = document.getElementById("fetchStatsBtn");
btn.disabled = true;
btn.innerHTML = ' Fetching...';
try {
const resp = await fetch("/api/lm-extra/fetch-stats", { method: "POST" });
const data = await resp.json();
if (data.success) {
btn.innerHTML = ` ${data.updated} updated`;
setTimeout(() => {
btn.innerHTML = ' Fetch Stats';
btn.disabled = false;
}, 3000);
// Trigger page reload to show new stats
const sortSelect = document.getElementById("sortSelect");
if (sortSelect) {
sortSelect.dispatchEvent(new Event("change"));
}
} else {
throw new Error(data.error || "Unknown error");
}
} catch (err) {
btn.innerHTML = ' Error';
console.error("[LM-Stats] Fetch failed:", err);
setTimeout(() => {
btn.innerHTML = ' Fetch Stats';
btn.disabled = false;
}, 3000);
}
});
}
// ── Observe DOM for card rendering ─────────────────────────────
function startObserver() {
injectStyles();
patchSortDropdown();
addFetchStatsButton();
patchCards();
const observer = new MutationObserver(() => {
patchCards();
patchSortDropdown();
addFetchStatsButton();
});
observer.observe(document.body, { childList: true, subtree: true });
}
// ── Init ───────────────────────────────────────────────────────
if (document.readyState === "loading") {
document.addEventListener("DOMContentLoaded", startObserver);
} else {
startObserver();
}
})();
```
**Step 2: Ensure the card rendering stores stats as dataset attributes**
The card badges read from `card.dataset.downloadCount` etc. The ModelCard.js in the original manager creates cards from API response data. Since we enrich the list response in the proxy, these fields will be in the JSON. But we need the card rendering to put them in `data-*` attributes.
Two approaches:
1. The original `ModelCard.js` already stores all API fields as dataset attributes on the card → check if this is the case
2. If not, add a small DOM patch in `lm_stats_ui.js` that reads from a data store
Check `ModelCard.js` — it stores specific fields in `dataset`. We need to patch the card creation to also store stats fields. Add this to the MutationObserver in `lm_stats_ui.js`:
The cleaner approach: instead of relying on dataset attributes, intercept the fetch response and maintain a JS-side lookup map. Update `lm_stats_ui.js` to also intercept the `/api/lm/loras/list` fetch response:
Add before the IIFE's `startObserver`:
```javascript
// ── Intercept list API to capture stats ────────────────────────
const _statsMap = {};
const _origFetch = window.fetch;
window.fetch = async function (...args) {
const response = await _origFetch.apply(this, args);
const url = typeof args[0] === "string" ? args[0] : args[0]?.url;
if (url && url.includes("/api/lm/") && url.includes("/list")) {
try {
const clone = response.clone();
const data = await clone.json();
(data.items || []).forEach((item) => {
if (item.sha256 && item.download_count != null) {
_statsMap[item.sha256] = {
download_count: item.download_count,
rating: item.rating,
thumbs_up_count: item.thumbs_up_count,
};
}
});
} catch (e) { /* ignore */ }
}
return response;
};
```
Then update `patchCards` to use `_statsMap` instead of dataset attributes:
```javascript
function patchCards() {
const cards = document.querySelectorAll(".model-card:not([data-stats-patched])");
cards.forEach((card) => {
card.setAttribute("data-stats-patched", "1");
const sha = card.dataset.sha256 || card.dataset.hash;
if (!sha || !_statsMap[sha]) return;
const stats = _statsMap[sha];
// ... create badges from stats.download_count, stats.rating, etc.
});
}
```
**Step 3: Commit**
```bash
cd custom_nodes/ComfyUI-LM-Remote
git add web/comfyui/lm_stats_ui.js
git commit -m "feat: add stats UI — card badges, sort dropdown, fetch button"
```
---
### Task 7: Wire Up Proxy Cleanup + Integration Test
**Files:**
- Modify: `custom_nodes/ComfyUI-LM-Remote/proxy.py` (cleanup hook)
**Step 1: Add stats DB cleanup to shutdown hook**
In `_cleanup_proxy_session`:
```python
async def _cleanup_proxy_session(app) -> None:
global _proxy_session, _stats_db
if _proxy_session and not _proxy_session.closed:
await _proxy_session.close()
_proxy_session = None
if _stats_db:
_stats_db.close()
_stats_db = None
```
**Step 2: Manual integration test**
1. Start ComfyUI with ComfyUI-LM-Remote configured
2. Open the LoRA Manager web UI (`/loras`)
3. Click "Fetch Stats" button in toolbar
4. Watch progress (check browser console and server logs)
5. After completion, verify cards show download/rating/thumbs-up badges
6. Test sort dropdown: select "Most downloaded" — cards should reorder
7. Refresh page — stats persist (stored in SQLite)
**Step 3: Commit**
```bash
cd custom_nodes/ComfyUI-LM-Remote
git add proxy.py
git commit -m "feat: wire up stats DB cleanup on shutdown"
```
---
## File Summary
| File | Action | Purpose |
|------|--------|---------|
| `stats_db.py` | CREATE | SQLite wrapper for civitai_stats.db |
| `stats_service.py` | CREATE | CivitAI API fetch + batch logic |
| `proxy.py` | MODIFY | Response enrichment, sort handling, new endpoints |
| `web/comfyui/lm_stats_ui.js` | CREATE | Card badges, sort dropdown, fetch button |
| `tests/test_stats_db.py` | CREATE | DB layer tests |
| `tests/test_stats_service.py` | CREATE | Service tests |
| `tests/test_proxy_enrichment.py` | CREATE | Enrichment tests |
| `tests/test_stats_sort.py` | CREATE | Sort tests |