feat(civitai): implement URL rewriting for Civitai previews and enhance download handling, fixes #499

tests/services/test_download_manager.py

@@ -394,3 +394,98 @@ async def test_execute_download_retries_urls(monkeypatch, tmp_path):
     assert result == {"success": True}
     assert [url for url, *_ in dummy_downloader.calls] == download_urls
     assert dummy_scanner.calls  # ensure cache updated
+
+
+@pytest.mark.asyncio
+async def test_execute_download_uses_rewritten_civitai_preview(monkeypatch, tmp_path):
+    manager = DownloadManager()
+    save_dir = tmp_path / "downloads"
+    save_dir.mkdir()
+    target_path = save_dir / "file.safetensors"
+
+    manager._active_downloads["dl"] = {}
+
+    class DummyMetadata:
+        def __init__(self, path: Path):
+            self.file_path = str(path)
+            self.sha256 = "sha256"
+            self.file_name = path.stem
+            self.preview_url = None
+            self.preview_nsfw_level = None
+
+        def generate_unique_filename(self, *_args, **_kwargs):
+            return os.path.basename(self.file_path)
+
+        def update_file_info(self, _path):
+            return None
+
+        def to_dict(self):
+            return {"file_path": self.file_path}
+
+    metadata = DummyMetadata(target_path)
+    version_info = {
+        "images": [
+            {
+                "url": "https://image.civitai.com/container/example/original=true/sample.jpeg",
+                "type": "image",
+                "nsfwLevel": 2,
+            }
+        ]
+    }
+    download_urls = ["https://example.invalid/file.safetensors"]
+
+    class DummyDownloader:
+        def __init__(self):
+            self.file_calls: list[tuple[str, str]] = []
+            self.memory_calls = 0
+
+        async def download_file(self, url, path, progress_callback=None, use_auth=None):
+            self.file_calls.append((url, path))
+            if url.endswith(".jpeg"):
+                Path(path).write_bytes(b"preview")
+                return True, None
+            if url.endswith(".safetensors"):
+                Path(path).write_bytes(b"model")
+                return True, None
+            return False, "unexpected url"
+
+        async def download_to_memory(self, *_args, **_kwargs):
+            self.memory_calls += 1
+            return False, b"", {}
+
+    dummy_downloader = DummyDownloader()
+    monkeypatch.setattr(download_manager, "get_downloader", AsyncMock(return_value=dummy_downloader))
+
+    optimize_called = {"value": False}
+
+    def fake_optimize_image(**_kwargs):
+        optimize_called["value"] = True
+        return b"", {}
+
+    monkeypatch.setattr(download_manager.ExifUtils, "optimize_image", staticmethod(fake_optimize_image))
+    monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
+
+    dummy_scanner = SimpleNamespace(add_model_to_cache=AsyncMock(return_value=None))
+    monkeypatch.setattr(DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner))
+
+    result = await manager._execute_download(
+        download_urls=download_urls,
+        save_dir=str(save_dir),
+        metadata=metadata,
+        version_info=version_info,
+        relative_path="",
+        progress_callback=None,
+        model_type="lora",
+        download_id="dl",
+    )
+
+    assert result == {"success": True}
+    preview_urls = [url for url, _ in dummy_downloader.file_calls if url.endswith(".jpeg")]
+    assert any("width=450,optimized=true" in url for url in preview_urls)
+    assert dummy_downloader.memory_calls == 0
+    assert optimize_called["value"] is False
+    assert metadata.preview_url.endswith(".jpeg")
+    assert metadata.preview_nsfw_level == 2
+    stored_preview = manager._active_downloads["dl"]["preview_path"]
+    assert stored_preview.endswith(".jpeg")
+    assert Path(stored_preview).exists()

tests/services/test_preview_asset_service.py

@@ -0,0 +1,182 @@
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from py.services.preview_asset_service import PreviewAssetService
+
+
+class StubMetadataManager:
+    async def save_metadata(self, *_args: Any, **_kwargs: Any) -> bool:  # pragma: no cover - helper
+        return True
+
+
+class RecordingExifUtils:
+    def __init__(self) -> None:
+        self.called = False
+
+    def optimize_image(self, **kwargs):
+        self.called = True
+        return kwargs["image_data"], {}
+
+
+@pytest.mark.asyncio
+async def test_ensure_preview_prefers_rewritten_civitai_image(tmp_path):
+    metadata_path = tmp_path / "model.metadata.json"
+    metadata_path.write_text("{}")
+    local_metadata: dict[str, Any] = {}
+
+    class Downloader:
+        def __init__(self):
+            self.file_calls: list[tuple[str, str]] = []
+            self.memory_calls = 0
+
+        async def download_file(self, url, path, use_auth=False):
+            self.file_calls.append((url, path))
+            if "width=450,optimized=true" in url:
+                Path(path).write_bytes(b"image-data")
+                return True, None
+            return False, "fail"
+
+        async def download_to_memory(self, *_args, **_kwargs):
+            self.memory_calls += 1
+            return False, b"", {}
+
+    downloader = Downloader()
+
+    async def downloader_factory():
+        return downloader
+
+    exif_utils = RecordingExifUtils()
+    service = PreviewAssetService(
+        metadata_manager=StubMetadataManager(),
+        downloader_factory=downloader_factory,
+        exif_utils=exif_utils,
+    )
+
+    images = [
+        {
+            "url": "https://image.civitai.com/container/example/original=true/sample.jpeg",
+            "type": "image",
+            "nsfwLevel": 3,
+        }
+    ]
+
+    await service.ensure_preview_for_metadata(str(metadata_path), local_metadata, images)
+
+    assert downloader.memory_calls == 0
+    assert exif_utils.called is False
+    assert len(downloader.file_calls) == 1
+    assert "width=450,optimized=true" in downloader.file_calls[0][0]
+    preview_path = Path(local_metadata["preview_url"])
+    assert preview_path.exists()
+    assert preview_path.suffix == ".jpeg"
+    assert local_metadata["preview_nsfw_level"] == 3
+
+
+@pytest.mark.asyncio
+async def test_ensure_preview_falls_back_to_webp_when_rewrite_fails(tmp_path):
+    metadata_path = tmp_path / "model.metadata.json"
+    metadata_path.write_text("{}")
+    local_metadata: dict[str, Any] = {}
+
+    class Downloader:
+        def __init__(self):
+            self.file_calls: list[tuple[str, str]] = []
+            self.memory_calls = 0
+
+        async def download_file(self, url, path, use_auth=False):
+            self.file_calls.append((url, path))
+            return False, "fail"
+
+        async def download_to_memory(self, *_args, **_kwargs):
+            self.memory_calls += 1
+            return True, b"raw-image", {}
+
+    downloader = Downloader()
+
+    async def downloader_factory():
+        return downloader
+
+    class ExifUtils:
+        def __init__(self):
+            self.calls = 0
+
+        def optimize_image(self, **kwargs):
+            self.calls += 1
+            return b"webp-data", {}
+
+    exif_utils = ExifUtils()
+
+    service = PreviewAssetService(
+        metadata_manager=StubMetadataManager(),
+        downloader_factory=downloader_factory,
+        exif_utils=exif_utils,
+    )
+
+    images = [
+        {
+            "url": "https://image.civitai.com/container/example/original=true/sample.png",
+            "type": "image",
+            "nsfwLevel": 1,
+        }
+    ]
+
+    await service.ensure_preview_for_metadata(str(metadata_path), local_metadata, images)
+
+    assert downloader.memory_calls == 1
+    assert exif_utils.calls == 1
+    preview_path = Path(local_metadata["preview_url"])
+    assert preview_path.exists()
+    assert preview_path.suffix == ".webp"
+
+
+@pytest.mark.asyncio
+async def test_ensure_preview_rewrites_civitai_video(tmp_path):
+    metadata_path = tmp_path / "model.metadata.json"
+    metadata_path.write_text("{}")
+    local_metadata: dict[str, Any] = {}
+
+    class Downloader:
+        def __init__(self):
+            self.file_calls: list[tuple[str, str]] = []
+
+        async def download_file(self, url, path, use_auth=False):
+            self.file_calls.append((url, path))
+            if "transcode=true,width=450,optimized=true" in url:
+                Path(path).write_bytes(b"video-data")
+                return True, None
+            if url.endswith(".mp4"):
+                return False, "fail"
+            return False, "unexpected"
+
+        async def download_to_memory(self, *_args, **_kwargs):
+            pytest.fail("download_to_memory should not be used for video previews")
+
+    downloader = Downloader()
+
+    async def downloader_factory():
+        return downloader
+
+    service = PreviewAssetService(
+        metadata_manager=StubMetadataManager(),
+        downloader_factory=downloader_factory,
+        exif_utils=RecordingExifUtils(),
+    )
+
+    images = [
+        {
+            "url": "https://image.civitai.com/container/example/original=true/sample.mp4",
+            "type": "video",
+            "nsfwLevel": 2,
+        }
+    ]
+
+    await service.ensure_preview_for_metadata(str(metadata_path), local_metadata, images)
+
+    assert len(downloader.file_calls) >= 1
+    assert any("transcode=true,width=450,optimized=true" in url for url, _ in downloader.file_calls)
+    preview_path = Path(local_metadata["preview_url"])
+    assert preview_path.exists()
+    assert preview_path.suffix == ".mp4"
+    assert local_metadata["preview_nsfw_level"] == 2