diff --git a/__init__.py b/__init__.py index 731ae63..2e49a88 100644 --- a/__init__.py +++ b/__init__.py @@ -19,38 +19,42 @@ import logging logger = logging.getLogger(__name__) -# ── Import node classes ──────────────────────────────────────────────── -from .nodes.lora_loader import LoraLoaderRemoteLM, LoraTextLoaderRemoteLM -from .nodes.lora_stacker import LoraStackerRemoteLM -from .nodes.lora_randomizer import LoraRandomizerRemoteLM -from .nodes.lora_cycler import LoraCyclerRemoteLM -from .nodes.lora_pool import LoraPoolRemoteLM -from .nodes.save_image import SaveImageRemoteLM -from .nodes.wanvideo import WanVideoLoraSelectRemoteLM, WanVideoLoraTextSelectRemoteLM +# 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 ──────────────────────────────────────────────── + from .nodes.lora_loader import LoraLoaderRemoteLM, LoraTextLoaderRemoteLM + from .nodes.lora_stacker import LoraStackerRemoteLM + from .nodes.lora_randomizer import LoraRandomizerRemoteLM + from .nodes.lora_cycler import LoraCyclerRemoteLM + from .nodes.lora_pool import LoraPoolRemoteLM + from .nodes.save_image import SaveImageRemoteLM + from .nodes.wanvideo import WanVideoLoraSelectRemoteLM, WanVideoLoraTextSelectRemoteLM -# ── NODE_CLASS_MAPPINGS (how ComfyUI discovers nodes) ────────────────── -NODE_CLASS_MAPPINGS = { - LoraLoaderRemoteLM.NAME: LoraLoaderRemoteLM, - LoraTextLoaderRemoteLM.NAME: LoraTextLoaderRemoteLM, - LoraStackerRemoteLM.NAME: LoraStackerRemoteLM, - LoraRandomizerRemoteLM.NAME: LoraRandomizerRemoteLM, - LoraCyclerRemoteLM.NAME: LoraCyclerRemoteLM, - LoraPoolRemoteLM.NAME: LoraPoolRemoteLM, - SaveImageRemoteLM.NAME: SaveImageRemoteLM, - WanVideoLoraSelectRemoteLM.NAME: WanVideoLoraSelectRemoteLM, - WanVideoLoraTextSelectRemoteLM.NAME: WanVideoLoraTextSelectRemoteLM, -} + # ── NODE_CLASS_MAPPINGS (how ComfyUI discovers nodes) ────────────────── + NODE_CLASS_MAPPINGS = { + LoraLoaderRemoteLM.NAME: LoraLoaderRemoteLM, + LoraTextLoaderRemoteLM.NAME: LoraTextLoaderRemoteLM, + LoraStackerRemoteLM.NAME: LoraStackerRemoteLM, + LoraRandomizerRemoteLM.NAME: LoraRandomizerRemoteLM, + LoraCyclerRemoteLM.NAME: LoraCyclerRemoteLM, + LoraPoolRemoteLM.NAME: LoraPoolRemoteLM, + SaveImageRemoteLM.NAME: SaveImageRemoteLM, + WanVideoLoraSelectRemoteLM.NAME: WanVideoLoraSelectRemoteLM, + WanVideoLoraTextSelectRemoteLM.NAME: WanVideoLoraTextSelectRemoteLM, + } -# ── WEB_DIRECTORY tells ComfyUI where to find our JS extensions ─────── -WEB_DIRECTORY = "./web/comfyui" + # ── WEB_DIRECTORY tells ComfyUI where to find our JS extensions ─────── + WEB_DIRECTORY = "./web/comfyui" -# ── Register proxy middleware ────────────────────────────────────────── -try: - from server import PromptServer # type: ignore - from .proxy import register_proxy + # ── Register proxy middleware ────────────────────────────────────────── + try: + from server import PromptServer # type: ignore + from .proxy import register_proxy - register_proxy(PromptServer.instance.app) -except Exception as exc: - logger.warning("[LM-Remote] Could not register proxy middleware: %s", exc) + register_proxy(PromptServer.instance.app) + except Exception as exc: + logger.warning("[LM-Remote] Could not register proxy middleware: %s", exc) -__all__ = ["NODE_CLASS_MAPPINGS", "WEB_DIRECTORY"] + __all__ = ["NODE_CLASS_MAPPINGS", "WEB_DIRECTORY"] diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..2b301ad --- /dev/null +++ b/conftest.py @@ -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"] diff --git a/stats_db.py b/stats_db.py new file mode 100644 index 0000000..1f2d66a --- /dev/null +++ b/stats_db.py @@ -0,0 +1,135 @@ +"""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 diff --git a/tests/__init__.py b/tests/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/test_stats_db.py b/tests/test_stats_db.py new file mode 100644 index 0000000..edef981 --- /dev/null +++ b/tests/test_stats_db.py @@ -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