"""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 # Known model patterns and their likely sources MODEL_PATTERNS = { # HuggingFace models r"gemma": ("huggingface", "google/gemma"), r"llama": ("huggingface", "meta-llama"), r"mistral": ("huggingface", "mistralai"), r"clip[_-]": ("huggingface", "openai/clip"), r"t5[_-]": ("huggingface", "google/t5"), r"bert": ("huggingface", "google-bert"), r"wav2vec": ("huggingface", "facebook/wav2vec2"), r"whisper": ("huggingface", "openai/whisper"), r"sam[_-]": ("huggingface", "facebook/sam"), r"dino": ("huggingface", "facebook/dinov2"), r"siglip": ("huggingface", "google/siglip"), r"flux": ("huggingface", "black-forest-labs/FLUX"), r"sd3": ("huggingface", "stabilityai/stable-diffusion-3"), r"sdxl": ("huggingface", "stabilityai/stable-diffusion-xl"), r"hunyuan": ("huggingface", "tencent/HunyuanVideo"), r"cogvideo": ("huggingface", "THUDM/CogVideoX"), r"ltx": ("huggingface", "Lightricks/LTX-Video"), r"mochi": ("huggingface", "genmo/mochi"), # GitHub/other sources r"yolo": ("github", "ultralytics/yolov5"), r"retinaface": ("github", "biubug6/Pytorch_Retinaface"), r"facexlib": ("github", "xinntao/facexlib"), r"gfpgan": ("github", "TencentARC/GFPGAN"), r"realesrgan": ("github", "xinntao/Real-ESRGAN"), r"codeformer": ("github", "sczhou/CodeFormer"), r"animatediff": ("huggingface", "guoyww/animatediff"), r"mm_sd": ("huggingface", "guoyww/animatediff"), # motion modules r"svd": ("huggingface", "stabilityai/stable-video-diffusion"), r"insightface": ("github", "deepinsight/insightface"), r"onnx.*det": ("github", "deepinsight/insightface"), } 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 def guess_source_from_filename(self, filename: str) -> Optional[Dict[str, Any]]: """ Guess the source of a model based on filename patterns. Args: filename: The model filename Returns: Dict with source_type and likely source info, or None """ filename_lower = filename.lower() for pattern, (source_type, source_hint) in self.MODEL_PATTERNS.items(): if re.search(pattern, filename_lower): return { "source_type": source_type, "source_hint": source_hint, "matched_pattern": pattern, } return None async def search_huggingface_by_name(self, filename: str) -> Optional[Dict[str, Any]]: """ Search HuggingFace for a model by filename. Args: filename: The model filename to search for Returns: Dict with source info if found, or None """ if not HF_AVAILABLE: return None # Clean up filename for search search_name = filename.replace(".safetensors", "").replace(".ckpt", "") search_name = search_name.replace(".pt", "").replace(".pth", "") search_name = search_name.replace(".bin", "").replace(".onnx", "") search_name = re.sub(r"[-_]", " ", search_name) try: # Search HuggingFace models API from huggingface_hub import HfApi api = HfApi() # Try to find models matching the name models = api.list_models( search=search_name, limit=5, sort="downloads", direction=-1 ) models_list = list(models) if models_list: best_match = models_list[0] return { "source_type": "huggingface", "source_url": f"https://huggingface.co/{best_match.modelId}", "model_name": best_match.modelId, "downloads": getattr(best_match, "downloads", 0), "confidence": "low", # Search-based match } except Exception as e: print(f"HuggingFace search error: {e}") return None async def smart_lookup(self, filename: str, file_hash: Optional[str] = None) -> Optional[Dict[str, Any]]: """ Smart lookup that tries multiple strategies to find a model's source. 1. Try CivitAI by hash (most reliable) 2. Try pattern matching for known model families 3. Try HuggingFace search by name Args: filename: The model filename file_hash: Optional SHA256 hash Returns: Dict with source info if found """ # Strategy 1: CivitAI hash lookup (if we have a hash) if file_hash: result = await self.lookup_civitai_by_hash(file_hash) if result: result["lookup_method"] = "civitai_hash" return result # Strategy 2: Pattern matching pattern_match = self.guess_source_from_filename(filename) if pattern_match: # For HuggingFace patterns, try to find the actual model if pattern_match["source_type"] == "huggingface": hf_result = await self.search_huggingface_by_name(filename) if hf_result: hf_result["lookup_method"] = "huggingface_search" hf_result["matched_pattern"] = pattern_match["matched_pattern"] return hf_result # Return pattern match as a hint even if we can't find exact URL return { "source_type": pattern_match["source_type"], "source_hint": pattern_match["source_hint"], "lookup_method": "pattern_match", "confidence": "medium", } # Strategy 3: Generic HuggingFace search hf_result = await self.search_huggingface_by_name(filename) if hf_result: hf_result["lookup_method"] = "huggingface_search" return hf_result return None