nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(huggingface): honor author index ttl during force search

Browse Source
This commit is contained in:
Dariusz L
2026-08-05 20:57:46 +02:00
parent f103b67b98
commit ff25791062
2 changed files with 229 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/sources/huggingface.py
+66 -32
View File
@@ -320,6 +320,7 @@ def _get_author_index(
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")
@@ -339,6 +340,13 @@ def _get_author_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
@@ -1051,12 +1059,6 @@ def search_huggingface_for_file(
)
return _search_cache[cache_key]
# Force Search may initialize a missing persistent index, but must not
# overwrite an index that already exists.
persist_author_index = not force_refresh or not os.path.isfile(
HF_AUTHOR_INDEX_CACHE_PATH
)
try:
def result_matches_requested_hash(result: Optional[Dict[str, Any]]) -> bool:
if not requested_sha256:
@@ -1074,6 +1076,51 @@ def search_huggingface_for_file(
log.info(f"HuggingFace start file={filename} exact={exact_only}")
def find_author_index_result(
author: str,
) -> tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
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))
@@ -1087,23 +1134,11 @@ def search_huggingface_for_file(
author_index=author_index,
author_count=total_authors,
)
index = _get_author_index(
author,
headers={},
force_refresh=force_refresh,
persist=persist_author_index,
)
index, result = find_author_index_result(author)
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(
@@ -1229,21 +1264,20 @@ def search_huggingface_for_file(
)
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 = _get_author_index(
author,
headers={},
force_refresh=force_refresh,
persist=persist_author_index,
)
index, result = find_author_index_result(author)
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):
_search_cache[cache_key] = result
tests/test_huggingface.py
+163 -41
View File
@@ -1,30 +1,33 @@
import unittest
import re
import json
import re
import tempfile
import time
import unittest
from pathlib import Path
from unittest.mock import patch
from core.matcher import build_filename_search_queries
from core.sources.huggingface import (
HF_AUTHOR_FALLBACKS,
_build_author_index_from_models,
_fetch_author_index,
_get_author_index,
_find_matching_file_in_author_index,
_get_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,
_read_persistent_author_indexes,
_write_persistent_author_index,
build_huggingface_custom_result,
clear_search_cache,
get_huggingface_download_url,
get_huggingface_model_details,
get_known_author_fallback_indexes_status,
parse_huggingface_url,
search_huggingface_for_file,
)
from core.matcher import build_filename_search_queries
from core.type_utils import extract_file_size
class HuggingFaceSourceTests(unittest.TestCase):
def test_known_author_fallbacks_include_comfy_org_and_kijai(self):
@@ -274,30 +277,33 @@ class HuggingFaceSourceTests(unittest.TestCase):
self.assertFalse(index_path.exists())
fetch_index.assert_called_once_with("Example", headers={})
def test_force_search_does_not_overwrite_existing_author_index(self):
sha256 = "d" * 64
index = {
"author": "Comfy-Org",
"updated_at": 123,
"repo_count": 1,
"file_count": 1,
"hash_count": 1,
"repos": ["Comfy-Org/example"],
"files": [
{
"repo_id": "Comfy-Org/example",
"path": "model.safetensors",
"filename": "model.safetensors",
"size": 123,
"sha256": sha256,
}
],
def test_force_search_uses_fresh_author_indexes_without_refreshing(self):
now = time.time()
indexes = {
author: {
"author": author,
"updated_at": now,
"repo_count": 1,
"file_count": 1,
"hash_count": 1,
"repos": [f"{author}/example"],
"files": [
{
"repo_id": f"{author}/example",
"path": "other.safetensors",
"filename": "other.safetensors",
"size": 123,
"sha256": "1" * 64,
}
],
}
for author in HF_AUTHOR_FALLBACKS
}
with tempfile.TemporaryDirectory() as tmpdir:
index_path = Path(tmpdir) / "huggingface-author-index.json"
index_path.write_text(
json.dumps({"version": 2, "authors": {"Comfy-Org": index}}),
json.dumps({"version": 2, "authors": indexes}),
encoding="utf-8",
)
with (
@@ -306,19 +312,20 @@ class HuggingFaceSourceTests(unittest.TestCase):
str(index_path),
),
patch(
"core.sources.huggingface._get_author_index",
return_value=index,
) as get_index,
"core.sources.huggingface._fetch_author_index"
) as fetch_index,
):
clear_search_cache()
result = search_huggingface_for_file(
"model.safetensors",
sha256=sha256,
"missing.safetensors",
sha256="2" * 64,
force_refresh=True,
use_api_search=False,
use_brave_fallback=False,
)
self.assertIsNotNone(result)
self.assertFalse(get_index.call_args.kwargs["persist"])
self.assertIsNone(result)
fetch_index.assert_not_called()
def test_force_search_can_initialize_missing_author_index(self):
sha256 = "e" * 64
@@ -348,19 +355,134 @@ class HuggingFaceSourceTests(unittest.TestCase):
str(index_path),
),
patch(
"core.sources.huggingface._get_author_index",
"core.sources.huggingface._fetch_author_index",
return_value=index,
) as get_index,
) as fetch_index,
):
clear_search_cache()
result = search_huggingface_for_file(
"model.safetensors",
sha256=sha256,
force_refresh=True,
use_api_search=False,
use_brave_fallback=False,
)
index_was_persisted = index_path.exists()
self.assertIsNotNone(result)
self.assertTrue(index_was_persisted)
fetch_index.assert_called_once_with("Comfy-Org", headers={})
def test_force_search_refreshes_stale_author_index_after_cache_miss(self):
sha256 = "f" * 64
stale_index = {
"author": "Comfy-Org",
"updated_at": 0,
"repo_count": 1,
"file_count": 1,
"hash_count": 0,
"repos": ["Comfy-Org/example"],
"files": [],
}
refreshed_index = {
**stale_index,
"updated_at": time.time(),
"hash_count": 1,
"files": [
{
"repo_id": "Comfy-Org/example",
"path": "model.safetensors",
"filename": "model.safetensors",
"size": 123,
"sha256": sha256,
}
],
}
with tempfile.TemporaryDirectory() as tmpdir:
index_path = Path(tmpdir) / "huggingface-author-index.json"
index_path.write_text(
json.dumps(
{
"version": 2,
"authors": {
"Comfy-Org": stale_index,
"Kijai": {**stale_index, "author": "Kijai"},
},
}
),
encoding="utf-8",
)
with (
patch(
"core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
str(index_path),
),
patch(
"core.sources.huggingface._fetch_author_index",
return_value=refreshed_index,
) as fetch_index,
):
clear_search_cache()
result = search_huggingface_for_file(
"renamed-model.safetensors",
sha256=sha256,
force_refresh=True,
use_api_search=False,
use_brave_fallback=False,
)
self.assertIsNotNone(result)
self.assertTrue(get_index.call_args.kwargs["persist"])
self.assertEqual("hash", result["match_type"])
fetch_index.assert_called_once_with("Comfy-Org", headers={})
def test_force_search_uses_stale_author_index_when_hash_matches(self):
sha256 = "7" * 64
stale_index = {
"author": "Comfy-Org",
"updated_at": 0,
"repo_count": 1,
"file_count": 1,
"hash_count": 1,
"repos": ["Comfy-Org/example"],
"files": [
{
"repo_id": "Comfy-Org/example",
"path": "model.safetensors",
"filename": "model.safetensors",
"size": 123,
"sha256": sha256,
}
],
}
with tempfile.TemporaryDirectory() as tmpdir:
index_path = Path(tmpdir) / "huggingface-author-index.json"
index_path.write_text(
json.dumps({"version": 2, "authors": {"Comfy-Org": stale_index}}),
encoding="utf-8",
)
with (
patch(
"core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
str(index_path),
),
patch(
"core.sources.huggingface._fetch_author_index"
) as fetch_index,
):
clear_search_cache()
result = search_huggingface_for_file(
"renamed-model.safetensors",
sha256=sha256,
force_refresh=True,
use_api_search=False,
use_brave_fallback=False,
)
self.assertIsNotNone(result)
self.assertEqual("hash", result["match_type"])
fetch_index.assert_not_called()
def test_explicit_author_index_refresh_still_persists(self):
fresh_index = {