From c084fb3da885b6c811798034c2cd2b594da99a0d Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 15 Mar 2026 13:49:47 +0100 Subject: [PATCH] docs: add CivitAI stats shim implementation plan 7-task plan covering: SQLite DB layer, fetch service, proxy enrichment, stats-based sorting, frontend UI (badges/sort/button), and cleanup. Co-Authored-By: Claude Opus 4.6 --- docs/plans/2026-03-15-civitai-stats-shim.md | 1383 +++++++++++++++++++ 1 file changed, 1383 insertions(+) create mode 100644 docs/plans/2026-03-15-civitai-stats-shim.md diff --git a/docs/plans/2026-03-15-civitai-stats-shim.md b/docs/plans/2026-03-15-civitai-stats-shim.md new file mode 100644 index 0000000..877b923 --- /dev/null +++ b/docs/plans/2026-03-15-civitai-stats-shim.md @@ -0,0 +1,1383 @@ +# 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 |