""" HuggingFace Source Module Search and download models from HuggingFace Hub. """ import json import os import re import threading import time import requests from typing import Dict, Any, Optional, List from urllib.parse import urlparse, quote from ..log_system.log_funcs import ( log_debug, log_info, log_warn, log_error, log_exception, ) HF_API_URL = "https://huggingface.co/api" HF_AUTHOR_FALLBACKS = ["Comfy-Org"] BRAVE_SEARCH_API_URL = "https://api.search.brave.com/res/v1/web/search" HF_AUTHOR_INDEX_CACHE_TTL_SECONDS = 24 * 60 * 60 HF_AUTHOR_INDEX_CACHE_VERSION = 1 METADATA_DIR = os.path.join( os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "metadata" ) HF_AUTHOR_INDEX_CACHE_PATH = os.path.join( METADATA_DIR, "huggingface-author-index.json" ) # Cache for search results _search_cache: Dict[str, Any] = {} _author_index_cache: Dict[str, Dict[str, Any]] = {} _author_index_lock = threading.RLock() def clear_search_cache(): """Clear cached HuggingFace search results and in-memory indexes.""" global _search_cache, _author_index_cache _search_cache.clear() with _author_index_lock: _author_index_cache.clear() def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]: """Check whether a HuggingFace access token is accepted.""" value = (token or "").strip() if not value: return { "success": False, "valid": False, "status": "missing", "message": "Paste a HuggingFace token first.", } headers = {"Authorization": f"Bearer {value}"} try: response = requests.get(f"{HF_API_URL}/whoami-v2", headers=headers, timeout=10) if response.status_code == 404: response = requests.get(f"{HF_API_URL}/whoami", headers=headers, timeout=10) if response.status_code == 200: username = "" try: data = response.json() username = data.get("name") or data.get("user", {}).get("name") or "" except Exception: username = "" message = "HuggingFace token is valid." if username: message = f"HuggingFace token is valid for {username}." return { "success": True, "valid": True, "status": "valid", "message": message, "username": username, } if response.status_code in {401, 403}: return { "success": True, "valid": False, "status": "invalid", "message": "HuggingFace token is not accepted.", "status_code": response.status_code, } return { "success": False, "valid": False, "status": "error", "message": f"HuggingFace returned HTTP {response.status_code}.", "status_code": response.status_code, } except requests.exceptions.Timeout: return { "success": False, "valid": False, "status": "timeout", "message": "HuggingFace did not respond before the timeout.", } except Exception as e: return { "success": False, "valid": False, "status": "error", "message": str(e), } def check_brave_search_api_key(api_key: Optional[str]) -> Dict[str, Any]: """Check whether a Brave Search API key is accepted.""" key = (api_key or "").strip() if not key: return { "success": False, "valid": False, "status": "missing", "message": "Paste a Brave Search API key first.", } headers = {"X-Subscription-Token": key} params = {"q": "test", "count": 1} try: response = requests.get( BRAVE_SEARCH_API_URL, headers=headers, params=params, timeout=10 ) if response.status_code == 200: return { "success": True, "valid": True, "status": "valid", "message": "Brave Search API key is valid.", } if response.status_code in {401, 403}: return { "success": True, "valid": False, "status": "invalid", "message": "Brave Search API key is not accepted.", "status_code": response.status_code, } if response.status_code == 429: return { "success": False, "valid": False, "status": "limited", "message": "Brave Search rate limit was reached.", "status_code": response.status_code, } return { "success": False, "valid": False, "status": "error", "message": f"Brave Search returned HTTP {response.status_code}.", "status_code": response.status_code, } except requests.exceptions.Timeout: return { "success": False, "valid": False, "status": "timeout", "message": "Brave Search did not respond before the timeout.", } except Exception as e: return { "success": False, "valid": False, "status": "error", "message": str(e), } def _author_index_cache_key(author: str, headers: Dict[str, str]) -> str: auth_key = "token" if headers.get("Authorization") else "public" return f"{author}::{auth_key}" def _is_author_index_fresh(index: Optional[Dict[str, Any]]) -> bool: if not index: return False updated_at = index.get("updated_at") if not isinstance(updated_at, (int, float)): return False return (time.time() - updated_at) < HF_AUTHOR_INDEX_CACHE_TTL_SECONDS def _read_persistent_author_indexes() -> Dict[str, Any]: if not os.path.exists(HF_AUTHOR_INDEX_CACHE_PATH): return {"version": HF_AUTHOR_INDEX_CACHE_VERSION, "authors": {}} try: with open(HF_AUTHOR_INDEX_CACHE_PATH, "r", encoding="utf-8") as f: data = json.load(f) if data.get("version") != HF_AUTHOR_INDEX_CACHE_VERSION: return {"version": HF_AUTHOR_INDEX_CACHE_VERSION, "authors": {}} if not isinstance(data.get("authors"), dict): data["authors"] = {} return data except Exception as e: log_debug(f"Error reading HuggingFace author index cache: {e}") return {"version": HF_AUTHOR_INDEX_CACHE_VERSION, "authors": {}} def _write_persistent_author_index(author: str, index: Dict[str, Any]) -> None: try: os.makedirs(METADATA_DIR, exist_ok=True) data = _read_persistent_author_indexes() data["authors"][author] = index tmp_path = f"{HF_AUTHOR_INDEX_CACHE_PATH}.tmp" with open(tmp_path, "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False, separators=(",", ":")) os.replace(tmp_path, HF_AUTHOR_INDEX_CACHE_PATH) except Exception as e: log_debug(f"Error writing HuggingFace author index cache: {e}") def _build_author_index_from_models( author: str, repos: List[Dict[str, Any]] ) -> Dict[str, Any]: files = [] repo_ids = [] for repo in repos: repo_id = repo.get("id", "") if not repo_id: continue repo_ids.append(repo_id) for sibling in repo.get("siblings") or []: file_path = sibling.get("rfilename") or sibling.get("path") or "" if not file_path: continue files.append( { "repo_id": repo_id, "path": file_path, "filename": os.path.basename(file_path), "size": sibling.get("size"), } ) return { "author": author, "updated_at": time.time(), "repo_count": len(repo_ids), "file_count": len(files), "repos": repo_ids, "files": files, } def _fetch_author_index(author: str, headers: Dict[str, str], limit: int = 200): try: response = requests.get( f"{HF_API_URL}/models", params={"author": author, "limit": limit, "full": "true"}, headers=headers, timeout=15, ) if response.status_code != 200: log_debug( f"HuggingFace author index returned {response.status_code} for author={author}" ) return None repos = response.json() if not isinstance(repos, list): return None index = _build_author_index_from_models(author, repos) log_info( f"HuggingFace author index refreshed for author={author}: repos={index['repo_count']}, files={index['file_count']}" ) return index except Exception as e: log_debug(f"Error refreshing HuggingFace author index for author={author}: {e}") return None def _get_author_index( author: str, headers: Dict[str, str], force_refresh: bool = False ) -> Optional[Dict[str, Any]]: cache_key = _author_index_cache_key(author, headers) can_persist = not headers.get("Authorization") with _author_index_lock: memory_index = _author_index_cache.get(cache_key) if not force_refresh and _is_author_index_fresh(memory_index): return memory_index persistent_index = None if can_persist: persistent_index = _read_persistent_author_indexes().get("authors", {}).get( author ) if not force_refresh and _is_author_index_fresh(persistent_index): _author_index_cache[cache_key] = persistent_index log_debug(f"HuggingFace author index cache hit for author={author}") return persistent_index fresh_index = _fetch_author_index(author, headers=headers) if fresh_index: _author_index_cache[cache_key] = fresh_index if can_persist: _write_persistent_author_index(author, fresh_index) return fresh_index if persistent_index: log_warn( f"Using stale HuggingFace author index for author={author}; refresh failed" ) _author_index_cache[cache_key] = persistent_index return persistent_index return None def get_author_fallback_index_status(author: str = "Comfy-Org") -> Dict[str, Any]: """Return persisted public author fallback index status for the options UI.""" with _author_index_lock: index = _read_persistent_author_indexes().get("authors", {}).get(author) updated_at = index.get("updated_at") if isinstance(index, dict) else None age_seconds = (time.time() - updated_at) if isinstance(updated_at, (int, float)) else None return { "author": author, "exists": bool(index), "updated_at": updated_at, "age_seconds": age_seconds, "ttl_seconds": HF_AUTHOR_INDEX_CACHE_TTL_SECONDS, "stale": bool(index) and not _is_author_index_fresh(index), "repo_count": index.get("repo_count", 0) if isinstance(index, dict) else 0, "file_count": index.get("file_count", 0) if isinstance(index, dict) else 0, } def refresh_author_fallback_index( token: Optional[str] = None, author: str = "Comfy-Org" ) -> Dict[str, Any]: """Force refresh the public HuggingFace author fallback index.""" headers = {} index = _get_author_index(author, headers=headers, force_refresh=True) status = get_author_fallback_index_status(author) status["success"] = bool(index) status["persisted"] = bool(index) return status def parse_huggingface_url(url: str) -> Optional[Dict[str, str]]: """ Parse a HuggingFace URL to extract repo and filename. Handles formats: - https://huggingface.co/user/repo/resolve/main/file.safetensors - https://huggingface.co/user/repo/blob/main/file.safetensors - hf://user/repo/file.safetensors Returns: Dictionary with 'repo' and 'filename' keys, or None if not HF URL """ if url.startswith("hf://"): parts = url[5:].split("/", 2) if len(parts) >= 3: return {"repo": f"{parts[0]}/{parts[1]}", "filename": parts[2]} return None parsed = urlparse(url) if "huggingface.co" not in parsed.netloc: return None match = re.match(r"^/([^/]+/[^/]+)/(resolve|blob)/([^/]+)/(.+)$", parsed.path) if match: return { "repo": match.group(1), "branch": match.group(3), "filename": match.group(4), } return None def get_huggingface_download_url(repo: str, filename: str, branch: str = "main") -> str: """Generate a direct download URL for a HuggingFace file.""" return f"https://huggingface.co/{repo}/resolve/{branch}/{quote(filename)}" def clean_filename_for_search(filename: str) -> str: """ Clean up filename for better search results. Remove common suffixes that might prevent matches. """ base = os.path.splitext(filename)[0] # Remove common precision/format suffixes base = re.sub( r"[-_]?(fp16|fp32|fp8|fp4|bf16|e4m3fn|scaled|pruned|emaonly|mixed|q4|q8).*$", "", base, flags=re.IGNORECASE, ) return base def _build_huggingface_result( repo_id: str, file_path: str, file_info: Dict[str, Any], match_type: str ) -> Dict[str, Any]: return { "source": "huggingface", "repo_id": repo_id, "filename": os.path.basename(file_path), "path": file_path, "url": get_huggingface_download_url(repo_id, file_path), "size": file_info.get("size"), "match_type": match_type, } def _get_repo_tree( repo_id: str, headers: Dict[str, str], branch: str = "main" ) -> Optional[List[Dict[str, Any]]]: files_url = f"{HF_API_URL}/models/{repo_id}/tree/{quote(branch, safe='')}" try: response = requests.get( files_url, params={"recursive": "1"}, headers=headers, timeout=15, ) if response.status_code != 200: log_debug( f"HuggingFace tree request returned {response.status_code} for repo={repo_id}, branch={branch}" ) return None return response.json() except Exception as e: log_debug(f"Error getting HuggingFace tree for {repo_id}@{branch}: {e}") return None def _find_matching_file_in_repo( repo_id: str, files: List[Dict[str, Any]], filename: str, exact_only: bool = False, ) -> Optional[Dict[str, Any]]: filename_lower = filename.lower() filename_base = os.path.splitext(filename_lower)[0] partial_match = None for file_info in files: file_path = file_info.get("path", "") if not file_path: continue file_name = os.path.basename(file_path) file_name_lower = file_name.lower() file_base = os.path.splitext(file_name_lower)[0] if file_name_lower == filename_lower or file_path.lower().endswith(filename_lower): return _build_huggingface_result(repo_id, file_path, file_info, "exact") if exact_only: return None for file_info in files: file_path = file_info.get("path", "") if not file_path: continue file_name = os.path.basename(file_path) file_name_lower = file_name.lower() file_base = os.path.splitext(file_name_lower)[0] if not exact_only: if filename_base in file_base or file_base in filename_base: if file_path.endswith(".safetensors") or file_path.endswith(".ckpt"): partial_match = _build_huggingface_result( repo_id, file_path, file_info, "partial" ) break return partial_match def _find_matching_file_in_author_index( index: Dict[str, Any], filename: str, exact_only: bool = False ) -> Optional[Dict[str, Any]]: files = index.get("files") or [] filename_lower = filename.lower() filename_base = os.path.splitext(filename_lower)[0] partial_match = None for file_info in files: repo_id = file_info.get("repo_id", "") file_path = file_info.get("path", "") if not repo_id or not file_path: continue file_name = os.path.basename(file_path) file_name_lower = file_name.lower() if file_name_lower == filename_lower or file_path.lower().endswith( filename_lower ): return _build_huggingface_result(repo_id, file_path, file_info, "exact") if exact_only: return None for file_info in files: repo_id = file_info.get("repo_id", "") file_path = file_info.get("path", "") if not repo_id or not file_path: continue file_name = os.path.basename(file_path) file_name_lower = file_name.lower() file_base = os.path.splitext(file_name_lower)[0] file_path_lower = file_path.lower() if filename_base in file_base or file_base in filename_base: if file_path_lower.endswith(".safetensors") or file_path_lower.endswith( ".ckpt" ): partial_match = _build_huggingface_result( repo_id, file_path, file_info, "partial" ) break return partial_match def _build_huggingface_search_queries(filename: str) -> List[str]: filename_base = os.path.splitext(os.path.basename(filename))[0].lower() cleaned_base = clean_filename_for_search(filename).lower() simplified_base = re.sub( r"[-_]?(fp16|fp32|fp8|fp4|bf16|e4m3fn|mixed|scaled|pruned|emaonly|q4|q8)$", "", filename_base, flags=re.IGNORECASE, ) queries = [] for query in [filename_base, cleaned_base, simplified_base]: query = (query or "").strip(" -_") if query and query not in queries: queries.append(query) return queries def _get_repos_by_author( author: str, headers: Dict[str, str], limit: int = 200 ) -> List[str]: try: response = requests.get( f"{HF_API_URL}/models", params={"author": author, "limit": limit}, headers=headers, timeout=15, ) if response.status_code != 200: log_debug( f"HuggingFace author search returned {response.status_code} for author={author}" ) return [] repos = response.json() return [repo.get("id", "") for repo in repos if repo.get("id")] except Exception as e: log_debug(f"Error getting HuggingFace repos for author={author}: {e}") return [] def _search_brave_for_huggingface_candidates( filename: str, brave_api_key: str ) -> List[Dict[str, str]]: if not brave_api_key: return [] query = f'"{filename}" site:huggingface.co' headers = { "Accept": "application/json", "Accept-Encoding": "gzip", "X-Subscription-Token": brave_api_key, } params = { "q": query, "count": 20, "safesearch": "off", "spellcheck": "false", } candidates = [] seen = set() try: response = requests.get( BRAVE_SEARCH_API_URL, headers=headers, params=params, timeout=15 ) if response.status_code != 200: log_warn( f"Brave search returned {response.status_code} for query={query}" ) return [] payload = response.json() results = ((payload or {}).get("web") or {}).get("results") or [] for result in results: candidate_url = result.get("url", "") if not candidate_url: continue parsed = parse_huggingface_url(candidate_url) if parsed: repo_id = parsed.get("repo", "") file_path = parsed.get("filename", "") branch = parsed.get("branch", "main") else: parsed_url = urlparse(candidate_url) if "huggingface.co" not in parsed_url.netloc: continue path_parts = [part for part in parsed_url.path.split("/") if part] if len(path_parts) < 2: continue repo_id = f"{path_parts[0]}/{path_parts[1]}" file_path = filename branch = "main" key = (repo_id, branch, file_path) if not repo_id or key in seen: continue seen.add(key) candidates.append( { "repo": repo_id, "branch": branch, "filename": file_path, "url": candidate_url, } ) except Exception as e: log_debug(f"Brave search error for filename={filename}: {e}") return candidates def search_huggingface_for_file( filename: str, token: Optional[str] = None, exact_only: bool = False, brave_api_key: Optional[str] = None, use_api_search: bool = True, use_comfy_org_fallback: bool = True, use_brave_fallback: bool = True, force_refresh: bool = False, ) -> Optional[Dict[str, Any]]: """ Search HuggingFace for a specific model file. Returns the first repo that contains a matching file. Args: filename: Filename to search for token: Optional HF token exact_only: If True, only return exact filename matches (for downloads). If False, also try partial matching (for local file resolution). Returns: Dict with url, repo, filename if found, None otherwise """ global _search_cache brave_key = "brave" if brave_api_key and use_brave_fallback else "nobrave" token_key = "token" if token else "notoken" methods_key = ( f"api{int(bool(use_api_search))}_" f"comfy{int(bool(use_comfy_org_fallback))}_" f"brave{int(bool(use_brave_fallback))}_" f"force{int(bool(force_refresh))}" ) cache_key = f"hf_{filename}_exact{exact_only}_{token_key}_{brave_key}_{methods_key}" if cache_key in _search_cache: log_debug(f"HuggingFace cache hit file={filename} exact={exact_only}") return _search_cache[cache_key] try: headers = {} if token: headers["Authorization"] = f"Bearer {token}" log_info(f"HuggingFace start file={filename} exact={exact_only}") repos_to_check = [] seen_repos = set() search_queries = _build_huggingface_search_queries(filename) if use_api_search: for search_query in search_queries: search_url = f"{HF_API_URL}/models?search={quote(search_query)}&limit=20" response = requests.get(search_url, headers=headers, timeout=10) if response.status_code != 200: log_warn( f"HuggingFace API status={response.status_code} query={search_query}" ) continue repos = response.json() log_info(f"HuggingFace API query={search_query} repos={len(repos)}") for repo in repos: repo_id = repo.get("id", "") if repo_id and repo_id not in seen_repos: seen_repos.add(repo_id) repos_to_check.append(repo_id) for repo_id in repos_to_check: log_debug(f"HuggingFace check repo={repo_id} file={filename}") files = _get_repo_tree(repo_id, headers=headers) if not files: continue log_debug(f"HuggingFace tree repo={repo_id} files={len(files)}") result = _find_matching_file_in_repo( repo_id, files, filename, exact_only=exact_only ) if result: _search_cache[cache_key] = result log_info( f"HuggingFace found match={result['match_type']} file={filename} repo={repo_id} path={result['path']}" ) return result author_fallback_repos = [] author_fallback_repo_count = 0 if use_comfy_org_fallback: for author in HF_AUTHOR_FALLBACKS: index = _get_author_index( author, headers={}, force_refresh=force_refresh, ) if index: result = _find_matching_file_in_author_index( index, filename, exact_only=exact_only ) author_fallback_repo_count += int(index.get("repo_count") or 0) if result: _search_cache[cache_key] = result log_info( f"HuggingFace found source=author_index match={result['match_type']} file={filename} repo={result['repo_id']} path={result['path']}" ) return result log_debug( f"HuggingFace author_index miss author={author} file={filename} repos={index.get('repo_count', 0)} files={index.get('file_count', 0)}" ) continue repos = _get_repos_by_author(author, headers=headers) log_info(f"HuggingFace author_fallback author={author} repos={len(repos)}") for repo_id in repos: if repo_id not in seen_repos: seen_repos.add(repo_id) author_fallback_repos.append(repo_id) author_fallback_repo_count += 1 for repo_id in author_fallback_repos: log_debug( f"HuggingFace author_fallback check repo={repo_id} file={filename}" ) files = _get_repo_tree(repo_id, headers=headers) if not files: continue result = _find_matching_file_in_repo( repo_id, files, filename, exact_only=exact_only ) if result: _search_cache[cache_key] = result log_info( f"HuggingFace found source=author_fallback match={result['match_type']} file={filename} repo={repo_id} path={result['path']}" ) return result brave_candidates = ( _search_brave_for_huggingface_candidates(filename, brave_api_key or "") if use_brave_fallback else [] ) log_info(f"HuggingFace brave_fallback file={filename} candidates={len(brave_candidates)}") for candidate in brave_candidates: repo_id = candidate.get("repo", "") branch = candidate.get("branch", "main") expected_file = candidate.get("filename", filename) files = _get_repo_tree(repo_id, headers=headers, branch=branch) if not files: continue result = _find_matching_file_in_repo( repo_id, files, expected_file, exact_only=True ) if result and result.get("filename", "").lower() == filename.lower(): _search_cache[cache_key] = result log_info( f"HuggingFace found source=brave_fallback match=exact file={filename} repo={repo_id} path={result['path']}" ) return result # Not found _search_cache[cache_key] = None log_info( f"HuggingFace miss file={filename} repos={len(repos_to_check)} author_repos={author_fallback_repo_count} brave={len(brave_candidates)}" ) return None except Exception as e: log_exception(f"HuggingFace search error for {filename}: {e}") return None def search_huggingface( query: str, model_type: Optional[str] = None, limit: int = 10, token: Optional[str] = None, ) -> List[Dict[str, Any]]: """ Search HuggingFace Hub for models (general search). Returns repos that might be relevant, not guaranteed to have exact file. """ results = [] try: headers = {} if token: headers["Authorization"] = f"Bearer {token}" params = {"search": query, "limit": limit, "full": "true"} response = requests.get( f"{HF_API_URL}/models", params=params, headers=headers, timeout=15 ) if response.status_code == 200: models = response.json() for model in models: repo_id = model.get("id", "") results.append( { "source": "huggingface", "repo": repo_id, "name": model.get("modelId", repo_id), "downloads": model.get("downloads", 0), "likes": model.get("likes", 0), "url": f"https://huggingface.co/{repo_id}", } ) except Exception as e: log_error(f"HuggingFace search error: {e}") return results def get_repo_files(repo: str, token: Optional[str] = None) -> List[Dict[str, Any]]: """Get list of model files in a HuggingFace repo.""" files = [] try: headers = {} if token: headers["Authorization"] = f"Bearer {token}" items = _get_repo_tree(repo, headers=headers) if items: for item in items: path = item.get("path", "") if path.endswith( (".safetensors", ".ckpt", ".pt", ".bin", ".pth", ".onnx", ".gguf") ): files.append( { "filename": os.path.basename(path), "path": path, "url": get_huggingface_download_url(repo, path), "size": item.get("size"), } ) except Exception as e: log_error(f"Error getting repo files: {e}") return files