Leonid PershinandCursor
6dfbf4e702
Local civitai-dataset launchers collect ~2000 prompt/params rows without images; search.jsonl is pushed on up for cheap example lookup. Co-authored-by: Cursor <cursoragent@cursor.com>
575 lines
20 KiB
Python
575 lines
20 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
|
||
|
||
|
||
def _auth_headers(token: str | None) -> dict[str, str]:
|
||
return {"Authorization": f"Bearer {token}"} if token else {}
|
||
|
||
|
||
def _get_json(
|
||
token: str | None,
|
||
host: str,
|
||
path: str,
|
||
params: dict[str, str | int | bool | None],
|
||
*,
|
||
timeout: float = 60.0,
|
||
retries_429: int = 2,
|
||
) -> tuple[str, dict | list]:
|
||
"""GET /api/v1/{path} with host failover and 429 backoff. Returns (host, json)."""
|
||
first = _normalize_host(host)
|
||
order = [first, other_host(first)]
|
||
cleaned = {k: v for k, v in params.items() if v is not None and v != ""}
|
||
last_error = "нет ответа"
|
||
seen: set[str] = set()
|
||
headers = _auth_headers(token)
|
||
for candidate in order:
|
||
if candidate in seen or candidate not in ALLOWED_HOSTS:
|
||
continue
|
||
seen.add(candidate)
|
||
url = f"https://{candidate}/api/v1/{path.lstrip('/')}"
|
||
for attempt in range(retries_429 + 1):
|
||
try:
|
||
with httpx.Client(timeout=timeout, follow_redirects=True) as client:
|
||
response = client.get(url, headers=headers, params=cleaned)
|
||
except httpx.HTTPError as exc:
|
||
last_error = str(exc)
|
||
break
|
||
if response.status_code == 429 and attempt < retries_429:
|
||
time.sleep(2.0 * (attempt + 1))
|
||
continue
|
||
if response.status_code == 200:
|
||
data = response.json()
|
||
if isinstance(data, (dict, list)):
|
||
return candidate, data
|
||
last_error = "неожиданный JSON"
|
||
break
|
||
detail = (response.text or response.reason_phrase or "")[:120].replace("\n", " ")
|
||
last_error = f"HTTP {response.status_code}" + (f" {detail}" if detail else "")
|
||
if response.status_code not in {404, 400}:
|
||
break
|
||
break
|
||
raise CloudError(f"Civitai {path}: {last_error} (хосты {', '.join(seen)})")
|
||
|
||
|
||
def list_models(
|
||
token: str | None,
|
||
host: str,
|
||
*,
|
||
types: str | None = None,
|
||
query: str | None = None,
|
||
sort: str = "Most Downloaded",
|
||
period: str = "AllTime",
|
||
limit: int = 100,
|
||
page: int | None = None,
|
||
cursor: str | None = None,
|
||
nsfw: bool | str | None = None,
|
||
timeout: float = 60.0,
|
||
) -> tuple[str, dict]:
|
||
"""GET /api/v1/models. Prefer cursor for deep pages; page*limit capped ~1000.
|
||
|
||
Do not pass ``nsfw=True`` — several hosts return HTTP 400 for boolean nsfw on /models.
|
||
"""
|
||
params: dict[str, str | int | bool | None] = {
|
||
"limit": max(1, min(int(limit), 100)),
|
||
"sort": sort,
|
||
"period": period,
|
||
"types": types,
|
||
"query": query,
|
||
}
|
||
if nsfw is not None:
|
||
# String form only; boolean True often 400s on /models.
|
||
params["nsfw"] = "true" if nsfw is True else ("false" if nsfw is False else nsfw)
|
||
if cursor:
|
||
params["cursor"] = cursor
|
||
elif page is not None and not query:
|
||
# Query search rejects page= — cursor only (or omit for first page).
|
||
params["page"] = int(page)
|
||
host_used, data = _get_json(token, host, "models", params, timeout=timeout)
|
||
if not isinstance(data, dict):
|
||
raise CloudError("Civitai models: ожидался object с items")
|
||
return host_used, data
|
||
|
||
|
||
def list_images(
|
||
token: str | None,
|
||
host: str,
|
||
*,
|
||
model_version_id: int | None = None,
|
||
model_id: int | None = None,
|
||
sort: str = "Most Reactions",
|
||
period: str = "AllTime",
|
||
limit: int = 100,
|
||
cursor: str | None = None,
|
||
page: int | None = None,
|
||
with_meta: bool = True,
|
||
nsfw: str | bool | None = "X",
|
||
timeout: float = 60.0,
|
||
) -> tuple[str, dict]:
|
||
"""GET /api/v1/images. Pass modelVersionId alone (not with modelId) so sort works."""
|
||
params: dict[str, str | int | bool | None] = {
|
||
"limit": max(1, min(int(limit), 200)),
|
||
"sort": sort,
|
||
"period": period,
|
||
"withMeta": "true" if with_meta else "false",
|
||
"nsfw": nsfw,
|
||
}
|
||
if model_version_id is not None:
|
||
params["modelVersionId"] = int(model_version_id)
|
||
elif model_id is not None:
|
||
params["modelId"] = int(model_id)
|
||
if cursor:
|
||
params["cursor"] = cursor
|
||
elif page is not None:
|
||
params["page"] = int(page)
|
||
host_used, data = _get_json(token, host, "images", params, timeout=timeout)
|
||
if not isinstance(data, dict):
|
||
raise CloudError("Civitai images: ожидался object с items")
|
||
return host_used, data
|
||
|
||
|
||
def iter_models_pages(
|
||
token: str | None,
|
||
host: str,
|
||
*,
|
||
types: str,
|
||
query: str,
|
||
sort: str = "Most Downloaded",
|
||
period: str = "AllTime",
|
||
limit: int = 100,
|
||
max_pages: int = 5,
|
||
nsfw: bool | str | None = None,
|
||
timeout: float = 60.0,
|
||
):
|
||
"""Yield model item dicts across page/cursor pagination."""
|
||
cursor: str | None = None
|
||
page = 1
|
||
for _ in range(max_pages):
|
||
_h, data = list_models(
|
||
token,
|
||
host,
|
||
types=types,
|
||
query=query,
|
||
sort=sort,
|
||
period=period,
|
||
limit=limit,
|
||
page=None if (cursor or query) else page,
|
||
cursor=cursor,
|
||
nsfw=nsfw,
|
||
timeout=timeout,
|
||
)
|
||
items = data.get("items") or []
|
||
if not isinstance(items, list):
|
||
break
|
||
for item in items:
|
||
if isinstance(item, dict):
|
||
yield item
|
||
meta = data.get("metadata") if isinstance(data.get("metadata"), dict) else {}
|
||
next_cursor = meta.get("nextCursor")
|
||
if next_cursor:
|
||
cursor = str(next_cursor)
|
||
continue
|
||
if len(items) < limit:
|
||
break
|
||
if cursor is None and not query:
|
||
page += 1
|
||
if page * limit > 1000:
|
||
break
|
||
else:
|
||
break
|
||
|
||
|
||
def iter_images_pages(
|
||
token: str | None,
|
||
host: str,
|
||
*,
|
||
model_version_id: int,
|
||
sort: str = "Most Reactions",
|
||
period: str = "AllTime",
|
||
limit: int = 100,
|
||
max_pages: int = 50,
|
||
with_meta: bool = True,
|
||
nsfw: str | bool | None = "X",
|
||
timeout: float = 60.0,
|
||
):
|
||
"""Yield image item dicts for one modelVersionId (cursor preferred)."""
|
||
cursor: str | None = None
|
||
page = 1
|
||
for _ in range(max_pages):
|
||
_h, data = list_images(
|
||
token,
|
||
host,
|
||
model_version_id=model_version_id,
|
||
sort=sort,
|
||
period=period,
|
||
limit=limit,
|
||
cursor=cursor,
|
||
page=None if cursor else page,
|
||
with_meta=with_meta,
|
||
nsfw=nsfw,
|
||
timeout=timeout,
|
||
)
|
||
items = data.get("items") or []
|
||
if not isinstance(items, list):
|
||
break
|
||
for item in items:
|
||
if isinstance(item, dict):
|
||
yield item
|
||
meta = data.get("metadata") if isinstance(data.get("metadata"), dict) else {}
|
||
next_cursor = meta.get("nextCursor")
|
||
if next_cursor:
|
||
cursor = str(next_cursor)
|
||
continue
|
||
if len(items) < limit:
|
||
break
|
||
if cursor is None:
|
||
page += 1
|
||
if page * limit > 1000:
|
||
break
|
||
else:
|
||
break
|