Support linked models from Civitai image URLs with page fallback

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

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'<script\b(?=[^>]*\bid=["\']__NEXT_DATA__["\'])[^>]*>(.*?)</script>',
+        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 = "",

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 '<script id="__NEXT_DATA__" type="application/json">' + json.dumps(payload) + '</script>'
+
+
+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='<script id="__NEXT_DATA__">invalid</script>')):
+        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) == []

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,

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",

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');