perf(example-images): parallelize downloads and recover transient failures

- Add a shared MediaDownloadLimiter (default 6 concurrent) with a per-run
  429 cooldown and AIMD recovery; drop the empirical inter-model delay.
- Request width-aware Civitai renditions (800px images, transcoded 450px
  videos) and map Genur signed/direct URLs to c.genur.art _small.webp or
  _large.webp.
- Return an ImageDownloadOutcome so offline/connect/timeout/5xx failures
  stay transient instead of being written as downloadFailed; track
  retryable_models and leave those models pending.
- Force retries flagged Genur images and clears downloadFailed/downloadError
  only after a successful write, through the normal metadata save/cache path.
- Force bypasses the populated-folder guard; success is false when flagged
  images were skipped, so models are no longer falsely marked recovered.
- Add settings, tests, and a design/measurement note.
This commit is contained in:
Will Miao
2026-11-08 18:05:02 +08:00
parent 9ed75e8fa4
commit 04b26b2e9f
10 changed files with 1671 additions and 293 deletions
@@ -0,0 +1,133 @@
# Example-image download throughput
## Problem
Downloading example media for a large library was effectively serial: one
Civitai CDN request per image/video, one model after another, plus an
empirical `delay` (0.2 s) between models. A model with 10–20 gallery items
therefore paid 10–20 network round-trips before the next model started.
The inter-model delay is not the cost. In a 10-image probe it accounted for
about 2 s of a 7 s run; the remaining time was the per-image round-trips.
## Is there a batch endpoint?
No. Civitai's public API exposes:
- `GET /api/v1/images?ids=...` — **metadata** for up to 100 images in one
request (`IMAGE_IDS_BATCH_MAX = 100`), not the binaries.
- One model-version response that already contains the image URLs the
downloader uses.
- One binary per `image.civitai.com` URL. Each URL 301-redirects to a
`blobs-b2.civitai.com` object.
The image cacher also has no multi-URL delivery endpoint. The practical
"batch" is therefore client-side pipelining: bounded concurrent requests, a
rendition that is cheaper to transfer, and a shared cooldown when the CDN
returns 429.
## Measurements (single sample, 2026-11)
Ten images from one model, cold process, Civitai CDN over the local link:
| Strategy | Wall time | Bytes |
| --- | ---: | ---: |
| Old: serial, `/optimized=true` | 7.07 s | 1669 KiB |
| New: 4 concurrent, `width=800` | 3.08 s | 834 KiB |
| New: 8 concurrent, `width=800` | 1.62 s | 834 KiB |
| New: 6 concurrent (default), `width=800` | 3.23 s | 834 KiB |
A 720×924 sample video:
| Rendition | Size |
| --- | ---: |
| Old: `/optimized=true` (original MP4) | 73.4 MB |
| New: `transcode=true,width=450,optimized=true` | 27.0 MB |
The video case was the largest single regression in the old path: the URL
helper only appended `optimized=true`, which the cacher serves as the original
video.
## Design
- `MediaDownloadLimiter` (`py/utils/example_images_processor.py`):
- bounds concurrent requests (default 6),
- shares a 429 cooldown across every worker instead of letting each task
sleep and retry independently,
- halves the limit on a 429 and adds one slot back after a run of successes
(AIMD), so a sustained limit settles below the threshold.
- `download_model_images_with_tracking` builds one task per pending item and
runs them through the limiter. The task result still drives the existing
per-model bookkeeping (`failed_models`, `rate_limited_models`, metadata
refresh on 404).
- URLs are built with `build_civitai_media_url`:
- images: `width=800,optimized=true` (the filename is preserved),
- videos: `transcode=true,width=450,optimized=true`,
- non-Civitai URLs: unchanged,
- `optimize=False`: the original URL.
- The legacy inter-model `delay` defaults to 0. API callers that still send it
keep the old behavior.
- Failure lists now store the *original* metadata URL, which is what
`_remove_failed_images_from_metadata` matches against.
## Settings (`DEFAULT_SETTINGS`)
| Key | Default | Meaning |
| --- | ---: | --- |
| `example_images_download_concurrency` | 6 | Simultaneous CDN requests per run |
| `example_images_max_width` | 800 | Image rendition width |
| `example_images_video_max_width` | 450 | Video transcode width |
| `example_images_rate_limit_cooldown` | 30 | Fallback cooldown when `Retry-After` is absent |
Widths are chosen from the image cacher's common-size ladder
(96/320/450/512/800/1200/1600/2200); off-ladder widths are snapped up by the
cacher, so they do not buy a smaller transfer.
## Follow-up: transient failures and Genur recovery
The first implementation still treated *any* non-404 failure as permanent. A
network outage (or the ConnectivityGuard's shared cooldown message) was
written to the sidecar as `downloadFailed: true, downloadError: "not_found"`,
which made a transient failure permanent and produced empty example folders
that were reprocessed every run. Genur was the visible victim: its signed
proxy host (`img.genur.art`) can be unreachable (fake-IP DNS / TLS reset)
while the direct host serves the same asset.
The follow-up changes:
- `download_model_images_with_tracking` returns an `ImageDownloadOutcome`
(`downloaded`, `failed`, `rate_limited`, `transient`, `recovered`,
`skipped_failed`) instead of a tuple. Offline/connect/timeout/5xx errors go
to `transient` and are **not** written to the sidecar; the model is left
pending (tracked in `retryable_models`) so the next run retries it.
- `retry_failed=True` (force / explicit targets) re-attempts Genur entries
that carry `downloadFailed`; the flag is cleared only after a file is
actually written, through `_clear_failed_images_from_metadata` and the
normal `save_metadata` + cache-update path. No bulk sidecar rewrite is
needed.
- Genur signed URLs are resolved at download time:
`img.genur.art/sig/.../<base64>` → `c.genur.art/<uuid>_small.webp`
(`_large.webp` when optimization is disabled). Direct
`c.genur.art/<uuid>` URLs get the same `_small.webp` suffix. The direct
renditions return 200 for the sampled Organic Sauce / Better Lego assets.
- Force now bypasses the model-level "populated folder" guard, so flagged
images in partially populated folders can be retried; per-image existence
checks still skip files that are already there.
Existing affected libraries recover by restarting the instance and running a
force download (global or per-model). The 398 Genur entries counted in the
user's sidecars are eligible; Civitai 404 entries stay flagged because those
assets are genuinely gone.
## Future work
- Pipeline models, not just images within a model: start model N+1's image
jobs while model N's metadata/failure bookkeeping finishes. This helps
models with only a few gallery items.
- Stream media to a temp file instead of `download_to_memory`; the current
pool bounds memory at `concurrency × item size`, which is fine for
transcoded videos but large if `optimize_example_images` is disabled.
- Cache the `image.civitai.com` → `blobs-b2.civitai.com` redirect when a
rendition is retried, to save one hop per retry.
- Expose the four settings in the settings UI if users need to tune them
without editing `settings.json`.
+9
View File
@@ -147,6 +147,15 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
"example_images_local_root": "",
"example_images_open_uri_template": "",
"optimize_example_images": True,
# Example-media downloads are one CDN request per image/video. Serial
# requests spend most of their wall time waiting on round-trips, so a
# small bounded pool plus a shared 429 cooldown is the throughput lever.
# Widths follow the Civitai image-cacher's common-size ladder; videos are
# capped lower because a video rendition is orders of magnitude larger.
"example_images_download_concurrency": 6,
"example_images_max_width": 800,
"example_images_video_max_width": 450,
"example_images_rate_limit_cooldown": 30,
"auto_download_example_images": False,
"blur_mature_content": True,
"mature_blur_level": "R",
+88 -4
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import posixpath
import re
from typing import Any, Dict, Iterable, Mapping, Sequence
from urllib.parse import parse_qs, urlparse, urlunparse
@@ -338,14 +339,40 @@ def resolve_license_info(
return payload, build_license_flags(payload)
def _normalize_media_width(width: int | None) -> int | None:
if width is None:
return None
try:
return max(1, int(width))
except (TypeError, ValueError):
return 450
def _media_rendition_params(media_type: str | None, width: int | None) -> str:
"""Return the Civitai image-cacher params for a media rendition."""
is_video = (media_type or "").lower() == "video"
normalized_width = _normalize_media_width(width)
if normalized_width is None:
return "transcode=true,optimized=true" if is_video else "optimized=true"
if is_video:
return f"transcode=true,width={normalized_width},optimized=true"
return f"width={normalized_width},optimized=true"
def rewrite_preview_url(
source_url: str | None, media_type: str | None = None
source_url: str | None,
media_type: str | None = None,
*,
width: int | None = 450,
) -> tuple[str | None, bool]:
"""Rewrite Civitai preview URLs to use optimized renditions.
Args:
source_url: Original preview URL from the Civitai API.
media_type: Optional media type hint (e.g. ``"image"`` or ``"video"``).
width: Rendition width. ``None`` keeps the legacy width-less
``optimized=true`` form.
Returns:
A tuple of the potentially rewritten URL and a flag indicating whether the
@@ -368,9 +395,7 @@ def rewrite_preview_url(
if hostname == "civitai.com" or not hostname.endswith(".civitai.com"):
return source_url, False
replacement = "/width=450,optimized=true"
if (media_type or "").lower() == "video":
replacement = "/transcode=true,width=450,optimized=true"
replacement = f"/{_media_rendition_params(media_type, width)}"
if "/original=true" not in parsed.path:
return source_url, False
@@ -383,8 +408,67 @@ def rewrite_preview_url(
return rewritten, True
def build_civitai_media_url(
source_url: str | None,
media_type: str | None = None,
*,
width: int | None = 450,
optimize: bool = True,
) -> str | None:
"""Build the image-cacher delivery URL for a Civitai media asset.
Unlike :func:`rewrite_preview_url`, this handles URLs that do not carry the
current ``/original=true`` params segment (older cached metadata, legacy
``.../<uuid>`` URLs) and replaces an existing rendition segment instead of
appending a second one. The filename is preserved: the image cacher uses it
to derive the output name, and dropping it makes the requested rendition
ambiguous.
"""
if not source_url or not optimize:
return source_url
try:
parsed = urlparse(source_url)
except ValueError:
return source_url
hostname = parsed.hostname
if hostname is None:
return source_url
hostname = hostname.lower()
if hostname == "civitai.com" or not hostname.endswith(".civitai.com"):
return source_url
rewritten, changed = rewrite_preview_url(source_url, media_type, width=width)
if changed:
return rewritten
segments = [segment for segment in parsed.path.split("/") if segment]
if not segments:
return source_url
params = _media_rendition_params(media_type, width)
last = segments[-1]
if "=" in last:
# Legacy ``.../<uuid>/optimized=true`` form (no filename).
segments[-1] = params
elif len(segments) >= 2 and "=" in segments[-2]:
# The rendition segment is immediately before the filename.
segments[-2] = params
elif posixpath.splitext(last)[1]:
# ``.../<uuid>/<filename>``: insert the rendition before the filename.
segments.insert(len(segments) - 1, params)
else:
# ``.../<uuid>``: no filename to preserve.
segments.append(params)
return urlunparse(parsed._replace(path="/" + "/".join(segments)))
__all__ = [
"build_civitai_image_page_url",
"build_civitai_media_url",
"build_license_flags",
"extract_civitai_image_id",
"extract_civitai_image_id_from_cdn_url",
+246 -80
View File
@@ -20,7 +20,7 @@ from ..utils.example_images_paths import (
uses_library_scoped_folders,
)
from ..utils.metadata_manager import MetadataManager
from .example_images_processor import ExampleImagesProcessor
from .example_images_processor import ExampleImagesProcessor, MediaDownloadLimiter
from .example_images_metadata import (
MetadataUpdater,
update_cache_from_metadata,
@@ -79,6 +79,7 @@ class _DownloadProgress(dict[str, Any]):
failed_models=set(),
reprocessed_models=set(),
rate_limited_models=set(),
retryable_models=set(),
)
def snapshot(self) -> Dict[str, Any]:
@@ -90,6 +91,7 @@ class _DownloadProgress(dict[str, Any]):
snapshot["failed_models"] = list(self["failed_models"])
snapshot["reprocessed_models"] = list(self.get("reprocessed_models", set()))
snapshot["rate_limited_models"] = list(self.get("rate_limited_models", set()))
snapshot["retryable_models"] = list(self.get("retryable_models", set()))
return snapshot
@@ -163,6 +165,38 @@ class DownloadManager:
return ""
return ensure_library_root_exists(library_name)
def _create_media_download_limiter(
self, optimize: bool = True
) -> MediaDownloadLimiter:
"""Build the run-wide limiter from settings.
The limiter is shared by every model in a run so a 429 seen while
downloading one model also pauses the next one instead of being
forgotten at the model boundary.
"""
settings = get_settings_manager()
try:
concurrency = int(
settings.get("example_images_download_concurrency", 6)
)
except (TypeError, ValueError):
concurrency = 6
try:
cooldown = float(
settings.get("example_images_rate_limit_cooldown", 30)
)
except (TypeError, ValueError):
cooldown = 30.0
if not optimize:
# Original renditions can be tens of megabytes per video; keep the
# memory/bandwidth spike bounded when optimization is disabled.
concurrency = min(concurrency, 2)
return MediaDownloadLimiter(
max_concurrency=max(1, concurrency),
default_cooldown=max(1.0, cooldown),
)
async def start_download(self, options: Dict[str, Any]):
"""Start downloading example images for models."""
@@ -171,7 +205,7 @@ class DownloadManager:
auto_mode = data.get("auto_mode", False)
optimize = data.get("optimize", True)
model_types = data.get("model_types", ["lora", "checkpoint"])
delay = float(data.get("delay", 0.2))
delay = float(data.get("delay", 0) or 0)
force = data.get("force", False)
model_hashes = data.get("model_hashes", [])
@@ -589,6 +623,7 @@ class DownloadManager:
"""Download example images for all models (or only the given hashes)."""
downloader = await get_downloader()
limiter = self._create_media_download_limiter(optimize)
try:
# Get scanners
@@ -654,6 +689,7 @@ class DownloadManager:
library_name,
force,
explicit_targets,
limiter=limiter,
)
# Update progress
@@ -671,9 +707,12 @@ class DownloadManager:
if should_stop:
break
# Only add delay after remote download of models, and not after processing the last model
# Pacing is handled by the shared media-download limiter. The
# legacy ``delay`` option is kept for API callers that still
# send it, but defaults to 0.
if (
was_remote_download
delay > 0
and was_remote_download
and i < len(all_models) - 1
and current_status == "running"
):
@@ -751,6 +790,7 @@ class DownloadManager:
library_name,
force: bool = False,
explicit_targets: bool = False,
limiter: MediaDownloadLimiter | None = None,
):
"""Process a single model download."""
@@ -784,10 +824,12 @@ class DownloadManager:
)
existing_files = _model_directory_has_files(model_dir)
# 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:
# Model-level guard: a populated folder counts as done. Force and
# explicitly targeted models bypass it so the per-image existence
# pre-check can fill individual gaps, retry flagged images, or
# clear a stale ``downloadFailed`` marker without re-fetching the
# files that are already present.
if not force and not explicit_targets:
# Skip if already processed AND directory exists with files
if model_hash in self._progress["processed_models"]:
if existing_files:
@@ -846,20 +888,30 @@ class DownloadManager:
if civitai_payload.get("images"):
images = civitai_payload.get("images", [])
(
success,
is_stale,
failed_images,
rate_limited_images,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash, model_name, images, model_dir, optimize, downloader
outcome = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
images,
model_dir,
optimize,
downloader,
limiter=limiter,
retry_failed=bool(force or explicit_targets),
)
failed_urls: Set[str] = set(failed_images)
rate_limited_urls: Set[str] = set(rate_limited_images)
failed_urls: Set[str] = set(outcome.failed)
rate_limited_urls: Set[str] = set(outcome.rate_limited)
transient_urls: Set[str] = set(outcome.transient)
recovered_urls: Set[str] = set(outcome.recovered)
final_success = outcome.success
if recovered_urls:
await self._clear_failed_images_from_metadata(
model_hash, model_name, recovered_urls, scanner
)
# If metadata is stale, try to refresh it
if is_stale and model_hash not in self._progress["refreshed_models"]:
if outcome.is_stale and model_hash not in self._progress["refreshed_models"]:
await MetadataUpdater.refresh_model_metadata(
model_hash, model_name, scanner_type, scanner, self._progress
)
@@ -876,26 +928,35 @@ class DownloadManager:
if updated_civitai.get("images"):
# Retry download with updated metadata
updated_images = updated_civitai.get("images", [])
(
success,
_,
additional_failed,
additional_rate_limited,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
updated_images,
model_dir,
optimize,
downloader,
retry_outcome = (
await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
updated_images,
model_dir,
optimize,
downloader,
limiter=limiter,
retry_failed=bool(force or explicit_targets),
)
)
failed_urls.update(additional_failed)
rate_limited_urls.update(additional_rate_limited)
# The refreshed metadata is now the source of truth:
# failures against the pre-refresh URLs are gone (or
# will fail again on this pass).
failed_urls = set(retry_outcome.failed)
rate_limited_urls = set(retry_outcome.rate_limited)
transient_urls = set(retry_outcome.transient)
recovered_urls.update(retry_outcome.recovered)
final_success = retry_outcome.success
if retry_outcome.recovered:
await self._clear_failed_images_from_metadata(
model_hash, model_name, retry_outcome.recovered, scanner
)
self._progress["refreshed_models"].add(model_hash)
# Separate permanent failures from rate-limited ones
# Separate permanent failures from retryable ones
permanent_failures = failed_urls - rate_limited_urls
if permanent_failures:
@@ -914,17 +975,25 @@ class DownloadManager:
len(rate_limited_urls),
model_name,
)
# Clear failed_models so non-force runs can retry
if transient_urls:
self._progress["retryable_models"].add(model_hash)
logger.warning(
"%d example images for %s hit a transient network error; will retry next run",
len(transient_urls),
model_name,
)
if rate_limited_urls or transient_urls:
# Leave the model pending: the next run retries the
# rate-limited/transient URLs. Permanent failures are
# still flagged above; do not mark the model processed.
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"
"Removed %s from failed_models for a retryable failure",
model_name,
)
if rate_limited_urls:
# Don't mark as failed or fully processed — rate-limited
# images will be retried next time.
pass
elif permanent_failures:
self._progress["failed_models"].add(model_hash)
self._progress["processed_models"].add(model_hash)
@@ -933,7 +1002,7 @@ class DownloadManager:
len(permanent_failures),
model_name,
)
elif success:
elif final_success:
self._progress["processed_models"].add(model_hash)
if (force or explicit_targets) and model_hash in self._progress["failed_models"]:
self._progress["failed_models"].discard(model_hash)
@@ -1035,7 +1104,7 @@ class DownloadManager:
model_hashes = data.get("model_hashes", [])
optimize = data.get("optimize", True)
model_types = data.get("model_types", ["lora", "checkpoint"])
delay = float(data.get("delay", 0.2))
delay = float(data.get("delay", 0) or 0)
if not model_hashes:
raise DownloadConfigurationError("Missing model_hashes parameter")
@@ -1106,6 +1175,7 @@ class DownloadManager:
"""Download example images for specific models only - synchronous version."""
downloader = await get_downloader()
limiter = self._create_media_download_limiter(optimize)
try:
# Get scanners
@@ -1160,6 +1230,7 @@ class DownloadManager:
optimize,
downloader,
library_name,
limiter=limiter,
)
if was_successful:
@@ -1181,9 +1252,12 @@ class DownloadManager:
if should_stop:
break
# Only add delay after remote download, and not after processing the last model
# Pacing is handled by the shared media-download limiter. The
# legacy ``delay`` option is kept for API callers that still
# send it, but defaults to 0.
if (
was_successful
delay > 0
and was_successful
and i < len(models_to_process) - 1
and current_status == "running"
):
@@ -1255,6 +1329,7 @@ class DownloadManager:
optimize,
downloader,
library_name,
limiter: MediaDownloadLimiter | None = None,
):
"""Process a specific model for forced download, ignoring previous download status."""
@@ -1313,20 +1388,30 @@ class DownloadManager:
if civitai_payload.get("images"):
images = civitai_payload.get("images", [])
(
success,
is_stale,
failed_images,
rate_limited_images,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash, model_name, images, model_dir, optimize, downloader
outcome = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
images,
model_dir,
optimize,
downloader,
limiter=limiter,
retry_failed=True,
)
failed_urls: Set[str] = set(failed_images)
rate_limited_urls: Set[str] = set(rate_limited_images)
failed_urls: Set[str] = set(outcome.failed)
rate_limited_urls: Set[str] = set(outcome.rate_limited)
transient_urls: Set[str] = set(outcome.transient)
recovered_urls: Set[str] = set(outcome.recovered)
final_success = outcome.success
if recovered_urls:
await self._clear_failed_images_from_metadata(
model_hash, model_name, recovered_urls, scanner
)
# If metadata is stale, try to refresh it
if is_stale and model_hash not in self._progress["refreshed_models"]:
if outcome.is_stale and model_hash not in self._progress["refreshed_models"]:
await MetadataUpdater.refresh_model_metadata(
model_hash, model_name, scanner_type, scanner, self._progress
)
@@ -1343,26 +1428,35 @@ class DownloadManager:
if updated_civitai.get("images"):
# Retry download with updated metadata
updated_images = updated_civitai.get("images", [])
(
success,
_,
additional_failed_images,
additional_rate_limited,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
updated_images,
model_dir,
optimize,
downloader,
retry_outcome = (
await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
updated_images,
model_dir,
optimize,
downloader,
limiter=limiter,
retry_failed=True,
)
)
failed_urls.update(additional_failed_images)
rate_limited_urls.update(additional_rate_limited)
# The refreshed metadata is now the source of truth:
# failures against the pre-refresh URLs are gone (or
# will fail again on this pass).
failed_urls = set(retry_outcome.failed)
rate_limited_urls = set(retry_outcome.rate_limited)
transient_urls = set(retry_outcome.transient)
recovered_urls.update(retry_outcome.recovered)
final_success = retry_outcome.success
if retry_outcome.recovered:
await self._clear_failed_images_from_metadata(
model_hash, model_name, retry_outcome.recovered, scanner
)
self._progress["refreshed_models"].add(model_hash)
# Separate permanent failures from rate-limited ones
# Separate permanent failures from retryable ones
permanent_failures = failed_urls - rate_limited_urls
# Only remove permanently failed images from metadata
@@ -1379,13 +1473,21 @@ class DownloadManager:
model_name,
)
# Mark as processed only when no rate-limited images remain
if rate_limited_urls:
pass
if transient_urls:
self._progress["retryable_models"].add(model_hash)
logger.warning(
"%d example images for %s hit a transient network error; will retry next run",
len(transient_urls),
model_name,
)
# Mark as processed only when nothing is left to retry.
if rate_limited_urls or transient_urls:
self._progress["failed_models"].discard(model_hash)
elif permanent_failures:
self._progress["processed_models"].add(model_hash)
self._progress["failed_models"].add(model_hash)
elif success:
elif final_success:
self._progress["processed_models"].add(model_hash)
return True # Return True to indicate a remote download happened
@@ -1401,6 +1503,76 @@ class DownloadManager:
self._progress["last_error"] = error_msg
return False # Return False on exception
async def _clear_failed_images_from_metadata(
self,
model_hash: str,
model_name: str,
recovered_images: Iterable[str],
scanner,
) -> None:
"""Clear stale ``downloadFailed`` flags for images that are now on disk.
Option B keeps the sidecar untouched until a retry actually succeeds;
this is the write half of that contract. Only the two failure fields
are removed; the URL and every other image field stay as they are.
"""
recovered_set: Set[str] = {url for url in recovered_images if url}
if not recovered_set:
return
try:
model_data = await MetadataUpdater.get_updated_model(model_hash, scanner)
if not model_data:
logger.warning(
"Could not find model data for %s to clear recovered images",
model_name,
)
return
civitai_payload = model_data.get("civitai") or {}
current_images = civitai_payload.get("images") or []
cleared = 0
for image in current_images:
if image.get("url") not in recovered_set:
continue
if not image.get("downloadFailed"):
continue
image.pop("downloadFailed", None)
image.pop("downloadError", None)
cleared += 1
if not cleared:
return
file_path = model_data.get("file_path")
if not file_path:
return
model_copy = model_data.copy()
model_copy.pop("folder", None)
await MetadataManager.save_metadata(file_path, model_copy)
try:
await update_cache_from_metadata(scanner, file_path, model_copy)
except Exception as cache_error:
logger.warning(
"Failed to update cache after clearing recovered flags for %s: %s",
model_name,
cache_error,
)
logger.info(
"Recovered %s previously failed example images for %s",
cleared,
model_name,
)
except Exception as e:
logger.error(
"Failed to clear recovered example image flags for %s: %s",
model_name,
e,
exc_info=True,
)
async def _remove_failed_images_from_metadata(
self,
model_hash: str,
@@ -1439,13 +1611,7 @@ class DownloadManager:
for image in current_images:
image_url = image.get("url")
optimized_url = (
ExampleImagesProcessor.get_civitai_optimized_url(image_url)
if image_url and "civitai.com" in image_url
else None
)
if image_url not in failed_set and optimized_url not in failed_set:
if image_url not in failed_set:
continue
if image.get("downloadFailed"):
+492 -199
View File
@@ -1,14 +1,25 @@
import asyncio
import base64
import logging
import os
import posixpath
import re
import random
import string
import time
from dataclasses import dataclass, field
from typing import Any
from urllib.parse import urlparse
from aiohttp import web
from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
from ..services.connectivity_guard import (
OFFLINE_COOLDOWN_ERROR,
OFFLINE_FRIENDLY_MESSAGE,
)
from ..services.service_registry import ServiceRegistry
from ..services.settings_manager import get_settings_manager
from ..utils.civitai_utils import build_civitai_media_url
from ..utils.example_images_paths import get_model_folder, get_model_relative_path
from .example_images_metadata import MetadataUpdater, update_cache_from_metadata
from .file_utils import detect_media_extension
@@ -16,6 +27,236 @@ from ..utils.metadata_manager import MetadataManager
logger = logging.getLogger(__name__)
MAX_DOWNLOAD_ATTEMPTS = 3
_GENUR_HOSTS = frozenset({"img.genur.art", "c.genur.art"})
_TRANSIENT_ERROR_MARKERS = (
"network offline",
"offline_cooldown",
"cannot connect to host",
"server disconnected",
"connection reset",
"connection refused",
"temporary failure",
"name or service not known",
"timed out",
"timeout",
)
def _write_media_file(path: str, content: bytes) -> None:
with open(path, "wb") as handle:
handle.write(content)
def is_genur_media_url(url: str) -> bool:
"""Return True for Genur image URLs handled by the fallback resolver."""
try:
hostname = urlparse(url or "").hostname
except ValueError:
return False
return bool(hostname) and hostname.lower() in _GENUR_HOSTS
def resolve_genur_media_url(
url: str,
media_type: str | None = None,
*,
size: str = "small",
) -> str:
"""Map Genur image URLs to a direct rendition that is currently reachable.
Genur's signed proxy (``img.genur.art/sig/.../<base64>``) embeds the direct
asset URL in the last path segment. The proxy host is not reachable from
every network (fake-IP DNS / TLS resets), while the direct host serves the
same image as ``<uuid>_small.webp`` and ``<uuid>_large.webp``. Returning
the direct rendition avoids a guaranteed failed request and the
ConnectivityGuard cooldown that follows it.
Non-Genur URLs, videos, and already-suffixed renditions are returned
unchanged.
"""
if not url or (media_type or "").lower() == "video":
return url
try:
parsed = urlparse(url)
except ValueError:
return url
hostname = (parsed.hostname or "").lower()
if hostname == "img.genur.art":
payload = parsed.path.rsplit("/", 1)[-1]
try:
padded = payload + "=" * (-len(payload) % 4)
decoded = base64.b64decode(padded).decode("utf-8")
except (ValueError, UnicodeDecodeError):
return url
try:
decoded_parsed = urlparse(decoded)
except ValueError:
return url
if (decoded_parsed.hostname or "").lower() != "c.genur.art":
return url
base = f"https://c.genur.art{decoded_parsed.path}"
elif hostname == "c.genur.art":
base = f"https://c.genur.art{parsed.path}"
else:
return url
base = base.rstrip("/")
if base.endswith(("_small.webp", "_large.webp")):
return base
if posixpath.splitext(base)[1].lower() in SUPPORTED_MEDIA_EXTENSIONS["images"]:
# Already a direct image path; do not invent a suffix.
return base
suffix = "_large.webp" if size == "large" else "_small.webp"
return f"{base}{suffix}"
@dataclass
class ImageDownloadOutcome:
"""Result of one model's example-media download pass.
``downloaded`` and ``recovered`` contain the *original* metadata URLs so
callers can match them against ``civitai.images``. ``success`` is false
when anything is still missing: a permanent failure, a rate-limit retry,
a transient network failure, or a previously flagged image that this pass
did not retry.
"""
downloaded: list[str] = field(default_factory=list)
failed: list[str] = field(default_factory=list)
rate_limited: list[str] = field(default_factory=list)
transient: list[str] = field(default_factory=list)
recovered: list[str] = field(default_factory=list)
skipped_failed: int = 0
is_stale: bool = False
@property
def success(self) -> bool:
return not (
self.failed
or self.rate_limited
or self.transient
or self.skipped_failed
)
class MediaDownloadLimiter:
"""Bound concurrent media requests and share 429 cooldowns across workers.
Civitai serves one binary per request and each request costs a CDN
round-trip, so serial downloads spend most of their wall clock waiting.
Concurrency is the useful lever, but it has to stay below the CDN's
threshold: a 429 is a signal for the whole destination, not for the one
task that happened to receive it. A cooldown armed here therefore parks
every worker, and the limit is halved on a 429 and grows back one slot at
a time after a run of successes (AIMD) instead of oscillating around the
limit.
All waits are bounded by ``max_cooldown``. The per-request retry budget
lives in the caller; this class only decides when a request may start.
"""
def __init__(
self,
max_concurrency: int = 6,
*,
min_cooldown: float = 1.0,
default_cooldown: float = 30.0,
max_cooldown: float = 300.0,
recovery_successes: int | None = None,
) -> None:
self.max_concurrency = max(1, int(max_concurrency))
self._limit = self.max_concurrency
self._in_flight = 0
self._cooldown_until = 0.0
self._min_cooldown = max(0.0, float(min_cooldown))
self._default_cooldown = max(self._min_cooldown, float(default_cooldown))
self._max_cooldown = max(self._default_cooldown, float(max_cooldown))
self._recovery_successes = max(
1, int(recovery_successes or max(8, self.max_concurrency * 2))
)
self._successes = 0
self._condition = asyncio.Condition()
@property
def limit(self) -> int:
"""Current concurrency limit (may be below ``max_concurrency``)."""
return self._limit
async def acquire(self) -> None:
"""Wait for the shared cooldown and a concurrency slot."""
async with self._condition:
while True:
remaining = self._cooldown_until - time.monotonic()
if remaining > 0:
try:
await asyncio.wait_for(
self._condition.wait(), timeout=remaining
)
except asyncio.TimeoutError:
pass
continue
if self._in_flight < self._limit:
self._in_flight += 1
return
await self._condition.wait()
async def release(self) -> None:
async with self._condition:
self._in_flight = max(0, self._in_flight - 1)
self._condition.notify()
async def report_rate_limit(self, retry_after: float | None = None) -> float:
"""Arm the shared cooldown and shrink the concurrency limit.
Returns the cooldown actually applied, after clamping.
"""
async with self._condition:
requested = (
retry_after
if retry_after is not None and retry_after > 0
else self._default_cooldown
)
cooldown = min(
max(self._min_cooldown, float(requested)), self._max_cooldown
)
self._cooldown_until = max(
self._cooldown_until, time.monotonic() + cooldown
)
self._limit = max(1, self._limit // 2)
self._successes = 0
self._condition.notify_all()
return cooldown
async def report_success(self) -> None:
"""Slowly restore concurrency after a run of successful requests."""
async with self._condition:
if self._limit >= self.max_concurrency:
return
self._successes += 1
if self._successes >= self._recovery_successes:
self._limit += 1
self._successes = 0
self._condition.notify()
def reset(self) -> None:
"""Test seam: restore the initial state."""
self._limit = self.max_concurrency
self._in_flight = 0
self._cooldown_until = 0.0
self._successes = 0
class ExampleImagesImportError(RuntimeError):
"""Base error for example image import operations."""
@@ -70,16 +311,64 @@ class ExampleImagesProcessor:
return ''.join(random.choice(chars) for _ in range(length))
@staticmethod
def get_civitai_optimized_url(media_url):
"""Convert Civitai media URL (image or video) to its optimized version"""
base_pattern = r'(https://image\.civitai\.com/[^/]+/[^/]+)'
match = re.match(base_pattern, media_url)
if match:
base_url = match.group(1)
return f"{base_url}/optimized=true"
return media_url
def _is_video_media(media_url: str, media_type_hint: str | None = None) -> bool:
if (media_type_hint or "").lower() == "video":
return True
path = urlparse(media_url or "").path
extension = posixpath.splitext(path)[1].lower()
return extension in SUPPORTED_MEDIA_EXTENSIONS["videos"]
@staticmethod
def _media_width(media_type: str | None) -> int:
"""Return the configured rendition width for example media."""
is_video = (media_type or "").lower() == "video"
key = "example_images_video_max_width" if is_video else "example_images_max_width"
default = 450 if is_video else 800
try:
width = int(get_settings_manager().get(key, default))
except (TypeError, ValueError):
return default
# The image cacher's ladder starts at 96px; a smaller value is a
# configuration mistake and would request an unusable thumbnail.
return width if width >= 96 else default
@staticmethod
def build_example_media_url(
media_url: str,
media_type_hint: str | None,
optimize: bool,
) -> str:
"""Return the URL to download for an example image or video.
Civitai URLs are rewritten to a resized rendition (images default to
800px, videos to 450px; both adjustable in settings). Genur signed
proxy URLs are mapped to the direct ``c.genur.art`` rendition, which
the broken proxy host never reaches. Other URLs are returned as-is.
"""
media_url = (media_url or "").strip()
if not media_url:
return media_url
media_type = (
"video"
if ExampleImagesProcessor._is_video_media(media_url, media_type_hint)
else "image"
)
if is_genur_media_url(media_url):
return resolve_genur_media_url(
media_url, media_type, size="small" if optimize else "large"
)
if not optimize:
return media_url
width = ExampleImagesProcessor._media_width(media_type)
return build_civitai_media_url(
media_url, media_type, width=width, optimize=True
) or media_url
@staticmethod
def _get_file_extension_from_content_or_headers(content, headers, fallback_url=None, media_type_hint=None):
@@ -120,234 +409,238 @@ class ExampleImagesProcessor:
@staticmethod
async def download_model_images(model_hash, model_name, model_images, model_dir, optimize, downloader):
"""Download images for a single model
"""Download example media for a single model.
Returns:
tuple: (success, is_stale_metadata) - whether download was successful, whether metadata is stale
tuple: (success, is_stale_metadata) - whether download was successful,
whether metadata is stale.
"""
model_success = True
for i, image in enumerate(model_images):
image_url = image.get('url')
if not image_url:
continue
if image.get('downloadFailed'):
logger.debug(
"Skipping example image %s for %s because it previously failed to download",
image_url,
model_name,
)
continue
outcome = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
model_images,
model_dir,
optimize,
downloader,
)
return outcome.success, outcome.is_stale
# Apply optimization for Civitai URLs if enabled
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}")
# Download using the unified downloader with headers
success, content, headers = await downloader.download_to_memory(
image_url,
use_auth=False, # Example images don't need auth
return_headers=True
)
if success:
# Determine file extension from content or headers
media_ext = ExampleImagesProcessor._get_file_extension_from_content_or_headers(
content, headers, original_url, image.get("type")
)
# Check if the detected file type is supported
is_image = media_ext in SUPPORTED_MEDIA_EXTENSIONS['images']
is_video = media_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
if not (is_image or is_video):
logger.debug(f"Skipping unsupported file type: {media_ext}")
continue
# Use 0-based indexing with the detected extension
save_filename = f"image_{i}{media_ext}"
save_path = os.path.join(model_dir, save_filename)
# Check if already downloaded
if os.path.exists(save_path):
logger.debug(f"File already exists: {save_path}")
continue
# Save the file
with open(save_path, 'wb') as f:
f.write(content)
elif ExampleImagesProcessor._is_not_found_error(content):
error_msg = f"Failed to download file: {image_url}, status code: 404 - Model metadata might be stale"
logger.warning(error_msg)
model_success = False # Mark the model as failed due to 404 error
# Return early to trigger metadata refresh attempt
return False, True # (success, is_metadata_stale)
else:
error_msg = f"Failed to download file: {image_url}, error: {content}"
logger.warning(error_msg)
model_success = False # Mark the model as failed
except Exception as e:
error_msg = f"Error downloading file {image_url}: {str(e)}"
logger.error(error_msg)
model_success = False # Mark the model as failed
return model_success, False # (success, is_metadata_stale)
@staticmethod
def _extract_retry_after(error_message: str) -> int:
def _extract_retry_after(error_message: str) -> float:
if not error_message:
return 60
match = re.search(r"retry after (\d+)s", str(error_message))
return 30.0
match = re.search(r"retry after ([\d.]+)s", str(error_message))
if match:
return max(1, int(match.group(1)))
return 60
return max(0.0, float(match.group(1)))
return 30.0
@staticmethod
async def download_model_images_with_tracking(model_hash, model_name, model_images, model_dir, optimize, downloader):
model_success = True
failed_images = []
rate_limited_images = []
any_successful_download = False
def _is_transient_failure(message: str) -> bool:
"""Return True for transport/offline/5xx errors that deserve a retry.
for i, image in enumerate(model_images):
image_url = image.get('url')
if not image_url:
A 404 is a statement about the asset and is handled separately; an
offline/shared-cooldown message describes the network path, so
blacklisting the image for it (the old behavior) is wrong.
"""
if not message:
return False
text = str(message)
if text in {OFFLINE_FRIENDLY_MESSAGE, OFFLINE_COOLDOWN_ERROR}:
return True
lower = text.lower()
if any(marker in lower for marker in _TRANSIENT_ERROR_MARKERS):
return True
match = re.search(r"status (\d{3})", text, re.IGNORECASE)
if match:
status = int(match.group(1))
return status in {408, 425, 429} or 500 <= status < 600
return False
@staticmethod
async def download_model_images_with_tracking(
model_hash,
model_name,
model_images,
model_dir,
optimize,
downloader,
limiter: MediaDownloadLimiter | None = None,
retry_failed: bool = False,
) -> ImageDownloadOutcome:
"""Download all pending media for one model concurrently.
Each image is still one CDN request (Civitai has no batch-binary
endpoint), but the requests are pipelined through ``limiter`` instead
of being serialized. A 429 arms the limiter's shared cooldown, so one
rate-limited response pauses every worker rather than letting each
URL sleep and retry on its own.
``retry_failed`` re-attempts image-level ``downloadFailed`` entries
for Genur URLs. Genur images are the ones affected by the old
network-error classification bug; retrying them lets the caller clear
the flag only after a real download succeeds.
Returns:
ImageDownloadOutcome: per-URL lists and a ``success`` property.
URL lists contain the *original* metadata URLs so callers can
match them against the model's ``civitai.images`` entries.
"""
limiter = limiter or MediaDownloadLimiter()
outcome = ImageDownloadOutcome()
candidates: list[tuple[int, dict[str, Any], str, bool]] = []
for index, image in enumerate(model_images):
original_url = image.get("url")
if not original_url:
continue
if image.get('downloadFailed'):
# A file already on disk is authoritative; a stale failure flag
# must not force a redundant request.
if ExampleImagesProcessor._example_image_file_exists(
model_dir, index, image.get("type")
):
logger.debug("File already exists, skipping download for %s", original_url)
if (
image.get("downloadFailed")
and retry_failed
and is_genur_media_url(original_url)
):
outcome.recovered.append(original_url)
continue
if image.get("downloadFailed"):
if retry_failed and is_genur_media_url(original_url):
logger.debug(
"Retrying previously failed Genur image %s for %s",
original_url,
model_name,
)
candidates.append((index, image, original_url, True))
continue
outcome.skipped_failed += 1
logger.debug(
"Skipping example image %s for %s because it previously failed to download",
image_url,
original_url,
model_name,
)
continue
# Apply optimization for Civitai URLs if enabled
original_url = image_url
if optimize and 'civitai.com' in image_url:
image_url = ExampleImagesProcessor.get_civitai_optimized_url(image_url)
candidates.append((index, image, original_url, False))
# 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 download_one(
index: int,
image: dict[str, Any],
original_url: str,
was_flagged: bool,
) -> None:
target_url = ExampleImagesProcessor.build_example_media_url(
original_url, image.get("type"), optimize
)
logger.debug(
"Downloading media file %s for %s from %s",
index,
model_name,
target_url,
)
async def _attempt_download() -> tuple[bool, Any, Any]:
logger.debug("Downloading media file %s for %s", i, model_name)
return await downloader.download_to_memory(
image_url,
use_auth=False,
return_headers=True,
)
try:
success, content, headers = await _attempt_download()
for attempt in range(1, MAX_DOWNLOAD_ATTEMPTS + 1):
await limiter.acquire()
try:
success, content, headers = await downloader.download_to_memory(
target_url,
use_auth=False,
return_headers=True,
)
except Exception as e:
logger.error("Error downloading file %s: %s", target_url, e)
outcome.failed.append(original_url)
return
finally:
await limiter.release()
if success:
await limiter.report_success()
media_ext = ExampleImagesProcessor._get_file_extension_from_content_or_headers(
content, headers, original_url, image.get("type")
)
is_image = media_ext in SUPPORTED_MEDIA_EXTENSIONS['images']
is_video = media_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
is_image = media_ext in SUPPORTED_MEDIA_EXTENSIONS["images"]
is_video = media_ext in SUPPORTED_MEDIA_EXTENSIONS["videos"]
if not (is_image or is_video):
logger.debug("Skipping unsupported file type: %s", media_ext)
continue
save_filename = f"image_{i}{media_ext}"
save_path = os.path.join(model_dir, save_filename)
return
save_path = os.path.join(model_dir, f"image_{index}{media_ext}")
if os.path.exists(save_path):
logger.debug("File already exists: %s", save_path)
continue
return
with open(save_path, 'wb') as f:
f.write(content)
any_successful_download = True
try:
await asyncio.to_thread(_write_media_file, save_path, content)
except Exception as e:
logger.error("Failed to save example image %s: %s", save_path, e)
outcome.failed.append(original_url)
return
elif ExampleImagesProcessor._is_not_found_error(content):
error_msg = f"Failed to download file: {image_url}, status code: 404 - Model metadata might be stale"
logger.warning(error_msg)
model_success = False
failed_images.append(image_url)
return False, True, failed_images, rate_limited_images
outcome.downloaded.append(original_url)
if was_flagged:
outcome.recovered.append(original_url)
return
elif "Rate limited (429)" in str(content):
max_attempts = 3
for attempt in range(1, max_attempts + 1):
wait = ExampleImagesProcessor._extract_retry_after(str(content)) * (2 ** (attempt - 1))
logger.warning(
"Rate limited (429) for %s, retry %d/%d after %ds",
image_url, attempt, max_attempts, wait,
)
await asyncio.sleep(wait)
message = str(content)
if ExampleImagesProcessor._is_not_found_error(message):
logger.warning(
"Failed to download file: %s, status code: 404 - Model metadata might be stale",
target_url,
)
outcome.failed.append(original_url)
outcome.is_stale = True
return
success, content, headers = await _attempt_download()
if success:
media_ext = ExampleImagesProcessor._get_file_extension_from_content_or_headers(
content, headers, original_url, image.get("type")
)
is_image = media_ext in SUPPORTED_MEDIA_EXTENSIONS['images']
is_video = media_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
if "Rate limited (429)" in message:
retry_after = ExampleImagesProcessor._extract_retry_after(message)
cooldown = await limiter.report_rate_limit(retry_after)
logger.warning(
"Rate limited (429) for %s, shared cooldown %.1fs (attempt %d/%d)",
target_url,
cooldown,
attempt,
MAX_DOWNLOAD_ATTEMPTS,
)
if attempt >= MAX_DOWNLOAD_ATTEMPTS:
outcome.rate_limited.append(original_url)
return
# Loop back into acquire(): it parks until the shared
# cooldown expires, so the retry is not a herd.
continue
if not (is_image or is_video):
logger.debug("Skipping unsupported file type: %s", media_ext)
break
if ExampleImagesProcessor._is_transient_failure(message):
logger.warning(
"Transient failure for %s: %s (will retry next run)",
target_url,
content,
)
outcome.transient.append(original_url)
return
save_filename = f"image_{i}{media_ext}"
save_path = os.path.join(model_dir, save_filename)
if os.path.exists(save_path):
logger.debug("File already exists: %s", save_path)
break
logger.warning("Failed to download file: %s, error: %s", target_url, content)
outcome.failed.append(original_url)
return
with open(save_path, 'wb') as f:
f.write(content)
any_successful_download = True
break
elif "Rate limited (429)" in str(content):
continue
elif ExampleImagesProcessor._is_not_found_error(content):
logger.warning("Failed to download file: %s, status code: 404", image_url)
model_success = False
failed_images.append(image_url)
break
else:
logger.warning("Failed to download file: %s, error: %s", image_url, content)
model_success = False
failed_images.append(image_url)
break
else:
logger.warning(
"Giving up on %s after %d retries due to rate limiting",
image_url, max_attempts,
)
rate_limited_images.append(image_url)
model_success = False
else:
error_msg = f"Failed to download file: {image_url}, error: {content}"
logger.warning(error_msg)
model_success = False
failed_images.append(image_url)
except Exception as e:
error_msg = f"Error downloading file {image_url}: {str(e)}"
logger.error(error_msg)
model_success = False
failed_images.append(image_url)
if candidates:
await asyncio.gather(
*(
download_one(index, image, url, was_flagged)
for index, image, url, was_flagged in candidates
)
)
return outcome
return any_successful_download or model_success, False, failed_images, rate_limited_images
@staticmethod
async def process_local_examples(model_file_path, model_file_name, model_name, model_dir, optimize):
"""Process local example images
@@ -10,6 +10,7 @@ import pytest
from py.services.settings_manager import SettingsManager, get_settings_manager
from py.utils import example_images_download_manager as download_module
from py.utils.example_images_processor import ImageDownloadOutcome
class RecordingWebSocketManager:
@@ -174,7 +175,7 @@ async def test_pause_resume_blocks_processing(
await first_release.wait()
else:
second_call_started.set()
return True, False, [], []
return ImageDownloadOutcome()
async def fake_get_downloader():
class _Downloader:
@@ -298,7 +299,7 @@ async def test_legacy_folder_migrated_and_skipped(
async def fake_download_model_images(*_args, **_kwargs):
nonlocal download_called
download_called = True
return True, False, [], []
return ImageDownloadOutcome()
async def fake_get_downloader():
class _Downloader:
@@ -1017,3 +1018,274 @@ def test_unblock_failed_example_image_models_without_progress_file(
monkeypatch.setitem(settings_manager.settings, "active_library", "default")
assert download_module.unblock_failed_example_image_models(["a" * 64]) == 0
async def _run_single_outcome(
monkeypatch: pytest.MonkeyPatch,
tmp_path,
settings_manager,
outcome: ImageDownloadOutcome,
*,
force: bool = False,
):
model_hash = "a" * 64
model_path = tmp_path / "model.safetensors"
model = {
"sha256": model_hash,
"model_name": "Outcome Model",
"file_path": str(model_path),
"file_name": "model.safetensors",
"civitai": {"images": [{"url": "https://img.genur.art/sig/x/y"}]},
}
_patch_scanner(monkeypatch, StubScanner([model]))
monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path))
async def fake_process_local_examples(*_args, **_kwargs):
return False
captured: dict[str, Any] = {}
async def fake_download(*_args, **kwargs):
captured.update(kwargs)
return outcome
async def fake_get_downloader():
return object()
monkeypatch.setattr(
download_module.ExampleImagesProcessor,
"process_local_examples",
staticmethod(fake_process_local_examples),
)
monkeypatch.setattr(
download_module.ExampleImagesProcessor,
"download_model_images_with_tracking",
staticmethod(fake_download),
)
monkeypatch.setattr(download_module, "get_downloader", fake_get_downloader)
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
await manager.start_download(
{"model_types": ["lora"], "delay": 0, "force": force}
)
if manager._download_task is not None:
await asyncio.wait_for(manager._download_task, timeout=2)
return manager, model_hash, captured
@pytest.mark.asyncio
async def test_transient_failure_leaves_model_retryable(
monkeypatch: pytest.MonkeyPatch, tmp_path, settings_manager
) -> None:
outcome = ImageDownloadOutcome(transient=["https://img.genur.art/sig/x/y"])
manager, model_hash, _ = await _run_single_outcome(
monkeypatch, tmp_path, settings_manager, outcome
)
assert model_hash in manager._progress["retryable_models"]
assert model_hash not in manager._progress["processed_models"]
assert model_hash not in manager._progress["failed_models"]
@pytest.mark.asyncio
async def test_skipped_flagged_image_is_not_reported_as_recovered(
monkeypatch: pytest.MonkeyPatch, tmp_path, settings_manager
) -> None:
outcome = ImageDownloadOutcome(skipped_failed=1)
manager, model_hash, _ = await _run_single_outcome(
monkeypatch, tmp_path, settings_manager, outcome
)
assert model_hash in manager._progress["failed_models"]
assert model_hash not in manager._progress["processed_models"]
@pytest.mark.asyncio
async def test_force_passes_retry_failed_to_processor(
monkeypatch: pytest.MonkeyPatch, tmp_path, settings_manager
) -> None:
manager, model_hash, captured = await _run_single_outcome(
monkeypatch,
tmp_path,
settings_manager,
ImageDownloadOutcome(),
force=True,
)
assert captured["retry_failed"] is True
assert model_hash in manager._progress["processed_models"]
@pytest.mark.asyncio
async def test_clear_recovered_flags_only_touches_matching_images(
monkeypatch: pytest.MonkeyPatch, tmp_path
) -> None:
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
recovered_url = "https://img.genur.art/sig/x/y"
other_url = "https://image.civitai.com/x/y/original=true/1.jpeg"
model_data = {
"file_path": str(tmp_path / "model.safetensors"),
"civitai": {
"images": [
{
"url": recovered_url,
"type": "image",
"downloadFailed": True,
"downloadError": "not_found",
},
{
"url": other_url,
"type": "image",
"downloadFailed": True,
"downloadError": "not_found",
},
]
},
}
async def fake_get_updated_model(*_args, **_kwargs):
return model_data
saved: list[dict[str, Any]] = []
cached: list[dict[str, Any]] = []
async def fake_save(_path, payload):
saved.append(payload)
return True
async def fake_update_cache(_scanner, _path, payload):
cached.append(payload)
return True
monkeypatch.setattr(
download_module.MetadataUpdater,
"get_updated_model",
staticmethod(fake_get_updated_model),
)
monkeypatch.setattr(
download_module.MetadataManager, "save_metadata", staticmethod(fake_save)
)
monkeypatch.setattr(
download_module, "update_cache_from_metadata", fake_update_cache
)
await manager._clear_failed_images_from_metadata(
"hash", "Model", {recovered_url}, object()
)
assert "downloadFailed" not in model_data["civitai"]["images"][0]
assert "downloadError" not in model_data["civitai"]["images"][0]
assert model_data["civitai"]["images"][1]["downloadFailed"] is True
assert saved and "folder" not in saved[0]
assert cached
@pytest.mark.asyncio
async def test_force_bypasses_populated_folder_guard(
monkeypatch: pytest.MonkeyPatch, tmp_path, settings_manager
) -> None:
model_hash = "b" * 64
model = {
"sha256": model_hash,
"model_name": "Partial Model",
"file_path": str(tmp_path / "partial.safetensors"),
"file_name": "partial.safetensors",
"civitai": {"images": [{"url": "https://img.genur.art/sig/x/y"}]},
}
scanner = StubScanner([model])
_patch_scanner(monkeypatch, scanner)
monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path))
monkeypatch.setitem(settings_manager.settings, "active_library", "comfyui")
model_dir = download_module.ExampleImagePathResolver.get_model_folder(
model_hash, "comfyui"
)
Path(model_dir).mkdir(parents=True, exist_ok=True)
Path(model_dir, "image_0.webp").write_bytes(b"existing")
async def fake_process_local_examples(*_args, **_kwargs):
return False
captured: dict[str, Any] = {}
async def fake_download(*_args, **kwargs):
captured.update(kwargs)
return ImageDownloadOutcome()
monkeypatch.setattr(
download_module.ExampleImagesProcessor,
"process_local_examples",
staticmethod(fake_process_local_examples),
)
monkeypatch.setattr(
download_module.ExampleImagesProcessor,
"download_model_images_with_tracking",
staticmethod(fake_download),
)
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
manager._progress["status"] = "running"
manager._progress["processed_models"].add(model_hash)
await manager._process_model(
"lora",
model,
scanner,
str(tmp_path),
False,
object(),
"comfyui",
force=True,
explicit_targets=False,
limiter=None,
)
assert captured["retry_failed"] is True
@pytest.mark.asyncio
async def test_clear_recovered_flags_roundtrips_through_sidecar(
tmp_path, settings_manager
) -> None:
model_path = tmp_path / "roundtrip.safetensors"
model_path.write_bytes(b"model")
model_hash = "c" * 64
recovered_url = "https://img.genur.art/sig/x/y"
model_data = {
"file_name": "roundtrip",
"model_name": "Roundtrip Model",
"file_path": str(model_path),
"sha256": model_hash,
"size": 1,
"modified": 0,
"base_model": "SDXL",
"preview_url": "",
"preview_nsfw_level": 0,
"from_civitai": True,
"civitai": {
"images": [
{
"url": recovered_url,
"type": "image",
"downloadFailed": True,
"downloadError": "not_found",
}
]
},
}
assert await download_module.MetadataManager.save_metadata(
str(model_path), model_data
)
scanner = StubScanner([{**model_data, "folder": ""}])
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
await manager._clear_failed_images_from_metadata(
model_hash, "Roundtrip Model", {recovered_url}, scanner
)
loaded, _ = await download_module.MetadataManager.load_metadata(str(model_path))
assert loaded is not None
images = loaded.to_dict()["civitai"]["images"]
assert images[0]["url"] == recovered_url
assert "downloadFailed" not in images[0]
assert "downloadError" not in images[0]
+3 -2
View File
@@ -5,6 +5,7 @@ from pathlib import Path
from typing import Any
from py.services.settings_manager import get_settings_manager
from py.utils import example_images_download_manager as download_module
from py.utils.example_images_processor import ImageDownloadOutcome
class RecordingWebSocketManager:
def __init__(self) -> None:
@@ -73,12 +74,12 @@ async def test_reprocessing_triggered_when_folder_missing(monkeypatch, tmp_path)
monkeypatch.setattr(download_module.ExampleImagesProcessor, "process_local_examples", fake_process_local_examples)
async def fake_download_model_images(*args):
async def fake_download_model_images(*args, **_kwargs):
# Create the directory so it's "fixed"
model_dir = args[3]
Path(model_dir).mkdir(parents=True, exist_ok=True)
(Path(model_dir) / "image_0.png").write_text("fixed")
return True, False, [], []
return ImageDownloadOutcome()
monkeypatch.setattr(download_module.ExampleImagesProcessor, "download_model_images_with_tracking", fake_download_model_images)
+70 -1
View File
@@ -2,7 +2,7 @@
import pytest
from py.utils.civitai_utils import rewrite_preview_url
from py.utils.civitai_utils import build_civitai_media_url, rewrite_preview_url
class TestRewritePreviewUrl:
@@ -148,3 +148,72 @@ class TestRewritePreviewUrl:
for url in test_cases:
result, was_rewritten = rewrite_preview_url(url, "image")
assert was_rewritten is True, f"Failed for URL: {url}"
def test_width_parameter_overrides_default(self):
url = "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg"
result, was_rewritten = rewrite_preview_url(url, "image", width=800)
assert was_rewritten is True
assert result == "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true/1.jpeg"
def test_width_parameter_applies_to_video_transcode(self):
url = "https://image.civitai.com/xG1nk/uuid/original=true/1.mp4"
result, was_rewritten = rewrite_preview_url(url, "video", width=450)
assert was_rewritten is True
assert (
result
== "https://image.civitai.com/xG1nk/uuid/transcode=true,width=450,optimized=true/1.mp4"
)
class TestBuildCivitaiMediaUrl:
"""Tests for the example-media download URL builder."""
def test_rewrites_original_true_and_keeps_filename(self):
url = "https://image.civitai.com/xG1nk/uuid/original=true/74821598.jpeg"
assert (
build_civitai_media_url(url, "image", width=800)
== "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true/74821598.jpeg"
)
def test_rewrites_videos_with_transcode(self):
url = "https://image.civitai.com/xG1nk/uuid/original=true/75044257.mp4"
assert (
build_civitai_media_url(url, "video", width=450)
== "https://image.civitai.com/xG1nk/uuid/transcode=true,width=450,optimized=true/75044257.mp4"
)
def test_replaces_legacy_params_segment(self):
url = "https://image.civitai.com/xG1nk/uuid/optimized=true"
assert (
build_civitai_media_url(url, "image", width=800)
== "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true"
)
def test_inserts_params_before_filename_without_params(self):
url = "https://image.civitai.com/xG1nk/uuid/74821598.jpeg"
assert (
build_civitai_media_url(url, "image", width=800)
== "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true/74821598.jpeg"
)
def test_appends_params_when_there_is_no_filename(self):
url = "https://image.civitai.com/xG1nk/uuid"
assert (
build_civitai_media_url(url, "image", width=800)
== "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true"
)
def test_handles_image_b2_subdomain(self):
url = "https://image-b2.civitai.com/file/cache/original=true/sample.png"
assert (
build_civitai_media_url(url, "image", width=800)
== "https://image-b2.civitai.com/file/cache/width=800,optimized=true/sample.png"
)
def test_leaves_non_civitai_urls_alone(self):
url = "https://example.com/original=true/sample.png"
assert build_civitai_media_url(url, "image", width=800) == url
def test_optimize_false_returns_source_url(self):
url = "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg"
assert build_civitai_media_url(url, "image", optimize=False) == url
@@ -1,7 +1,9 @@
from __future__ import annotations
import asyncio
import json
import os
import time
from pathlib import Path
from types import SimpleNamespace
from typing import Any, Dict, Generator, Tuple
@@ -134,14 +136,14 @@ async def test_download_model_images_with_tracking_skips_existing_files(tmp_path
return True, b"\xff\xd8\xff" + b"data", {}
downloader = RecordingDownloader()
success, is_stale, failed, rate_limited = await proc.download_model_images_with_tracking(
outcome = 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 == []
assert outcome.success is True
assert outcome.is_stale is False
assert outcome.failed == []
assert outcome.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]
@@ -464,3 +466,276 @@ async def test_resolve_hash_for_file_path_unknown_file_returns_empty(
)
== ""
)
@pytest.mark.asyncio
async def test_download_model_images_uses_bounded_renditions(
monkeypatch: pytest.MonkeyPatch, tmp_path
) -> None:
"""Images use the configured width and videos are transcoded at their cap."""
settings_manager = get_settings_manager()
settings_manager.settings["example_images_max_width"] = 800
settings_manager.settings["example_images_video_max_width"] = 450
images = [
{
"url": "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg",
"type": "image",
},
{
"url": "https://image.civitai.com/xG1nk/uuid/original=true/2.mp4",
"type": "video",
},
]
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)
if url.endswith(".mp4"):
return True, b"\x00\x00\x00\x18ftypmp42data", {"content-type": "video/mp4"}
return True, b"\xff\xd8\xff" + b"data", {"content-type": "image/jpeg"}
downloader = RecordingDownloader()
limiter = processor_module.MediaDownloadLimiter(max_concurrency=2)
outcome = await (
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
"hash", "model", images, str(tmp_path), True, downloader, limiter=limiter
)
)
assert outcome.success is True
assert outcome.is_stale is False
assert outcome.failed == []
assert outcome.rate_limited == []
assert any("width=800,optimized=true/1.jpeg" in url for url in downloader.calls)
assert any(
"transcode=true,width=450,optimized=true/2.mp4" in url
for url in downloader.calls
)
assert Path(tmp_path, "image_0.jpg").exists()
assert Path(tmp_path, "image_1.mp4").exists()
@pytest.mark.asyncio
async def test_download_model_images_runs_requests_concurrently(tmp_path) -> None:
"""The limiter's pool overlaps network waits instead of serializing them."""
images = [
{
"url": f"https://image.civitai.com/xG1nk/uuid/original=true/{i}.jpeg",
"type": "image",
}
for i in range(6)
]
class TrackingDownloader:
def __init__(self) -> None:
self.active = 0
self.max_active = 0
async def download_to_memory(self, url, use_auth=False, return_headers=False):
self.active += 1
self.max_active = max(self.max_active, self.active)
try:
await asyncio.sleep(0.02)
return True, b"\xff\xd8\xff" + b"data", {"content-type": "image/jpeg"}
finally:
self.active -= 1
downloader = TrackingDownloader()
limiter = processor_module.MediaDownloadLimiter(max_concurrency=3)
outcome = await (
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
"hash", "model", images, str(tmp_path), False, downloader, limiter=limiter
)
)
assert outcome.success is True
assert outcome.failed == []
assert outcome.rate_limited == []
assert 1 < downloader.max_active <= 3
assert len(list(Path(tmp_path).glob("image_*.jpg"))) == 6
@pytest.mark.asyncio
async def test_download_model_images_shares_429_cooldown(tmp_path) -> None:
"""A 429 retries after the shared cooldown, not a per-task 60s sleep."""
images = [
{
"url": f"https://image.civitai.com/xG1nk/uuid/original=true/{i}.jpeg",
"type": "image",
}
for i in range(2)
]
class RateLimitedThenOk:
def __init__(self) -> None:
self.attempts: dict[str, int] = {}
async def download_to_memory(self, url, use_auth=False, return_headers=False):
attempt = self.attempts.get(url, 0) + 1
self.attempts[url] = attempt
if attempt == 1:
return False, "Rate limited (429), retry after 0.01s", None
return True, b"\xff\xd8\xff" + b"data", {"content-type": "image/jpeg"}
downloader = RateLimitedThenOk()
limiter = processor_module.MediaDownloadLimiter(
max_concurrency=2, min_cooldown=0.01, default_cooldown=0.01
)
started = time.monotonic()
outcome = await (
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
"hash", "model", images, str(tmp_path), False, downloader, limiter=limiter
)
)
elapsed = time.monotonic() - started
assert outcome.success is True
assert outcome.failed == []
assert outcome.rate_limited == []
assert all(attempts == 2 for attempts in downloader.attempts.values())
assert elapsed < 1.0
@pytest.mark.asyncio
async def test_download_model_images_marks_404_stale_with_original_url(tmp_path) -> None:
image = {
"url": "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg",
"type": "image",
}
class NotFoundDownloader:
async def download_to_memory(self, url, use_auth=False, return_headers=False):
return False, "File not found", None
outcome = await (
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
"hash", "model", [image], str(tmp_path), True, NotFoundDownloader()
)
)
assert outcome.success is False
assert outcome.is_stale is True
assert outcome.failed == [image["url"]]
assert outcome.rate_limited == []
def test_resolve_genur_signed_url_uses_direct_small_rendition() -> None:
signed = (
"https://img.genur.art/sig/width:450/quality:85/"
"aHR0cHM6Ly9jLmdlbnVyLmFydC83MjIwZGMwYy05N2YyLTQ0YzEtYWNkZi00ZDBhMzkxY2ZjNmQ="
)
assert (
processor_module.resolve_genur_media_url(signed)
== "https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d_small.webp"
)
def test_resolve_genur_direct_url_adds_small_suffix() -> None:
assert (
processor_module.resolve_genur_media_url(
"https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d"
)
== "https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d_small.webp"
)
existing = "https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d_large.webp"
assert processor_module.resolve_genur_media_url(existing) == existing
def test_resolve_genur_leaves_other_urls_alone() -> None:
url = "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg"
assert processor_module.resolve_genur_media_url(url) == url
@pytest.mark.asyncio
async def test_download_model_images_classifies_network_errors_as_transient(tmp_path) -> None:
image = {
"url": "https://img.genur.art/sig/width:450/quality:85/"
"aHR0cHM6Ly9jLmdlbnVyLmFydC83MjIwZGMwYy05N2YyLTQ0YzEtYWNkZi00ZDBhMzkxY2ZjNmQ=",
"type": "image",
}
class OfflineDownloader:
async def download_to_memory(self, url, use_auth=False, return_headers=False):
return False, "Network offline, will retry automatically later", None
outcome = await (
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
"hash", "model", [image], str(tmp_path), True, OfflineDownloader()
)
)
assert outcome.success is False
assert outcome.transient == [image["url"]]
assert outcome.failed == []
assert outcome.rate_limited == []
@pytest.mark.asyncio
async def test_download_model_images_retries_flagged_genur_image(tmp_path) -> None:
original = (
"https://img.genur.art/sig/width:450/quality:85/"
"aHR0cHM6Ly9jLmdlbnVyLmFydC83MjIwZGMwYy05N2YyLTQ0YzEtYWNkZi00ZDBhMzkxY2ZjNmQ="
)
image = {
"url": original,
"type": "image",
"downloadFailed": True,
"downloadError": "not_found",
}
class GenurDownloader:
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"RIFF\x00\x00\x00\x00WEBPdata", {"content-type": "image/webp"}
downloader = GenurDownloader()
outcome = await (
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
"hash",
"model",
[image],
str(tmp_path),
True,
downloader,
retry_failed=True,
)
)
assert outcome.success is True
assert outcome.recovered == [original]
assert outcome.downloaded == [original]
assert downloader.calls == [
"https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d_small.webp"
]
assert Path(tmp_path, "image_0.webp").exists()
@pytest.mark.asyncio
async def test_download_model_images_does_not_retry_flag_without_retry_option(tmp_path) -> None:
image = {
"url": "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg",
"type": "image",
"downloadFailed": True,
"downloadError": "not_found",
}
class ShouldNotRun:
async def download_to_memory(self, url, use_auth=False, return_headers=False):
raise AssertionError("flagged image must not be requested")
outcome = await (
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
"hash", "model", [image], str(tmp_path), True, ShouldNotRun()
)
)
assert outcome.success is False
assert outcome.skipped_failed == 1
assert outcome.downloaded == []
@@ -0,0 +1,76 @@
from __future__ import annotations
import asyncio
import time
import pytest
from py.utils.example_images_processor import MediaDownloadLimiter
@pytest.mark.asyncio
async def test_limiter_bounds_in_flight_requests() -> None:
limiter = MediaDownloadLimiter(max_concurrency=2)
await limiter.acquire()
await limiter.acquire()
waiting = asyncio.create_task(limiter.acquire())
await asyncio.sleep(0)
assert waiting.done() is False
await limiter.release()
await asyncio.wait_for(waiting, timeout=1)
await limiter.release()
await limiter.release()
@pytest.mark.asyncio
async def test_limiter_parks_every_worker_during_shared_cooldown() -> None:
limiter = MediaDownloadLimiter(
max_concurrency=2, min_cooldown=0.0, default_cooldown=0.01
)
await limiter.report_rate_limit(0.05)
started = time.monotonic()
await limiter.acquire()
elapsed = time.monotonic() - started
assert elapsed >= 0.04
await limiter.release()
@pytest.mark.asyncio
async def test_limiter_halves_on_429_and_recovers_one_slot_at_a_time() -> None:
limiter = MediaDownloadLimiter(
max_concurrency=4,
min_cooldown=0.0,
default_cooldown=0.01,
recovery_successes=2,
)
await limiter.report_rate_limit(0.01)
assert limiter.limit == 2
await limiter.report_success()
assert limiter.limit == 2
await limiter.report_success()
assert limiter.limit == 3
await limiter.report_success()
await limiter.report_success()
assert limiter.limit == 4
def test_limiter_reset_restores_initial_limit() -> None:
limiter = MediaDownloadLimiter(max_concurrency=3)
limiter._limit = 1
limiter._in_flight = 1
limiter._cooldown_until = time.monotonic() + 30
limiter.reset()
assert limiter.limit == 3
assert limiter._in_flight == 0
assert limiter._cooldown_until == 0.0