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>

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, <lora:x:0.7>", 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")