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.
2026-08-21 07:34:15 +03:00
parent 7343fb0e83
commit 618e6e4806
12 changed files with 642 additions and 96 deletions
@@ -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)
@@ -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: