diff --git a/modules/civitai/download_civitai.py b/modules/civitai/download_civitai.py index c7733711e..d784c6234 100644 --- a/modules/civitai/download_civitai.py +++ b/modules/civitai/download_civitai.py @@ -327,6 +327,9 @@ class DownloadManager: if code == 200: log.info(f'CivitAI preview saved: id={item.id}') break + if code == 304 and backfill_preview_parameters(final_file, img.url, img.meta): + log.info(f'CivitAI preview backfilled: id={item.id}') + break except Exception as e: log.warning(f'CivitAI preview fetch failed: id={item.id} {e}') @@ -442,6 +445,64 @@ def embed_preview_parameters(preview_file: str, parameters: str) -> bool: return False +def preview_has_parameters(preview_file: str) -> bool: + """Check whether `preview_file` already carries an embedded parameters string. + + Mirrors the lookup `modules/image/metadata.py:read_info_from_image` + performs: PNG `image.info["parameters"]` or JPEG/WEBP EXIF `UserComment`. + """ + if not preview_file or not os.path.exists(preview_file): + return False + try: + from PIL import Image + img = Image.open(preview_file) + try: + info = img.info or {} + if info.get('parameters'): + return True + if info.get('UserComment'): + return True + exif = info.get('exif') + if exif: + import piexif + try: + parsed = piexif.load(exif) + if parsed.get('Exif', {}).get(piexif.ExifIFD.UserComment): + return True + except Exception: + pass + finally: + img.close() + except Exception: + pass + return False + + +def backfill_preview_parameters(model_path: str, preview_url: str, meta: dict | None) -> bool: + """Embed Civitai meta into an existing preview file when it lacks parameters. + + Used by the rescan path to retroactively populate preview metadata + without re-downloading bytes. Returns True only if a new chunk was + written; False for no-op (file missing, no meta, already populated, + embed failed). + """ + if not meta or not model_path or not preview_url: + return False + ext = os.path.splitext(preview_url)[1] + preview_file = os.path.splitext(model_path)[0] + ext + if not os.path.exists(preview_file): + return False + if preview_has_parameters(preview_file): + return False + parameters = civitai_meta_to_parameters(meta) + if not parameters: + return False + if embed_preview_parameters(preview_file, parameters): + log.info(f'CivitAI preview backfill: file="{preview_file}"') + return True + return False + + # ---- Legacy compatibility functions ---- def download_civit_meta(model_path: str, model_id): diff --git a/modules/civitai/metadata_civitai.py b/modules/civitai/metadata_civitai.py index 7d7a8b1c1..edd519708 100644 --- a/modules/civitai/metadata_civitai.py +++ b/modules/civitai/metadata_civitai.py @@ -104,7 +104,7 @@ def civit_update_metadata(raw: bool = False): def atomic_civit_search_metadata(item, results): - from modules.civitai.download_civitai import download_civit_preview, download_civit_meta + from modules.civitai.download_civitai import download_civit_preview, download_civit_meta, backfill_preview_parameters if item is None: return try: @@ -141,6 +141,10 @@ def atomic_civit_search_metadata(item, results): results.append({**result, 'code': code, 'size': size, 'note': note, 'type': 'preview'}) found = True break + if code == 304 and backfill_preview_parameters(item['filename'], img.url, img.meta): + results.append({**result, 'code': 200, 'size': '', 'note': 'metadata embedded', 'type': 'preview'}) + found = True + break else: result['code'] = 404 time.sleep(0.25) # rate limiting @@ -162,6 +166,10 @@ def atomic_civit_search_metadata(item, results): results.append({**result, 'code': code, 'size': size, 'note': note, 'type': 'preview'}) found = True break + if code == 304 and backfill_preview_parameters(item['filename'], img.url, img.meta): + results.append({**result, 'code': 200, 'size': '', 'note': 'metadata embedded', 'type': 'preview'}) + found = True + break else: result['code'] = 404 time.sleep(0.25) # rate limiting