mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-25 07:05:43 -03:00
Merge pull request #689 from willmiao/civitai-deletion-logic
feat(metadata): improve civitai deletion detection logic, see #670
This commit is contained in:
@@ -214,6 +214,7 @@ class MetadataSyncService:
|
|||||||
metadata_provider: Optional[MetadataProviderProtocol] = None
|
metadata_provider: Optional[MetadataProviderProtocol] = None
|
||||||
provider_used: Optional[str] = None
|
provider_used: Optional[str] = None
|
||||||
last_error: Optional[str] = None
|
last_error: Optional[str] = None
|
||||||
|
civitai_api_not_found = False
|
||||||
|
|
||||||
for provider_name, provider in provider_attempts:
|
for provider_name, provider in provider_attempts:
|
||||||
try:
|
try:
|
||||||
@@ -228,19 +229,24 @@ class MetadataSyncService:
|
|||||||
if provider_name == "sqlite":
|
if provider_name == "sqlite":
|
||||||
sqlite_attempted = True
|
sqlite_attempted = True
|
||||||
|
|
||||||
|
is_default_provider = provider_name is None
|
||||||
|
|
||||||
if civitai_metadata_candidate:
|
if civitai_metadata_candidate:
|
||||||
civitai_metadata = civitai_metadata_candidate
|
civitai_metadata = civitai_metadata_candidate
|
||||||
metadata_provider = provider
|
metadata_provider = provider
|
||||||
provider_used = provider_name
|
provider_used = provider_name
|
||||||
break
|
break
|
||||||
|
|
||||||
|
if is_default_provider and error == "Model not found":
|
||||||
|
civitai_api_not_found = True
|
||||||
|
|
||||||
last_error = error or last_error
|
last_error = error or last_error
|
||||||
|
|
||||||
if civitai_metadata is None or metadata_provider is None:
|
if civitai_metadata is None or metadata_provider is None:
|
||||||
if sqlite_attempted:
|
if sqlite_attempted:
|
||||||
model_data["db_checked"] = True
|
model_data["db_checked"] = True
|
||||||
|
|
||||||
if last_error == "Model not found":
|
if civitai_api_not_found:
|
||||||
model_data["from_civitai"] = False
|
model_data["from_civitai"] = False
|
||||||
model_data["civitai_deleted"] = True
|
model_data["civitai_deleted"] = True
|
||||||
model_data["db_checked"] = sqlite_attempted or (enable_archive and model_data.get("db_checked", False))
|
model_data["db_checked"] = sqlite_attempted or (enable_archive and model_data.get("db_checked", False))
|
||||||
@@ -266,7 +272,10 @@ class MetadataSyncService:
|
|||||||
return False, error_msg
|
return False, error_msg
|
||||||
|
|
||||||
model_data["from_civitai"] = True
|
model_data["from_civitai"] = True
|
||||||
model_data["civitai_deleted"] = civitai_metadata.get("source") == "archive_db" or civitai_metadata.get("source") == "civarchive"
|
if provider_used is None:
|
||||||
|
model_data["civitai_deleted"] = False
|
||||||
|
elif civitai_api_not_found:
|
||||||
|
model_data["civitai_deleted"] = True
|
||||||
model_data["db_checked"] = enable_archive and (
|
model_data["db_checked"] = enable_archive and (
|
||||||
civitai_metadata.get("source") == "archive_db" or sqlite_attempted
|
civitai_metadata.get("source") == "archive_db" or sqlite_attempted
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -179,6 +179,39 @@ async def test_fetch_and_update_model_success_updates_cache(tmp_path):
|
|||||||
update_cache.assert_awaited_once()
|
update_cache.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fetch_and_update_model_keeps_deleted_flag_false_for_archive_source(tmp_path):
|
||||||
|
helpers = build_service()
|
||||||
|
|
||||||
|
civitai_payload = {
|
||||||
|
"source": "archive_db",
|
||||||
|
"model": {"name": "Recovered", "description": "", "tags": ["tag"]},
|
||||||
|
"images": [],
|
||||||
|
"baseModel": "sd15",
|
||||||
|
}
|
||||||
|
helpers.default_provider.get_model_by_hash.return_value = (civitai_payload, None)
|
||||||
|
|
||||||
|
model_path = tmp_path / "model.safetensors"
|
||||||
|
model_data = {
|
||||||
|
"model_name": "Local",
|
||||||
|
"folder": "root",
|
||||||
|
"file_path": str(model_path),
|
||||||
|
}
|
||||||
|
update_cache = AsyncMock(return_value=True)
|
||||||
|
|
||||||
|
ok, error = await helpers.service.fetch_and_update_model(
|
||||||
|
sha256="abc",
|
||||||
|
file_path=str(model_path),
|
||||||
|
model_data=model_data,
|
||||||
|
update_cache_func=update_cache,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert ok and error is None
|
||||||
|
assert model_data["metadata_source"] == "archive_db"
|
||||||
|
assert model_data["civitai_deleted"] is False
|
||||||
|
update_cache.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_fetch_and_update_model_handles_missing_remote_metadata(tmp_path):
|
async def test_fetch_and_update_model_handles_missing_remote_metadata(tmp_path):
|
||||||
helpers = build_service()
|
helpers = build_service()
|
||||||
@@ -301,6 +334,7 @@ async def test_fetch_and_update_model_prefers_civarchive_for_deleted_models(tmp_
|
|||||||
civarchive_provider.get_model_by_hash.assert_awaited_once_with("deadbeef")
|
civarchive_provider.get_model_by_hash.assert_awaited_once_with("deadbeef")
|
||||||
update_cache.assert_awaited()
|
update_cache.assert_awaited()
|
||||||
assert model_data["metadata_source"] == "civarchive"
|
assert model_data["metadata_source"] == "civarchive"
|
||||||
|
assert model_data["civitai_deleted"] is True
|
||||||
helpers.metadata_manager.save_metadata.assert_awaited()
|
helpers.metadata_manager.save_metadata.assert_awaited()
|
||||||
|
|
||||||
|
|
||||||
@@ -359,6 +393,7 @@ async def test_fetch_and_update_model_falls_back_to_sqlite_after_civarchive_fail
|
|||||||
assert sqlite_provider.get_model_by_hash.await_count == 1
|
assert sqlite_provider.get_model_by_hash.await_count == 1
|
||||||
assert model_data["metadata_source"] == "archive_db"
|
assert model_data["metadata_source"] == "archive_db"
|
||||||
assert model_data["db_checked"] is True
|
assert model_data["db_checked"] is True
|
||||||
|
assert model_data["civitai_deleted"] is True
|
||||||
assert provider_selector.await_args_list[0].args == ("civarchive_api",)
|
assert provider_selector.await_args_list[0].args == ("civarchive_api",)
|
||||||
assert provider_selector.await_args_list[1].args == ("sqlite",)
|
assert provider_selector.await_args_list[1].args == ("sqlite",)
|
||||||
update_cache.assert_awaited()
|
update_cache.assert_awaited()
|
||||||
|
|||||||
Reference in New Issue
Block a user