kenyonxu
47193575d6
- Cache manager with JSON read/write, TTL, corruption handling - Civitai client with SHA256 hash (sync+async) and API fetch - KaiCivitaiLoRALoader node with trigger words and JSON output - Node registration, HTTP routes, and secure image serving - Frontend extension with image preview, pagination, and NSFW blur - 49 unit tests covering all modules
111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
import json
|
|
import os
|
|
from datetime import datetime, timezone
|
|
|
|
|
|
def _import_compute_file_hash_sync():
|
|
"""Lazy import to avoid circular deps at import time."""
|
|
from .civitai_client import compute_file_hash_sync
|
|
return compute_file_hash_sync
|
|
|
|
|
|
class CacheManager:
|
|
"""Manages local JSON metadata cache and preview image storage."""
|
|
|
|
def __init__(self, cache_dir: str):
|
|
self.cache_dir = cache_dir
|
|
self._ensure_dir()
|
|
|
|
def _ensure_dir(self):
|
|
os.makedirs(self.cache_dir, exist_ok=True)
|
|
|
|
def _json_path(self, hash_val: str) -> str:
|
|
return os.path.join(self.cache_dir, f"{hash_val}.json")
|
|
|
|
def _ensure_image_dir(self, hash_val: str) -> str:
|
|
"""Ensure image directory exists and return its path."""
|
|
d = os.path.join(self.cache_dir, hash_val)
|
|
os.makedirs(d, exist_ok=True)
|
|
return d
|
|
|
|
def read_cache(self, hash_val: str) -> dict | None:
|
|
"""Read cached metadata. Auto-deletes corrupted JSON files."""
|
|
path = self._json_path(hash_val)
|
|
if not os.path.exists(path):
|
|
return None
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, OSError):
|
|
try:
|
|
os.remove(path)
|
|
except OSError:
|
|
pass
|
|
return None
|
|
|
|
def write_cache(self, hash_val: str, data: dict):
|
|
"""Write metadata to cache JSON."""
|
|
self._ensure_dir()
|
|
path = self._json_path(hash_val)
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
def is_stale(self, data: dict) -> bool:
|
|
"""Check if cached data is older than 7 days."""
|
|
fetched_at = data.get("fetched_at")
|
|
if not fetched_at:
|
|
return True
|
|
try:
|
|
# Parse ISO format with Z suffix (Python 3.10 compatible)
|
|
if fetched_at.endswith("Z"):
|
|
fetched_at = fetched_at[:-1]
|
|
dt = datetime.strptime(fetched_at, "%Y-%m-%dT%H:%M:%S")
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
age = datetime.now(timezone.utc) - dt
|
|
return age.total_seconds() > 7 * 24 * 3600
|
|
except (ValueError, TypeError):
|
|
return True
|
|
|
|
def clear_cache(self, hash_val: str):
|
|
"""Delete cached metadata and images for a hash."""
|
|
json_path = self._json_path(hash_val)
|
|
if os.path.exists(json_path):
|
|
try:
|
|
os.remove(json_path)
|
|
except OSError:
|
|
pass
|
|
img_dir = os.path.join(self.cache_dir, hash_val)
|
|
if os.path.isdir(img_dir):
|
|
try:
|
|
import shutil
|
|
shutil.rmtree(img_dir)
|
|
except OSError:
|
|
pass
|
|
|
|
def get_cache_for_file_sync(self, lora_path: str) -> dict | None:
|
|
"""Synchronous version for node execution (runs in worker thread).
|
|
|
|
Returns stale cache data if present. Node execution does not auto-fetch;
|
|
staleness is only checked by the frontend "Refresh" button flow.
|
|
"""
|
|
compute_file_hash_sync = _import_compute_file_hash_sync()
|
|
try:
|
|
file_hash = compute_file_hash_sync(lora_path)
|
|
except (OSError, FileNotFoundError):
|
|
return None
|
|
cache = self.read_cache(file_hash)
|
|
if cache is None:
|
|
return None
|
|
if "hash" not in cache:
|
|
cache["hash"] = file_hash
|
|
return cache
|
|
|
|
def get_image_path(self, hash_val: str, index: int) -> str:
|
|
"""Get path for a cached preview image."""
|
|
return os.path.join(self.cache_dir, hash_val, f"preview_{index}.jpg")
|
|
|
|
|
|
# Module-level singleton
|
|
_CACHE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cache", "civitai")
|
|
cache_manager = CacheManager(_CACHE_DIR)
|