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>
345 lines
11 KiB
Python
345 lines
11 KiB
Python
"""Version tracking for HuggingFace and CivitAI models."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import asyncio
|
|
import aiohttp
|
|
from datetime import datetime
|
|
from typing import Optional, Dict, Any, Tuple, List
|
|
from urllib.parse import urlparse, unquote
|
|
|
|
try:
|
|
from huggingface_hub import HfApi, hf_hub_url
|
|
HF_AVAILABLE = True
|
|
except ImportError:
|
|
HF_AVAILABLE = False
|
|
|
|
|
|
class VersionTracker:
|
|
"""Track model versions from HuggingFace and CivitAI."""
|
|
|
|
CIVITAI_API_BASE = "https://civitai.com/api/v1"
|
|
|
|
def __init__(self):
|
|
"""Initialize the version tracker."""
|
|
if HF_AVAILABLE:
|
|
self.hf_api = HfApi()
|
|
else:
|
|
self.hf_api = None
|
|
self._session: Optional[aiohttp.ClientSession] = None
|
|
|
|
async def _get_session(self) -> aiohttp.ClientSession:
|
|
"""Get or create aiohttp session."""
|
|
if self._session is None or self._session.closed:
|
|
self._session = aiohttp.ClientSession()
|
|
return self._session
|
|
|
|
async def close(self):
|
|
"""Close the aiohttp session."""
|
|
if self._session and not self._session.closed:
|
|
await self._session.close()
|
|
|
|
def detect_source_type(self, url: str) -> str:
|
|
"""
|
|
Detect the source type from a URL.
|
|
|
|
Args:
|
|
url: The source URL
|
|
|
|
Returns:
|
|
Source type: 'huggingface', 'civitai', 'github', or 'unknown'
|
|
"""
|
|
if not url:
|
|
return "unknown"
|
|
|
|
parsed = urlparse(url)
|
|
domain = parsed.netloc.lower()
|
|
|
|
if "huggingface.co" in domain or "hf.co" in domain:
|
|
return "huggingface"
|
|
elif "civitai.com" in domain:
|
|
return "civitai"
|
|
elif "github.com" in domain or "raw.githubusercontent.com" in domain:
|
|
return "github"
|
|
else:
|
|
return "unknown"
|
|
|
|
def parse_huggingface_url(self, url: str) -> Optional[Tuple[str, str]]:
|
|
"""
|
|
Parse a HuggingFace URL to extract repo_id and file path.
|
|
|
|
Supports:
|
|
- https://huggingface.co/{repo}/resolve/main/{path}
|
|
- https://huggingface.co/{repo}/blob/main/{path}
|
|
- https://hf.co/{repo}/resolve/main/{path}
|
|
|
|
Args:
|
|
url: HuggingFace URL
|
|
|
|
Returns:
|
|
Tuple of (repo_id, file_path) or None if not parseable
|
|
"""
|
|
patterns = [
|
|
# https://huggingface.co/org/repo/resolve/main/path/to/file.safetensors
|
|
r"(?:https?://)?(?:huggingface\.co|hf\.co)/([^/]+/[^/]+)/(?:resolve|blob)/([^/]+)/(.+)",
|
|
# https://huggingface.co/org/repo/path/to/file.safetensors (direct)
|
|
r"(?:https?://)?(?:huggingface\.co|hf\.co)/([^/]+/[^/]+)/(?!resolve|blob)(.+)",
|
|
]
|
|
|
|
for pattern in patterns:
|
|
match = re.match(pattern, url)
|
|
if match:
|
|
groups = match.groups()
|
|
if len(groups) == 3:
|
|
repo_id = groups[0]
|
|
file_path = unquote(groups[2])
|
|
else:
|
|
repo_id = groups[0]
|
|
file_path = unquote(groups[1])
|
|
return repo_id, file_path
|
|
|
|
return None
|
|
|
|
def parse_civitai_url(self, url: str) -> Optional[Dict[str, str]]:
|
|
"""
|
|
Parse a CivitAI URL to extract model and version IDs.
|
|
|
|
Supports:
|
|
- https://civitai.com/models/{id}
|
|
- https://civitai.com/models/{id}?modelVersionId={versionId}
|
|
- https://civitai.com/api/download/models/{versionId}
|
|
|
|
Args:
|
|
url: CivitAI URL
|
|
|
|
Returns:
|
|
Dict with model_id and/or version_id, or None if not parseable
|
|
"""
|
|
patterns = [
|
|
# API download URL
|
|
(r"civitai\.com/api/download/models/(\d+)", "version_id"),
|
|
# Model page with version
|
|
(r"civitai\.com/models/(\d+)\?modelVersionId=(\d+)", "both"),
|
|
# Model page
|
|
(r"civitai\.com/models/(\d+)", "model_id"),
|
|
]
|
|
|
|
for pattern, result_type in patterns:
|
|
match = re.search(pattern, url)
|
|
if match:
|
|
if result_type == "version_id":
|
|
return {"version_id": match.group(1)}
|
|
elif result_type == "both":
|
|
return {"model_id": match.group(1), "version_id": match.group(2)}
|
|
else:
|
|
return {"model_id": match.group(1)}
|
|
|
|
return None
|
|
|
|
async def check_huggingface(self, url: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Check HuggingFace for file metadata.
|
|
|
|
Args:
|
|
url: HuggingFace URL
|
|
|
|
Returns:
|
|
Dict with hash, size, modified, or None on failure
|
|
"""
|
|
if not HF_AVAILABLE:
|
|
return None
|
|
|
|
parsed = self.parse_huggingface_url(url)
|
|
if not parsed:
|
|
return None
|
|
|
|
repo_id, file_path = parsed
|
|
|
|
try:
|
|
# Use HfApi to get file info
|
|
file_info = self.hf_api.get_paths_info(
|
|
repo_id=repo_id,
|
|
paths=[file_path],
|
|
repo_type="model"
|
|
)
|
|
|
|
if not file_info:
|
|
return None
|
|
|
|
info = file_info[0]
|
|
|
|
result = {
|
|
"source_type": "huggingface",
|
|
"repo_id": repo_id,
|
|
"file_path": file_path,
|
|
}
|
|
|
|
# Get LFS info if available (for large model files)
|
|
if hasattr(info, "lfs") and info.lfs:
|
|
result["hash"] = info.lfs.sha256
|
|
result["size"] = info.lfs.size
|
|
else:
|
|
result["size"] = getattr(info, "size", None)
|
|
|
|
# Get last commit info
|
|
if hasattr(info, "last_commit") and info.last_commit:
|
|
result["modified"] = info.last_commit.date.isoformat()
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
print(f"Error checking HuggingFace: {e}")
|
|
return None
|
|
|
|
async def check_civitai(self, url: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Check CivitAI for model version info.
|
|
|
|
Args:
|
|
url: CivitAI URL
|
|
|
|
Returns:
|
|
Dict with hash, size, version info, or None on failure
|
|
"""
|
|
parsed = self.parse_civitai_url(url)
|
|
if not parsed:
|
|
return None
|
|
|
|
try:
|
|
session = await self._get_session()
|
|
|
|
if "version_id" in parsed:
|
|
# Get model version directly
|
|
version_url = f"{self.CIVITAI_API_BASE}/model-versions/{parsed['version_id']}"
|
|
async with session.get(version_url) as resp:
|
|
if resp.status != 200:
|
|
return None
|
|
version_data = await resp.json()
|
|
else:
|
|
# Get model and find latest version
|
|
model_url = f"{self.CIVITAI_API_BASE}/models/{parsed['model_id']}"
|
|
async with session.get(model_url) as resp:
|
|
if resp.status != 200:
|
|
return None
|
|
model_data = await resp.json()
|
|
|
|
if not model_data.get("modelVersions"):
|
|
return None
|
|
|
|
# Get the first (latest) version
|
|
version_data = model_data["modelVersions"][0]
|
|
|
|
# Extract file info from primary file
|
|
files = version_data.get("files", [])
|
|
if not files:
|
|
return None
|
|
|
|
# Find primary file (or first file)
|
|
primary_file = next(
|
|
(f for f in files if f.get("primary", False)),
|
|
files[0]
|
|
)
|
|
|
|
result = {
|
|
"source_type": "civitai",
|
|
"model_id": version_data.get("modelId"),
|
|
"version_id": version_data.get("id"),
|
|
"version_name": version_data.get("name"),
|
|
"size": primary_file.get("sizeKB", 0) * 1024,
|
|
"modified": version_data.get("publishedAt"),
|
|
}
|
|
|
|
# Get hash info
|
|
hashes = primary_file.get("hashes", {})
|
|
if "SHA256" in hashes:
|
|
result["hash"] = hashes["SHA256"]
|
|
elif "BLAKE3" in hashes:
|
|
result["hash"] = hashes["BLAKE3"]
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
print(f"Error checking CivitAI: {e}")
|
|
return None
|
|
|
|
async def check_version(self, url: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Check for version info from any supported source.
|
|
|
|
Args:
|
|
url: Source URL (HuggingFace, CivitAI, etc.)
|
|
|
|
Returns:
|
|
Dict with version info or None
|
|
"""
|
|
source_type = self.detect_source_type(url)
|
|
|
|
if source_type == "huggingface":
|
|
return await self.check_huggingface(url)
|
|
elif source_type == "civitai":
|
|
return await self.check_civitai(url)
|
|
else:
|
|
return None
|
|
|
|
def is_update_available(
|
|
self,
|
|
local_hash: Optional[str],
|
|
local_size: Optional[int],
|
|
remote_hash: Optional[str],
|
|
remote_size: Optional[int],
|
|
size_tolerance: int = 1024
|
|
) -> bool:
|
|
"""
|
|
Determine if an update is available by comparing local and remote info.
|
|
|
|
Args:
|
|
local_hash: SHA256 hash of local file
|
|
local_size: Size of local file in bytes
|
|
remote_hash: SHA256 hash from remote source
|
|
remote_size: Size from remote source in bytes
|
|
size_tolerance: Byte tolerance for size comparison
|
|
|
|
Returns:
|
|
True if update appears available
|
|
"""
|
|
# Compare by hash (most reliable)
|
|
if local_hash and remote_hash:
|
|
return local_hash.lower() != remote_hash.lower()
|
|
|
|
# Fallback: compare by size
|
|
if local_size and remote_size:
|
|
return abs(local_size - remote_size) > size_tolerance
|
|
|
|
# Can't determine
|
|
return False
|
|
|
|
async def batch_check(
|
|
self,
|
|
models: List[Dict[str, Any]],
|
|
max_concurrent: int = 5
|
|
) -> Dict[int, Dict[str, Any]]:
|
|
"""
|
|
Check multiple models for updates concurrently.
|
|
|
|
Args:
|
|
models: List of model dicts with 'id' and 'source_url' keys
|
|
max_concurrent: Maximum concurrent requests
|
|
|
|
Returns:
|
|
Dict mapping model id to version info
|
|
"""
|
|
semaphore = asyncio.Semaphore(max_concurrent)
|
|
results = {}
|
|
|
|
async def check_one(model: Dict[str, Any]):
|
|
async with semaphore:
|
|
url = model.get("source_url")
|
|
if url:
|
|
info = await self.check_version(url)
|
|
if info:
|
|
results[model["id"]] = info
|
|
|
|
await asyncio.gather(*[check_one(m) for m in models])
|
|
return results
|