diff --git a/__init__.py b/__init__.py index a1755b8..43187a4 100644 --- a/__init__.py +++ b/__init__.py @@ -101,6 +101,8 @@ class ModelLinkerExtension: search_huggingface_for_file, get_author_fallback_index_status, refresh_author_fallback_index, + check_huggingface_token, + check_brave_search_api_key, clear_search_cache as clear_huggingface_search_cache, ) from .core.sources.civitai import ( @@ -108,6 +110,8 @@ class ModelLinkerExtension: search_civitai, get_civitai_download_url, resolve_urn, + check_civitai_api_key, + check_civitai_session_token, clear_search_cache as clear_civitai_search_cache, ) from .core.sources.civarchive import ( @@ -1270,6 +1274,60 @@ class ModelLinkerExtension: log_exception(f"Clear search cache error: {e}") return web.json_response({"error": str(e)}, status=500) + @routes.post("/model_linker/civitai/session-token/check") + async def civitai_session_token_check_route(request): + """Check whether a CivitAI browser session token is valid.""" + try: + data = await request.json() + token = data.get("civitai_session_token", "") + result = await asyncio.to_thread( + check_civitai_session_token, token + ) + return web.json_response(result) + except Exception as e: + log_exception(f"CivitAI session token check error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @routes.post("/model_linker/civitai/api-key/check") + async def civitai_api_key_check_route(request): + """Check whether a CivitAI API key is valid.""" + try: + data = await request.json() + api_key = data.get("civitai_key", "") + result = await asyncio.to_thread( + check_civitai_api_key, api_key + ) + return web.json_response(result) + except Exception as e: + log_exception(f"CivitAI API key check error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @routes.post("/model_linker/huggingface/token/check") + async def huggingface_token_check_route(request): + """Check whether a HuggingFace token is valid.""" + try: + data = await request.json() + token = data.get("hf_token", "") + result = await asyncio.to_thread(check_huggingface_token, token) + return web.json_response(result) + except Exception as e: + log_exception(f"HuggingFace token check error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @routes.post("/model_linker/brave/api-key/check") + async def brave_api_key_check_route(request): + """Check whether a Brave Search API key is valid.""" + try: + data = await request.json() + api_key = data.get("brave_search_api_key", "") + result = await asyncio.to_thread( + check_brave_search_api_key, api_key + ) + return web.json_response(result) + except Exception as e: + log_exception(f"Brave Search API key check error: {e}") + return web.json_response({"error": str(e)}, status=500) + @routes.get("/model_linker/huggingface/author-index/status") async def huggingface_author_index_status_route(request): """Return local HuggingFace author fallback index status.""" diff --git a/core/sources/civitai.py b/core/sources/civitai.py index d1be6f6..5ac6880 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -50,6 +50,170 @@ def clear_search_cache(): _search_cache.clear() +def check_civitai_session_token(session_token: Optional[str]) -> Dict[str, Any]: + """Check whether a CivitAI browser session token is accepted by civitai.com.""" + token = (session_token or "").strip() + if not token: + return { + "success": False, + "valid": False, + "status": "missing", + "message": "Paste a CivitAI session token first.", + } + + headers = { + "accept": "application/json", + "Cookie": f"__Secure-civitai-token={token}", + "user-agent": ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/147.0.0.0 Safari/537.36" + ), + } + + try: + response = requests.get( + "https://civitai.com/api/v1/me", headers=headers, timeout=10 + ) + if response.status_code == 200: + username = "" + try: + data = response.json() + username = ( + data.get("username") + or data.get("name") + or data.get("email") + or "" + ) + except Exception: + username = "" + + message = "Session token is valid." + if username: + message = f"Session token is valid for {username}." + + return { + "success": True, + "valid": True, + "status": "valid", + "message": message, + "username": username, + } + + if response.status_code in {401, 403}: + return { + "success": True, + "valid": False, + "status": "invalid", + "message": "Session token is not accepted by CivitAI.", + "status_code": response.status_code, + } + + return { + "success": False, + "valid": False, + "status": "error", + "message": f"CivitAI returned HTTP {response.status_code}.", + "status_code": response.status_code, + } + except requests.exceptions.Timeout: + return { + "success": False, + "valid": False, + "status": "timeout", + "message": "CivitAI did not respond before the timeout.", + } + except Exception as e: + return { + "success": False, + "valid": False, + "status": "error", + "message": str(e), + } + + +def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]: + """Check whether a CivitAI API key is accepted by civitai.com.""" + key = (api_key or "").strip() + if not key: + return { + "success": False, + "valid": False, + "status": "missing", + "message": "Paste a CivitAI API key first.", + } + + headers = { + "accept": "application/json", + "Authorization": f"Bearer {key}", + "user-agent": ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/147.0.0.0 Safari/537.36" + ), + } + + try: + response = requests.get( + "https://civitai.com/api/v1/me", headers=headers, timeout=10 + ) + if response.status_code == 200: + username = "" + try: + data = response.json() + username = ( + data.get("username") + or data.get("name") + or data.get("email") + or "" + ) + except Exception: + username = "" + + message = "CivitAI API key is valid." + if username: + message = f"CivitAI API key is valid for {username}." + + return { + "success": True, + "valid": True, + "status": "valid", + "message": message, + "username": username, + } + + if response.status_code in {401, 403}: + return { + "success": True, + "valid": False, + "status": "invalid", + "message": "CivitAI API key is not accepted.", + "status_code": response.status_code, + } + + return { + "success": False, + "valid": False, + "status": "error", + "message": f"CivitAI returned HTTP {response.status_code}.", + "status_code": response.status_code, + } + except requests.exceptions.Timeout: + return { + "success": False, + "valid": False, + "status": "timeout", + "message": "CivitAI did not respond before the timeout.", + } + except Exception as e: + return { + "success": False, + "valid": False, + "status": "error", + "message": str(e), + } + + def _build_civitai_result_from_version( model_id: int, model_name: str, diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index a2ee326..0503b2f 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -45,6 +45,141 @@ def clear_search_cache(): _search_cache.clear() +def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]: + """Check whether a HuggingFace access token is accepted.""" + value = (token or "").strip() + if not value: + return { + "success": False, + "valid": False, + "status": "missing", + "message": "Paste a HuggingFace token first.", + } + + headers = {"Authorization": f"Bearer {value}"} + try: + response = requests.get(f"{HF_API_URL}/whoami-v2", headers=headers, timeout=10) + if response.status_code == 404: + response = requests.get(f"{HF_API_URL}/whoami", headers=headers, timeout=10) + + if response.status_code == 200: + username = "" + try: + data = response.json() + username = data.get("name") or data.get("user", {}).get("name") or "" + except Exception: + username = "" + + message = "HuggingFace token is valid." + if username: + message = f"HuggingFace token is valid for {username}." + + return { + "success": True, + "valid": True, + "status": "valid", + "message": message, + "username": username, + } + + if response.status_code in {401, 403}: + return { + "success": True, + "valid": False, + "status": "invalid", + "message": "HuggingFace token is not accepted.", + "status_code": response.status_code, + } + + return { + "success": False, + "valid": False, + "status": "error", + "message": f"HuggingFace returned HTTP {response.status_code}.", + "status_code": response.status_code, + } + except requests.exceptions.Timeout: + return { + "success": False, + "valid": False, + "status": "timeout", + "message": "HuggingFace did not respond before the timeout.", + } + except Exception as e: + return { + "success": False, + "valid": False, + "status": "error", + "message": str(e), + } + + +def check_brave_search_api_key(api_key: Optional[str]) -> Dict[str, Any]: + """Check whether a Brave Search API key is accepted.""" + key = (api_key or "").strip() + if not key: + return { + "success": False, + "valid": False, + "status": "missing", + "message": "Paste a Brave Search API key first.", + } + + headers = {"X-Subscription-Token": key} + params = {"q": "test", "count": 1} + try: + response = requests.get( + BRAVE_SEARCH_API_URL, headers=headers, params=params, timeout=10 + ) + if response.status_code == 200: + return { + "success": True, + "valid": True, + "status": "valid", + "message": "Brave Search API key is valid.", + } + + if response.status_code in {401, 403}: + return { + "success": True, + "valid": False, + "status": "invalid", + "message": "Brave Search API key is not accepted.", + "status_code": response.status_code, + } + + if response.status_code == 429: + return { + "success": False, + "valid": False, + "status": "limited", + "message": "Brave Search rate limit was reached.", + "status_code": response.status_code, + } + + return { + "success": False, + "valid": False, + "status": "error", + "message": f"Brave Search returned HTTP {response.status_code}.", + "status_code": response.status_code, + } + except requests.exceptions.Timeout: + return { + "success": False, + "valid": False, + "status": "timeout", + "message": "Brave Search did not respond before the timeout.", + } + except Exception as e: + return { + "success": False, + "valid": False, + "status": "error", + "message": str(e), + } + + def _author_index_cache_key(author: str, headers: Dict[str, str]) -> str: auth_key = "token" if headers.get("Authorization") else "public" return f"{author}::{auth_key}" diff --git a/web/css/linker-main.css b/web/css/linker-main.css index 8c903b8..efb0ff9 100644 --- a/web/css/linker-main.css +++ b/web/css/linker-main.css @@ -3542,6 +3542,8 @@ flex-direction: column; gap: 0; align-items: stretch; + padding-bottom: 14px; + border-bottom: 1px solid rgba(255,255,255,0.1); } .ml-options-help, #model-linker-modal .ml-options-help { @@ -3579,6 +3581,32 @@ color: var(--ml-text); font-size: var(--ml-font-xs); } +.ml-options-token-check, +#model-linker-modal .ml-options-token-check { + display: flex; + align-items: center; + flex-wrap: wrap; + gap: 10px; + margin-top: 10px; +} +.ml-options-token-check-status, +#model-linker-modal .ml-options-token-check-status { + color: var(--ml-text-muted); + font-size: var(--ml-font-sm); + font-weight: 600; +} +.ml-options-token-check-status.is-pending, +#model-linker-modal .ml-options-token-check-status.is-pending { + color: var(--ml-text-secondary); +} +.ml-options-token-check-status.is-valid, +#model-linker-modal .ml-options-token-check-status.is-valid { + color: var(--ml-color-success); +} +.ml-options-token-check-status.is-invalid, +#model-linker-modal .ml-options-token-check-status.is-invalid { + color: var(--ml-color-error-light); +} .ml-options-label { display: flex; align-items: center; @@ -3650,8 +3678,7 @@ grid-template-columns: minmax(150px, 1fr) minmax(0, 380px) 42px; align-items: center; gap: 10px; - padding: 14px 0; - border-bottom: 1px solid rgba(255,255,255,0.1); + padding: 14px 0 0; min-width: 0; } .ml-options-input-row > .ml-options-label { diff --git a/web/linker.js b/web/linker.js index 5f01752..92d440c 100644 --- a/web/linker.js +++ b/web/linker.js @@ -4051,6 +4051,10 @@ class LinkerManagerDialog extends ComfyDialog { ${getSvgIcon('eye')} +