combuddy 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (54) hide show
  1. combuddy-0.1.0/LICENSE +21 -0
  2. combuddy-0.1.0/PKG-INFO +89 -0
  3. combuddy-0.1.0/README.md +59 -0
  4. combuddy-0.1.0/combuddy/__init__.py +1 -0
  5. combuddy-0.1.0/combuddy/__main__.py +25 -0
  6. combuddy-0.1.0/combuddy/api.py +129 -0
  7. combuddy-0.1.0/combuddy/civitai.py +111 -0
  8. combuddy-0.1.0/combuddy/config.py +68 -0
  9. combuddy-0.1.0/combuddy/db.py +111 -0
  10. combuddy-0.1.0/combuddy/hashes.py +102 -0
  11. combuddy-0.1.0/combuddy/headers.py +209 -0
  12. combuddy-0.1.0/combuddy/norm.py +11 -0
  13. combuddy-0.1.0/combuddy/queries.py +61 -0
  14. combuddy-0.1.0/combuddy/resolver.py +50 -0
  15. combuddy-0.1.0/combuddy/roles.py +23 -0
  16. combuddy-0.1.0/combuddy/scan_service.py +70 -0
  17. combuddy-0.1.0/combuddy/scanner.py +70 -0
  18. combuddy-0.1.0/combuddy/stats.py +29 -0
  19. combuddy-0.1.0/combuddy/trash.py +58 -0
  20. combuddy-0.1.0/combuddy/web/assets/Inter-italic.var-DhD-tpjY.woff2 +0 -0
  21. combuddy-0.1.0/combuddy/web/assets/Inter-roman.var-C-r5W2Hj.woff2 +0 -0
  22. combuddy-0.1.0/combuddy/web/assets/index-DDQng-X8.js +17 -0
  23. combuddy-0.1.0/combuddy/web/assets/index-DlRCX3ps.css +1 -0
  24. combuddy-0.1.0/combuddy/web/index.html +5 -0
  25. combuddy-0.1.0/combuddy/workflows.py +39 -0
  26. combuddy-0.1.0/combuddy.egg-info/PKG-INFO +89 -0
  27. combuddy-0.1.0/combuddy.egg-info/SOURCES.txt +52 -0
  28. combuddy-0.1.0/combuddy.egg-info/dependency_links.txt +1 -0
  29. combuddy-0.1.0/combuddy.egg-info/entry_points.txt +2 -0
  30. combuddy-0.1.0/combuddy.egg-info/requires.txt +6 -0
  31. combuddy-0.1.0/combuddy.egg-info/top_level.txt +1 -0
  32. combuddy-0.1.0/pyproject.toml +44 -0
  33. combuddy-0.1.0/setup.cfg +4 -0
  34. combuddy-0.1.0/tests/test_api.py +24 -0
  35. combuddy-0.1.0/tests/test_api_v11.py +62 -0
  36. combuddy-0.1.0/tests/test_civitai.py +126 -0
  37. combuddy-0.1.0/tests/test_config.py +47 -0
  38. combuddy-0.1.0/tests/test_db.py +11 -0
  39. combuddy-0.1.0/tests/test_db_migrate.py +57 -0
  40. combuddy-0.1.0/tests/test_enrich.py +71 -0
  41. combuddy-0.1.0/tests/test_extract_facts.py +38 -0
  42. combuddy-0.1.0/tests/test_hashes.py +85 -0
  43. combuddy-0.1.0/tests/test_headers.py +43 -0
  44. combuddy-0.1.0/tests/test_headers_gguf.py +28 -0
  45. combuddy-0.1.0/tests/test_main.py +8 -0
  46. combuddy-0.1.0/tests/test_norm.py +7 -0
  47. combuddy-0.1.0/tests/test_queries.py +77 -0
  48. combuddy-0.1.0/tests/test_resolver.py +47 -0
  49. combuddy-0.1.0/tests/test_roles.py +13 -0
  50. combuddy-0.1.0/tests/test_scan_service.py +116 -0
  51. combuddy-0.1.0/tests/test_scanner.py +78 -0
  52. combuddy-0.1.0/tests/test_stats.py +54 -0
  53. combuddy-0.1.0/tests/test_trash.py +72 -0
  54. combuddy-0.1.0/tests/test_workflows.py +43 -0
combuddy-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 peilinok
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,89 @@
1
+ Metadata-Version: 2.4
2
+ Name: combuddy
3
+ Version: 0.1.0
4
+ Summary: A local model & workflow dependency manager for ComfyUI
5
+ Author-email: peilinok <peilinok@gmail.com>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/peilinok/combuddy
8
+ Project-URL: Repository, https://github.com/peilinok/combuddy
9
+ Project-URL: Issues, https://github.com/peilinok/combuddy/issues
10
+ Keywords: comfyui,stable-diffusion,models,workflows,dependency-management
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: End Users/Desktop
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3.11
15
+ Classifier: Programming Language :: Python :: 3.12
16
+ Classifier: Programming Language :: Python :: 3.13
17
+ Classifier: Operating System :: OS Independent
18
+ Classifier: Environment :: Web Environment
19
+ Classifier: Framework :: FastAPI
20
+ Classifier: Topic :: Multimedia :: Graphics
21
+ Requires-Python: >=3.11
22
+ Description-Content-Type: text/markdown
23
+ License-File: LICENSE
24
+ Requires-Dist: fastapi>=0.110
25
+ Requires-Dist: uvicorn[standard]>=0.29
26
+ Provides-Extra: dev
27
+ Requires-Dist: pytest>=8; extra == "dev"
28
+ Requires-Dist: httpx>=0.27; extra == "dev"
29
+ Dynamic: license-file
30
+
31
+ # combuddy
32
+
33
+ [![CI](https://github.com/peilinok/combuddy/actions/workflows/ci.yml/badge.svg)](https://github.com/peilinok/combuddy/actions/workflows/ci.yml)
34
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
35
+
36
+ A model & workflow dependency manager for [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Point it at your local model library and workflow files, and it shows you how they depend on each other — which workflows use which models, which models nothing uses, and which models a workflow is missing. **Local-first**: model identity (base architecture, precision, parameters) is computed from file headers; optional Civitai enrichment by content hash is on by default and toggleable in settings.
37
+
38
+ ## Why
39
+
40
+ Community workflows reference models by bare filename — no download link, no hash. Over time you lose track of what a model is, which workflows need it, and what's safe to delete. combuddy answers those from what's already on your disk:
41
+
42
+ - **Dashboard** — total models & size, base-architecture coverage, and how many models nothing references.
43
+ - **Model library** — search/filter every model; click one to see its details and the **workflows that reference it** (reverse dependencies).
44
+ - **Workflow resolution** — pick a workflow, see each referenced model marked **hit / ambiguous / missing**.
45
+ - **Cleanup** — the models no workflow uses, with reclaimable space, moved to a **recoverable trash** (never a hard delete; only 0-reference models can go).
46
+ - **Offline identity** — base architecture (SD1.5/SDXL/Flux/…, incl. GGUF), role labels (text encoder, VAE, ControlNet, …), and precision, all read straight from file headers.
47
+ - **Civitai identity** — for models found on Civitai by content hash: real name, base model, trigger words, and a cached preview thumbnail with **HD zoom**. Enrichment sends **only the hash**, is on by default, and is toggleable in settings.
48
+
49
+ ## Install & run
50
+
51
+ Requires Python 3.11+.
52
+
53
+ ```bash
54
+ pip install -e .
55
+ combuddy
56
+ ```
57
+
58
+ This starts a local server on `http://127.0.0.1:8511` and opens your browser. On first run, point it at your ComfyUI **models** directory and **workflows** directory (e.g. `.../user/default/workflows`); it scans and populates the Dashboard in seconds.
59
+
60
+ ## How it works
61
+
62
+ - Scans model directories (skipping noise and its own trash), and parses ComfyUI workflow JSON for model references.
63
+ - Matches references to local files by **directory type + normalized relative path** (handling subfolders, case, Unicode, and backslashes) — not by fragile basename guessing.
64
+ - Computes content hashes for identity; optional Civitai enrichment sends only the hash to look up real names and metadata.
65
+ - Stores everything in a single local SQLite index; the UI reads it live.
66
+
67
+ ## Tech stack
68
+
69
+ - **Backend:** Python 3.11+, FastAPI, standard-library `sqlite3` (no ORM). One command (`combuddy`) runs Uvicorn and serves both the API and the UI.
70
+ - **Frontend:** Vue 3, Vite, PrimeVue, Tailwind — built into `combuddy/web/` and served by the backend, so the whole app installs as one package.
71
+
72
+ ## Development
73
+
74
+ ```bash
75
+ pip install -e ".[dev]" # backend + test deps
76
+ pytest -q # backend tests
77
+
78
+ cd frontend
79
+ npm install
80
+ npm test # frontend tests (Vitest)
81
+ npm run dev # dev server, proxies /api to :8511
82
+ npm run build # rebuild the bundle into ../combuddy/web
83
+ ```
84
+
85
+ After changing anything in `frontend/src`, run `npm run build` — the packaged app serves the committed `combuddy/web/` bundle.
86
+
87
+ ## Status & roadmap
88
+
89
+ v1.2 — local-first with optional Civitai enrichment (real names, preview images, trigger words by content hash, default on and toggleable). Planned next: download center for missing models, and dependency pinning for shareable, self-healing workflows.
@@ -0,0 +1,59 @@
1
+ # combuddy
2
+
3
+ [![CI](https://github.com/peilinok/combuddy/actions/workflows/ci.yml/badge.svg)](https://github.com/peilinok/combuddy/actions/workflows/ci.yml)
4
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
5
+
6
+ A model & workflow dependency manager for [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Point it at your local model library and workflow files, and it shows you how they depend on each other — which workflows use which models, which models nothing uses, and which models a workflow is missing. **Local-first**: model identity (base architecture, precision, parameters) is computed from file headers; optional Civitai enrichment by content hash is on by default and toggleable in settings.
7
+
8
+ ## Why
9
+
10
+ Community workflows reference models by bare filename — no download link, no hash. Over time you lose track of what a model is, which workflows need it, and what's safe to delete. combuddy answers those from what's already on your disk:
11
+
12
+ - **Dashboard** — total models & size, base-architecture coverage, and how many models nothing references.
13
+ - **Model library** — search/filter every model; click one to see its details and the **workflows that reference it** (reverse dependencies).
14
+ - **Workflow resolution** — pick a workflow, see each referenced model marked **hit / ambiguous / missing**.
15
+ - **Cleanup** — the models no workflow uses, with reclaimable space, moved to a **recoverable trash** (never a hard delete; only 0-reference models can go).
16
+ - **Offline identity** — base architecture (SD1.5/SDXL/Flux/…, incl. GGUF), role labels (text encoder, VAE, ControlNet, …), and precision, all read straight from file headers.
17
+ - **Civitai identity** — for models found on Civitai by content hash: real name, base model, trigger words, and a cached preview thumbnail with **HD zoom**. Enrichment sends **only the hash**, is on by default, and is toggleable in settings.
18
+
19
+ ## Install & run
20
+
21
+ Requires Python 3.11+.
22
+
23
+ ```bash
24
+ pip install -e .
25
+ combuddy
26
+ ```
27
+
28
+ This starts a local server on `http://127.0.0.1:8511` and opens your browser. On first run, point it at your ComfyUI **models** directory and **workflows** directory (e.g. `.../user/default/workflows`); it scans and populates the Dashboard in seconds.
29
+
30
+ ## How it works
31
+
32
+ - Scans model directories (skipping noise and its own trash), and parses ComfyUI workflow JSON for model references.
33
+ - Matches references to local files by **directory type + normalized relative path** (handling subfolders, case, Unicode, and backslashes) — not by fragile basename guessing.
34
+ - Computes content hashes for identity; optional Civitai enrichment sends only the hash to look up real names and metadata.
35
+ - Stores everything in a single local SQLite index; the UI reads it live.
36
+
37
+ ## Tech stack
38
+
39
+ - **Backend:** Python 3.11+, FastAPI, standard-library `sqlite3` (no ORM). One command (`combuddy`) runs Uvicorn and serves both the API and the UI.
40
+ - **Frontend:** Vue 3, Vite, PrimeVue, Tailwind — built into `combuddy/web/` and served by the backend, so the whole app installs as one package.
41
+
42
+ ## Development
43
+
44
+ ```bash
45
+ pip install -e ".[dev]" # backend + test deps
46
+ pytest -q # backend tests
47
+
48
+ cd frontend
49
+ npm install
50
+ npm test # frontend tests (Vitest)
51
+ npm run dev # dev server, proxies /api to :8511
52
+ npm run build # rebuild the bundle into ../combuddy/web
53
+ ```
54
+
55
+ After changing anything in `frontend/src`, run `npm run build` — the packaged app serves the committed `combuddy/web/` bundle.
56
+
57
+ ## Status & roadmap
58
+
59
+ v1.2 — local-first with optional Civitai enrichment (real names, preview images, trigger words by content hash, default on and toggleable). Planned next: download center for missing models, and dependency pinning for shareable, self-healing workflows.
@@ -0,0 +1 @@
1
+ __version__ = "0.1.0"
@@ -0,0 +1,25 @@
1
+ import os, sys, threading, webbrowser
2
+ from .api import create_app
3
+
4
+ def default_db_path() -> str:
5
+ d = os.path.join(os.path.expanduser("~"), ".combuddy")
6
+ os.makedirs(d, exist_ok=True)
7
+ return os.path.join(d, "combuddy.sqlite")
8
+
9
+ def _static_dir() -> str:
10
+ return os.path.join(os.path.dirname(__file__), "web")
11
+
12
+ def build_app(db_path: str | None = None, static_dir: str | None = None):
13
+ return create_app(db_path or default_db_path(),
14
+ static_dir if static_dir is not None else _static_dir())
15
+
16
+ def main(argv: list[str] | None = None) -> int:
17
+ import uvicorn
18
+ host, port = "127.0.0.1", 8511
19
+ app = build_app()
20
+ threading.Timer(1.0, lambda: webbrowser.open(f"http://{host}:{port}")).start()
21
+ uvicorn.run(app, host=host, port=port)
22
+ return 0
23
+
24
+ if __name__ == "__main__":
25
+ sys.exit(main())
@@ -0,0 +1,129 @@
1
+ import os, re, threading
2
+ from fastapi import FastAPI
3
+ from fastapi.responses import JSONResponse, FileResponse
4
+ from fastapi.staticfiles import StaticFiles
5
+ from . import db as dbm, config, stats, scan_service, queries, trash
6
+
7
+ def _sniff_image(path: str) -> str:
8
+ with open(path, "rb") as f:
9
+ head = f.read(12)
10
+ if head[:3] == b"\xff\xd8\xff":
11
+ return "image/jpeg"
12
+ if head[:8] == b"\x89PNG\r\n\x1a\n":
13
+ return "image/png"
14
+ if head[:4] == b"RIFF" and head[8:12] == b"WEBP":
15
+ return "image/webp"
16
+ return "application/octet-stream"
17
+
18
+ def create_app(db_path: str, static_dir: str | None = None) -> FastAPI:
19
+ app = FastAPI(title="combuddy")
20
+
21
+ def conn():
22
+ c = dbm.connect(db_path); dbm.init_schema(c); return c
23
+
24
+ @app.get("/api/stats")
25
+ def get_stats():
26
+ c = conn()
27
+ s = stats.get_stats(c)
28
+ s["scanning"] = scan_service.STATUS["running"]
29
+ s["scan"] = dict(scan_service.STATUS)
30
+ c.close()
31
+ return s
32
+
33
+ @app.get("/api/roots")
34
+ def get_roots():
35
+ c = conn()
36
+ rows = [dict(r) for r in config.get_roots(c)]
37
+ c.close()
38
+ return {"roots": rows}
39
+
40
+ @app.post("/api/roots")
41
+ def post_roots(body: dict):
42
+ c = conn(); config.set_roots(c, body.get("roots", [])); c.close()
43
+ return {"ok": True}
44
+
45
+ @app.get("/api/unreferenced")
46
+ def get_unreferenced():
47
+ c = conn(); rows = stats.get_unreferenced(c); c.close()
48
+ return {"models": rows}
49
+
50
+ @app.post("/api/scan")
51
+ def post_scan():
52
+ if scan_service.STATUS["running"]:
53
+ return JSONResponse({"started": False, "reason": "already running"}, status_code=409)
54
+ def _bg():
55
+ c = conn()
56
+ try:
57
+ scan_service.run_scan(c)
58
+ finally:
59
+ c.close()
60
+ threading.Thread(target=_bg, daemon=True).start()
61
+ return {"started": True}
62
+
63
+ @app.post("/api/scan/cancel")
64
+ def post_scan_cancel():
65
+ scan_service.STATUS["cancel"] = True
66
+ return {"ok": True}
67
+
68
+ @app.get("/api/settings")
69
+ def get_settings():
70
+ c = conn(); s = config.get_settings(c); c.close()
71
+ return s
72
+
73
+ @app.post("/api/settings")
74
+ def post_settings(body: dict):
75
+ c = conn(); config.set_settings(c, body or {}); s = config.get_settings(c); c.close()
76
+ return s
77
+
78
+ @app.get("/api/models")
79
+ def api_models(search: str = "", type: str = "", flag: str = ""):
80
+ c = conn(); rows = queries.list_models(c, search, type, flag); c.close()
81
+ return {"models": rows}
82
+
83
+ @app.get("/api/models/{model_id}")
84
+ def api_model(model_id: int):
85
+ c = conn(); d = queries.get_model_detail(c, model_id); c.close()
86
+ if d is None:
87
+ return JSONResponse({"error": "not found"}, status_code=404)
88
+ return d
89
+
90
+ @app.get("/api/workflows")
91
+ def api_workflows():
92
+ c = conn(); rows = queries.list_workflows(c); c.close()
93
+ return {"workflows": rows}
94
+
95
+ @app.get("/api/workflows/{workflow_id}")
96
+ def api_workflow(workflow_id: int):
97
+ c = conn(); d = queries.get_workflow_resolution(c, workflow_id); c.close()
98
+ if d is None:
99
+ return JSONResponse({"error": "not found"}, status_code=404)
100
+ return d
101
+
102
+ @app.post("/api/cleanup/trash")
103
+ def api_trash(body: dict):
104
+ c = conn(); res = trash.move_to_trash(c, body.get("model_ids", [])); c.close()
105
+ return res
106
+
107
+ @app.get("/api/cleanup/trash")
108
+ def api_list_trash():
109
+ c = conn(); rows = trash.list_trash(c); c.close()
110
+ return {"trash": rows}
111
+
112
+ @app.post("/api/cleanup/restore")
113
+ def api_restore(body: dict):
114
+ c = conn(); res = trash.restore(c, body.get("trash_ids", [])); c.close()
115
+ return res
116
+
117
+ @app.get("/api/preview/{sha256}")
118
+ def api_preview(sha256: str, hd: int = 0):
119
+ if not re.fullmatch(r"[0-9a-fA-F]{64}", sha256):
120
+ return JSONResponse({"error": "bad hash"}, status_code=404)
121
+ name = sha256 + ("_hd.jpg" if hd else ".jpg")
122
+ path = os.path.join(os.path.dirname(db_path), "previews", name)
123
+ if not os.path.isfile(path):
124
+ return JSONResponse({"error": "not found"}, status_code=404)
125
+ return FileResponse(path, media_type=_sniff_image(path))
126
+
127
+ if static_dir and os.path.isdir(static_dir):
128
+ app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")
129
+ return app
@@ -0,0 +1,111 @@
1
+ import json, os, time, urllib.request, urllib.error
2
+
3
+ _API = "https://civitai.com/api/v1/model-versions/by-hash/"
4
+ _UA = "combuddy (model identity lookup)"
5
+ _TIMEOUT = 10
6
+ _DELAY = 0.3 # 礼貌延迟(秒),顺序请求
7
+
8
+ def _pick_preview(imgs: list) -> dict:
9
+ """预览优先图片;没有图片再退视频;都没有 type 标记则取第一张。"""
10
+ for want in ("image", "video"):
11
+ for im in imgs:
12
+ if im.get("type") == want:
13
+ return im
14
+ return imgs[0] if imgs else {}
15
+
16
+ def parse_version(data: dict) -> dict:
17
+ model = data.get("model") or {}
18
+ img = _pick_preview(data.get("images") or [])
19
+ return {
20
+ "name": model.get("name"),
21
+ "version_name": data.get("name"),
22
+ "base_model": data.get("baseModel"),
23
+ "model_type": model.get("type"),
24
+ "trigger_words": json.dumps(data.get("trainedWords") or []),
25
+ "nsfw_level": img.get("nsfwLevel"),
26
+ "civitai_url": f"https://civitai.com/models/{data.get('modelId')}?modelVersionId={data.get('id')}",
27
+ "image_url": img.get("url"),
28
+ }
29
+
30
+ def fetch_by_hash(sha256: str):
31
+ req = urllib.request.Request(_API + sha256, headers={"User-Agent": _UA})
32
+ try:
33
+ with urllib.request.urlopen(req, timeout=_TIMEOUT) as r:
34
+ data = json.loads(r.read().decode())
35
+ return ("found", parse_version(data))
36
+ except urllib.error.HTTPError as e:
37
+ return ("notfound", None) if e.code == 404 else ("skip", None)
38
+ except Exception:
39
+ return ("skip", None)
40
+
41
+ def _sized(url: str, width: int) -> str:
42
+ """把 Civitai 图片 URL 的变换段改成 anim=false,width=N(视频取静态帧 + 限宽)。"""
43
+ parts = url.rsplit("/", 2)
44
+ if len(parts) == 3 and "civitai" in parts[0]:
45
+ return f"{parts[0]}/anim=false,width={width}/{parts[2]}"
46
+ return url
47
+
48
+ def download_image(url: str, dest: str) -> bool:
49
+ try:
50
+ req = urllib.request.Request(url, headers={"User-Agent": _UA})
51
+ with urllib.request.urlopen(req, timeout=_TIMEOUT) as r:
52
+ data = r.read()
53
+ with open(dest, "wb") as f:
54
+ f.write(data)
55
+ return True
56
+ except Exception:
57
+ return False
58
+
59
+ def _previews_dir(conn) -> str:
60
+ for r in conn.execute("PRAGMA database_list"):
61
+ if r["name"] == "main" and r["file"]:
62
+ return os.path.join(os.path.dirname(r["file"]), "previews")
63
+ return "previews"
64
+
65
+ def _upsert(conn, model_id, sha256, found, ident=None, image_path=None):
66
+ ident = ident or {}
67
+ conn.execute(
68
+ """INSERT INTO civitai(model_id,sha256,found,name,version_name,base_model,model_type,
69
+ trigger_words,nsfw_level,civitai_url,image_path,checked_at)
70
+ VALUES(?,?,?,?,?,?,?,?,?,?,?,?)
71
+ ON CONFLICT(model_id) DO UPDATE SET sha256=excluded.sha256, found=excluded.found,
72
+ name=excluded.name, version_name=excluded.version_name, base_model=excluded.base_model,
73
+ model_type=excluded.model_type, trigger_words=excluded.trigger_words,
74
+ nsfw_level=excluded.nsfw_level, civitai_url=excluded.civitai_url,
75
+ image_path=excluded.image_path, checked_at=excluded.checked_at""",
76
+ (model_id, sha256, found, ident.get("name"), ident.get("version_name"),
77
+ ident.get("base_model"), ident.get("model_type"), ident.get("trigger_words"),
78
+ ident.get("nsfw_level"), ident.get("civitai_url"), image_path, time.time()))
79
+
80
+ def enrich_models(conn, progress=None, should_cancel=None, fetch=fetch_by_hash, download=download_image) -> dict:
81
+ pv = _previews_dir(conn); os.makedirs(pv, exist_ok=True)
82
+ rows = conn.execute(
83
+ """SELECT m.id, m.sha256 FROM models m LEFT JOIN civitai c ON c.model_id = m.id
84
+ WHERE m.sha256 IS NOT NULL AND (c.model_id IS NULL OR c.sha256 != m.sha256)
85
+ ORDER BY m.id""").fetchall()
86
+ total = len(rows); done = found = 0
87
+ if progress:
88
+ progress(0, total)
89
+ for r in rows:
90
+ if should_cancel and should_cancel():
91
+ break
92
+ kind, ident = fetch(r["sha256"])
93
+ if kind == "found":
94
+ image_path = None
95
+ url = ident.get("image_url")
96
+ if url:
97
+ sha = r["sha256"]
98
+ if download(_sized(url, 256), os.path.join(pv, sha + ".jpg")):
99
+ image_path = sha + ".jpg"
100
+ download(_sized(url, 1024), os.path.join(pv, sha + "_hd.jpg")) # HD 尽力而为
101
+ _upsert(conn, r["id"], r["sha256"], 1, ident, image_path)
102
+ conn.commit(); found += 1
103
+ elif kind == "notfound":
104
+ _upsert(conn, r["id"], r["sha256"], 0)
105
+ conn.commit()
106
+ # kind == "skip": 不写,下次重试
107
+ done += 1
108
+ if progress:
109
+ progress(done, total)
110
+ time.sleep(_DELAY)
111
+ return {"found": found, "checked": done, "total": total}
@@ -0,0 +1,68 @@
1
+ import os, sqlite3
2
+
3
+ # 只 stat 这些已知精确路径的相对形态;绝不递归遍历。
4
+ def _classify(path: str) -> str | None:
5
+ base = os.path.basename(path.rstrip("/"))
6
+ if base == "workflows":
7
+ return "workflow"
8
+ if base == "models":
9
+ return "model"
10
+ # models 目录也可能直接被指认;有典型子目录则认作 model root
11
+ if os.path.isdir(os.path.join(path, "checkpoints")) or os.path.isdir(os.path.join(path, "loras")):
12
+ return "model"
13
+ return None
14
+
15
+ def detect_candidates(explicit: list[str] | None = None) -> list[dict]:
16
+ out: list[dict] = []
17
+ for p in (explicit or []):
18
+ rp = os.path.realpath(p)
19
+ if not os.path.isdir(rp):
20
+ continue
21
+ kind = _classify(rp)
22
+ if kind is None:
23
+ continue
24
+ out.append({"kind": kind, "path": rp, "label": os.path.basename(rp) or rp, "source": "manual"})
25
+ return out
26
+
27
+ def set_roots(conn: sqlite3.Connection, roots: list[dict]) -> None:
28
+ for r in roots:
29
+ conn.execute(
30
+ "INSERT OR IGNORE INTO roots(kind,path,label,source,enabled) VALUES(?,?,?,?,1)",
31
+ (r["kind"], os.path.realpath(r["path"]), r.get("label"), r.get("source", "manual")),
32
+ )
33
+ conn.commit()
34
+
35
+ def get_roots(conn: sqlite3.Connection, kind: str | None = None) -> list[sqlite3.Row]:
36
+ if kind:
37
+ return conn.execute("SELECT * FROM roots WHERE kind=? AND enabled=1", (kind,)).fetchall()
38
+ return conn.execute("SELECT * FROM roots WHERE enabled=1").fetchall()
39
+
40
+ def get_settings(conn: sqlite3.Connection) -> dict:
41
+ have = {r["key"]: r["value"] for r in conn.execute(
42
+ """SELECT key,value FROM meta WHERE key IN
43
+ ('auto_hash','hash_workers','hash_max_mbps','online_enrich','nsfw_blur_threshold')""")}
44
+ return {
45
+ "auto_hash": have.get("auto_hash", "1") == "1",
46
+ "hash_workers": int(have.get("hash_workers", "1")),
47
+ "hash_max_mbps": int(have.get("hash_max_mbps", "0")),
48
+ "online_enrich": have.get("online_enrich", "1") == "1",
49
+ "nsfw_blur_threshold": int(have.get("nsfw_blur_threshold", "1")),
50
+ }
51
+
52
+ def set_settings(conn: sqlite3.Connection, values: dict) -> None:
53
+ if "auto_hash" in values:
54
+ conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('auto_hash',?)",
55
+ ("1" if values["auto_hash"] else "0",))
56
+ if "hash_workers" in values:
57
+ w = max(1, min(int(values["hash_workers"]), 8))
58
+ conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('hash_workers',?)", (str(w),))
59
+ if "hash_max_mbps" in values:
60
+ m = max(0, int(values["hash_max_mbps"]))
61
+ conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('hash_max_mbps',?)", (str(m),))
62
+ if "online_enrich" in values:
63
+ conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('online_enrich',?)",
64
+ ("1" if values["online_enrich"] else "0",))
65
+ if "nsfw_blur_threshold" in values:
66
+ t = max(0, min(int(values["nsfw_blur_threshold"]), 32))
67
+ conn.execute("INSERT OR REPLACE INTO meta(key,value) VALUES('nsfw_blur_threshold',?)", (str(t),))
68
+ conn.commit()
@@ -0,0 +1,111 @@
1
+ import sqlite3
2
+
3
+ SCHEMA_VERSION = 3
4
+
5
+ def connect(db_path: str) -> sqlite3.Connection:
6
+ conn = sqlite3.connect(db_path, check_same_thread=False)
7
+ conn.row_factory = sqlite3.Row
8
+ conn.execute("PRAGMA journal_mode=WAL")
9
+ conn.execute("PRAGMA foreign_keys=ON")
10
+ return conn
11
+
12
+ def init_schema(conn: sqlite3.Connection) -> None:
13
+ conn.executescript(
14
+ """
15
+ CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT);
16
+
17
+ CREATE TABLE IF NOT EXISTS roots (
18
+ id INTEGER PRIMARY KEY,
19
+ kind TEXT NOT NULL, -- 'model' | 'workflow'
20
+ path TEXT NOT NULL UNIQUE,
21
+ label TEXT,
22
+ source TEXT,
23
+ enabled INTEGER NOT NULL DEFAULT 1
24
+ );
25
+
26
+ CREATE TABLE IF NOT EXISTS models (
27
+ id INTEGER PRIMARY KEY,
28
+ root_id INTEGER NOT NULL REFERENCES roots(id),
29
+ path TEXT NOT NULL UNIQUE,
30
+ rel_path TEXT NOT NULL,
31
+ dir_type TEXT NOT NULL,
32
+ rel_in_type TEXT NOT NULL,
33
+ filename TEXT NOT NULL,
34
+ ext TEXT NOT NULL,
35
+ size INTEGER NOT NULL,
36
+ mtime REAL NOT NULL,
37
+ base_arch TEXT,
38
+ base_source TEXT,
39
+ header_meta TEXT, -- reserved for future enrichment, unpopulated in v1 (same as sha256)
40
+ sha256 TEXT,
41
+ precision TEXT,
42
+ param_count INTEGER,
43
+ display_name TEXT,
44
+ match_key TEXT NOT NULL,
45
+ name_key TEXT NOT NULL,
46
+ first_seen REAL NOT NULL,
47
+ last_scanned REAL NOT NULL
48
+ );
49
+ CREATE INDEX IF NOT EXISTS ix_models_match ON models(dir_type, match_key);
50
+ CREATE INDEX IF NOT EXISTS ix_models_name ON models(name_key);
51
+
52
+ CREATE TABLE IF NOT EXISTS workflows (
53
+ id INTEGER PRIMARY KEY,
54
+ root_id INTEGER NOT NULL REFERENCES roots(id),
55
+ path TEXT NOT NULL UNIQUE,
56
+ filename TEXT NOT NULL,
57
+ mtime REAL NOT NULL,
58
+ ref_count INTEGER NOT NULL DEFAULT 0,
59
+ parse_error TEXT,
60
+ last_scanned REAL NOT NULL
61
+ );
62
+
63
+ CREATE TABLE IF NOT EXISTS edges (
64
+ id INTEGER PRIMARY KEY,
65
+ workflow_id INTEGER NOT NULL REFERENCES workflows(id),
66
+ ref_string TEXT NOT NULL,
67
+ ref_key TEXT NOT NULL,
68
+ ref_dir_type TEXT,
69
+ node_type TEXT,
70
+ model_id INTEGER REFERENCES models(id),
71
+ match_kind TEXT,
72
+ UNIQUE(workflow_id, ref_string, node_type)
73
+ );
74
+ CREATE INDEX IF NOT EXISTS ix_edges_model ON edges(model_id);
75
+
76
+ CREATE TABLE IF NOT EXISTS trash (
77
+ id INTEGER PRIMARY KEY,
78
+ model_path TEXT NOT NULL,
79
+ rel_path TEXT,
80
+ dir_type TEXT,
81
+ size INTEGER,
82
+ trashed_at REAL NOT NULL,
83
+ trash_path TEXT NOT NULL
84
+ );
85
+
86
+ CREATE TABLE IF NOT EXISTS civitai (
87
+ model_id INTEGER PRIMARY KEY REFERENCES models(id) ON DELETE CASCADE,
88
+ sha256 TEXT NOT NULL,
89
+ found INTEGER NOT NULL,
90
+ name TEXT,
91
+ version_name TEXT,
92
+ base_model TEXT,
93
+ model_type TEXT,
94
+ trigger_words TEXT,
95
+ nsfw_level INTEGER,
96
+ civitai_url TEXT,
97
+ image_path TEXT,
98
+ checked_at REAL NOT NULL
99
+ );
100
+ """
101
+ )
102
+ # migrate existing (pre-v2) dbs: add columns the CREATE-IF-NOT-EXISTS above can't add
103
+ have = {r["name"] for r in conn.execute("PRAGMA table_info(models)")}
104
+ for col, decl in (("precision", "TEXT"), ("param_count", "INTEGER"), ("display_name", "TEXT")):
105
+ if col not in have:
106
+ conn.execute(f"ALTER TABLE models ADD COLUMN {col} {decl}")
107
+ conn.execute(
108
+ "INSERT OR REPLACE INTO meta(key,value) VALUES('schema_version',?)",
109
+ (str(SCHEMA_VERSION),),
110
+ )
111
+ conn.commit()