Preserve full CivitAI metadata in LoRA Manager sidecars
This commit is contained in:
+10
-1
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user