diff --git a/docs/plans/example-images-download-throughput.md b/docs/plans/example-images-download-throughput.md new file mode 100644 index 00000000..1908ec3a --- /dev/null +++ b/docs/plans/example-images-download-throughput.md @@ -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/.../` → `c.genur.art/_small.webp` + (`_large.webp` when optimization is disabled). Direct + `c.genur.art/` 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`. diff --git a/py/services/settings_manager.py b/py/services/settings_manager.py index b28c84a9..b4c3acae 100644 --- a/py/services/settings_manager.py +++ b/py/services/settings_manager.py @@ -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", diff --git a/py/utils/civitai_utils.py b/py/utils/civitai_utils.py index f40f7338..223b2754 100644 --- a/py/utils/civitai_utils.py +++ b/py/utils/civitai_utils.py @@ -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 + ``.../`` 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 ``...//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]: + # ``...//``: insert the rendition before the filename. + segments.insert(len(segments) - 1, params) + else: + # ``.../``: 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", diff --git a/py/utils/example_images_download_manager.py b/py/utils/example_images_download_manager.py index a083bdd9..2f681e99 100644 --- a/py/utils/example_images_download_manager.py +++ b/py/utils/example_images_download_manager.py @@ -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"): diff --git a/py/utils/example_images_processor.py b/py/utils/example_images_processor.py index dc354c10..705c9964 100644 --- a/py/utils/example_images_processor.py +++ b/py/utils/example_images_processor.py @@ -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/.../``) 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 ``_small.webp`` and ``_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 diff --git a/tests/services/test_example_images_download_manager_async.py b/tests/services/test_example_images_download_manager_async.py index c9738ca0..3a4895a8 100644 --- a/tests/services/test_example_images_download_manager_async.py +++ b/tests/services/test_example_images_download_manager_async.py @@ -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] diff --git a/tests/services/test_issue_760_repro.py b/tests/services/test_issue_760_repro.py index 874be1a0..4bb23b0f 100644 --- a/tests/services/test_issue_760_repro.py +++ b/tests/services/test_issue_760_repro.py @@ -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) diff --git a/tests/utils/test_civitai_utils_rewrite.py b/tests/utils/test_civitai_utils_rewrite.py index c65598d1..a735215e 100644 --- a/tests/utils/test_civitai_utils_rewrite.py +++ b/tests/utils/test_civitai_utils_rewrite.py @@ -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 diff --git a/tests/utils/test_example_images_processor_unit.py b/tests/utils/test_example_images_processor_unit.py index 02da9d21..eca15be5 100644 --- a/tests/utils/test_example_images_processor_unit.py +++ b/tests/utils/test_example_images_processor_unit.py @@ -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 == [] diff --git a/tests/utils/test_media_download_limiter.py b/tests/utils/test_media_download_limiter.py new file mode 100644 index 00000000..d910f510 --- /dev/null +++ b/tests/utils/test_media_download_limiter.py @@ -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