diff --git a/.gitignore b/.gitignore index 88b4e9b..80c5e40 100644 --- a/.gitignore +++ b/.gitignore @@ -27,6 +27,13 @@ Output/** !Output/**/.gitkeep !Output/README.md +# Civitai Krea2 scrape (meta only — train.jsonl / search.jsonl) +datasets/** +!datasets/.gitkeep +!datasets/README.md +!datasets/**/.gitkeep +!datasets/**/README.md + # Python .venv/ __pycache__/ diff --git a/README.md b/README.md index 7060176..7f09e0c 100644 --- a/README.md +++ b/README.md @@ -59,10 +59,13 @@ Unix: `./gpu-rent.sh …` (один раз `chmod +x gpu-rent.sh`). | `up --llm-only --ollama` | только Ollama, без SwarmUI | | `tunnel` / `open` | Снова UI / браузер | | `hold` / `status` | Пауза killer / состояние | +| `debug` | Read-only Debug API на `:17821` (агент / Assistent) | +| `seed-personas` | Overlay `assistent-personas/` → VM | +| `watchdog` | Локальный safety-net (мёртвый туннель → stop) | | `stop` / `destroy --i-understand-data-loss` | Стоп GPU / + диски | | `push` / `pull-output` / `seed-*` / `capture` | Файлы ↔ VM; capture = ссылки в yaml | -Полный список: [docs/cli.md](docs/cli.md). +Полный список: [docs/cli.md](docs/cli.md). Civitai scrape (без картинок) → Assistent FTS: [datasets/README.md](datasets/README.md). --- @@ -87,8 +90,9 @@ Unix: `./gpu-rent.sh …` (один раз `chmod +x gpu-rent.sh`). | --- | --- | | [docs/setup.md](docs/setup.md) | **Пошаговая подготовка** до первого `up` | | [docs/cli.md](docs/cli.md) | Все команды и переменные | -| [docs/llm.md](docs/llm.md) | Ollama (opt-in LLM) | +| [docs/llm.md](docs/llm.md) | Ollama + Assistent (opt-in LLM) | | [docs/models.md](docs/models.md) | Civitai + папка `Models/` | +| [datasets/README.md](datasets/README.md) | Civitai Krea2 scrape → FTS Assistent | | [docs/spike-notes.md](docs/spike-notes.md) | Чеклист первого живого прогона | | [docs/README.md](docs/README.md) | Оглавление всего `docs/` | diff --git a/civitai-dataset.bat b/civitai-dataset.bat new file mode 100644 index 0000000..fd93f44 --- /dev/null +++ b/civitai-dataset.bat @@ -0,0 +1,15 @@ +@echo off +setlocal EnableExtensions +cd /d "%~dp0" +chcp 65001 >nul + +rem cmd.exe / Explorer: bypass ExecutionPolicy. Logic lives in civitai-dataset.ps1. +if "%~1"=="" ( + powershell.exe -NoProfile -ExecutionPolicy Bypass -File "%~dp0civitai-dataset.ps1" + echo. + pause + exit /b %ERRORLEVEL% +) + +powershell.exe -NoProfile -ExecutionPolicy Bypass -File "%~dp0civitai-dataset.ps1" %* +exit /b %ERRORLEVEL% diff --git a/civitai-dataset.ps1 b/civitai-dataset.ps1 new file mode 100644 index 0000000..3863c47 --- /dev/null +++ b/civitai-dataset.ps1 @@ -0,0 +1,122 @@ +#Requires -Version 5.1 +# Windows launcher: same .venv as gpu-rent, runs python -m gpu_rent.civitai_dataset +# Default (no args): discover + scrape --target 2000 + split +Set-StrictMode -Version Latest +$ErrorActionPreference = "Stop" + +$Root = Split-Path -Parent $MyInvocation.MyCommand.Path +Set-Location -LiteralPath $Root +try { + $utf8 = New-Object System.Text.UTF8Encoding $false + [Console]::OutputEncoding = $utf8 + $script:OutputEncoding = $utf8 +} catch { +} + +function Import-GpuRentVars { + param([Parameter(Mandatory = $true)][string]$Path) + if (-not (Test-Path -LiteralPath $Path)) { return } + Get-Content -LiteralPath $Path -Encoding UTF8 | ForEach-Object { + $line = $_.Trim() + if (-not $line -or $line.StartsWith("#")) { return } + if ($line -match '^(?i)export\s+') { + $line = $line.Substring($Matches[0].Length).Trim() + } + $eq = $line.IndexOf("=") + if ($eq -lt 1) { return } + $key = $line.Substring(0, $eq).Trim() + $val = $line.Substring($eq + 1).Trim() + if ($val.Length -ge 2) { + $q = $val[0] + if (($q -eq '"' -or $q -eq "'") -and $val[-1] -eq $q) { + $val = $val.Substring(1, $val.Length - 2) + } + } + if (-not $key) { return } + $existing = [Environment]::GetEnvironmentVariable($key, "Process") + if ([string]::IsNullOrEmpty($existing)) { + Set-Item -Path "Env:$key" -Value $val + } + } +} + +function Test-Python311 { + param( + [Parameter(Mandatory = $true)][string]$Exe, + [string[]]$Prefix = @() + ) + $cmd = Get-Command $Exe -ErrorAction SilentlyContinue + if (-not $cmd) { + return $false + } + try { + & $Exe @Prefix -c "import sys; raise SystemExit(0 if sys.version_info >= (3, 11) else 1)" 2>$null | Out-Null + return ($LASTEXITCODE -eq 0) + } catch { + return $false + } +} + +$PyExe = $null +$PyPrefix = @() +$found = $false +foreach ($row in @( + @{ Exe = "py"; Prefix = @("-3.12") }, + @{ Exe = "py"; Prefix = @("-3.11") }, + @{ Exe = "py"; Prefix = @("-3") }, + @{ Exe = "python3"; Prefix = @() }, + @{ Exe = "python"; Prefix = @() } + )) { + if (Test-Python311 -Exe $row.Exe -Prefix $row.Prefix) { + $PyExe = $row.Exe + $PyPrefix = $row.Prefix + $found = $true + break + } +} + +if (-not $found) { + Write-Host "civitai-dataset: нужен Python 3.11+." + exit 1 +} + +$VenvPy = Join-Path $Root ".venv\Scripts\python.exe" +if (-not (Test-Path -LiteralPath $VenvPy)) { + Write-Host "civitai-dataset: создаю .venv" + & $PyExe @PyPrefix -m venv (Join-Path $Root ".venv") +} + +$Stamp = Join-Path $Root ".venv\.gpu-rent-installed" +$Pyproject = Join-Path $Root "pyproject.toml" +$needInstall = -not (Test-Path -LiteralPath $Stamp) +if (-not $needInstall -and (Test-Path -LiteralPath $Pyproject)) { + $needInstall = (Get-Item -LiteralPath $Pyproject).LastWriteTime -gt (Get-Item -LiteralPath $Stamp).LastWriteTime +} +if ($needInstall) { + Write-Host "civitai-dataset: ставлю пакет в .venv" + & $VenvPy -m pip install -U pip + if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } + & $VenvPy -m pip install -e $Root + if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } + Get-Date -Format o | Set-Content -LiteralPath $Stamp -Encoding ascii +} + +$EnvFile = Join-Path $Root ".env" +$Example = Join-Path $Root "env.example" +if (-not (Test-Path -LiteralPath $EnvFile)) { + if (Test-Path -LiteralPath $Example) { + Copy-Item -LiteralPath $Example -Destination $EnvFile + Write-Host "civitai-dataset: created .env - fill CIVITAI_API_TOKEN" + } +} + +Import-GpuRentVars (Join-Path $Root "gpu-rent.vars") + +$InvokeArgs = @($args) +if ($InvokeArgs.Count -eq 0) { + $InvokeArgs = @("all") +} + +$ErrorActionPreference = "Continue" +& $VenvPy -m gpu_rent.civitai_dataset @InvokeArgs +exit $LASTEXITCODE diff --git a/civitai-dataset.sh b/civitai-dataset.sh new file mode 100644 index 0000000..a88bba3 --- /dev/null +++ b/civitai-dataset.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +# Linux / macOS / Git Bash: same .venv as gpu-rent → python -m gpu_rent.civitai_dataset +# Default (no args): discover + scrape --target 2000 + split +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$ROOT" + +load_vars_file() { + local file="$1" + [[ -f "$file" ]] || return 0 + while IFS= read -r line || [[ -n "$line" ]]; do + line="${line#"${line%%[![:space:]]*}"}" + line="${line%"${line##*[![:space:]]}"}" + [[ -z "$line" || "$line" == \#* ]] && continue + if [[ "$line" == [eE][xX][pP][oO][rR][tT][[:space:]]* ]]; then + line="${line#*[eE][xX][pP][oO][rR][tT]}" + line="${line#"${line%%[![:space:]]*}"}" + fi + [[ "$line" != *=* ]] && continue + local key="${line%%=*}" + local val="${line#*=}" + key="${key%"${key##*[![:space:]]}"}" + key="${key#"${key%%[![:space:]]*}"}" + val="${val#"${val%%[![:space:]]*}"}" + val="${val%"${val##*[![:space:]]}"}" + if [[ ${#val} -ge 2 ]]; then + local q="${val:0:1}" + if [[ ( "$q" == '"' || "$q" == "'" ) && "${val: -1}" == "$q" ]]; then + val="${val:1:${#val}-2}" + fi + fi + [[ -z "$key" ]] && continue + if [[ -z "${!key+x}" || -z "${!key}" ]]; then + export "$key=$val" + fi + done <"$file" +} + +ok_py() { + local exe="$1" + shift || true + command -v "$exe" >/dev/null 2>&1 || return 1 + "$exe" "$@" -c "import sys; raise SystemExit(0 if sys.version_info >= (3, 11) else 1)" >/dev/null 2>&1 +} + +if ok_py python3.12; then + PY=(python3.12) +elif ok_py python3.11; then + PY=(python3.11) +elif ok_py python3; then + PY=(python3) +elif ok_py python; then + PY=(python) +elif ok_py py -3.12; then + PY=(py -3.12) +elif ok_py py -3.11; then + PY=(py -3.11) +else + echo "civitai-dataset: нужен Python 3.11+" >&2 + exit 1 +fi + +if [[ -x "$ROOT/.venv/Scripts/python.exe" ]]; then + VENV_PY="$ROOT/.venv/Scripts/python.exe" +elif [[ -x "$ROOT/.venv/bin/python" ]]; then + VENV_PY="$ROOT/.venv/bin/python" +else + echo "civitai-dataset: создаю .venv" + "${PY[@]}" -m venv "$ROOT/.venv" + if [[ -x "$ROOT/.venv/Scripts/python.exe" ]]; then + VENV_PY="$ROOT/.venv/Scripts/python.exe" + else + VENV_PY="$ROOT/.venv/bin/python" + fi +fi + +STAMP="$ROOT/.venv/.gpu-rent-installed" +if [[ ! -f "$STAMP" || "$ROOT/pyproject.toml" -nt "$STAMP" ]]; then + echo "civitai-dataset: ставлю пакет в .venv" + "$VENV_PY" -m pip install -U pip + "$VENV_PY" -m pip install -e "$ROOT" + date -u +"%Y-%m-%dT%H:%M:%SZ" >"$STAMP" +fi + +if [[ ! -f "$ROOT/.env" && -f "$ROOT/env.example" ]]; then + cp "$ROOT/env.example" "$ROOT/.env" + echo "civitai-dataset: created .env - fill CIVITAI_API_TOKEN" +fi + +load_vars_file "$ROOT/gpu-rent.vars" + +ARGS=("$@") +if [[ ${#ARGS[@]} -eq 0 ]]; then + ARGS=(all) +fi + +exec "$VENV_PY" -m gpu_rent.civitai_dataset "${ARGS[@]}" diff --git a/datasets/.gitkeep b/datasets/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/datasets/README.md b/datasets/README.md new file mode 100644 index 0000000..6dfaff3 --- /dev/null +++ b/datasets/README.md @@ -0,0 +1,41 @@ +# Civitai Krea2 datasets (local) + +Scraped generation metadata for Assistent training + cheap FTS search. **No image files.** + +Not part of `gpu-rent` Typer CLI and **not** `models.yaml` / `gpu-rent push`. Separate entrypoint: + +```text +civitai-dataset.bat # Windows (or .\civitai-dataset.ps1 / ./civitai-dataset.sh) +# default = all → discover + scrape + split +python -m gpu_rent.civitai_dataset all +``` + +Needs `CIVITAI_API_TOKEN` and preferably `CIVITAI_API_HOST=civitai.red` in `.env`. + +## Commands + +| Command | Flags | Role | +| --- | --- | --- | +| `discover` | `--top-checkpoints` (20), `--top-loras` (30) | Popular Krea2 ckpt/LoRAs (+ versions from `models.yaml`) → `catalog/models.json` | +| `scrape` | `--target` (2000), `--per-version`, `--min-score` | Gallery metadata only; resume by id; NSFW→rating; skip minor tags | +| `split` | — | `by_kind/` / `by_rating/` / `train.jsonl` / `search.jsonl` | +| `all` | same as discover+scrape | discover → scrape → split (launcher default) | + +Global: `--out` = app root (writes under `datasets/civitai`). + +## Layout + +| Path | Role | +| --- | --- | +| `civitai/catalog/models.json` | Discovered Krea2 checkpoints / LoRAs | +| `civitai/catalog/images.jsonl` | Canonical rows (resume by id) | +| `civitai/by_kind/` | checkpoint vs lora | +| `civitai/by_rating/` | pg / pg13 / r / x | +| `civitai/train.jsonl` | Alpaca instruction/output for Assistent Обучение | +| `civitai/search.jsonl` | Compact FTS index → pushed to VM as `Assistent/civitai-examples.jsonl` on `up` / `seed-personas` | + +Push of `search.jsonl` is skipped if the file is missing. Weights stay in [models.md](../docs/models.md) (Civitai seed / `Models/` / `capture`). + +Git ignores the data; this README stays. + +See also: [docs/cli.md](../docs/cli.md) (civitai-dataset block), [docs/extensions.md](../docs/extensions.md), [docs/local-folders.md](../docs/local-folders.md). diff --git a/docs/README.md b/docs/README.md index cca8c47..ab9a5e5 100644 --- a/docs/README.md +++ b/docs/README.md @@ -16,12 +16,15 @@ CLI поднимает прерываемый GPU в Selectel, держит Swar | --- | --- | | Понять зачем и сколько стоит | [concept.md](concept.md) | | Все команды и `.env` | [cli.md](cli.md) | +| Debug API / Assistent curl | [cli.md](cli.md) (раздел Debug API) | | Модели с Civitai / `Models/` | [models.md](models.md) | -| Git-расширения SwarmUI/Comfy | [extensions.md](extensions.md) | +| Civitai scrape → FTS Assistent | [../datasets/README.md](../datasets/README.md) | +| Git-расширения / personas | [extensions.md](extensions.md) (`seed-personas`) | | Word-list промптов | [autocomplete.md](autocomplete.md) | | Push/pull папок | [local-folders.md](local-folders.md) | -| Ollama | [llm.md](llm.md) | +| Ollama + Assistent | [llm.md](llm.md) | | Capture ссылок с VM | [cli.md](cli.md) (`capture`) + [models.md](models.md) | +| Баланс ₽ (X-Token панели) | [cli.md](cli.md) / [decisions.md](decisions.md) (`SELECTEL_API_TOKEN`) | | Как устроены диски и killer | [architecture.md](architecture.md) | | Контракт Selectel | [selectel.md](selectel.md) | | Нативный SwarmUI на VM | [swarmui.md](swarmui.md) | @@ -29,11 +32,11 @@ CLI поднимает прерываемый GPU в Selectel, держит Swar | Статус реализации | [roadmap.md](roadmap.md) | | Что ещё проверить на spike | [open-questions.md](open-questions.md) | -Ревью кода: `docs/reviews/` (для разработки, не для первого запуска). +Ревью кода: `docs/reviews/` (для разработки, не для первого запуска). Это **замороженные снимки** от 21.08.2026 — могут противоречить текущему коду (не источник правды). ## Источники (Selectel / SwarmUI / Civitai) -Проверено 21 августа 2026: +Проверено 23 августа 2026: - [Selectel: preemptible](https://docs.selectel.ru/en/cloud-servers/about/preemptible-servers/) - [Selectel: GPU](https://docs.selectel.ru/en/cloud-servers/create/gpus/) diff --git a/docs/autocomplete.md b/docs/autocomplete.md index 7d87caf..75ef1a6 100644 --- a/docs/autocomplete.md +++ b/docs/autocomplete.md @@ -32,15 +32,13 @@ Swarm Assistent индексирует тот же csv в FTS-таблицу `ta ## Куда класть -На data volume (уже примонтирован как `/SwarmUI/Data`): +На data volume (bind → `/opt/swarmui/Data/…`): ```text /mnt/swarm_data/Data/Autocompletions/danbooru.csv /mnt/swarm_data/Data/Autocompletions/danbooru.csv.gpu-rent-meta.json ``` -На data volume: `/mnt/swarm_data/Data/Autocompletions/danbooru.csv` (bind → `/opt/swarmui/Data/…`). - Sidecar meta (не отдавать в Output): ```json diff --git a/docs/concept.md b/docs/concept.md index ec2c1cd..f1b3b83 100644 --- a/docs/concept.md +++ b/docs/concept.md @@ -24,7 +24,10 @@ GPU в облаке дорогой. Веса для генерации карт - Первый clone git-реп расширений SwarmUI и ComfyUI nodes из `extensions.yaml`. - Autocomplete: word-list в `Data/Autocompletions` до старта UI, на каждом `up` проверка новой версии. - `doctor` до create; фоллбек flavor; интерактивный выбор flavor/диска; `hold`; toast Idle; `open` на 17801. -- Opt-in LLM (Ollama) через `setup` / флаги / `LLM_RUNTIME`; `capture` ссылок с VM. +- Opt-in LLM (Ollama) через `setup` / флаги / `LLM_RUNTIME`; расширение swarm-assistent поверх Swarm+Ollama; `capture` ссылок с VM. +- Debug sidecar на `127.0.0.1:17821` (агент / Assistent diagnose) на `up` / `tunnel` / `gpu-rent debug`. +- Локальный scrape Civitai (метаданные Krea2, без картинок) → FTS Assistent на `up` ([datasets/README](../datasets/README.md)). +- Опциональный local-watchdog: мёртвый туннель без `stop` → через grace удалить compute. - `status`: диск, окно preempt 24 ч, killer / LLM. - Snapshot boot после первого удачного bootstrap. diff --git a/docs/llm.md b/docs/llm.md index 32ec116..553a4ad 100644 --- a/docs/llm.md +++ b/docs/llm.md @@ -148,7 +148,8 @@ Busy (не гасить GPU): По умолчанию install тянет upstream без pin (`WARN` в логе). Жёстче: ```bash -OLLAMA_VERSION=0.6.5 +# Pin must be ≥ 0.12.7 if you use Qwen3-VL (see above) +OLLAMA_VERSION=0.12.7 OLLAMA_SHA256= ``` diff --git a/docs/roadmap.md b/docs/roadmap.md index 0595fe6..491f169 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -65,14 +65,14 @@ - [x] Clone `extensions.yaml` в Extensions / DLNodes до старта UI - [x] Autocomplete: danbooru.csv + Settings.fds; на каждом up — GitHub sha - [x] Civitai seed по манифесту; без токена — дефолт SwarmUI -- [x] systemd idle-killer + application credential; льгота 90 мин; hold-файл; очередь/loading = busy; затем 60 мин пустой очереди; вооружение после seed/start (качалка UI = hold) +- [x] systemd idle-killer + application credential; льгота `IDLE_GRACE_MINUTES=45`; hold-файл; очередь/loading = busy; затем `IDLE_MINUTES=30` пустой очереди; вооружение после seed/start (качалка UI = hold) - [x] Один snapshot `gpu-rent-boot-ok` после первого backend Idle; следующие create — из snapshot если есть - [x] `gpu-rent hold` / `hold --clear` по SSH - [x] `ready` по HTTP, затем backend Idle ## 5. Туннель -- [x] `gpu-rent tunnel` → 17801, Ctrl+C не делает `stop` +- [x] `gpu-rent tunnel` → 17801; **Ctrl+C / Ctrl+D** вызывают `stop` (диски остаются) - [x] `gpu-rent open` / `tunnel --open` - [x] EXPIRED → unshelve + reconnect, пока туннель жив - [x] Сниппет MCP в stdout @@ -87,6 +87,15 @@ - [x] Цветной CLI-лог; `capture` ссылок с VM - [x] Local-watchdog; access-card +## Уже в коде после v1 + +- Debug API на `127.0.0.1:17821` (`gpu-rent debug` / sidecar на `up`/`tunnel`); Assistent session / chat-eval +- `civitai-dataset` → `datasets/civitai/search.jsonl` → FTS push на `up` / `seed-personas` +- `seed-personas`, `capture wanted`, `logs --unit swarm|ollama|killer|cloud-init` +- Balance notify (`SELECTEL_API_TOKEN` + watchdog) + +Spike в §0 по-прежнему нужно закрыть на живом GPU. + ## Вне скоупа v1 - `generate` в CLI, S3 как источник моделей diff --git a/docs/spike-notes.md b/docs/spike-notes.md index bc48609..1ed73c1 100644 --- a/docs/spike-notes.md +++ b/docs/spike-notes.md @@ -122,7 +122,7 @@ - Имя GPU-образа в пуле: - Реальные flavor id (не в git): -- Хватает ли `IDLE_GRACE_MINUTES=90`: +- Хватает ли `IDLE_GRACE_MINUTES=45` (дефолт; поднять, если мало на первый Comfy): - Application credential с узкими access_rules: ок / ошибка (fail closed): - Прочее: diff --git a/src/gpu_rent/civitai.py b/src/gpu_rent/civitai.py index b05153d..8b0684c 100644 --- a/src/gpu_rent/civitai.py +++ b/src/gpu_rent/civitai.py @@ -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 diff --git a/src/gpu_rent/civitai_dataset.py b/src/gpu_rent/civitai_dataset.py new file mode 100644 index 0000000..1f1e37b --- /dev/null +++ b/src/gpu_rent/civitai_dataset.py @@ -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()) diff --git a/src/gpu_rent/provision.py b/src/gpu_rent/provision.py index b3e5d69..3b5a8ed 100644 --- a/src/gpu_rent/provision.py +++ b/src/gpu_rent/provision.py @@ -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: diff --git a/tests/test_civitai_dataset.py b/tests/test_civitai_dataset.py new file mode 100644 index 0000000..1b2ed55 --- /dev/null +++ b/tests/test_civitai_dataset.py @@ -0,0 +1,258 @@ +"""Unit tests for Civitai Krea2 dataset scrape / split (no live network).""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from gpu_rent.civitai import list_images, list_models +from gpu_rent.civitai_dataset import ( + cmd_split, + looks_minor, + normalize_image, + rating_from_nsfw, + reaction_score, + search_row, + tags_from_image, + train_row, +) +from gpu_rent.errors import CloudError + + +def test_rating_from_nsfw(): + assert rating_from_nsfw("None") == "pg" + assert rating_from_nsfw("Soft") == "pg13" + assert rating_from_nsfw("Mature") == "r" + assert rating_from_nsfw("X") == "x" + assert rating_from_nsfw(True) == "x" + assert rating_from_nsfw(False) == "pg" + + +def test_reaction_score(): + assert reaction_score({"likeCount": 2, "heartCount": 3, "laughCount": 1}) == 6 + assert reaction_score(None) == 0 + + +def test_looks_minor(): + assert looks_minor("loli school", []) + assert looks_minor("portrait", ["shota"]) + assert not looks_minor("adult woman redhead", ["stockings"]) + + +def test_tags_from_prompt_fallback(): + tags = tags_from_image({}, "woman, redhead, cinematic lighting, ", cap=5) + assert "woman" in tags + assert "redhead" in tags + assert not any("lora:" in t.lower() for t in tags) + + +def test_normalize_image_skips_no_prompt(): + item = { + "id": 1, + "type": "image", + "nsfwLevel": "None", + "stats": {"likeCount": 10}, + "meta": {}, + } + assert normalize_image(item, kind="checkpoint", model_id=1, model_version_id=2, ours=False) is None + + +def test_normalize_image_ok(): + item = { + "id": 42, + "type": "image", + "url": "https://example/x.jpg", + "username": "u", + "createdAt": "2026-01-01T00:00:00Z", + "nsfwLevel": "Soft", + "width": 832, + "height": 1216, + "stats": {"likeCount": 10, "heartCount": 2}, + "meta": { + "prompt": "a woman in soft light", + "negativePrompt": "blur", + "steps": 8, + "cfgScale": 1, + "sampler": "euler", + "civitaiResources": [{"type": "lora", "modelVersionId": 9, "weight": 0.7}], + }, + "tags": ["woman", "portrait"], + } + row = normalize_image( + item, kind="checkpoint", model_id=100, model_version_id=200, ours=True, min_score=5 + ) + assert row is not None + assert row["id"] == 42 + assert row["rating"] == "pg13" + assert row["score"] == 12 + assert row["ours"] is True + assert row["params"]["steps"] == 8 + assert "woman" in row["tags"] + + +def test_normalize_skips_low_score(): + item = { + "id": 3, + "type": "image", + "nsfwLevel": "None", + "stats": {"likeCount": 1}, + "meta": {"prompt": "hello world"}, + } + assert ( + normalize_image(item, kind="lora", model_id=1, model_version_id=2, ours=False, min_score=5) + is None + ) + + +def test_train_and_search_rows(): + row = { + "id": 7, + "rating": "r", + "score": 50, + "kind": "checkpoint", + "modelVersionId": 3231611, + "tags": ["woman", "cinematic"], + "prompt": "A woman in cinematic light", + "negativePrompt": "blur", + "params": {"steps": 8, "cfgScale": 1, "sampler": "euler"}, + "resources": [{"type": "lora", "modelVersionId": 1, "weight": 0.5}], + } + tr = train_row(row) + assert "Tags: woman, cinematic" in tr["instruction"] + assert "Rating: r" in tr["instruction"] + assert "A woman" in tr["output"] + assert "cfg: 1" in tr["output"] + sr = search_row(row) + assert sr["id"] == 7 + assert sr["loras"] == [{"versionId": 1, "weight": 0.5}] + assert sr["params"]["cfg"] == 1 + + +def test_cmd_split_writes_artifacts(tmp_path: Path): + cat = tmp_path / "datasets" / "civitai" / "catalog" + cat.mkdir(parents=True) + rows = [ + { + "id": 1, + "kind": "checkpoint", + "rating": "pg", + "score": 10, + "modelVersionId": 1, + "tags": ["a"], + "prompt": "p1", + "negativePrompt": "", + "params": {"steps": 4}, + "resources": [], + }, + { + "id": 2, + "kind": "lora", + "rating": "x", + "score": 20, + "modelVersionId": 2, + "tags": ["b"], + "prompt": "p2", + "negativePrompt": "n", + "params": {"cfgScale": 1}, + "resources": [], + }, + ] + with (cat / "images.jsonl").open("w", encoding="utf-8") as fh: + for r in rows: + fh.write(json.dumps(r) + "\n") + counts = cmd_split(out_root=tmp_path, log=lambda m: None) + assert counts["train"] == 2 + assert counts["search"] == 2 + assert counts["kind:checkpoint"] == 1 + assert counts["kind:lora"] == 1 + assert counts["rating:pg"] == 1 + assert counts["rating:x"] == 1 + search = (tmp_path / "datasets" / "civitai" / "search.jsonl").read_text(encoding="utf-8") + assert '"id": 1' in search + train = (tmp_path / "datasets" / "civitai" / "train.jsonl").read_text(encoding="utf-8") + assert "Write a Krea 2 prompt" in train + + +def _mock_response(payload: dict, status: int = 200) -> httpx.Response: + return httpx.Response(status, json=payload, request=httpx.Request("GET", "https://civitai.red/api/v1/x")) + + +def test_list_models_failover(monkeypatch): + calls = [] + + class FakeClient: + def __init__(self, *a, **k): + pass + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def get(self, url, headers=None, params=None): + calls.append((url, dict(params or {}))) + if "civitai.com" in url: + return _mock_response({"items": [{"id": 1, "name": "Krea"}], "metadata": {}}) + return _mock_response({}, status=404) + + monkeypatch.setattr("gpu_rent.civitai.httpx.Client", FakeClient) + host, data = list_models("tok", "civitai.red", types="Checkpoint", query="krea2", limit=10) + assert host == "civitai.com" + assert data["items"][0]["id"] == 1 + assert any("civitai.red" in u for u, _ in calls) + # query search must not send page= + assert all("page" not in p for _, p in calls) + + +def test_list_images_cursor(monkeypatch): + class FakeClient: + def __init__(self, *a, **k): + pass + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def get(self, url, headers=None, params=None): + assert params.get("modelVersionId") == 3231611 + assert params.get("withMeta") == "true" + assert "modelId" not in params + return _mock_response( + { + "items": [{"id": 9, "meta": {"prompt": "x"}}], + "metadata": {"nextCursor": "abc"}, + } + ) + + monkeypatch.setattr("gpu_rent.civitai.httpx.Client", FakeClient) + host, data = list_images( + "tok", "civitai.red", model_version_id=3231611, with_meta=True, limit=50 + ) + assert host == "civitai.red" + assert data["metadata"]["nextCursor"] == "abc" + + +def test_list_models_raises_on_fail(monkeypatch): + class FakeClient: + def __init__(self, *a, **k): + pass + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def get(self, url, headers=None, params=None): + return _mock_response({}, status=500) + + monkeypatch.setattr("gpu_rent.civitai.httpx.Client", FakeClient) + with pytest.raises(CloudError): + list_models("tok", "civitai.red", types="Checkpoint", query="krea")