From 63ed2b8ca734577de22cb415c0cd9b9988f939db Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sun, 26 Jul 2026 23:39:09 +0200 Subject: [PATCH] feat: support video previews in model downloads, tooltips, and info dialogs --- __init__.py | 16 +- core/downloader.py | 233 ++++++++++++++++---- core/sources/civitai.py | 7 + core/type_utils.py | 6 + tests/test_downloader_metadata.py | 60 +++++ tests/test_downloads_tab_workflow_route.mjs | 93 ++++++++ tests/test_search_orchestration.py | 9 +- tests/test_unified_helpers.py | 22 ++ web/css/resolver-main.css | 21 +- web/resolver/search/search_panel.js | 67 +++++- web/resolver/views/model_info_methods.js | 91 ++++++-- 11 files changed, 531 insertions(+), 94 deletions(-) diff --git a/__init__.py b/__init__.py index d37d0b4..8144fd3 100644 --- a/__init__.py +++ b/__init__.py @@ -967,7 +967,7 @@ class ModelResolverExtension: return web.FileResponse( preview_path, - headers={"Cache-Control": "private, max-age=300"}, + headers={"Cache-Control": "private, no-cache"}, ) @routes.post("/model_resolver/workflow-model-hashes") @@ -2665,17 +2665,9 @@ class ModelResolverExtension: ): preview_url = result.get("preview_url") or "" if not preview_url: - model_base_path, _model_ext = _os.path.splitext(file_path) - for preview_ext in ( - ".preview.png", - ".preview.jpg", - ".preview.jpeg", - ".preview.webp", - ): - preview_path = f"{model_base_path}{preview_ext}" - if _os.path.exists(preview_path): - preview_url = preview_path.replace("\\", "/") - break + preview_path = get_existing_model_preview_path(file_path) + if preview_path: + preview_url = preview_path.replace("\\", "/") metadata_payload = { "source": "metadata_import", "details_source": result.get("source") or "metadata", diff --git a/core/downloader.py b/core/downloader.py index db5bcbb..715d309 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -16,7 +16,7 @@ import time from collections import deque from io import BytesIO from pathlib import Path -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, Tuple from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import requests @@ -89,7 +89,12 @@ MODEL_PREVIEW_WIDTH = 480 MODEL_PREVIEW_MAX_HEIGHT = 4096 MODEL_PREVIEW_QUALITY = 85 MODEL_PREVIEW_MAX_DOWNLOAD_BYTES = 32 * 1024 * 1024 +MODEL_PREVIEW_VIDEO_MAX_DOWNLOAD_BYTES = 128 * 1024 * 1024 MODEL_PREVIEW_EXTENSIONS = ( + ".mp4", + ".webm", + ".preview.mp4", + ".preview.webm", ".webp", ".jpeg", ".jpg", @@ -720,7 +725,7 @@ def build_lora_manager_metadata( def get_existing_model_preview_path(dest_path: str) -> str: - """Return the first adjacent preview image available for a model file.""" + """Return the first adjacent preview media file available for a model.""" model_base_path, _model_ext = os.path.splitext(dest_path) for extension in MODEL_PREVIEW_EXTENSIONS: candidate = f"{model_base_path}{extension}" @@ -729,13 +734,16 @@ def get_existing_model_preview_path(dest_path: str) -> str: return "" -def _first_model_preview_url(metadata: Dict[str, Any]) -> str: - candidates = [ - metadata.get("preview_url"), - metadata.get("previewUrl"), - metadata.get("thumbnail_url"), - metadata.get("thumbnailUrl"), - ] +def _preview_media_type(url: str, declared_type: Any = "") -> str: + media_type = str(declared_type or "").strip().lower() + if media_type == "video": + return "video" + extension = os.path.splitext(urlparse(str(url or "")).path)[1].lower() + return "video" if extension in {".mp4", ".webm"} else "image" + + +def _first_model_preview_asset(metadata: Dict[str, Any]) -> Tuple[str, str]: + candidates = [] image_sources = [ metadata.get("images"), _as_dict(metadata.get("civitai")).get("images"), @@ -749,24 +757,64 @@ def _first_model_preview_url(metadata: Dict[str, Any]) -> str: for image in _as_list(images): if not isinstance(image, dict): continue - if str(image.get("type") or "").strip().lower() == "video": - continue + declared_type = image.get("type") or image.get("mediaType") candidates.extend( - [ + (value, declared_type) + for value in ( image.get("url"), image.get("imageUrl"), image.get("src"), - ] + ) ) + candidates.extend( + [ + (metadata.get("preview_url"), metadata.get("preview_type")), + (metadata.get("previewUrl"), metadata.get("previewType")), + (metadata.get("thumbnail_url"), metadata.get("thumbnail_type")), + (metadata.get("thumbnailUrl"), metadata.get("thumbnailType")), + ] + ) - for value in candidates: + for value, declared_type in candidates: url = str(value or "").strip() if url.startswith(("http://", "https://")): - return url - return "" + return url, _preview_media_type(url, declared_type) + return "", "" -def _download_preview_image(url: str) -> bytes: +def _first_model_preview_url(metadata: Dict[str, Any]) -> str: + """Return the first preview URL, retained for compatibility with callers.""" + return _first_model_preview_asset(metadata)[0] + + +def _rewrite_civitai_preview_url(url: str, media_type: str) -> str: + parsed = urlparse(str(url or "")) + hostname = (parsed.hostname or "").lower() + if not (hostname == "civitai.com" or hostname.endswith(".civitai.com")): + return url + if "/original=true" not in parsed.path: + return url + + transform = ( + "/transcode=true,width=450,optimized=true" + if media_type == "video" + else "/width=450,optimized=true" + ) + return urlunparse(parsed._replace(path=parsed.path.replace("/original=true", transform, 1))) + + +def _download_preview_asset(url: str, media_type: str = "image") -> bytes: + max_bytes = ( + MODEL_PREVIEW_VIDEO_MAX_DOWNLOAD_BYTES + if media_type == "video" + else MODEL_PREVIEW_MAX_DOWNLOAD_BYTES + ) + accept_header = ( + "video/mp4,video/webm,video/*,*/*;q=0.8" + if media_type == "video" + else "image/avif,image/webp,image/*,*/*;q=0.8" + ) + response = None try: response, _final_url, _final_headers = request_public_url( @@ -774,15 +822,15 @@ def _download_preview_image(url: str) -> bytes: url, headers={ "User-Agent": DOWNLOAD_USER_AGENT, - "Accept": "image/avif,image/webp,image/*,*/*;q=0.8", + "Accept": accept_header, }, timeout=20, stream=True, ) response.raise_for_status() content_length = int(response.headers.get("Content-Length") or 0) - if content_length > MODEL_PREVIEW_MAX_DOWNLOAD_BYTES: - raise ValueError("Preview image exceeds the download size limit") + if content_length > max_bytes: + raise ValueError("Preview media exceeds the download size limit") chunks = [] downloaded = 0 @@ -790,18 +838,31 @@ def _download_preview_image(url: str) -> bytes: if not chunk: continue downloaded += len(chunk) - if downloaded > MODEL_PREVIEW_MAX_DOWNLOAD_BYTES: - raise ValueError("Preview image exceeds the download size limit") + if downloaded > max_bytes: + raise ValueError("Preview media exceeds the download size limit") chunks.append(chunk) return b"".join(chunks) except requests.exceptions.SSLError: - return _download_preview_image_with_system_trust(url) + if media_type == "image": + return _download_preview_image_with_system_trust(url) + return _download_preview_asset_with_system_trust( + url, + media_type=media_type, + ) finally: if response is not None: response.close() -def _download_preview_image_with_system_trust(url: str) -> bytes: +def _download_preview_image(url: str) -> bytes: + """Download an image preview, retained for compatibility with callers.""" + return _download_preview_asset(url, media_type="image") + + +def _download_preview_asset_with_system_trust( + url: str, + media_type: str = "image", +) -> bytes: import ssl from urllib.request import ( HTTPRedirectHandler, @@ -810,6 +871,17 @@ def _download_preview_image_with_system_trust(url: str) -> bytes: build_opener, ) + max_bytes = ( + MODEL_PREVIEW_VIDEO_MAX_DOWNLOAD_BYTES + if media_type == "video" + else MODEL_PREVIEW_MAX_DOWNLOAD_BYTES + ) + accept_header = ( + "video/mp4,video/webm,video/*,*/*;q=0.8" + if media_type == "video" + else "image/avif,image/webp,image/*,*/*;q=0.8" + ) + class ValidatedRedirectHandler(HTTPRedirectHandler): max_redirections = 5 @@ -829,7 +901,7 @@ def _download_preview_image_with_system_trust(url: str) -> bytes: validated_url, headers={ "User-Agent": DOWNLOAD_USER_AGENT, - "Accept": "image/avif,image/webp,image/*,*/*;q=0.8", + "Accept": accept_header, }, method="GET", ) @@ -839,8 +911,8 @@ def _download_preview_image_with_system_trust(url: str) -> bytes: ) with opener.open(request, timeout=20) as response: content_length = int(response.headers.get("Content-Length") or 0) - if content_length > MODEL_PREVIEW_MAX_DOWNLOAD_BYTES: - raise ValueError("Preview image exceeds the download size limit") + if content_length > max_bytes: + raise ValueError("Preview media exceeds the download size limit") chunks = [] downloaded = 0 @@ -849,12 +921,44 @@ def _download_preview_image_with_system_trust(url: str) -> bytes: if not chunk: break downloaded += len(chunk) - if downloaded > MODEL_PREVIEW_MAX_DOWNLOAD_BYTES: - raise ValueError("Preview image exceeds the download size limit") + if downloaded > max_bytes: + raise ValueError("Preview media exceeds the download size limit") chunks.append(chunk) return b"".join(chunks) +def _download_preview_image_with_system_trust(url: str) -> bytes: + """Download an image through the Windows trust store fallback.""" + return _download_preview_asset_with_system_trust(url, media_type="image") + + +def _save_preview_video(video_data: bytes, preview_path: str) -> None: + if not video_data: + raise ValueError("Preview video is empty") + + temp_path = "" + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=os.path.dirname(preview_path), + prefix=f".{os.path.basename(preview_path)}.", + suffix=".tmp", + delete=False, + ) as temp_file: + temp_path = temp_file.name + temp_file.write(video_data) + temp_file.flush() + os.fsync(temp_file.fileno()) + os.replace(temp_path, preview_path) + temp_path = "" + finally: + if temp_path and os.path.exists(temp_path): + try: + os.remove(temp_path) + except OSError: + pass + + def _save_optimized_jpeg(image_data: bytes, preview_path: str) -> None: from PIL import Image, ImageOps @@ -917,30 +1021,65 @@ def create_model_preview( dest_path: str, metadata: Optional[Dict[str, Any]] = None, ) -> str: - """Create a compact LoRA Manager-compatible JPEG preview when one is missing.""" + """Create adjacent preview media using LoRA Manager-compatible behavior.""" existing_path = get_existing_model_preview_path(dest_path) - if existing_path: + source = metadata if isinstance(metadata, dict) else {} + preview_url, media_type = _first_model_preview_asset(source) + if not preview_url: return existing_path - source = metadata if isinstance(metadata, dict) else {} - preview_url = _first_model_preview_url(source) - if not preview_url: - return "" + existing_type = _preview_media_type(existing_path) if existing_path else "" + if existing_path and not (media_type == "video" and existing_type != "video"): + return existing_path model_base_path, _model_ext = os.path.splitext(dest_path) + if media_type == "video": + source_extension = os.path.splitext(urlparse(preview_url).path)[1].lower() + preview_extension = source_extension if source_extension in {".mp4", ".webm"} else ".mp4" + preview_path = f"{model_base_path}{preview_extension}" + attempts = [] + for candidate in ( + _rewrite_civitai_preview_url(preview_url, media_type), + preview_url, + ): + if candidate and candidate not in attempts: + attempts.append(candidate) + + for candidate in attempts: + try: + video_data = _download_preview_asset(candidate, media_type="video") + _save_preview_video(video_data, preview_path) + log.info(f"Video preview saved: {preview_path}") + return preview_path + except Exception as exc: + log.warning( + "Could not download video preview " + f"for {get_filename_from_path(dest_path)} from {candidate}: {exc}" + ) + return existing_path + preview_path = f"{model_base_path}.jpeg" - try: - image_data = _download_preview_image(preview_url) - if not image_data: - return "" - _save_optimized_jpeg(image_data, preview_path) - log.info(f"Preview saved: {preview_path}") - return preview_path - except Exception as exc: - log.warning( - f"Could not create preview for {get_filename_from_path(dest_path)}: {exc}" - ) - return "" + attempts = [] + for candidate in ( + _rewrite_civitai_preview_url(preview_url, media_type), + preview_url, + ): + if candidate and candidate not in attempts: + attempts.append(candidate) + for candidate in attempts: + try: + image_data = _download_preview_image(candidate) + if not image_data: + continue + _save_optimized_jpeg(image_data, preview_path) + log.info(f"Preview saved: {preview_path}") + return preview_path + except Exception as exc: + log.warning( + "Could not create image preview " + f"for {get_filename_from_path(dest_path)} from {candidate}: {exc}" + ) + return existing_path def write_lora_manager_metadata( diff --git a/core/sources/civitai.py b/core/sources/civitai.py index ea23126..734c3b4 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -1968,6 +1968,13 @@ def _metadata_to_model_info(metadata: Dict[str, Any]) -> Dict[str, Any]: "tags": as_list(metadata.get("tags") or model_info.get("tags")), "trained_words": trained_words, "images": images, + "preview_url": first_non_empty( + metadata.get("preview_url"), + metadata.get("previewUrl"), + metadata.get("thumbnail_url"), + metadata.get("thumbnailUrl"), + ) + or "", "clip_skip": civitai_data.get("clipSkip"), "description": first_non_empty( metadata.get("modelDescription"), diff --git a/core/type_utils.py b/core/type_utils.py index b94fce9..4a13f20 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -813,6 +813,12 @@ def normalize_model_image(image_data: Dict[str, Any], default_civitai_url: str = return { "url": url, + "type": ( + image_data.get("type") + or image_data.get("mediaType") + or meta.get("type") + or "" + ), "civitaiUrl": civitai_url, "seed": image_data.get("seed") or meta.get("seed"), "steps": image_data.get("steps") or meta.get("steps"), diff --git a/tests/test_downloader_metadata.py b/tests/test_downloader_metadata.py index 9ef88b7..4bf009b 100644 --- a/tests/test_downloader_metadata.py +++ b/tests/test_downloader_metadata.py @@ -269,6 +269,66 @@ class DownloaderMetadataSidecarTests(unittest.TestCase): "https://image.civitai.com/example.jpeg" ) + def test_writes_optimized_civitai_video_preview_and_prefers_it_over_jpeg(self): + video_data = b"\x00\x00\x00\x18ftypmp42preview-video" + response = MagicMock() + response.headers = {"Content-Length": str(len(video_data))} + response.iter_content.return_value = [video_data] + + with tempfile.TemporaryDirectory() as tmpdir: + model_path = os.path.join(tmpdir, "video-model.safetensors") + old_image_path = os.path.join(tmpdir, "video-model.jpeg") + with open(model_path, "wb") as handle: + handle.write(b"model") + with open(old_image_path, "wb") as handle: + handle.write(b"old-image") + + with patch( + "core.downloader.request_public_url", + return_value=( + response, + "https://image.civitai.com/example.mp4", + {}, + ), + ) as request_preview: + metadata_path = write_lora_manager_metadata( + model_path, + { + "images": [ + { + "url": ( + "https://image.civitai.com/x/" + "original=true/12345.mp4" + ), + "type": "video", + } + ], + }, + category="loras", + create_preview=True, + ) + + preview_path = os.path.join(tmpdir, "video-model.mp4") + self.assertTrue(os.path.isfile(preview_path)) + with open(preview_path, "rb") as handle: + self.assertEqual(video_data, handle.read()) + with open(metadata_path, "r", encoding="utf-8") as handle: + payload = json.load(handle) + + self.assertEqual( + preview_path, + downloader.get_existing_model_preview_path(model_path), + ) + + requested_url = request_preview.call_args.args[1] + self.assertIn( + "/transcode=true,width=450,optimized=true/", + requested_url, + ) + self.assertEqual(preview_path.replace(os.sep, "/"), payload["preview_url"]) + response.raise_for_status.assert_called_once_with() + response.close.assert_called_once_with() + def test_existing_model_preview_path_finds_adjacent_preview(self): with tempfile.TemporaryDirectory() as tmpdir: model_path = os.path.join(tmpdir, "model.safetensors") diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index fb77bcf..66fb320 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -121,6 +121,99 @@ test('model tooltip stays hidden while the context menu is visible', () => { ); }); +test('model tooltip revalidates and detects video previews from the preview route', async () => { + const getTooltipPreviewMediaType = eval(`(${extractMethod(searchPanelMethodsSource, 'getTooltipPreviewMediaType')})`); + const previousFetch = globalThis.fetch; + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return { + ok: true, + headers: { + get(name) { + return name.toLowerCase() === 'content-type' ? 'video/mp4' : ''; + }, + }, + }; + }; + + try { + const dialog = {}; + const previewUrl = '/model_resolver/model-preview?path=model.safetensors'; + assert.equal( + await getTooltipPreviewMediaType.call(dialog, previewUrl), + 'video' + ); + assert.equal( + await getTooltipPreviewMediaType.call(dialog, previewUrl), + 'video' + ); + assert.equal(fetchCalls, 2); + } finally { + globalThis.fetch = previousFetch; + } +}); + +test('Show info uses the local preview route and renders video media', () => { + const isInfoPreviewVideo = eval(`(${extractMethod(modelInfoMethodsSource, 'isInfoPreviewVideo')})`); + const getInfoDialogMedia = eval(`(${extractMethod(modelInfoMethodsSource, 'getInfoDialogMedia')})`); + const renderInfoPreviewMedia = eval(`(${extractMethod(modelInfoMethodsSource, 'renderInfoPreviewMedia')})`); + const previousApi = globalThis.api; + globalThis.api = { + apiURL(route) { + return `/comfy${route}`; + }, + }; + + try { + const dialog = { + isInfoPreviewVideo, + getModelInfoResolvedPath(data) { + return data.resolved_path || ''; + }, + escapeHtml(value) { + return String(value).replace(/&/g, '&').replace(/"/g, '"'); + }, + }; + const media = getInfoDialogMedia.call(dialog, { + resolved_path: 'E:\\Models\\video-model.safetensors', + preview_url: 'E:/Models/video-model.mp4', + images: [{ + url: 'https://image.civitai.com/original=true/12345.mp4', + type: 'video', + seed: 123, + }], + }); + + assert.equal(media.length, 1); + assert.equal(media[0].type, 'video'); + assert.equal(media[0].seed, 123); + assert.match( + media[0].url, + /^\/comfy\/model_resolver\/model-preview\?path=E%3A%5CModels%5Cvideo-model\.safetensors$/ + ); + assert.equal(isInfoPreviewVideo({ url: media[0].url, type: 'video' }), true); + + const thumbnail = renderInfoPreviewMedia.call(dialog, media[0], { thumbnail: true }); + const fullPreview = renderInfoPreviewMedia.call(dialog, media[0]); + assert.match(thumbnail, /^