Dariusz L
a8ba99800c
Support both Comfy-Org and Kijai author indexes as known HuggingFace fallbacks, aggregate their status in the backend, refresh the cached author index, and update the frontend UI and tests.
318 lines
12 KiB
Python
318 lines
12 KiB
Python
import unittest
|
|
import re
|
|
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
from core.sources.huggingface import (
|
|
HF_AUTHOR_FALLBACKS,
|
|
_fetch_author_index,
|
|
_is_author_index_fresh,
|
|
_read_persistent_author_indexes,
|
|
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,
|
|
_write_persistent_author_index,
|
|
)
|
|
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):
|
|
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": 1,
|
|
"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):
|
|
with patch(
|
|
"core.sources.huggingface.execute_provider_json_request",
|
|
return_value=[
|
|
{
|
|
"id": "Comfy-Org/example",
|
|
"siblings": [
|
|
{
|
|
"rfilename": "text_encoders/model.safetensors",
|
|
"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(
|
|
"text_encoders/model.safetensors",
|
|
index["files"][0]["path"],
|
|
)
|
|
|
|
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": 1,
|
|
"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_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_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)
|
|
|
|
# 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({}))
|
|
|
|
def test_normalize_huggingface_size_probe_url(self):
|
|
url = "https://huggingface.co/user/repo/blob/main/model.safetensors"
|
|
normalized = _normalize_huggingface_size_probe_url(url)
|
|
self.assertEqual(normalized, "https://huggingface.co/user/repo/resolve/main/model.safetensors")
|
|
|
|
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)
|
|
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):
|
|
with patch(
|
|
"core.sources.huggingface.fetch_remote_file_size_cached",
|
|
return_value=456,
|
|
):
|
|
result = build_huggingface_custom_result(url)
|
|
|
|
self.assertIsNotNone(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):
|
|
with 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:
|
|
with 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"])
|