alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-07-13 20:21:16 -03:00
Code Issues Packages Projects Releases Wiki Activity

fix(download): retry on Civitai 429 rate limit instead of removing images from metadata

Browse Source 
When Civitai returns 429 (Too Many Requests) during example image
downloads, the previous behavior treated all failures identically and
permanently removed the corresponding images from model metadata —
making them impossible to retry.

This commit adds:
- 429 detection + Retry-After header parsing in download_to_memory
- Exponential backoff retry (up to 3 attempts) in
  download_model_images_with_tracking
- Separate tracking of rate-limited vs permanently failed URLs
- rate_limited_models progress tracking persisted to disk
- Rate-limited models are NOT added to failed_models/processed_models
  so they are automatically retried on subsequent download runs
- Force mode clears failed_models when rate-limited images exist
This commit is contained in:
Will Miao
2026-07-06 11:58:19 +08:00
parent 8bee8f4069
commit 00228deaaa
5 changed files with 213 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
py/services/downloader.py
View File

@@ -46,6 +46,30 @@ def is_ssl_cert_verify_error(exc: BaseException) -> bool:
return "CERTIFICATE_VERIFY_FAILED" in str(exc)
def _parse_retry_after(value: str) -> int:
"""Parse a Retry-After header value into seconds.
Supports both integer seconds and HTTP-date formats.
Returns a default of 60 seconds on invalid/missing input.
"""
if not value or not value.strip():
return 60
value = value.strip()
try:
return max(1, int(value))
except ValueError:
pass
try:
parsed = parsedate_to_datetime(value)
now = datetime.now().astimezone()
delta = (parsed - now).total_seconds()
return max(1, int(delta))
except (ValueError, OverflowError, OSError):
return 60
@dataclass(frozen=True)
class DownloadProgress:
"""Snapshot of a download transfer at a moment in time."""
@@ -911,6 +935,19 @@ class Downloader:
elif response.status == 404:
error_msg = "File not found"
return False, error_msg, None
elif response.status == 429:
raw_retry_after = response.headers.get("Retry-After")
retry_after = _parse_retry_after(raw_retry_after or "")
if raw_retry_after:
logger.warning(
"Rate limited (429) for %s, Retry-After: %ss", url, retry_after
)
else:
logger.warning(
"Rate limited (429) for %s, no Retry-After header; defaulting to %ss",
url, retry_after,
)
return False, f"Rate limited (429), retry after {retry_after}s", None
else:
error_msg = f"Download failed with status {response.status}"
return False, error_msg, None

105
py/utils/example_images_download_manager.py
View File

@@ -72,6 +72,7 @@ class _DownloadProgress(dict):
refreshed_models=set(),
failed_models=set(),
reprocessed_models=set(),
rate_limited_models=set(),
)
def snapshot(self) -> dict:
@@ -82,6 +83,7 @@ class _DownloadProgress(dict):
snapshot["refreshed_models"] = list(self["refreshed_models"])
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()))
return snapshot
@@ -153,13 +155,15 @@ class DownloadManager:
# Step 3: Load progress file (I/O operation, done outside lock)
processed_models = set()
failed_models = set()
rate_limited_models = set()
try:
progress_file, processed_models, failed_models = await self._load_progress_file(output_dir)
progress_file, processed_models, failed_models, rate_limited_models = await self._load_progress_file(output_dir)
logger.debug(
"Loaded previous progress, %s models already processed, %s models marked as failed",
"Loaded previous progress, %s models already processed, %s models marked as failed, %s models rate-limited",
len(processed_models),
len(failed_models),
len(rate_limited_models),
)
except Exception as e:
logger.error(f"Failed to load progress file: {e}")
@@ -175,6 +179,7 @@ class DownloadManager:
self._progress.reset()
self._progress["processed_models"] = processed_models
self._progress["failed_models"] = failed_models
self._progress["rate_limited_models"] = rate_limited_models
self._stop_requested = False
self._progress["status"] = "running"
self._progress["start_time"] = time.time()
@@ -242,8 +247,8 @@ class DownloadManager:
"status": self._progress.snapshot(),
}
async def _load_progress_file(self, output_dir: str) -> tuple[str, set, set]:
"""Load progress file from disk. Returns (progress_file_path, processed_models, failed_models).
async def _load_progress_file(self, output_dir: str) -> tuple[str, set, set, set]:
"""Load progress file from disk. Returns (progress_file_path, processed_models, failed_models, rate_limited_models).
This is a separate async method to allow running in executor to avoid blocking event loop.
"""
@@ -252,8 +257,12 @@ class DownloadManager:
None, self._load_progress_file_sync, output_dir
)
def _load_progress_file_sync(self, output_dir: str) -> tuple[str, set, set]:
"""Synchronous implementation of progress file loading."""
def _load_progress_file_sync(self, output_dir: str) -> tuple[str, set, set, set]:
"""Synchronous implementation of progress file loading.
Returns:
tuple: (progress_file_path, processed_models, failed_models, rate_limited_models)
"""
progress_file = os.path.join(output_dir, ".download_progress.json")
progress_source = progress_file
@@ -289,6 +298,7 @@ class DownloadManager:
processed_models = set()
failed_models = set()
rate_limited_models = set()
if os.path.exists(progress_source):
try:
@@ -296,11 +306,11 @@ class DownloadManager:
saved_progress = json.load(f)
processed_models = set(saved_progress.get("processed_models", []))
failed_models = set(saved_progress.get("failed_models", []))
rate_limited_models = set(saved_progress.get("rate_limited_models", []))
except Exception:
# Return empty sets on error
pass
return progress_file, processed_models, failed_models
return progress_file, processed_models, failed_models, rate_limited_models
def _load_progress_sets_sync(self, progress_file: str) -> tuple[set, set]:
"""Load only the processed and failed model sets from progress file.
@@ -732,11 +742,13 @@ class DownloadManager:
success,
is_stale,
failed_images,
rate_limited_images,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash, model_name, images, model_dir, optimize, downloader
)
failed_urls: Set[str] = set(failed_images)
rate_limited_urls: Set[str] = set(rate_limited_images)
# If metadata is stale, try to refresh it
if is_stale and model_hash not in self._progress["refreshed_models"]:
@@ -760,6 +772,7 @@ class DownloadManager:
success,
_,
additional_failed,
additional_rate_limited,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
@@ -770,29 +783,50 @@ class DownloadManager:
)
failed_urls.update(additional_failed)
rate_limited_urls.update(additional_rate_limited)
self._progress["refreshed_models"].add(model_hash)
if failed_urls:
# Separate permanent failures from rate-limited ones
permanent_failures = failed_urls - rate_limited_urls
if permanent_failures:
await self._remove_failed_images_from_metadata(
model_hash,
model_name,
model_dir,
failed_urls,
permanent_failures,
scanner,
)
if failed_urls:
if rate_limited_urls:
self._progress["rate_limited_models"].add(model_hash)
logger.warning(
"%d example images for %s are rate-limited (429), will retry next time",
len(rate_limited_urls),
model_name,
)
# Clear failed_models so non-force runs can retry
if force 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"
)
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)
logger.info(
"Removed %s failed example images for %s",
len(failed_urls),
len(permanent_failures),
model_name,
)
elif success:
self._progress["processed_models"].add(model_hash)
# Remove from failed_models if force mode enabled and model was previously failed
if force and model_hash in self._progress["failed_models"]:
self._progress["failed_models"].discard(model_hash)
logger.info(
@@ -850,6 +884,7 @@ class DownloadManager:
"processed_models": list(self._progress["processed_models"]),
"refreshed_models": list(self._progress["refreshed_models"]),
"failed_models": list(self._progress["failed_models"]),
"rate_limited_models": list(self._progress.get("rate_limited_models", set())),
"completed": self._progress["completed"],
"total": self._progress["total"],
"last_update": time.time(),
@@ -1155,11 +1190,13 @@ class DownloadManager:
success,
is_stale,
failed_images,
rate_limited_images,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash, model_name, images, model_dir, optimize, downloader
)
failed_urls: Set[str] = set(failed_images)
rate_limited_urls: Set[str] = set(rate_limited_images)
# If metadata is stale, try to refresh it
if is_stale and model_hash not in self._progress["refreshed_models"]:
@@ -1183,6 +1220,7 @@ class DownloadManager:
success,
_,
additional_failed_images,
additional_rate_limited,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
@@ -1192,21 +1230,35 @@ class DownloadManager:
downloader,
)
# Combine failed images from both attempts
failed_urls.update(additional_failed_images)
rate_limited_urls.update(additional_rate_limited)
self._progress["refreshed_models"].add(model_hash)
# For forced downloads, remove failed images from metadata
if failed_urls:
# Separate permanent failures from rate-limited ones
permanent_failures = failed_urls - rate_limited_urls
# Only remove permanently failed images from metadata
if permanent_failures:
await self._remove_failed_images_from_metadata(
model_hash, model_name, model_dir, failed_urls, scanner
model_hash, model_name, model_dir, permanent_failures, scanner
)
# Mark as processed
if (
success or failed_urls
): # Mark as processed if we successfully downloaded some images or removed failed ones
if rate_limited_urls:
self._progress["rate_limited_models"].add(model_hash)
logger.warning(
"%d example images for %s are rate-limited (429), will retry next time",
len(rate_limited_urls),
model_name,
)
# Mark as processed only when no rate-limited images remain
if rate_limited_urls:
pass
elif permanent_failures:
self._progress["processed_models"].add(model_hash)
self._progress["failed_models"].add(model_hash)
elif success:
self._progress["processed_models"].add(model_hash)
return True # Return True to indicate a remote download happened
@@ -1229,15 +1281,20 @@ class DownloadManager:
model_dir: str,
failed_images: Iterable[str],
scanner,
error_type: str = "not_found",
) -> None:
"""Mark failed images in model metadata so they won't be retried."""
"""Mark failed images in model metadata so they won't be retried.
Args:
error_type: Reason string stored in the image's ``downloadError`` field
(default ``"not_found"``).
"""
failed_set: Set[str] = {url for url in failed_images if url}
if not failed_set:
return
try:
# Get current model data
model_data = await MetadataUpdater.get_updated_model(model_hash, scanner)
if not model_data:
logger.warning(
@@ -1268,7 +1325,7 @@ class DownloadManager:
continue
image["downloadFailed"] = True
image.setdefault("downloadError", "not_found")
image.setdefault("downloadError", error_type)
logger.debug(
"Marked example image %s for %s as failed due to missing remote asset",
image_url,

131
py/utils/example_images_processor.py
View File

@@ -1,3 +1,4 @@
import asyncio
import logging
import os
import re
@@ -194,16 +195,22 @@ class ExampleImagesProcessor:
return model_success, False # (success, is_metadata_stale)
@staticmethod
def _extract_retry_after(error_message: str) -> int:
if not error_message:
return 60
match = re.search(r"retry after (\d+)s", str(error_message))
if match:
return max(1, int(match.group(1)))
return 60
@staticmethod
async def download_model_images_with_tracking(model_hash, model_name, model_images, model_dir, optimize, downloader):
"""Download images for a single model with tracking of failed image URLs
Returns:
tuple: (success, is_stale_metadata, failed_images) - whether download was successful, whether metadata is stale, list of failed image URLs
"""
model_success = True
failed_images = []
rate_limited_images = []
any_successful_download = False
for i, image in enumerate(model_images):
image_url = image.get('url')
if not image_url:
@@ -221,64 +228,110 @@ class ExampleImagesProcessor:
original_url = image_url
if optimize and 'civitai.com' in image_url:
image_url = ExampleImagesProcessor.get_civitai_optimized_url(image_url)
# 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(
async def _attempt_download() -> tuple:
logger.debug("Downloading media file %s for %s", i, model_name)
return await downloader.download_to_memory(
image_url,
use_auth=False, # Example images don't need auth
return_headers=True
use_auth=False,
return_headers=True,
)
try:
success, content, headers = await _attempt_download()
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}")
logger.debug("Skipping unsupported file type: %s", 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}")
logger.debug("File already exists: %s", save_path)
continue
# Save the file
with open(save_path, 'wb') as f:
f.write(content)
any_successful_download = True
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
failed_images.append(image_url) # Track failed URL
# Return early to trigger metadata refresh attempt
return False, True, failed_images # (success, is_metadata_stale, failed_images)
model_success = False
failed_images.append(image_url)
return False, True, failed_images, rate_limited_images
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)
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 not (is_image or is_video):
logger.debug("Skipping unsupported file type: %s", media_ext)
break
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
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 # Mark the model as failed
failed_images.append(image_url) # Track failed URL
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 # Mark the model as failed
failed_images.append(image_url) # Track failed URL
return model_success, False, failed_images # (success, is_metadata_stale, failed_images)
model_success = False
failed_images.append(image_url)
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):

4
tests/services/test_example_images_download_manager_async.py
View File

@@ -164,7 +164,7 @@ async def test_pause_resume_blocks_processing(
await first_release.wait()
else:
second_call_started.set()
return True, False, []
return True, False, [], []
async def fake_get_downloader():
class _Downloader:
@@ -288,7 +288,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 True, False, [], []
async def fake_get_downloader():
class _Downloader:

2
tests/services/test_issue_760_repro.py
View File

@@ -77,7 +77,7 @@ async def test_reprocessing_triggered_when_folder_missing(monkeypatch, tmp_path)
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 True, False, [], []
monkeypatch.setattr(download_module.ExampleImagesProcessor, "download_model_images_with_tracking", fake_download_model_images)