Leonid PershinandCursor
07bb521b70
Introduce books/ with civitai-krea2 and HF fictext builders, sha-diff seed to VM, and drop train.jsonl from the civitai scrape pipeline. Co-authored-by: Cursor <cursoragent@cursor.com>
253 lines
7.6 KiB
Python
253 lines
7.6 KiB
Python
"""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,
|
|
)
|
|
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_search_row():
|
|
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}],
|
|
}
|
|
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["search"] == 2
|
|
assert counts["kind:checkpoint"] == 1
|
|
assert counts["kind:lora"] == 1
|
|
assert counts["rating:pg"] == 1
|
|
assert counts["rating:x"] == 1
|
|
assert counts.get("book:civitai-krea2") == 2
|
|
search = (tmp_path / "datasets" / "civitai" / "search.jsonl").read_text(encoding="utf-8")
|
|
assert '"id": 1' in search
|
|
book = tmp_path / "books" / "civitai-krea2" / "search.jsonl"
|
|
assert book.is_file()
|
|
|
|
|
|
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")
|