nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Normalize CivitAI search queries and simplify tRPC search behavior

Browse Source
This commit is contained in:
Dariusz L
2026-07-18 13:43:51 +02:00
parent c39eea4f51
commit 89a721fc62
4 changed files with 58 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/sources/civitai.py
+37 -36
View File
@@ -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",
core/type_utils.py
+4
View File
@@ -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",
tests/test_source_internals.py
+16 -29
View File
@@ -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", ""))
tests/test_unified_helpers.py
+1 -1
View File
@@ -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):