"""Civitai site API. Bearer only on civitai.com / .red / .green.""" from __future__ import annotations from dataclasses import dataclass import httpx from gpu_rent.errors import CloudError ALLOWED_HOSTS = ("civitai.com", "civitai.red", "civitai.green") @dataclass class CivitaiProbe: host: str ok: bool status: int | None detail: str def _normalize_host(host: str) -> str: h = host.strip().lower().removeprefix("https://").removeprefix("http://").split("/")[0] if h.startswith("www."): h = h[4:] return h def other_host(host: str) -> str: h = _normalize_host(host) if h.endswith(".red") or h == "civitai.red": return "civitai.com" return "civitai.red" def probe_me(token: str, host: str, timeout: float = 15.0) -> CivitaiProbe: host = _normalize_host(host) if host not in ALLOWED_HOSTS: return CivitaiProbe(host=host, ok=False, status=None, detail="хост не из allow-list") url = f"https://{host}/api/v1/me" 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 CivitaiProbe(host=host, ok=False, status=None, detail=str(exc)) if response.status_code == 200: return CivitaiProbe(host=host, ok=True, status=200, detail="токен принят") if response.status_code in {401, 403}: return CivitaiProbe( host=host, ok=False, status=response.status_code, detail="токен отвергнут — перевыпусти ключ на civitai.com/user/account", ) return CivitaiProbe( host=host, ok=False, status=response.status_code, detail=response.text[:200] or response.reason_phrase, ) def pick_primary_file(version: dict) -> dict | None: files = version.get("files") or [] if not isinstance(files, list): return None for item in files: if isinstance(item, dict) and item.get("primary"): return item for item in files: if not isinstance(item, dict): continue name = str(item.get("name") or "").lower() if "safetensor" in name or name.endswith(".safetensors"): return item for item in files: if isinstance(item, dict): return item return None def fetch_model_version(token: str, host: str, version_id: int, timeout: float = 30.0) -> tuple[str, dict]: """GET /api/v1/model-versions/{id}; one retry on the other Civitai host.""" first = _normalize_host(host) order = [first, other_host(first)] last_error = "нет ответа" seen: set[str] = set() for candidate in order: if candidate in seen or candidate not in ALLOWED_HOSTS: continue seen.add(candidate) url = f"https://{candidate}/api/v1/model-versions/{version_id}" try: with httpx.Client(timeout=timeout, follow_redirects=True) as client: headers = {"Authorization": f"Bearer {token}"} if token else {} response = client.get(url, headers=headers) except httpx.HTTPError as exc: last_error = str(exc) continue if response.status_code == 200: data = response.json() if isinstance(data, dict) and pick_primary_file(data): return candidate, data last_error = "пустой files[]" continue last_error = f"HTTP {response.status_code}" if response.status_code not in {404, 400}: break raise CloudError(f"Civitai version {version_id}: {last_error} (хосты {', '.join(seen)})") def fetch_model_version_by_hash( token: str | None, host: str, sha256: str, timeout: float = 30.0, ) -> tuple[str, dict]: """GET /api/v1/model-versions/by-hash/{sha}; public, token optional for NSFW/region.""" digest = sha256.strip().lower() if len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest): raise CloudError(f"Civitai by-hash: неверный SHA256 ({sha256[:16]}…)") first = _normalize_host(host) order = [first, other_host(first)] last_error = "нет ответа" seen: set[str] = set() headers = {"Authorization": f"Bearer {token}"} if token else {} for candidate in order: if candidate in seen or candidate not in ALLOWED_HOSTS: continue seen.add(candidate) url = f"https://{candidate}/api/v1/model-versions/by-hash/{digest}" try: with httpx.Client(timeout=timeout, follow_redirects=True) as client: response = client.get(url, headers=headers) except httpx.HTTPError as exc: last_error = str(exc) continue if response.status_code == 200: data = response.json() if isinstance(data, dict) and data.get("id") is not None: return candidate, data last_error = "пустой ответ" continue last_error = f"HTTP {response.status_code}" if response.status_code not in {404, 400}: break raise CloudError(f"Civitai by-hash {digest[:12]}…: {last_error} (хосты {', '.join(seen)})") def civitai_model_url(model_id: int, version_id: int, host: str = "civitai.red") -> str: """Canonical manifest URL (links only — no download).""" h = _normalize_host(host) if h not in ALLOWED_HOSTS: h = "civitai.red" if h == "civitai.green": h = "civitai.com" return f"https://{h}/models/{int(model_id)}?modelVersionId={int(version_id)}" def version_ids_from_payload(version: dict) -> tuple[int | None, int | None]: """Extract (version_id, model_id) from a Civitai version JSON object.""" try: vid = int(version["id"]) if version.get("id") is not None else None except (TypeError, ValueError, KeyError): vid = None mid = version.get("modelId") if mid is None and isinstance(version.get("model"), dict): mid = version["model"].get("id") try: model_id = int(mid) if mid is not None else None except (TypeError, ValueError): model_id = None return vid, model_id