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:
Will Miao
2026-09-25 18:44:06 +08:00
parent 067e605e75
commit 8b7ba59263
23 changed files with 446 additions and 23 deletions
@@ -27,6 +27,7 @@
]),
'settings': dict({
'civitai_api_key_set': True,
'huggingface_api_key_set': False,
'language': 'en',
'llm_api_key_set': False,
'other_models_paths_available': False,
@@ -1148,3 +1148,54 @@ async def test_download_hydrates_the_card_from_the_site(tmp_path, monkeypatch):
assert scanner.update_single_model_cache.await_count == 1
cached = scanner.update_single_model_cache.await_args.args[2]
assert cached["model_name"] == "Krea-2-LORA"
@pytest.mark.asyncio
async def test_download_model_source_sends_hf_token_as_custom_headers(
tmp_path, monkeypatch
):
"""A gated/private HF repo needs the configured token on the download."""
captured = _stub_download_backend(monkeypatch)
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", AsyncMock())
monkeypatch.setattr(
"py.services.model_sources.huggingface._hf_token", lambda: "hf_secret"
)
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "huggingface",
"repo": "user/repo",
"filename": "f.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
assert captured["custom_headers"] == {"Authorization": "Bearer hf_secret"}
@pytest.mark.asyncio
async def test_download_model_source_sends_no_headers_without_hf_token(
tmp_path, monkeypatch
):
captured = _stub_download_backend(monkeypatch)
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", AsyncMock())
monkeypatch.setattr(
"py.services.model_sources.huggingface._hf_token", lambda: ""
)
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "huggingface",
"repo": "user/repo",
"filename": "f.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
assert captured["custom_headers"] is None
+61
View File
@@ -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]
+97
View File
@@ -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)