Add Civitai Krea2 metadata scrape for train JSONL and Assistent FTS search.

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>
2026-08-23 08:13:53 +03:00
parent b83d1d1e9c
commit 6dfbf4e702
17 changed files with 1548 additions and 14 deletions
@@ -345,3 +345,230 @@ def version_ids_from_payload(version: dict) -> tuple[int | None, int | 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