Dariusz L
dca6bc93fb
Use order-insensitive token comparison in calculate_filename_confidence so variant names with reordered tokens (e.g., int8_convrot vs convrot_int8) score higher. Replace first-match partial matching in HuggingFace source with best-match selection based on filename confidence, and surface confidence in search results.
887 lines
31 KiB
Python
887 lines
31 KiB
Python
import json
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from core.contracts import SearchResult
|
|
from core.matcher import build_filename_search_queries
|
|
from core.sources.common import probe_remote_file_size
|
|
from core.sources.huggingface import (
|
|
HF_AUTHOR_FALLBACKS,
|
|
_build_author_index_from_models,
|
|
_fetch_author_index,
|
|
_find_matching_file_in_author_index,
|
|
_find_matching_file_in_repo,
|
|
_get_author_index,
|
|
_is_author_index_fresh,
|
|
_read_persistent_author_indexes,
|
|
_write_persistent_author_index,
|
|
build_huggingface_custom_result,
|
|
clear_search_cache,
|
|
get_huggingface_download_url,
|
|
get_huggingface_file_sha256,
|
|
get_huggingface_model_details,
|
|
get_known_author_fallback_indexes_status,
|
|
parse_huggingface_url,
|
|
search_huggingface,
|
|
search_huggingface_for_file,
|
|
)
|
|
from core.type_utils import extract_file_size
|
|
|
|
|
|
def _result_dict(result):
|
|
assert isinstance(result, SearchResult)
|
|
return result.to_dict()
|
|
|
|
|
|
class HuggingFaceSourceTests(unittest.TestCase):
|
|
|
|
@patch("core.sources.huggingface.execute_provider_json_request")
|
|
def test_general_search_returns_typed_results(self, request_json):
|
|
request_json.return_value = [
|
|
{
|
|
"id": "Comfy-Org/example",
|
|
"modelId": "Comfy-Org/example",
|
|
"downloads": 12,
|
|
"likes": 3,
|
|
}
|
|
]
|
|
|
|
results = search_huggingface("example", limit=1)
|
|
|
|
self.assertEqual(1, len(results))
|
|
self.assertIsInstance(results[0], SearchResult)
|
|
self.assertEqual("Comfy-Org/example", results[0].name)
|
|
self.assertEqual("Comfy-Org/example", results[0].extra_value("repo"))
|
|
self.assertEqual(12, results[0].extra_value("downloads"))
|
|
self.assertEqual("https://huggingface.co/Comfy-Org/example", results[0].url)
|
|
|
|
def test_known_author_fallbacks_include_comfy_org_and_kijai(self):
|
|
self.assertEqual(["Comfy-Org", "Kijai"], HF_AUTHOR_FALLBACKS)
|
|
|
|
def test_known_author_status_reports_each_author_file_count(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
index_path = Path(tmpdir) / "huggingface-author-index.json"
|
|
index_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"version": 2,
|
|
"authors": {
|
|
"Comfy-Org": {
|
|
"author": "Comfy-Org",
|
|
"updated_at": 1,
|
|
"repo_count": 2,
|
|
"file_count": 3,
|
|
"repos": [],
|
|
"files": [],
|
|
},
|
|
"Kijai": {
|
|
"author": "Kijai",
|
|
"updated_at": 2,
|
|
"repo_count": 4,
|
|
"file_count": 5,
|
|
"repos": [],
|
|
"files": [],
|
|
},
|
|
},
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
with patch(
|
|
"core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
|
|
str(index_path),
|
|
):
|
|
status = get_known_author_fallback_indexes_status()
|
|
|
|
self.assertEqual(
|
|
{"Comfy-Org": 3, "Kijai": 5},
|
|
{item["author"]: item["file_count"] for item in status["authors"]},
|
|
)
|
|
self.assertEqual(8, status["file_count"])
|
|
self.assertTrue(status["fully_cached"])
|
|
|
|
def test_fetch_author_index_builds_index_from_provider_repo_list(self):
|
|
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={})
|
|
|
|
self.assertIsInstance(index, dict)
|
|
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)
|
|
result = _result_dict(result)
|
|
self.assertEqual("hash", result["match_type"])
|
|
self.assertEqual(sha256, result["sha256"])
|
|
self.assertEqual("models/example.safetensors", result["path"])
|
|
self.assertEqual(100.0, result["confidence"])
|
|
|
|
def test_huggingface_matches_reordered_variant_with_confidence(self):
|
|
files = [
|
|
{
|
|
"path": "checkpoints/yue2_3b_bf16.safetensors",
|
|
"size": 100,
|
|
},
|
|
{
|
|
"path": "checkpoints/yue2_3b_int8_convrot.safetensors",
|
|
"size": 100,
|
|
},
|
|
]
|
|
|
|
repo_result = _find_matching_file_in_repo(
|
|
"Comfy-Org/YuE2",
|
|
files,
|
|
"yue2_convrot_int8.safetensors",
|
|
)
|
|
index_result = _find_matching_file_in_author_index(
|
|
{
|
|
"files": [
|
|
{
|
|
"repo_id": "Comfy-Org/YuE2",
|
|
"path": file["path"],
|
|
"size": file["size"],
|
|
}
|
|
for file in files
|
|
]
|
|
},
|
|
"yue2_convrot_int8.safetensors",
|
|
)
|
|
|
|
for result in (repo_result, index_result):
|
|
self.assertIsNotNone(result)
|
|
result = _result_dict(result)
|
|
self.assertEqual(
|
|
"checkpoints/yue2_3b_int8_convrot.safetensors", result["path"]
|
|
)
|
|
self.assertEqual("partial", result["match_type"])
|
|
self.assertGreater(result["confidence"], 0.0)
|
|
self.assertLess(result["confidence"], 100.0)
|
|
|
|
def test_huggingface_variant_match_stays_strict_for_exact_only(self):
|
|
result = _find_matching_file_in_repo(
|
|
"Comfy-Org/YuE2",
|
|
[
|
|
{
|
|
"path": "checkpoints/yue2_3b_int8_convrot.safetensors",
|
|
"size": 100,
|
|
}
|
|
],
|
|
"yue2_convrot_int8.safetensors",
|
|
exact_only=True,
|
|
)
|
|
|
|
self.assertIsNone(result)
|
|
|
|
def test_huggingface_variant_match_keeps_precision_alternatives(self):
|
|
result = _find_matching_file_in_author_index(
|
|
{
|
|
"files": [
|
|
{
|
|
"repo_id": "Comfy-Org/YuE2",
|
|
"path": "checkpoints/yue2_3b_bf16.safetensors",
|
|
"size": 100,
|
|
}
|
|
]
|
|
},
|
|
"yue2_convrot_int8.safetensors",
|
|
)
|
|
|
|
self.assertIsNotNone(result)
|
|
result = _result_dict(result)
|
|
self.assertEqual(
|
|
"checkpoints/yue2_3b_bf16.safetensors", result["path"]
|
|
)
|
|
self.assertGreater(result["confidence"], 0.0)
|
|
|
|
def test_file_search_returns_variant_match_from_author_index(self):
|
|
index = {
|
|
"files": [
|
|
{
|
|
"repo_id": "Comfy-Org/YuE2",
|
|
"path": "checkpoints/yue2_3b_int8_convrot.safetensors",
|
|
"size": 100,
|
|
}
|
|
]
|
|
}
|
|
|
|
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(
|
|
"yue2_convrot_int8.safetensors",
|
|
use_api_search=False,
|
|
use_brave_fallback=False,
|
|
)
|
|
|
|
self.assertIsNotNone(result)
|
|
result = _result_dict(result)
|
|
self.assertEqual("Comfy-Org/YuE2", result["repo_id"])
|
|
self.assertEqual(
|
|
"checkpoints/yue2_3b_int8_convrot.safetensors", result["path"]
|
|
)
|
|
self.assertGreater(result["confidence"], 0.0)
|
|
get_index.assert_called_once()
|
|
request.assert_not_called()
|
|
|
|
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)
|
|
result = _result_dict(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"}]))
|
|
|
|
def test_persistent_cache_discards_raw_provider_list_entries(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
index_path = Path(tmpdir) / "huggingface-author-index.json"
|
|
index_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"version": 2,
|
|
"authors": {
|
|
"Comfy-Org": [{"id": "Comfy-Org/example"}],
|
|
"Valid": {
|
|
"author": "Valid",
|
|
"updated_at": 1,
|
|
"repo_count": 0,
|
|
"file_count": 0,
|
|
"repos": [],
|
|
"files": [],
|
|
},
|
|
},
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
with patch(
|
|
"core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
|
|
str(index_path),
|
|
):
|
|
data = _read_persistent_author_indexes()
|
|
|
|
self.assertNotIn("Comfy-Org", data["authors"])
|
|
self.assertIn("Valid", data["authors"])
|
|
|
|
def test_persistent_author_index_is_pretty_printed(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
index_path = Path(tmpdir) / "huggingface-author-index.json"
|
|
with patch(
|
|
"core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
|
|
str(index_path),
|
|
):
|
|
_write_persistent_author_index(
|
|
"Example",
|
|
{
|
|
"author": "Example",
|
|
"updated_at": 1,
|
|
"repo_count": 0,
|
|
"file_count": 0,
|
|
"repos": [],
|
|
"files": [],
|
|
},
|
|
)
|
|
|
|
content = index_path.read_text(encoding="utf-8")
|
|
self.assertIn('\n "version":', content)
|
|
self.assertIn('\n "authors": {', content)
|
|
self.assertTrue(content.endswith("\n"))
|
|
self.assertEqual("Example", json.loads(content)["authors"]["Example"]["author"])
|
|
|
|
def test_non_persistent_author_index_refresh_does_not_persist(self):
|
|
fresh_index = {
|
|
"author": "Example",
|
|
"updated_at": 123,
|
|
"repo_count": 1,
|
|
"file_count": 1,
|
|
"hash_count": 1,
|
|
"repos": ["Example/repo"],
|
|
"files": [],
|
|
}
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
index_path = Path(tmpdir) / "huggingface-author-index.json"
|
|
with (
|
|
patch(
|
|
"core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
|
|
str(index_path),
|
|
),
|
|
patch(
|
|
"core.sources.huggingface._fetch_author_index",
|
|
return_value=fresh_index,
|
|
) as fetch_index,
|
|
):
|
|
clear_search_cache()
|
|
result = _get_author_index(
|
|
"Example",
|
|
headers={},
|
|
force_refresh=True,
|
|
persist=False,
|
|
)
|
|
|
|
self.assertIs(fresh_index, result)
|
|
self.assertFalse(index_path.exists())
|
|
fetch_index.assert_called_once_with("Example", headers={})
|
|
|
|
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": indexes}),
|
|
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(
|
|
"missing.safetensors",
|
|
sha256="2" * 64,
|
|
force_refresh=True,
|
|
use_api_search=False,
|
|
use_brave_fallback=False,
|
|
)
|
|
|
|
self.assertIsNone(result)
|
|
fetch_index.assert_not_called()
|
|
|
|
def test_force_search_can_initialize_missing_author_index(self):
|
|
sha256 = "e" * 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,
|
|
}
|
|
],
|
|
}
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
index_path = Path(tmpdir) / "huggingface-author-index.json"
|
|
with (
|
|
patch(
|
|
"core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
|
|
str(index_path),
|
|
),
|
|
patch(
|
|
"core.sources.huggingface._fetch_author_index",
|
|
return_value=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)
|
|
result = _result_dict(result)
|
|
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)
|
|
result = _result_dict(result)
|
|
self.assertEqual("hash", result["match_type"])
|
|
fetch_index.assert_not_called()
|
|
|
|
def test_explicit_author_index_refresh_still_persists(self):
|
|
fresh_index = {
|
|
"author": "Example",
|
|
"updated_at": 123,
|
|
"repo_count": 1,
|
|
"file_count": 1,
|
|
"hash_count": 1,
|
|
"repos": ["Example/repo"],
|
|
"files": [],
|
|
}
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
index_path = Path(tmpdir) / "huggingface-author-index.json"
|
|
with (
|
|
patch(
|
|
"core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
|
|
str(index_path),
|
|
),
|
|
patch(
|
|
"core.sources.huggingface._fetch_author_index",
|
|
return_value=fresh_index,
|
|
),
|
|
):
|
|
clear_search_cache()
|
|
result = _get_author_index(
|
|
"Example",
|
|
headers={},
|
|
force_refresh=True,
|
|
)
|
|
|
|
self.assertIs(fresh_index, result)
|
|
self.assertTrue(index_path.exists())
|
|
|
|
def test_parse_huggingface_url_valid_http(self):
|
|
url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors"
|
|
result = parse_huggingface_url(url)
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.repo, "runwayml/stable-diffusion-v1-5")
|
|
self.assertEqual(result.branch, "main")
|
|
self.assertEqual(result.filename, "v1-5-pruned-emaonly.safetensors")
|
|
|
|
def test_parse_huggingface_url_valid_hf_protocol(self):
|
|
url = "hf://stabilityai/stable-diffusion-xl-base-1.0/sd_xl_base_1.0.safetensors"
|
|
result = parse_huggingface_url(url)
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.repo, "stabilityai/stable-diffusion-xl-base-1.0")
|
|
self.assertEqual(result.filename, "sd_xl_base_1.0.safetensors")
|
|
|
|
def test_parse_huggingface_url_invalid(self):
|
|
self.assertIsNone(parse_huggingface_url("https://example.com/not-hf"))
|
|
self.assertIsNone(parse_huggingface_url("hf://not_enough_slashes"))
|
|
|
|
def test_get_huggingface_file_sha256_matches_repo_path(self):
|
|
sha256 = "a" * 64
|
|
repo_tree = [
|
|
{
|
|
"path": "vae/minimax_h3_audio_vae_fp32.safetensors",
|
|
"lfs": {"oid": f"sha256:{sha256}"},
|
|
}
|
|
]
|
|
with patch(
|
|
"core.sources.huggingface._get_repo_tree",
|
|
return_value=repo_tree,
|
|
) as get_repo_tree:
|
|
result = get_huggingface_file_sha256(
|
|
"https://huggingface.co/Comfy-Org/MiniMax-H3/resolve/main/"
|
|
"vae/minimax_h3_audio_vae_fp32.safetensors",
|
|
headers={"Authorization": "Bearer test"},
|
|
)
|
|
|
|
self.assertEqual(sha256, result)
|
|
get_repo_tree.assert_called_once_with(
|
|
"Comfy-Org/MiniMax-H3",
|
|
headers={"Authorization": "Bearer test"},
|
|
branch="main",
|
|
)
|
|
|
|
def test_get_huggingface_download_url(self):
|
|
repo = "runwayml/stable-diffusion-v1-5"
|
|
filename = "v1-5-pruned-emaonly.safetensors"
|
|
url = get_huggingface_download_url(repo, filename)
|
|
self.assertEqual(
|
|
url,
|
|
"https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors"
|
|
)
|
|
|
|
def test_extract_file_size(self):
|
|
# Directly nested key
|
|
self.assertEqual(extract_file_size({"size": 100}), 100)
|
|
self.assertEqual(extract_file_size({"sizeBytes": "200"}), 200)
|
|
self.assertEqual(extract_file_size({"sizeKB": 2}), 2048)
|
|
self.assertEqual(extract_file_size({"sizeKB": 0}), 0)
|
|
|
|
# LFS nested key
|
|
self.assertEqual(extract_file_size({"lfs": {"size": 300}}), 300)
|
|
|
|
# None cases
|
|
self.assertIsNone(extract_file_size(None))
|
|
self.assertIsNone(extract_file_size({}))
|
|
|
|
@patch(
|
|
"core.sources.common.fetch_remote_file_size_cached",
|
|
return_value=None,
|
|
)
|
|
def test_remote_size_probe_normalizes_huggingface_blob_url(self, mock_fetch):
|
|
url = "https://huggingface.co/user/repo/blob/main/model.safetensors"
|
|
self.assertIsNone(probe_remote_file_size(url))
|
|
mock_fetch.assert_called_once_with(
|
|
"https://huggingface.co/user/repo/resolve/main/model.safetensors",
|
|
headers=None,
|
|
timeout=15,
|
|
)
|
|
|
|
@patch(
|
|
"core.sources.common.fetch_remote_file_size_cached",
|
|
return_value=789,
|
|
)
|
|
def test_remote_size_probe_preserves_headers_and_timeout(self, mock_fetch):
|
|
url = "https://huggingface.co/user/repo/blob/main/model.safetensors"
|
|
headers = {"Authorization": "Bearer test"}
|
|
|
|
size = probe_remote_file_size(url, headers=headers, timeout=7)
|
|
|
|
self.assertEqual(size, 789)
|
|
mock_fetch.assert_called_once_with(
|
|
"https://huggingface.co/user/repo/resolve/main/model.safetensors",
|
|
headers=headers,
|
|
timeout=7,
|
|
)
|
|
|
|
def test_custom_result_uses_sha256_from_exact_huggingface_file(self):
|
|
sha256 = "a" * 64
|
|
url = (
|
|
"https://huggingface.co/DreamFast/Qwen3-VL-4b-Heretic-ComfyUI/"
|
|
"blob/main/qwen3-vl-4b-heretic_int8.safetensors"
|
|
)
|
|
with patch(
|
|
"core.sources.huggingface._get_repo_tree",
|
|
return_value=[
|
|
{
|
|
"path": "qwen3-vl-4b-heretic_int8.safetensors",
|
|
"size": 123,
|
|
"lfs": {"oid": f"sha256:{sha256}", "size": 123},
|
|
}
|
|
],
|
|
):
|
|
result = build_huggingface_custom_result(url)
|
|
|
|
self.assertIsNotNone(result)
|
|
result = _result_dict(result)
|
|
self.assertEqual(sha256, result["sha256"])
|
|
self.assertEqual({"SHA256": sha256}, result["hashes"])
|
|
self.assertEqual(123, result["size"])
|
|
self.assertEqual("huggingface", result["source"])
|
|
self.assertTrue(result["custom_url"])
|
|
|
|
def test_custom_result_leaves_hash_empty_when_huggingface_does_not_provide_it(self):
|
|
url = "https://huggingface.co/user/repo/blob/main/model.safetensors"
|
|
with (
|
|
patch("core.sources.huggingface._get_repo_tree", return_value=None),
|
|
patch(
|
|
"core.sources.common.fetch_remote_file_size_cached",
|
|
return_value=456,
|
|
),
|
|
):
|
|
result = build_huggingface_custom_result(url)
|
|
|
|
self.assertIsNotNone(result)
|
|
result = _result_dict(result)
|
|
self.assertEqual("", result["sha256"])
|
|
self.assertEqual({}, result["hashes"])
|
|
self.assertEqual(456, result["size"])
|
|
|
|
def test_build_huggingface_search_queries(self):
|
|
queries = build_filename_search_queries("some_model_bf16.safetensors")
|
|
self.assertIn("some_model", queries)
|
|
self.assertIn("some_model_bf16", queries)
|
|
|
|
def test_model_details_lists_only_variants_from_matched_folder(self):
|
|
sha256 = "b" * 64
|
|
tree = [
|
|
{
|
|
"path": "diffusion_models/model_fp8_scaled.safetensors",
|
|
"size": 100,
|
|
"lfs": {"oid": f"sha256:{sha256}", "size": 100},
|
|
},
|
|
{
|
|
"path": "diffusion_models/model_int8_convrot.safetensors",
|
|
"size": 80,
|
|
},
|
|
{
|
|
"path": "text_encoders/encoder_fp8.safetensors",
|
|
"size": 50,
|
|
},
|
|
{"path": "diffusion_models/README.md", "size": 10},
|
|
]
|
|
with (
|
|
patch("core.sources.huggingface._get_repo_tree", return_value=tree),
|
|
patch(
|
|
"core.sources.huggingface.execute_provider_json_request",
|
|
return_value={
|
|
"modelId": "Comfy-Org/example",
|
|
"author": "Comfy-Org",
|
|
"downloads": 123,
|
|
"likes": 4,
|
|
"tags": ["comfyui"],
|
|
},
|
|
),
|
|
):
|
|
details = get_huggingface_model_details(
|
|
"Comfy-Org/example",
|
|
"diffusion_models/model_fp8_scaled.safetensors",
|
|
)
|
|
|
|
self.assertIsNotNone(details)
|
|
self.assertEqual("huggingface", details["source"])
|
|
self.assertEqual("diffusion_models", details["folder"])
|
|
files = details["selected_version"]["files"]
|
|
self.assertEqual(
|
|
[
|
|
"model_fp8_scaled.safetensors",
|
|
"model_int8_convrot.safetensors",
|
|
],
|
|
[file_info["name"] for file_info in files],
|
|
)
|
|
self.assertTrue(files[0]["primary"])
|
|
self.assertEqual("FP8 scaled", files[0]["metadata"]["fp"])
|
|
self.assertEqual(sha256, files[0]["sha256"])
|
|
self.assertEqual("INT8 convrot", files[1]["metadata"]["fp"])
|
|
|
|
def test_model_details_passes_huggingface_token_to_tree_and_metadata(self):
|
|
tree = [{"path": "model.safetensors", "size": 100}]
|
|
with (
|
|
patch(
|
|
"core.sources.huggingface._get_repo_tree",
|
|
return_value=tree,
|
|
) as mock_tree,
|
|
patch(
|
|
"core.sources.huggingface.execute_provider_json_request",
|
|
return_value={},
|
|
) as mock_request,
|
|
):
|
|
details = get_huggingface_model_details(
|
|
"private/repo",
|
|
"model.safetensors",
|
|
token="hf_secret",
|
|
)
|
|
|
|
self.assertIsNotNone(details)
|
|
expected_headers = {"Authorization": "Bearer hf_secret"}
|
|
mock_tree.assert_called_once_with(
|
|
"private/repo",
|
|
headers=expected_headers,
|
|
branch="main",
|
|
)
|
|
self.assertEqual(expected_headers, mock_request.call_args.kwargs["headers"])
|