Add Hugging Face support and enhance model resolution logic
- Introduced support for Hugging Face API integration, allowing fallback model resolution when Civitai fails. - Updated configuration to include `HF_TOKEN` and `HF_TOKEN_PATH` for authentication. - Enhanced model capture logic to differentiate between Civitai and Hugging Face sources. - Improved error handling for model downloads, providing clearer messages for authentication issues. - Updated documentation to reflect new environment variables and usage instructions for Hugging Face integration. - Added tests to validate the new fallback mechanism and ensure robust model resolution.
This commit is contained in:
Leonid Pershin
@@ -0,0 +1,203 @@
|
||||
"""Hugging Face Hub: auth probe, metadata lookup (capture fallback), URL helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from urllib.parse import quote, urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
from gpu_rent.errors import CloudError
|
||||
|
||||
HF_API = "https://huggingface.co/api"
|
||||
HF_HOST = "huggingface.co"
|
||||
|
||||
|
||||
@dataclass
|
||||
class HfProbe:
|
||||
ok: bool
|
||||
name: str = ""
|
||||
detail: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class HfFileHit:
|
||||
repo_id: str
|
||||
filename: str
|
||||
url: str
|
||||
sha256: str = ""
|
||||
title: str = ""
|
||||
|
||||
|
||||
def is_huggingface_url(url: str) -> bool:
|
||||
host = (urlparse(url).hostname or "").lower()
|
||||
return host == HF_HOST or host.endswith(".huggingface.co")
|
||||
|
||||
|
||||
def hf_resolve_url(repo_id: str, filename: str, *, revision: str = "main") -> str:
|
||||
repo = repo_id.strip().strip("/")
|
||||
name = filename.lstrip("/")
|
||||
rev = revision.strip() or "main"
|
||||
return f"https://{HF_HOST}/{repo}/resolve/{quote(rev, safe='')}/{quote(name, safe='/')}"
|
||||
|
||||
|
||||
def probe_whoami(token: str, timeout: float = 15.0) -> HfProbe:
|
||||
token = (token or "").strip()
|
||||
if not token:
|
||||
return HfProbe(ok=False, detail="токена нет")
|
||||
url = f"{HF_API}/whoami-v2"
|
||||
try:
|
||||
with httpx.Client(timeout=timeout, follow_redirects=True) as client:
|
||||
response = client.get(url, headers={"Authorization": f"Bearer {token}"})
|
||||
except httpx.HTTPError as exc:
|
||||
return HfProbe(ok=False, detail=str(exc)[:160])
|
||||
if response.status_code == 200:
|
||||
data = response.json() if response.content else {}
|
||||
name = ""
|
||||
if isinstance(data, dict):
|
||||
name = str(data.get("name") or data.get("fullname") or "")
|
||||
return HfProbe(ok=True, name=name, detail="токен принят")
|
||||
if response.status_code in {401, 403}:
|
||||
return HfProbe(
|
||||
ok=False,
|
||||
detail="токен отвергнут — https://huggingface.co/settings/tokens",
|
||||
)
|
||||
return HfProbe(
|
||||
ok=False,
|
||||
detail=f"HTTP {response.status_code}: {(response.text or '')[:120]}",
|
||||
)
|
||||
|
||||
|
||||
def _headers(token: str | None) -> dict[str, str]:
|
||||
h = {"User-Agent": "gpu-rent/1"}
|
||||
if token:
|
||||
h["Authorization"] = f"Bearer {token}"
|
||||
return h
|
||||
|
||||
|
||||
def _sha_from_entry(entry: dict) -> str:
|
||||
lfs = entry.get("lfs")
|
||||
if isinstance(lfs, dict):
|
||||
for key in ("sha256", "oid", "sha"):
|
||||
val = lfs.get(key)
|
||||
if val:
|
||||
return str(val).strip().lower()
|
||||
for key in ("sha256", "oid", "blob_id"):
|
||||
val = entry.get(key)
|
||||
if val and len(str(val)) >= 40:
|
||||
return str(val).strip().lower()
|
||||
return ""
|
||||
|
||||
|
||||
def _tree_match(
|
||||
client: httpx.Client,
|
||||
repo_id: str,
|
||||
want_sha: str,
|
||||
*,
|
||||
filename_hint: str | None,
|
||||
token: str | None,
|
||||
) -> HfFileHit | None:
|
||||
url = f"{HF_API}/models/{repo_id}/tree/main"
|
||||
response = client.get(
|
||||
url,
|
||||
params={"recursive": "true"},
|
||||
headers=_headers(token),
|
||||
)
|
||||
if response.status_code != 200:
|
||||
return None
|
||||
data = response.json()
|
||||
if not isinstance(data, list):
|
||||
return None
|
||||
hint = (filename_hint or "").replace("\\", "/").split("/")[-1].lower()
|
||||
for entry in data:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
if entry.get("type") and entry.get("type") != "file":
|
||||
continue
|
||||
path = str(entry.get("path") or entry.get("rfilename") or "")
|
||||
if not path:
|
||||
continue
|
||||
sha = _sha_from_entry(entry)
|
||||
if sha and sha == want_sha:
|
||||
return HfFileHit(
|
||||
repo_id=repo_id,
|
||||
filename=path,
|
||||
url=hf_resolve_url(repo_id, path),
|
||||
sha256=sha,
|
||||
title=repo_id,
|
||||
)
|
||||
# Filename fallback when tree has no sha (rare) — only exact name match
|
||||
if hint and path.replace("\\", "/").split("/")[-1].lower() == hint and not sha:
|
||||
return HfFileHit(
|
||||
repo_id=repo_id,
|
||||
filename=path,
|
||||
url=hf_resolve_url(repo_id, path),
|
||||
sha256="",
|
||||
title=repo_id,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def lookup_by_sha256(
|
||||
token: str | None,
|
||||
sha256: str,
|
||||
*,
|
||||
filename: str | None = None,
|
||||
limit_repos: int = 12,
|
||||
timeout: float = 45.0,
|
||||
) -> HfFileHit | None:
|
||||
"""Best-effort: search Hub by filename, match LFS sha256 in repo tree.
|
||||
|
||||
Hugging Face has no public global by-hash index; this is the practical fallback
|
||||
after Civitai miss. Token optional for public repos; gated need HF_TOKEN.
|
||||
"""
|
||||
digest = sha256.strip().lower()
|
||||
if len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest):
|
||||
raise CloudError(f"HF by-hash: неверный SHA256 ({sha256[:16]}…)")
|
||||
|
||||
hint = (filename or "").replace("\\", "/").split("/")[-1]
|
||||
stem = hint.rsplit(".", 1)[0] if hint else ""
|
||||
queries: list[str] = []
|
||||
if hint:
|
||||
queries.append(hint)
|
||||
if stem and stem.lower() != hint.lower():
|
||||
queries.append(stem)
|
||||
if not queries:
|
||||
return None
|
||||
|
||||
try:
|
||||
with httpx.Client(timeout=timeout, follow_redirects=True) as client:
|
||||
seen: set[str] = set()
|
||||
for q in queries:
|
||||
response = client.get(
|
||||
f"{HF_API}/models",
|
||||
params={"search": q, "limit": limit_repos},
|
||||
headers=_headers(token),
|
||||
)
|
||||
if response.status_code in {401, 403}:
|
||||
raise CloudError(
|
||||
"HF search: 401/403 — нужен валидный HF_TOKEN "
|
||||
"(https://huggingface.co/settings/tokens)"
|
||||
)
|
||||
if response.status_code != 200:
|
||||
continue
|
||||
models = response.json()
|
||||
if not isinstance(models, list):
|
||||
continue
|
||||
for model in models:
|
||||
if not isinstance(model, dict):
|
||||
continue
|
||||
repo_id = str(model.get("modelId") or model.get("id") or "").strip()
|
||||
if not repo_id or repo_id in seen:
|
||||
continue
|
||||
seen.add(repo_id)
|
||||
hit = _tree_match(
|
||||
client, repo_id, digest, filename_hint=hint, token=token
|
||||
)
|
||||
if hit:
|
||||
return hit
|
||||
if len(seen) >= limit_repos:
|
||||
return None
|
||||
except httpx.HTTPError as exc:
|
||||
raise CloudError(f"HF lookup: {exc}") from exc
|
||||
return None
|
||||
Reference in New Issue
Block a user