Compare commits

..

13 Commits

Author SHA1 Message Date
Ethanfel b39ea7a647 fix: address bug review findings
- Debounce MutationObserver (200ms) and scope to #modelGrid when available
- Sanitize data.updated via parseInt before innerHTML to prevent XSS
- Guard pagination against page_size=0 (division by zero) and negative values
- Reject non-POST requests on /api/lm-extra/fetch-stats (state-mutating)
- Early return in upsert_batch for empty rows list

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-15 14:48:14 +01:00
Ethanfel f6e39df157 fix: address code review findings
- Add retry (2 attempts) on CivitAI 429 rate limit instead of silently skipping
- Use check_same_thread=False for SQLite connection (aiohttp context)
- Add debug logging to silent exception handlers in proxy

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-15 14:44:12 +01:00
Ethanfel 63ddaaa70d chore: add stats DB cleanup and gitignore
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-15 14:40:32 +01:00
Ethanfel 467aea0f47 feat: add stats API endpoints and frontend JS with proxy injection
- Add /api/lm-extra/fetch-stats (POST) endpoint that bulk-fetches CivitAI
  stats for all remote models and stores them in the local stats DB with
  optional WebSocket progress broadcasting
- Add /api/lm-extra/stats-status (GET) endpoint to check DB population
- Add /lm-extra/stats-ui.js handler that serves the client-side JS file
- Create static/lm_stats_ui.js: intercepts list API responses to capture
  stats, patches model cards with download/rating/likes badges, adds
  CivitAI Stats sort options to the dropdown, and adds a Fetch Stats
  toolbar button
- Inject stats-ui.js script tag into proxied HTML page responses

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-15 14:39:56 +01:00
Ethanfel 1ce1e7d48a feat: add stats-based sorting (downloads, rating, thumbsup)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-15 14:36:55 +01:00
Ethanfel c3cd35add7 feat: enrich list API responses with CivitAI stats from local DB
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-15 14:34:51 +01:00
Ethanfel c1e124bcc8 feat: add CivitAI stats fetch service
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-15 14:32:52 +01:00
Ethanfel a8308874d2 feat: add civitai_api_key to config
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-15 14:31:10 +01:00
Ethanfel 44088649d8 feat: add SQLite stats database layer for CivitAI stats
Implements StatsDB class with upsert, batch upsert, get_by_hashes,
get_all, count, and close methods backed by a WAL-mode SQLite database.
Also guards __init__.py relative imports behind __package__ check so
pytest can run without ComfyUI context, and removes empty tests/__init__.py
to prevent pytest Package collector from traversing into the plugin root.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-15 14:29:59 +01:00
Ethanfel aefb129e07 chore: add test infrastructure with conftest for import paths
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-15 14:14:08 +01:00
Ethanfel be6d8bcef7 chore: add .gitignore with worktrees exclusion 2026-03-15 14:12:21 +01:00
Ethanfel c084fb3da8 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 <noreply@anthropic.com>
2026-03-15 13:49:47 +01:00
Ethanfel afb3840105 docs: add CivitAI extended stats shim design
Design for a shim that stores CivitAI stats (downloads, rating,
thumbs up) in a separate SQLite DB and enriches the list API
response via the proxy middleware. Includes bulk fetch button,
card badges, and new sort options.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-15 13:41:51 +01:00
18 changed files with 2732 additions and 37 deletions
+2
View File
@@ -0,0 +1,2 @@
.worktrees
civitai_stats.db
+4
View File
@@ -19,6 +19,10 @@ import logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Guard relative imports: only execute when loaded as a proper package
# (i.e. inside ComfyUI). Skipped when the file is imported standalone by
# pytest or other tools that add the directory directly to sys.path.
if __package__:
# ── Import node classes ──────────────────────────────────────────────── # ── Import node classes ────────────────────────────────────────────────
from .nodes.lora_loader import LoraLoaderRemoteLM, LoraTextLoaderRemoteLM from .nodes.lora_loader import LoraLoaderRemoteLM, LoraTextLoaderRemoteLM
from .nodes.lora_stacker import LoraStackerRemoteLM from .nodes.lora_stacker import LoraStackerRemoteLM
+2 -1
View File
@@ -1,5 +1,6 @@
{ {
"remote_url": "http://192.168.1.3:8188", "remote_url": "http://192.168.1.3:8188",
"timeout": 30, "timeout": 30,
"path_mappings": {} "path_mappings": {},
"civitai_api_key": ""
} }
+6
View File
@@ -19,6 +19,7 @@ class RemoteConfig:
self.remote_url: str = "" self.remote_url: str = ""
self.timeout: int = 30 self.timeout: int = 30
self.path_mappings: dict[str, str] = {} self.path_mappings: dict[str, str] = {}
self.civitai_api_key: str = ""
self._load() self._load()
# ------------------------------------------------------------------ # ------------------------------------------------------------------
@@ -35,6 +36,7 @@ class RemoteConfig:
self.remote_url = data.get("remote_url", "") self.remote_url = data.get("remote_url", "")
self.timeout = int(data.get("timeout", 30)) self.timeout = int(data.get("timeout", 30))
self.path_mappings = data.get("path_mappings", {}) self.path_mappings = data.get("path_mappings", {})
self.civitai_api_key = data.get("civitai_api_key", "")
except Exception as exc: except Exception as exc:
logger.warning("[LM-Remote] Failed to read config.json: %s", exc) logger.warning("[LM-Remote] Failed to read config.json: %s", exc)
@@ -44,6 +46,10 @@ class RemoteConfig:
if env_timeout: if env_timeout:
self.timeout = int(env_timeout) self.timeout = int(env_timeout)
env_api_key = os.environ.get("LM_CIVITAI_API_KEY", "")
if env_api_key:
self.civitai_api_key = env_api_key
# Strip trailing slash # Strip trailing slash
self.remote_url = self.remote_url.rstrip("/") self.remote_url = self.remote_url.rstrip("/")
+3
View File
@@ -0,0 +1,3 @@
# Root-level conftest: prevent pytest from collecting the package __init__.py
# as a test module (it uses relative imports that require ComfyUI context).
collect_ignore = ["__init__.py"]
@@ -0,0 +1,157 @@
# CivitAI Extended Stats Shim — Design
## Problem
The LoRA Manager fetches metadata from CivitAI but discards stats (download count, rating, thumbs up). Users want to sort and filter loras by popularity/quality, and see these stats on model cards — like CivitAI's own UI.
## Constraints
- No modifications to the original LoRA Manager codebase (avoid merge conflicts on upstream sync)
- Lives inside ComfyUI-LM-Remote, which already has proxy middleware intercepting API responses
- Must handle CivitAI API rate limits gracefully when bulk-fetching stats
- Stats should persist across restarts (SQLite)
## Architecture
### Hybrid Shim Approach
- **Separate SQLite DB** (`civitai_stats.db`) for extended stats — avoids schema conflicts with the manager's own cache
- **Proxy enrichment** — the existing middleware intercepts `/api/lm/loras/list` responses from the remote and merges stats into each item before forwarding to the browser
- **New API endpoints** on LM-Remote for the bulk fetch button and status checks
- **Frontend JS** injected via LM-Remote's `web/comfyui/` directory, patches sort dropdown and card rendering
### Data Flow
```
List request:
Browser → GET /api/lm/loras/list
→ proxy fetches response from remote
→ enriches each item with stats from local civitai_stats.db
→ returns merged response to browser
Bulk fetch (initial):
Browser → POST /api/lm-extra/fetch-stats
→ handler reads models from remote list API
→ extracts civitai_model_id + sha256 for each
→ batches CivitAI API calls (/api/v1/models?ids=id1,id2,...)
→ stores stats in civitai_stats.db
→ WebSocket progress updates
→ frontend auto-refreshes grid when done
Piggyback (ongoing):
When proxy intercepts metadata sync responses that contain stats,
opportunistically store them in civitai_stats.db
```
## SQLite Schema
File: `civitai_stats.db` (stored alongside LM-Remote's config)
```sql
CREATE TABLE model_stats (
sha256 TEXT PRIMARY KEY,
civitai_model_id INTEGER,
civitai_version_id INTEGER,
download_count INTEGER,
rating REAL,
rating_count INTEGER,
thumbs_up_count INTEGER,
fetched_at REAL
);
```
Keyed by `sha256` — the common identifier between the manager's cache and CivitAI. Joins naturally regardless of file paths or renames.
## API Endpoints (LM-Remote)
| Method | Path | Purpose |
|--------|------|---------|
| POST | `/api/lm-extra/fetch-stats` | Bulk fetch stats for all models (button) |
| GET | `/api/lm-extra/stats-status` | Check if stats DB is populated, staleness info |
## Proxy Response Enrichment
The middleware intercepts the `/api/lm/loras/list` response and:
1. Parses the JSON response body
2. Collects all `sha256` values from items
3. Batch-queries `civitai_stats.db` (single SELECT with IN clause)
4. Merges `download_count`, `rating`, `rating_count`, `thumbs_up_count` into each item
5. Returns the enriched response
## Sorting
### New Sort Options
Added to the sort dropdown alongside existing options (name, date, size, usage):
| Value | Label |
|-------|-------|
| `downloads:desc` | Most downloaded |
| `downloads:asc` | Least downloaded |
| `rating:desc` | Highest rated |
| `rating:asc` | Lowest rated |
| `thumbsup:desc` | Most liked |
| `thumbsup:asc` | Least liked |
### Sort Implementation
When the proxy sees these sort values in the list request, sorting is handled locally:
1. Fetch the full (unpaginated) list from remote
2. Join with stats DB
3. Sort locally
4. Apply pagination
5. Return result
Items without stats sort to the end.
## Frontend Changes
All in `ComfyUI-LM-Remote/web/comfyui/`.
### Card Badges
Small badges on each model card (styled like the existing base model badge):
- Download icon + count (e.g., "12.3k")
- Star icon + rating (e.g., "4.8")
- Thumbs-up icon + count (e.g., "89")
Numbers use compact formatting (1000 → "1k", 12345 → "12.3k").
### Toolbar Button
A "Fetch Stats" button in the toolbar (next to existing Fetch/Download buttons):
- Chart-bar icon
- Triggers POST to `/api/lm-extra/fetch-stats`
- Shows progress via WebSocket
- Auto-refreshes grid on completion
### Sort Dropdown Patch
JS extension adds the new sort options to the existing `#sortSelect` dropdown on page load.
## Bulk Fetch Strategy
CivitAI's `/api/v1/models?ids=` endpoint accepts multiple model IDs per request.
1. Collect all unique `civitai_model_id` values from the lora list
2. Batch into groups (e.g., 50 IDs per request)
3. For each batch, fetch model data and extract `stats` from the matching version
4. Rate-limit: 1-2 requests/second with backoff on 429 responses
5. Broadcast progress via WebSocket (e.g., "Fetching stats: 150/300 models")
## Piggyback Strategy
When the existing metadata sync runs (triggered by the manager's refresh/fetch), the proxy can intercept the CivitAI API response if it flows through, and opportunistically extract stats. This covers new models added after the initial bulk fetch.
## Files to Create/Modify
| File | Action |
|------|--------|
| `stats_db.py` | NEW — SQLite wrapper for civitai_stats.db |
| `stats_service.py` | NEW — bulk fetch logic, CivitAI API calls, rate limiting |
| `proxy.py` | MODIFY — intercept list response, enrich with stats; add sort handling |
| `web/comfyui/lm_stats_ui.js` | NEW — card badges, toolbar button, sort dropdown patch |
File diff suppressed because it is too large Load Diff
+294 -3
View File
@@ -11,13 +11,22 @@ to the local ComfyUI frontend (the remote instance has no connected browsers).
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import json
import logging import logging
from pathlib import Path
import aiohttp import aiohttp
from aiohttp import web, WSMsgType from aiohttp import web, WSMsgType
try:
from .config import remote_config from .config import remote_config
except ImportError:
from config import remote_config # type: ignore[no-redef]
try:
from .remote_client import RemoteLoraClient from .remote_client import RemoteLoraClient
except ImportError:
from remote_client import RemoteLoraClient # type: ignore[no-redef]
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -178,6 +187,244 @@ _SEND_SYNC_HANDLERS = {
"/api/lm/register-nodes": _handle_register_nodes, "/api/lm/register-nodes": _handle_register_nodes,
} }
# ---------------------------------------------------------------------------
# Stats enrichment
# ---------------------------------------------------------------------------
_STATS_JS_PATH = Path(__file__).resolve().parent / "static" / "lm_stats_ui.js"
_stats_db = None
def _get_stats_db():
"""Lazy-initialize the stats database."""
global _stats_db
if _stats_db is None:
try:
from .stats_db import StatsDB as _StatsDBClass
except ImportError:
from stats_db import StatsDB as _StatsDBClass
_stats_db = _StatsDBClass()
_stats_db.init()
return _stats_db
def _enrich_list_response(response_data: dict, db=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
# ---------------------------------------------------------------------------
# Stats-based sorting
# ---------------------------------------------------------------------------
_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 True if sort_key is a stats-based sort."""
return sort_key in _STATS_SORT_KEYS
def _sort_items_by_stats(items, sort_key, order, db=None):
"""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 get_value(item):
sha = item.get("sha256")
if sha and sha in stats_map:
return stats_map[sha].get(column, 0)
return None
# Split items into those with stats and those without, then sort only
# the former so that items without stats always end up last.
with_stats = [i for i in items if get_value(i) is not None]
without_stats = [i for i in items if get_value(i) is None]
with_stats.sort(key=get_value, reverse=reverse)
return with_stats + without_stats
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."""
from urllib.parse import urlencode
page = max(1, int(request.query.get("page", "1")))
page_size = max(1, int(request.query.get("page_size", "20")))
sort_key, _, order = sort_by.partition(":")
order = order or "desc"
# Build query to fetch all items from remote with default sort
query = dict(request.query)
query["page"] = "1"
query["page_size"] = "9999"
query["sort_by"] = "name:asc"
remote_url = f"{remote_config.remote_url}{request.path}?{urlencode(query)}"
# Forward relevant headers
headers = {}
skip = {"host", "transfer-encoding", "connection", "keep-alive", "upgrade"}
for k, v in request.headers.items():
if k.lower() not in skip:
headers[k] = v
session = await _get_proxy_session()
try:
async with session.get(remote_url, headers=headers) 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 locally
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,
})
# ---------------------------------------------------------------------------
# Extra API endpoints (stats)
# ---------------------------------------------------------------------------
async def _handle_fetch_stats(request: web.Request) -> web.Response:
"""Bulk fetch CivitAI stats for all models."""
if request.method != "POST":
return web.json_response({"error": "POST required"}, status=405)
try:
from .stats_service import StatsFetchService
except ImportError:
from stats_service import StatsFetchService
db = _get_stats_db()
# Get the lora list from remote
session = await _get_proxy_session()
remote_url = f"{remote_config.remote_url}/api/lm/loras/list?page=1&page_size=9999"
try:
async with session.get(remote_url) as resp:
if resp.status != 200:
return web.json_response(
{"success": False, "error": "Failed to fetch lora list"},
status=500,
)
data = await resp.json()
except Exception as exc:
return web.json_response(
{"success": False, "error": f"Failed to fetch lora list: {exc}"},
status=500,
)
items = data.get("items", [])
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"})
service = StatsFetchService(db, api_key=remote_config.civitai_api_key or None)
try:
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()
return web.json_response({"success": True, "count": db.count()})
async def _handle_stats_ui_js(request: web.Request) -> web.Response:
"""Serve the stats UI JavaScript file."""
try:
content = _STATS_JS_PATH.read_text(encoding="utf-8")
return web.Response(
text=content,
content_type="application/javascript",
)
except Exception as exc:
return web.Response(status=404, text=str(exc))
_EXTRA_API_HANDLERS = {
"/api/lm-extra/fetch-stats": _handle_fetch_stats,
"/api/lm-extra/stats-status": _handle_stats_status,
"/lm-extra/stats-ui.js": _handle_stats_ui_js,
}
# Shared HTTP session for proxied requests (connection pooling) # Shared HTTP session for proxied requests (connection pooling)
_proxy_session: aiohttp.ClientSession | None = None _proxy_session: aiohttp.ClientSession | None = None
@@ -330,20 +577,64 @@ async def lm_remote_proxy_middleware(request: web.Request, handler):
if _is_ws_route(path): if _is_ws_route(path):
return await _proxy_ws(request) return await _proxy_ws(request)
# Extra API endpoints (stats, etc.)
extra_handler = _EXTRA_API_HANDLERS.get(path)
if extra_handler is not None:
return await extra_handler(request)
# Stats-based sorting: intercept list requests, fetch all from remote,
# sort locally, paginate locally
if (_should_proxy(path) and 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)
# Regular proxy routes # Regular proxy routes
if _should_proxy(path): if _should_proxy(path):
return await _proxy_http(request) response = await _proxy_http(request)
# Enrich model list responses with stats from local DB
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 as exc:
logger.debug("[LM-Remote] Stats enrichment failed: %s", exc)
# Inject stats UI script into HTML pages
if (path in _PROXY_PAGE_ROUTES or path.rstrip("/") in _PROXY_PAGE_ROUTES):
if response.status == 200:
content_type = response.headers.get("Content-Type", "")
if "text/html" in content_type:
try:
html = response.body.decode("utf-8")
inject = '<script src="/lm-extra/stats-ui.js"></script>'
html = html.replace("</body>", f"{inject}\n</body>")
return web.Response(
text=html,
content_type="text/html",
charset="utf-8",
status=response.status,
)
except Exception as exc:
logger.debug("[LM-Remote] HTML injection failed: %s", exc)
return response
# Not a LoRA Manager route — fall through # Not a LoRA Manager route — fall through
return await handler(request) return await handler(request)
async def _cleanup_proxy_session(app) -> None: async def _cleanup_proxy_session(app) -> None:
"""Shutdown hook to close the shared proxy session.""" """Shutdown hook to close the shared proxy session and stats DB."""
global _proxy_session global _proxy_session, _stats_db
if _proxy_session and not _proxy_session.closed: if _proxy_session and not _proxy_session.closed:
await _proxy_session.close() await _proxy_session.close()
_proxy_session = None _proxy_session = None
if _stats_db:
_stats_db.close()
_stats_db = None
def register_proxy(app) -> None: def register_proxy(app) -> None:
+2
View File
@@ -0,0 +1,2 @@
[pytest]
testpaths = tests
+3
View File
@@ -8,7 +8,10 @@ from typing import Any
import aiohttp import aiohttp
try:
from .config import remote_config from .config import remote_config
except ImportError:
from config import remote_config # type: ignore[no-redef]
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
+230
View File
@@ -0,0 +1,230 @@
// static/lm_stats_ui.js
/**
* CivitAI Stats UI — card badges, sort dropdown options, fetch button.
*
* Injected into LoRA Manager pages by the LM-Remote proxy.
* Reads stats from enriched /api/lm/loras/list responses and
* patches the DOM to show badges, sort options, and a fetch button.
*/
(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);
}
// ── 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,
rating_count: item.rating_count,
thumbs_up_count: item.thumbs_up_count,
};
}
});
} catch (e) { /* ignore parse errors */ }
}
return response;
};
// ── 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 sha = card.dataset.sha256;
if (!sha || !_statsMap[sha]) return;
const stats = _statsMap[sha];
const container = document.createElement("div");
container.className = "lm-stat-badges";
const dlBadge = createStatBadge("download", stats.download_count, "Downloads");
const ratingBadge = createStatBadge("star",
stats.rating ? Number(stats.rating.toFixed(1)) : null, "Rating");
const thumbsBadge = createStatBadge("thumbs-up", stats.thumbs_up_count, "Likes");
[dlBadge, ratingBadge, thumbsBadge].forEach((b) => {
if (b) container.appendChild(b);
});
if (container.children.length > 0) {
// Insert inside .model-info, after .model-name
const modelInfo = card.querySelector(".model-info");
if (modelInfo) {
modelInfo.appendChild(container);
}
}
});
}
// ── 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.closest(".control-group")) {
toolbar.insertBefore(group, bulkBtn.closest(".control-group"));
} 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 _origFetch("/api/lm-extra/fetch-stats", { method: "POST" });
const data = await resp.json();
if (data.success) {
const count = parseInt(data.updated, 10) || 0;
btn.innerHTML = `<i class="fas fa-check"></i> <span>${count} 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();
let debounceTimer = null;
const observer = new MutationObserver(() => {
if (debounceTimer) return;
debounceTimer = setTimeout(() => {
debounceTimer = null;
patchCards();
patchSortDropdown();
addFetchStatsButton();
}, 200);
});
const grid = document.getElementById("modelGrid");
observer.observe(grid || document.body, { childList: true, subtree: true });
}
// ── Init ───────────────────────────────────────────────────────
if (document.readyState === "loading") {
document.addEventListener("DOMContentLoaded", startObserver);
} else {
startObserver();
}
})();
+137
View File
@@ -0,0 +1,137 @@
"""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), check_same_thread=False)
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."""
if not rows:
return
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
+145
View File
@@ -0,0 +1,145 @@
"""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
try:
from .stats_db import StatsDB
except ImportError:
from stats_db import StatsDB
logger = logging.getLogger(__name__)
_CIVITAI_API = "https://civitai.com/api/v1"
_RATE_LIMIT_DELAY = 1.5 # seconds between requests
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, _retries: int = 2) -> dict | None:
"""Fetch a single model's data from CivitAI API."""
url = f"{_CIVITAI_API}/models/{model_id}"
session = await self._get_session()
for attempt in range(_retries + 1):
try:
async with session.get(url) as resp:
if resp.status == 429:
logger.warning("[LM-Remote] CivitAI rate limited, backing off (attempt %d)",
attempt + 1)
await asyncio.sleep(5 * (attempt + 1))
continue
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
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 model versions 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
+9
View File
@@ -0,0 +1,9 @@
"""Test configuration — adds package root to sys.path so bare imports work."""
import sys
from pathlib import Path
# Add the package root so `from stats_db import StatsDB` works in tests
# without requiring the full package to be installed.
_PACKAGE_ROOT = Path(__file__).resolve().parent.parent
if str(_PACKAGE_ROOT) not in sys.path:
sys.path.insert(0, str(_PACKAGE_ROOT))
+41
View File
@@ -0,0 +1,41 @@
"""Tests for proxy response enrichment with stats."""
import pytest
@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 no extra fields
assert "download_count" not in enriched["items"][1]
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"] == []
+99
View File
@@ -0,0 +1,99 @@
"""Tests for the CivitAI stats database layer."""
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
+121
View File
@@ -0,0 +1,121 @@
"""Tests for the CivitAI stats fetch service."""
import pytest
from unittest.mock import patch
@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}},
],
}
assert extract_version_stats(model_response) == []
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
+61
View File
@@ -0,0 +1,61 @@
"""Tests for stats-based sorting."""
import pytest
@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"}, {"sha256": "h2"}]
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"}, {"sha256": "h1"}]
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