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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user