From 618e6e4806c18e0eb8f8b47327b44aad80ede647 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Fri, 21 Aug 2026 07:34:15 +0300 Subject: [PATCH] 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. --- docs/llm.md | 2 +- docs/models.md | 8 + env.example | 5 +- src/gpu_rent/capture.py | 171 +++++++++++++++++----- src/gpu_rent/config.py | 4 + src/gpu_rent/doctor.py | 53 +++++++ src/gpu_rent/huggingface.py | 203 ++++++++++++++++++++++++++ src/gpu_rent/provision.py | 100 +++++++++---- src/gpu_rent/remote/civitai_fetch.py | 68 +++++++-- src/gpu_rent/remote/llamacpp_fetch.py | 34 ++++- tests/test_capture_merge.py | 75 ++++++++-- tests/test_huggingface.py | 15 ++ 12 files changed, 642 insertions(+), 96 deletions(-) create mode 100644 src/gpu_rent/huggingface.py create mode 100644 tests/test_huggingface.py diff --git a/docs/llm.md b/docs/llm.md index 52f3f32..944fe8b 100644 --- a/docs/llm.md +++ b/docs/llm.md @@ -129,7 +129,7 @@ Unit `gpu-rent-ollama` читает `/mnt/swarm_data/.gpu-rent-gpu.json`: | 4 | empty | только runtime | | — | keep | не трогать yaml | -`-ngl` / `-c` — по GPU probe. Опционально `HF_TOKEN` для gated HF. Вручную: положи GGUF в models и `systemctl restart gpu-rent-llamacpp`. +`-ngl` / `-c` — по GPU probe. Для recommended abliterate GGUF нужен **`HF_TOKEN`** в `.env` (иначе 401). Вручную: положи GGUF в models и `systemctl restart gpu-rent-llamacpp`. --- diff --git a/docs/models.md b/docs/models.md index 573727b..973085d 100644 --- a/docs/models.md +++ b/docs/models.md @@ -189,6 +189,14 @@ REST не раздвоился: те же `/api/v1/...` на обоих хост ```env CIVITAI_API_TOKEN= CIVITAI_API_HOST=civitai.red # полный каталог; civitai.com = только SFW +HF_TOKEN= # или HUGGING_FACE_HUB_TOKEN — HF GGUF / gated / capture fallback MODELS_MANIFEST= # пусто = /models.yaml LOCAL_MODELS_DIR= # пусто = <корень приложения>/Models ``` + +## Hugging Face (резерв) + +- **Скачивание:** `HF_TOKEN` в `.env` → GGUF (llama.cpp) и строки `models.yaml` с `huggingface.co/…/resolve/…`. +- **Capture:** если Civitai by-hash не нашёл файл → поиск на Hub по имени + LFS sha256 → URL в манифест. +- **SwarmUI:** тот же токен прокидывается как `huggingface_api` (Model Downloader). +- Abliterated / gated репозитории без токена почти всегда дают **401**. diff --git a/env.example b/env.example index 50369c0..b0199ef 100644 --- a/env.example +++ b/env.example @@ -46,7 +46,10 @@ OLLAMA_LOCAL_PORT=17811 LLAMACPP_LOCAL_PORT=17812 # OLLAMA_MODELS_MANIFEST= # LLAMACPP_MODELS_MANIFEST= -# HF_TOKEN= # optional, gated GGUF on Hugging Face +# CIVITAI_API_TOKEN= +# CIVITAI_API_HOST=civitai.red +# Hugging Face (GGUF / gated HF URLs / capture fallback metadata): +# HF_TOKEN= # or HUGGING_FACE_HUB_TOKEN — https://huggingface.co/settings/tokens # git pull SwarmUI + extensions on each up (default true). CLI: --no-update UPDATE_GIT=true diff --git a/src/gpu_rent/capture.py b/src/gpu_rent/capture.py index 4122c4a..1d062db 100644 --- a/src/gpu_rent/capture.py +++ b/src/gpu_rent/capture.py @@ -22,6 +22,7 @@ from gpu_rent.civitai import ( ) from gpu_rent.config import Config from gpu_rent.errors import CloudError, GpuRentError +from gpu_rent.huggingface import lookup_by_sha256 from gpu_rent.manifests import ( MODEL_TYPES, extract_version_id, @@ -49,11 +50,12 @@ def strip_git_auth(url: str) -> str: @dataclass class ModelCaptureItem: kind: str - version_id: int - model_id: int url: str title: str = "" rel: str = "" + version_id: int | None = None + model_id: int | None = None + source: str = "civitai" # civitai | huggingface @dataclass @@ -115,6 +117,7 @@ def resolve_model_item( token: str, api_host: str, link_host: str, + hf_token: str | None = None, ) -> ResolveOutcome: kind = str(raw.get("kind") or "") if kind not in MODEL_TYPES: @@ -133,21 +136,50 @@ def resolve_model_item( except (TypeError, ValueError): mid = None - def _ok(v: int, m: int, name: str = title) -> ResolveOutcome: + def _ok_civitai(v: int, m: int, name: str = title) -> ResolveOutcome: return ResolveOutcome( item=ModelCaptureItem( kind=kind, - version_id=v, - model_id=m, url=civitai_model_url(m, v, link_host), title=name, rel=rel, + version_id=v, + model_id=m, + source="civitai", + ), + status="ok", + ) + + def _try_hf(sha: str, detail_prefix: str) -> ResolveOutcome | None: + """Civitai miss → Hugging Face search by filename + LFS sha.""" + try: + hit = lookup_by_sha256( + hf_token, + sha, + filename=str(raw.get("name") or rel), + ) + except CloudError as exc: + return ResolveOutcome( + status="api_error", + detail=f"{detail_prefix} HF: {exc}", + ) + if hit is None: + return None + return ResolveOutcome( + item=ModelCaptureItem( + kind=kind, + url=hit.url, + title=hit.title or title, + rel=rel, + version_id=None, + model_id=None, + source="huggingface", ), status="ok", ) if vid is not None and mid is not None: - return _ok(vid, mid) + return _ok_civitai(vid, mid) # Partial sidecar: have version id → GET /model-versions/{id} for modelId. if vid is not None and mid is None: @@ -167,7 +199,7 @@ def resolve_model_item( vid2, mid2 = version_ids_from_payload(version) if vid2 is not None and mid2 is not None: name = str(version.get("name") or title) - return _ok(vid2, mid2, name) + return _ok_civitai(vid2, mid2, name) return ResolveOutcome( status="unknown", detail=f"{rel} version_id={vid} (нет modelId в ответе)", @@ -183,11 +215,14 @@ def resolve_model_item( _host, version = fetch_model_version_by_hash(token or None, api_host, str(sha)) except CloudError as exc: msg = str(exc) - # 404 = genuinely not on Civitai; other → api_error + # 404 = genuinely not on Civitai → try HF if "HTTP 404" in msg or msg.rstrip().endswith("404"): + hf_out = _try_hf(str(sha), f"{rel} sha={str(sha)[:12]}…") + if hf_out is not None: + return hf_out return ResolveOutcome( status="unknown", - detail=f"{rel} sha={str(sha)[:12]}…", + detail=f"{rel} sha={str(sha)[:12]}… (нет на Civitai/HF)", ) return ResolveOutcome( status="api_error", @@ -195,12 +230,15 @@ def resolve_model_item( ) vid2, mid2 = version_ids_from_payload(version) if vid2 is None or mid2 is None: + hf_out = _try_hf(str(sha), f"{rel} sha={str(sha)[:12]}…") + if hf_out is not None: + return hf_out return ResolveOutcome( status="unknown", detail=f"{rel} sha={str(sha)[:12]}… (пустой payload)", ) name = str(version.get("name") or title) - return _ok(vid2, mid2, name) + return _ok_civitai(vid2, mid2, name) def _backup(path: Path) -> None: @@ -218,29 +256,42 @@ def _keep_model_entry(it: dict) -> bool: return vid not in (None, "", 0, "0") +def _model_dedupe_key(kind: str, *, version_id: int | None, url: str | None) -> tuple: + u = (url or "").rstrip("/").lower() + if u and ("huggingface.co" in u or "hf.co/" in u): + return ("hf", kind, u) + if version_id is not None: + return ("civitai", kind, int(version_id)) + return ("url", kind, u or "?") + + def merge_models_yaml( path: Path, new_items: list[ModelCaptureItem], *, dry_run: bool, ) -> tuple[list[ModelCaptureItem], list[str]]: - """Return (actually_new, skip_msgs). Dedupe by (kind, version_id).""" + """Return (actually_new, skip_msgs). Dedupe by Civitai version_id or HF url.""" existing = parse_models(path) if path.is_file() else [] - have: set[tuple[str, int]] = set() + have: set[tuple] = set() for e in existing: vid = e.version_id if vid is None and e.url: vid = extract_version_id(e.url) - if vid is not None: - have.add((e.kind, vid)) + have.add(_model_dedupe_key(e.kind, version_id=vid, url=e.url)) added: list[ModelCaptureItem] = [] skipped: list[str] = [] - seen_new: set[tuple[str, int]] = set() + seen_new: set[tuple] = set() for item in new_items: - key = (item.kind, item.version_id) + key = _model_dedupe_key(item.kind, version_id=item.version_id, url=item.url) if key in have or key in seen_new: - skipped.append(f"{item.kind} {item.title} modelVersionId={item.version_id}") + label = ( + f"modelVersionId={item.version_id}" + if item.version_id is not None + else item.url + ) + skipped.append(f"{item.kind} {item.title} {label}") continue seen_new.add(key) added.append(item) @@ -374,6 +425,7 @@ def capture_models( link_host = cfg.civitai_api_host or "civitai.red" token = cfg.civitai_api_token api_host = cfg.civitai_api_host + hf_token = cfg.hf_token or None for raw in raw_models: if not isinstance(raw, dict): @@ -395,18 +447,30 @@ def capture_models( # Full sidecar / ids → no API. Partial vid → GET version. Else batch by-hash. if vid_i is not None and mid_i is not None: outcome = resolve_model_item( - raw, token=token, api_host=api_host, link_host=link_host + raw, + token=token, + api_host=api_host, + link_host=link_host, + hf_token=hf_token, ) elif vid_i is not None: outcome = resolve_model_item( - raw, token=token, api_host=api_host, link_host=link_host + raw, + token=token, + api_host=api_host, + link_host=link_host, + hf_token=hf_token, ) elif sha: need_hash.append(raw) continue else: outcome = resolve_model_item( - raw, token=token, api_host=api_host, link_host=link_host + raw, + token=token, + api_host=api_host, + link_host=link_host, + hf_token=hf_token, ) if outcome.status == "ok" and outcome.item is not None: @@ -430,31 +494,62 @@ def capture_models( report.models_api_errors.append(f"{rel} sha={sha}… (batch failed)") need_hash = [] + hf_fallback = 0 for raw in need_hash: sha = str(raw.get("sha256") or "").strip().lower() rel = str(raw.get("rel") or raw.get("name") or "?") version = by_hash.get(sha) - if not version: + if version: + vid2, mid2 = version_ids_from_payload(version) + if vid2 is not None and mid2 is not None: + kind = str(raw.get("kind") or "") + if kind not in MODEL_TYPES: + continue + title = str( + version.get("name") or Path(str(raw.get("name") or rel)).stem + ) + resolved.append( + ModelCaptureItem( + kind=kind, + url=civitai_model_url(mid2, vid2, link_host), + title=title, + rel=rel, + version_id=vid2, + model_id=mid2, + source="civitai", + ) + ) + continue + # Civitai miss → Hugging Face + try: + hit = lookup_by_sha256( + hf_token, + sha, + filename=str(raw.get("name") or rel), + ) + except CloudError as exc: + report.models_api_errors.append(f"{rel} HF: {exc}") report.models_unknown.append(f"{rel} sha={sha[:12]}…") continue - vid2, mid2 = version_ids_from_payload(version) - if vid2 is None or mid2 is None: - report.models_unknown.append(f"{rel} sha={sha[:12]}… (пустой payload)") - continue - kind = str(raw.get("kind") or "") - if kind not in MODEL_TYPES: - continue - title = str(version.get("name") or Path(str(raw.get("name") or rel)).stem) - resolved.append( - ModelCaptureItem( - kind=kind, - version_id=vid2, - model_id=mid2, - url=civitai_model_url(mid2, vid2, link_host), - title=title, - rel=rel, + if hit: + kind = str(raw.get("kind") or "") + if kind not in MODEL_TYPES: + continue + hf_fallback += 1 + resolved.append( + ModelCaptureItem( + kind=kind, + url=hit.url, + title=hit.title or Path(str(raw.get("name") or rel)).stem, + rel=rel, + source="huggingface", + ) ) - ) + else: + report.models_unknown.append(f"{rel} sha={sha[:12]}… (нет на Civitai/HF)") + + if hf_fallback: + log(f"capture: Hugging Face fallback — {hf_fallback} файл(ов)") added, skipped = merge_models_yaml( cfg.models_manifest, resolved, dry_run=dry_run diff --git a/src/gpu_rent/config.py b/src/gpu_rent/config.py index 5cc76c1..750f559 100644 --- a/src/gpu_rent/config.py +++ b/src/gpu_rent/config.py @@ -83,6 +83,7 @@ class Config: civitai_api_token: str civitai_api_host: str + hf_token: str models_manifest: Path extensions_manifest: Path git_token: str @@ -215,6 +216,9 @@ def load_config(*, require_auth: bool = True) -> Config: boot_snapshot_name=(os.environ.get("BOOT_SNAPSHOT_NAME") or "gpu-rent-boot-ok").strip(), civitai_api_token=(os.environ.get("CIVITAI_API_TOKEN") or "").strip(), civitai_api_host=(os.environ.get("CIVITAI_API_HOST") or "civitai.red").strip().lower(), + hf_token=( + os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or "" + ).strip(), models_manifest=models_manifest, extensions_manifest=extensions_manifest, git_token=(os.environ.get("GIT_TOKEN") or "").strip(), diff --git a/src/gpu_rent/doctor.py b/src/gpu_rent/doctor.py index cd25ee2..cda197e 100644 --- a/src/gpu_rent/doctor.py +++ b/src/gpu_rent/doctor.py @@ -236,6 +236,7 @@ def run_doctor() -> list[Check]: ) _civitai(cfg, checks) + _huggingface(cfg, checks) _local_manifests(cfg, checks) _local_folders(cfg, checks) return checks @@ -279,6 +280,58 @@ def _civitai(cfg: Config, checks: list[Check]) -> None: ) +def _huggingface(cfg: Config, checks: list[Check]) -> None: + from gpu_rent.huggingface import is_huggingface_url, probe_whoami + from gpu_rent.llm_runtime import normalize_runtime, parse_llamacpp_models + from gpu_rent.manifests import parse_models + + needs_hf = False + try: + for e in parse_models(cfg.models_manifest): + if e.url and is_huggingface_url(e.url): + needs_hf = True + break + except Exception: + pass + try: + if normalize_runtime(cfg.llm_runtime) == "llamacpp": + for e in parse_llamacpp_models(cfg.llamacpp_models_manifest): + if e.url and is_huggingface_url(e.url): + needs_hf = True + break + except Exception: + pass + + if not cfg.hf_token: + checks.append( + Check( + "Hugging Face", + True, + False, + ( + "HF_TOKEN нет — gated GGUF / HF в models.yaml дадут 401. " + "https://huggingface.co/settings/tokens" + if needs_hf + else "токена нет (опционально для HF URL / capture fallback)" + ), + ) + ) + return + probe = probe_whoami(cfg.hf_token) + if probe.ok: + who = f" ({probe.name})" if probe.name else "" + checks.append(Check("Hugging Face", True, True, f"токен ок{who}")) + else: + checks.append( + Check( + "Hugging Face", + True, + False, + f"токен не принят: {probe.detail}", + ) + ) + + def _local_manifests(cfg: Config, checks: list[Check]) -> None: try: models = parse_models(cfg.models_manifest) diff --git a/src/gpu_rent/huggingface.py b/src/gpu_rent/huggingface.py new file mode 100644 index 0000000..a3cab31 --- /dev/null +++ b/src/gpu_rent/huggingface.py @@ -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 diff --git a/src/gpu_rent/provision.py b/src/gpu_rent/provision.py index 0a6ed4f..8f7572e 100644 --- a/src/gpu_rent/provision.py +++ b/src/gpu_rent/provision.py @@ -234,16 +234,10 @@ def seed_swarmui_api_keys(cfg: Config, host: str, log: Log) -> None: Swarm stores them in Users.ldb GenericData — needed for Model Downloader in the UI. Call after SwarmUI HTTP is up (after wait_backend / verify). """ - import os - keys: dict[str, str] = {} if cfg.civitai_api_token: keys["civitai_api"] = cfg.civitai_api_token - hf = ( - os.environ.get("HF_TOKEN") - or os.environ.get("HUGGING_FACE_HUB_TOKEN") - or "" - ).strip() + hf = (cfg.hf_token or "").strip() if hf: keys["huggingface_api"] = hf if not keys: @@ -273,21 +267,52 @@ def seed_swarmui_api_keys(cfg: Config, host: str, log: Log) -> None: def seed_civitai(cfg: Config, host: str, log: Log) -> None: + from gpu_rent.huggingface import is_huggingface_url + entries = parse_models(cfg.models_manifest) - if not cfg.civitai_api_token: - log("Civitai-seed пропущен: нет CIVITAI_API_TOKEN — дефолт SwarmUI") - return if not entries: - log("Civitai-seed пропущен: манифест пуст — дефолт SwarmUI") + log("model-seed пропущен: манифест пуст — дефолт SwarmUI") return - jobs = [] + + jobs: list[dict] = [] for entry in entries: + url = (entry.url or "").strip() + if url and is_huggingface_url(url): + name = url.rstrip("/").rsplit("/", 1)[-1].split("?", 1)[0] or "model.safetensors" + folder = MODEL_DIRS.get(entry.kind, entry.kind) + dest = f"{DATA}/Models/{folder}/{name}" + stem = Path(name).stem + swarm = { + "name": stem, + "title": stem, + "description": f"Hugging Face: {url}", + "trigger_phrase": "", + "author": "", + "tags": ["huggingface"], + } + jobs.append( + { + "dest": dest, + "url": url, + "sha256": "", + "auth": "hf", + "sidecars": { + f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2), + }, + } + ) + continue + + if not cfg.civitai_api_token: + continue vid = entry.version_id or (extract_version_id(entry.url) if entry.url else None) if not vid: - log(f"пропуск {entry.kind}: нет version_id") + log(f"пропуск {entry.kind}: нет version_id / HF url") continue try: - api_host, version = fetch_model_version(cfg.civitai_api_token, cfg.civitai_api_host, vid) + api_host, version = fetch_model_version( + cfg.civitai_api_token, cfg.civitai_api_host, vid + ) except CloudError as exc: log(str(exc)) continue @@ -309,7 +334,11 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None: "title": version.get("name") or stem, "description": (version.get("description") or "")[:2000], "trigger_phrase": phrase, - "author": ((version.get("model") or {}) if isinstance(version.get("model"), dict) else {}).get("name"), + "author": ( + ((version.get("model") or {}) if isinstance(version.get("model"), dict) else {}).get( + "name" + ) + ), "tags": version.get("tags") or [], } jobs.append( @@ -317,22 +346,43 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None: "dest": dest, "url": _download_url(api_host, vid, info), "sha256": sha, + "auth": "civitai", "sidecars": { f"{stem}.civitai.json": civitai_json, f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2), }, } ) + if not jobs: - log("Civitai-seed: ни одной скачиваемой строки") + if not cfg.civitai_api_token: + log("Civitai-seed пропущен: нет CIVITAI_API_TOKEN — дефолт SwarmUI") + else: + log("Civitai-seed: ни одной скачиваемой строки") return + + hf_n = sum(1 for j in jobs if j.get("auth") == "hf") + civ_n = len(jobs) - hf_n + if hf_n and not cfg.hf_token: + log( + "⚠ в манифесте есть Hugging Face URL, но нет HF_TOKEN — " + "gated/abliterated файлы дадут 401. Токен: huggingface.co/settings/tokens" + ) + if civ_n and not cfg.civitai_api_token: + log("⚠ Civitai-строки пропущены: нет CIVITAI_API_TOKEN") + if not any(e.kind == "checkpoint" for e in entries): log("в манифесте нет checkpoint — генерация может не стартовать") put_text(cfg, host, "/tmp/gpu-rent-civitai-jobs.json", json.dumps(jobs, indent=2)) - put_text(cfg, host, "/tmp/gpu-rent-civitai.token", cfg.civitai_api_token + "\n", mode=0o600) + if cfg.civitai_api_token: + put_text( + cfg, host, "/tmp/gpu-rent-civitai.token", cfg.civitai_api_token + "\n", mode=0o600 + ) + if cfg.hf_token: + put_text(cfg, host, "/tmp/gpu-rent-hf.token", cfg.hf_token + "\n", mode=0o600) log( - f"Civitai: {len(jobs)} в манифесте — на VM качаю отсутствующие " - f"(уже есть + sha → skip; прогресс [N/{len(jobs)}])" + f"model-seed: {len(jobs)} файл(ов) " + f"(civitai={civ_n}, huggingface={hf_n}) — прогресс [N/{len(jobs)}]" ) run_python( cfg, @@ -419,7 +469,6 @@ def provision_llm(cfg: Config, host: str, log: Log) -> None: gguf_filename_from_url, parse_llamacpp_models, ) - import os entries = parse_llamacpp_models(cfg.llamacpp_models_manifest) defaults = [e for e in entries if e.default] @@ -439,13 +488,14 @@ def provision_llm(cfg: Config, host: str, log: Log) -> None: put_text( cfg, host, "/tmp/gpu-rent-llamacpp-models.json", json.dumps(jobs, indent=2) ) - hf = ( - os.environ.get("HF_TOKEN") - or os.environ.get("HUGGING_FACE_HUB_TOKEN") - or "" - ).strip() + hf = (cfg.hf_token or "").strip() if hf: put_text(cfg, host, "/tmp/gpu-rent-hf.token", hf + "\n", mode=0o600) + else: + log( + "⚠ HF_TOKEN не задан — gated GGUF (abliterated и др.) часто дают 401. " + "Добавь HF_TOKEN=hf_… в .env → https://huggingface.co/settings/tokens" + ) log(f"llama.cpp: скачиваю {len(jobs)} GGUF из манифеста") run_python( cfg, diff --git a/src/gpu_rent/remote/civitai_fetch.py b/src/gpu_rent/remote/civitai_fetch.py index ef4c040..e966b12 100644 --- a/src/gpu_rent/remote/civitai_fetch.py +++ b/src/gpu_rent/remote/civitai_fetch.py @@ -11,6 +11,7 @@ import urllib.request from pathlib import Path TOKEN_PATH = Path("/tmp/gpu-rent-civitai.token") +HF_TOKEN_PATH = Path("/tmp/gpu-rent-hf.token") JOBS_PATH = Path("/tmp/gpu-rent-civitai-jobs.json") MARKER = Path("/mnt/swarm_data/.gpu-rent-models-seeded") @@ -103,7 +104,7 @@ def should_skip(dest: Path, expect_sha: str) -> tuple[bool, str]: return False, "sha не совпал — перекачиваю" -def download(url: str, dest: Path, token: str, *, label: str) -> None: +def download(url: str, dest: Path, token: str, *, label: str, auth_host: str) -> None: dest.parent.mkdir(parents=True, exist_ok=True) partial = dest.with_suffix(dest.suffix + ".partial") @@ -115,15 +116,26 @@ def download(url: str, dest: Path, token: str, *, label: str) -> None: if new is None: return None host = (urllib.parse.urlparse(new.full_url).hostname or "").lower() - if host.endswith("civitai.com") or host.endswith("civitai.red") or host.endswith("civitai.green"): - return new - # presigned S3 / CDN: token must not leave civitai - return urllib.request.Request(new.full_url, headers={"User-Agent": "gpu-rent/0.1"}) + if auth_host == "civitai": + if ( + host.endswith("civitai.com") + or host.endswith("civitai.red") + or host.endswith("civitai.green") + ): + return new + elif auth_host == "hf": + # Keep Bearer only on hub host; CDN (cdn-lfs.*) is signed — no auth. + if host in {"huggingface.co", "hf.co"}: + return new + return urllib.request.Request( + new.full_url, headers={"User-Agent": "gpu-rent/0.1"} + ) opener = urllib.request.build_opener(StripAuthRedirect) - req = urllib.request.Request( - url, headers={"Authorization": f"Bearer {token}", "User-Agent": "gpu-rent/0.1"} - ) + headers = {"User-Agent": "gpu-rent/0.1"} + if token: + headers["Authorization"] = f"Bearer {token}" + req = urllib.request.Request(url, headers=headers) with opener.open(req, timeout=600) as response, partial.open("wb") as out: total = response.headers.get("Content-Length") try: @@ -142,10 +154,14 @@ def download(url: str, dest: Path, token: str, *, label: str) -> None: def main() -> int: - if not TOKEN_PATH.is_file(): - print("нет токена", file=sys.stderr) - return 1 - token = TOKEN_PATH.read_text(encoding="utf-8").strip() + civitai_token = ( + TOKEN_PATH.read_text(encoding="utf-8").strip() if TOKEN_PATH.is_file() else "" + ) + hf_token = ( + HF_TOKEN_PATH.read_text(encoding="utf-8").strip() + if HF_TOKEN_PATH.is_file() + else "" + ) jobs = json.loads(JOBS_PATH.read_text(encoding="utf-8")) total = len(jobs) failed = 0 @@ -155,12 +171,18 @@ def main() -> int: dest = Path(job["dest"]) expect = (job.get("sha256") or "").lower() prefix = f"[{index}/{total}]" + auth = str(job.get("auth") or "civitai").lower() + if auth == "hf": + token = hf_token + auth_host = "hf" + else: + token = civitai_token + auth_host = "civitai" skip, reason = should_skip(dest, expect) if skip: skipped += 1 size = dest.stat().st_size if dest.is_file() else 0 print(f"{prefix} {reason}: {dest.name} ({fmt_bytes(size)})") - # refresh sidecars even on skip for extra_name, extra_text in (job.get("sidecars") or {}).items(): extra = dest.parent / extra_name extra.parent.mkdir(parents=True, exist_ok=True) @@ -168,9 +190,19 @@ def main() -> int: continue if reason: print(f"{prefix} {reason}: {dest.name}") + if auth_host == "civitai" and not token: + failed += 1 + print(f"{prefix} FAIL {dest.name}: нет CIVITAI токена", file=sys.stderr) + continue try: print(f"{prefix} качаю: {dest.name}", flush=True) - download(job["url"], dest, token, label=f"{prefix} {dest.name}") + download( + job["url"], + dest, + token, + label=f"{prefix} {dest.name}", + auth_host=auth_host, + ) downloaded += 1 if expect: got = sha256_path(dest).lower() @@ -183,9 +215,13 @@ def main() -> int: print(f"{prefix} ok {dest.name} ({fmt_bytes(dest.stat().st_size)})") except Exception as exc: failed += 1 - print(f"{prefix} FAIL {dest.name}: {exc}", file=sys.stderr) + msg = str(exc) + if "401" in msg and auth_host == "hf": + msg += " — нужен HF_TOKEN в .env (huggingface.co/settings/tokens)" + print(f"{prefix} FAIL {dest.name}: {msg}", file=sys.stderr) TOKEN_PATH.unlink(missing_ok=True) - print(f"Civitai итог: скачано {downloaded}, пропущено {skipped}, ошибок {failed} (из {total})") + HF_TOKEN_PATH.unlink(missing_ok=True) + print(f"seed итог: скачано {downloaded}, пропущено {skipped}, ошибок {failed} (из {total})") if failed: return 1 MARKER.parent.mkdir(parents=True, exist_ok=True) diff --git a/src/gpu_rent/remote/llamacpp_fetch.py b/src/gpu_rent/remote/llamacpp_fetch.py index e47b735..eb2340f 100644 --- a/src/gpu_rent/remote/llamacpp_fetch.py +++ b/src/gpu_rent/remote/llamacpp_fetch.py @@ -7,6 +7,7 @@ import os import sys import time import urllib.error +import urllib.parse import urllib.request from pathlib import Path @@ -77,8 +78,25 @@ class DownloadProgress: def download(url: str, dest: Path, headers: dict[str, str], *, label: str) -> None: partial = dest.with_suffix(dest.suffix + ".partial") + + class StripAuthRedirect(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers_resp, newurl): + new = urllib.request.HTTPRedirectHandler.redirect_request( + self, req, fp, code, msg, headers_resp, newurl + ) + if new is None: + return None + host = (urllib.parse.urlparse(new.full_url).hostname or "").lower() + # Hub needs Bearer; CDN (cdn-lfs.*) is pre-signed — drop Authorization. + if host in {"huggingface.co", "hf.co"}: + return new + return urllib.request.Request( + new.full_url, headers={"User-Agent": headers.get("User-Agent", "gpu-rent/1")} + ) + + opener = urllib.request.build_opener(StripAuthRedirect) req = urllib.request.Request(url, headers=headers) - with urllib.request.urlopen(req, timeout=600) as resp, partial.open("wb") as out: + with opener.open(req, timeout=600) as resp, partial.open("wb") as out: cl = resp.headers.get("Content-Length") try: total_n = int(cl) if cl else None @@ -137,7 +155,19 @@ def main() -> int: print(f"{prefix} ok {name} ({fmt_bytes(dest.stat().st_size)})") except (urllib.error.URLError, urllib.error.HTTPError, OSError, TimeoutError) as exc: failed += 1 - print(f"FAIL {name}: {exc}", file=sys.stderr) + msg = str(exc) + if "401" in msg or "403" in msg: + if not token: + msg += ( + " — нет HF_TOKEN: добавь в .env " + "(https://huggingface.co/settings/tokens) и прими условия репо" + ) + else: + msg += ( + " — токен есть, но отказано: проверь scopes / " + "Accept license на странице модели" + ) + print(f"FAIL {name}: {msg}", file=sys.stderr) try: dest.with_suffix(dest.suffix + ".partial").unlink(missing_ok=True) except OSError: diff --git a/tests/test_capture_merge.py b/tests/test_capture_merge.py index 7ed2841..eb45b43 100644 --- a/tests/test_capture_merge.py +++ b/tests/test_capture_merge.py @@ -79,6 +79,9 @@ def test_resolve_by_hash_404_is_unknown(): with patch( "gpu_rent.capture.fetch_model_version_by_hash", side_effect=CloudError("Civitai by-hash abc: HTTP 404 (хосты civitai.red)"), + ), patch( + "gpu_rent.capture.lookup_by_sha256", + return_value=None, ): out = resolve_model_item( { @@ -95,6 +98,40 @@ def test_resolve_by_hash_404_is_unknown(): assert out.item is None +def test_resolve_by_hash_404_hf_fallback(): + from gpu_rent.huggingface import HfFileHit + + with patch( + "gpu_rent.capture.fetch_model_version_by_hash", + side_effect=CloudError("Civitai by-hash abc: HTTP 404 (хосты civitai.red)"), + ), patch( + "gpu_rent.capture.lookup_by_sha256", + return_value=HfFileHit( + repo_id="org/model", + filename="m.safetensors", + url="https://huggingface.co/org/model/resolve/main/m.safetensors", + sha256="a" * 64, + title="org/model", + ), + ): + out = resolve_model_item( + { + "kind": "lora", + "rel": "Lora/m.safetensors", + "name": "m.safetensors", + "sha256": "a" * 64, + }, + token="", + api_host="civitai.red", + link_host="civitai.red", + hf_token="hf_x", + ) + assert out.status == "ok" + assert out.item is not None + assert out.item.source == "huggingface" + assert "huggingface.co" in out.item.url + + def test_resolve_by_hash_network_is_api_error(): with patch( "gpu_rent.capture.fetch_model_version_by_hash", @@ -125,17 +162,25 @@ def test_merge_models_dedupe(tmp_path: Path): ) items = [ ModelCaptureItem( - "lora", 100, 1, "https://civitai.red/models/1?modelVersionId=100", "old" + kind="lora", + version_id=100, + model_id=1, + url="https://civitai.red/models/1?modelVersionId=100", + title="old", ), ModelCaptureItem( - "lora", 200, 2, "https://civitai.red/models/2?modelVersionId=200", "new" + kind="lora", + version_id=200, + model_id=2, + url="https://civitai.red/models/2?modelVersionId=200", + title="new", ), ModelCaptureItem( - "checkpoint", - 300, - 3, - "https://civitai.red/models/3?modelVersionId=300", - "ckpt", + kind="checkpoint", + version_id=300, + model_id=3, + url="https://civitai.red/models/3?modelVersionId=300", + title="ckpt", ), ] added, skipped = merge_models_yaml(path, items, dry_run=False) @@ -153,7 +198,11 @@ def test_merge_models_dry_run(tmp_path: Path): before = path.read_text(encoding="utf-8") items = [ ModelCaptureItem( - "lora", 1, 1, "https://civitai.red/models/1?modelVersionId=1", "x" + kind="lora", + version_id=1, + model_id=1, + url="https://civitai.red/models/1?modelVersionId=1", + title="x", ), ] added, _ = merge_models_yaml(path, items, dry_run=True) @@ -170,11 +219,11 @@ def test_merge_models_kind_scoped_dedupe(tmp_path: Path): ) items = [ ModelCaptureItem( - "checkpoint", - 100, - 1, - "https://civitai.red/models/1?modelVersionId=100", - "as-ckpt", + kind="checkpoint", + version_id=100, + model_id=1, + url="https://civitai.red/models/1?modelVersionId=100", + title="as-ckpt", ), ] added, skipped = merge_models_yaml(path, items, dry_run=False) diff --git a/tests/test_huggingface.py b/tests/test_huggingface.py new file mode 100644 index 0000000..26f4f9a --- /dev/null +++ b/tests/test_huggingface.py @@ -0,0 +1,15 @@ +from gpu_rent.huggingface import hf_resolve_url, is_huggingface_url + + +def test_is_huggingface_url(): + assert is_huggingface_url( + "https://huggingface.co/org/model/resolve/main/a.gguf" + ) + assert not is_huggingface_url("https://civitai.red/models/1") + + +def test_hf_resolve_url(): + assert ( + hf_resolve_url("bartowski/foo", "bar.gguf") + == "https://huggingface.co/bartowski/foo/resolve/main/bar.gguf" + )