mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-12 12:32:11 -03:00
Will Miao
fd40cd952b
`with sqlite3.Connection` ends the transaction but keeps the file handle open, so every cache class leaked one handle per instance until the garbage collector happened to run. On Windows that makes the database file undeletable while the process lives (WinError 32), which broke 30 cache tests in teardown and would also block moving or clearing a cache at runtime. All eight `with <connection>` sites now go through `closing_connection()`, and a regression test asserts the schema init leaves no usable connection behind.
985 lines
43 KiB
Python
985 lines
43 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple
|
|
|
|
from ..utils.cache_db import closing_connection, connect_cache_db
|
|
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
|
|
from ..utils.file_lock import exclusive_lock
|
|
from ..utils.root_paths import normalize_root_identity
|
|
from .model_sources import normalize_metadata_source
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _like_prefix_pattern(identity: str) -> str:
|
|
"""SQLite LIKE pattern matching every path inside *identity*.
|
|
|
|
``%``/``_`` are legal characters in a path and ``\\`` is the escape character,
|
|
so all three are escaped. The separator is added only when the identity does not
|
|
already end in one, so a filesystem root (``/``) still matches everything below
|
|
it instead of only itself.
|
|
"""
|
|
|
|
escaped = (
|
|
identity.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
|
)
|
|
prefix = escaped if escaped.endswith("/") else f"{escaped}/"
|
|
return f"{prefix}%"
|
|
|
|
|
|
@dataclass
|
|
class PersistedCacheData:
|
|
"""Lightweight structure returned by the persistent cache."""
|
|
|
|
raw_data: List[Dict[str, Any]]
|
|
hash_rows: List[Tuple[str, str]]
|
|
excluded_models: List[str]
|
|
autov3_hash_rows: List[Tuple[str, str]] = field(default_factory=list)
|
|
# Every directory under the model roots (including empty ones), or None
|
|
# when the snapshot predates folder recording.
|
|
all_folders: Optional[List[str]] = None
|
|
|
|
|
|
DEFAULT_LICENSE_FLAGS = 127 # 127 (0b1111111) encodes default CivitAI permissions with all commercial modes enabled.
|
|
|
|
|
|
class PersistentModelCache:
|
|
"""Persist core model metadata and hash index data in SQLite."""
|
|
|
|
_DEFAULT_FILENAME = "model_cache.sqlite"
|
|
_MODEL_COLUMNS: Tuple[str, ...] = (
|
|
"model_type",
|
|
"file_path",
|
|
"file_name",
|
|
"model_name",
|
|
"folder",
|
|
"size",
|
|
"modified",
|
|
"sha256",
|
|
"autov3",
|
|
"base_model",
|
|
"preview_url",
|
|
"preview_nsfw_level",
|
|
"from_civitai",
|
|
"favorite",
|
|
"notes",
|
|
"usage_tips",
|
|
"metadata_source",
|
|
"civitai_id",
|
|
"civitai_model_id",
|
|
"civitai_model_type",
|
|
"civitai_name",
|
|
"civitai_creator_username",
|
|
"trained_words",
|
|
"license_flags",
|
|
"civitai_deleted",
|
|
"skip_metadata_refresh",
|
|
"exclude",
|
|
"db_checked",
|
|
"last_checked_at",
|
|
"hash_status",
|
|
"source_platform",
|
|
"source_url",
|
|
"hf_url",
|
|
"source_model_id",
|
|
"source_version_id",
|
|
)
|
|
_MODEL_UPDATE_COLUMNS: Tuple[str, ...] = _MODEL_COLUMNS[2:]
|
|
_instances: Dict[str, "PersistentModelCache"] = {}
|
|
_instance_lock = threading.Lock()
|
|
|
|
def __init__(self, library_name: str = "default", db_path: Optional[str] = None) -> None:
|
|
self._library_name = library_name or "default"
|
|
self._db_path = db_path or self._resolve_default_path(self._library_name)
|
|
self._db_lock = threading.Lock()
|
|
self._schema_initialized = False
|
|
directory = os.path.dirname(self._db_path)
|
|
try:
|
|
if directory:
|
|
os.makedirs(directory, exist_ok=True)
|
|
except Exception as exc: # pragma: no cover - defensive guard
|
|
logger.warning("Could not create cache directory %s: %s", directory, exc)
|
|
if self.is_enabled():
|
|
self._initialize_schema()
|
|
|
|
@classmethod
|
|
def get_default(cls, library_name: Optional[str] = None) -> "PersistentModelCache":
|
|
name = (library_name or "default")
|
|
with cls._instance_lock:
|
|
if name not in cls._instances:
|
|
cls._instances[name] = cls(name)
|
|
return cls._instances[name]
|
|
|
|
def is_enabled(self) -> bool:
|
|
return os.environ.get("LORA_MANAGER_DISABLE_PERSISTENT_CACHE", "0") != "1"
|
|
|
|
def get_database_path(self) -> str:
|
|
"""Expose the resolved SQLite database path."""
|
|
|
|
return self._db_path
|
|
|
|
def load_cache(self, model_type: str) -> Optional[PersistedCacheData]:
|
|
if not self.is_enabled():
|
|
return None
|
|
if not self._schema_initialized:
|
|
self._initialize_schema()
|
|
if not self._schema_initialized:
|
|
return None
|
|
try:
|
|
with self._db_lock:
|
|
conn = self._connect(readonly=True)
|
|
try:
|
|
model_columns_sql = ", ".join(self._MODEL_COLUMNS[1:])
|
|
rows = conn.execute(
|
|
f"SELECT {model_columns_sql} FROM models WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
|
|
if not rows:
|
|
return None
|
|
|
|
tags = self._load_tags(conn, model_type)
|
|
hash_rows = conn.execute(
|
|
"SELECT sha256, file_path FROM hash_index WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
autov3_rows = conn.execute(
|
|
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
excluded = conn.execute(
|
|
"SELECT file_path FROM excluded_models WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
folder_rows = conn.execute(
|
|
"SELECT path FROM folders WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
folders_recorded = conn.execute(
|
|
"SELECT value FROM cache_meta WHERE key = ?",
|
|
(f"folders_recorded:{model_type}",),
|
|
).fetchone()
|
|
finally:
|
|
conn.close()
|
|
except Exception as exc:
|
|
logger.warning("Failed to load persisted cache for %s: %s", model_type, exc)
|
|
return None
|
|
|
|
raw_data: List[Dict[str, Any]] = []
|
|
for row in rows:
|
|
file_path: str = row["file_path"]
|
|
trained_words = []
|
|
if row["trained_words"]:
|
|
try:
|
|
trained_words = json.loads(row["trained_words"])
|
|
except json.JSONDecodeError:
|
|
trained_words = []
|
|
|
|
creator_username = row["civitai_creator_username"]
|
|
civitai: Optional[Dict[str, Any]] = None
|
|
civitai_has_data = any(
|
|
row[col] is not None
|
|
for col in ("civitai_id", "civitai_model_id", "civitai_model_type", "civitai_name")
|
|
) or trained_words or creator_username
|
|
if civitai_has_data:
|
|
civitai = {}
|
|
if row["civitai_id"] is not None:
|
|
civitai["id"] = row["civitai_id"]
|
|
if row["civitai_model_id"] is not None:
|
|
civitai["modelId"] = row["civitai_model_id"]
|
|
if row["civitai_name"]:
|
|
civitai["name"] = row["civitai_name"]
|
|
if trained_words:
|
|
civitai["trainedWords"] = trained_words
|
|
if creator_username:
|
|
civitai.setdefault("creator", {})["username"] = creator_username
|
|
model_type_value = row["civitai_model_type"]
|
|
if model_type_value:
|
|
civitai.setdefault("model", {})["type"] = model_type_value
|
|
|
|
license_value = row["license_flags"]
|
|
if license_value is None:
|
|
license_value = DEFAULT_LICENSE_FLAGS
|
|
|
|
item = {
|
|
"file_path": file_path,
|
|
"file_name": row["file_name"] or "",
|
|
"model_name": row["model_name"] or "",
|
|
"folder": row["folder"] or "",
|
|
"size": row["size"] or 0,
|
|
"modified": row["modified"] or 0.0,
|
|
"sha256": row["sha256"] or "",
|
|
"base_model": row["base_model"] or "",
|
|
"preview_url": row["preview_url"] or "",
|
|
"preview_nsfw_level": row["preview_nsfw_level"] or 0,
|
|
"from_civitai": bool(row["from_civitai"]),
|
|
"favorite": bool(row["favorite"]),
|
|
"notes": row["notes"] or "",
|
|
"usage_tips": row["usage_tips"] or "",
|
|
"metadata_source": row["metadata_source"] or None,
|
|
"exclude": bool(row["exclude"]),
|
|
"db_checked": bool(row["db_checked"]),
|
|
"last_checked_at": row["last_checked_at"] or 0.0,
|
|
"tags": tags.get(file_path, []),
|
|
"civitai": civitai,
|
|
"civitai_deleted": bool(row["civitai_deleted"]),
|
|
"skip_metadata_refresh": bool(row["skip_metadata_refresh"]),
|
|
"license_flags": int(license_value),
|
|
"hash_status": row["hash_status"] or "completed",
|
|
"source_platform": row["source_platform"] or "",
|
|
"source_url": row["source_url"] or "",
|
|
"hf_url": row["hf_url"] or "",
|
|
"source_model_id": row["source_model_id"] or "",
|
|
"source_version_id": row["source_version_id"] or "",
|
|
}
|
|
# Legacy rows only carry `hf_url`; derive the canonical pair so
|
|
# every consumer sees the same shape.
|
|
normalize_metadata_source(item)
|
|
if row["autov3"] is not None:
|
|
item["autov3"] = (row["autov3"] or "").lower()
|
|
raw_data.append(item)
|
|
|
|
hash_pairs = [(entry["sha256"].lower(), entry["file_path"]) for entry in hash_rows if entry["sha256"]]
|
|
if not hash_pairs:
|
|
# Fall back to hashes stored on the model rows
|
|
for item in raw_data:
|
|
sha_value = item.get("sha256")
|
|
if sha_value:
|
|
hash_pairs.append((sha_value.lower(), item["file_path"]))
|
|
|
|
autov3_pairs = [
|
|
(entry["autov3"].lower(), entry["file_path"])
|
|
for entry in autov3_rows
|
|
if entry["autov3"]
|
|
]
|
|
|
|
excluded_paths = [row["file_path"] for row in excluded]
|
|
all_folders: Optional[List[str]] = None
|
|
if folders_recorded is not None:
|
|
all_folders = sorted(
|
|
(row["path"] for row in folder_rows), key=lambda x: x.lower()
|
|
)
|
|
return PersistedCacheData(
|
|
raw_data=raw_data,
|
|
hash_rows=hash_pairs,
|
|
excluded_models=excluded_paths,
|
|
autov3_hash_rows=autov3_pairs,
|
|
all_folders=all_folders,
|
|
)
|
|
|
|
def roots_with_cached_entries(
|
|
self, model_type: str, roots: Sequence[str]
|
|
) -> Set[str]:
|
|
"""Return the subset of *roots* that hold at least one cached entry.
|
|
|
|
Used by the settings sync to tell a root that was registered but never
|
|
scanned -- ComfyUI registers ``<output>/loras`` for its save nodes whether
|
|
or not that directory exists -- from a real root whose drive happens to be
|
|
switched off: only the latter must keep its place in the saved library
|
|
paths (issue #1146). Paths are matched on their forward-slash identity, so
|
|
the host's own slash style does not matter.
|
|
|
|
Raises when the answer is unknown (a disabled cache, a missing database, an
|
|
unreadable file): the caller keeps every root in that case.
|
|
"""
|
|
|
|
if not self.is_enabled():
|
|
raise RuntimeError("persistent model cache is disabled")
|
|
|
|
candidates: Dict[str, str] = {}
|
|
for root in roots:
|
|
identity = normalize_root_identity(root)
|
|
if identity:
|
|
candidates.setdefault(identity, root)
|
|
if not candidates:
|
|
return set()
|
|
|
|
matched: Set[str] = set()
|
|
# One scan per model type rather than one per root: the probe runs at
|
|
# startup for every configured root that is missing right now, which on
|
|
# Windows includes all of ComfyUI's `<output>/<type>` folders.
|
|
clause = "replace(file_path, '\\', '/') LIKE ? ESCAPE '\\'"
|
|
sql = (
|
|
"SELECT DISTINCT replace(file_path, '\\', '/') FROM models "
|
|
f"WHERE model_type = ? AND ({' OR '.join([clause] * len(candidates))})"
|
|
)
|
|
params: List[Any] = [model_type]
|
|
params.extend(_like_prefix_pattern(identity) for identity in candidates)
|
|
|
|
with self._db_lock:
|
|
conn = self._connect(readonly=True)
|
|
try:
|
|
rows = conn.execute(sql, params).fetchall()
|
|
finally:
|
|
conn.close()
|
|
|
|
for (cached_path,) in rows:
|
|
folded = normalize_root_identity(cached_path)
|
|
for identity, root in candidates.items():
|
|
prefix = identity if identity.endswith("/") else f"{identity}/"
|
|
if folded.startswith(prefix):
|
|
matched.add(root)
|
|
break
|
|
return matched
|
|
|
|
def save_cache(self, model_type: str, raw_data: Sequence[Dict[str, Any]], hash_index: Dict[str, List[str]], excluded_models: Sequence[str], autov3_hash_index: Optional[Dict[str, List[str]]] = None, all_folders: Optional[Sequence[str]] = None) -> None:
|
|
if not self.is_enabled():
|
|
return
|
|
if not self._schema_initialized:
|
|
self._initialize_schema()
|
|
if not self._schema_initialized:
|
|
return
|
|
try:
|
|
with self._db_lock:
|
|
# Cross-process serialization: another LoRA Manager instance may
|
|
# share this settings directory, and the read-merge-write below
|
|
# spans several statements.
|
|
with exclusive_lock(self._db_path):
|
|
conn = self._connect()
|
|
try:
|
|
conn.execute("PRAGMA foreign_keys = ON")
|
|
conn.execute("BEGIN")
|
|
|
|
model_rows = [self._prepare_model_row(model_type, item) for item in raw_data]
|
|
model_map: Dict[str, Tuple[Any, ...]] = {
|
|
row[1]: row for row in model_rows if row[1] # row[1] is file_path
|
|
}
|
|
|
|
existing_models = conn.execute(
|
|
"SELECT "
|
|
+ ", ".join(self._MODEL_COLUMNS[1:])
|
|
+ " FROM models WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
existing_model_map: Dict[str, sqlite3.Row] = {
|
|
row["file_path"]: row for row in existing_models
|
|
}
|
|
|
|
to_remove_models = [
|
|
(model_type, path)
|
|
for path in existing_model_map.keys()
|
|
if path not in model_map
|
|
]
|
|
if to_remove_models:
|
|
conn.executemany(
|
|
"DELETE FROM models WHERE model_type = ? AND file_path = ?",
|
|
to_remove_models,
|
|
)
|
|
conn.executemany(
|
|
"DELETE FROM model_tags WHERE model_type = ? AND file_path = ?",
|
|
to_remove_models,
|
|
)
|
|
conn.executemany(
|
|
"DELETE FROM hash_index WHERE model_type = ? AND file_path = ?",
|
|
to_remove_models,
|
|
)
|
|
conn.executemany(
|
|
"DELETE FROM autov3_index WHERE model_type = ? AND file_path = ?",
|
|
to_remove_models,
|
|
)
|
|
conn.executemany(
|
|
"DELETE FROM excluded_models WHERE model_type = ? AND file_path = ?",
|
|
to_remove_models,
|
|
)
|
|
|
|
insert_rows: List[Tuple[Any, ...]] = []
|
|
update_rows: List[Tuple[Any, ...]] = []
|
|
|
|
for file_path, row in model_map.items():
|
|
existing = existing_model_map.get(file_path)
|
|
if existing is None:
|
|
insert_rows.append(row)
|
|
continue
|
|
|
|
existing_values = tuple(
|
|
existing[column] for column in self._MODEL_COLUMNS[1:]
|
|
)
|
|
current_values = row[1:]
|
|
if existing_values != current_values:
|
|
update_rows.append(row[2:] + (model_type, file_path))
|
|
|
|
if insert_rows:
|
|
conn.executemany(self._insert_model_sql(), insert_rows)
|
|
|
|
if update_rows:
|
|
set_clause = ", ".join(
|
|
f"{column} = ?"
|
|
for column in self._MODEL_UPDATE_COLUMNS
|
|
)
|
|
update_sql = (
|
|
f"UPDATE models SET {set_clause} WHERE model_type = ? AND file_path = ?"
|
|
)
|
|
conn.executemany(update_sql, update_rows)
|
|
|
|
existing_tags_rows = conn.execute(
|
|
"SELECT file_path, tag FROM model_tags WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
existing_tags: Dict[str, set[str]] = {}
|
|
for row in existing_tags_rows:
|
|
existing_tags.setdefault(row["file_path"], set()).add(row["tag"])
|
|
|
|
new_tags: Dict[str, set[str]] = {}
|
|
for item in raw_data:
|
|
file_path = item.get("file_path")
|
|
if not file_path:
|
|
continue
|
|
tags = set(item.get("tags") or [])
|
|
if tags:
|
|
new_tags[file_path] = tags
|
|
|
|
tag_inserts: List[Tuple[str, str, str]] = []
|
|
tag_deletes: List[Tuple[str, str, str]] = []
|
|
|
|
all_tag_paths = set(existing_tags.keys()) | set(new_tags.keys())
|
|
for path in all_tag_paths:
|
|
existing_set = existing_tags.get(path, set())
|
|
new_set = new_tags.get(path, set())
|
|
to_add = new_set - existing_set
|
|
to_remove = existing_set - new_set
|
|
|
|
for tag in to_add:
|
|
tag_inserts.append((model_type, path, tag))
|
|
for tag in to_remove:
|
|
tag_deletes.append((model_type, path, tag))
|
|
|
|
if tag_deletes:
|
|
conn.executemany(
|
|
"DELETE FROM model_tags WHERE model_type = ? AND file_path = ? AND tag = ?",
|
|
tag_deletes,
|
|
)
|
|
if tag_inserts:
|
|
conn.executemany(
|
|
"INSERT INTO model_tags (model_type, file_path, tag) VALUES (?, ?, ?)",
|
|
tag_inserts,
|
|
)
|
|
|
|
existing_hash_rows = conn.execute(
|
|
"SELECT sha256, file_path FROM hash_index WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
existing_hash_map: Dict[str, set[str]] = {}
|
|
for row in existing_hash_rows:
|
|
sha_value = (row["sha256"] or "").lower()
|
|
if not sha_value:
|
|
continue
|
|
existing_hash_map.setdefault(sha_value, set()).add(row["file_path"])
|
|
|
|
new_hash_map: Dict[str, set[str]] = {}
|
|
for sha_value, paths in hash_index.items():
|
|
normalized_sha = (sha_value or "").lower()
|
|
if not normalized_sha:
|
|
continue
|
|
bucket = new_hash_map.setdefault(normalized_sha, set())
|
|
for path in paths:
|
|
if path:
|
|
bucket.add(path)
|
|
|
|
hash_inserts: List[Tuple[str, str, str]] = []
|
|
hash_deletes: List[Tuple[str, str, str]] = []
|
|
|
|
all_shas = set(existing_hash_map.keys()) | set(new_hash_map.keys())
|
|
for sha_value in all_shas:
|
|
existing_paths = existing_hash_map.get(sha_value, set())
|
|
new_paths = new_hash_map.get(sha_value, set())
|
|
|
|
for path in existing_paths - new_paths:
|
|
hash_deletes.append((model_type, sha_value, path))
|
|
for path in new_paths - existing_paths:
|
|
hash_inserts.append((model_type, sha_value, path))
|
|
|
|
if hash_deletes:
|
|
conn.executemany(
|
|
"DELETE FROM hash_index WHERE model_type = ? AND sha256 = ? AND file_path = ?",
|
|
hash_deletes,
|
|
)
|
|
if hash_inserts:
|
|
conn.executemany(
|
|
"INSERT OR IGNORE INTO hash_index (model_type, sha256, file_path) VALUES (?, ?, ?)",
|
|
hash_inserts,
|
|
)
|
|
|
|
if autov3_hash_index is not None:
|
|
existing_autov3_rows = conn.execute(
|
|
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
existing_autov3_map: Dict[str, set[str]] = {}
|
|
for row in existing_autov3_rows:
|
|
autov3_value = (row["autov3"] or "").lower()
|
|
if not autov3_value:
|
|
continue
|
|
existing_autov3_map.setdefault(autov3_value, set()).add(row["file_path"])
|
|
|
|
new_autov3_map: Dict[str, set[str]] = {}
|
|
for autov3_value, paths in autov3_hash_index.items():
|
|
normalized_autov3 = (autov3_value or "").lower()
|
|
if not normalized_autov3:
|
|
continue
|
|
bucket = new_autov3_map.setdefault(normalized_autov3, set())
|
|
for path in paths:
|
|
if path:
|
|
bucket.add(path)
|
|
|
|
autov3_inserts: List[Tuple[str, str, str]] = []
|
|
autov3_deletes: List[Tuple[str, str, str]] = []
|
|
|
|
all_autov3 = set(existing_autov3_map.keys()) | set(new_autov3_map.keys())
|
|
for autov3_value in all_autov3:
|
|
existing_paths = existing_autov3_map.get(autov3_value, set())
|
|
new_paths = new_autov3_map.get(autov3_value, set())
|
|
|
|
for path in existing_paths - new_paths:
|
|
autov3_deletes.append((model_type, autov3_value, path))
|
|
for path in new_paths - existing_paths:
|
|
autov3_inserts.append((model_type, autov3_value, path))
|
|
|
|
if autov3_deletes:
|
|
conn.executemany(
|
|
"DELETE FROM autov3_index WHERE model_type = ? AND autov3 = ? AND file_path = ?",
|
|
autov3_deletes,
|
|
)
|
|
if autov3_inserts:
|
|
conn.executemany(
|
|
"INSERT OR IGNORE INTO autov3_index (model_type, autov3, file_path) VALUES (?, ?, ?)",
|
|
autov3_inserts,
|
|
)
|
|
|
|
existing_excluded_rows = conn.execute(
|
|
"SELECT file_path FROM excluded_models WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
existing_excluded = {row["file_path"] for row in existing_excluded_rows}
|
|
new_excluded = {path for path in excluded_models if path}
|
|
|
|
excluded_deletes = [
|
|
(model_type, path)
|
|
for path in existing_excluded - new_excluded
|
|
]
|
|
excluded_inserts = [
|
|
(model_type, path)
|
|
for path in new_excluded - existing_excluded
|
|
]
|
|
|
|
if excluded_deletes:
|
|
conn.executemany(
|
|
"DELETE FROM excluded_models WHERE model_type = ? AND file_path = ?",
|
|
excluded_deletes,
|
|
)
|
|
if excluded_inserts:
|
|
conn.executemany(
|
|
"INSERT OR IGNORE INTO excluded_models (model_type, file_path) VALUES (?, ?)",
|
|
excluded_inserts,
|
|
)
|
|
|
|
if all_folders is not None:
|
|
conn.execute(
|
|
"DELETE FROM folders WHERE model_type = ?",
|
|
(model_type,),
|
|
)
|
|
folder_inserts = [
|
|
(model_type, path) for path in all_folders if path
|
|
]
|
|
if folder_inserts:
|
|
conn.executemany(
|
|
"INSERT OR IGNORE INTO folders (model_type, path) VALUES (?, ?)",
|
|
folder_inserts,
|
|
)
|
|
# Mark the snapshot as having folder data even when the
|
|
# library has no subfolders, so an empty list is not
|
|
# mistaken for "never recorded" on load.
|
|
conn.execute(
|
|
"INSERT OR REPLACE INTO cache_meta (key, value) VALUES (?, ?)",
|
|
(f"folders_recorded:{model_type}", "1"),
|
|
)
|
|
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
except Exception as exc:
|
|
logger.warning("Failed to persist cache for %s: %s", model_type, exc)
|
|
|
|
# Internal helpers -------------------------------------------------
|
|
|
|
def _resolve_default_path(self, library_name: str) -> str:
|
|
env_override = os.environ.get("LORA_MANAGER_CACHE_DB")
|
|
return resolve_cache_path_with_migration(
|
|
CacheType.MODEL,
|
|
library_name=library_name,
|
|
env_override=env_override,
|
|
)
|
|
|
|
def _initialize_schema(self) -> None:
|
|
with self._db_lock:
|
|
if self._schema_initialized:
|
|
return
|
|
try:
|
|
with closing_connection(self._connect()) as conn:
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
conn.execute("PRAGMA foreign_keys = ON")
|
|
conn.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS models (
|
|
model_type TEXT NOT NULL,
|
|
file_path TEXT NOT NULL,
|
|
file_name TEXT,
|
|
model_name TEXT,
|
|
folder TEXT,
|
|
size INTEGER,
|
|
modified REAL,
|
|
sha256 TEXT,
|
|
autov3 TEXT,
|
|
base_model TEXT,
|
|
preview_url TEXT,
|
|
preview_nsfw_level INTEGER,
|
|
from_civitai INTEGER,
|
|
favorite INTEGER,
|
|
notes TEXT,
|
|
usage_tips TEXT,
|
|
metadata_source TEXT,
|
|
civitai_id INTEGER,
|
|
civitai_model_id INTEGER,
|
|
civitai_model_type TEXT,
|
|
civitai_name TEXT,
|
|
civitai_creator_username TEXT,
|
|
trained_words TEXT,
|
|
civitai_deleted INTEGER,
|
|
exclude INTEGER,
|
|
db_checked INTEGER,
|
|
last_checked_at REAL,
|
|
hash_status TEXT,
|
|
source_platform TEXT DEFAULT '',
|
|
source_url TEXT DEFAULT '',
|
|
hf_url TEXT DEFAULT '',
|
|
source_model_id TEXT DEFAULT '',
|
|
source_version_id TEXT DEFAULT '',
|
|
PRIMARY KEY (model_type, file_path)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS model_tags (
|
|
model_type TEXT NOT NULL,
|
|
file_path TEXT NOT NULL,
|
|
tag TEXT NOT NULL,
|
|
PRIMARY KEY (model_type, file_path, tag)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS hash_index (
|
|
model_type TEXT NOT NULL,
|
|
sha256 TEXT NOT NULL,
|
|
file_path TEXT NOT NULL,
|
|
PRIMARY KEY (model_type, sha256, file_path)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS autov3_index (
|
|
model_type TEXT NOT NULL,
|
|
autov3 TEXT NOT NULL,
|
|
file_path TEXT NOT NULL,
|
|
PRIMARY KEY (model_type, autov3, file_path)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS excluded_models (
|
|
model_type TEXT NOT NULL,
|
|
file_path TEXT NOT NULL,
|
|
PRIMARY KEY (model_type, file_path)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS folders (
|
|
model_type TEXT NOT NULL,
|
|
path TEXT NOT NULL,
|
|
PRIMARY KEY (model_type, path)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS cache_meta (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT
|
|
);
|
|
"""
|
|
)
|
|
self._ensure_additional_model_columns(conn)
|
|
conn.commit()
|
|
self._schema_initialized = True
|
|
except Exception as exc: # pragma: no cover - defensive guard
|
|
logger.warning("Failed to initialize persistent cache schema: %s", exc)
|
|
|
|
def _ensure_additional_model_columns(self, conn: sqlite3.Connection) -> None:
|
|
try:
|
|
existing_columns = {
|
|
row["name"]
|
|
for row in conn.execute("PRAGMA table_info(models)").fetchall()
|
|
}
|
|
except Exception: # pragma: no cover - defensive guard
|
|
return
|
|
|
|
required_columns = {
|
|
"metadata_source": "TEXT",
|
|
"civitai_creator_username": "TEXT",
|
|
"civitai_model_type": "TEXT",
|
|
"civitai_deleted": "INTEGER DEFAULT 0",
|
|
"skip_metadata_refresh": "INTEGER DEFAULT 0",
|
|
# Persisting without explicit flags should assume CivitAI's documented defaults (0b111001 == 57).
|
|
"license_flags": f"INTEGER DEFAULT {DEFAULT_LICENSE_FLAGS}",
|
|
"hash_status": "TEXT DEFAULT 'completed'",
|
|
"source_platform": "TEXT DEFAULT ''",
|
|
"source_url": "TEXT DEFAULT ''",
|
|
"hf_url": "TEXT DEFAULT ''",
|
|
"source_model_id": "TEXT DEFAULT ''",
|
|
"source_version_id": "TEXT DEFAULT ''",
|
|
"autov3": "TEXT",
|
|
}
|
|
|
|
for column, definition in required_columns.items():
|
|
if column not in existing_columns:
|
|
conn.execute(f"ALTER TABLE models ADD COLUMN {column} {definition}")
|
|
|
|
def _connect(self, readonly: bool = False) -> sqlite3.Connection:
|
|
if readonly and not os.path.exists(self._db_path):
|
|
raise FileNotFoundError(self._db_path)
|
|
return connect_cache_db(
|
|
self._db_path,
|
|
readonly=readonly,
|
|
detect_types=sqlite3.PARSE_DECLTYPES,
|
|
row_factory=sqlite3.Row,
|
|
)
|
|
|
|
def _prepare_model_row(self, model_type: str, item: Dict[str, Any]) -> Tuple[Any, ...]:
|
|
# Keep `source_*` and the legacy `hf_url` alias consistent no matter
|
|
# which caller populated the item.
|
|
normalize_metadata_source(item)
|
|
civitai = item.get("civitai") or {}
|
|
trained_words = civitai.get("trainedWords")
|
|
if isinstance(trained_words, str):
|
|
trained_words_json = trained_words
|
|
elif trained_words is None:
|
|
trained_words_json = None
|
|
else:
|
|
trained_words_json = json.dumps(trained_words)
|
|
|
|
metadata_source = item.get("metadata_source") or None
|
|
creator_username = None
|
|
creator_data = civitai.get("creator") if isinstance(civitai, dict) else None
|
|
if isinstance(creator_data, dict):
|
|
creator_username = creator_data.get("username") or None
|
|
model_type_value = None
|
|
if isinstance(civitai, Mapping):
|
|
civitai_model_info = civitai.get("model")
|
|
if isinstance(civitai_model_info, Mapping):
|
|
candidate_type = civitai_model_info.get("type")
|
|
if candidate_type not in (None, "", []):
|
|
model_type_value = candidate_type
|
|
|
|
license_flags = item.get("license_flags")
|
|
if license_flags is None:
|
|
license_flags = DEFAULT_LICENSE_FLAGS
|
|
|
|
autov3_value = item.get("autov3")
|
|
if autov3_value is None:
|
|
autov3_column = None
|
|
else:
|
|
autov3_column = (autov3_value or "").lower()
|
|
|
|
return (
|
|
model_type,
|
|
item.get("file_path"),
|
|
item.get("file_name") or "",
|
|
item.get("model_name") or "",
|
|
item.get("folder") or "",
|
|
int(item.get("size") or 0),
|
|
float(item.get("modified") or 0.0),
|
|
(item.get("sha256") or "").lower() or None,
|
|
autov3_column,
|
|
item.get("base_model") or "",
|
|
item.get("preview_url") or "",
|
|
int(item.get("preview_nsfw_level") or 0),
|
|
1 if item.get("from_civitai", True) else 0,
|
|
1 if item.get("favorite") else 0,
|
|
item.get("notes") or "",
|
|
item.get("usage_tips") or "",
|
|
metadata_source,
|
|
civitai.get("id"),
|
|
civitai.get("modelId"),
|
|
model_type_value,
|
|
civitai.get("name"),
|
|
creator_username,
|
|
trained_words_json,
|
|
int(license_flags),
|
|
1 if item.get("civitai_deleted") else 0,
|
|
1 if item.get("skip_metadata_refresh") else 0,
|
|
1 if item.get("exclude") else 0,
|
|
1 if item.get("db_checked") else 0,
|
|
float(item.get("last_checked_at") or 0.0),
|
|
item.get("hash_status", "completed"),
|
|
item.get("source_platform") or "",
|
|
item.get("source_url") or "",
|
|
item.get("hf_url") or "",
|
|
item.get("source_model_id") or "",
|
|
item.get("source_version_id") or "",
|
|
)
|
|
|
|
def _insert_model_sql(self) -> str:
|
|
columns = ", ".join(self._MODEL_COLUMNS)
|
|
placeholders = ", ".join(["?"] * len(self._MODEL_COLUMNS))
|
|
return f"INSERT INTO models ({columns}) VALUES ({placeholders})"
|
|
|
|
def update_single_model(
|
|
self,
|
|
model_type: str,
|
|
new_item: Dict[str, Any],
|
|
old_item: Optional[Dict[str, Any]] = None,
|
|
) -> None:
|
|
"""Update a single model row in the persistent cache.
|
|
|
|
A lightweight alternative to :meth:`save_cache` that performs a targeted
|
|
DELETE + INSERT for the model row and computes incremental tag / hash-index
|
|
deltas from *old_item*. When *old_item* is omitted the previous tags and
|
|
hash are not cleaned up (callers should only omit it for brand-new entries).
|
|
|
|
All operations run inside a single transaction so readers see a consistent
|
|
view.
|
|
"""
|
|
if not self.is_enabled():
|
|
return
|
|
if not self._schema_initialized:
|
|
self._initialize_schema()
|
|
if not self._schema_initialized:
|
|
return
|
|
|
|
file_path: Optional[str] = new_item.get("file_path")
|
|
if not file_path:
|
|
return
|
|
|
|
try:
|
|
with self._db_lock:
|
|
conn = self._connect()
|
|
try:
|
|
conn.execute("PRAGMA foreign_keys = ON")
|
|
conn.execute("BEGIN")
|
|
|
|
# --- model row (DELETE + INSERT = upsert) ---
|
|
conn.execute(
|
|
"DELETE FROM models WHERE model_type = ? AND file_path = ?",
|
|
(model_type, file_path),
|
|
)
|
|
row = self._prepare_model_row(model_type, new_item)
|
|
conn.execute(self._insert_model_sql(), row)
|
|
|
|
# --- tags ---
|
|
new_tags: set[str] = set(new_item.get("tags") or [])
|
|
old_tags: set[str] = set(old_item.get("tags") or []) if old_item else set()
|
|
tags_to_delete = old_tags - new_tags
|
|
tags_to_insert = new_tags - old_tags
|
|
|
|
if tags_to_delete:
|
|
conn.executemany(
|
|
"DELETE FROM model_tags WHERE model_type = ? AND file_path = ? AND tag = ?",
|
|
[(model_type, file_path, t) for t in tags_to_delete],
|
|
)
|
|
if tags_to_insert:
|
|
conn.executemany(
|
|
"INSERT INTO model_tags (model_type, file_path, tag) VALUES (?, ?, ?)",
|
|
[(model_type, file_path, t) for t in tags_to_insert],
|
|
)
|
|
|
|
# --- hash_index ---
|
|
new_sha: Optional[str] = (new_item.get("sha256") or "").lower() or None
|
|
old_sha: Optional[str] = (
|
|
(old_item.get("sha256") or "").lower() or None
|
|
) if old_item else None
|
|
if new_sha != old_sha:
|
|
if old_sha:
|
|
conn.execute(
|
|
"DELETE FROM hash_index WHERE model_type = ? AND sha256 = ? AND file_path = ?",
|
|
(model_type, old_sha, file_path),
|
|
)
|
|
if new_sha:
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO hash_index (model_type, sha256, file_path) VALUES (?, ?, ?)",
|
|
(model_type, new_sha, file_path),
|
|
)
|
|
|
|
# --- autov3_index ---
|
|
new_autov3: Optional[str] = new_item.get("autov3")
|
|
if new_autov3 is not None:
|
|
new_autov3 = (new_autov3 or "").lower()
|
|
old_autov3: Optional[str] = (old_item.get("autov3") if old_item else None)
|
|
if old_autov3 is not None:
|
|
old_autov3 = (old_autov3 or "").lower()
|
|
if new_autov3 != old_autov3:
|
|
if old_autov3:
|
|
conn.execute(
|
|
"DELETE FROM autov3_index WHERE model_type = ? AND autov3 = ? AND file_path = ?",
|
|
(model_type, old_autov3, file_path),
|
|
)
|
|
if new_autov3:
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO autov3_index (model_type, autov3, file_path) VALUES (?, ?, ?)",
|
|
(model_type, new_autov3, file_path),
|
|
)
|
|
|
|
conn.execute("COMMIT")
|
|
except Exception:
|
|
conn.execute("ROLLBACK")
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"Failed to update single model in persistent cache (%s): %s",
|
|
file_path,
|
|
exc,
|
|
)
|
|
|
|
def get_models_missing_autov3(self, model_type: str) -> List[str]:
|
|
"""Return file paths whose models lack an AutoV3 checked state.
|
|
|
|
Only rows with a completed sha256 and a NULL autov3 column qualify —
|
|
rows with '' (checked-unavailable) or a value are never returned, so
|
|
the backfill query self-terminates.
|
|
"""
|
|
if not self.is_enabled():
|
|
return []
|
|
if not self._schema_initialized:
|
|
self._initialize_schema()
|
|
if not self._schema_initialized:
|
|
return []
|
|
try:
|
|
with self._db_lock:
|
|
conn = self._connect(readonly=True)
|
|
try:
|
|
rows = conn.execute(
|
|
"SELECT file_path FROM models "
|
|
"WHERE model_type = ? AND autov3 IS NULL "
|
|
"AND sha256 IS NOT NULL AND sha256 != ''",
|
|
(model_type,),
|
|
).fetchall()
|
|
finally:
|
|
conn.close()
|
|
return [row["file_path"] for row in rows]
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"Failed to query models missing autov3 for %s: %s",
|
|
model_type,
|
|
exc,
|
|
)
|
|
return []
|
|
|
|
def _load_tags(self, conn: sqlite3.Connection, model_type: str) -> Dict[str, List[str]]:
|
|
tag_rows = conn.execute(
|
|
"SELECT file_path, tag FROM model_tags WHERE model_type = ?",
|
|
(model_type,),
|
|
).fetchall()
|
|
result: Dict[str, List[str]] = {}
|
|
for row in tag_rows:
|
|
result.setdefault(row["file_path"], []).append(row["tag"])
|
|
return result
|
|
|
|
|
|
def get_persistent_cache() -> PersistentModelCache:
|
|
from .settings_manager import get_settings_manager # Local import to avoid cycles
|
|
|
|
library_name = get_settings_manager().get_active_library_name()
|
|
return PersistentModelCache.get_default(library_name)
|