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
+26
View File
@@ -1,5 +1,6 @@
from pathlib import Path
from gpu_rent.civitai import pick_preview_image
from gpu_rent.remote.civitai_fetch import fmt_bytes, progress_line, should_skip
@@ -54,3 +55,28 @@ def test_progress_line_unknown_total():
line = progress_line("file", 1024 * 1024, None, 100_000)
assert "%" not in line
assert "/s" in line
def test_pick_preview_image_first_jpeg():
version = {
"images": [
{"type": "video", "url": "https://example.com/x.mp4"},
{"type": "image", "url": "https://image.civitai.com/x/width=450/abc.jpeg"},
]
}
got = pick_preview_image(version)
assert got is not None
url, suffix = got
assert "abc.jpeg" in url
assert suffix == ".preview.jpg"
def test_pick_preview_image_png_and_default():
assert pick_preview_image({"images": [{"url": "https://cdn.example/a.png"}]})[1] == (
".preview.png"
)
assert pick_preview_image({"images": [{"url": "https://cdn.example/hash/width=450"}]})[
1
] == ".preview.jpg"
assert pick_preview_image({"images": []}) is None
assert pick_preview_image({}) is None