Add file size conversion and enhance Civitai job processing
- Introduced the `file_size_bytes` function to convert Civitai model sizes from kilobytes to bytes, improving data handling. - Updated the `seed_civitai` function to include file size in job definitions, enhancing model processing efficiency. - Enhanced the `should_skip` function to utilize expected size for faster decision-making during job processing. - Added tests for new functionality, ensuring robustness in file size handling and job processing logic.
This commit is contained in:
@@ -1,7 +1,12 @@
|
||||
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
|
||||
from gpu_rent.civitai import file_size_bytes, pick_preview_image
|
||||
from gpu_rent.remote.civitai_fetch import (
|
||||
fmt_bytes,
|
||||
progress_line,
|
||||
should_skip,
|
||||
write_sha_sidecar,
|
||||
)
|
||||
|
||||
|
||||
def test_should_skip_missing(tmp_path: Path):
|
||||
@@ -17,21 +22,52 @@ def test_should_skip_exists_no_sha(tmp_path: Path):
|
||||
assert "уже есть" in reason
|
||||
|
||||
|
||||
def test_should_skip_sha_match(tmp_path: Path):
|
||||
def test_should_skip_size_ok_fast(tmp_path: Path):
|
||||
"""Re-up must not hash multi-GB weights when size matches."""
|
||||
dest = tmp_path / "big.safetensors"
|
||||
data = b"x" * 10_000
|
||||
dest.write_bytes(data)
|
||||
expect = "deadbeef" * 8
|
||||
skip, reason = should_skip(dest, expect, expect_size=len(data))
|
||||
assert skip is True
|
||||
assert "size ok" in reason
|
||||
# sidecar written for next run
|
||||
assert (tmp_path / "big.safetensors.sha256").is_file()
|
||||
|
||||
|
||||
def test_should_skip_sha_sidecar(tmp_path: Path):
|
||||
dest = tmp_path / "a.safetensors"
|
||||
dest.write_bytes(b"weights")
|
||||
expect = "abcd" * 16
|
||||
write_sha_sidecar(dest, expect)
|
||||
skip, reason = should_skip(dest, expect, expect_size=7)
|
||||
assert skip is True
|
||||
assert "sidecar" in reason
|
||||
|
||||
|
||||
def test_should_skip_sha_match_slow(tmp_path: Path):
|
||||
dest = tmp_path / "a.safetensors"
|
||||
dest.write_bytes(b"weights")
|
||||
import hashlib
|
||||
|
||||
digest = hashlib.sha256(b"weights").hexdigest()
|
||||
skip, reason = should_skip(dest, digest)
|
||||
skip, reason = should_skip(dest, digest, verify_sha=True)
|
||||
assert skip is True
|
||||
assert "sha ok" in reason
|
||||
|
||||
|
||||
def test_should_not_skip_bad_size(tmp_path: Path):
|
||||
dest = tmp_path / "a.safetensors"
|
||||
dest.write_bytes(b"weights")
|
||||
skip, reason = should_skip(dest, "deadbeef" * 8, expect_size=99999)
|
||||
assert skip is False
|
||||
assert "размер" in reason
|
||||
|
||||
|
||||
def test_should_not_skip_bad_sha(tmp_path: Path):
|
||||
dest = tmp_path / "a.safetensors"
|
||||
dest.write_bytes(b"weights")
|
||||
skip, reason = should_skip(dest, "deadbeef" * 8)
|
||||
skip, reason = should_skip(dest, "deadbeef" * 8, verify_sha=True)
|
||||
assert skip is False
|
||||
assert "перекачиваю" in reason
|
||||
|
||||
@@ -80,3 +116,9 @@ def test_pick_preview_image_png_and_default():
|
||||
] == ".preview.jpg"
|
||||
assert pick_preview_image({"images": []}) is None
|
||||
assert pick_preview_image({}) is None
|
||||
|
||||
|
||||
def test_file_size_bytes():
|
||||
assert file_size_bytes({"sizeKB": 1.0}) == 1024
|
||||
assert file_size_bytes({"sizeKB": 6775430.35}) == int(round(6775430.35 * 1024))
|
||||
assert file_size_bytes({}) is None
|
||||
|
||||
Reference in New Issue
Block a user