Add preview image handling for Civitai models and enhance job processing

- Introduced the `pick_preview_image` function to extract the first usable preview image URL and its suffix from model version data.
- Updated the `seed_civitai` function to include preview image URLs and destinations in job definitions, improving model handling.
- Implemented the `ensure_preview` function to download missing preview images during job processing, enhancing user experience.
- Added tests for `pick_preview_image` to ensure correct functionality across various scenarios, ensuring robustness in image handling.
This commit is contained in:
Leonid Pershin
2026-08-21 10:44:52 +03:00
parent dd0ffbf69b
commit 506369548c
5 changed files with 134 additions and 23 deletions
+26
View File
@@ -153,6 +153,30 @@ def download(url: str, dest: Path, token: str, *, label: str, auth_host: str) ->
partial.replace(dest)
def ensure_preview(job: dict, token: str, *, prefix: str, auth_host: str) -> None:
"""Download SwarmUI sidecar preview if missing (even when weight was skipped)."""
preview_url = str(job.get("preview_url") or "").strip()
preview_dest = str(job.get("preview_dest") or "").strip()
if not preview_url or not preview_dest:
return
dest = Path(preview_dest)
if dest.is_file() and dest.stat().st_size > 0:
return
try:
print(f"{prefix} preview: {dest.name}", flush=True)
# Image CDN usually needs no Bearer; still pass token for civitai hosts.
download(
preview_url,
dest,
token if auth_host == "civitai" else "",
label=f"{prefix} {dest.name}",
auth_host=auth_host if auth_host == "civitai" else "civitai",
)
print(f"{prefix} preview ok {dest.name} ({fmt_bytes(dest.stat().st_size)})")
except Exception as exc:
print(f"{prefix} WARN preview {dest.name}: {exc}")
def main() -> int:
civitai_token = (
TOKEN_PATH.read_text(encoding="utf-8").strip() if TOKEN_PATH.is_file() else ""
@@ -187,6 +211,7 @@ def main() -> int:
extra = dest.parent / extra_name
extra.parent.mkdir(parents=True, exist_ok=True)
extra.write_text(extra_text, encoding="utf-8")
ensure_preview(job, token, prefix=prefix, auth_host=auth_host)
continue
if reason:
print(f"{prefix} {reason}: {dest.name}")
@@ -212,6 +237,7 @@ def main() -> int:
for extra_name, extra_text in (job.get("sidecars") or {}).items():
extra = dest.parent / extra_name
extra.write_text(extra_text, encoding="utf-8")
ensure_preview(job, token, prefix=prefix, auth_host=auth_host)
print(f"{prefix} ok {dest.name} ({fmt_bytes(dest.stat().st_size)})")
except Exception as exc:
failed += 1