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>
This commit is contained in:
Leonid Pershin
Cursor
@@ -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
|
||||
|
||||
@@ -0,0 +1,731 @@
|
||||
"""Local Civitai Krea2 gallery scrape → train.jsonl + search.jsonl (no image files)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import Callable, Iterable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from gpu_rent.civitai import (
|
||||
fetch_model_version,
|
||||
iter_images_pages,
|
||||
iter_models_pages,
|
||||
list_models,
|
||||
)
|
||||
from gpu_rent.config import load_config
|
||||
from gpu_rent.errors import CloudError, GpuRentError
|
||||
from gpu_rent.manifests import parse_models
|
||||
from gpu_rent.paths import app_root, models_manifest_path
|
||||
|
||||
Log = Callable[[str], None]
|
||||
|
||||
DEFAULT_TARGET = 2000
|
||||
DEFAULT_TOP_CHECKPOINTS = 20
|
||||
DEFAULT_TOP_LORAS = 30
|
||||
DEFAULT_PER_VERSION = 150
|
||||
DEFAULT_MIN_SCORE = 5
|
||||
KREA_QUERIES = ("krea2", "krea-2", "krea")
|
||||
MINOR_TAG_RE = re.compile(
|
||||
r"\b(loli|shota|lolicon|shotacon|underage|child\b|preteen|under.?18)\b",
|
||||
re.I,
|
||||
)
|
||||
NSFW_TO_RATING = {
|
||||
"none": "pg",
|
||||
"0": "pg",
|
||||
"pg": "pg",
|
||||
"soft": "pg13",
|
||||
"1": "pg13",
|
||||
"pg13": "pg13",
|
||||
"mature": "r",
|
||||
"2": "r",
|
||||
"r": "r",
|
||||
"x": "x",
|
||||
"4": "x",
|
||||
"xxx": "xxx",
|
||||
"8": "xxx",
|
||||
"16": "xxx",
|
||||
}
|
||||
|
||||
|
||||
def dataset_root(root: Path | None = None) -> Path:
|
||||
return (root or app_root()) / "datasets" / "civitai"
|
||||
|
||||
|
||||
def catalog_dir(root: Path | None = None) -> Path:
|
||||
return dataset_root(root) / "catalog"
|
||||
|
||||
|
||||
def _log(msg: str) -> None:
|
||||
print(msg, flush=True)
|
||||
|
||||
|
||||
def _is_krea_version(version: dict, model_name: str = "") -> bool:
|
||||
base = str(version.get("baseModel") or "").lower()
|
||||
name = str(version.get("name") or "").lower()
|
||||
model = (model_name or "").lower()
|
||||
blob = f"{base} {name} {model}"
|
||||
return "krea" in blob
|
||||
|
||||
|
||||
def _ours_version_ids(manifest: Path | None) -> set[int]:
|
||||
path = manifest or models_manifest_path()
|
||||
if not path.is_file():
|
||||
example = app_root() / "models.example.yaml"
|
||||
path = example if example.is_file() else path
|
||||
if not path.is_file():
|
||||
return set()
|
||||
try:
|
||||
entries = parse_models(path)
|
||||
except Exception:
|
||||
return set()
|
||||
out: set[int] = set()
|
||||
for e in entries:
|
||||
if e.version_id:
|
||||
out.add(int(e.version_id))
|
||||
return out
|
||||
|
||||
|
||||
def discover_krea_models(
|
||||
token: str,
|
||||
host: str,
|
||||
*,
|
||||
top_checkpoints: int = DEFAULT_TOP_CHECKPOINTS,
|
||||
top_loras: int = DEFAULT_TOP_LORAS,
|
||||
ours: set[int] | None = None,
|
||||
log: Log = _log,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Find popular Krea2 checkpoints and LoRAs; return catalog model rows."""
|
||||
ours = ours if ours is not None else _ours_version_ids(None)
|
||||
found: dict[int, dict[str, Any]] = {}
|
||||
|
||||
def _ingest(kind: str, item: dict, cap: int) -> None:
|
||||
if sum(1 for r in found.values() if r["kind"] == kind) >= cap:
|
||||
return
|
||||
model_id = item.get("id")
|
||||
name = str(item.get("name") or "")
|
||||
versions = item.get("modelVersions") or []
|
||||
if not isinstance(versions, list):
|
||||
return
|
||||
model_is_krea = "krea" in name.lower()
|
||||
for ver in versions:
|
||||
if not isinstance(ver, dict):
|
||||
continue
|
||||
if not model_is_krea and not _is_krea_version(ver, name):
|
||||
continue
|
||||
try:
|
||||
vid = int(ver["id"])
|
||||
mid = int(model_id if model_id is not None else ver.get("modelId") or 0)
|
||||
except (TypeError, ValueError, KeyError):
|
||||
continue
|
||||
if mid <= 0 or vid <= 0:
|
||||
continue
|
||||
if vid in found:
|
||||
continue
|
||||
if sum(1 for r in found.values() if r["kind"] == kind) >= cap:
|
||||
return
|
||||
stats = item.get("stats") if isinstance(item.get("stats"), dict) else {}
|
||||
found[vid] = {
|
||||
"kind": kind,
|
||||
"modelId": mid,
|
||||
"modelVersionId": vid,
|
||||
"name": name,
|
||||
"versionName": str(ver.get("name") or ""),
|
||||
"baseModel": str(ver.get("baseModel") or ""),
|
||||
"downloadCount": int(stats.get("downloadCount") or 0),
|
||||
"ours": vid in ours,
|
||||
"trainedWords": list(ver.get("trainedWords") or [])
|
||||
if isinstance(ver.get("trainedWords"), list)
|
||||
else [],
|
||||
}
|
||||
# One version per model is enough for gallery diversity.
|
||||
break
|
||||
|
||||
for kind, types, cap in (
|
||||
("checkpoint", "Checkpoint", top_checkpoints),
|
||||
("lora", "LORA", top_loras),
|
||||
):
|
||||
for query in KREA_QUERIES:
|
||||
if sum(1 for r in found.values() if r["kind"] == kind) >= cap:
|
||||
break
|
||||
log(f"discover: {types} query={query!r}")
|
||||
try:
|
||||
for item in iter_models_pages(
|
||||
token,
|
||||
host,
|
||||
types=types,
|
||||
query=query,
|
||||
sort="Most Downloaded",
|
||||
limit=100,
|
||||
max_pages=3,
|
||||
nsfw=None,
|
||||
):
|
||||
_ingest(kind, item, cap)
|
||||
if sum(1 for r in found.values() if r["kind"] == kind) >= cap:
|
||||
break
|
||||
except CloudError as exc:
|
||||
log(f"discover warn: {exc}")
|
||||
# Also try Highest Rated once
|
||||
try:
|
||||
_h, data = list_models(
|
||||
token,
|
||||
host,
|
||||
types=types,
|
||||
query=query,
|
||||
sort="Highest Rated",
|
||||
limit=50,
|
||||
page=None,
|
||||
nsfw=None,
|
||||
)
|
||||
for item in data.get("items") or []:
|
||||
if isinstance(item, dict):
|
||||
_ingest(kind, item, cap)
|
||||
except CloudError as exc:
|
||||
log(f"discover warn (rated): {exc}")
|
||||
|
||||
# Always include version_ids from local models.yaml / models.example.yaml.
|
||||
for vid in sorted(ours):
|
||||
if vid in found:
|
||||
found[vid]["ours"] = True
|
||||
continue
|
||||
try:
|
||||
_h, ver = fetch_model_version(token, host, vid)
|
||||
except CloudError as exc:
|
||||
log(f"discover ours warn version {vid}: {exc}")
|
||||
continue
|
||||
mid = ver.get("modelId")
|
||||
model = ver.get("model") if isinstance(ver.get("model"), dict) else {}
|
||||
name = str(model.get("name") or ver.get("name") or vid)
|
||||
mtype = str(model.get("type") or "").lower()
|
||||
kind = "lora" if mtype in {"lora", "locon", "dora"} else "checkpoint"
|
||||
try:
|
||||
mid_i = int(mid) if mid is not None else 0
|
||||
except (TypeError, ValueError):
|
||||
mid_i = 0
|
||||
if mid_i <= 0:
|
||||
continue
|
||||
found[vid] = {
|
||||
"kind": kind,
|
||||
"modelId": mid_i,
|
||||
"modelVersionId": vid,
|
||||
"name": name,
|
||||
"versionName": str(ver.get("name") or ""),
|
||||
"baseModel": str(ver.get("baseModel") or ""),
|
||||
"downloadCount": 0,
|
||||
"ours": True,
|
||||
"trainedWords": list(ver.get("trainedWords") or [])
|
||||
if isinstance(ver.get("trainedWords"), list)
|
||||
else [],
|
||||
}
|
||||
|
||||
rows = sorted(
|
||||
found.values(),
|
||||
key=lambda r: (-int(r.get("downloadCount") or 0), r["modelVersionId"]),
|
||||
)
|
||||
log(
|
||||
f"discover: {sum(1 for r in rows if r['kind']=='checkpoint')} ckpt, "
|
||||
f"{sum(1 for r in rows if r['kind']=='lora')} lora "
|
||||
f"({sum(1 for r in rows if r.get('ours'))} ours)"
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def rating_from_nsfw(level: Any) -> str:
|
||||
if level is None:
|
||||
return "pg"
|
||||
if isinstance(level, bool):
|
||||
return "x" if level else "pg"
|
||||
key = str(level).strip().lower()
|
||||
return NSFW_TO_RATING.get(key, "pg13" if key not in {"false", ""} else "pg")
|
||||
|
||||
|
||||
def reaction_score(stats: dict | None) -> int:
|
||||
if not isinstance(stats, dict):
|
||||
return 0
|
||||
total = 0
|
||||
for key in ("likeCount", "heartCount", "laughCount", "cryCount"):
|
||||
try:
|
||||
total += int(stats.get(key) or 0)
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
return total
|
||||
|
||||
|
||||
def tags_from_image(item: dict, prompt: str, *, cap: int = 20) -> list[str]:
|
||||
tags: list[str] = []
|
||||
raw = item.get("tags")
|
||||
if isinstance(raw, list):
|
||||
for t in raw:
|
||||
if isinstance(t, str) and t.strip():
|
||||
tags.append(t.strip())
|
||||
elif isinstance(t, dict):
|
||||
name = str(t.get("name") or t.get("tag") or "").strip()
|
||||
if name:
|
||||
tags.append(name)
|
||||
if not tags and prompt:
|
||||
for part in re.split(r"[,.\n]", prompt):
|
||||
word = part.strip().strip("<>()[]{}\"'")
|
||||
if len(word) < 2 or len(word) > 48:
|
||||
continue
|
||||
if word.lower().startswith("lora:"):
|
||||
continue
|
||||
if word not in tags:
|
||||
tags.append(word)
|
||||
if len(tags) >= cap:
|
||||
break
|
||||
return tags[:cap]
|
||||
|
||||
|
||||
def looks_minor(prompt: str, tags: Iterable[str]) -> bool:
|
||||
blob = " ".join([prompt or "", *tags]).lower()
|
||||
return bool(MINOR_TAG_RE.search(blob))
|
||||
|
||||
|
||||
def normalize_image(
|
||||
item: dict,
|
||||
*,
|
||||
kind: str,
|
||||
model_id: int,
|
||||
model_version_id: int,
|
||||
ours: bool,
|
||||
min_score: int = DEFAULT_MIN_SCORE,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Map API image → catalog row, or None if skip."""
|
||||
try:
|
||||
image_id = int(item["id"])
|
||||
except (TypeError, ValueError, KeyError):
|
||||
return None
|
||||
media = str(item.get("type") or "image").lower()
|
||||
if media and media not in {"image", "img", ""}:
|
||||
return None
|
||||
meta = item.get("meta") if isinstance(item.get("meta"), dict) else {}
|
||||
prompt = str(meta.get("prompt") or "").strip()
|
||||
if not prompt:
|
||||
return None
|
||||
tags = tags_from_image(item, prompt)
|
||||
if looks_minor(prompt, tags):
|
||||
return None
|
||||
stats = item.get("stats") if isinstance(item.get("stats"), dict) else {}
|
||||
score = reaction_score(stats)
|
||||
if score < min_score:
|
||||
return None
|
||||
neg = str(meta.get("negativePrompt") or meta.get("negative_prompt") or "").strip()
|
||||
params: dict[str, Any] = {}
|
||||
for src, dst in (
|
||||
("steps", "steps"),
|
||||
("cfgScale", "cfgScale"),
|
||||
("sampler", "sampler"),
|
||||
("seed", "seed"),
|
||||
("Size", "size"),
|
||||
("clipSkip", "clipSkip"),
|
||||
("scheduler", "scheduler"),
|
||||
):
|
||||
if meta.get(src) is not None:
|
||||
params[dst] = meta[src]
|
||||
w, h = item.get("width"), item.get("height")
|
||||
if w is not None and h is not None:
|
||||
params.setdefault("width", w)
|
||||
params.setdefault("height", h)
|
||||
resources = meta.get("civitaiResources")
|
||||
if not isinstance(resources, list):
|
||||
resources = []
|
||||
return {
|
||||
"id": image_id,
|
||||
"url": str(item.get("url") or ""),
|
||||
"username": str(item.get("username") or ""),
|
||||
"createdAt": str(item.get("createdAt") or ""),
|
||||
"kind": kind,
|
||||
"modelId": model_id,
|
||||
"modelVersionId": model_version_id,
|
||||
"ours": bool(ours),
|
||||
"nsfwLevel": item.get("nsfwLevel"),
|
||||
"rating": rating_from_nsfw(item.get("nsfwLevel")),
|
||||
"stats": {
|
||||
"likeCount": int(stats.get("likeCount") or 0),
|
||||
"heartCount": int(stats.get("heartCount") or 0),
|
||||
"laughCount": int(stats.get("laughCount") or 0),
|
||||
"cryCount": int(stats.get("cryCount") or 0),
|
||||
"commentCount": int(stats.get("commentCount") or 0),
|
||||
},
|
||||
"score": score,
|
||||
"prompt": prompt,
|
||||
"negativePrompt": neg,
|
||||
"params": params,
|
||||
"resources": resources,
|
||||
"tags": tags,
|
||||
}
|
||||
|
||||
|
||||
def load_existing_ids(images_path: Path) -> set[int]:
|
||||
ids: set[int] = set()
|
||||
if not images_path.is_file():
|
||||
return ids
|
||||
with images_path.open(encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
row = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(row, dict) and row.get("id") is not None:
|
||||
try:
|
||||
ids.add(int(row["id"]))
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
return ids
|
||||
|
||||
|
||||
def append_jsonl(path: Path, rows: Iterable[dict]) -> int:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
n = 0
|
||||
with path.open("a", encoding="utf-8") as fh:
|
||||
for row in rows:
|
||||
fh.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
def read_jsonl(path: Path) -> list[dict[str, Any]]:
|
||||
if not path.is_file():
|
||||
return []
|
||||
out: list[dict[str, Any]] = []
|
||||
with path.open(encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
row = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(row, dict):
|
||||
out.append(row)
|
||||
return out
|
||||
|
||||
|
||||
def cmd_discover(
|
||||
*,
|
||||
out_root: Path | None = None,
|
||||
top_checkpoints: int = DEFAULT_TOP_CHECKPOINTS,
|
||||
top_loras: int = DEFAULT_TOP_LORAS,
|
||||
log: Log = _log,
|
||||
) -> Path:
|
||||
cfg = load_config()
|
||||
token = (cfg.civitai_api_token or "").strip()
|
||||
if not token:
|
||||
raise GpuRentError("Нужен CIVITAI_API_TOKEN в .env")
|
||||
host = cfg.civitai_api_host or "civitai.red"
|
||||
rows = discover_krea_models(
|
||||
token,
|
||||
host,
|
||||
top_checkpoints=top_checkpoints,
|
||||
top_loras=top_loras,
|
||||
log=log,
|
||||
)
|
||||
cat = catalog_dir(out_root)
|
||||
cat.mkdir(parents=True, exist_ok=True)
|
||||
path = cat / "models.json"
|
||||
path.write_text(json.dumps(rows, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
log(f"wrote {path} ({len(rows)} models)")
|
||||
return path
|
||||
|
||||
|
||||
def cmd_scrape(
|
||||
*,
|
||||
out_root: Path | None = None,
|
||||
target: int = DEFAULT_TARGET,
|
||||
per_version: int = DEFAULT_PER_VERSION,
|
||||
min_score: int = DEFAULT_MIN_SCORE,
|
||||
log: Log = _log,
|
||||
) -> int:
|
||||
cfg = load_config()
|
||||
token = (cfg.civitai_api_token or "").strip()
|
||||
if not token:
|
||||
raise GpuRentError("Нужен CIVITAI_API_TOKEN в .env")
|
||||
host = cfg.civitai_api_host or "civitai.red"
|
||||
cat = catalog_dir(out_root)
|
||||
models_path = cat / "models.json"
|
||||
if not models_path.is_file():
|
||||
raise GpuRentError(f"Нет {models_path} — сначала discover")
|
||||
models = json.loads(models_path.read_text(encoding="utf-8"))
|
||||
if not isinstance(models, list):
|
||||
raise GpuRentError("models.json: ожидался list")
|
||||
images_path = cat / "images.jsonl"
|
||||
seen = load_existing_ids(images_path)
|
||||
have = len(seen)
|
||||
log(f"scrape: already {have}/{target} in {images_path}")
|
||||
if have >= target:
|
||||
log("scrape: target reached")
|
||||
return have
|
||||
|
||||
for m in models:
|
||||
if not isinstance(m, dict):
|
||||
continue
|
||||
if have >= target:
|
||||
break
|
||||
try:
|
||||
vid = int(m["modelVersionId"])
|
||||
mid = int(m["modelId"])
|
||||
except (TypeError, ValueError, KeyError):
|
||||
continue
|
||||
kind = str(m.get("kind") or "checkpoint")
|
||||
ours = bool(m.get("ours"))
|
||||
name = str(m.get("name") or vid)
|
||||
budget = min(per_version, target - have)
|
||||
log(f"scrape: {kind} {name} version={vid} (need {budget})")
|
||||
batch: list[dict] = []
|
||||
taken = 0
|
||||
try:
|
||||
for item in iter_images_pages(
|
||||
token,
|
||||
host,
|
||||
model_version_id=vid,
|
||||
limit=100,
|
||||
max_pages=max(1, (per_version // 100) + 2),
|
||||
with_meta=True,
|
||||
nsfw="X",
|
||||
):
|
||||
if have + len(batch) >= target or taken >= per_version:
|
||||
break
|
||||
try:
|
||||
iid = int(item["id"])
|
||||
except (TypeError, ValueError, KeyError):
|
||||
continue
|
||||
if iid in seen:
|
||||
continue
|
||||
row = normalize_image(
|
||||
item,
|
||||
kind=kind,
|
||||
model_id=mid,
|
||||
model_version_id=vid,
|
||||
ours=ours,
|
||||
min_score=min_score,
|
||||
)
|
||||
if row is None:
|
||||
continue
|
||||
seen.add(iid)
|
||||
batch.append(row)
|
||||
taken += 1
|
||||
except CloudError as exc:
|
||||
log(f"scrape warn version {vid}: {exc}")
|
||||
continue
|
||||
if batch:
|
||||
append_jsonl(images_path, batch)
|
||||
have += len(batch)
|
||||
log(f"scrape: +{len(batch)} -> {have}/{target}")
|
||||
log(f"scrape done: {have} rows")
|
||||
return have
|
||||
|
||||
|
||||
def train_row(row: dict[str, Any]) -> dict[str, Any]:
|
||||
tags = row.get("tags") or []
|
||||
tag_s = ", ".join(str(t) for t in tags[:12])
|
||||
rating = row.get("rating") or "pg"
|
||||
instruction = f"Write a Krea 2 prompt.\nTags: {tag_s}\nRating: {rating}"
|
||||
params = row.get("params") if isinstance(row.get("params"), dict) else {}
|
||||
parts = [str(row.get("prompt") or "").strip()]
|
||||
neg = str(row.get("negativePrompt") or "").strip()
|
||||
if neg:
|
||||
parts.append(f"Negative: {neg}")
|
||||
for key, label in (
|
||||
("steps", "steps"),
|
||||
("cfgScale", "cfg"),
|
||||
("sampler", "sampler"),
|
||||
("seed", "seed"),
|
||||
("size", "size"),
|
||||
):
|
||||
if params.get(key) is not None:
|
||||
parts.append(f"{label}: {params[key]}")
|
||||
return {"instruction": instruction, "output": "\n".join(parts)}
|
||||
|
||||
|
||||
def search_row(row: dict[str, Any]) -> dict[str, Any]:
|
||||
resources = row.get("resources") if isinstance(row.get("resources"), list) else []
|
||||
loras: list[dict[str, Any]] = []
|
||||
for r in resources:
|
||||
if not isinstance(r, dict):
|
||||
continue
|
||||
rtype = str(r.get("type") or "").lower()
|
||||
if rtype and rtype not in {"lora", "locon", "dora"}:
|
||||
continue
|
||||
try:
|
||||
vid = int(r.get("modelVersionId"))
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
entry: dict[str, Any] = {"versionId": vid}
|
||||
if r.get("weight") is not None:
|
||||
entry["weight"] = r["weight"]
|
||||
loras.append(entry)
|
||||
params = row.get("params") if isinstance(row.get("params"), dict) else {}
|
||||
compact_params: dict[str, Any] = {}
|
||||
for key in ("steps", "cfgScale", "sampler", "seed", "width", "height", "size"):
|
||||
if params.get(key) is not None:
|
||||
compact_params["cfg" if key == "cfgScale" else key] = params[key]
|
||||
return {
|
||||
"id": row.get("id"),
|
||||
"rating": row.get("rating") or "pg",
|
||||
"score": int(row.get("score") or 0),
|
||||
"kind": row.get("kind"),
|
||||
"modelVersionId": row.get("modelVersionId"),
|
||||
"tags": list(row.get("tags") or []),
|
||||
"prompt": row.get("prompt") or "",
|
||||
"negative": row.get("negativePrompt") or "",
|
||||
"params": compact_params,
|
||||
"loras": loras,
|
||||
}
|
||||
|
||||
|
||||
def write_jsonl(path: Path, rows: Iterable[dict]) -> int:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
n = 0
|
||||
with path.open("w", encoding="utf-8") as fh:
|
||||
for row in rows:
|
||||
fh.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
def cmd_split(*, out_root: Path | None = None, log: Log = _log) -> dict[str, int]:
|
||||
root = dataset_root(out_root)
|
||||
cat = catalog_dir(out_root)
|
||||
images_path = cat / "images.jsonl"
|
||||
rows = read_jsonl(images_path)
|
||||
if not rows:
|
||||
raise GpuRentError(f"Пустой каталог {images_path} — сначала scrape")
|
||||
|
||||
by_kind: dict[str, list] = {"checkpoint": [], "lora": []}
|
||||
by_rating: dict[str, list] = {}
|
||||
train: list[dict] = []
|
||||
search: list[dict] = []
|
||||
for row in rows:
|
||||
kind = str(row.get("kind") or "checkpoint")
|
||||
by_kind.setdefault(kind, []).append(row)
|
||||
rating = str(row.get("rating") or "pg")
|
||||
by_rating.setdefault(rating, []).append(row)
|
||||
train.append(train_row(row))
|
||||
search.append(search_row(row))
|
||||
|
||||
counts: dict[str, int] = {}
|
||||
for kind, items in by_kind.items():
|
||||
path = root / "by_kind" / f"{kind}.jsonl"
|
||||
counts[f"kind:{kind}"] = write_jsonl(path, items)
|
||||
for rating, items in by_rating.items():
|
||||
safe = re.sub(r"[^a-z0-9]+", "", rating.lower()) or "pg"
|
||||
path = root / "by_rating" / f"{safe}.jsonl"
|
||||
counts[f"rating:{safe}"] = write_jsonl(path, items)
|
||||
counts["train"] = write_jsonl(root / "train.jsonl", train)
|
||||
counts["search"] = write_jsonl(root / "search.jsonl", search)
|
||||
log(
|
||||
"split: "
|
||||
+ ", ".join(f"{k}={v}" for k, v in sorted(counts.items()))
|
||||
+ f" → {root}"
|
||||
)
|
||||
return counts
|
||||
|
||||
|
||||
def cmd_all(
|
||||
*,
|
||||
out_root: Path | None = None,
|
||||
top_checkpoints: int = DEFAULT_TOP_CHECKPOINTS,
|
||||
top_loras: int = DEFAULT_TOP_LORAS,
|
||||
target: int = DEFAULT_TARGET,
|
||||
per_version: int = DEFAULT_PER_VERSION,
|
||||
min_score: int = DEFAULT_MIN_SCORE,
|
||||
log: Log = _log,
|
||||
) -> None:
|
||||
cmd_discover(
|
||||
out_root=out_root,
|
||||
top_checkpoints=top_checkpoints,
|
||||
top_loras=top_loras,
|
||||
log=log,
|
||||
)
|
||||
cmd_scrape(
|
||||
out_root=out_root,
|
||||
target=target,
|
||||
per_version=per_version,
|
||||
min_score=min_score,
|
||||
log=log,
|
||||
)
|
||||
cmd_split(out_root=out_root, log=log)
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
p = argparse.ArgumentParser(
|
||||
prog="python -m gpu_rent.civitai_dataset",
|
||||
description="Scrape Civitai Krea2 galleries → train.jsonl + search.jsonl",
|
||||
)
|
||||
p.add_argument(
|
||||
"--out",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="App root (default: detect). Writes datasets/civitai under it.",
|
||||
)
|
||||
sub = p.add_subparsers(dest="cmd")
|
||||
|
||||
d = sub.add_parser("discover", help="Find Krea2 checkpoints and LoRAs")
|
||||
d.add_argument("--top-checkpoints", type=int, default=DEFAULT_TOP_CHECKPOINTS)
|
||||
d.add_argument("--top-loras", type=int, default=DEFAULT_TOP_LORAS)
|
||||
|
||||
s = sub.add_parser("scrape", help="Scrape galleries until --target rows")
|
||||
s.add_argument("--target", type=int, default=DEFAULT_TARGET)
|
||||
s.add_argument("--per-version", type=int, default=DEFAULT_PER_VERSION)
|
||||
s.add_argument("--min-score", type=int, default=DEFAULT_MIN_SCORE)
|
||||
|
||||
sub.add_parser("split", help="Write by_kind / by_rating / train / search")
|
||||
|
||||
a = sub.add_parser("all", help="discover + scrape + split (default)")
|
||||
a.add_argument("--top-checkpoints", type=int, default=DEFAULT_TOP_CHECKPOINTS)
|
||||
a.add_argument("--top-loras", type=int, default=DEFAULT_TOP_LORAS)
|
||||
a.add_argument("--target", type=int, default=DEFAULT_TARGET)
|
||||
a.add_argument("--per-version", type=int, default=DEFAULT_PER_VERSION)
|
||||
a.add_argument("--min-score", type=int, default=DEFAULT_MIN_SCORE)
|
||||
return p
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
argv = list(sys.argv[1:] if argv is None else argv)
|
||||
parser = build_parser()
|
||||
if not argv:
|
||||
argv = ["all"]
|
||||
args = parser.parse_args(argv)
|
||||
out = args.out
|
||||
try:
|
||||
if args.cmd == "discover":
|
||||
cmd_discover(
|
||||
out_root=out,
|
||||
top_checkpoints=args.top_checkpoints,
|
||||
top_loras=args.top_loras,
|
||||
)
|
||||
elif args.cmd == "scrape":
|
||||
cmd_scrape(
|
||||
out_root=out,
|
||||
target=args.target,
|
||||
per_version=args.per_version,
|
||||
min_score=args.min_score,
|
||||
)
|
||||
elif args.cmd == "split":
|
||||
cmd_split(out_root=out)
|
||||
else:
|
||||
# all (explicit or default via empty → all)
|
||||
kw = {
|
||||
"out_root": out,
|
||||
"top_checkpoints": getattr(args, "top_checkpoints", DEFAULT_TOP_CHECKPOINTS),
|
||||
"top_loras": getattr(args, "top_loras", DEFAULT_TOP_LORAS),
|
||||
"target": getattr(args, "target", DEFAULT_TARGET),
|
||||
"per_version": getattr(args, "per_version", DEFAULT_PER_VERSION),
|
||||
"min_score": getattr(args, "min_score", DEFAULT_MIN_SCORE),
|
||||
}
|
||||
cmd_all(**kw)
|
||||
except GpuRentError as exc:
|
||||
print(f"error: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1006,6 +1006,23 @@ def seed_assistent_personas(cfg: Config, host: str, log: Log) -> None:
|
||||
)
|
||||
ctx_note = f", num_ctx={num_ctx}" if num_ctx else ""
|
||||
log(f"assistent-personas → overlay personas/{written} (default={default_id}{ctx_note})")
|
||||
seed_civitai_examples(cfg, host, log)
|
||||
|
||||
|
||||
def seed_civitai_examples(cfg: Config, host: str, log: Log) -> None:
|
||||
"""Push datasets/civitai/search.jsonl → Assistent/civitai-examples.jsonl (FTS, no embed)."""
|
||||
from gpu_rent.paths import app_root
|
||||
from gpu_rent.ssh_ops import put_file
|
||||
|
||||
local = Path(getattr(cfg, "app_root", None) or app_root()) / "datasets" / "civitai" / "search.jsonl"
|
||||
if not local.is_file():
|
||||
log("civitai-examples: нет datasets/civitai/search.jsonl — skip")
|
||||
return
|
||||
remote = f"{DATA}/Assistent/civitai-examples.jsonl"
|
||||
run_ssh(cfg, host, f"mkdir -p {shlex.quote(DATA + '/Assistent')}", check=False)
|
||||
put_file(cfg, host, local, remote)
|
||||
size_kb = max(1, local.stat().st_size // 1024)
|
||||
log(f"civitai-examples -> {remote} ({size_kb} KB)")
|
||||
|
||||
|
||||
def count_wanted_models_on_vm(cfg: Config, host: str) -> int:
|
||||
|
||||
Reference in New Issue
Block a user