""" HuggingFace Source Module Search and download models from HuggingFace Hub. """ import os import posixpath import re import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Callable, Dict, List, Optional from urllib.parse import quote, unquote, urlparse from ..contracts import HuggingFaceFileReference, SearchResult from ..log_system import create_module_logger from ..progress import get_progress_reporter log = create_module_logger(__name__) from ..matcher import ( build_filename_search_queries, calculate_filename_confidence, ) from ..network_utils import execute_provider_json_request, host_matches_domain, request_source_response from ..path_utils import METADATA_DIR, get_filename_from_path, read_json_safe, write_json_atomic from ..type_utils import ( build_model_result, check_credential_http, clear_remote_size_cache, extract_file_size, looks_like_model_file, normalize_sha256, ) from .common import probe_remote_file_size, resolve_file_size HF_API_URL = "https://huggingface.co/api" HF_AUTHOR_FALLBACKS = ["Comfy-Org", "Kijai"] 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 = 2 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() clear_remote_size_cache() with _author_index_lock: _author_index_cache.clear() _report_progress = get_progress_reporter("HuggingFace progress callback") def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]: """Check whether a HuggingFace access token is accepted.""" from ..type_utils import check_credential_preconditions precheck = check_credential_preconditions(token, "HuggingFace token") if precheck: return precheck value = (token or "").strip() headers = {"Authorization": f"Bearer {value}"} def get_user(data): return data.get("name") or data.get("user", {}).get("name") or "" result = check_credential_http( f"{HF_API_URL}/whoami-v2", headers=headers, success_message="HuggingFace token is valid.", get_username=get_user, error_msg_401_403="HuggingFace token is not accepted.", ) if result.get("status_code") == 404: result = check_credential_http( f"{HF_API_URL}/whoami", headers=headers, success_message="HuggingFace token is valid.", get_username=get_user, error_msg_401_403="HuggingFace token is not accepted.", ) return result def check_brave_search_api_key(api_key: Optional[str]) -> Dict[str, Any]: """Check whether a Brave Search API key is accepted.""" from ..type_utils import check_credential_preconditions precheck = check_credential_preconditions(api_key, "Brave Search API key") if precheck: return precheck key = (api_key or "").strip() headers = {"X-Subscription-Token": key} params = {"q": "test", "count": 1} def custom_429(response): return { "success": False, "valid": False, "status": "limited", "message": "Brave Search rate limit was reached.", "status_code": 429, } return check_credential_http( BRAVE_SEARCH_API_URL, headers=headers, params=params, success_message="Brave Search API key is valid.", error_msg_401_403="Brave Search API key is not accepted.", custom_429_handler=custom_429, ) 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: Any) -> bool: if not isinstance(index, dict) or 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]: default_val = {"version": HF_AUTHOR_INDEX_CACHE_VERSION, "authors": {}} data = read_json_safe(HF_AUTHOR_INDEX_CACHE_PATH, default_val) if not isinstance(data, dict) or 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"] = {} else: # A short-lived regression stored the raw HuggingFace response list in # place of the built author index. Ignore malformed entries so they are # fetched and rebuilt normally. data["authors"] = { author: index for author, index in data["authors"].items() if isinstance(index, dict) } return data def _write_persistent_author_index(author: str, index: Dict[str, Any]) -> None: try: data = _read_persistent_author_indexes() data["authors"][author] = index write_json_atomic( HF_AUTHOR_INDEX_CACHE_PATH, data, indent=2, ) 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 = [] hash_count = 0 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 sha256 = _extract_huggingface_file_sha256(sibling) if sha256: hash_count += 1 files.append( { "repo_id": repo_id, "path": file_path, "filename": get_filename_from_path(file_path), "size": sibling.get("size"), "sha256": sha256, } ) return { "author": author, "updated_at": time.time(), "repo_count": len(repo_ids), "file_count": len(files), "hash_count": hash_count, "repos": repo_ids, "files": files, } def _is_huggingface_hash_candidate_path(file_path: str) -> bool: return str(file_path or "").lower().endswith( ( ".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf", ".onnx", ) ) def _enrich_author_index_with_repo_trees( index: Dict[str, Any], headers: Dict[str, str] ) -> Dict[str, Any]: """Add SHA-256 values from repository tree responses to an author index.""" files_by_repo_path = {} for file_info in index.get("files") or []: repo_id = file_info.get("repo_id", "") file_path = file_info.get("path", "") if ( repo_id and file_path and _is_huggingface_hash_candidate_path(file_path) and not _extract_huggingface_file_sha256(file_info) ): files_by_repo_path[(repo_id, file_path)] = file_info repo_ids = sorted({repo_id for repo_id, _ in files_by_repo_path}) if not repo_ids: index["hash_count"] = sum( 1 for file_info in index.get("files") or [] if _extract_huggingface_file_sha256(file_info) ) return index def fetch_repo_tree(repo_id: str): return repo_id, _get_repo_tree(repo_id, headers=headers) max_workers = min(8, len(repo_ids)) with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = { executor.submit(fetch_repo_tree, repo_id): repo_id for repo_id in repo_ids } for future in as_completed(futures): repo_id = futures[future] try: _, tree = future.result() except Exception as e: log.debug( f"HuggingFace author index tree failed repo={repo_id}: {e}" ) continue for tree_file in tree or []: file_path = tree_file.get("path", "") file_info = files_by_repo_path.get((repo_id, file_path)) if not file_info: continue sha256 = _extract_huggingface_file_sha256(tree_file) if sha256: file_info["sha256"] = sha256 if file_info.get("size") is None: file_info["size"] = extract_file_size(tree_file) index["hash_count"] = sum( 1 for file_info in index.get("files") or [] if _extract_huggingface_file_sha256(file_info) ) return index def _fetch_author_index(author: str, headers: Dict[str, str], limit: int = 200): try: repos = execute_provider_json_request( "HuggingFace author index", f"{HF_API_URL}/models", params={"author": author, "limit": limit, "full": "true"}, headers=headers, timeout=15, ) if not isinstance(repos, list): return None index = _build_author_index_from_models(author, repos) index = _enrich_author_index_with_repo_trees(index, headers=headers) log.info( f"HuggingFace author index refreshed for author={author}: " f"repos={index['repo_count']}, files={index['file_count']}, " f"hashes={index['hash_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, persist: bool = True, allow_stale: 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 if allow_stale and not force_refresh: stale_index = persistent_index or memory_index if isinstance(stale_index, dict): _author_index_cache[cache_key] = stale_index log.debug(f"HuggingFace stale author index cache hit for author={author}") return stale_index fresh_index = _fetch_author_index(author, headers=headers) if fresh_index: _author_index_cache[cache_key] = fresh_index if can_persist and persist: _write_persistent_author_index(author, fresh_index) return fresh_index if persistent_index: log.warning( 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, "hash_count": index.get("hash_count", 0) if isinstance(index, dict) else 0, } def get_known_author_fallback_indexes_status() -> Dict[str, Any]: """Return aggregate status for all known HuggingFace author indexes.""" author_statuses = [ get_author_fallback_index_status(author) for author in HF_AUTHOR_FALLBACKS ] updated_times = [ status["updated_at"] for status in author_statuses if isinstance(status.get("updated_at"), (int, float)) ] cached_author_count = sum(bool(status["exists"]) for status in author_statuses) fully_cached = cached_author_count == len(author_statuses) return { "authors": author_statuses, "author_count": len(author_statuses), "cached_author_count": cached_author_count, "exists": cached_author_count > 0, "fully_cached": fully_cached, "updated_at": max(updated_times) if updated_times else None, "stale": not fully_cached or any(bool(status["stale"]) for status in author_statuses), "repo_count": sum(int(status["repo_count"]) for status in author_statuses), "file_count": sum(int(status["file_count"]) for status in author_statuses), "hash_count": sum(int(status.get("hash_count") or 0) for status in author_statuses), } 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 refresh_known_author_fallback_indexes( token: Optional[str] = None, ) -> Dict[str, Any]: """Force refresh all known public HuggingFace author indexes.""" refresh_results = [ refresh_author_fallback_index(token, author) for author in HF_AUTHOR_FALLBACKS ] status = get_known_author_fallback_indexes_status() failed_authors = [ result["author"] for result in refresh_results if not result.get("success") ] status["success"] = not failed_authors status["persisted"] = all( bool(result.get("persisted")) for result in refresh_results ) status["failed_authors"] = failed_authors return status def parse_huggingface_url(url: str) -> Optional[HuggingFaceFileReference]: """ 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: Typed repository-file identity, or None if not an HF URL """ if url.startswith("hf://"): parts = url[5:].split("/", 2) if len(parts) >= 3: return HuggingFaceFileReference( repo=f"{parts[0]}/{parts[1]}", filename=parts[2], ) return None parsed = urlparse(url) if not host_matches_domain(parsed.hostname, "huggingface.co"): return None match = re.match(r"^/([^/]+/[^/]+)/(resolve|blob)/([^/]+)/(.+)$", parsed.path) if match: return HuggingFaceFileReference( 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/" f"{quote(branch, safe='')}/{quote(filename)}" ) def _normalize_huggingface_path(value: Any) -> str: return str(value or "").replace("\\", "/").strip("/") def _get_huggingface_variant_label(filename: str) -> str: """Return a concise precision/quantization label inferred from a filename.""" stem = os.path.splitext(get_filename_from_path(filename))[0].lower() patterns = ( ("nvfp4_mixed", "NVFP4 mixed"), ("fp8_scaled", "FP8 scaled"), ("int8_convrot", "INT8 convrot"), ("nvfp4", "NVFP4"), ("fp8", "FP8"), ("bf16", "BF16"), ("fp16", "FP16"), ("int8", "INT8"), ("int4", "INT4"), ) for token, label in patterns: if re.search(rf"(?:^|[_.-]){re.escape(token)}(?:$|[_.-])", stem): return label return "" def _normalize_huggingface_details_file( repo_id: str, branch: str, file_info: Dict[str, Any], target_path: str, ) -> Optional[Dict[str, Any]]: file_path = _normalize_huggingface_path(file_info.get("path")) if not file_path: return None download_url = get_huggingface_download_url(repo_id, file_path, branch) if not looks_like_model_file(download_url): return None filename = get_filename_from_path(file_path) sha256 = _extract_huggingface_file_sha256(file_info) extension = os.path.splitext(filename)[1].lstrip(".") metadata = { "format": extension.upper() if extension else "", "fp": _get_huggingface_variant_label(filename), } metadata = {key: value for key, value in metadata.items() if value} page_url = ( f"https://huggingface.co/{repo_id}/blob/" f"{quote(branch, safe='')}/{quote(file_path, safe='/')}" ) return { "name": filename, "filename": filename, "path": file_path, "size": extract_file_size(file_info), "sha256": sha256, "hash": sha256, "hashes": {"SHA256": sha256} if sha256 else {}, "download_url": download_url, "url": page_url, "primary": bool(target_path and file_path == target_path), "metadata": metadata, } def get_huggingface_model_details( repo_id: str, file_path: str = "", branch: str = "main", token: Optional[str] = None, ) -> Optional[Dict[str, Any]]: """Fetch model-file variants from the folder containing a matched HF file.""" normalized_repo_id = str(repo_id or "").strip().strip("/") if not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", normalized_repo_id): return None normalized_branch = str(branch or "main").strip() or "main" if len(normalized_branch) > 200 or any( character in normalized_branch for character in "\r\n?#" ): return None target_path = _normalize_huggingface_path(file_path) target_directory = posixpath.dirname(target_path) headers = {"Authorization": f"Bearer {token.strip()}"} if token and token.strip() else {} repo_tree = _get_repo_tree( normalized_repo_id, headers=headers, branch=normalized_branch, ) if not repo_tree: return None normalized_files = [] for file_info in repo_tree: candidate_path = _normalize_huggingface_path(file_info.get("path")) if not candidate_path: continue if target_path and posixpath.dirname(candidate_path) != target_directory: continue normalized_file = _normalize_huggingface_details_file( normalized_repo_id, normalized_branch, file_info, target_path, ) if normalized_file: normalized_files.append(normalized_file) if not normalized_files: return None normalized_files.sort( key=lambda item: ( not item.get("primary", False), str(item.get("name") or "").lower(), ) ) model_data: Dict[str, Any] = {} try: response_data = execute_provider_json_request( "HuggingFace model details", f"{HF_API_URL}/models/{normalized_repo_id}", headers=headers, timeout=15, ) if isinstance(response_data, dict): model_data = response_data except Exception as e: log.debug( f"Error getting HuggingFace model metadata for {normalized_repo_id}: {e}" ) card_data = ( model_data.get("cardData") if isinstance(model_data.get("cardData"), dict) else {} ) tags = model_data.get("tags") if isinstance(model_data.get("tags"), list) else [] repo_url = f"https://huggingface.co/{normalized_repo_id}" tree_url = ( f"{repo_url}/tree/{quote(normalized_branch, safe='')}" + ( f"/{quote(target_directory, safe='/')}" if target_directory else "" ) ) version = { "id": normalized_branch, "name": normalized_branch, "published_at": model_data.get("createdAt") or "", "updated_at": model_data.get("lastModified") or "", "description": "", "trained_words": [], "stats": { "downloads": model_data.get("downloads", 0), "thumbsUpCount": model_data.get("likes", 0), }, "files": normalized_files, "images": [], "url": tree_url, } owner = normalized_repo_id.split("/", 1)[0] description = ( model_data.get("description") or card_data.get("description") or ( f"Select a model file from `{target_directory}` in this Hugging Face " "repository." if target_directory else "Select a model file from this Hugging Face repository." ) ) return { "source": "huggingface", "details_source": "huggingface", "model_id": normalized_repo_id, "version_id": normalized_branch, "name": model_data.get("modelId") or normalized_repo_id, "type": model_data.get("pipeline_tag") or model_data.get("library_name") or "Model repository", "description": description, "creator": {"username": model_data.get("author") or owner}, "stats": { "downloads": model_data.get("downloads", 0), "thumbsUpCount": model_data.get("likes", 0), }, "tags": tags, "url": repo_url, "version_url": tree_url, "folder": target_directory, "branch": normalized_branch, "versions": [version], "selected_version": version, "images": [], } # _extract_file_info_size removed in favor of extract_file_size from type_utils def _build_huggingface_result( repo_id: str, file_path: str, file_info: Dict[str, Any], match_type: str, headers: Optional[Dict[str, str]] = None, confidence: float = 0.0, ) -> SearchResult: sha256 = _extract_huggingface_file_sha256(file_info) download_url = get_huggingface_download_url(repo_id, file_path) size = resolve_file_size( file_info, [download_url], probe=lambda url: probe_remote_file_size(url, headers=headers), ) return build_model_result( "huggingface", model_id=repo_id, version_id=None, filename=get_filename_from_path(file_path), url=download_url, download_url=download_url, size=size, match_type=match_type, confidence=confidence, sha256=sha256, normalize_hashes=True, repo_id=repo_id, path=file_path, branch="main", ) def _extract_huggingface_file_sha256(file_info: Dict[str, Any]) -> str: lfs_info = file_info.get("lfs") if isinstance(file_info.get("lfs"), dict) else {} return normalize_sha256( file_info.get("sha256") or file_info.get("hash") or lfs_info.get("sha256") or lfs_info.get("oid") ) 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: data = execute_provider_json_request( "HuggingFace tree", files_url, params={"recursive": "1"}, headers=headers, timeout=15, ) if isinstance(data, list): return data return None except Exception as e: log.debug(f"Error getting HuggingFace tree for {repo_id}@{branch}: {e}") return None def get_huggingface_file_sha256( url: str, headers: Optional[Dict[str, str]] = None, ) -> str: """Return the regular file SHA256 for a Hugging Face file URL.""" parsed = parse_huggingface_url(url) if not parsed: return "" repo_id = parsed.repo branch = parsed.branch target_path = _normalize_huggingface_path( unquote(parsed.filename) ) if not repo_id or not target_path: return "" repo_tree = _get_repo_tree(repo_id, headers=headers or {}, branch=branch) for file_info in repo_tree or []: candidate_path = _normalize_huggingface_path(file_info.get("path")) if candidate_path == target_path: return _extract_huggingface_file_sha256(file_info) return "" def _find_matching_file_in_repo( repo_id: str, files: List[Dict[str, Any]], filename: str, exact_only: bool = False, headers: Optional[Dict[str, str]] = None, ) -> Optional[SearchResult]: filename_lower = filename.lower() filename_base = os.path.splitext(filename_lower)[0] for file_info in files: file_path = file_info.get("path", "") if not file_path: continue file_name = get_filename_from_path(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", headers=headers, confidence=100.0, ) if exact_only: return None best_match: Optional[SearchResult] = None for file_info in files: file_path = file_info.get("path", "") if not file_path: continue file_name = get_filename_from_path(file_path) file_name_lower = file_name.lower() file_base = os.path.splitext(file_name_lower)[0] if not file_path.lower().endswith((".safetensors", ".ckpt")): continue confidence = calculate_filename_confidence(filename, file_name) is_partial_filename = ( filename_base in file_base or file_base in filename_base ) if not is_partial_filename and confidence < 70.0: continue candidate = _build_huggingface_result( repo_id, file_path, file_info, "partial", headers=headers, confidence=confidence, ) if best_match is None or candidate.confidence > best_match.confidence: best_match = candidate return best_match def _find_matching_file_in_author_index( index: Dict[str, Any], filename: str, exact_only: bool = False, headers: Optional[Dict[str, str]] = None, sha256: Optional[str] = None, ) -> Optional[SearchResult]: files = index.get("files") or [] requested_sha256 = normalize_sha256(sha256) if requested_sha256: 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 or _extract_huggingface_file_sha256(file_info) != requested_sha256 ): continue return _build_huggingface_result( repo_id, file_path, file_info, "hash", headers=headers, confidence=100.0, ) return None filename_lower = filename.lower() filename_base = os.path.splitext(filename_lower)[0] 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 = get_filename_from_path(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", headers=headers, confidence=100.0, ) if exact_only: return None best_match: Optional[SearchResult] = 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 = get_filename_from_path(file_path) file_name_lower = file_name.lower() file_base = os.path.splitext(file_name_lower)[0] file_path_lower = file_path.lower() if not file_path_lower.endswith((".safetensors", ".ckpt")): continue confidence = calculate_filename_confidence(filename, file_name) is_partial_filename = ( filename_base in file_base or file_base in filename_base ) if not is_partial_filename and confidence < 70.0: continue candidate = _build_huggingface_result( repo_id, file_path, file_info, "partial", headers=headers, confidence=confidence, ) if best_match is None or candidate.confidence > best_match.confidence: best_match = candidate return best_match def _get_repos_by_author( author: str, headers: Dict[str, str], limit: int = 200 ) -> List[str]: try: response = request_source_response( f"{HF_API_URL}/models", params={"author": author, "limit": limit}, headers=headers, timeout=15, log_name="HuggingFace repos by author" ) 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 = request_source_response( BRAVE_SEARCH_API_URL, headers=headers, params=params, timeout=15, log_name="Brave search" ) if response.status_code != 200: log.warning( 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.repo file_path = parsed.filename branch = parsed.branch 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, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, sha256: Optional[str] = None, ) -> Optional[SearchResult]: """ 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: A typed SearchResult if a matching file is found, otherwise None. """ global _search_cache requested_sha256 = normalize_sha256(sha256) 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}_sha{requested_sha256}_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}") _report_progress( progress_callback, "cache", "Using HuggingFace cache", 86, ) return _search_cache[cache_key] try: def result_matches_requested_hash(result: Optional[SearchResult]) -> bool: if not requested_sha256: return True return normalize_sha256( result.sha256 or result.extra_value("hash") or (result.hashes or {}).get("SHA256") or (result.hashes or {}).get("sha256") ) == requested_sha256 headers = {} if token: headers["Authorization"] = f"Bearer {token}" log.info(f"HuggingFace start file={filename} exact={exact_only}") def find_author_index_result( author: str, ) -> tuple[Optional[Dict[str, Any]], Optional[SearchResult]]: index = _get_author_index( author, headers={}, allow_stale=force_refresh, ) result = ( _find_matching_file_in_author_index( index, filename, exact_only=exact_only, headers={}, sha256=requested_sha256, ) if index else None ) if ( force_refresh and index and not _is_author_index_fresh(index) and result is None ): index = _get_author_index( author, headers={}, force_refresh=True, ) result = ( _find_matching_file_in_author_index( index, filename, exact_only=exact_only, headers={}, sha256=requested_sha256, ) if index else None ) return index, result cached_author_indexes: Dict[str, Optional[Dict[str, Any]]] = {} if requested_sha256 and use_comfy_org_fallback: total_authors = max(1, len(HF_AUTHOR_FALLBACKS)) for author_index, author in enumerate(HF_AUTHOR_FALLBACKS, start=1): _report_progress( progress_callback, "author_hash_index", f"Checking HF hash index {author_index}/{total_authors}", 28 + ((author_index - 1) / total_authors) * 6, author=author, author_index=author_index, author_count=total_authors, ) index, result = find_author_index_result(author) cached_author_indexes[author] = index if not index: continue if result: _search_cache[cache_key] = result _report_progress( progress_callback, "found", "Found HuggingFace SHA-256 match", 92, repo=result.extra_value("repo_id"), match_type=result.match_type, ) log.info( f"HuggingFace found source=author_index " f"match=hash file={result.filename} " f"repo={result.extra_value('repo_id')} path={result.extra_value('path')}" ) return result repos_to_check = [] seen_repos = set() search_queries = build_filename_search_queries(filename) if use_api_search: total_queries = max(1, len(search_queries)) for query_index, search_query in enumerate(search_queries, start=1): _report_progress( progress_callback, "api_search", f"HuggingFace API query {query_index}/{total_queries}", 32 + ((query_index - 1) / total_queries) * 10, query=search_query, query_index=query_index, query_count=total_queries, ) search_url = f"{HF_API_URL}/models?search={quote(search_query)}&limit=20" repos = execute_provider_json_request( "HuggingFace models search", search_url, api_key=token, timeout=10, ) if repos is None: log.warning( f"HuggingFace API query={search_query} failed" ) continue if not repos: continue 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) _report_progress( progress_callback, "api_candidates", f"HuggingFace API candidates: {len(repos_to_check)} repos", 44, candidate_count=len(repos_to_check), ) else: _report_progress( progress_callback, "api_skip", "Skipping HuggingFace API search", 44, ) repo_count = len(repos_to_check) for repo_index, repo_id in enumerate(repos_to_check, start=1): _report_progress( progress_callback, "repo_files", f"Checking HF repo {repo_index}/{repo_count}", 48 + (repo_index / max(1, repo_count)) * 14, repo=repo_id, candidate_index=repo_index, candidate_count=repo_count, ) 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, headers=headers, ) if result and result_matches_requested_hash(result): _search_cache[cache_key] = result _report_progress( progress_callback, "found", "Found HuggingFace match", 92, repo=repo_id, match_type=result.match_type, ) log.info( f"HuggingFace found match={result.match_type} file={filename} repo={repo_id} path={result.extra_value('path')}" ) return result author_fallback_repos = [] author_fallback_repo_count = 0 if use_comfy_org_fallback: total_authors = max(1, len(HF_AUTHOR_FALLBACKS)) for author_index, author in enumerate(HF_AUTHOR_FALLBACKS, start=1): _report_progress( progress_callback, "author_index", f"Checking HF fallback index {author_index}/{total_authors}", 64 + ((author_index - 1) / total_authors) * 6, author=author, author_index=author_index, author_count=total_authors, ) if author in cached_author_indexes: index = cached_author_indexes[author] result = ( _find_matching_file_in_author_index( index, filename, exact_only=exact_only, headers={}, sha256=requested_sha256, ) if index else None ) else: index, result = find_author_index_result(author) if index: author_fallback_repo_count += int(index.get("repo_count") or 0) if result and result_matches_requested_hash(result): _search_cache[cache_key] = result _report_progress( progress_callback, "found", "Found HuggingFace fallback match", 92, repo=result.extra_value("repo_id"), match_type=result.match_type, ) log.info( f"HuggingFace found source=author_index match={result.match_type} file={filename} repo={result.extra_value('repo_id')} path={result.extra_value('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 else: _report_progress( progress_callback, "author_skip", "Skipping HuggingFace fallback index", 70, ) fallback_repo_count = len(author_fallback_repos) for repo_index, repo_id in enumerate(author_fallback_repos, start=1): _report_progress( progress_callback, "author_repo", f"Checking HF fallback repo {repo_index}/{fallback_repo_count}", 70 + (repo_index / max(1, fallback_repo_count)) * 8, repo=repo_id, candidate_index=repo_index, candidate_count=fallback_repo_count, ) 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, headers=headers, ) if result and result_matches_requested_hash(result): _search_cache[cache_key] = result _report_progress( progress_callback, "found", "Found HuggingFace fallback match", 92, repo=repo_id, match_type=result.match_type, ) log.info( f"HuggingFace found source=author_fallback match={result.match_type} file={filename} repo={repo_id} path={result.extra_value('path')}" ) return result if use_brave_fallback: _report_progress( progress_callback, "brave_search", "Searching Brave fallback", 80, ) 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)}") _report_progress( progress_callback, "brave_candidates", f"Brave candidates: {len(brave_candidates)}", 84, candidate_count=len(brave_candidates), ) brave_count = len(brave_candidates) for candidate_index, candidate in enumerate(brave_candidates, start=1): repo_id = candidate.get("repo", "") branch = candidate.get("branch", "main") expected_file = candidate.get("filename", filename) _report_progress( progress_callback, "brave_repo", f"Checking Brave candidate {candidate_index}/{brave_count}", 84 + (candidate_index / max(1, brave_count)) * 6, repo=repo_id, candidate_index=candidate_index, candidate_count=brave_count, ) 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, headers=headers, ) if ( result and result.filename.lower() == filename.lower() and result_matches_requested_hash(result) ): _search_cache[cache_key] = result _report_progress( progress_callback, "found", "Found HuggingFace Brave match", 92, repo=repo_id, match_type=result.match_type, ) log.info( f"HuggingFace found source=brave_fallback match=exact file={filename} repo={repo_id} path={result.extra_value('path')}" ) return result # Not found _search_cache[cache_key] = None _report_progress( progress_callback, "done", "HuggingFace checked", 92, candidate_count=len(repos_to_check), author_repo_count=author_fallback_repo_count, brave_candidate_count=len(brave_candidates), ) 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}") _report_progress( progress_callback, "error", f"HuggingFace search error: {e}", 100, status="error", ) return None def search_huggingface( query: str, model_type: Optional[str] = None, limit: int = 10, token: Optional[str] = None, ) -> List[SearchResult]: """ Search HuggingFace Hub for models (general search). Returns repos that might be relevant, not guaranteed to have exact file. """ results = [] try: params = {"search": query, "limit": limit, "full": "true"} data = execute_provider_json_request( "HuggingFace general search", f"{HF_API_URL}/models", params=params, api_key=token, timeout=15, ) if data: models = data for model in models: repo_id = model.get("id", "") results.append( build_model_result( "huggingface", name=model.get("modelId", repo_id), url=f"https://huggingface.co/{repo_id}", repo=repo_id, downloads=model.get("downloads", 0), likes=model.get("likes", 0), ) ) except Exception as e: log.error(f"HuggingFace search error: {e}") return results def build_huggingface_custom_result( url: str, expected_filename: str = "", token: Optional[str] = None, ) -> Optional[SearchResult]: from urllib.parse import quote, unquote parsed = parse_huggingface_url(url) if not parsed: return None repo_id = parsed.repo branch = parsed.branch file_path = unquote(parsed.filename) if not repo_id or not file_path: return None filename = get_filename_from_path(file_path) download_url = get_huggingface_download_url(repo_id, file_path, branch) if not looks_like_model_file(download_url, expected_filename or filename): return None headers = {} if token: headers["Authorization"] = f"Bearer {token}" file_info: Dict[str, Any] = {} repo_tree = _get_repo_tree(repo_id, headers, branch) normalized_file_path = file_path.replace("\\", "/").strip("/") for candidate in repo_tree or []: candidate_path = str(candidate.get("path") or "").replace("\\", "/").strip("/") if candidate_path == normalized_file_path: file_info = candidate break size = resolve_file_size( file_info, [download_url], probe=lambda url: probe_remote_file_size(url, headers=headers, timeout=10), ) sha256 = _extract_huggingface_file_sha256(file_info) def quote_url_path(val): return quote(str(val or "").replace("\\", "/"), safe="/") page_url = ( f"https://huggingface.co/{repo_id}/blob/{branch}/{quote_url_path(file_path)}" ) return build_model_result( "huggingface", model_id=repo_id, version_id=None, name=repo_id, filename=filename, url=download_url, download_url=download_url, size=size, sha256=sha256, hashes={"SHA256": sha256} if sha256 else {}, match_type="custom_url", details_source="huggingface", repo_id=repo_id, repo=repo_id, path=file_path, branch=branch, page_url=page_url, version_url=page_url, custom_url=True, )