From 8b7ba59263526114f0f67a443bd13996c3a16d30 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Fri, 25 Sep 2026 18:44:06 +0800 Subject: [PATCH] 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 --- locales/de.json | 6 ++ locales/en.json | 6 ++ locales/es.json | 6 ++ locales/fr.json | 6 ++ locales/he.json | 6 ++ locales/ja.json | 6 ++ locales/ko.json | 6 ++ locales/ru.json | 6 ++ locales/zh-CN.json | 6 ++ locales/zh-TW.json | 6 ++ py/routes/handlers/misc_handlers.py | 3 + py/routes/handlers/model_source_handlers.py | 6 ++ py/services/aria2_downloader.py | 22 +++-- py/services/model_sources/base.py | 25 ++++- py/services/model_sources/huggingface.py | 39 +++++++- py/services/settings_manager.py | 10 ++ static/js/managers/SettingsManager.js | 55 +++++++++-- static/js/state/index.js | 2 + .../components/modals/settings/general.html | 37 +++++++ .../__snapshots__/test_api_snapshots.ambr | 1 + tests/routes/test_model_source_handlers.py | 51 ++++++++++ tests/services/test_aria2_downloader.py | 61 ++++++++++++ tests/services/test_model_sources.py | 97 +++++++++++++++++++ 23 files changed, 446 insertions(+), 23 deletions(-) diff --git a/locales/de.json b/locales/de.json index 17e01e22..38c89488 100644 --- a/locales/de.json +++ b/locales/de.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "Konfiguriert", "civitaiApiKeyNotConfigured": "Nicht konfiguriert", "civitaiApiKeySet": "Einrichten", + "huggingfaceApiKey": "[TODO: Translate] Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "[TODO: Translate] Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "[TODO: Translate] Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "[TODO: Translate] Configured", + "huggingfaceApiKeyNotConfigured": "[TODO: Translate] Not configured", + "huggingfaceApiKeySet": "[TODO: Translate] Set up", "civitaiHost": { "label": "CivitAI-Host", "help": "Wählen Sie aus, welche CivitAI-Seite geöffnet wird, wenn Sie „View on CivitAI“-Links verwenden.", diff --git a/locales/en.json b/locales/en.json index 9be1f8c5..de13e421 100644 --- a/locales/en.json +++ b/locales/en.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "Configured", "civitaiApiKeyNotConfigured": "Not configured", "civitaiApiKeySet": "Set up", + "huggingfaceApiKey": "Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "Configured", + "huggingfaceApiKeyNotConfigured": "Not configured", + "huggingfaceApiKeySet": "Set up", "civitaiHost": { "label": "CivitAI host", "help": "Choose which CivitAI site opens when using View on CivitAI links.", diff --git a/locales/es.json b/locales/es.json index 3d1cc465..dc574770 100644 --- a/locales/es.json +++ b/locales/es.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "Configurado", "civitaiApiKeyNotConfigured": "No configurado", "civitaiApiKeySet": "Configurar", + "huggingfaceApiKey": "[TODO: Translate] Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "[TODO: Translate] Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "[TODO: Translate] Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "[TODO: Translate] Configured", + "huggingfaceApiKeyNotConfigured": "[TODO: Translate] Not configured", + "huggingfaceApiKeySet": "[TODO: Translate] Set up", "civitaiHost": { "label": "Host de CivitAI", "help": "Elige qué sitio de CivitAI se abre al usar los enlaces de \"View on CivitAI\".", diff --git a/locales/fr.json b/locales/fr.json index c67ee45d..7c4c8e22 100644 --- a/locales/fr.json +++ b/locales/fr.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "Configuré", "civitaiApiKeyNotConfigured": "Non configuré", "civitaiApiKeySet": "Configurer", + "huggingfaceApiKey": "[TODO: Translate] Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "[TODO: Translate] Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "[TODO: Translate] Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "[TODO: Translate] Configured", + "huggingfaceApiKeyNotConfigured": "[TODO: Translate] Not configured", + "huggingfaceApiKeySet": "[TODO: Translate] Set up", "civitaiHost": { "label": "Hôte CivitAI", "help": "Choisissez quel site CivitAI s'ouvre lorsque vous utilisez les liens « View on CivitAI ».", diff --git a/locales/he.json b/locales/he.json index a2119f1d..43e1d3d8 100644 --- a/locales/he.json +++ b/locales/he.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "מוגדר", "civitaiApiKeyNotConfigured": "לא מוגדר", "civitaiApiKeySet": "הגדר", + "huggingfaceApiKey": "[TODO: Translate] Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "[TODO: Translate] Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "[TODO: Translate] Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "[TODO: Translate] Configured", + "huggingfaceApiKeyNotConfigured": "[TODO: Translate] Not configured", + "huggingfaceApiKeySet": "[TODO: Translate] Set up", "civitaiHost": { "label": "מארח CivitAI", "help": "בחר איזה אתר של CivitAI ייפתח בעת שימוש בקישורי \"View on CivitAI\".", diff --git a/locales/ja.json b/locales/ja.json index 84f0a8e1..f80860da 100644 --- a/locales/ja.json +++ b/locales/ja.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "設定済み", "civitaiApiKeyNotConfigured": "未設定", "civitaiApiKeySet": "設定", + "huggingfaceApiKey": "[TODO: Translate] Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "[TODO: Translate] Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "[TODO: Translate] Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "[TODO: Translate] Configured", + "huggingfaceApiKeyNotConfigured": "[TODO: Translate] Not configured", + "huggingfaceApiKeySet": "[TODO: Translate] Set up", "civitaiHost": { "label": "CivitAI ホスト", "help": "「View on CivitAI」リンクを使うときに開く CivitAI サイトを選択します。", diff --git a/locales/ko.json b/locales/ko.json index 8225e7a1..61ca7716 100644 --- a/locales/ko.json +++ b/locales/ko.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "설정됨", "civitaiApiKeyNotConfigured": "설정되지 않음", "civitaiApiKeySet": "설정", + "huggingfaceApiKey": "[TODO: Translate] Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "[TODO: Translate] Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "[TODO: Translate] Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "[TODO: Translate] Configured", + "huggingfaceApiKeyNotConfigured": "[TODO: Translate] Not configured", + "huggingfaceApiKeySet": "[TODO: Translate] Set up", "civitaiHost": { "label": "CivitAI 호스트", "help": "\"View on CivitAI\" 링크를 사용할 때 어떤 CivitAI 사이트를 열지 선택합니다.", diff --git a/locales/ru.json b/locales/ru.json index bf39642c..309e913a 100644 --- a/locales/ru.json +++ b/locales/ru.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "Настроен", "civitaiApiKeyNotConfigured": "Не настроен", "civitaiApiKeySet": "Настроить", + "huggingfaceApiKey": "[TODO: Translate] Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "[TODO: Translate] Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "[TODO: Translate] Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "[TODO: Translate] Configured", + "huggingfaceApiKeyNotConfigured": "[TODO: Translate] Not configured", + "huggingfaceApiKeySet": "[TODO: Translate] Set up", "civitaiHost": { "label": "Хост CivitAI", "help": "Выберите, какой сайт CivitAI будет открываться при использовании ссылок «View on CivitAI».", diff --git a/locales/zh-CN.json b/locales/zh-CN.json index accc51a1..82b74e00 100644 --- a/locales/zh-CN.json +++ b/locales/zh-CN.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "已配置", "civitaiApiKeyNotConfigured": "未配置", "civitaiApiKeySet": "设置", + "huggingfaceApiKey": "[TODO: Translate] Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "[TODO: Translate] Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "[TODO: Translate] Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "[TODO: Translate] Configured", + "huggingfaceApiKeyNotConfigured": "[TODO: Translate] Not configured", + "huggingfaceApiKeySet": "[TODO: Translate] Set up", "civitaiHost": { "label": "CivitAI 站点", "help": "选择使用“在 CivitAI 中查看”时默认打开的 CivitAI 站点。", diff --git a/locales/zh-TW.json b/locales/zh-TW.json index c5b2557f..643fa965 100644 --- a/locales/zh-TW.json +++ b/locales/zh-TW.json @@ -325,6 +325,12 @@ "civitaiApiKeyConfigured": "已設定", "civitaiApiKeyNotConfigured": "未設定", "civitaiApiKeySet": "設定", + "huggingfaceApiKey": "[TODO: Translate] Hugging Face Access Token", + "huggingfaceApiKeyPlaceholder": "[TODO: Translate] Enter your Hugging Face access token", + "huggingfaceApiKeyHelp": "[TODO: Translate] Required to download from gated or private Hugging Face repositories. Create a read-only token at huggingface.co/settings/tokens, and accept the repository's terms on its page first.", + "huggingfaceApiKeyConfigured": "[TODO: Translate] Configured", + "huggingfaceApiKeyNotConfigured": "[TODO: Translate] Not configured", + "huggingfaceApiKeySet": "[TODO: Translate] Set up", "civitaiHost": { "label": "CivitAI 站點", "help": "選擇使用「在 CivitAI 中查看」時預設開啟的 CivitAI 站點。", diff --git a/py/routes/handlers/misc_handlers.py b/py/routes/handlers/misc_handlers.py index 8faa37f9..f3021935 100644 --- a/py/routes/handlers/misc_handlers.py +++ b/py/routes/handlers/misc_handlers.py @@ -1506,6 +1506,7 @@ class SettingsHandler: # Sensitive — never expose the actual value to the frontend; # frontend receives a boolean instead (*_set). "civitai_api_key", + "huggingface_api_key", "llm_api_key", } ) @@ -1564,6 +1565,8 @@ class SettingsHandler: # Sensitive fields: only expose a boolean indicating whether set raw_key = self._settings.get("civitai_api_key") response_data["civitai_api_key_set"] = bool(raw_key) + raw_hf_key = self._settings.get("huggingface_api_key") + response_data["huggingface_api_key_set"] = bool(raw_hf_key) raw_llm_key = self._settings.get("llm_api_key") response_data["llm_api_key_set"] = bool(raw_llm_key) # Derived capability flag (not persisted): whether the host exposes diff --git a/py/routes/handlers/model_source_handlers.py b/py/routes/handlers/model_source_handlers.py index a5fdb5f7..5bca9ed6 100644 --- a/py/routes/handlers/model_source_handlers.py +++ b/py/routes/handlers/model_source_handlers.py @@ -580,6 +580,10 @@ class ModelSourceHandler: get_settings_manager().get("download_backend", "default") ) + # Site-specific credentials (e.g. a Hugging Face access token for + # gated/private repositories); empty for anonymous downloads. + auth_headers = source.auth_headers() + if download_backend == "aria2": aria2 = await Aria2Downloader.get_instance() aid = download_id or f"{source.platform}_{repo}_{filename}" @@ -589,6 +593,7 @@ class ModelSourceHandler: save_path=dest_path, download_id=aid, progress_callback=progress_callback, + headers=auth_headers or None, ) if ok: await _save_source_metadata( @@ -618,6 +623,7 @@ class ModelSourceHandler: use_auth=False, allow_resume=True, progress_callback=progress_callback, + custom_headers=auth_headers or None, ) if success: await _save_source_metadata( diff --git a/py/services/aria2_downloader.py b/py/services/aria2_downloader.py index 88e5eb42..435ca660 100644 --- a/py/services/aria2_downloader.py +++ b/py/services/aria2_downloader.py @@ -81,6 +81,14 @@ CIVITAI_DOWNLOAD_URL_PREFIXES = ( "https://civitai.red/api/download/", ) +#: Hosts whose authenticated downloads redirect to a signed CDN URL. aria2 +#: forwards custom headers to redirect targets, so for these hosts the +#: redirect is resolved first and the signed URL is handed to aria2 without +#: the credentials. +AUTH_REDIRECT_DOWNLOAD_URL_PREFIXES = CIVITAI_DOWNLOAD_URL_PREFIXES + ( + "https://huggingface.co/", +) + def _is_no_uri_available_error(message: str) -> bool: """Return True for aria2's "No URI available" transfer failure. @@ -308,12 +316,12 @@ class Aria2Downloader: resolved_url = url request_headers = headers - if headers and url.startswith(CIVITAI_DOWNLOAD_URL_PREFIXES): + if headers and url.startswith(AUTH_REDIRECT_DOWNLOAD_URL_PREFIXES): resolved_url = await self._resolve_authenticated_redirect_url(url, headers) if resolved_url != url: request_headers = None logger.debug( - "Resolved Civitai download %s to signed URL for aria2", + "Resolved authenticated download %s to signed URL for aria2", download_id, ) @@ -341,7 +349,7 @@ class Aria2Downloader: ] logger.debug( - "Submitting aria2 download %s -> %s (auth=%s, civitai_signed=%s)", + "Submitting aria2 download %s -> %s (auth=%s, signed_url=%s)", download_id, save_path, bool(request_headers), @@ -732,7 +740,7 @@ class Aria2Downloader: if location: return location raise Aria2Error( - "Authenticated Civitai redirect did not include a Location header" + "Authenticated redirect did not include a Location header" ) if response.status == 200: @@ -740,12 +748,12 @@ class Aria2Downloader: body = await response.text() raise Aria2Error( - f"Failed to resolve authenticated Civitai redirect: status={response.status} body={body[:300]}" + f"Failed to resolve authenticated redirect: status={response.status} body={body[:300]}" ) except aiohttp.ClientError as exc: if is_ssl_cert_verify_error(exc): logger.error( - "SSL certificate verification failed during Civitai redirect " + "SSL certificate verification failed during authenticated redirect " "resolution for %s. This is usually caused by an outdated CA " "certificate bundle. Recommended fixes:\n" " 1. pip install --upgrade certifi\n" @@ -753,7 +761,7 @@ class Aria2Downloader: url, ) raise Aria2Error( - f"Failed to resolve authenticated Civitai redirect: {exc}" + f"Failed to resolve authenticated redirect: {exc}" ) from exc async def _ensure_process(self) -> None: diff --git a/py/services/model_sources/base.py b/py/services/model_sources/base.py index 011e676b..6d097914 100644 --- a/py/services/model_sources/base.py +++ b/py/services/model_sources/base.py @@ -193,7 +193,9 @@ def is_valid_source_id(source_id: str) -> bool: ) -async def fetch_text(url: str, *, timeout: int = HTTP_TIMEOUT) -> str: +async def fetch_text( + url: str, *, timeout: int = HTTP_TIMEOUT, headers: Optional[Dict[str, str]] = None +) -> str: """Fetch *url* and return its body as text, or ``""`` on any failure. Network problems are expected (offline installs, rate limits, dead @@ -202,8 +204,11 @@ async def fetch_text(url: str, *, timeout: int = HTTP_TIMEOUT) -> str: """ try: + request_headers = {"User-Agent": USER_AGENT} + if headers: + request_headers.update(headers) async with aiohttp.ClientSession( - headers={"User-Agent": USER_AGENT}, + headers=request_headers, timeout=aiohttp.ClientTimeout(total=timeout), ) as session: async with session.get(url) as resp: @@ -216,7 +221,7 @@ async def fetch_text(url: str, *, timeout: int = HTTP_TIMEOUT) -> str: async def fetch_json( - url: str, *, timeout: int = HTTP_TIMEOUT + url: str, *, timeout: int = HTTP_TIMEOUT, headers: Optional[Dict[str, str]] = None ) -> tuple[int, Any]: """Fetch *url* and return ``(status, parsed_body)``. @@ -227,8 +232,11 @@ async def fetch_json( """ try: + request_headers = {"User-Agent": USER_AGENT} + if headers: + request_headers.update(headers) async with aiohttp.ClientSession( - headers={"User-Agent": USER_AGENT}, + headers=request_headers, timeout=aiohttp.ClientTimeout(total=timeout), ) as session: async with session.get(url) as resp: @@ -381,6 +389,15 @@ class ModelSource: return [] + def auth_headers(self) -> Dict[str, str]: + """Extra request headers this site needs for API and file downloads. + + Empty by default; sites with gated/private content (Hugging Face) + override it to attach the user's access token when one is configured. + """ + + return {} + def file_download_url( self, source_id: str, filename: str, revision: str = "" ) -> str: diff --git a/py/services/model_sources/huggingface.py b/py/services/model_sources/huggingface.py index 4fad9086..f67f0b0f 100644 --- a/py/services/model_sources/huggingface.py +++ b/py/services/model_sources/huggingface.py @@ -27,6 +27,18 @@ _STRICT_URL_PATTERN = re.compile( ) +def _hf_token() -> str: + """Return the configured Hugging Face access token, or ``""``.""" + + try: + from ..settings_manager import get_settings_manager + + token = get_settings_manager().get("huggingface_api_key", "") + except Exception: # pragma: no cover - settings must never break downloads + return "" + return token.strip() if isinstance(token, str) else "" + + class HuggingFaceSource(ModelSource): """Hugging Face Hub (``huggingface.co``).""" @@ -45,12 +57,20 @@ class HuggingFaceSource(ModelSource): def asset_base_url(self, source_id: str, revision: str = "") -> str: return f"https://huggingface.co/{source_id}/resolve/{self.resolve_revision(revision)}" + def auth_headers(self) -> dict[str, str]: + """Bearer header for gated/private repositories, when a token is set.""" + + token = _hf_token() + return {"Authorization": f"Bearer {token}"} if token else {} + async def fetch_model_card(self, source_id: str) -> str: """Fetch ``README.md`` from Hugging Face (tries ``main``, then ``master``).""" + headers = self.auth_headers() for branch in ("main", "master"): text = await fetch_text( - f"https://huggingface.co/{source_id}/raw/{branch}/README.md" + f"https://huggingface.co/{source_id}/raw/{branch}/README.md", + headers=headers, ) if text: return text @@ -67,11 +87,26 @@ class HuggingFaceSource(ModelSource): revision = self.resolve_revision(revision) status, payload = await fetch_json( - f"https://huggingface.co/api/models/{source_id}/tree/{revision}" + f"https://huggingface.co/api/models/{source_id}/tree/{revision}", + headers=self.auth_headers(), ) if status == 404: raise ModelSourceError(f"Repository '{source_id}' not found", status=404) + if status in (401, 403): + if _hf_token(): + raise ModelSourceError( + f"Access to '{source_id}' was denied (HTTP {status}). For a gated " + "repository you must accept its terms on the Hugging Face page, " + "and the configured token needs read permission for it.", + status=403, + ) + raise ModelSourceError( + f"'{source_id}' requires a Hugging Face access token (gated or " + "private repository). Configure one in Settings → Hugging Face " + "Access Token, and accept the repository's terms on its page.", + status=401, + ) if status != 200 or not isinstance(payload, list): raise ModelSourceError( f"Hugging Face API error while listing '{source_id}' (HTTP {status})" diff --git a/py/services/settings_manager.py b/py/services/settings_manager.py index e12fa9b7..4f25ef0e 100644 --- a/py/services/settings_manager.py +++ b/py/services/settings_manager.py @@ -67,6 +67,7 @@ DEFAULT_KEYS_CLEANUP_THRESHOLD = 10 DEFAULT_SETTINGS: Dict[str, Any] = { "civitai_api_key": "", + "huggingface_api_key": "", "civitai_host": "civitai.com", "download_backend": "python", "aria2c_path": "", @@ -1124,6 +1125,15 @@ class SettingsManager: self.settings["civitai_api_key"] = env_api_key self._save_settings() + # Hugging Face accepts either of its conventional variable names + env_hf_token = os.environ.get("HF_TOKEN") or os.environ.get( + "HUGGING_FACE_HUB_TOKEN" + ) + if env_hf_token: + logger.info("Found HF_TOKEN environment variable") + self.settings["huggingface_api_key"] = env_hf_token + self._save_settings() + # LLM provider overrides llm_env_map = { "LLM_API_KEY": "llm_api_key", diff --git a/static/js/managers/SettingsManager.js b/static/js/managers/SettingsManager.js index 996e4f38..7ca0d22e 100644 --- a/static/js/managers/SettingsManager.js +++ b/static/js/managers/SettingsManager.js @@ -928,6 +928,7 @@ export class SettingsManager { // Update API key status display (do NOT pre-fill the input) this.updateApiKeyStatus(); + this.updateHfApiKeyStatus(); this.updateLlmApiKeyStatus(); // ── AI Provider settings ────────────────────────────────────── @@ -4350,6 +4351,28 @@ export class SettingsManager { } } + updateHfApiKeyStatus() { + const hasKey = !!(state.global.settings.huggingface_api_key_set || + state.global.settings.huggingface_api_key); + const statusText = document.getElementById('huggingfaceApiKeyStatusText'); + const actionBtn = document.getElementById('huggingfaceApiKeyActionBtn'); + if (!statusText || !actionBtn) return; + + if (hasKey) { + statusText.classList.remove('api-key-status--unconfigured'); + statusText.classList.add('api-key-status--configured'); + statusText.innerHTML = ' ' + + translate('settings.huggingfaceApiKeyConfigured', {}, 'Configured'); + actionBtn.textContent = translate('common.actions.change', {}, 'Change'); + } else { + statusText.classList.remove('api-key-status--configured'); + statusText.classList.add('api-key-status--unconfigured'); + statusText.innerHTML = ' ' + + translate('settings.huggingfaceApiKeyNotConfigured', {}, 'Not configured'); + actionBtn.textContent = translate('settings.huggingfaceApiKeySet', {}, 'Set up'); + } + } + updateLlmApiKeyStatus() { const hasKey = !!(state.global.settings.llm_api_key_set || state.global.settings.llm_api_key); const statusText = document.getElementById('llmApiKeyStatusText'); @@ -4397,9 +4420,17 @@ export class SettingsManager { const input = document.getElementById(inputId); if (input) input.value = ''; if (!silent) { - if (inputId === 'civitaiApiKey') { - this.updateApiKeyStatus(); - } + this.refreshApiKeyStatus(inputId); + } + } + + refreshApiKeyStatus(inputId) { + if (inputId === 'civitaiApiKey') { + this.updateApiKeyStatus(); + } else if (inputId === 'huggingfaceApiKey') { + this.updateHfApiKeyStatus(); + } else if (inputId === 'llmApiKey') { + this.updateLlmApiKeyStatus(); } } @@ -4409,11 +4440,16 @@ export class SettingsManager { const value = input.value.trim(); + const labelNames = { + civitai_api_key: 'CivitAI API Key', + huggingface_api_key: 'Hugging Face Access Token', + llm_api_key: 'LLM API Key', + }; + try { await this.saveSetting(settingsKey, value); - const labelName = settingsKey === 'civitai_api_key' ? 'CivitAI API Key' : 'LLM API Key'; showToast('toast.settings.settingsUpdated', - { setting: labelName }, 'success'); + { setting: labelNames[settingsKey] || 'API Key' }, 'success'); } catch (error) { showToast('toast.settings.settingSaveFailed', { message: error.message }, 'error'); @@ -4421,13 +4457,12 @@ export class SettingsManager { } // Update the in-memory flag so the UI reflects the change - if (settingsKey === 'civitai_api_key') { - state.global.settings.civitai_api_key_set = !!value; + const setFlagKey = `${settingsKey}_set`; + if (setFlagKey in state.global.settings) { + state.global.settings[setFlagKey] = !!value; } this.cancelEditApiKey(true, inputId); - if (inputId === 'civitaiApiKey') { - this.updateApiKeyStatus(); - } + this.refreshApiKeyStatus(inputId); } toggleInputVisibility(button) { diff --git a/static/js/state/index.js b/static/js/state/index.js index 04c582d8..20eabd7c 100644 --- a/static/js/state/index.js +++ b/static/js/state/index.js @@ -6,6 +6,8 @@ import { DEFAULT_PATH_TEMPLATES, DEFAULT_FILENAME_TEMPLATES, DEFAULT_PRIORITY_TA const DEFAULT_SETTINGS_BASE = Object.freeze({ civitai_api_key: '', civitai_api_key_set: false, + huggingface_api_key: '', + huggingface_api_key_set: false, civitai_host: 'civitai.com', download_backend: 'python', aria2c_path: '', diff --git a/templates/components/modals/settings/general.html b/templates/components/modals/settings/general.html index 8a79afb0..f29282f9 100644 --- a/templates/components/modals/settings/general.html +++ b/templates/components/modals/settings/general.html @@ -68,6 +68,43 @@ +
+
+
+ + +
+
+ +
+ + + {{ t('settings.huggingfaceApiKeyNotConfigured') }} + + +
+ + +
+
+
+ {{ sm.setting_select('civitaiHost', 'civitai_host', 'settings.civitaiHost.label', [ ('civitai.com', 'settings.civitaiHost.options.com'), ('civitai.red', 'settings.civitaiHost.options.red'), diff --git a/tests/routes/__snapshots__/test_api_snapshots.ambr b/tests/routes/__snapshots__/test_api_snapshots.ambr index a5f81cd7..c2bab914 100644 --- a/tests/routes/__snapshots__/test_api_snapshots.ambr +++ b/tests/routes/__snapshots__/test_api_snapshots.ambr @@ -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, diff --git a/tests/routes/test_model_source_handlers.py b/tests/routes/test_model_source_handlers.py index 45f25feb..afbb99c5 100644 --- a/tests/routes/test_model_source_handlers.py +++ b/tests/routes/test_model_source_handlers.py @@ -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 diff --git a/tests/services/test_aria2_downloader.py b/tests/services/test_aria2_downloader.py index dff0aedc..23dc4295 100644 --- a/tests/services/test_aria2_downloader.py +++ b/tests/services/test_aria2_downloader.py @@ -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] diff --git a/tests/services/test_model_sources.py b/tests/services/test_model_sources.py index beda1508..9e720c4d 100644 --- a/tests/services/test_model_sources.py +++ b/tests/services/test_model_sources.py @@ -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)