combuddy 0.1.0__py3-none-any.whl

combuddy/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

combuddy/__main__.py

@@ -0,0 +1,25 @@
+import os, sys, threading, webbrowser
+from .api import create_app
+
+def default_db_path() -> str:
+    d = os.path.join(os.path.expanduser("~"), ".combuddy")
+    os.makedirs(d, exist_ok=True)
+    return os.path.join(d, "combuddy.sqlite")
+
+def _static_dir() -> str:
+    return os.path.join(os.path.dirname(__file__), "web")
+
+def build_app(db_path: str | None = None, static_dir: str | None = None):
+    return create_app(db_path or default_db_path(),
+                      static_dir if static_dir is not None else _static_dir())
+
+def main(argv: list[str] | None = None) -> int:
+    import uvicorn
+    host, port = "127.0.0.1", 8511
+    app = build_app()
+    threading.Timer(1.0, lambda: webbrowser.open(f"http://{host}:{port}")).start()
+    uvicorn.run(app, host=host, port=port)
+    return 0
+
+if __name__ == "__main__":
+    sys.exit(main())

combuddy/api.py

@@ -0,0 +1,129 @@
+import os, re, threading
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
+from . import db as dbm, config, stats, scan_service, queries, trash
+
+def _sniff_image(path: str) -> str:
+    with open(path, "rb") as f:
+        head = f.read(12)
+    if head[:3] == b"\xff\xd8\xff":
+        return "image/jpeg"
+    if head[:8] == b"\x89PNG\r\n\x1a\n":
+        return "image/png"
+    if head[:4] == b"RIFF" and head[8:12] == b"WEBP":
+        return "image/webp"
+    return "application/octet-stream"
+
+def create_app(db_path: str, static_dir: str | None = None) -> FastAPI:
+    app = FastAPI(title="combuddy")
+
+    def conn():
+        c = dbm.connect(db_path); dbm.init_schema(c); return c
+
+    @app.get("/api/stats")
+    def get_stats():
+        c = conn()
+        s = stats.get_stats(c)
+        s["scanning"] = scan_service.STATUS["running"]
+        s["scan"] = dict(scan_service.STATUS)
+        c.close()
+        return s
+
+    @app.get("/api/roots")
+    def get_roots():
+        c = conn()
+        rows = [dict(r) for r in config.get_roots(c)]
+        c.close()
+        return {"roots": rows}
+
+    @app.post("/api/roots")
+    def post_roots(body: dict):
+        c = conn(); config.set_roots(c, body.get("roots", [])); c.close()
+        return {"ok": True}
+
+    @app.get("/api/unreferenced")
+    def get_unreferenced():
+        c = conn(); rows = stats.get_unreferenced(c); c.close()
+        return {"models": rows}
+
+    @app.post("/api/scan")
+    def post_scan():
+        if scan_service.STATUS["running"]:
+            return JSONResponse({"started": False, "reason": "already running"}, status_code=409)
+        def _bg():
+            c = conn()
+            try:
+                scan_service.run_scan(c)
+            finally:
+                c.close()
+        threading.Thread(target=_bg, daemon=True).start()
+        return {"started": True}
+
+    @app.post("/api/scan/cancel")
+    def post_scan_cancel():
+        scan_service.STATUS["cancel"] = True
+        return {"ok": True}
+
+    @app.get("/api/settings")
+    def get_settings():
+        c = conn(); s = config.get_settings(c); c.close()
+        return s
+
+    @app.post("/api/settings")
+    def post_settings(body: dict):
+        c = conn(); config.set_settings(c, body or {}); s = config.get_settings(c); c.close()
+        return s
+
+    @app.get("/api/models")
+    def api_models(search: str = "", type: str = "", flag: str = ""):
+        c = conn(); rows = queries.list_models(c, search, type, flag); c.close()
+        return {"models": rows}
+
+    @app.get("/api/models/{model_id}")
+    def api_model(model_id: int):
+        c = conn(); d = queries.get_model_detail(c, model_id); c.close()
+        if d is None:
+            return JSONResponse({"error": "not found"}, status_code=404)
+        return d
+
+    @app.get("/api/workflows")
+    def api_workflows():
+        c = conn(); rows = queries.list_workflows(c); c.close()
+        return {"workflows": rows}
+
+    @app.get("/api/workflows/{workflow_id}")
+    def api_workflow(workflow_id: int):
+        c = conn(); d = queries.get_workflow_resolution(c, workflow_id); c.close()
+        if d is None:
+            return JSONResponse({"error": "not found"}, status_code=404)
+        return d
+
+    @app.post("/api/cleanup/trash")
+    def api_trash(body: dict):
+        c = conn(); res = trash.move_to_trash(c, body.get("model_ids", [])); c.close()
+        return res
+
+    @app.get("/api/cleanup/trash")
+    def api_list_trash():
+        c = conn(); rows = trash.list_trash(c); c.close()
+        return {"trash": rows}
+
+    @app.post("/api/cleanup/restore")
+    def api_restore(body: dict):
+        c = conn(); res = trash.restore(c, body.get("trash_ids", [])); c.close()
+        return res
+
+    @app.get("/api/preview/{sha256}")
+    def api_preview(sha256: str, hd: int = 0):
+        if not re.fullmatch(r"[0-9a-fA-F]{64}", sha256):
+            return JSONResponse({"error": "bad hash"}, status_code=404)
+        name = sha256 + ("_hd.jpg" if hd else ".jpg")
+        path = os.path.join(os.path.dirname(db_path), "previews", name)
+        if not os.path.isfile(path):
+            return JSONResponse({"error": "not found"}, status_code=404)
+        return FileResponse(path, media_type=_sniff_image(path))
+
+    if static_dir and os.path.isdir(static_dir):
+        app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")
+    return app

combuddy/civitai.py

@@ -0,0 +1,111 @@
+import json, os, time, urllib.request, urllib.error
+
+_API = "https://civitai.com/api/v1/model-versions/by-hash/"
+_UA = "combuddy (model identity lookup)"
+_TIMEOUT = 10
+_DELAY = 0.3          # 礼貌延迟(秒),顺序请求
+
+def _pick_preview(imgs: list) -> dict:
+    """预览优先图片;没有图片再退视频;都没有 type 标记则取第一张。"""
+    for want in ("image", "video"):
+        for im in imgs:
+            if im.get("type") == want:
+                return im
+    return imgs[0] if imgs else {}
+
+def parse_version(data: dict) -> dict:
+    model = data.get("model") or {}
+    img = _pick_preview(data.get("images") or [])
+    return {
+        "name": model.get("name"),
+        "version_name": data.get("name"),
+        "base_model": data.get("baseModel"),
+        "model_type": model.get("type"),
+        "trigger_words": json.dumps(data.get("trainedWords") or []),
+        "nsfw_level": img.get("nsfwLevel"),
+        "civitai_url": f"https://civitai.com/models/{data.get('modelId')}?modelVersionId={data.get('id')}",
+        "image_url": img.get("url"),
+    }
+
+def fetch_by_hash(sha256: str):
+    req = urllib.request.Request(_API + sha256, headers={"User-Agent": _UA})
+    try:
+        with urllib.request.urlopen(req, timeout=_TIMEOUT) as r:
+            data = json.loads(r.read().decode())
+        return ("found", parse_version(data))
+    except urllib.error.HTTPError as e:
+        return ("notfound", None) if e.code == 404 else ("skip", None)
+    except Exception:
+        return ("skip", None)
+
+def _sized(url: str, width: int) -> str:
+    """把 Civitai 图片 URL 的变换段改成 anim=false,width=N(视频取静态帧 + 限宽)。"""
+    parts = url.rsplit("/", 2)
+    if len(parts) == 3 and "civitai" in parts[0]:
+        return f"{parts[0]}/anim=false,width={width}/{parts[2]}"
+    return url
+
+def download_image(url: str, dest: str) -> bool:
+    try:
+        req = urllib.request.Request(url, headers={"User-Agent": _UA})
+        with urllib.request.urlopen(req, timeout=_TIMEOUT) as r:
+            data = r.read()
+        with open(dest, "wb") as f:
+            f.write(data)
+        return True
+    except Exception:
+        return False
+
+def _previews_dir(conn) -> str:
+    for r in conn.execute("PRAGMA database_list"):
+        if r["name"] == "main" and r["file"]:
+            return os.path.join(os.path.dirname(r["file"]), "previews")
+    return "previews"
+
+def _upsert(conn, model_id, sha256, found, ident=None, image_path=None):
+    ident = ident or {}
+    conn.execute(
+        """INSERT INTO civitai(model_id,sha256,found,name,version_name,base_model,model_type,
+             trigger_words,nsfw_level,civitai_url,image_path,checked_at)
+           VALUES(?,?,?,?,?,?,?,?,?,?,?,?)
+           ON CONFLICT(model_id) DO UPDATE SET sha256=excluded.sha256, found=excluded.found,
+             name=excluded.name, version_name=excluded.version_name, base_model=excluded.base_model,
+             model_type=excluded.model_type, trigger_words=excluded.trigger_words,
+             nsfw_level=excluded.nsfw_level, civitai_url=excluded.civitai_url,
+             image_path=excluded.image_path, checked_at=excluded.checked_at""",
+        (model_id, sha256, found, ident.get("name"), ident.get("version_name"),
+         ident.get("base_model"), ident.get("model_type"), ident.get("trigger_words"),
+         ident.get("nsfw_level"), ident.get("civitai_url"), image_path, time.time()))
+
+def enrich_models(conn, progress=None, should_cancel=None, fetch=fetch_by_hash, download=download_image) -> dict:
+    pv = _previews_dir(conn); os.makedirs(pv, exist_ok=True)
+    rows = conn.execute(
+        """SELECT m.id, m.sha256 FROM models m LEFT JOIN civitai c ON c.model_id = m.id
+           WHERE m.sha256 IS NOT NULL AND (c.model_id IS NULL OR c.sha256 != m.sha256)
+           ORDER BY m.id""").fetchall()
+    total = len(rows); done = found = 0
+    if progress:
+        progress(0, total)
+    for r in rows:
+        if should_cancel and should_cancel():
+            break
+        kind, ident = fetch(r["sha256"])
+        if kind == "found":
+            image_path = None
+            url = ident.get("image_url")
+            if url:
+                sha = r["sha256"]
+                if download(_sized(url, 256), os.path.join(pv, sha + ".jpg")):
+                    image_path = sha + ".jpg"
+                    download(_sized(url, 1024), os.path.join(pv, sha + "_hd.jpg"))  # HD 尽力而为
+            _upsert(conn, r["id"], r["sha256"], 1, ident, image_path)
+            conn.commit(); found += 1
+        elif kind == "notfound":
+            _upsert(conn, r["id"], r["sha256"], 0)
+            conn.commit()
+        # kind == "skip": 不写,下次重试
+        done += 1
+        if progress:
+            progress(done, total)
+        time.sleep(_DELAY)
+    return {"found": found, "checked": done, "total": total}

combuddy/config.py

@@ -0,0 +1,68 @@
+import os, sqlite3
+
+# 只 stat 这些已知精确路径的相对形态;绝不递归遍历。
+def _classify(path: str) -> str | None:
+    base = os.path.basename(path.rstrip("/"))
+    if base == "workflows":
+        return "workflow"
+    if base == "models":
+        return "model"
+    # models 目录也可能直接被指认;有典型子目录则认作 model root
+    if os.path.isdir(os.path.join(path, "checkpoints")) or os.path.isdir(os.path.join(path, "loras")):
+        return "model"
+    return None
+
+def detect_candidates(explicit: list[str] | None = None) -> list[dict]:
+    out: list[dict] = []
+    for p in (explicit or []):
+        rp = os.path.realpath(p)
+        if not os.path.isdir(rp):
+            continue
+        kind = _classify(rp)
+        if kind is None:
+            continue
+        out.append({"kind": kind, "path": rp, "label": os.path.basename(rp) or rp, "source": "manual"})
+    return out
+
+def set_roots(conn: sqlite3.Connection, roots: list[dict]) -> None:
+    for r in roots:
+        conn.execute(
+            "INSERT OR IGNORE INTO roots(kind,path,label,source,enabled) VALUES(?,?,?,?,1)",
+            (r["kind"], os.path.realpath(r["path"]), r.get("label"), r.get("source", "manual")),
+        )
+    conn.commit()
+
+def get_roots(conn: sqlite3.Connection, kind: str | None = None) -> list[sqlite3.Row]:
+    if kind:
+        return conn.execute("SELECT * FROM roots WHERE kind=? AND enabled=1", (kind,)).fetchall()
+    return conn.execute("SELECT * FROM roots WHERE enabled=1").fetchall()
+
+def get_settings(conn: sqlite3.Connection) -> dict:
+    have = {r["key"]: r["value"] for r in conn.execute(
+        """SELECT key,value FROM meta WHERE key IN
+           ('auto_hash','hash_workers','hash_max_mbps','online_enrich','nsfw_blur_threshold')""")}
+    return {
+        "auto_hash": have.get("auto_hash", "1") == "1",
+        "hash_workers": int(have.get("hash_workers", "1")),
+        "hash_max_mbps": int(have.get("hash_max_mbps", "0")),
+        "online_enrich": have.get("online_enrich", "1") == "1",
+        "nsfw_blur_threshold": int(have.get("nsfw_blur_threshold", "1")),
+    }
+
+def set_settings(conn: sqlite3.Connection, values: dict) -> None:
+    if "auto_hash" in values:
+        conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('auto_hash',?)",
+                     ("1" if values["auto_hash"] else "0",))
+    if "hash_workers" in values:
+        w = max(1, min(int(values["hash_workers"]), 8))
+        conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('hash_workers',?)", (str(w),))
+    if "hash_max_mbps" in values:
+        m = max(0, int(values["hash_max_mbps"]))
+        conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('hash_max_mbps',?)", (str(m),))
+    if "online_enrich" in values:
+        conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('online_enrich',?)",
+                     ("1" if values["online_enrich"] else "0",))
+    if "nsfw_blur_threshold" in values:
+        t = max(0, min(int(values["nsfw_blur_threshold"]), 32))
+        conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('nsfw_blur_threshold',?)", (str(t),))
+    conn.commit()

combuddy/db.py

@@ -0,0 +1,111 @@
+import sqlite3
+
+SCHEMA_VERSION = 3
+
+def connect(db_path: str) -> sqlite3.Connection:
+    conn = sqlite3.connect(db_path, check_same_thread=False)
+    conn.row_factory = sqlite3.Row
+    conn.execute("PRAGMA journal_mode=WAL")
+    conn.execute("PRAGMA foreign_keys=ON")
+    return conn
+
+def init_schema(conn: sqlite3.Connection) -> None:
+    conn.executescript(
+        """
+        CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT);
+
+        CREATE TABLE IF NOT EXISTS roots (
+            id INTEGER PRIMARY KEY,
+            kind TEXT NOT NULL,           -- 'model' | 'workflow'
+            path TEXT NOT NULL UNIQUE,
+            label TEXT,
+            source TEXT,
+            enabled INTEGER NOT NULL DEFAULT 1
+        );
+
+        CREATE TABLE IF NOT EXISTS models (
+            id INTEGER PRIMARY KEY,
+            root_id INTEGER NOT NULL REFERENCES roots(id),
+            path TEXT NOT NULL UNIQUE,
+            rel_path TEXT NOT NULL,
+            dir_type TEXT NOT NULL,
+            rel_in_type TEXT NOT NULL,
+            filename TEXT NOT NULL,
+            ext TEXT NOT NULL,
+            size INTEGER NOT NULL,
+            mtime REAL NOT NULL,
+            base_arch TEXT,
+            base_source TEXT,
+            header_meta TEXT,             -- reserved for future enrichment, unpopulated in v1 (same as sha256)
+            sha256 TEXT,
+            precision TEXT,
+            param_count INTEGER,
+            display_name TEXT,
+            match_key TEXT NOT NULL,
+            name_key TEXT NOT NULL,
+            first_seen REAL NOT NULL,
+            last_scanned REAL NOT NULL
+        );
+        CREATE INDEX IF NOT EXISTS ix_models_match ON models(dir_type, match_key);
+        CREATE INDEX IF NOT EXISTS ix_models_name ON models(name_key);
+
+        CREATE TABLE IF NOT EXISTS workflows (
+            id INTEGER PRIMARY KEY,
+            root_id INTEGER NOT NULL REFERENCES roots(id),
+            path TEXT NOT NULL UNIQUE,
+            filename TEXT NOT NULL,
+            mtime REAL NOT NULL,
+            ref_count INTEGER NOT NULL DEFAULT 0,
+            parse_error TEXT,
+            last_scanned REAL NOT NULL
+        );
+
+        CREATE TABLE IF NOT EXISTS edges (
+            id INTEGER PRIMARY KEY,
+            workflow_id INTEGER NOT NULL REFERENCES workflows(id),
+            ref_string TEXT NOT NULL,
+            ref_key TEXT NOT NULL,
+            ref_dir_type TEXT,
+            node_type TEXT,
+            model_id INTEGER REFERENCES models(id),
+            match_kind TEXT,
+            UNIQUE(workflow_id, ref_string, node_type)
+        );
+        CREATE INDEX IF NOT EXISTS ix_edges_model ON edges(model_id);
+
+        CREATE TABLE IF NOT EXISTS trash (
+            id INTEGER PRIMARY KEY,
+            model_path TEXT NOT NULL,
+            rel_path TEXT,
+            dir_type TEXT,
+            size INTEGER,
+            trashed_at REAL NOT NULL,
+            trash_path TEXT NOT NULL
+        );
+
+        CREATE TABLE IF NOT EXISTS civitai (
+            model_id     INTEGER PRIMARY KEY REFERENCES models(id) ON DELETE CASCADE,
+            sha256       TEXT NOT NULL,
+            found        INTEGER NOT NULL,
+            name         TEXT,
+            version_name TEXT,
+            base_model   TEXT,
+            model_type   TEXT,
+            trigger_words TEXT,
+            nsfw_level   INTEGER,
+            civitai_url  TEXT,
+            image_path   TEXT,
+            checked_at   REAL NOT NULL
+        );
+        """
+    )
+    # migrate existing (pre-v2) dbs: add columns the CREATE-IF-NOT-EXISTS above can't add
+    have = {r["name"] for r in conn.execute("PRAGMA table_info(models)")}
+    for col, decl in (("precision", "TEXT"), ("param_count", "INTEGER"), ("display_name", "TEXT")):
+        if col not in have:
+            conn.execute(f"ALTER TABLE models ADD COLUMN {col} {decl}")
+    conn.execute(
+        "INSERT OR REPLACE INTO meta(key,value) VALUES('schema_version',?)",
+        (str(SCHEMA_VERSION),),
+    )
+    conn.commit()

combuddy/hashes.py

@@ -0,0 +1,102 @@
+import hashlib, threading, time
+
+_CHUNK = 8 << 20     # 8 MiB，与实测吞吐一致
+_MAX_WORKERS = 8
+_CANCELLED = object()
+
+class _Throttle:
+    """全局令牌桶：把累计已读字节配速到 max_mbps。max_mbps<=0 表示不限速。"""
+    def __init__(self, max_mbps: int):
+        self.bps = int(max_mbps) * 1_000_000 if max_mbps and max_mbps > 0 else 0
+        self.lock = threading.Lock()
+        self.start = time.monotonic()
+        self.bytes = 0
+
+    def account(self, n: int) -> None:
+        if not self.bps:
+            return
+        with self.lock:
+            self.bytes += n
+            target = self.start + self.bytes / self.bps
+            wait = target - time.monotonic()
+        if wait > 0:
+            time.sleep(wait)
+
+def sha256_file(path, throttle=None, should_cancel=None):
+    """流式 sha256。返回摘要 / None(读失败)/ _CANCELLED(仅当传入 should_cancel 且中途取消)。"""
+    h = hashlib.sha256()
+    try:
+        with open(path, "rb", buffering=0) as f:
+            while True:
+                if should_cancel and should_cancel():
+                    return _CANCELLED
+                b = f.read(_CHUNK)
+                if not b:
+                    break
+                h.update(b)
+                if throttle:
+                    throttle.account(len(b))
+    except OSError:
+        return None
+    return h.hexdigest()
+
+def compute_hashes(conn, workers=1, max_mbps=0, progress=None, should_cancel=None) -> dict:
+    rows = conn.execute(
+        "SELECT id, path FROM models WHERE sha256 IS NULL ORDER BY size").fetchall()
+    total = len(rows)
+    if progress:
+        progress(0, total)
+    workers = max(1, min(int(workers or 1), _MAX_WORKERS))
+    throttle = _Throttle(max_mbps)
+    write_lock = threading.Lock()
+    state_lock = threading.Lock()
+    idx_lock = threading.Lock()
+    state = {"done": 0, "hashed": 0, "errors": 0, "cancelled": False}
+    idx = {"i": 0}
+
+    def next_row():
+        with idx_lock:
+            i = idx["i"]
+            if i >= total:
+                return None
+            idx["i"] = i + 1
+            return rows[i]
+
+    def worker():
+        while True:
+            if should_cancel and should_cancel():
+                with state_lock:
+                    state["cancelled"] = True
+                return
+            r = next_row()
+            if r is None:
+                return
+            res = sha256_file(r["path"], throttle, should_cancel)
+            if res is _CANCELLED:
+                with state_lock:
+                    state["cancelled"] = True
+                return
+            with write_lock:
+                if res is not None:
+                    conn.execute("UPDATE models SET sha256=? WHERE id=?", (res, r["id"]))
+                    conn.commit()               # 逐文件提交 → 断点续跑
+            with state_lock:
+                state["done"] += 1
+                if res is None:
+                    state["errors"] += 1
+                else:
+                    state["hashed"] += 1
+                d = state["done"]
+            if progress:
+                progress(d, total)
+
+    if workers == 1:
+        worker()
+    else:
+        threads = [threading.Thread(target=worker) for _ in range(workers)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+    return {"hashed": state["hashed"], "errors": state["errors"],
+            "total": total, "cancelled": state["cancelled"]}