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