diff --git a/py/services/download_manager.py b/py/services/download_manager.py index fae40360..628555bd 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -18,7 +18,7 @@ from ..utils.constants import ( VALID_LORA_TYPES, ) from ..utils.civitai_utils import normalize_civitai_download_url, rewrite_preview_url -from ..utils.file_utils import calculate_sha256 +from ..utils.file_utils import calculate_sha256, calculate_autov3 from ..utils.preview_selection import resolve_mature_threshold, select_preview_media from ..utils.utils import sanitize_folder_name from ..utils.exif_utils import ExifUtils @@ -2160,6 +2160,10 @@ class DownloadManager: "error": f"Zip archive does not contain any supported model files ({supported_text})", } actual_file_paths = extracted_paths + # The archive entry's AutoV3 (if any) describes the zip itself, + # not the extracted models; clear it so per-file header + # resolution applies to every extracted model. + metadata.autov3 = None try: os.remove(save_path) except OSError as exc: @@ -2374,6 +2378,16 @@ class DownloadManager: sha256 = await calculate_sha256(file_path) if sha256: entry.sha256 = sha256.lower() + # AutoV3: the Civitai-reported value for the downloaded file (set + # by from_civitai_info) takes precedence. Only the un-checked + # state (None) triggers a header read; '' (checked-unavailable) + # is never re-read, honoring the three-state contract so rows + # marked at download time stay untouched by later passes. + if entry.autov3 is None: + autov3 = await asyncio.get_running_loop().run_in_executor( + None, calculate_autov3, file_path + ) + entry.autov3 = (autov3 or "").lower() entries.append(entry) return entries diff --git a/py/utils/models.py b/py/utils/models.py index 9e029661..587b243c 100644 --- a/py/utils/models.py +++ b/py/utils/models.py @@ -6,15 +6,29 @@ from .constants import INVALID_AUTOV3_EMPTY_HASH from .model_utils import determine_base_model +def normalize_autov3(value: Any) -> Optional[str]: + """Normalize a raw Civitai AutoV3 value to the canonical 12-char form. + + Returns the first 12 characters, lowercased, when ``value`` is a string + of at least 12 characters. The empty-string SHA256 placeholder + (``e3b0c44298fc``) is rejected — it is a repackaging-tool artifact, not a + real hash. + + Returns ``None`` when the value is unusable. + """ + if isinstance(value, str) and len(value) >= 12: + candidate = value[:12].lower() + if candidate != INVALID_AUTOV3_EMPTY_HASH: + return candidate + return None + + def autov3_from_civitai_files(civitai_data: Optional[Dict], sha256: str) -> Optional[str]: """Extract the AutoV3 hash from Civitai metadata for the matching file. Civitai versions can ship multiple files; the AutoV3 hash is only valid for the file whose ``hashes.SHA256`` equals the local model's sha256. - Matching is case-insensitive. The value is the first 12 characters of - Civitai's AutoV3 hash, lowercased. The empty-string SHA256 placeholder - (``e3b0c44298fc``) is rejected — it is a repackaging-tool artifact, not a - real hash. + Matching is case-insensitive. Returns ``None`` when no Civitai data, no matching file, or no usable AutoV3 hash is available. @@ -28,11 +42,7 @@ def autov3_from_civitai_files(civitai_data: Optional[Dict], sha256: str) -> Opti hashes = file_info.get("hashes") or {} file_sha = (hashes.get("SHA256") or "").lower() if file_sha and file_sha == target_sha: - auto_v3 = hashes.get("AutoV3") - if isinstance(auto_v3, str) and len(auto_v3) >= 12: - candidate = auto_v3[:12].lower() - if candidate != INVALID_AUTOV3_EMPTY_HASH: - return candidate + return normalize_autov3(hashes.get("AutoV3")) return None @@ -264,7 +274,8 @@ class LoraMetadata(BaseModelMetadata): civitai=version_info, tags=tags, modelDescription=description, - autov3=autov3_from_civitai_files(version_info, sha256_value), + # Direct read: the downloaded file IS file_info, no SHA256 matching. + autov3=normalize_autov3((file_info.get("hashes") or {}).get("AutoV3")), ) @@ -308,7 +319,8 @@ class CheckpointMetadata(BaseModelMetadata): sub_type=sub_type, tags=tags, modelDescription=description, - autov3=autov3_from_civitai_files(version_info, sha256_value), + # Direct read: the downloaded file IS file_info, no SHA256 matching. + autov3=normalize_autov3((file_info.get("hashes") or {}).get("AutoV3")), ) @@ -352,5 +364,6 @@ class EmbeddingMetadata(BaseModelMetadata): sub_type=sub_type, tags=tags, modelDescription=description, - autov3=autov3_from_civitai_files(version_info, sha256_value), + # Direct read: the downloaded file IS file_info, no SHA256 matching. + autov3=normalize_autov3((file_info.get("hashes") or {}).get("AutoV3")), ) diff --git a/tests/services/test_download_manager_autov3.py b/tests/services/test_download_manager_autov3.py new file mode 100644 index 00000000..e6cbdfcd --- /dev/null +++ b/tests/services/test_download_manager_autov3.py @@ -0,0 +1,100 @@ +"""AutoV3 resolution in the download completion path (``_build_metadata_entries``). + +Covers the Civitai-first / local-header fallback contract: a downloaded file +whose Civitai file_info reports no AutoV3 gets the embedded safetensors header +hash resolved right away (instead of waiting for the next startup's backfill), +and a file with no usable header is marked ``''`` (checked but unavailable). +""" + +import json +import struct + +import pytest + +from py.services.download_manager import DownloadManager +from py.utils.models import LoraMetadata + + +def _write_safetensors(path, metadata, payload=b"payload-bytes"): + """Write a minimal real safetensors file: 8-byte little-endian header + length, a JSON header containing ``__metadata__``, then arbitrary payload.""" + header = json.dumps({"__metadata__": metadata}).encode("utf-8") + path.write_bytes(struct.pack("