Ben FinkleaandClaude Opus 4.5
69197802b6
When CivitAI hash lookup fails, now tries: 1. Pattern matching for known model families (Gemma, CLIP, SAM, YOLO, etc.) 2. HuggingFace API search by filename This identifies sources for models from HuggingFace, GitHub, and other repositories that aren't on CivitAI. Also fixed lookup getting stuck on hint results by properly filtering out already-checked models regardless of source_type. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
595 lines
20 KiB
Python
595 lines
20 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
|
|
|
|
# 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
|