mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-07 14:30:15 -03:00
feat(example-images): add missing-only download path and skip existing files
Split the single-model and bulk context menu actions into 'Download Missing Example Images' (regular endpoint, skips already-processed models) and 'Re-process Example Images' (force endpoint, retries failed models). - start_download accepts model_hashes so a selected subset can be processed with the progress-aware skip logic; explicitly targeted models bypass the failed/processed model-level guards so per-image gaps are filled - pre-download existence check in the processor skips network requests for image files already on disk across all download paths - force download retries previously failed models and clears their failed status on success - add i18n keys for the new menu items across all locales
This commit is contained in:
@@ -63,7 +63,7 @@ async def test_start_download_bootstraps_progress_and_task(
|
||||
release = asyncio.Event()
|
||||
|
||||
async def fake_download(
|
||||
self, output_dir, optimize, model_types, delay, library_name, force=False
|
||||
self, output_dir, optimize, model_types, delay, library_name, force=False, model_hashes=None
|
||||
):
|
||||
started.set()
|
||||
await release.wait()
|
||||
@@ -93,6 +93,44 @@ async def test_start_download_bootstraps_progress_and_task(
|
||||
assert manager._progress["status"] == "completed"
|
||||
|
||||
|
||||
async def test_start_download_forwards_model_hashes(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path
|
||||
) -> None:
|
||||
settings_manager = get_settings_manager()
|
||||
settings_manager.settings["example_images_path"] = str(tmp_path)
|
||||
settings_manager.settings["libraries"] = {"default": {}}
|
||||
settings_manager.settings["active_library"] = "default"
|
||||
|
||||
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
|
||||
|
||||
received: Dict[str, Any] = {}
|
||||
|
||||
async def fake_download(
|
||||
self, output_dir, optimize, model_types, delay, library_name, force=False, model_hashes=None
|
||||
):
|
||||
received["model_hashes"] = model_hashes
|
||||
async with self._state_lock:
|
||||
self._is_downloading = False
|
||||
self._download_task = None
|
||||
self._progress["status"] = "completed"
|
||||
|
||||
monkeypatch.setattr(
|
||||
download_module.DownloadManager,
|
||||
"_download_all_example_images",
|
||||
fake_download,
|
||||
)
|
||||
|
||||
result = await manager.start_download(
|
||||
{"model_types": ["lora"], "delay": 0, "model_hashes": ["abc123", "def456"]}
|
||||
)
|
||||
assert result["success"] is True
|
||||
|
||||
task = manager._download_task
|
||||
assert task is not None
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
assert received["model_hashes"] == ["abc123", "def456"]
|
||||
|
||||
|
||||
async def test_pause_and_resume_flow(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None:
|
||||
settings_manager = get_settings_manager()
|
||||
settings_manager.settings["example_images_path"] = str(tmp_path)
|
||||
|
||||
@@ -100,6 +100,54 @@ def test_get_file_extension_media_type_hint_low_priority() -> None:
|
||||
assert ext == ".mp4"
|
||||
|
||||
|
||||
def test_example_image_file_exists_checks_plausible_extensions(tmp_path) -> None:
|
||||
proc = processor_module.ExampleImagesProcessor
|
||||
assert proc._example_image_file_exists(str(tmp_path), 0) is False
|
||||
Path(tmp_path, "image_0.webp").write_bytes(b"x")
|
||||
assert proc._example_image_file_exists(str(tmp_path), 0) is True
|
||||
assert proc._example_image_file_exists(str(tmp_path), 1) is False
|
||||
|
||||
|
||||
def test_example_image_file_exists_video_hint_only_checks_video_extensions(tmp_path) -> None:
|
||||
proc = processor_module.ExampleImagesProcessor
|
||||
Path(tmp_path, "image_2.jpg").write_bytes(b"x")
|
||||
# An existing image file must not satisfy a video-hinted lookup
|
||||
assert proc._example_image_file_exists(str(tmp_path), 2, "video") is False
|
||||
Path(tmp_path, "image_2.mp4").write_bytes(b"x")
|
||||
assert proc._example_image_file_exists(str(tmp_path), 2, "video") is True
|
||||
|
||||
|
||||
async def test_download_model_images_with_tracking_skips_existing_files(tmp_path) -> None:
|
||||
proc = processor_module.ExampleImagesProcessor
|
||||
images = [
|
||||
{"url": "https://image.civitai.com/a/b", "type": "image"},
|
||||
{"url": "https://image.civitai.com/c/d", "type": "image"},
|
||||
]
|
||||
Path(tmp_path, "image_0.jpg").write_bytes(b"existing")
|
||||
|
||||
class RecordingDownloader:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[str] = []
|
||||
|
||||
async def download_to_memory(self, url, use_auth=False, return_headers=False):
|
||||
self.calls.append(url)
|
||||
return True, b"\xff\xd8\xff" + b"data", {}
|
||||
|
||||
downloader = RecordingDownloader()
|
||||
success, is_stale, failed, rate_limited = await proc.download_model_images_with_tracking(
|
||||
"hash", "model", images, str(tmp_path), False, downloader
|
||||
)
|
||||
|
||||
assert success is True
|
||||
assert is_stale is False
|
||||
assert failed == []
|
||||
assert rate_limited == []
|
||||
# Only the missing image is requested; the existing one is skipped without a network call
|
||||
assert len(downloader.calls) == 1
|
||||
assert "c/d" in downloader.calls[0]
|
||||
assert Path(tmp_path, "image_1.jpg").exists()
|
||||
|
||||
|
||||
class StubScanner:
|
||||
def __init__(self, models: list[Dict[str, Any]]) -> None:
|
||||
self._cache = SimpleNamespace(raw_data=models)
|
||||
|
||||
Reference in New Issue
Block a user