Ben FinkleaandClaude Opus 4.5
9256a1c87c
A custom node extension for ComfyUI that provides: - Model version tracking for HuggingFace and CivitAI - Automatic model updates with configurable scheduling - Multi-GPU synchronization with VRAM-based filtering - Metadata management with hash verification Features: - SQLite database for model metadata storage - REST API endpoints for all operations - Web UI panel integrated into ComfyUI menu - Background scheduler for automatic update checks - rsync-based sync to remote machines - VRAM estimation for model filtering Tests: - 17 unit tests covering all library modules - Playwright integration tests for UI verification Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
347 lines
12 KiB
Python
347 lines
12 KiB
Python
"""SQLite database operations for model metadata storage."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
import os
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import Optional, List, Dict, Any
|
|
from contextlib import contextmanager
|
|
|
|
|
|
class MetadataStore:
|
|
"""SQLite-based storage for model metadata and sync configuration."""
|
|
|
|
SCHEMA_VERSION = 1
|
|
|
|
def __init__(self, db_path: Optional[str] = None):
|
|
"""
|
|
Initialize the metadata store.
|
|
|
|
Args:
|
|
db_path: Path to SQLite database file. If None, uses default location.
|
|
"""
|
|
if db_path is None:
|
|
# Default to storing in the custom node directory
|
|
db_path = Path(__file__).parent.parent / "modelsync.db"
|
|
self.db_path = Path(db_path)
|
|
self._init_database()
|
|
|
|
@contextmanager
|
|
def _get_connection(self):
|
|
"""Context manager for database connections."""
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
conn.row_factory = sqlite3.Row
|
|
try:
|
|
yield conn
|
|
conn.commit()
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
def _init_database(self):
|
|
"""Initialize database schema."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# Models table - tracks installed models
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS models (
|
|
id INTEGER PRIMARY KEY,
|
|
path TEXT UNIQUE NOT NULL,
|
|
filename TEXT NOT NULL,
|
|
model_type TEXT NOT NULL,
|
|
source_url TEXT,
|
|
source_type TEXT DEFAULT 'unknown',
|
|
|
|
-- Local file info
|
|
file_hash TEXT,
|
|
file_size INTEGER,
|
|
install_date TEXT,
|
|
|
|
-- Remote version info (cached)
|
|
remote_hash TEXT,
|
|
remote_size INTEGER,
|
|
remote_modified TEXT,
|
|
last_checked TEXT,
|
|
|
|
-- Status
|
|
update_available INTEGER DEFAULT 0,
|
|
auto_update INTEGER DEFAULT 0,
|
|
|
|
-- Multi-GPU sync
|
|
vram_mb INTEGER,
|
|
exclude_from_sync INTEGER DEFAULT 0
|
|
)
|
|
""")
|
|
|
|
# Sync targets table - other machines to sync to
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS sync_targets (
|
|
id INTEGER PRIMARY KEY,
|
|
name TEXT UNIQUE NOT NULL,
|
|
ssh_host TEXT,
|
|
model_path TEXT NOT NULL,
|
|
max_vram_mb INTEGER,
|
|
enabled INTEGER DEFAULT 1,
|
|
last_sync TEXT
|
|
)
|
|
""")
|
|
|
|
# Configuration table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS config (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT
|
|
)
|
|
""")
|
|
|
|
# Schema version tracking
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS schema_version (
|
|
version INTEGER PRIMARY KEY
|
|
)
|
|
""")
|
|
|
|
# Insert default config if not exists
|
|
default_config = {
|
|
"auto_update_enabled": "false",
|
|
"check_interval_hours": "24",
|
|
"cache_ttl_hours": "6",
|
|
"verify_downloads": "true",
|
|
"backup_before_update": "true",
|
|
"max_concurrent_downloads": "2",
|
|
}
|
|
for key, value in default_config.items():
|
|
cursor.execute(
|
|
"INSERT OR IGNORE INTO config (key, value) VALUES (?, ?)",
|
|
(key, value)
|
|
)
|
|
|
|
# Model CRUD operations
|
|
|
|
def add_model(self, model_data: Dict[str, Any]) -> int:
|
|
"""Add a new model to the database."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# Build column list from provided data
|
|
columns = []
|
|
values = []
|
|
for key in ["path", "filename", "model_type", "source_url", "source_type",
|
|
"file_hash", "file_size", "install_date", "vram_mb"]:
|
|
if key in model_data:
|
|
columns.append(key)
|
|
values.append(model_data[key])
|
|
|
|
placeholders = ",".join(["?" for _ in values])
|
|
col_str = ",".join(columns)
|
|
|
|
cursor.execute(
|
|
f"INSERT INTO models ({col_str}) VALUES ({placeholders})",
|
|
values
|
|
)
|
|
return cursor.lastrowid
|
|
|
|
def get_model(self, model_id: int) -> Optional[Dict[str, Any]]:
|
|
"""Get a model by ID."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT * FROM models WHERE id = ?", (model_id,))
|
|
row = cursor.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def get_model_by_path(self, path: str) -> Optional[Dict[str, Any]]:
|
|
"""Get a model by file path."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT * FROM models WHERE path = ?", (path,))
|
|
row = cursor.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def get_all_models(self, model_type: Optional[str] = None) -> List[Dict[str, Any]]:
|
|
"""Get all models, optionally filtered by type."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
if model_type:
|
|
cursor.execute(
|
|
"SELECT * FROM models WHERE model_type = ? ORDER BY filename",
|
|
(model_type,)
|
|
)
|
|
else:
|
|
cursor.execute("SELECT * FROM models ORDER BY model_type, filename")
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_models_with_updates(self) -> List[Dict[str, Any]]:
|
|
"""Get all models that have updates available."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"SELECT * FROM models WHERE update_available = 1 ORDER BY filename"
|
|
)
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_auto_update_models(self) -> List[Dict[str, Any]]:
|
|
"""Get all models with auto-update enabled."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"SELECT * FROM models WHERE auto_update = 1 ORDER BY filename"
|
|
)
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def update_model(self, model_id: int, updates: Dict[str, Any]) -> bool:
|
|
"""Update model fields."""
|
|
if not updates:
|
|
return False
|
|
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
set_clauses = []
|
|
values = []
|
|
for key, value in updates.items():
|
|
set_clauses.append(f"{key} = ?")
|
|
values.append(value)
|
|
values.append(model_id)
|
|
|
|
cursor.execute(
|
|
f"UPDATE models SET {','.join(set_clauses)} WHERE id = ?",
|
|
values
|
|
)
|
|
return cursor.rowcount > 0
|
|
|
|
def delete_model(self, model_id: int) -> bool:
|
|
"""Delete a model from the database."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("DELETE FROM models WHERE id = ?", (model_id,))
|
|
return cursor.rowcount > 0
|
|
|
|
def delete_model_by_path(self, path: str) -> bool:
|
|
"""Delete a model by file path."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("DELETE FROM models WHERE path = ?", (path,))
|
|
return cursor.rowcount > 0
|
|
|
|
# Sync target operations
|
|
|
|
def add_sync_target(self, target_data: Dict[str, Any]) -> int:
|
|
"""Add a new sync target."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"""INSERT INTO sync_targets (name, ssh_host, model_path, max_vram_mb, enabled)
|
|
VALUES (?, ?, ?, ?, ?)""",
|
|
(
|
|
target_data["name"],
|
|
target_data.get("ssh_host"),
|
|
target_data["model_path"],
|
|
target_data.get("max_vram_mb"),
|
|
target_data.get("enabled", 1)
|
|
)
|
|
)
|
|
return cursor.lastrowid
|
|
|
|
def get_sync_target(self, target_id: int) -> Optional[Dict[str, Any]]:
|
|
"""Get a sync target by ID."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT * FROM sync_targets WHERE id = ?", (target_id,))
|
|
row = cursor.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def get_all_sync_targets(self, enabled_only: bool = False) -> List[Dict[str, Any]]:
|
|
"""Get all sync targets."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
if enabled_only:
|
|
cursor.execute(
|
|
"SELECT * FROM sync_targets WHERE enabled = 1 ORDER BY name"
|
|
)
|
|
else:
|
|
cursor.execute("SELECT * FROM sync_targets ORDER BY name")
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def update_sync_target(self, target_id: int, updates: Dict[str, Any]) -> bool:
|
|
"""Update sync target fields."""
|
|
if not updates:
|
|
return False
|
|
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
set_clauses = []
|
|
values = []
|
|
for key, value in updates.items():
|
|
set_clauses.append(f"{key} = ?")
|
|
values.append(value)
|
|
values.append(target_id)
|
|
|
|
cursor.execute(
|
|
f"UPDATE sync_targets SET {','.join(set_clauses)} WHERE id = ?",
|
|
values
|
|
)
|
|
return cursor.rowcount > 0
|
|
|
|
def delete_sync_target(self, target_id: int) -> bool:
|
|
"""Delete a sync target."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("DELETE FROM sync_targets WHERE id = ?", (target_id,))
|
|
return cursor.rowcount > 0
|
|
|
|
# Configuration operations
|
|
|
|
def get_config(self, key: str, default: Optional[str] = None) -> Optional[str]:
|
|
"""Get a configuration value."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT value FROM config WHERE key = ?", (key,))
|
|
row = cursor.fetchone()
|
|
return row["value"] if row else default
|
|
|
|
def set_config(self, key: str, value: str):
|
|
"""Set a configuration value."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)",
|
|
(key, value)
|
|
)
|
|
|
|
def get_all_config(self) -> Dict[str, str]:
|
|
"""Get all configuration values."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT key, value FROM config")
|
|
return {row["key"]: row["value"] for row in cursor.fetchall()}
|
|
|
|
# Utility methods
|
|
|
|
def get_model_count(self) -> int:
|
|
"""Get total number of tracked models."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT COUNT(*) as count FROM models")
|
|
return cursor.fetchone()["count"]
|
|
|
|
def get_update_count(self) -> int:
|
|
"""Get number of models with updates available."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"SELECT COUNT(*) as count FROM models WHERE update_available = 1"
|
|
)
|
|
return cursor.fetchone()["count"]
|
|
|
|
def clear_all_models(self):
|
|
"""Remove all models from the database. Use with caution."""
|
|
with self._get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("DELETE FROM models")
|