ComfyUI-LM-Remote/docs/plans/2026-03-15-civitai-stats-shim.md

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

# 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

# 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

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

# 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

# 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

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

# 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):

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:

# 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

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:

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:

_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:

    # 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

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

# 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:

_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:

    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:

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

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

// 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 = `<i class="fas fa-${icon}"></i> ${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 = `
            <button id="fetchStatsBtn" data-action="fetch-stats"
                    title="Fetch CivitAI stats (downloads, ratings, likes)">
                <i class="fas fa-chart-bar"></i> <span>Fetch Stats</span>
            </button>
        `;

        // 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 = '<i class="fas fa-spinner fa-spin"></i> <span>Fetching...</span>';

            try {
                const resp = await fetch("/api/lm-extra/fetch-stats", { method: "POST" });
                const data = await resp.json();
                if (data.success) {
                    btn.innerHTML = `<i class="fas fa-check"></i> <span>${data.updated} updated</span>`;
                    setTimeout(() => {
                        btn.innerHTML = '<i class="fas fa-chart-bar"></i> <span>Fetch Stats</span>';
                        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 = '<i class="fas fa-exclamation-triangle"></i> <span>Error</span>';
                console.error("[LM-Stats] Fetch failed:", err);
                setTimeout(() => {
                    btn.innerHTML = '<i class="fas fa-chart-bar"></i> <span>Fetch Stats</span>';
                    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:

    // ── 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:

    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

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:

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

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