diff --git a/__init__.py b/__init__.py index 2cb60e8..d37d0b4 100644 --- a/__init__.py +++ b/__init__.py @@ -424,6 +424,7 @@ class ModelResolverExtension: get_all_progress, get_aria2_status, get_download_directory, + get_existing_model_preview_path, get_progress, is_allowed_model_download_filename, normalize_download_category, @@ -934,6 +935,41 @@ class ModelResolverExtension: get_local_model_hash_metadata(normalized_path, model=model) ) + @routes.get("/model_resolver/model-preview") + async def get_model_preview(request): + """Serve an adjacent model preview from a configured model directory.""" + model_path = str(request.query.get("path") or "").strip() + if not model_path: + return web.Response(text="path is required", status=400) + + try: + normalized_path = os.path.realpath( + os.path.abspath(os.path.normpath(model_path)) + ) + except (OSError, TypeError, ValueError): + return web.Response(text="invalid model path", status=400) + if not os.path.isfile(normalized_path): + return web.Response(text="model file does not exist", status=404) + if not is_path_in_configured_model_roots(normalized_path): + return web.Response( + text="path is outside configured model directories", + status=403, + ) + + preview_path = get_existing_model_preview_path(normalized_path) + if not preview_path: + return web.Response(text="preview not found", status=404) + if not is_path_in_configured_model_roots(preview_path): + return web.Response( + text="preview is outside configured model directories", + status=403, + ) + + return web.FileResponse( + preview_path, + headers={"Cache-Control": "private, max-age=300"}, + ) + @routes.post("/model_resolver/workflow-model-hashes") @json_api_endpoint("workflow-model-hashes") async def workflow_model_hashes(request): diff --git a/core/downloader.py b/core/downloader.py index f541208..db5bcbb 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -719,7 +719,8 @@ def build_lora_manager_metadata( return payload -def _existing_model_preview_path(dest_path: str) -> str: +def get_existing_model_preview_path(dest_path: str) -> str: + """Return the first adjacent preview image available for a model file.""" model_base_path, _model_ext = os.path.splitext(dest_path) for extension in MODEL_PREVIEW_EXTENSIONS: candidate = f"{model_base_path}{extension}" @@ -917,7 +918,7 @@ def create_model_preview( metadata: Optional[Dict[str, Any]] = None, ) -> str: """Create a compact LoRA Manager-compatible JPEG preview when one is missing.""" - existing_path = _existing_model_preview_path(dest_path) + existing_path = get_existing_model_preview_path(dest_path) if existing_path: return existing_path diff --git a/tests/test_downloader_metadata.py b/tests/test_downloader_metadata.py index 9afb5f0..9ef88b7 100644 --- a/tests/test_downloader_metadata.py +++ b/tests/test_downloader_metadata.py @@ -269,6 +269,19 @@ class DownloaderMetadataSidecarTests(unittest.TestCase): "https://image.civitai.com/example.jpeg" ) + def test_existing_model_preview_path_finds_adjacent_preview(self): + with tempfile.TemporaryDirectory() as tmpdir: + model_path = os.path.join(tmpdir, "model.safetensors") + preview_path = os.path.join(tmpdir, "model.jpeg") + with open(model_path, "wb") as model_file: + model_file.write(b"model") + with open(preview_path, "wb") as preview_file: + preview_file.write(b"preview") + + result = downloader.get_existing_model_preview_path(model_path) + + self.assertEqual(preview_path, result) + def test_existing_file_with_same_hash_is_marked_already_downloaded(self): content = b"existing model" expected_sha256 = hashlib.sha256(content).hexdigest() diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index ac806cb..f4a3460 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -64,6 +64,40 @@ const renderFormatMethodsSource = fs.readFileSync( 'utf8' ); +test('local and loaded model tooltips include preview image routes above full names', () => { + const getModelPreviewTooltipAttrs = eval(`(${extractMethod(renderFormatMethodsSource, 'getModelPreviewTooltipAttrs')})`); + const previousApi = globalThis.api; + globalThis.api = { + apiURL(route) { + return `/comfy${route}`; + }, + }; + + try { + const attrs = getModelPreviewTooltipAttrs.call({ + escapeHtml(value) { + return String(value).replace(/&/g, '&').replace(/"/g, '"'); + }, + }, { + resolved_path: 'E:\\AI Models\\KREA2\\model.safetensors', + }, 'KREA2\\model.safetensors'); + + assert.match(attrs, /data-tooltip="KREA2\\model\.safetensors"/); + assert.match( + attrs, + /data-tooltip-image="\/comfy\/model_resolver\/model-preview\?path=E%3A%5CAI%20Models%5CKREA2%5Cmodel\.safetensors"/ + ); + assert.match(searchPanelMethodsSource, /mr-match-filename[^]*getModelPreviewTooltipAttrs/); + assert.match(tabsLoadedMethodsSource, /mr-model-chip[^]*getModelPreviewTooltipAttrs/); + } finally { + if (previousApi === undefined) { + delete globalThis.api; + } else { + globalThis.api = previousApi; + } + } +}); + test('CivitAI Comfy workflow metadata is extracted as pasteable JSON', () => { const workflow = { last_node_id: 2, diff --git a/tests/test_search_orchestration.py b/tests/test_search_orchestration.py index e5e205f..9c397f1 100644 --- a/tests/test_search_orchestration.py +++ b/tests/test_search_orchestration.py @@ -21,6 +21,7 @@ search_handler = None civitai_search_handler = None local_matches_by_hash_handler = None custom_url_handler = None +model_preview_handler = None def post_decorator(path): def decorator(func): @@ -38,6 +39,16 @@ def post_decorator(path): mock_routes.post = post_decorator +def get_decorator(path): + def decorator(func): + global model_preview_handler + if path == "/model_resolver/model-preview": + model_preview_handler = func + return func + return decorator + +mock_routes.get = get_decorator + # Make sure the parent package directory is in sys.path parent_dir = r"e:\AI\AI\ComfyUI\ComfyUI_Windows_portable\ComfyUI\custom_nodes" if parent_dir not in sys.path: @@ -305,6 +316,36 @@ class SearchOrchestrationTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(body["local_hash_matches"][0]["hash_lookup_source"], "civitai") self.assertEqual(body["local_hash_matches"][0]["hash_lookup_sha256"], file_hash) + async def test_model_preview_route_serves_only_adjacent_preview(self): + with tempfile.TemporaryDirectory() as temp_dir: + model_path = os.path.join(temp_dir, "model.safetensors") + preview_path = os.path.join(temp_dir, "model.jpeg") + with open(model_path, "wb") as model_file: + model_file.write(b"model") + with open(preview_path, "wb") as preview_file: + preview_file.write(b"preview") + + mock_request = MagicMock() + mock_request.query = {"path": model_path} + + with ( + patch( + "folder_paths.folder_names_and_paths", + {"loras": ([temp_dir], {".safetensors"})}, + create=True, + ), + patch( + "folder_paths.get_folder_paths", + return_value=[temp_dir], + create=True, + ), + ): + response = await model_preview_handler(mock_request) + + self.assertIsInstance(response, web.FileResponse) + self.assertEqual(os.path.realpath(preview_path), os.path.realpath(response._path)) + self.assertEqual("private, max-age=300", response.headers["Cache-Control"]) + async def test_search_collects_local_hash_matches_for_any_remote_result_with_hash(self): file_hash = "c" * 64 with tempfile.TemporaryDirectory() as temp_dir: diff --git a/web/css/resolver-main.css b/web/css/resolver-main.css index a6e939a..a5f529d 100644 --- a/web/css/resolver-main.css +++ b/web/css/resolver-main.css @@ -7563,6 +7563,27 @@ a.mr-image-preview-resource-name:hover { z-index: var(--mr-z-tooltip); } +.mr-global-tooltip.mr-global-tooltip-with-image { + width: min(300px, calc(100vw - 24px)); + padding: 8px; + white-space: normal; +} + +.mr-tooltip-preview { + display: block; + width: 100%; + max-height: min(360px, 55vh); + object-fit: contain; + border-radius: 7px; + background: rgba(0, 0, 0, 0.28); +} + +.mr-tooltip-label { + padding: 8px 3px 2px; + white-space: pre-line; + overflow-wrap: anywhere; +} + .mr-options-input-row { display: grid; grid-template-columns: minmax(150px, 1fr) minmax(0, 380px) 42px; diff --git a/web/resolver/search/search_panel.js b/web/resolver/search/search_panel.js index f76c96a..ac49b0a 100644 --- a/web/resolver/search/search_panel.js +++ b/web/resolver/search/search_panel.js @@ -794,9 +794,44 @@ export const searchPanelMethods = { const text = target.getAttribute('data-tooltip'); if (!text) return; - this.tooltipElement.textContent = text; - this.tooltipElement.style.display = 'block'; + const imageUrl = target.getAttribute('data-tooltip-image'); + this._tooltipTarget = target; + this.tooltipElement.replaceChildren(); + this.tooltipElement.classList.toggle('mr-global-tooltip-with-image', Boolean(imageUrl)); + if (imageUrl) { + const image = document.createElement('img'); + image.className = 'mr-tooltip-preview'; + image.alt = ''; + image.decoding = 'async'; + + const label = document.createElement('div'); + label.className = 'mr-tooltip-label'; + label.textContent = text; + + image.addEventListener('load', () => { + if (this._tooltipTarget === target) { + this.positionTooltip(target); + } + }); + image.addEventListener('error', () => { + if (this._tooltipTarget !== target) return; + image.remove(); + this.tooltipElement.classList.remove('mr-global-tooltip-with-image'); + this.positionTooltip(target); + }); + image.src = imageUrl; + this.tooltipElement.append(image, label); + } else { + this.tooltipElement.textContent = text; + } + this.tooltipElement.style.display = 'block'; + this.positionTooltip(target); + this.tooltipElement.setAttribute('data-visible', 'true'); + }, + + positionTooltip(target) { + if (!target || !this.tooltipElement || this._tooltipTarget !== target) return; const rect = target.getBoundingClientRect(); const tooltipRect = this.tooltipElement.getBoundingClientRect(); const margin = 12; @@ -812,12 +847,14 @@ export const searchPanelMethods = { this.tooltipElement.style.left = `${Math.round(left)}px`; this.tooltipElement.style.top = `${Math.round(top)}px`; - this.tooltipElement.setAttribute('data-visible', 'true'); }, hideTooltip() { if (!this.tooltipElement) return; + this._tooltipTarget = null; this.tooltipElement.style.display = 'none'; + this.tooltipElement.classList.remove('mr-global-tooltip-with-image'); + this.tooltipElement.replaceChildren(); this.tooltipElement.removeAttribute('data-visible'); }, @@ -2741,7 +2778,7 @@ export const searchPanelMethods = { html += `