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