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
+33
View File
@@ -80,6 +80,39 @@ def pick_primary_file(version: dict) -> dict | None:
return None
def pick_preview_image(version: dict) -> tuple[str, str] | None:
"""First usable Civitai preview → (url, sidecar_suffix) for SwarmUI.
SwarmUI looks for ``stem.preview.jpg`` / ``.preview.png`` / ``.jpg`` etc.
next to the weight (T2IModelHandler AutoImageFormatSuffixes).
"""
images = version.get("images") or []
if not isinstance(images, list):
return None
for item in images:
if not isinstance(item, dict):
continue
kind = str(item.get("type") or "image").lower()
if kind and kind not in {"image", "img", ""}:
continue
url = str(item.get("url") or "").strip()
if not url.startswith("http"):
continue
path = url.split("?", 1)[0].lower()
if path.endswith(".png"):
suffix = ".preview.png"
elif path.endswith(".jpeg") or path.endswith(".jpg"):
suffix = ".preview.jpg"
elif path.endswith(".webp"):
# SwarmUI auto-format list has no .preview.webp — use .jpg name;
# most Civitai CDN URLs are jpeg without extension.
suffix = ".preview.jpg"
else:
suffix = ".preview.jpg"
return url, suffix
return None
def fetch_model_version(token: str, host: str, version_id: int, timeout: float = 30.0) -> tuple[str, dict]:
"""GET /api/v1/model-versions/{id}; one retry on the other Civitai host."""
first = _normalize_host(host)