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:
Will Miao
2026-08-03 20:52:46 +08:00
parent 191c4e03cd
commit 9087b4b07c
23 changed files with 20450 additions and 20060 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -772,7 +772,8 @@
"unfavorite": "Remove from Favorites",
"deleteAll": "Delete Selected",
"downloadMissingLoras": "Download Missing LoRAs",
"downloadExamples": "Download Example Images",
"downloadMissingExamples": "Download Missing Example Images",
"reprocessExamples": "Re-process Example Images",
"clear": "Clear Selection",
"skipMetadataRefreshCount": "Skip ({count} models)",
"resumeMetadataRefreshCount": "Resume ({count} models)",
@@ -807,7 +808,8 @@
"sendToWorkflowAppend": "Send to Workflow (Append)",
"sendToWorkflowReplace": "Send to Workflow (Replace)",
"openExamples": "Open Examples Folder",
"downloadExamples": "Download Example Images",
"downloadExamples": "Download Missing Example Images",
"reprocessExamples": "Re-process Example Images",
"replacePreview": "Replace Preview",
"setContentRating": "Set Content Rating",
"moveToFolder": "Move to Folder",

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -172,6 +172,7 @@ class DownloadManager:
model_types = data.get("model_types", ["lora", "checkpoint"])
delay = float(data.get("delay", 0.2))
force = data.get("force", False)
model_hashes = data.get("model_hashes", [])
# Step 2: Validate configuration (fast lookup)
settings_manager = get_settings_manager()
@@ -241,6 +242,7 @@ class DownloadManager:
delay,
active_library,
force,
model_hashes,
)
)
@@ -577,8 +579,9 @@ class DownloadManager:
delay,
library_name,
force: bool = False,
model_hashes: list[str] | None = None,
):
"""Download example images for all models."""
"""Download example images for all models (or only the given hashes)."""
downloader = await get_downloader()
@@ -606,6 +609,18 @@ class DownloadManager:
if model.get("sha256"):
all_models.append((scanner_type, model, scanner))
# Restrict to the requested hashes when provided (empty = all models).
# Explicit targets are a directed user request, so previously failed
# models are retried instead of skipped.
explicit_targets = bool(model_hashes)
if model_hashes:
hash_set = {h.lower() for h in model_hashes}
all_models = [
(scanner_type, model, scanner)
for scanner_type, model, scanner in all_models
if model.get("sha256", "").lower() in hash_set
]
# Update total count
self._progress["total"] = len(all_models)
logger.debug(f"Found {self._progress['total']} models to process")
@@ -629,6 +644,7 @@ class DownloadManager:
downloader,
library_name,
force,
explicit_targets,
)
# Update progress
@@ -725,6 +741,7 @@ class DownloadManager:
downloader,
library_name,
force: bool = False,
explicit_targets: bool = False,
):
"""Process a single model download."""
@@ -747,8 +764,9 @@ class DownloadManager:
self._progress["current_model"] = f"{model_name} ({model_hash[:8]})"
await self._broadcast_progress(status="running")
# Skip if already in failed models (unless force mode is enabled)
if not force and model_hash in self._progress["failed_models"]:
# Skip if already in failed models (unless force mode is enabled or
# the model was explicitly targeted by hash)
if not force and not explicit_targets and model_hash in self._progress["failed_models"]:
logger.debug(f"Skipping known failed model: {model_name}")
return False
@@ -757,30 +775,34 @@ class DownloadManager:
)
existing_files = _model_directory_has_files(model_dir)
# Skip if already processed AND directory exists with files
if model_hash in self._progress["processed_models"]:
if existing_files:
logger.debug(f"Skipping already processed model: {model_name}")
# Model-level guard: a populated folder counts as done. Explicitly
# targeted models bypass it so the per-image existence pre-check can
# fill individual gaps without re-fetching existing files.
if not explicit_targets:
# Skip if already processed AND directory exists with files
if model_hash in self._progress["processed_models"]:
if existing_files:
logger.debug(f"Skipping already processed model: {model_name}")
return False
logger.debug(
"Model %s (%s) marked as processed but folder empty or missing, reprocessing triggered",
model_name,
model_hash,
)
# Track that we are reprocessing this model for summary logging
self._progress["reprocessed_models"].add(model_hash)
# Remove from processed models since we need to reprocess
self._progress["processed_models"].discard(model_hash)
if existing_files and model_hash not in self._progress["processed_models"]:
logger.debug(
"Model folder already populated for %s, marking as processed without download",
model_name,
)
self._progress["processed_models"].add(model_hash)
return False
logger.debug(
"Model %s (%s) marked as processed but folder empty or missing, reprocessing triggered",
model_name,
model_hash,
)
# Track that we are reprocessing this model for summary logging
self._progress["reprocessed_models"].add(model_hash)
# Remove from processed models since we need to reprocess
self._progress["processed_models"].discard(model_hash)
if existing_files and model_hash not in self._progress["processed_models"]:
logger.debug(
"Model folder already populated for %s, marking as processed without download",
model_name,
)
self._progress["processed_models"].add(model_hash)
return False
if not model_dir:
logger.warning(
"Unable to resolve example images folder for model %s (%s)",
@@ -884,7 +906,7 @@ class DownloadManager:
model_name,
)
# Clear failed_models so non-force runs can retry
if force and model_hash in self._progress["failed_models"]:
if (force or explicit_targets) and model_hash in self._progress["failed_models"]:
self._progress["failed_models"].discard(model_hash)
logger.info(
f"Removed {model_name} from failed_models after force retry with rate-limited images"
@@ -904,7 +926,7 @@ class DownloadManager:
)
elif success:
self._progress["processed_models"].add(model_hash)
if force and model_hash in self._progress["failed_models"]:
if (force or explicit_targets) and model_hash in self._progress["failed_models"]:
self._progress["failed_models"].discard(model_hash)
logger.info(
f"Removed {model_name} from failed_models after successful force retry"

View File

@@ -113,6 +113,26 @@ class ExampleImagesProcessor:
message = str(error).lower()
return '404' in message or 'file not found' in message
@staticmethod
def _example_image_file_exists(model_dir: str, index: int, media_type_hint: str | None = None) -> bool:
"""Return True when the file that would be written for a media index already exists.
The final filename (``image_{index}{extension}``) depends on the downloaded
content, so the extension cannot be known ahead of time. The post-download
check skips the write when the exact target file exists; this pre-check
approximates that with the candidate extensions for the media type (videos
only when the metadata hints at a video) so the network request is avoided
for files that already exist on disk.
"""
if media_type_hint == "video":
extensions = SUPPORTED_MEDIA_EXTENSIONS['videos']
else:
extensions = SUPPORTED_MEDIA_EXTENSIONS['images']
return any(
os.path.exists(os.path.join(model_dir, f"image_{index}{ext}"))
for ext in extensions
)
@staticmethod
async def download_model_images(model_hash, model_name, model_images, model_dir, optimize, downloader):
"""Download images for a single model
@@ -139,7 +159,12 @@ class ExampleImagesProcessor:
original_url = image_url
if optimize and 'civitai.com' in image_url:
image_url = ExampleImagesProcessor.get_civitai_optimized_url(image_url)
# Skip the download when the file already exists on disk
if ExampleImagesProcessor._example_image_file_exists(model_dir, i, image.get("type")):
logger.debug("File already exists, skipping download for %s", image_url)
continue
# Download the file first to determine the actual file type
try:
logger.debug(f"Downloading media file {i} for {model_name}")
@@ -229,6 +254,11 @@ class ExampleImagesProcessor:
if optimize and 'civitai.com' in image_url:
image_url = ExampleImagesProcessor.get_civitai_optimized_url(image_url)
# Skip the download when the file already exists on disk
if ExampleImagesProcessor._example_image_file_exists(model_dir, i, image.get("type")):
logger.debug("File already exists, skipping download for %s", image_url)
continue
async def _attempt_download() -> tuple:
logger.debug("Downloading media file %s for %s", i, model_name)
return await downloader.download_to_memory(

View File

@@ -184,7 +184,8 @@ export const DOWNLOAD_ENDPOINTS = {
downloadGet: '/api/lm/download-model-get',
cancelGet: '/api/lm/cancel-download-get',
progress: '/api/lm/download-progress',
exampleImages: '/api/lm/force-download-example-images' // New endpoint for downloading example images
exampleImages: '/api/lm/force-download-example-images', // Re-process example images ignoring previous status
exampleImagesMissing: '/api/lm/download-example-images' // Download only missing example images
};
// Hugging Face API endpoints

View File

@@ -1641,7 +1641,7 @@ export class BaseModelApiClient {
}
}
async downloadExampleImages(modelHashes, modelTypes = null) {
async downloadExampleImages(modelHashes, modelTypes = null, { force = true } = {}) {
let ws = null;
await state.loadingManager.showWithProgress(async (loading) => {
@@ -1700,8 +1700,13 @@ export class BaseModelApiClient {
// Determine optimize setting
const optimize = state.global?.settings?.optimize_example_images ?? true;
// force=false routes to the regular endpoint, which skips already-processed models
const endpoint = force
? DOWNLOAD_ENDPOINTS.exampleImages
: DOWNLOAD_ENDPOINTS.exampleImagesMissing;
// Make the API request to start the download process
const response = await fetch(DOWNLOAD_ENDPOINTS.exampleImages, {
const response = await fetch(endpoint, {
method: 'POST',
headers: {
'Content-Type': 'application/json'
@@ -1710,6 +1715,7 @@ export class BaseModelApiClient {
model_hashes: modelHashes,
output_dir: outputDir,
optimize: optimize,
force: force,
model_types: modelTypes || [this.apiConfig.config.singularName]
})
});

View File

@@ -144,6 +144,12 @@ export class BulkContextMenu extends BaseContextMenu {
downloadExampleImagesItem.style.display = modelPages.includes(currentModelType) ? 'flex' : 'none';
}
const downloadMissingExampleImagesItem = this.menu.querySelector('[data-action="download-missing-example-images"]');
if (downloadMissingExampleImagesItem) {
// Show on model pages (loras, checkpoints, embeddings), hide on recipes
downloadMissingExampleImagesItem.style.display = ['loras', 'checkpoints', 'embeddings'].includes(currentModelType) ? 'flex' : 'none';
}
const skipMetadataRefreshItem = this.menu.querySelector('[data-action="skip-metadata-refresh"]');
const resumeMetadataRefreshItem = this.menu.querySelector('[data-action="resume-metadata-refresh"]');
@@ -294,8 +300,11 @@ export class BulkContextMenu extends BaseContextMenu {
case 'download-missing-loras':
this.handleDownloadMissingLoras();
break;
case 'download-missing-example-images':
this.handleDownloadExampleImages({ force: false });
break;
case 'download-example-images':
this.handleDownloadExampleImages();
this.handleDownloadExampleImages({ force: true });
break;
case 'clear':
bulkManager.clearSelection();
@@ -340,7 +349,7 @@ export class BulkContextMenu extends BaseContextMenu {
await bulkMissingLoraDownloadManager.downloadMissingLoras(selectedRecipes);
}
async handleDownloadExampleImages() {
async handleDownloadExampleImages({ force = true } = {}) {
if (state.selectedModels.size === 0) {
return;
}
@@ -361,7 +370,7 @@ export class BulkContextMenu extends BaseContextMenu {
try {
const apiClient = getModelApiClient();
await apiClient.downloadExampleImages([...hashes]);
await apiClient.downloadExampleImages([...hashes], null, { force });
} catch (error) {
console.error('Bulk download example images failed:', error);
}

View File

@@ -347,7 +347,10 @@ export const ModelContextMenuMixin = {
openExampleImagesFolder(this.currentCard.dataset.sha256);
return true;
case 'download-examples':
this.downloadExampleImages();
this.downloadExampleImages(false);
return true;
case 'download-examples-force':
this.downloadExampleImages(true);
return true;
case 'civitai':
if (this.currentCard.dataset.from_civitai === 'true') {
@@ -378,7 +381,7 @@ export const ModelContextMenuMixin = {
},
// Download example images method
async downloadExampleImages() {
async downloadExampleImages(force = false) {
const modelHash = this.currentCard.dataset.sha256;
if (!modelHash) {
showToast('toast.contextMenu.missingHash', {}, 'error');
@@ -387,7 +390,7 @@ export const ModelContextMenuMixin = {
try {
const apiClient = getModelApiClient();
await apiClient.downloadExampleImages([modelHash]);
await apiClient.downloadExampleImages([modelHash], null, { force });
} catch (error) {
console.error('Error downloading example images:', error);
}

View File

@@ -33,6 +33,7 @@
<!-- Media / Preview -->
<div class="context-menu-item" data-action="preview"><i class="fas fa-folder-open"></i> {{ t('loras.contextMenu.openExamples') }}</div>
<div class="context-menu-item" data-action="download-examples"><i class="fas fa-download"></i> {{ t('loras.contextMenu.downloadExamples') }}</div>
<div class="context-menu-item" data-action="download-examples-force"><i class="fas fa-redo-alt"></i> {{ t('loras.contextMenu.reprocessExamples') }}</div>
<div class="context-menu-item" data-action="replace-preview"><i class="fas fa-image"></i> {{ t('loras.contextMenu.replacePreview') }}</div>
<div class="context-menu-separator menu-section-break"></div>
<!-- Attributes -->

View File

@@ -47,6 +47,9 @@
<div class="context-menu-item" data-action="download-examples">
<i class="fas fa-download"></i> <span>{{ t('loras.contextMenu.downloadExamples') }}</span>
</div>
<div class="context-menu-item" data-action="download-examples-force">
<i class="fas fa-redo-alt"></i> <span>{{ t('loras.contextMenu.reprocessExamples') }}</span>
</div>
<div class="context-menu-item" data-action="replace-preview">
<i class="fas fa-image"></i> <span>{{ t('loras.contextMenu.replacePreview') }}</span>
</div>
@@ -136,8 +139,11 @@
</div>
<div class="context-menu-section" data-section="download">
<div class="context-menu-section-header">{{ t('loras.bulkOperations.sections.download') }}</div>
<div class="context-menu-item" data-action="download-missing-example-images">
<i class="fas fa-download"></i> <span>{{ t('loras.bulkOperations.downloadMissingExamples') }}</span>
</div>
<div class="context-menu-item" data-action="download-example-images">
<i class="fas fa-download"></i> <span>{{ t('loras.bulkOperations.downloadExamples') }}</span>
<i class="fas fa-redo-alt"></i> <span>{{ t('loras.bulkOperations.reprocessExamples') }}</span>
</div>
<div class="context-menu-item" data-action="download-missing-loras">
<i class="fas fa-download"></i> <span>{{ t('loras.bulkOperations.downloadMissingLoras') }}</span>

View File

@@ -33,6 +33,7 @@
<!-- Media / Preview -->
<div class="context-menu-item" data-action="preview"><i class="fas fa-folder-open"></i> {{ t('loras.contextMenu.openExamples') }}</div>
<div class="context-menu-item" data-action="download-examples"><i class="fas fa-download"></i> {{ t('loras.contextMenu.downloadExamples') }}</div>
<div class="context-menu-item" data-action="download-examples-force"><i class="fas fa-redo-alt"></i> {{ t('loras.contextMenu.reprocessExamples') }}</div>
<div class="context-menu-item" data-action="replace-preview"><i class="fas fa-image"></i> {{ t('loras.contextMenu.replacePreview') }}</div>
<div class="context-menu-separator menu-section-break"></div>
<!-- Attributes -->

View File

@@ -2155,4 +2155,31 @@ describe('Interaction-level regression coverage', () => {
excludedItem.dispatchEvent(new Event('click', { bubbles: true }));
expect(window.pageControls.enterExcludedView).toHaveBeenCalledTimes(1);
});
it('routes single-model example downloads to missing-only and force paths', async () => {
document.body.innerHTML = `
<div id="loraContextMenu" class="context-menu">
<div class="context-menu-item" data-action="download-examples"></div>
<div class="context-menu-item" data-action="download-examples-force"></div>
</div>
`;
const { LoraContextMenu } = await import('../../../static/js/components/ContextMenu/LoraContextMenu.js');
const contextMenu = new LoraContextMenu();
const card = document.createElement('div');
card.className = 'model-card';
card.dataset.filepath = '/models/test.safetensors';
card.dataset.sha256 = 'abc123hash';
document.body.appendChild(card);
contextMenu.showMenu(100, 100, card);
document.querySelector('[data-action="download-examples"]').dispatchEvent(new Event('click', { bubbles: true }));
expect(downloadExampleImagesApiMock).toHaveBeenCalledWith(['abc123hash'], null, { force: false });
contextMenu.showMenu(100, 100, card);
document.querySelector('[data-action="download-examples-force"]').dispatchEvent(new Event('click', { bubbles: true }));
expect(downloadExampleImagesApiMock).toHaveBeenCalledWith(['abc123hash'], null, { force: true });
});
});

View File

@@ -529,7 +529,8 @@ async def test_not_found_example_images_are_cleaned(
model_dir = images_root / model_hash
model_dir.mkdir(parents=True, exist_ok=True)
(model_dir / "image_0.png").write_bytes(b"first")
# Pre-existing file collides with the valid image index (1) so the
# pre-download existence check must skip it without a network request
(model_dir / "image_1.png").write_bytes(b"second")
async def fake_process_local_examples(*_args, **_kwargs):
@@ -608,11 +609,188 @@ async def test_not_found_example_images_are_cleaned(
]
files = sorted(p.name for p in model_dir.iterdir())
assert files == ["image_0.png", "image_1.png"]
assert (model_dir / "image_0.png").read_bytes() == b"first"
assert files == ["image_1.png"]
assert (model_dir / "image_1.png").read_bytes() == b"second"
async def test_failed_models_retried_when_explicitly_targeted(
monkeypatch: pytest.MonkeyPatch,
tmp_path,
settings_manager,
):
ws_manager = RecordingWebSocketManager()
manager = download_module.DownloadManager(ws_manager=ws_manager)
images_root = tmp_path / "examples"
monkeypatch.setitem(settings_manager.settings, "example_images_path", str(images_root))
model_hash = "a" * 64
model_path = tmp_path / "model.safetensors"
model_path.write_text("data", encoding="utf-8")
model_metadata = {
"sha256": model_hash,
"model_name": "Failed Example",
"file_path": str(model_path),
"file_name": "model.safetensors",
"civitai": {"images": [{"url": "https://example.com/valid.png"}]},
}
scanner = StubScanner([model_metadata.copy()])
_patch_scanner(monkeypatch, scanner)
# Persist a previous failure so the skip path is exercised
images_root.mkdir(parents=True, exist_ok=True)
(images_root / ".download_progress.json").write_text(
json.dumps(
{
"failed_models": [model_hash],
"processed_models": [],
"rate_limited_models": [],
}
),
encoding="utf-8",
)
async def fake_process_local_examples(*_args, **_kwargs):
return False
async def fake_get_updated_model(model_hash_arg, _scanner):
return model_metadata
class DownloaderStub:
def __init__(self):
self.calls: list[str] = []
async def download_to_memory(self, url, *_args, **_kwargs):
self.calls.append(url)
return True, b"\x89PNG\r\n\x1a\n", {"content-type": "image/png"}
downloader = DownloaderStub()
async def fake_get_downloader():
return downloader
monkeypatch.setattr(
download_module.ExampleImagesProcessor,
"process_local_examples",
staticmethod(fake_process_local_examples),
)
monkeypatch.setattr(
download_module.MetadataUpdater,
"get_updated_model",
staticmethod(fake_get_updated_model),
)
monkeypatch.setattr(download_module, "get_downloader", fake_get_downloader)
# Without explicit hashes the previously failed model is skipped
skipped_manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
result = await skipped_manager.start_download({"model_types": ["lora"], "delay": 0})
assert result["success"] is True
if skipped_manager._download_task is not None:
await asyncio.wait_for(skipped_manager._download_task, timeout=1)
assert downloader.calls == []
# With explicit hashes the previously failed model is retried and cleared
result = await manager.start_download(
{"model_types": ["lora"], "delay": 0, "model_hashes": [model_hash]}
)
assert result["success"] is True
if manager._download_task is not None:
await asyncio.wait_for(manager._download_task, timeout=1)
assert downloader.calls == ["https://example.com/valid.png"]
assert manager._progress["failed_models"] == set()
assert model_hash in manager._progress["processed_models"]
async def test_explicit_targets_fill_partial_example_gaps(
monkeypatch: pytest.MonkeyPatch,
tmp_path,
settings_manager,
):
ws_manager = RecordingWebSocketManager()
images_root = tmp_path / "examples"
monkeypatch.setitem(settings_manager.settings, "example_images_path", str(images_root))
model_hash = "b" * 64
model_path = tmp_path / "model.safetensors"
model_path.write_text("data", encoding="utf-8")
model_metadata = {
"sha256": model_hash,
"model_name": "Partial Example",
"file_path": str(model_path),
"file_name": "model.safetensors",
"civitai": {
"images": [
{"url": "https://example.com/first.png"},
{"url": "https://example.com/second.png"},
]
},
}
scanner = StubScanner([model_metadata.copy()])
_patch_scanner(monkeypatch, scanner)
# Simulate a partially populated folder: index 0 already downloaded
model_dir = images_root / model_hash
model_dir.mkdir(parents=True, exist_ok=True)
(model_dir / "image_0.png").write_bytes(b"existing")
async def fake_process_local_examples(*_args, **_kwargs):
return False
async def fake_get_updated_model(model_hash_arg, _scanner):
return model_metadata
class DownloaderStub:
def __init__(self):
self.calls: list[str] = []
async def download_to_memory(self, url, *_args, **_kwargs):
self.calls.append(url)
return True, b"\x89PNG\r\n\x1a\n", {"content-type": "image/png"}
downloader = DownloaderStub()
async def fake_get_downloader():
return downloader
monkeypatch.setattr(
download_module.ExampleImagesProcessor,
"process_local_examples",
staticmethod(fake_process_local_examples),
)
monkeypatch.setattr(
download_module.MetadataUpdater,
"get_updated_model",
staticmethod(fake_get_updated_model),
)
monkeypatch.setattr(download_module, "get_downloader", fake_get_downloader)
# Untargeted run treats the populated folder as done
untargeted = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
result = await untargeted.start_download({"model_types": ["lora"], "delay": 0})
assert result["success"] is True
if untargeted._download_task is not None:
await asyncio.wait_for(untargeted._download_task, timeout=1)
assert downloader.calls == []
# Explicitly targeted run fills only the missing index, skipping the
# existing file without a network request
targeted = download_module.DownloadManager(ws_manager=ws_manager)
result = await targeted.start_download(
{"model_types": ["lora"], "delay": 0, "model_hashes": [model_hash]}
)
assert result["success"] is True
if targeted._download_task is not None:
await asyncio.wait_for(targeted._download_task, timeout=1)
assert downloader.calls == ["https://example.com/second.png"]
assert (model_dir / "image_1.png").exists()
assert (model_dir / "image_0.png").read_bytes() == b"existing"
@pytest.fixture
def settings_manager():
return get_settings_manager()

View File

@@ -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)

View File

@@ -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)