diff --git a/__init__.py b/__init__.py index f280bbf..2cb60e8 100644 --- a/__init__.py +++ b/__init__.py @@ -2587,6 +2587,7 @@ class ModelResolverExtension: or result.get("platform_url") or result.get("download_url") or "", + create_preview=True, ) or "" metadata_saved = bool(metadata_path) if metadata_path: @@ -2688,6 +2689,7 @@ class ModelResolverExtension: or result.get("url") or result.get("download_url") or "", + create_preview=True, ) or "" metadata_saved = bool(metadata_path) if metadata_path: diff --git a/core/downloader.py b/core/downloader.py index f4453b2..f541208 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -10,9 +10,11 @@ import secrets import shutil import socket import subprocess +import tempfile import threading import time from collections import deque +from io import BytesIO from pathlib import Path from typing import Any, Callable, Dict, List, Optional from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse @@ -83,6 +85,20 @@ ARIA2_STATUS_RPC_RETRIES = 4 ARIA2_STATUS_RPC_RETRY_DELAY = 0.15 ARIA2_IDLE_STOP_SECONDS = 5 * 60 DOWNLOAD_USER_AGENT = DEFAULT_BROWSER_USER_AGENT +MODEL_PREVIEW_WIDTH = 480 +MODEL_PREVIEW_MAX_HEIGHT = 4096 +MODEL_PREVIEW_QUALITY = 85 +MODEL_PREVIEW_MAX_DOWNLOAD_BYTES = 32 * 1024 * 1024 +MODEL_PREVIEW_EXTENSIONS = ( + ".webp", + ".jpeg", + ".jpg", + ".png", + ".preview.webp", + ".preview.jpeg", + ".preview.jpg", + ".preview.png", +) MANAGED_ARIA2_ROOT = Path(__file__).resolve().parents[1] / "tools" / "aria2" HF_XET_ARIA2_AUTH_HOSTS = { "cas-bridge.xethub.hf.co", @@ -703,17 +719,249 @@ def build_lora_manager_metadata( return payload +def _existing_model_preview_path(dest_path: str) -> str: + model_base_path, _model_ext = os.path.splitext(dest_path) + for extension in MODEL_PREVIEW_EXTENSIONS: + candidate = f"{model_base_path}{extension}" + if os.path.isfile(candidate): + return candidate + 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"), + ] + image_sources = [ + metadata.get("images"), + _as_dict(metadata.get("civitai")).get("images"), + _as_dict(metadata.get("selected_version")).get("images"), + _as_dict(metadata.get("civitai_details")).get("images"), + _as_dict(_as_dict(metadata.get("civitai_details")).get("civitai")).get( + "images" + ), + ] + for images in image_sources: + for image in _as_list(images): + if not isinstance(image, dict): + continue + if str(image.get("type") or "").strip().lower() == "video": + continue + candidates.extend( + [ + image.get("url"), + image.get("imageUrl"), + image.get("src"), + ] + ) + + for value in candidates: + url = str(value or "").strip() + if url.startswith(("http://", "https://")): + return url + return "" + + +def _download_preview_image(url: str) -> bytes: + response = None + try: + response, _final_url, _final_headers = request_public_url( + "GET", + url, + headers={ + "User-Agent": DOWNLOAD_USER_AGENT, + "Accept": "image/avif,image/webp,image/*,*/*;q=0.8", + }, + 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") + + chunks = [] + downloaded = 0 + for chunk in response.iter_content(chunk_size=64 * 1024): + if not chunk: + continue + downloaded += len(chunk) + if downloaded > MODEL_PREVIEW_MAX_DOWNLOAD_BYTES: + raise ValueError("Preview image exceeds the download size limit") + chunks.append(chunk) + return b"".join(chunks) + except requests.exceptions.SSLError: + return _download_preview_image_with_system_trust(url) + finally: + if response is not None: + response.close() + + +def _download_preview_image_with_system_trust(url: str) -> bytes: + import ssl + from urllib.request import ( + HTTPRedirectHandler, + HTTPSHandler, + Request, + build_opener, + ) + + class ValidatedRedirectHandler(HTTPRedirectHandler): + max_redirections = 5 + + def redirect_request(self, req, fp, code, msg, headers, newurl): + validated_url = validate_public_http_url(newurl) + return super().redirect_request( + req, + fp, + code, + msg, + headers, + validated_url, + ) + + validated_url = validate_public_http_url(url) + request = Request( + validated_url, + headers={ + "User-Agent": DOWNLOAD_USER_AGENT, + "Accept": "image/avif,image/webp,image/*,*/*;q=0.8", + }, + method="GET", + ) + opener = build_opener( + HTTPSHandler(context=ssl.create_default_context()), + ValidatedRedirectHandler(), + ) + 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") + + chunks = [] + downloaded = 0 + while True: + chunk = response.read(64 * 1024) + if not chunk: + break + downloaded += len(chunk) + if downloaded > MODEL_PREVIEW_MAX_DOWNLOAD_BYTES: + raise ValueError("Preview image exceeds the download size limit") + chunks.append(chunk) + return b"".join(chunks) + + +def _save_optimized_jpeg(image_data: bytes, preview_path: str) -> None: + from PIL import Image, ImageOps + + temp_path = "" + try: + with Image.open(BytesIO(image_data)) as source_image: + source_image.seek(0) + oriented_image = ImageOps.exif_transpose(source_image) + width, height = oriented_image.size + if width <= 0 or height <= 0: + raise ValueError("Preview image has invalid dimensions") + + resize_scale = min( + MODEL_PREVIEW_WIDTH / width, + MODEL_PREVIEW_MAX_HEIGHT / height, + ) + target_width = max(1, round(width * resize_scale)) + target_height = max(1, round(height * resize_scale)) + resampling = getattr(Image, "Resampling", Image) + resized_image = oriented_image.resize( + (target_width, target_height), + resampling.LANCZOS, + ) + + if resized_image.mode in {"RGBA", "LA"} or ( + resized_image.mode == "P" and "transparency" in resized_image.info + ): + rgba_image = resized_image.convert("RGBA") + output_image = Image.new("RGB", rgba_image.size, "white") + output_image.paste(rgba_image, mask=rgba_image.getchannel("A")) + else: + output_image = resized_image.convert("RGB") + + 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 + output_image.save( + temp_file, + format="JPEG", + quality=MODEL_PREVIEW_QUALITY, + optimize=True, + progressive=True, + ) + 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 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.""" + existing_path = _existing_model_preview_path(dest_path) + if existing_path: + return existing_path + + source = metadata if isinstance(metadata, dict) else {} + preview_url = _first_model_preview_url(source) + if not preview_url: + return "" + + model_base_path, _model_ext = os.path.splitext(dest_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 "" + + def write_lora_manager_metadata( dest_path: str, metadata: Optional[Dict[str, Any]] = None, category: str = "", source_url: str = "", + create_preview: bool = False, ) -> Optional[str]: """Write the LoRA Manager-compatible sidecar metadata next to a model file.""" metadata_path = get_metadata_sidecar_path(dest_path) try: payload = build_lora_manager_metadata(dest_path, metadata, category, source_url) + if create_preview: + preview_source = { + **payload, + **(metadata if isinstance(metadata, dict) else {}), + } + preview_path = create_model_preview(dest_path, preview_source) + if preview_path: + payload["preview_url"] = _normalise_metadata_file_path(preview_path) write_json_atomic(metadata_path, payload, indent=2) log.info(f"Metadata saved: {metadata_path}") return metadata_path @@ -1635,6 +1883,7 @@ def _download_huggingface_xet( metadata or {}, category, validated_url, + create_preview=True, ) with download_lock: state = download_progress.get(download_id) @@ -1845,6 +2094,7 @@ def download_file_with_aria2( metadata or {}, category, url, + create_preview=True, ) with download_lock: download_progress[download_id].update( @@ -2159,6 +2409,7 @@ def download_file( metadata or {}, category, url, + create_preview=True, ) if metadata_path: result["metadata_path"] = metadata_path @@ -2354,6 +2605,7 @@ def download_model( metadata or {}, category, url, + create_preview=True, ) or "" size = os.path.getsize(dest_path) with download_lock: diff --git a/tests/test_downloader_metadata.py b/tests/test_downloader_metadata.py index c4b7bd6..9afb5f0 100644 --- a/tests/test_downloader_metadata.py +++ b/tests/test_downloader_metadata.py @@ -1,9 +1,12 @@ import json import hashlib +from io import BytesIO import os import tempfile import unittest -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +from PIL import Image from core import downloader from core.downloader import ( @@ -192,6 +195,80 @@ class DownloaderMetadataSidecarTests(unittest.TestCase): self.assertEqual(10, civitai["stats"]["downloadCount"]) self.assertNotIn("token=", civitai["downloadUrl"]) + def test_writes_optimized_jpeg_preview_from_first_model_image(self): + source_image = Image.new("RGBA", (960, 1440), (20, 40, 60, 128)) + source_buffer = BytesIO() + source_image.save(source_buffer, format="PNG") + response = MagicMock() + response.headers = {"Content-Length": str(len(source_buffer.getvalue()))} + response.iter_content.return_value = [source_buffer.getvalue()] + + with tempfile.TemporaryDirectory() as tmpdir: + model_path = os.path.join(tmpdir, "preview-model.safetensors") + with open(model_path, "wb") as handle: + handle.write(b"abc") + + with patch( + "core.downloader.request_public_url", + return_value=( + response, + "https://image.civitai.com/example.png", + {}, + ), + ): + metadata_path = write_lora_manager_metadata( + model_path, + { + "source": "civitai", + "images": [ + { + "url": "https://image.civitai.com/example.png", + "type": "image", + } + ], + }, + category="loras", + create_preview=True, + ) + + preview_path = os.path.join(tmpdir, "preview-model.jpeg") + self.assertTrue(os.path.isfile(preview_path)) + with Image.open(preview_path) as preview_image: + self.assertEqual("JPEG", preview_image.format) + self.assertEqual((480, 720), preview_image.size) + self.assertEqual("RGB", preview_image.mode) + self.assertFalse(preview_image.getexif()) + with open(metadata_path, "r", encoding="utf-8") as handle: + payload = json.load(handle) + + 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_preview_download_uses_system_trust_after_requests_ssl_error(self): + ssl_error = downloader.requests.exceptions.SSLError("certificate error") + with ( + patch( + "core.downloader.request_public_url", + side_effect=ssl_error, + ), + patch( + "core.downloader._download_preview_image_with_system_trust", + return_value=b"preview-data", + ) as fallback, + ): + result = downloader._download_preview_image( + "https://image.civitai.com/example.jpeg" + ) + + self.assertEqual(b"preview-data", result) + fallback.assert_called_once_with( + "https://image.civitai.com/example.jpeg" + ) + def test_existing_file_with_same_hash_is_marked_already_downloaded(self): content = b"existing model" expected_sha256 = hashlib.sha256(content).hexdigest()