Leonid Pershin
ef743a6e6d
- Introduced the `file_size_bytes` function to convert Civitai model sizes from kilobytes to bytes, improving data handling. - Updated the `seed_civitai` function to include file size in job definitions, enhancing model processing efficiency. - Enhanced the `should_skip` function to utilize expected size for faster decision-making during job processing. - Added tests for new functionality, ensuring robustness in file size handling and job processing logic.
348 lines
13 KiB
Python
348 lines
13 KiB
Python
"""Civitai site API. Bearer only on civitai.com / .red / .green."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
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 file_size_bytes(info: dict) -> int | None:
|
|
"""Civitai ``sizeKB`` (float) → bytes, or None."""
|
|
raw = info.get("sizeKB")
|
|
if raw is None:
|
|
raw = info.get("sizeKb")
|
|
if raw is None:
|
|
return None
|
|
try:
|
|
return int(round(float(raw) * 1024.0))
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def pick_preview_image(version: dict) -> tuple[str, str] | None:
|
|
"""First usable Civitai preview → (url, sidecar_suffix) for SwarmUI.
|
|
|
|
SwarmUI looks for ``stem.preview.jpg`` / ``.preview.png`` / ``.jpg`` etc.
|
|
next to the weight (T2IModelHandler AutoImageFormatSuffixes).
|
|
"""
|
|
images = version.get("images") or []
|
|
if not isinstance(images, list):
|
|
return None
|
|
for item in images:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
kind = str(item.get("type") or "image").lower()
|
|
if kind and kind not in {"image", "img", ""}:
|
|
continue
|
|
url = str(item.get("url") or "").strip()
|
|
if not url.startswith("http"):
|
|
continue
|
|
path = url.split("?", 1)[0].lower()
|
|
if path.endswith(".png"):
|
|
suffix = ".preview.png"
|
|
elif path.endswith(".jpeg") or path.endswith(".jpg"):
|
|
suffix = ".preview.jpg"
|
|
elif path.endswith(".webp"):
|
|
# SwarmUI auto-format list has no .preview.webp — use .jpg name;
|
|
# most Civitai CDN URLs are jpeg without extension.
|
|
suffix = ".preview.jpg"
|
|
else:
|
|
suffix = ".preview.jpg"
|
|
return url, suffix
|
|
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 _sha_from_version(version: dict) -> list[str]:
|
|
"""All SHA256 hashes listed on a version's files (lower)."""
|
|
out: list[str] = []
|
|
for f in version.get("files") or []:
|
|
if not isinstance(f, dict):
|
|
continue
|
|
hashes = f.get("hashes") or {}
|
|
if not isinstance(hashes, dict):
|
|
continue
|
|
sha = hashes.get("SHA256") or hashes.get("sha256")
|
|
if sha:
|
|
out.append(str(sha).strip().lower())
|
|
return out
|
|
|
|
|
|
def fetch_model_versions_by_hashes(
|
|
token: str | None,
|
|
host: str,
|
|
sha256_list: list[str],
|
|
*,
|
|
timeout: float = 60.0,
|
|
chunk_size: int = 100,
|
|
) -> dict[str, dict]:
|
|
"""POST /api/v1/model-versions/by-hash (≤100). Map sha256(lower) → version.
|
|
|
|
Unmatched hashes are omitted (caller treats as unknown). Retries once on 429.
|
|
If response versions lack files[].hashes, falls back to /by-hash/ids + GET version.
|
|
"""
|
|
cleaned: list[str] = []
|
|
seen_h: set[str] = set()
|
|
for raw in sha256_list:
|
|
digest = str(raw).strip().lower()
|
|
if len(digest) == 64 and all(c in "0123456789abcdef" for c in digest):
|
|
if digest not in seen_h:
|
|
cleaned.append(digest)
|
|
seen_h.add(digest)
|
|
if not cleaned:
|
|
return {}
|
|
|
|
first = _normalize_host(host)
|
|
order = [first, other_host(first)]
|
|
headers: dict[str, str] = {"Content-Type": "application/json"}
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
result: dict[str, dict] = {}
|
|
|
|
def _post(candidate: str, path: str, body: list[str]) -> httpx.Response | None:
|
|
url = f"https://{candidate}/api/v1/{path}"
|
|
last: httpx.Response | None = None
|
|
for attempt in range(2):
|
|
try:
|
|
with httpx.Client(timeout=timeout, follow_redirects=True) as client:
|
|
last = client.post(url, headers=headers, json=body)
|
|
except httpx.HTTPError:
|
|
return None
|
|
if last.status_code == 429 and attempt == 0:
|
|
time.sleep(2.0)
|
|
continue
|
|
return last
|
|
return last
|
|
|
|
for i in range(0, len(cleaned), chunk_size):
|
|
chunk = cleaned[i : i + chunk_size]
|
|
last_error = "нет ответа"
|
|
ok = False
|
|
for candidate in order:
|
|
if candidate not in ALLOWED_HOSTS:
|
|
continue
|
|
response = _post(candidate, "model-versions/by-hash", chunk)
|
|
if response is None:
|
|
last_error = "network"
|
|
continue
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
if isinstance(data, list):
|
|
for version in data:
|
|
if not isinstance(version, dict) or version.get("id") is None:
|
|
continue
|
|
for sha in _sha_from_version(version):
|
|
if sha in seen_h:
|
|
result.setdefault(sha, version)
|
|
ok = True
|
|
break
|
|
last_error = "не list"
|
|
continue
|
|
last_error = f"HTTP {response.status_code}"
|
|
if response.status_code not in {404, 400}:
|
|
continue
|
|
if not ok:
|
|
raise CloudError(
|
|
f"Civitai by-hash batch ({len(chunk)}): {last_error} "
|
|
f"(хосты {', '.join(order)})"
|
|
)
|
|
|
|
missing = [h for h in chunk if h not in result]
|
|
if not missing:
|
|
continue
|
|
# Fallback: ids endpoint maps hash → versionId, then GET version for modelId.
|
|
for candidate in order:
|
|
if candidate not in ALLOWED_HOSTS:
|
|
continue
|
|
response = _post(candidate, "model-versions/by-hash/ids", missing)
|
|
if response is None or response.status_code != 200:
|
|
continue
|
|
pairs = response.json()
|
|
if not isinstance(pairs, list):
|
|
continue
|
|
for pair in pairs:
|
|
if not isinstance(pair, dict):
|
|
continue
|
|
h = str(pair.get("hash") or "").strip().lower()
|
|
try:
|
|
vid = int(pair.get("modelVersionId"))
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if h not in seen_h:
|
|
continue
|
|
try:
|
|
_host, version = fetch_model_version(token or "", candidate, vid)
|
|
except CloudError:
|
|
continue
|
|
result.setdefault(h, version)
|
|
break
|
|
|
|
return result
|
|
|
|
|
|
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
|