diff --git a/core/services/custom_url_service.py b/core/services/custom_url_service.py index 453aea9..88f11a5 100644 --- a/core/services/custom_url_service.py +++ b/core/services/custom_url_service.py @@ -8,6 +8,7 @@ from ..request_utils import ( read_text_field, ) from ..routes.context import RouteContext +from ..sources.civitai import get_civitai_image_version_ids, parse_civitai_image_url from .model_utils import CustomUrlDependencies, ModelServiceDependencies @@ -175,6 +176,8 @@ class CustomUrlService(ModelServiceDependencies): ) result = None + results = [] + unresolved_versions = [] source = "" try: civitai_parsed = parse_civitai_url(normalized_url) @@ -185,7 +188,31 @@ class CustomUrlService(ModelServiceDependencies): except Exception: civarchive_parsed = None - if civitai_parsed: + image_id = parse_civitai_image_url(normalized_url) + if image_id: + source = "civitai" + version_ids = await asyncio.to_thread( + get_civitai_image_version_ids, image_id, civitai_key or None, + ) + for version_id in version_ids: + image_result = await asyncio.to_thread( + resolve_civitai_version_custom_result, version_id, "", civitai_key or None, + ) + if isinstance(image_result, SearchResult): + results.append(image_result) + else: + unresolved_versions.append(version_id) + if not results: + return web.json_response( + {"error": ( + "No downloadable linked models could be resolved for this CivitAI image." + if version_ids else + "CivitAI did not provide any linked model versions for this image." + )}, + status=400, + ) + result = results[0] + elif civitai_parsed: source = "civitai" model_id = civitai_parsed.model_id version_id = civitai_parsed.version_id @@ -297,7 +324,7 @@ class CustomUrlService(ModelServiceDependencies): { "error": ( "Unsupported or unresolved URL. Use a HuggingFace file URL, " - "CivitAI model/download URL, or CivArchive model/hash URL." + "CivitAI model/image/download URL, or CivArchive model/hash URL." ) }, status=400, @@ -310,15 +337,18 @@ class CustomUrlService(ModelServiceDependencies): status=400, ) - result = result.with_extra( + results = results or [result] + results = [item.with_extra( provided_url=normalized_url, url_source="custom", - searched_at=result.extra_value("searched_at") + searched_at=item.extra_value("searched_at") or _custom_result_timestamp(), - category=result.extra_value("category") or category, - ) + category=item.extra_value("category") or ("" if image_id else category), + ) for item in results] + result = results[0] if expected_filename and not result.filename: result = result.with_updates(filename=expected_filename) + results[0] = result if not (result.download_url or result.url): return web.json_response( @@ -327,16 +357,18 @@ class CustomUrlService(ModelServiceDependencies): ) source = source or result.source or "custom" - local_hash_matches = _collect_custom_url_local_hash_matches( - result, - category, - ) + local_hash_matches = [] + for item in results: + local_hash_matches.extend(_collect_custom_url_local_hash_matches( + item, "" if image_id else category, + )) result_mapping = result.to_dict() response = { "success": True, "source": source, "result": result_mapping, - "custom": [result_mapping], + "custom": [item.to_dict() for item in results], + "unresolved_version_ids": unresolved_versions, "searched_sources": ["custom"], "local_hash_matches": [ match.to_dict() for match in local_hash_matches diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 3ffc3b5..c7bdf1d 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -8,7 +8,7 @@ import json import os import re from typing import Any, Callable, Dict, List, Optional -from urllib.parse import quote, urlencode +from urllib.parse import quote, urlencode, urlparse import requests @@ -2569,6 +2569,94 @@ def build_civitai_custom_result( ) +def parse_civitai_image_url(url: str) -> Optional[int]: + """Recognize image pages without accepting lookalike provider hosts.""" + parsed = urlparse(url) + if parsed.scheme not in {"http", "https"} or parsed.hostname not in { + "civitai.com", "www.civitai.com", "civitai.red", "www.civitai.red", + }: + return None + match = re.fullmatch(r"/images/([1-9]\d*)/?", parsed.path) + return int(match.group(1)) if match else None + + +def _extract_civitai_image_page_version_ids(html_text: str, image_id: int) -> List[int]: + """Read only generation resources belonging to the requested image.""" + match = re.search( + r']*\bid=["\']__NEXT_DATA__["\'])[^>]*>(.*?)', + html_text, re.DOTALL | re.IGNORECASE, + ) + if not match: + return [] + try: + data = as_dict(json.loads(match.group(1))) + except (TypeError, ValueError): + return [] + page = as_dict(as_dict(data.get("props")).get("pageProps")) + state = as_dict(page.get("trpcState")) + state = as_dict(state.get("json")) or state + versions = [] + for query in as_list(state.get("queries")): + query = as_dict(query) + key = as_list(query.get("queryKey")) + if len(key) != 2 or key[0] != ["image", "getGenerationData"]: + continue + requested_id = as_dict(as_dict(key[1]).get("input")).get("id") + if str(requested_id) != str(image_id): + continue + generation = as_dict(as_dict(query.get("state")).get("data")) + for resource in as_list(generation.get("resources")): + resource = as_dict(resource) + if str(resource.get("imageId", image_id)) != str(image_id): + continue + value = resource.get("modelVersionId") or resource.get("versionId") + if isinstance(value, bool): + continue + try: + version_id = int(value) + except (TypeError, ValueError): + continue + if version_id > 0 and version_id not in versions: + versions.append(version_id) + return versions + + +def get_civitai_image_version_ids(image_id: int, api_key: Optional[str] = None) -> List[int]: + """Read the exact linked versions from the public image API.""" + data = execute_provider_json_request( + "CivitAI image resources", f"{CIVITAI_API_URL}/images", + params={"imageId": image_id}, api_key=api_key, timeout=20, + ) + for image in as_list(as_dict(data).get("items")): + if not isinstance(image, dict) or str(image.get("id")) != str(image_id): + continue + versions = [] + for value in as_list(image.get("modelVersionIds")): + if isinstance(value, bool): + continue + try: + version_id = int(value) + except (ValueError, TypeError): + continue + if version_id > 0 and version_id not in versions: + versions.append(version_id) + if versions: + return versions + break + # Public image results can omit manually attached resources even though the + # image page includes them in its server-rendered generation data. + headers = {"User-Agent": DEFAULT_BROWSER_USER_AGENT, "Accept": "text/html"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + response = request_source_response( + f"https://civitai.com/images/{image_id}", headers=headers, + timeout=20, log_name="CivitAI image page", + ) + if response is not None and response.status_code == 200: + return _extract_civitai_image_page_version_ids(response.text, image_id) + return [] + + def resolve_civitai_version_custom_result( version_id: int, expected_filename: str = "", diff --git a/tests/test_civitai_image_urls.py b/tests/test_civitai_image_urls.py new file mode 100644 index 0000000..72e0e02 --- /dev/null +++ b/tests/test_civitai_image_urls.py @@ -0,0 +1,52 @@ +import json +from unittest.mock import Mock, patch + +from core.sources.civitai import get_civitai_image_version_ids, parse_civitai_image_url + + +def test_image_url_host_and_path_validation(): + assert parse_civitai_image_url("https://civitai.com/images/144201545?foo=bar") == 144201545 + assert parse_civitai_image_url("https://civitai.red/images/123/") == 123 + for url in ("https://civitai.com.evil.test/images/123", "https://civitai.com/models/123", "https://civitai.com/images/0"): + assert parse_civitai_image_url(url) is None + + +def test_image_resources_require_matching_image_and_unique_valid_versions(): + payload = {"items": [ + {"id": 9, "modelVersionIds": [999]}, + {"id": 144201545, "meta": None, "modelVersionIds": [3352534, 3356151, 3352534, None, True, -1, "invalid"]}, + ]} + with patch("core.sources.civitai.execute_provider_json_request", return_value=payload), \ + patch("core.sources.civitai.request_source_response", return_value=None): + assert get_civitai_image_version_ids(144201545) == [3352534, 3356151] + assert get_civitai_image_version_ids(100) == [] + + +def image_page(image_id, resources): + query = { + "queryKey": [["image", "getGenerationData"], {"input": {"id": image_id}}], + "state": {"data": {"resources": resources}}, + } + payload = {"props": {"pageProps": {"trpcState": {"json": {"queries": [query]}}}}} + return '' + + +def test_empty_public_image_resources_fall_back_to_page_generation_data(): + page = image_page(143748745, [ + {"imageId": 143748745, "modelVersionId": 3356151}, + {"imageId": 143748745, "versionId": 3356151}, + {"imageId": 999, "modelVersionId": 123}, + {"modelVersionId": True}, + ]) + with patch("core.sources.civitai.execute_provider_json_request", return_value={ + "items": [{"id": 143748745, "meta": None, "modelVersionIds": []}], + }), patch("core.sources.civitai.request_source_response", return_value=Mock(status_code=200, text=page)): + assert get_civitai_image_version_ids(143748745) == [3356151] + assert get_civitai_image_version_ids(100) == [] + + +def test_image_page_fallback_ignores_malformed_or_unavailable_page(): + for response in (None, Mock(status_code=403, text="blocked"), Mock(status_code=200, text='')): + with patch("core.sources.civitai.execute_provider_json_request", return_value=None), \ + patch("core.sources.civitai.request_source_response", return_value=response): + assert get_civitai_image_version_ids(143748745) == [] diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index 42ac2b9..529afc8 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -2,6 +2,34 @@ import fs from 'node:fs'; import path from 'node:path'; import test from 'node:test'; import assert from 'node:assert/strict'; + +test('image URL adds every linked version and deduplicates repeated additions', async () => { + const methodStart = resolveDownloadMethodsSource.indexOf('async addCustomUrlResult('); + const methodEnd = resolveDownloadMethodsSource.indexOf('\n },', methodStart); + const addCustomUrlResult = eval(`(${resolveDownloadMethodsSource.slice(methodStart, methodEnd + 6).replace('async addCustomUrlResult', 'async function addCustomUrlResult')})`); + const state = { results: { custom: [] } }; + const models = [ + { source: 'civitai', version_id: 3352534, filename: 'qwen.safetensors' }, + { source: 'civitai', version_id: 3356151, filename: 'adapter.safetensors' }, + ]; + const dialog = { + getWorkflowScopedQueueKey: () => 'workflow', + getSearchStateForWorkflow: () => state, + getStoredTokens: () => ({}), + getFilenameFromPath: value => value, + fetchJson: async () => ({ result: models[0], custom: models }), + getSearchResultSignature: result => String(result.version_id), + persistSearchStateForWorkflow() {}, + refreshSearchUiForMissing() {}, + }; + const missing = { original_path: 'missing.safetensors' }; + const input = { value: 'https://civitai.com/images/144201545' }; + await addCustomUrlResult.call(dialog, missing, input, null); + input.value = 'https://civitai.com/images/144201545'; + await addCustomUrlResult.call(dialog, missing, input, null); + assert.deepEqual(state.results.custom.map(result => result.version_id), [3352534, 3356151]); + assert.equal(input.disabled, false); +}); import { Window } from 'happy-dom'; import { html, diff --git a/tests/test_search_orchestration.py b/tests/test_search_orchestration.py index 27d5b57..5cc1151 100644 --- a/tests/test_search_orchestration.py +++ b/tests/test_search_orchestration.py @@ -192,6 +192,33 @@ class SearchOrchestrationTests(unittest.IsolatedAsyncioTestCase): # Verify response structure and status self.assertIsInstance(response, web.Response) + async def test_custom_url_image_keeps_all_versions_and_reports_partial_failure(self): + service = importlib.import_module("comfyui-model-resolver.core.services.custom_url_service") + first = typed_contracts.SearchResult( + source="civitai", model_id=2954443, version_id=3352534, + filename="qwen.safetensors", download_url="https://civitai.com/api/download/models/3352534", + ) + second = typed_contracts.SearchResult( + source="civitai", model_id=2962344, version_id=3356151, + filename="adapter.safetensors", download_url="https://civitai.com/api/download/models/3356151", + ) + request = MagicMock() + request.json = AsyncMock(return_value={ + "url": "https://civitai.com/images/144201545", + "filename": "unrelated.safetensors", "category": "checkpoints", + }) + with patch.object(service, "get_civitai_image_version_ids", return_value=[3352534, 3356151, 99]), \ + patch.object(civitai_sources, "resolve_civitai_version_custom_result", side_effect=[first, second, None]) as resolver: + self.ext.routes_setup = False + self.ext.setup_routes() + response = await custom_url_handler(request) + body = json.loads(response.text) + self.assertEqual(response.status, 200, body) + self.assertEqual([item["version_id"] for item in body["custom"]], [3352534, 3356151]) + self.assertEqual([item["filename"] for item in body["custom"]], ["qwen.safetensors", "adapter.safetensors"]) + self.assertEqual(body["unresolved_version_ids"], [99]) + self.assertEqual(resolver.call_args_list[0].args[1], "") + async def test_custom_url_route_resolves_civitai_model_link(self): self.mock_civitai_details.return_value = { "source": "civitai", diff --git a/web/resolver/actions/resolve_download_methods.js b/web/resolver/actions/resolve_download_methods.js index b96a4e4..22699a0 100644 --- a/web/resolver/actions/resolve_download_methods.js +++ b/web/resolver/actions/resolve_download_methods.js @@ -2523,22 +2523,28 @@ export const resolveDownloadMethods = { silent: true }, 'Add custom URL'); - const result = data?.result || (Array.isArray(data?.custom) ? data.custom[0] : null); - if (!result || typeof result !== 'object') { + const results = (Array.isArray(data?.custom) ? data.custom : [data?.result]) + .filter(result => result && typeof result === 'object'); + if (!results.length) { throw new Error('The URL did not resolve to a downloadable model.'); } const searchedAt = new Date().toISOString(); - const customResult = this.withSearchResultTimestamp?.(result, result.searched_at || searchedAt) || result; + const customResults = results.map(result => ( + this.withSearchResultTimestamp?.(result, result.searched_at || searchedAt) || result + )); + const customResult = customResults[0]; const currentResults = state.results || this.createEmptySearchState().results; const existingCustom = Array.isArray(currentResults.custom) ? currentResults.custom : []; - const nextSignature = this.getSearchResultSignature(customResult) || (customResult.provided_url || url); - const nextCustom = [ - ...existingCustom.filter(existing => ( - (this.getSearchResultSignature(existing) || existing.provided_url || '') !== nextSignature - )), - customResult - ]; + const nextCustom = [...existingCustom]; + for (const item of customResults) { + const signature = this.getSearchResultSignature(item) || item.version_url || item.download_url || url; + const existingIndex = nextCustom.findIndex(existing => ( + (this.getSearchResultSignature(existing) || existing.version_url || existing.download_url || '') === signature + )); + if (existingIndex >= 0) nextCustom[existingIndex] = item; + else nextCustom.push(item); + } const localHashMatches = this.mergeLocalMatches ? this.mergeLocalMatches( Array.isArray(currentResults.local_hash_matches) ? currentResults.local_hash_matches : [], @@ -2576,7 +2582,13 @@ export const resolveDownloadMethods = { this.refreshSearchUiForMissing(missing, state, { workflowKey }); this.applySearchResultSuggestion?.(missing); if (inputEl) inputEl.value = ''; - this.showNotification?.('Link added.', 'success'); + const unresolvedCount = data?.unresolved_version_ids?.length || 0; + this.showNotification?.( + unresolvedCount + ? `Added ${customResults.length} models. ${unresolvedCount} linked versions could not be resolved.` + : customResults.length > 1 ? `Added ${customResults.length} models.` : 'Link added.', + unresolvedCount ? 'warning' : 'success' + ); } catch (error) { console.error('Model Resolver: custom URL add failed:', error); this.showNotification?.(error.message || 'Failed to add link.', 'error');