"""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 async def lookup_civitai_by_hash(self, file_hash: str) -> Optional[Dict[str, Any]]: """ Look up a model on CivitAI by its file hash. Args: file_hash: SHA256 hash of the model file Returns: Dict with model info including source_url, or None if not found """ if not file_hash: return None try: session = await self._get_session() url = f"{self.CIVITAI_API_BASE}/model-versions/by-hash/{file_hash}" async with session.get(url) as resp: if resp.status == 404: return None if resp.status != 200: print(f"CivitAI hash lookup failed: {resp.status}") return None data = await resp.json() # Build the source URL model_id = data.get("modelId") version_id = data.get("id") if not model_id or not version_id: return None source_url = f"https://civitai.com/models/{model_id}?modelVersionId={version_id}" # Get file info files = data.get("files", []) primary_file = next( (f for f in files if f.get("primary", False)), files[0] if files else {} ) result = { "source_type": "civitai", "source_url": source_url, "model_id": model_id, "version_id": version_id, "model_name": data.get("model", {}).get("name", "Unknown"), "version_name": data.get("name"), "size": primary_file.get("sizeKB", 0) * 1024, "modified": data.get("publishedAt"), } # Get hash info hashes = primary_file.get("hashes", {}) if "SHA256" in hashes: result["remote_hash"] = hashes["SHA256"] return result except Exception as e: print(f"Error looking up hash on CivitAI: {e}") return None async def batch_lookup_by_hash( self, models: List[Dict[str, Any]], max_concurrent: int = 3 ) -> Dict[int, Dict[str, Any]]: """ Look up multiple models by hash on CivitAI concurrently. Args: models: List of model dicts with 'id' and 'file_hash' keys max_concurrent: Maximum concurrent requests (keep low to avoid rate limits) Returns: Dict mapping model id to lookup result """ semaphore = asyncio.Semaphore(max_concurrent) results = {} async def lookup_one(model: Dict[str, Any]): async with semaphore: file_hash = model.get("file_hash") if file_hash and not model.get("source_url"): # Small delay to avoid rate limiting await asyncio.sleep(0.5) info = await self.lookup_civitai_by_hash(file_hash) if info: results[model["id"]] = info await asyncio.gather(*[lookup_one(m) for m in models]) return results