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
+25 -13
View File
@@ -11,7 +11,7 @@ from pathlib import Path
import httpx
from gpu_rent.civitai import fetch_model_version, pick_primary_file
from gpu_rent.civitai import fetch_model_version, pick_preview_image, pick_primary_file
from gpu_rent.config import Config
from gpu_rent.errors import CloudError, GpuRentError
from gpu_rent.manifests import (
@@ -126,6 +126,14 @@ def seed_extensions(cfg: Config, host: str, log: Log, *, update: bool = True) ->
else:
log("extensions.yaml пуст — стоковый SwarmUI")
return False
if not repos and update:
if not all_repos:
log("extensions.yaml пуст — только update уже установленных на data (если есть)")
else:
log(
"extensions: yaml отфильтрован по requires "
f"(LLM_RUNTIME={runtime}) — только update установленных"
)
jobs = []
for repo in repos:
jobs.append(
@@ -462,18 +470,22 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None:
),
"tags": version.get("tags") or [],
}
jobs.append(
{
"dest": dest,
"url": _download_url(api_host, vid, info),
"sha256": sha,
"auth": "civitai",
"sidecars": {
f"{stem}.civitai.json": civitai_json,
f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2),
},
}
)
job: dict = {
"dest": dest,
"url": _download_url(api_host, vid, info),
"sha256": sha,
"auth": "civitai",
"sidecars": {
f"{stem}.civitai.json": civitai_json,
f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2),
},
}
preview = pick_preview_image(version)
if preview:
preview_url, preview_suffix = preview
job["preview_url"] = preview_url
job["preview_dest"] = f"{DATA}/Models/{folder}/{stem}{preview_suffix}"
jobs.append(job)
if not jobs:
if not cfg.civitai_api_token: