Preserve full CivitAI metadata in LoRA Manager sidecars

This commit is contained in:
Dariusz L
2026-07-26 19:56:29 +02:00
parent c6ff826028
commit 58e1aeff29
4 changed files with 190 additions and 1 deletions
+10 -1
View File
@@ -602,7 +602,16 @@ def build_lora_manager_metadata(
and bool(source.get("civitai") or source.get("civitai_details"))
)
civitai_payload = (
dict(_as_dict(source.get("civitai"))) if is_civitai_source else {}
dict(
_as_dict(
_first_present(
source.get("civitai"),
details.get("civitai"),
)
)
)
if is_civitai_source
else {}
)
if is_civitai_source:
if model_id and "modelId" not in civitai_payload:
+41
View File
@@ -1208,6 +1208,36 @@ def _normalize_civitai_file(
)
def _build_civitai_sidecar_payload(
model_data: Dict[str, Any],
version_data: Dict[str, Any],
model_id: int,
) -> Dict[str, Any]:
"""Preserve the selected CivitAI version for LoRA Manager sidecars."""
if not isinstance(version_data, dict):
return {}
payload = dict(version_data)
payload.setdefault("modelId", model_id)
model_payload = {
key: value
for key, value in model_data.items()
if key not in {"id", "modelVersions", "creator", "stats"}
}
existing_model = payload.get("model")
if isinstance(existing_model, dict):
model_payload.update(existing_model)
if model_payload:
payload["model"] = model_payload
creator = model_data.get("creator")
if isinstance(creator, dict) and "creator" not in payload:
payload["creator"] = creator
return payload
def get_civitai_model_details(
model_id: int,
version_id: Optional[int] = None,
@@ -1228,11 +1258,14 @@ def get_civitai_model_details(
versions = data.get("modelVersions") or []
normalized_versions = []
selected_version = None
selected_raw_version = None
raw_versions = []
for version in versions:
if not isinstance(version, dict):
continue
raw_versions.append(version)
current_version_id = version.get("id")
files = [
_normalize_civitai_file(file_info, model_id, current_version_id, api_key)
@@ -1256,11 +1289,18 @@ def get_civitai_model_details(
normalized_versions.append(normalized)
if version_id and str(current_version_id) == str(version_id):
selected_version = normalized
selected_raw_version = version
if not selected_version and normalized_versions:
selected_version = normalized_versions[0]
selected_raw_version = raw_versions[0]
selected_images = selected_version.get("images", []) if selected_version else []
civitai_payload = _build_civitai_sidecar_payload(
data,
selected_raw_version or {},
model_id,
)
return {
"source": "civitai",
"model_id": model_id,
@@ -1276,6 +1316,7 @@ def get_civitai_model_details(
"versions": normalized_versions,
"selected_version": selected_version,
"images": selected_images,
"civitai": civitai_payload,
}
+62
View File
@@ -130,6 +130,68 @@ class DownloaderMetadataSidecarTests(unittest.TestCase):
self.assertIn("type=Model", payload["civitai"]["downloadUrl"])
self.assertNotIn("secret", json.dumps(payload))
def test_uses_full_civitai_payload_from_fetched_details(self):
with tempfile.TemporaryDirectory() as tmpdir:
model_path = os.path.join(tmpdir, "example.safetensors")
with open(model_path, "wb") as handle:
handle.write(b"abc")
payload = build_lora_manager_metadata(
model_path,
{
"source": "civitai",
"details_source": "civitai",
"model_id": 123,
"version_id": 456,
"download_url": (
"https://civitai.com/api/download/models/456?token=secret"
),
"civitai_details": {
"source": "civitai",
"model_id": 123,
"version_id": 456,
"civitai": {
"modelId": 123,
"id": 456,
"name": "v1",
"images": [
{
"url": "https://image.civitai.com/example.jpeg",
"nsfwLevel": 4,
"meta": {
"seed": 42,
"prompt": "Full prompt metadata",
},
}
],
"files": [
{
"name": "example.safetensors",
"virusScanResult": "Success",
}
],
"model": {
"name": "Example Model",
"allowCommercialUse": ["Image"],
},
"stats": {"downloadCount": 10},
},
},
},
category="loras",
)
civitai = payload["civitai"]
self.assertEqual(4, civitai["images"][0]["nsfwLevel"])
self.assertEqual(
"Full prompt metadata",
civitai["images"][0]["meta"]["prompt"],
)
self.assertEqual("Success", civitai["files"][0]["virusScanResult"])
self.assertEqual(["Image"], civitai["model"]["allowCommercialUse"])
self.assertEqual(10, civitai["stats"]["downloadCount"])
self.assertNotIn("token=", civitai["downloadUrl"])
def test_existing_file_with_same_hash_is_marked_already_downloaded(self):
content = b"existing model"
expected_sha256 = hashlib.sha256(content).hexdigest()
+77
View File
@@ -49,6 +49,7 @@ from core.sources.civitai import (
_search_civitai_trpc_candidates,
build_civitai_session_cookie,
clear_search_cache,
get_civitai_model_details,
get_model_info_by_hash,
get_model_info_for_file,
parse_civitai_url,
@@ -100,6 +101,82 @@ class CivitaiResultBuilderTests(unittest.TestCase):
self.assertEqual(10 * 1024, result["size"])
class CivitaiModelDetailsTests(unittest.TestCase):
def test_preserves_full_selected_version_for_metadata_sidecar(self):
response = {
"id": 123,
"name": "Alt-Girl/E-Girl",
"type": "LORA",
"description": "Model description",
"tags": ["style", "e-girl"],
"allowNoCredit": True,
"allowCommercialUse": ["Image", "Sell"],
"creator": {"username": "TheAINovice", "image": None},
"stats": {"downloadCount": 1000},
"modelVersions": [
{
"id": 456,
"name": "Old version",
"baseModel": "Krea 1",
"images": [],
"files": [],
},
{
"id": 789,
"name": "Krea2 v1.0",
"baseModel": "Krea 2",
"createdAt": "2026-07-14T05:52:54.321Z",
"stats": {"downloadCount": 1705, "thumbsUpCount": 106},
"images": [
{
"url": "https://image.civitai.com/example.jpeg",
"nsfwLevel": 4,
"hash": "blurhash",
"meta": {
"seed": 884670593557080,
"prompt": "AltGirl prompt",
"steps": 8,
},
}
],
"files": [
{
"id": 321,
"name": "AltGirlKrea.safetensors",
"type": "Model",
"pickleScanResult": "Success",
"virusScanResult": "Success",
"hashes": {"SHA256": "A" * 64},
"primary": True,
}
],
},
],
}
with patch(
"core.sources.civitai.execute_provider_json_request",
return_value=response,
):
details = get_civitai_model_details(123, 789)
self.assertIsNotNone(details)
civitai = details["civitai"]
self.assertEqual(789, civitai["id"])
self.assertEqual(123, civitai["modelId"])
self.assertEqual(4, civitai["images"][0]["nsfwLevel"])
self.assertEqual("AltGirl prompt", civitai["images"][0]["meta"]["prompt"])
self.assertEqual("Success", civitai["files"][0]["virusScanResult"])
self.assertEqual(1705, civitai["stats"]["downloadCount"])
self.assertEqual(
["Image", "Sell"],
civitai["model"]["allowCommercialUse"],
)
self.assertEqual("TheAINovice", civitai["creator"]["username"])
self.assertNotIn("modelVersions", civitai["model"])
class CivitaiHashLookupCacheTests(unittest.TestCase):
def setUp(self):