From ff16bd869f492e974073445c3f4bdfe575bb0335 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Tue, 4 Aug 2026 06:21:08 +0200 Subject: [PATCH] Add SHA-256 hash matching support for HuggingFace author index Enable matching files in the HuggingFace author index by SHA-256 in addition to filename. Enrich the author index with hashes from repo tree responses, store them in the index metadata, and use them as a primary lookup path during search_huggingface_for_file before falling back to filename-based search. Also update the author index cache version to 2 and add tests for the new behavior. --- core/sources/huggingface.py | 179 ++++++++++++++++++++++++++++++++++-- tests/test_huggingface.py | 117 ++++++++++++++++++++--- 2 files changed, 274 insertions(+), 22 deletions(-) diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 2b38df4..4920bda 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -9,6 +9,7 @@ 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, urlparse @@ -36,7 +37,7 @@ 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 = 1 +HF_AUTHOR_INDEX_CACHE_VERSION = 2 HF_AUTHOR_INDEX_CACHE_PATH = os.path.join( METADATA_DIR, "huggingface-author-index.json" @@ -174,6 +175,7 @@ def _build_author_index_from_models( ) -> Dict[str, Any]: files = [] repo_ids = [] + hash_count = 0 for repo in repos: repo_id = repo.get("id", "") @@ -186,12 +188,17 @@ def _build_author_index_from_models( 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, } ) @@ -200,11 +207,90 @@ def _build_author_index_from_models( "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( @@ -218,8 +304,11 @@ def _fetch_author_index(author: str, headers: Dict[str, str], limit: int = 200): 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}: repos={index['repo_count']}, files={index['file_count']}" + 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: @@ -281,6 +370,7 @@ def get_author_fallback_index_status(author: str = "Comfy-Org") -> Dict[str, Any "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, } @@ -308,6 +398,7 @@ def get_known_author_fallback_indexes_status() -> Dict[str, Any]: 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), } @@ -733,8 +824,32 @@ def _find_matching_file_in_author_index( filename: str, exact_only: bool = False, headers: Optional[Dict[str, str]] = None, + sha256: Optional[str] = None, ) -> Optional[Dict[str, Any]]: 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, + ) + + return None + filename_lower = filename.lower() filename_base = os.path.splitext(filename_lower)[0] partial_match = None @@ -950,6 +1065,52 @@ def search_huggingface_for_file( log.info(f"HuggingFace start file={filename} exact={exact_only}") + 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 = _get_author_index( + author, + headers={}, + force_refresh=force_refresh, + ) + cached_author_indexes[author] = index + if not index: + continue + + result = _find_matching_file_in_author_index( + index, + filename, + exact_only=exact_only, + headers={}, + sha256=requested_sha256, + ) + if result: + _search_cache[cache_key] = result + _report_progress( + progress_callback, + "found", + "Found HuggingFace SHA-256 match", + 92, + repo=result.get("repo_id"), + match_type=result.get("match_type"), + ) + log.info( + f"HuggingFace found source=author_index " + f"match=hash file={result['filename']} " + f"repo={result['repo_id']} path={result['path']}" + ) + return result + repos_to_check = [] seen_repos = set() search_queries = build_filename_search_queries(filename) @@ -1056,17 +1217,21 @@ def search_huggingface_for_file( author_index=author_index, author_count=total_authors, ) - index = _get_author_index( - author, - headers={}, - force_refresh=force_refresh, - ) + if author in cached_author_indexes: + index = cached_author_indexes[author] + else: + 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, headers={}, + sha256=requested_sha256, ) author_fallback_repo_count += int(index.get("repo_count") or 0) if result and result_matches_requested_hash(result): diff --git a/tests/test_huggingface.py b/tests/test_huggingface.py index 8560d35..e456972 100644 --- a/tests/test_huggingface.py +++ b/tests/test_huggingface.py @@ -6,15 +6,19 @@ from pathlib import Path from unittest.mock import patch from core.sources.huggingface import ( HF_AUTHOR_FALLBACKS, + _build_author_index_from_models, _fetch_author_index, + _find_matching_file_in_author_index, _is_author_index_fresh, _read_persistent_author_indexes, + clear_search_cache, build_huggingface_custom_result, get_known_author_fallback_indexes_status, get_huggingface_model_details, parse_huggingface_url, get_huggingface_download_url, _normalize_huggingface_size_probe_url, + search_huggingface_for_file, _write_persistent_author_index, ) from core.matcher import build_filename_search_queries @@ -31,7 +35,7 @@ class HuggingFaceSourceTests(unittest.TestCase): index_path.write_text( json.dumps( { - "version": 1, + "version": 2, "authors": { "Comfy-Org": { "author": "Comfy-Org", @@ -68,19 +72,32 @@ class HuggingFaceSourceTests(unittest.TestCase): self.assertTrue(status["fully_cached"]) def test_fetch_author_index_builds_index_from_provider_repo_list(self): - with patch( - "core.sources.huggingface.execute_provider_json_request", - return_value=[ - { - "id": "Comfy-Org/example", - "siblings": [ - { - "rfilename": "text_encoders/model.safetensors", - "size": 123, - } - ], - } - ], + sha256 = "a" * 64 + with ( + patch( + "core.sources.huggingface.execute_provider_json_request", + return_value=[ + { + "id": "Comfy-Org/example", + "siblings": [ + { + "rfilename": "text_encoders/model.safetensors", + "size": 123, + } + ], + } + ], + ), + patch( + "core.sources.huggingface._get_repo_tree", + return_value=[ + { + "path": "text_encoders/model.safetensors", + "size": 123, + "lfs": {"oid": f"sha256:{sha256}", "size": 123}, + } + ], + ), ): index = _fetch_author_index("Comfy-Org", headers={}) @@ -88,10 +105,80 @@ class HuggingFaceSourceTests(unittest.TestCase): self.assertEqual("Comfy-Org", index["author"]) self.assertEqual(1, index["repo_count"]) self.assertEqual(1, index["file_count"]) + self.assertEqual(1, index["hash_count"]) self.assertEqual( "text_encoders/model.safetensors", index["files"][0]["path"], ) + self.assertEqual(sha256, index["files"][0]["sha256"]) + + def test_author_index_matches_sha256_directly(self): + sha256 = "b" * 64 + index = _build_author_index_from_models( + "Comfy-Org", + [ + { + "id": "Comfy-Org/example", + "siblings": [ + { + "rfilename": "models/example.safetensors", + "size": 456, + "lfs": {"oid": f"sha256:{sha256}", "size": 456}, + } + ], + } + ], + ) + + result = _find_matching_file_in_author_index( + index, + "different-name.safetensors", + sha256=sha256, + ) + + self.assertIsNotNone(result) + self.assertEqual("hash", result["match_type"]) + self.assertEqual(sha256, result["sha256"]) + self.assertEqual("models/example.safetensors", result["path"]) + + def test_hash_search_uses_author_index_before_filename_search(self): + sha256 = "c" * 64 + index = _build_author_index_from_models( + "Comfy-Org", + [ + { + "id": "Comfy-Org/example", + "siblings": [ + { + "rfilename": "models/example.safetensors", + "size": 789, + "lfs": {"oid": f"sha256:{sha256}", "size": 789}, + } + ], + } + ], + ) + + clear_search_cache() + with ( + patch( + "core.sources.huggingface._get_author_index", + return_value=index, + ) as get_index, + patch( + "core.sources.huggingface.execute_provider_json_request" + ) as request, + ): + result = search_huggingface_for_file( + "different-name.safetensors", + sha256=sha256, + ) + + self.assertIsNotNone(result) + self.assertEqual("hash", result["match_type"]) + self.assertEqual(sha256, result["sha256"]) + get_index.assert_called_once() + request.assert_not_called() def test_malformed_list_author_index_is_not_treated_as_fresh(self): self.assertFalse(_is_author_index_fresh([{"id": "Comfy-Org/example"}])) @@ -102,7 +189,7 @@ class HuggingFaceSourceTests(unittest.TestCase): index_path.write_text( json.dumps( { - "version": 1, + "version": 2, "authors": { "Comfy-Org": [{"id": "Comfy-Org/example"}], "Valid": {