mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
feat(civitai): backfill preview meta on rescan
When download_civit_preview short-circuits at 304 (file exists), the rescan path now falls back to backfill_preview_parameters: checks whether the file already carries an embedded parameters chunk and, if not, embeds img.meta in place. Existing libraries get retroactive metadata via the existing "Search metadata" rescan, no model re-download required.
This commit is contained in:
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user