From 89a721fc6224ddbd8753348b80f84c5e5633eb57 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sat, 18 Jul 2026 13:43:51 +0200 Subject: [PATCH] Normalize CivitAI search queries and simplify tRPC search behavior --- core/sources/civitai.py | 73 +++++++++++++++++----------------- core/type_utils.py | 4 ++ tests/test_source_internals.py | 45 ++++++++------------- tests/test_unified_helpers.py | 2 +- 4 files changed, 58 insertions(+), 66 deletions(-) diff --git a/core/sources/civitai.py b/core/sources/civitai.py index e43a1b1..8b31301 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -371,11 +371,14 @@ def _search_civitai_trpc_candidates( ) -> List[Dict[str, Optional[int]]]: """Try CivitAI.red tRPC search endpoint and log the raw outcome for diagnostics.""" civitai_type = CIVITAI_API_TYPE_MAP.get(str(model_type).lower()) if model_type else None + search_query = get_filename_from_path(filename) + if _has_known_model_extension(search_query): + search_query = os.path.splitext(search_query)[0] headers = { "accept": "application/json", "content-type": "application/json", - "referer": f"https://civitai.red/models?query={quote(filename)}", + "referer": f"https://civitai.red/models?query={quote(search_query)}", "user-agent": DEFAULT_BROWSER_USER_AGENT, "x-client": "web", "x-client-version": "5.0.1657", @@ -389,7 +392,7 @@ def _search_civitai_trpc_candidates( "period": "Month", "periodMode": "stats", "sort": "Highest Rated", - "query": filename, + "query": search_query, "pending": False, "browsingLevel": 28, "excludedTagIds": [ @@ -423,43 +426,35 @@ def _search_civitai_trpc_candidates( + quote(json.dumps(input_payload, separators=(",", ":"))) ) - type_filters = [civitai_type] if civitai_type else [None] - if civitai_type: - type_filters.append(None) + url = build_url(civitai_type) + log.info( + f"CivitAI tRPC search start: filename={filename}, model_type={model_type}, type_filter={civitai_type or 'none'}, session_token={'yes' if session_token else 'no'}, url={url}" + ) - for type_filter in type_filters: - url = build_url(type_filter) - log.info( - f"CivitAI tRPC search start: filename={filename}, model_type={model_type}, type_filter={type_filter or 'none'}, session_token={'yes' if session_token else 'no'}, url={url}" - ) + response = request_source_response(url, headers=headers, timeout=timeout, log_name="CivitAI tRPC") + if response is None: + log.warning(f"CivitAI tRPC request failed for filename={filename}: no response") + return [] - response = request_source_response(url, headers=headers, timeout=timeout, log_name="CivitAI tRPC") - if response is None: - log.warning(f"CivitAI tRPC request failed for filename={filename}: no response") - return [] + text_preview = response.text[:800].replace("\n", " ").replace("\r", " ") + log.info( + f"CivitAI tRPC response: status={response.status_code}, content_type={response.headers.get('content-type')}, text_preview={text_preview}" + ) - text_preview = response.text[:800].replace("\n", " ").replace("\r", " ") - log.info( - f"CivitAI tRPC response: status={response.status_code}, content_type={response.headers.get('content-type')}, text_preview={text_preview}" - ) + if response.status_code != 200: + return [] - if response.status_code != 200: - return [] + try: + payload = response.json() + except Exception as e: + log.warning(f"CivitAI tRPC JSON parse failed for filename={filename}: {e}") + return [] - try: - payload = response.json() - except Exception as e: - log.warning(f"CivitAI tRPC JSON parse failed for filename={filename}: {e}") - return [] - - candidates = _extract_trpc_model_candidates(payload, limit=limit) - log.info( - f"CivitAI tRPC extracted {len(candidates)} candidates for filename={filename}: {candidates}" - ) - if candidates or not type_filter: - return candidates - - return [] + candidates = _extract_trpc_model_candidates(payload, limit=limit) + log.info( + f"CivitAI tRPC extracted {len(candidates)} candidates for filename={filename}: {candidates}" + ) + return candidates def _search_civitai_red_candidates( @@ -470,9 +465,12 @@ def _search_civitai_red_candidates( limit: int = 5, ) -> List[Dict[str, int]]: """Search civitai.red by full filename and return model/version candidates.""" + search_query = get_filename_from_path(filename) + if _has_known_model_extension(search_query): + search_query = os.path.splitext(search_query)[0] params = { "sortBy": "models_v9", - "query": filename, + "query": search_query, } civitai_type = CIVITAI_API_TYPE_MAP.get(str(model_type).lower()) if model_type else None if civitai_type: @@ -547,8 +545,11 @@ def _search_civitai_public_api_candidates( limit: int = 5, ) -> List[Dict[str, Optional[int]]]: """Search the public CivitAI API by model name and return model/version candidates.""" + search_query = get_filename_from_path(filename) + if _has_known_model_extension(search_query): + search_query = os.path.splitext(search_query)[0] params = { - "query": filename, + "query": search_query, "limit": max(1, min(int(limit), MAX_CIVITAI_CANDIDATE_LIMIT)), "sort": "Highest Rated", "period": "AllTime", diff --git a/core/type_utils.py b/core/type_utils.py index 4ee5be6..3b3b132 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -213,6 +213,10 @@ CATEGORY_MAP = { CIVITAI_API_TYPE_MAP = { "checkpoint": "Checkpoint", "checkpoints": "Checkpoint", + "diffusion_model": "Checkpoint", + "diffusion_models": "Checkpoint", + "unet": "Checkpoint", + "unet_gguf": "Checkpoint", "lora": "LORA", "loras": "LORA", "vae": "VAE", diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index 910d1d3..0730bf6 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -568,19 +568,22 @@ class CivitaiRedSearchTests(unittest.TestCase): response.text = "" mock_get.return_value = response - _search_civitai_red_candidates("ae.safetensors", model_type="vae") + _search_civitai_red_candidates( + r"models\diffusion_models\Anima_baseV10.safetensors", + model_type="diffusion_models", + ) url = mock_get.call_args.args[0] self.assertTrue(url.startswith("https://civitai.red/search/models?")) self.assertIn("sortBy=models_v9", url) - self.assertIn("query=ae.safetensors", url) - self.assertIn("modelType=VAE", url) + self.assertIn("query=Anima_baseV10", url) + self.assertIn("modelType=Checkpoint", url) class CivitaiTrpcSearchTests(unittest.TestCase): @patch("core.sources.civitai.requests.get") - def test_retries_without_type_filter_when_typed_search_is_empty(self, mock_get): + def test_does_not_retry_without_type_filter_when_typed_search_is_empty(self, mock_get): empty_response = MagicMock() empty_response.status_code = 200 empty_response.headers = {"content-type": "application/json"} @@ -589,35 +592,19 @@ class CivitaiTrpcSearchTests(unittest.TestCase): "result": {"data": {"json": {"items": []}}} } - fallback_response = MagicMock() - fallback_response.status_code = 200 - fallback_response.headers = {"content-type": "application/json"} - fallback_response.text = '{"result":{"data":{"json":{"items":[{"id":10,"version":{"id":20}}]}}}}' - fallback_response.json.return_value = { - "result": { - "data": { - "json": { - "items": [ - {"id": 10, "version": {"id": 20}}, - ] - } - } - } - } - mock_get.side_effect = [empty_response, fallback_response] + mock_get.return_value = empty_response result = _search_civitai_trpc_candidates( - "ae.safetensors", - model_type="vae", + "Anima_baseV10.safetensors", + model_type="diffusion_models", limit=5, ) - self.assertEqual(result, [{"model_id": 10, "version_id": 20}]) - self.assertEqual(mock_get.call_count, 2) - first_url = mock_get.call_args_list[0].args[0] - second_url = mock_get.call_args_list[1].args[0] - self.assertIn("types", first_url) - self.assertNotIn("types", second_url) + self.assertEqual(result, []) + mock_get.assert_called_once() + url = mock_get.call_args.args[0] + self.assertIn("Anima_baseV10", url) + self.assertIn("Checkpoint", url) def test_html_fallback_not_used_to_top_up_trpc_candidates(self): clear_search_cache() @@ -713,7 +700,7 @@ class CivitaiTrpcSearchTests(unittest.TestCase): self.assertEqual(result, [{"model_id": 101, "version_id": 202}]) params = mock_get.call_args.kwargs.get("params", {}) headers = mock_get.call_args.kwargs.get("headers", {}) - self.assertEqual(params.get("query"), "realistic.safetensors") + self.assertEqual(params.get("query"), "realistic") self.assertEqual(params.get("types"), "Checkpoint") self.assertEqual(headers.get("Authorization"), "Bearer api-key") self.assertIn("__Secure-civitai-token=session-token", headers.get("Cookie", "")) diff --git a/tests/test_unified_helpers.py b/tests/test_unified_helpers.py index 5dfd506..fac5ad4 100644 --- a/tests/test_unified_helpers.py +++ b/tests/test_unified_helpers.py @@ -1195,6 +1195,7 @@ class UnifiedHelpersTests(unittest.TestCase): self.assertEqual(resolve_model_category("checkpoints", target_format="folder"), "checkpoints") self.assertEqual(resolve_model_category("lora", target_format="folder"), "loras") self.assertEqual(resolve_model_category("checkpoint", target_format="civitai"), "Checkpoint") + self.assertEqual(resolve_model_category("diffusion_models", target_format="civitai"), "Checkpoint") self.assertEqual(resolve_model_category("loras", target_format="civitai"), "LORA") self.assertEqual(resolve_model_category("locon", target_format="civarchive"), "LoCon") self.assertEqual(resolve_model_category("loras", target_format="urn"), "loras") @@ -1227,4 +1228,3 @@ class UnifiedHelpersTests(unittest.TestCase): -