diff --git a/core/services/civitai_search_service.py b/core/services/civitai_search_service.py index 0bb6185..471463c 100644 --- a/core/services/civitai_search_service.py +++ b/core/services/civitai_search_service.py @@ -6,6 +6,42 @@ from ..sources.common import collect_download_urls from .model_utils import CivitAISearchDependencies, ModelServiceDependencies +def _build_metadata_sidecar_payload( + *, + source, + details_source, + filename, + category, + model_name, + model_type, + sha256, + size, + result, + description, + version_description, + extra_fields=None, +): + payload = { + "source": source, + "details_source": details_source, + "filename": filename, + "category": category, + "model_name": model_name, + "name": model_name, + "model_type": model_type, + "sha256": sha256, + "size": size, + "base_model": result.get("base_model"), + "base_model_source": result.get("base_model_source"), + "base_model_inferred": bool(result.get("base_model_inferred")), + "description": description, + "model_description": description, + "version_description": version_description, + } + payload.update(extra_fields or {}) + return payload + + class CivitAISearchService(ModelServiceDependencies): """Execute exact-match CivitAI and model metadata searches.""" @@ -434,51 +470,51 @@ class CivitAISearchService(ModelServiceDependencies): model_description, version_description = ( split_result_descriptions(result) ) - metadata_payload = { - "source": source_name, - "details_source": result.get("details_source") - or source_name, - "filename": filename, - "category": category, - "model_name": result.get("model_name", clean_name), - "name": result.get("model_name", clean_name), - "model_type": result.get("model_type", "") - or result.get("type", ""), - "type": result.get("model_type", "") - or result.get("type", ""), - "model_id": result.get("model_id"), - "version_id": result.get("version_id"), - "version_name": result.get("version_name", ""), - "sha256": result.get("sha256") or provided_hash, - "size": result.get("size"), - "base_model": result.get("base_model"), - "base_model_source": result.get("base_model_source"), - "base_model_inferred": bool(result.get("base_model_inferred")), - "tags": result.get("tags", []), - "trained_words": result.get("trained_words", []), - "images": result.get("images", []), - "clip_skip": result.get("clip_skip"), - "description": model_description, - "model_description": model_description, - "version_description": version_description, - "download_url": result.get("download_url"), - "source_url": result.get("version_url") or result.get("url"), - "version_url": result.get("version_url") or result.get("url"), - "model_url": result.get("url"), - "url": result.get("version_url") or result.get("url"), - "platform_url": result.get("platform_url"), - "repo_id": result.get("repo_id"), - "path": result.get("path"), - "path_metadata": { - "filename": filename, - "category": category, - "source": source_name, + model_type = result.get("model_type", "") or result.get( + "type", "" + ) + metadata_payload = _build_metadata_sidecar_payload( + source=source_name, + details_source=result.get("details_source") or source_name, + filename=filename, + category=category, + model_name=result.get("model_name", clean_name), + model_type=model_type, + sha256=result.get("sha256") or provided_hash, + size=result.get("size"), + result=result, + description=model_description, + version_description=version_description, + extra_fields={ + "type": model_type, "model_id": result.get("model_id"), "version_id": result.get("version_id"), + "version_name": result.get("version_name", ""), + "tags": result.get("tags", []), + "trained_words": result.get("trained_words", []), + "images": result.get("images", []), + "clip_skip": result.get("clip_skip"), + "download_url": result.get("download_url"), + "source_url": result.get("version_url") + or result.get("url"), + "version_url": result.get("version_url") + or result.get("url"), + "model_url": result.get("url"), + "url": result.get("version_url") or result.get("url"), + "platform_url": result.get("platform_url"), "repo_id": result.get("repo_id"), "path": result.get("path"), + "path_metadata": { + "filename": filename, + "category": category, + "source": source_name, + "model_id": result.get("model_id"), + "version_id": result.get("version_id"), + "repo_id": result.get("repo_id"), + "path": result.get("path"), + }, }, - } + ) add_local_metadata_fields(metadata_payload, result) metadata_path = write_model_resolver_metadata( file_path, @@ -539,47 +575,47 @@ class CivitAISearchService(ModelServiceDependencies): model_description, version_description = ( split_result_descriptions(result) ) - metadata_payload = { - "source": "metadata_import", - "details_source": result.get("source") or "metadata", - "filename": filename, - "category": category, - "model_name": result.get("model_name", clean_name), - "name": result.get("model_name", clean_name), - "model_type": result.get("model_type", ""), - "type": result.get("model_type", ""), - "model_id": result.get("model_id"), - "version_id": result.get("version_id"), - "version_name": result.get("version_name", ""), - "sha256": result.get("sha256"), - "size": result.get("size"), - "base_model": result.get("base_model"), - "base_model_source": result.get("base_model_source"), - "base_model_inferred": bool(result.get("base_model_inferred")), - "tags": result.get("tags", []), - "trained_words": result.get("trained_words", []), - "images": result.get("images", []), - "clip_skip": result.get("clip_skip"), - "description": model_description, - "model_description": model_description, - "version_description": version_description, - "download_url": result.get("download_url"), - "preview_url": preview_url, - "source_url": result.get("version_url") - or result.get("url"), - "version_url": result.get("version_url") - or result.get("url"), - "model_url": result.get("url"), - "url": result.get("version_url") or result.get("url"), - "path_metadata": { - "filename": filename, - "category": category, - "source": "metadata_import", + model_type = result.get("model_type", "") + metadata_payload = _build_metadata_sidecar_payload( + source="metadata_import", + details_source=result.get("source") or "metadata", + filename=filename, + category=category, + model_name=result.get("model_name", clean_name), + model_type=model_type, + sha256=result.get("sha256"), + size=result.get("size"), + result=result, + description=model_description, + version_description=version_description, + extra_fields={ + "type": model_type, "model_id": result.get("model_id"), "version_id": result.get("version_id"), - "imported_from": source_metadata_path, + "version_name": result.get("version_name", ""), + "tags": result.get("tags", []), + "trained_words": result.get("trained_words", []), + "images": result.get("images", []), + "clip_skip": result.get("clip_skip"), + "download_url": result.get("download_url"), + "preview_url": preview_url, + "source_url": result.get("version_url") + or result.get("url"), + "version_url": result.get("version_url") + or result.get("url"), + "model_url": result.get("url"), + "url": result.get("version_url") + or result.get("url"), + "path_metadata": { + "filename": filename, + "category": category, + "source": "metadata_import", + "model_id": result.get("model_id"), + "version_id": result.get("version_id"), + "imported_from": source_metadata_path, + }, }, - } + ) add_local_metadata_fields(metadata_payload, result) metadata_path = write_model_resolver_metadata( file_path, @@ -805,34 +841,29 @@ class CivitAISearchService(ModelServiceDependencies): response["url"] = None if force_refresh and file_path and _os.path.exists(file_path): try: - metadata_payload = { - "filename": filename, - "category": category, - "model_name": response.get("model_name") or clean_name, - "name": response.get("model_name") or clean_name, - "model_type": response.get("model_type") + model_description = response.get("model_description") or response.get( + "description" + ) or "" + metadata_payload = _build_metadata_sidecar_payload( + source="local", + details_source=response.get("details_source") or "", + filename=filename, + category=category, + model_name=response.get("model_name") or clean_name, + model_type=response.get("model_type") or infer_model_type_from_category(category), - "sha256": response.get("sha256") or provided_hash, - "size": response.get("size") or _os.path.getsize(file_path), - "base_model": response.get("base_model"), - "base_model_source": response.get("base_model_source"), - "base_model_inferred": bool(response.get("base_model_inferred")), - "description": response.get("model_description") - or response.get("description") + sha256=response.get("sha256") or provided_hash, + size=response.get("size") or _os.path.getsize(file_path), + result=response, + description=model_description, + version_description=response.get("version_description") or "", - "model_description": response.get("model_description") - or response.get("description") - or "", - "version_description": response.get( - "version_description" - ) - or "", - "civitai_deleted": True, - "civitai_checked": True, - "remote_metadata_missing": True, - "source": "local", - "details_source": response.get("details_source") or "", - } + extra_fields={ + "civitai_deleted": True, + "civitai_checked": True, + "remote_metadata_missing": True, + }, + ) add_local_metadata_fields(metadata_payload, response) metadata_path = write_model_resolver_metadata( file_path, diff --git a/tests/test_refactoring_targets.py b/tests/test_refactoring_targets.py index c1a4df4..741de5e 100644 --- a/tests/test_refactoring_targets.py +++ b/tests/test_refactoring_targets.py @@ -1384,6 +1384,244 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): payload = json.loads(api_response.text) self.assertEqual(payload["download_url"], mirror_url) + async def test_civitai_service_preserves_remote_sidecar_payload_fields(self): + expected_hash = "b" * 64 + remote_result = { + "source": "civitai", + "model_name": "Remote Model", + "model_type": "Checkpoint", + "model_id": 12, + "version_id": 34, + "version_name": "v1", + "sha256": expected_hash, + "size": 123, + "base_model": "SDXL", + "base_model_source": "civitai", + "tags": ["tag"], + "trained_words": ["trigger"], + "images": [{"url": "https://example.com/preview.png"}], + "clip_skip": 2, + "description": "Model description", + "version_description": "Version description", + "url": "https://civitai.com/models/12", + "version_url": "https://civitai.com/models/12?modelVersionId=34", + "download_url": "https://civitai.com/api/download/models/34", + "author": "Author", + } + write_metadata = MagicMock(return_value="remote.modelresolver.json") + + def to_bool(value, default=False): + return default if value is None else bool(value) + + service = self._build_model_service_for_unit_test( + service_module="civitai_search_service", + service_class="CivitAISearchService", + download_available=True, + extract_sha256_from_metadata=lambda value: ( + value.get("sha256", "") if isinstance(value, dict) else "" + ), + find_external_metadata_sidecar_path=lambda path: "", + get_filename_from_path=lambda value: str(value or "").replace( + "\\", "/" + ).rsplit("/", 1)[-1], + is_path_in_configured_model_roots=lambda path: True, + normalize_category_to_model_type=lambda value: "Checkpoint", + normalize_sha256=lambda value: str(value or ""), + read_json_safe=lambda path, default=None: default, + to_bool=to_bool, + write_model_resolver_metadata=write_metadata, + ) + with tempfile.NamedTemporaryFile( + suffix=".safetensors", delete=False + ) as handle: + model_path = handle.name + request = AsyncMock() + request.json.return_value = { + "filename": "remote.safetensors", + "category": "checkpoints", + "resolved_path": model_path, + "sha256": expected_hash, + } + + try: + civitai_module = importlib.import_module( + "comfyui-model-resolver.core.sources.civitai" + ) + with patch.object( + civitai_module, + "get_model_info_by_hash", + return_value=remote_result, + ): + response = await service.civitai_search(request) + finally: + os.unlink(model_path) + + self.assertEqual(response.status, 200) + metadata_payload = write_metadata.call_args.args[1] + self.assertEqual("civitai", metadata_payload["source"]) + self.assertEqual("Remote Model", metadata_payload["model_name"]) + self.assertEqual(expected_hash, metadata_payload["sha256"]) + self.assertEqual("Author", metadata_payload["author"]) + self.assertEqual( + { + "filename": "remote.safetensors", + "category": "checkpoints", + "source": "civitai", + "model_id": 12, + "version_id": 34, + "repo_id": None, + "path": None, + }, + metadata_payload["path_metadata"], + ) + + async def test_civitai_service_preserves_metadata_import_sidecar_fields(self): + expected_hash = "c" * 64 + imported_from = "external.metadata.json" + imported_result = { + "source": "civitai", + "from_metadata": True, + "metadata_path": imported_from, + "model_name": "Imported Model", + "model_type": "Lora", + "model_id": 56, + "version_id": 78, + "version_name": "v2", + "sha256": expected_hash, + "size": 456, + "description": "Imported description", + "preview_url": "https://example.com/imported.png", + "url": "https://civitai.com/models/56", + "version_url": "https://civitai.com/models/56?modelVersionId=78", + } + write_metadata = MagicMock(return_value="imported.modelresolver.json") + + def to_bool(value, default=False): + return default if value is None else bool(value) + + service = self._build_model_service_for_unit_test( + service_module="civitai_search_service", + service_class="CivitAISearchService", + download_available=True, + extract_sha256_from_metadata=lambda value: ( + value.get("sha256", "") if isinstance(value, dict) else "" + ), + find_external_metadata_sidecar_path=lambda path: "", + get_existing_model_preview_path=lambda path: "", + get_filename_from_path=lambda value: str(value or "").replace( + "\\", "/" + ).rsplit("/", 1)[-1], + get_model_resolver_sidecar_path=lambda path: f"{path}.modelresolver.json", + is_path_in_configured_model_roots=lambda path: True, + normalize_category_to_model_type=lambda value: "Lora", + normalize_sha256=lambda value: str(value or ""), + read_json_safe=lambda path, default=None: default, + to_bool=to_bool, + write_model_resolver_metadata=write_metadata, + ) + with tempfile.NamedTemporaryFile( + suffix=".safetensors", delete=False + ) as handle: + model_path = handle.name + request = AsyncMock() + request.json.return_value = { + "filename": "imported.safetensors", + "category": "loras", + "resolved_path": model_path, + "local_only": True, + } + + try: + civitai_module = importlib.import_module( + "comfyui-model-resolver.core.sources.civitai" + ) + with patch.object( + civitai_module, + "get_model_info_for_file", + return_value=imported_result, + ): + response = await service.civitai_search(request) + finally: + os.unlink(model_path) + + self.assertEqual(response.status, 200) + metadata_payload = write_metadata.call_args.args[1] + self.assertEqual("metadata_import", metadata_payload["source"]) + self.assertEqual("https://example.com/imported.png", metadata_payload["preview_url"]) + self.assertEqual( + imported_from, + metadata_payload["path_metadata"]["imported_from"], + ) + self.assertEqual( + "imported.modelresolver.json", + json.loads(response.text)["metadata_path"], + ) + + async def test_civitai_service_preserves_missing_remote_metadata_flags(self): + local_result = { + "source": "local", + "base_model": "SDXL", + "base_model_source": "safetensors_header", + "base_model_inferred": True, + "description": "Local description", + "sha256": "d" * 64, + } + write_metadata = MagicMock(return_value="deleted.modelresolver.json") + + def to_bool(value, default=False): + return default if value is None else bool(value) + + service = self._build_model_service_for_unit_test( + service_module="civitai_search_service", + service_class="CivitAISearchService", + download_available=True, + extract_sha256_from_metadata=lambda value: ( + value.get("sha256", "") if isinstance(value, dict) else "" + ), + find_external_metadata_sidecar_path=lambda path: "", + get_filename_from_path=lambda value: str(value or "").replace( + "\\", "/" + ).rsplit("/", 1)[-1], + is_path_in_configured_model_roots=lambda path: True, + normalize_category_to_model_type=lambda value: "Checkpoint", + normalize_sha256=lambda value: str(value or ""), + read_json_safe=lambda path, default=None: default, + search_huggingface_for_file=MagicMock(return_value=None), + to_bool=to_bool, + write_model_resolver_metadata=write_metadata, + ) + with tempfile.NamedTemporaryFile( + suffix=".safetensors", delete=False + ) as handle: + model_path = handle.name + request = AsyncMock() + request.json.return_value = { + "filename": "deleted.safetensors", + "category": "checkpoints", + "resolved_path": model_path, + "force_refresh": True, + } + + try: + civitai_module = importlib.import_module( + "comfyui-model-resolver.core.sources.civitai" + ) + with patch.object( + civitai_module, + "get_model_info_for_file", + return_value=local_result, + ): + response = await service.civitai_search(request) + finally: + os.unlink(model_path) + + self.assertEqual(response.status, 200) + metadata_payload = write_metadata.call_args.args[1] + self.assertEqual("local", metadata_payload["source"]) + self.assertTrue(metadata_payload["civitai_deleted"]) + self.assertTrue(metadata_payload["civitai_checked"]) + self.assertTrue(metadata_payload["remote_metadata_missing"]) + async def test_custom_url_service_rejects_unsafe_url(self): class TestUnsafeUrlError(Exception): pass