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