Enhance CLI and documentation for capturing VM inventory
- Introduced new `capture` commands in the CLI to allow users to merge VM inventory into local manifests without downloading weights. - Updated `README.md` and `cli.md` to include detailed instructions for the new capture functionality, including options for models and extensions. - Enhanced `decisions.md` to clarify the role of captured links in the manifest files. - Improved `extensions.md` to document the process of capturing installed extensions back to the local configuration. - Added new functions in `civitai.py` to support fetching model versions by hash and generating canonical URLs for models.
This commit is contained in:
Leonid Pershin
@@ -0,0 +1,199 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Scan VM Models/ + git Extensions/DLNodes. Stdlib only. Writes JSON inventory."""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
DATA = Path("/mnt/swarm_data")
|
||||
MODELS = DATA / "Models"
|
||||
OUT = Path("/tmp/gpu-rent-inventory.json")
|
||||
|
||||
WEIGHT_SUFFIXES = {".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf", ".sft", ".onnx"}
|
||||
|
||||
# SwarmUI folder name → models.yaml kind
|
||||
FOLDER_TO_KIND = {
|
||||
"Stable-Diffusion": "checkpoint",
|
||||
"Lora": "lora",
|
||||
"VAE": "vae",
|
||||
"Embeddings": "embedding",
|
||||
"controlnet": "controlnet",
|
||||
"upscale_models": "upscaler",
|
||||
"clip": "clip",
|
||||
}
|
||||
|
||||
EXT_ROOTS = (
|
||||
("swarmui", DATA / "Extensions"),
|
||||
("comfy", DATA / "DLNodes"),
|
||||
)
|
||||
|
||||
|
||||
def strip_auth(url: str) -> str:
|
||||
parts = urlsplit(url)
|
||||
host = parts.hostname or ""
|
||||
if parts.port:
|
||||
host = f"{host}:{parts.port}"
|
||||
return urlunsplit((parts.scheme, host, parts.path, parts.query, parts.fragment))
|
||||
|
||||
|
||||
def sha256_file(path: Path, chunk: int = 1024 * 1024) -> str:
|
||||
h = hashlib.sha256()
|
||||
with path.open("rb") as fh:
|
||||
while True:
|
||||
block = fh.read(chunk)
|
||||
if not block:
|
||||
break
|
||||
h.update(block)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def read_sidecar_ids(weight: Path) -> tuple[int | None, int | None]:
|
||||
"""Return (version_id, model_id) from {stem}.civitai.json if present."""
|
||||
sidecar = weight.parent / f"{weight.stem}.civitai.json"
|
||||
if not sidecar.is_file():
|
||||
return None, None
|
||||
try:
|
||||
data = json.loads(sidecar.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None, None
|
||||
if not isinstance(data, dict):
|
||||
return None, None
|
||||
vid = data.get("id")
|
||||
mid = data.get("modelId")
|
||||
try:
|
||||
version_id = int(vid) if vid is not None else None
|
||||
except (TypeError, ValueError):
|
||||
version_id = None
|
||||
try:
|
||||
model_id = int(mid) if mid is not None else None
|
||||
except (TypeError, ValueError):
|
||||
model_id = None
|
||||
return version_id, model_id
|
||||
|
||||
|
||||
def scan_models() -> list[dict]:
|
||||
items: list[dict] = []
|
||||
if not MODELS.is_dir():
|
||||
return items
|
||||
for folder, kind in FOLDER_TO_KIND.items():
|
||||
root = MODELS / folder
|
||||
if not root.is_dir():
|
||||
continue
|
||||
for path in sorted(root.rglob("*")):
|
||||
if not path.is_file():
|
||||
continue
|
||||
if path.suffix.lower() not in WEIGHT_SUFFIXES:
|
||||
continue
|
||||
if path.name.startswith("."):
|
||||
continue
|
||||
rel = path.relative_to(MODELS).as_posix()
|
||||
version_id, model_id = read_sidecar_ids(path)
|
||||
try:
|
||||
digest = sha256_file(path)
|
||||
except OSError as exc:
|
||||
items.append(
|
||||
{
|
||||
"kind": kind,
|
||||
"rel": rel,
|
||||
"name": path.name,
|
||||
"sha256": None,
|
||||
"version_id": version_id,
|
||||
"model_id": model_id,
|
||||
"error": str(exc),
|
||||
}
|
||||
)
|
||||
continue
|
||||
items.append(
|
||||
{
|
||||
"kind": kind,
|
||||
"rel": rel,
|
||||
"name": path.name,
|
||||
"sha256": digest,
|
||||
"version_id": version_id,
|
||||
"model_id": model_id,
|
||||
}
|
||||
)
|
||||
return items
|
||||
|
||||
|
||||
def git_out(args: list[str], cwd: Path) -> str | None:
|
||||
try:
|
||||
return subprocess.check_output(
|
||||
["git", "-C", str(cwd), *args],
|
||||
text=True,
|
||||
stderr=subprocess.DEVNULL,
|
||||
).strip()
|
||||
except (subprocess.CalledProcessError, FileNotFoundError, OSError):
|
||||
return None
|
||||
|
||||
|
||||
def scan_extensions() -> list[dict]:
|
||||
items: list[dict] = []
|
||||
for kind, root in EXT_ROOTS:
|
||||
if not root.is_dir():
|
||||
continue
|
||||
for child in sorted(root.iterdir()):
|
||||
if not child.is_dir():
|
||||
continue
|
||||
if not (child / ".git").exists():
|
||||
items.append(
|
||||
{
|
||||
"kind": kind,
|
||||
"dir": child.name,
|
||||
"url": None,
|
||||
"ref": None,
|
||||
"unknown": True,
|
||||
"reason": "no .git",
|
||||
}
|
||||
)
|
||||
continue
|
||||
origin = git_out(["remote", "get-url", "origin"], child)
|
||||
if not origin:
|
||||
items.append(
|
||||
{
|
||||
"kind": kind,
|
||||
"dir": child.name,
|
||||
"url": None,
|
||||
"ref": None,
|
||||
"unknown": True,
|
||||
"reason": "no origin",
|
||||
}
|
||||
)
|
||||
continue
|
||||
url = strip_auth(origin)
|
||||
branch = git_out(["rev-parse", "--abbrev-ref", "HEAD"], child)
|
||||
if not branch or branch == "HEAD":
|
||||
sha = git_out(["rev-parse", "--short", "HEAD"], child)
|
||||
ref = sha or "main"
|
||||
else:
|
||||
ref = branch
|
||||
items.append(
|
||||
{
|
||||
"kind": kind,
|
||||
"dir": child.name,
|
||||
"url": url,
|
||||
"ref": ref,
|
||||
"unknown": False,
|
||||
}
|
||||
)
|
||||
return items
|
||||
|
||||
|
||||
def main() -> int:
|
||||
payload = {
|
||||
"models": scan_models(),
|
||||
"extensions": scan_extensions(),
|
||||
}
|
||||
OUT.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||||
# One-line marker for local parsers; full JSON is in OUT.
|
||||
print(f"inventory ok models={len(payload['models'])} extensions={len(payload['extensions'])}")
|
||||
print(f"INVENTORY_PATH={OUT}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user