mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-26 13:34:08 -03:00
feat: support gated/private Hugging Face repos via access token
Add a huggingface_api_key setting (Settings UI, HF_TOKEN / HUGGING_FACE_HUB_TOKEN env override) and attach it as a Bearer token to Hugging Face file listing, model card fetching and downloads, so gated and private repositories can be downloaded once the user has accepted the repo terms. - fetch_json/fetch_text accept custom headers; ModelSource gains an auth_headers() hook so handlers stay platform-agnostic - 401/403 from the tree API now explain how to fix (configure token / accept gated terms) - aria2 pre-resolves huggingface.co redirects and strips credentials before handing the signed CDN URL to aria2, mirroring the CivitAI handling so the token never leaks to the CDN - settings API exposes huggingface_api_key_set only; the raw key joins _NO_SYNC_KEYS
This commit is contained in:
@@ -1281,3 +1281,64 @@ async def test_download_file_does_not_refresh_url_for_other_errors(
|
||||
assert "Download aborted" in result
|
||||
assert add_uri_count["n"] == 1
|
||||
assert downloader._transfers == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_file_preresolves_huggingface_redirect_and_strips_token(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
"""aria2 forwards custom headers to redirect targets, so the HF Bearer
|
||||
token must never leave huggingface.co: the /resolve/ redirect is resolved
|
||||
first and the signed CDN URL is handed to aria2 without headers."""
|
||||
downloader = Aria2Downloader()
|
||||
downloader._rpc_url = "http://127.0.0.1/jsonrpc"
|
||||
downloader._rpc_secret = "secret"
|
||||
|
||||
save_path = tmp_path / "downloads" / "model.safetensors"
|
||||
rpc_calls = []
|
||||
statuses = iter(
|
||||
[
|
||||
{
|
||||
"gid": "gid-1",
|
||||
"status": "complete",
|
||||
"completedLength": "10",
|
||||
"totalLength": "10",
|
||||
"downloadSpeed": "0",
|
||||
"files": [{"path": str(save_path)}],
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
async def fake_rpc_call(method, params, **_kwargs):
|
||||
rpc_calls.append((method, params))
|
||||
if method == "aria2.addUri":
|
||||
return "gid-1"
|
||||
if method == "aria2.tellStatus":
|
||||
return next(statuses)
|
||||
raise AssertionError(f"Unexpected RPC method: {method}")
|
||||
|
||||
monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
|
||||
monkeypatch.setattr(
|
||||
downloader,
|
||||
"_resolve_authenticated_redirect_url",
|
||||
AsyncMock(
|
||||
return_value="https://cdn-lfs.huggingface.co/signed/model.safetensors?sig=abc"
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
|
||||
monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
|
||||
|
||||
success, result = await downloader.download_file(
|
||||
"https://huggingface.co/user/repo/resolve/main/model.safetensors",
|
||||
str(save_path),
|
||||
download_id="download-1",
|
||||
headers={"Authorization": "Bearer hf_secret"},
|
||||
)
|
||||
|
||||
assert success is True
|
||||
assert result == str(save_path)
|
||||
assert rpc_calls[0][0] == "aria2.addUri"
|
||||
assert rpc_calls[0][1][0] == [
|
||||
"https://cdn-lfs.huggingface.co/signed/model.safetensors?sig=abc"
|
||||
]
|
||||
assert "header" not in rpc_calls[0][1][1]
|
||||
|
||||
@@ -1139,3 +1139,100 @@ class TestHashBasedVersionMatching:
|
||||
)
|
||||
|
||||
assert mock_ctx.call_args.kwargs["sha256"] == "c" * 64
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Hugging Face authentication (gated / private repositories)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestHuggingFaceAuth:
|
||||
def test_auth_headers_empty_without_token(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface._hf_token", lambda: ""
|
||||
)
|
||||
|
||||
assert HuggingFaceSource().auth_headers() == {}
|
||||
|
||||
def test_auth_headers_bearer_with_token(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface._hf_token", lambda: "hf_secret"
|
||||
)
|
||||
|
||||
assert HuggingFaceSource().auth_headers() == {
|
||||
"Authorization": "Bearer hf_secret"
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_files_sends_token_to_tree_api(self, monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
async def fake_fetch_json(url, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return 200, []
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface.fetch_json", fake_fetch_json
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface._hf_token", lambda: "hf_secret"
|
||||
)
|
||||
|
||||
await HuggingFaceSource().list_files("u/r")
|
||||
|
||||
assert captured["headers"] == {"Authorization": "Bearer hf_secret"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_card_sends_token(self, monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
async def fake_fetch_text(url, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return "# card"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface.fetch_text", fake_fetch_text
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface._hf_token", lambda: "hf_secret"
|
||||
)
|
||||
|
||||
await HuggingFaceSource().fetch_model_card("u/r")
|
||||
|
||||
assert captured["headers"] == {"Authorization": "Bearer hf_secret"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unauthorised_without_token_explains_how_to_fix(self, monkeypatch):
|
||||
async def fake_fetch_json(url, **_kwargs):
|
||||
return 401, None
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface.fetch_json", fake_fetch_json
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface._hf_token", lambda: ""
|
||||
)
|
||||
|
||||
with pytest.raises(ModelSourceError) as excinfo:
|
||||
await HuggingFaceSource().list_files("u/r")
|
||||
|
||||
assert excinfo.value.status == 401
|
||||
assert "access token" in str(excinfo.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_denied_with_token_points_at_repo_terms(self, monkeypatch):
|
||||
async def fake_fetch_json(url, **_kwargs):
|
||||
return 403, None
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface.fetch_json", fake_fetch_json
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.huggingface._hf_token", lambda: "hf_secret"
|
||||
)
|
||||
|
||||
with pytest.raises(ModelSourceError) as excinfo:
|
||||
await HuggingFaceSource().list_files("u/r")
|
||||
|
||||
assert excinfo.value.status == 403
|
||||
assert "accept its terms" in str(excinfo.value)
|
||||
|
||||
Reference in New Issue
Block a user