feat(metadata): add CivitAI AutoV3 hash support across all storage layers

- Three-state autov3 field (not-checked / checked-unavailable / 12-hex value)
  in .metadata.json sidecars, in-memory ModelHashIndex, and SQLite
  (models.autov3 column + autov3_index table) with column-presence migration
- Background self-terminating backfill for legacy rows: per-model-type
  concurrency guard, executor-offloaded I/O, Civitai-first resolution
  (SHA256-matched version file) falling back to the embedded safetensors
  header hash
- Civitai-first propagation on metadata refresh, scan, and download paths;
  reject the empty-string SHA256 placeholder and strip OneTrainer 0x prefix
- List API hash filters and hash index lookups accept 12-char AutoV3
- Cap safetensors header reads at 64 MiB to prevent crafted-file allocation
- Prevent stale AutoV3 mappings on file replacement while preserving them on
  same-file re-registration (lazy-hash completion)
This commit is contained in:
Will Miao
2026-08-08 14:30:34 +08:00
parent 4bf9a4b640
commit 97b9b1f62b
23 changed files with 1918 additions and 50 deletions

View File

@@ -39,6 +39,7 @@ These fields are present in all model metadata files.
| `metadata_source` | string\|null | ❌ No | ✅ Yes | Last provider that supplied metadata (see below) |
| `last_checked_at` | float | ❌ No (default: `0`) | ✅ Yes | Unix timestamp of last metadata check |
| `hash_status` | string | ❌ No (default: `"completed"`) | ✅ Yes | Hash calculation status: `"pending"`, `"calculating"`, `"completed"`, `"failed"` |
| `autov3` | string\|null | ❌ No | ✅ Yes | CivitAI AutoV3 hash (first 12 chars, lowercase hex) sourced from the safetensors embedded metadata (`sshs_model_hash` / `modelspec.hash_sha256`). **Absent** = not yet checked (may be backfilled later); **`null`** = checked but unavailable (header has no recognized hash); **12-char hex string** = value |
---
@@ -287,6 +288,7 @@ These fields are automatically synchronized with the filesystem:
- `preview_url` — Updated if preview file is moved/removed
- `sha256` — Updated during hash calculation (when `hash_status="pending"`)
- `hash_status` — Updated during hash calculation
- `autov3` — Set when metadata is first created (from safetensors header); may be backfilled later for entries where it is absent
- `last_checked_at` — Timestamp of scan
- `metadata_source` — Set based on metadata provider
@@ -345,6 +347,7 @@ These fields can be edited by users at any time through the Lora Manager UI or b
| `metadata_source` | `null` |
| `last_checked_at` | `0` |
| `hash_status` | `"completed"` |
| `autov3` | absent (not checked) or `null` (checked, no value) |
| `usage_tips` | `"{}"` (LoRA only) |
| `model_type` | `"checkpoint"` or `"embedding"` (not present in LoRA models) |
@@ -354,6 +357,7 @@ These fields can be edited by users at any time through the Lora Manager UI or b
| Version | Date | Changes |
|---------|------|---------|
| 1.1 | 2026-08 | Added `autov3` field (CivitAI AutoV3 hash with three-state semantics) |
| 1.0 | 2026-03 | Initial schema documentation |
---

View File

@@ -2119,20 +2119,25 @@ class RecipeManagementHandler:
for item in getattr(parent_cache_data, "raw_data", []):
if item.get("sha256", "").lower() == model_hash.lower():
local_cache[model_hash.lower()] = item
# Compute AutoV3 so the parser can also match on
# that hash type (CivitAI metadata resources use
# AutoV3).
file_path = item.get("file_path")
if file_path and os.path.exists(file_path):
try:
from ...utils.file_utils import (
calculate_autov3,
)
autov3 = calculate_autov3(file_path)
if autov3:
local_cache[autov3.lower()] = item
except Exception:
pass
# Register the AutoV3 hash so the parser can also
# match on that hash type (CivitAI metadata
# resources use AutoV3). Prefer the stored cache
# field; only compute it when the entry has none.
autov3 = (item.get("autov3") or "").lower()
if not autov3:
file_path = item.get("file_path")
if file_path and os.path.exists(file_path):
try:
from ...utils.file_utils import (
calculate_autov3,
)
autov3 = (
calculate_autov3(file_path) or ""
).lower()
except Exception:
pass
if autov3:
local_cache[autov3] = item
break
except Exception:
pass

View File

@@ -0,0 +1,140 @@
"""Backfill the AutoV3 checked state for models loaded from a persisted snapshot.
The SQLite persistent cache predates the AutoV3 feature, so entries hydrated
from it have a NULL ``autov3`` column (the "not checked yet" state). This
service computes the embedded AutoV3 hash for each such model — once per
process — and persists it through the scanner's single write path
(:meth:`ModelScanner.update_autov3_for_model`), marking every visited row so a
subsequent run finds nothing left to do.
Three-state contract honored here:
- ``NULL`` (sqlite) / absent (dict) = not checked yet → backfill computes it
- ``''`` (sqlite/dict) / JSON null = checked, no value available → never recompute
- 12-char lowercase hex = value → never recompute
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import threading
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING: # pragma: no cover - type-check only; runtime imports are local
from .model_scanner import ModelScanner
logger = logging.getLogger(__name__)
def _resolve_autov3(file_path: str) -> str:
"""Resolve the AutoV3 hash for a model file.
Prefers the Civitai AutoV3 reported for the file whose SHA256 matches
(the authoritative value for recipe matching); falls back to the embedded
safetensors header hash. Returns ``''`` when neither is available.
"""
try:
metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json"
if os.path.exists(metadata_path):
with open(metadata_path, "r", encoding="utf-8") as handle:
payload = json.load(handle)
if isinstance(payload, dict):
from ..utils.models import autov3_from_civitai_files # local import avoids cycles
sha256 = (payload.get("sha256") or "").lower()
civitai_autov3 = autov3_from_civitai_files(payload.get("civitai"), sha256)
if civitai_autov3:
return civitai_autov3
except Exception:
pass
from ..utils.file_utils import calculate_autov3 # local import avoids cycles
return calculate_autov3(file_path) or ""
class Autov3BackfillService:
"""Compute and persist AutoV3 hashes for models missing a checked state."""
_instance: Optional["Autov3BackfillService"] = None
_instance_lock = threading.Lock()
def __init__(self) -> None:
# Re-entrancy guard per model type: scanners for different model types
# initialize concurrently (lora_manager.py), so a global guard would
# silently skip every type but the first to start. Each model type
# runs its own backfill; a duplicate trigger for the same type no-ops.
self._running_types: set = set()
@classmethod
def get_instance(cls) -> "Autov3BackfillService":
"""Return the process-wide singleton instance."""
if cls._instance is None:
with cls._instance_lock:
if cls._instance is None:
cls._instance = cls()
return cls._instance
async def backfill(self, scanner: "ModelScanner") -> int:
"""Compute AutoV3 for every un-checked model of ``scanner.model_type``.
Each candidate file is read once via :func:`~py.utils.file_utils.calculate_autov3`
(cheap: safetensors header only) and the result is persisted through
``scanner.update_autov3_for_model``. Files that no longer exist on
disk are skipped — they are intentionally NOT marked, because scanner
cleanup removes the stale row later.
Returns:
The number of models successfully updated. Never raises; on any
failure a warning is logged and ``0`` is returned. A duplicate
trigger for a model type that is already being backfilled returns
``0`` immediately; different model types run concurrently.
"""
model_type = scanner.model_type
if model_type in self._running_types:
return 0
self._running_types.add(model_type)
try:
# Local imports avoid import cycles at module load time.
from .persistent_model_cache import get_persistent_cache
from ..utils.file_utils import calculate_autov3
persistent = getattr(scanner, "_persistent_cache", None) or get_persistent_cache()
paths = persistent.get_models_missing_autov3(model_type)
loop = asyncio.get_running_loop()
count = 0
for path in paths:
# A file that no longer exists must not be marked; scanner
# cleanup removes the stale row later. The existence check and
# hash resolution run in the executor so the loop stays
# responsive to API requests while the backfill iterates a
# large library.
if not await loop.run_in_executor(None, os.path.exists, path):
continue
autov3 = await loop.run_in_executor(None, _resolve_autov3, path)
if await scanner.update_autov3_for_model(model_type, path, autov3):
count += 1
if paths:
logger.info(
"AutoV3 backfill: updated %d/%d models for %s",
count,
len(paths),
model_type,
)
else:
# Steady state after the first run: nothing left to backfill.
logger.debug("AutoV3 backfill: nothing to process for %s", model_type)
return count
except Exception as exc:
logger.warning(
"AutoV3 backfill failed for %s: %s",
getattr(scanner, "model_type", "?"),
exc,
)
return 0
finally:
self._running_types.discard(model_type)

View File

@@ -446,20 +446,31 @@ class BaseModelService(ABC):
async def _apply_hash_filters(
self, data: List[Dict], hash_filters: Dict
) -> List[Dict]:
"""Apply hash-based filtering"""
"""Apply hash-based filtering (SHA256 and AutoV3)."""
def matches_hash_set(item: Dict, hash_set: set) -> bool:
"""Check whether an item matches any hash in the set.
Compares the item's ``sha256`` field and its non-empty ``autov3``
field, both case-insensitively.
"""
if item.get("sha256", "").lower() in hash_set:
return True
autov3 = item.get("autov3", "")
return bool(autov3) and autov3.lower() in hash_set
single_hash = hash_filters.get("single_hash")
multiple_hashes = hash_filters.get("multiple_hashes")
if single_hash:
# Filter by single hash
single_hash = single_hash.lower()
# Filter by single hash (SHA256 or AutoV3)
return [
item for item in data if item.get("sha256", "").lower() == single_hash
item for item in data if matches_hash_set(item, {single_hash.lower()})
]
elif multiple_hashes:
# Filter by multiple hashes
hash_set = set(hash.lower() for hash in multiple_hashes)
return [item for item in data if item.get("sha256", "").lower() in hash_set]
# Filter by multiple hashes (SHA256 or AutoV3)
hash_set = {hash.lower() for hash in multiple_hashes}
return [item for item in data if matches_hash_set(item, hash_set)]
return data

View File

@@ -59,6 +59,7 @@ class CacheEntryValidator:
'notes': ('', False),
'usage_tips': ('', False),
'hash_status': ('completed', False),
'autov3': (None, False),
}
@classmethod
@@ -119,8 +120,13 @@ class CacheEntryValidator:
if is_required:
errors.append(f"Required field '{field_name}' is missing or None")
if auto_repair:
working_entry[field_name] = cls._get_default_copy(default_value)
repaired = True
# A missing optional field whose default is None is already
# semantically equal to its default (e.g. autov3: absent
# means "not checked") — writing None back is a no-op, not
# a repair.
if default_value is not None:
working_entry[field_name] = cls._get_default_copy(default_value)
repaired = True
continue
# Validate field type and value
@@ -175,6 +181,15 @@ class CacheEntryValidator:
# that invalidates the entry, but we also don't mark it repaired.
pass
# Normalize autov3 to lowercase if needed (optional field, never stripped).
autov3 = working_entry.get('autov3')
if isinstance(autov3, str) and autov3:
normalized_autov3 = autov3.lower()
if normalized_autov3 != autov3:
if auto_repair:
working_entry['autov3'] = normalized_autov3
repaired = True
# Determine if entry is valid
# Entry is valid if no critical required field errors remain after repair
# Critical fields are file_path and sha256
@@ -242,6 +257,19 @@ class CacheEntryValidator:
"""
expected_type = type(default_value)
# Special case: autov3 is optional with a three-state contract.
# None = not checked, "" = checked but unavailable, otherwise a
# 12-character hex string (case-insensitive here; normalized to
# lowercase separately).
if field_name == 'autov3':
if value is None or value == "":
return None
if not isinstance(value, str):
return f"Field 'autov3' should be string or None, got {type(value).__name__}"
if len(value) != 12 or any(c not in '0123456789abcdefABCDEF' for c in value):
return "Field 'autov3' should be a 12-character hex string"
return None
# Special handling for numeric types
if expected_type == int:
if not isinstance(value, (int, float)):

View File

@@ -6,7 +6,7 @@ from datetime import datetime
from typing import Any, Dict, List, Optional
from ..utils.models import CheckpointMetadata
from ..utils.file_utils import find_preview_file, normalize_path
from ..utils.file_utils import find_preview_file, normalize_path, calculate_autov3
from ..utils.metadata_manager import MetadataManager
from ..config import config
from .model_scanner import ModelScanner
@@ -62,6 +62,11 @@ class CheckpointScanner(ModelScanner):
# Find preview image
preview_url = find_preview_file(base_name, dir_path)
# AutoV3 reads only the safetensors header, so it is cheap even for
# large checkpoints; record the checked state at creation time ("" =
# checked but unavailable).
autov3 = calculate_autov3(real_path)
# Create metadata WITHOUT calculating hash
metadata = CheckpointMetadata(
file_name=base_name,
@@ -77,6 +82,7 @@ class CheckpointScanner(ModelScanner):
sub_type="checkpoint",
from_civitai=False, # Mark as local model since no hash yet
hash_status="pending", # Mark hash as pending
autov3=autov3 or "",
)
# Save the created metadata
@@ -120,7 +126,11 @@ class CheckpointScanner(ModelScanner):
# that queries get_hash_by_filename first) will miss on every
# lookup and keep calling back into this method, creating a
# tight loop that never populates the index.
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
self._hash_index.add_entry(
metadata.sha256.lower(),
file_path,
getattr(metadata, "autov3", None) or None,
)
return metadata.sha256
async with self._hash_calculation_lock:
@@ -132,7 +142,11 @@ class CheckpointScanner(ModelScanner):
and metadata.hash_status == "completed"
and metadata.sha256
):
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
self._hash_index.add_entry(
metadata.sha256.lower(),
file_path,
getattr(metadata, "autov3", None) or None,
)
return metadata.sha256
task = self._hash_calculation_tasks.get(real_path)
@@ -185,7 +199,11 @@ class CheckpointScanner(ModelScanner):
if metadata.hash_status == "completed" and metadata.sha256:
# Populate the in-memory hash index even for pre-computed
# hashes, mirroring the fix in calculate_hash_for_model.
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
self._hash_index.add_entry(
metadata.sha256.lower(),
file_path,
getattr(metadata, "autov3", None) or None,
)
return metadata.sha256
# Update status to calculating
@@ -202,7 +220,11 @@ class CheckpointScanner(ModelScanner):
await MetadataManager.save_metadata(file_path, metadata)
# Update hash index
self._hash_index.add_entry(sha256.lower(), file_path)
self._hash_index.add_entry(
sha256.lower(),
file_path,
getattr(metadata, "autov3", None) or None,
)
# Update the in-memory cache entry so that subsequent
# _persist_current_cache / _save_persistent_cache calls

View File

@@ -11,6 +11,7 @@ from typing import Any, Awaitable, Callable, Dict, Iterable, Optional
from ..services.settings_manager import SettingsManager
from ..utils.civitai_utils import resolve_license_payload
from ..utils.model_utils import determine_base_model
from ..utils.models import autov3_from_civitai_files
from .connectivity_guard import OFFLINE_FRIENDLY_MESSAGE, is_expected_offline_error
from .errors import RateLimitError
@@ -152,6 +153,18 @@ class MetadataSyncService:
civitai_metadata.get("baseModel")
)
# Civitai-first AutoV3 propagation: the freshly fetched version
# metadata may report an AutoV3 for the file whose SHA256 matches the
# local model. Persist it now so recipe matching sees it immediately —
# no full rescan or restart required (the header is never re-read to
# upgrade the checked-unavailable '' state).
sha256_value = (local_metadata.get("sha256") or "").lower()
civitai_autov3 = autov3_from_civitai_files(
local_metadata.get("civitai"), sha256_value
)
if civitai_autov3:
local_metadata["autov3"] = civitai_autov3
await self._preview_service.ensure_preview_for_metadata(
metadata_path, local_metadata, civitai_metadata.get("images", [])
)

View File

@@ -8,11 +8,12 @@ class ModelHashIndex:
self._hash_to_path: Dict[str, str] = {}
self._filename_to_hash: Dict[str, str] = {}
self._autov2_to_path: Dict[str, str] = {}
self._autov3_to_path: Dict[str, str] = {}
# New data structures for tracking duplicates
self._duplicate_hashes: Dict[str, List[str]] = {} # sha256 -> list of paths
self._duplicate_filenames: Dict[str, List[str]] = {} # filename -> list of paths
def add_entry(self, sha256: str, file_path: str) -> None:
def add_entry(self, sha256: str, file_path: str, autov3: Optional[str] = None) -> None:
"""Add or update hash index entry"""
if not sha256 or not file_path:
return
@@ -33,9 +34,13 @@ class ModelHashIndex:
self._duplicate_hashes.setdefault(sha256, []).append(file_path)
# Track duplicates by filename - FIXED LOGIC
is_re_registration = False
if filename in self._filename_to_hash:
existing_hash = self._filename_to_hash[filename]
existing_path = self._hash_to_path.get(existing_hash)
# Same path registered again (e.g. a file replaced in place with
# new content) — used below to drop its stale autov3 mapping.
is_re_registration = existing_path == file_path
# If this is a different file with the same filename
if existing_path and existing_path != file_path:
@@ -67,6 +72,30 @@ class ModelHashIndex:
# AutoV2 = first 10 chars of SHA256
if len(sha256) >= 10:
self._autov2_to_path[sha256[:10]] = file_path
# AutoV3 is an independent hash (not derived from SHA256), stored as-is.
# Drop stale mappings for a path when it is re-registered with a NEW
# sha256 (file replaced in place) or with an explicit new autov3 value
# (correction). Re-registering the SAME file with the same sha256 and
# no autov3 (e.g. lazy-hash completion) must never clear its existing
# mapping. First-time registrations stay O(1).
if autov3:
autov3 = autov3.lower()
if is_re_registration and (existing_hash != sha256 or autov3):
stale_autov3_keys = [
key for key, mapped_path in self._autov3_to_path.items()
if mapped_path == file_path and key != autov3
]
for key in stale_autov3_keys:
del self._autov3_to_path[key]
if autov3:
self._autov3_to_path[autov3] = file_path
def add_autov3(self, autov3: str, file_path: str) -> None:
"""Add or update an AutoV3-only index entry (used when only AutoV3 is known)"""
if not autov3:
return
autov3 = autov3.lower()
self._autov3_to_path[autov3] = file_path
def _get_filename_from_path(self, file_path: str) -> str:
"""Extract filename without extension from path"""
@@ -167,6 +196,11 @@ class ModelHashIndex:
for k in autov2_keys_to_remove:
del self._autov2_to_path[k]
# Remove from AutoV3 index
autov3_keys_to_remove = [k for k, v in self._autov3_to_path.items() if v == file_path]
for k in autov3_keys_to_remove:
del self._autov3_to_path[k]
def remove_by_hash(self, sha256: str) -> None:
"""Remove entry by hash"""
sha256 = sha256.lower()
@@ -189,6 +223,11 @@ class ModelHashIndex:
autov2_key = sha256[:10]
if autov2_key in self._autov2_to_path:
del self._autov2_to_path[autov2_key]
# Remove AutoV3 entries pointing to any removed path
autov3_keys_to_remove = [k for k, v in self._autov3_to_path.items() if v in paths_to_remove]
for k in autov3_keys_to_remove:
del self._autov3_to_path[k]
# Update filename-to-hash and duplicate filenames for all paths
for path_to_remove in paths_to_remove:
@@ -209,22 +248,26 @@ class ModelHashIndex:
del self._duplicate_filenames[fname]
def has_hash(self, hash_value: str) -> bool:
"""Check if hash exists in index (SHA256 or AutoV2)"""
"""Check if hash exists in index (SHA256, AutoV2, or AutoV3)"""
normalized = hash_value.lower()
if normalized in self._hash_to_path:
return True
if len(normalized) == 10:
return normalized in self._autov2_to_path
if len(normalized) == 12:
return normalized in self._autov3_to_path
return False
def get_path(self, hash_value: str) -> Optional[str]:
"""Get file path for a hash (SHA256 or AutoV2)"""
"""Get file path for a hash (SHA256, AutoV2, or AutoV3)"""
normalized = hash_value.lower()
path = self._hash_to_path.get(normalized)
if path is not None:
return path
if len(normalized) == 10:
return self._autov2_to_path.get(normalized)
if len(normalized) == 12:
return self._autov3_to_path.get(normalized)
return None
def get_hash(self, file_path: str) -> Optional[str]:
@@ -243,6 +286,7 @@ class ModelHashIndex:
self._hash_to_path.clear()
self._filename_to_hash.clear()
self._autov2_to_path.clear()
self._autov3_to_path.clear()
self._duplicate_hashes.clear()
self._duplicate_filenames.clear()
@@ -253,6 +297,10 @@ class ModelHashIndex:
def get_all_filenames(self) -> Set[str]:
"""Get all filenames in the index"""
return set(self._filename_to_hash.keys())
def get_all_autov3(self) -> Dict[str, str]:
"""Get a snapshot of all AutoV3 hashes mapped to their file paths"""
return dict(self._autov3_to_path)
def get_duplicate_hashes(self) -> Dict[str, List[str]]:
"""Get dictionary of duplicate hashes and their paths"""

View File

@@ -7,9 +7,9 @@ import shutil
from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Set, Type, Union
from ..utils.models import BaseModelMetadata
from ..utils.models import BaseModelMetadata, autov3_from_civitai_files
from ..config import config
from ..utils.file_utils import find_preview_file, get_preview_extension, calculate_sha256
from ..utils.file_utils import find_preview_file, get_preview_extension, calculate_sha256, calculate_autov3
from ..utils.metadata_manager import MetadataManager
from ..utils.civitai_utils import resolve_license_info
from .model_cache import ModelCache
@@ -86,6 +86,7 @@ class ModelScanner:
self._persistent_cache = get_persistent_cache()
self._name_display_mode = self._resolve_name_display_mode()
self._cancel_requested = False # Flag for cancellation
self._autov3_backfill_scheduled = False # One-time AutoV3 backfill trigger per process
try:
loop = asyncio.get_running_loop()
except RuntimeError:
@@ -225,6 +226,19 @@ class ModelScanner:
if not isinstance(notes, str):
notes = str(notes)
# AutoV3 three-state contract: absent key / None = "not checked yet",
# "" = "checked but unavailable" (never re-read the header), else the
# 12-char lowercase hex value. A metadata object already follows the
# contract and is passed through unchanged; a payload dict only carries
# an explicit checked state when the key is present.
if is_mapping:
if 'autov3' in source:
entry_autov3 = source['autov3'] or ''
else:
entry_autov3 = None
else:
entry_autov3 = get_value('autov3', None)
entry: Dict[str, Any] = {
'file_path': normalized_path,
# file_name is always stored WITHOUT extension (e.g. "OWSMianne_ANIMA_V1",
@@ -238,6 +252,7 @@ class ModelScanner:
'size': int(get_value('size', 0) or 0),
'modified': float(get_value('modified', 0.0) or 0.0),
'sha256': (get_value('sha256', '') or '').lower(),
'autov3': entry_autov3,
'base_model': get_value('base_model', '') or '',
'preview_url': preview_url,
'preview_nsfw_level': int(get_value('preview_nsfw_level', 0) or 0),
@@ -473,6 +488,13 @@ class ModelScanner:
if sha_value and path:
hash_index.add_entry(sha_value.lower(), path)
# Rebuild the AutoV3 index from the persisted autov3_index rows. These
# cover every known autov3 -> path mapping regardless of whether a
# sha256 row also exists for the same file.
for autov3_value, path in persisted.autov3_hash_rows:
if autov3_value and path:
hash_index.add_autov3(autov3_value.lower(), path)
tags_count: Dict[str, int] = {}
adjusted_raw_data: List[Dict[str, Any]] = []
for item in persisted.raw_data:
@@ -541,8 +563,30 @@ class ModelScanner:
'scanner_type': self.model_type,
'pageType': page_type
})
# Schedule the one-time AutoV3 backfill task (at most once per process)
# so entries loaded from a persisted snapshot that predates autov3 get
# their checked state computed in the background. The task never blocks
# or crashes the load path.
if not self._autov3_backfill_scheduled:
self._autov3_backfill_scheduled = True
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = None
if loop is not None:
loop.create_task(self._run_autov3_backfill())
return True
async def _run_autov3_backfill(self) -> None:
"""Backfill autov3 for entries loaded from the persisted cache that lack it."""
try:
from ..services.autov3_backfill_service import Autov3BackfillService # lazy import (module created by another unit)
await Autov3BackfillService.get_instance().backfill(self)
except Exception as exc:
logger.warning("AutoV3 backfill failed: %s", exc)
async def _save_persistent_cache(self, scan_result: CacheBuildResult) -> None:
if not scan_result or not getattr(self, '_persistent_cache', None):
return
@@ -555,6 +599,7 @@ class ModelScanner:
return
hash_snapshot = self._build_hash_index_snapshot(scan_result.hash_index)
autov3_snapshot = self._build_autov3_index_snapshot(scan_result.hash_index)
loop = asyncio.get_event_loop()
try:
await loop.run_in_executor(
@@ -563,7 +608,8 @@ class ModelScanner:
self.model_type,
list(scan_result.raw_data),
hash_snapshot,
list(scan_result.excluded_models)
list(scan_result.excluded_models),
autov3_snapshot,
)
except Exception as exc:
logger.warning("%s Scanner: Failed to persist cache: %s", self.model_type.capitalize(), exc)
@@ -589,6 +635,20 @@ class ModelScanner:
bucket.append(path)
return snapshot
def _build_autov3_index_snapshot(self, hash_index: Optional[ModelHashIndex]) -> Dict[str, List[str]]:
"""Build the autov3 -> [paths] snapshot for the persisted cache."""
snapshot: Dict[str, List[str]] = {}
if not hash_index:
return snapshot
for autov3_value, path in hash_index.get_all_autov3().items():
if not autov3_value or not path:
continue
bucket = snapshot.setdefault(autov3_value.lower(), [])
if path not in bucket:
bucket.append(path)
return snapshot
async def _persist_current_cache(self) -> None:
if self._cache is None or not getattr(self, '_persistent_cache', None):
return
@@ -880,7 +940,11 @@ class ModelScanner:
# Update hash index if available
if 'sha256' in model_data and 'file_path' in model_data:
self._hash_index.add_entry(model_data['sha256'].lower(), model_data['file_path'])
self._hash_index.add_entry(
model_data['sha256'].lower(),
model_data['file_path'],
model_data.get('autov3') or None
)
# Update tags count
if 'tags' in model_data and model_data['tags']:
@@ -1130,6 +1194,36 @@ class ModelScanner:
except Exception as e:
logger.error(f"Failed to compute SHA256 for {file_path}: {e}")
# AutoV3 resolution: prefer the Civitai AutoV3 reported for the file
# whose SHA256 matches (authoritative for recipe matching), falling
# back to the embedded safetensors header hash only for models never
# checked before (autov3 is None). A checked-unavailable state ('')
# is only upgraded by Civitai data — the header is never re-read.
current_autov3 = model_data.get('autov3')
if current_autov3 in (None, ''):
try:
civitai_data = None
if isinstance(metadata, BaseModelMetadata):
civitai_data = metadata.civitai
elif isinstance(metadata, dict):
civitai_data = metadata.get("civitai")
autov3 = autov3_from_civitai_files(
civitai_data, model_data.get("sha256") or ""
) or ""
if not autov3 and current_autov3 is None:
autov3 = (calculate_autov3(os.path.realpath(file_path)) or '').lower()
if autov3 != current_autov3:
model_data['autov3'] = autov3
if isinstance(metadata, BaseModelMetadata):
metadata.autov3 = autov3
await MetadataManager.save_metadata(file_path, metadata)
elif isinstance(metadata, dict):
# Dict payload: JSON null encodes the checked-unavailable state.
metadata['autov3'] = autov3 or None
await MetadataManager.save_metadata(file_path, metadata)
except Exception as e:
logger.error(f"Failed to resolve AutoV3 for {file_path}: {e}")
# Skip excluded models
if model_data.get('exclude', False):
excluded_models.append(model_data['file_path'])
@@ -1322,7 +1416,7 @@ class ModelScanner:
sha_value = result.get('sha256')
model_path = result.get('file_path')
if sha_value and model_path:
hash_index.add_entry(sha_value.lower(), model_path)
hash_index.add_entry(sha_value.lower(), model_path, result.get('autov3') or None)
for tag in result.get('tags') or []:
tags_count[tag] = tags_count.get(tag, 0) + 1
@@ -1391,7 +1485,11 @@ class ModelScanner:
await self._cache.resort()
# Update the hash index
self._hash_index.add_entry(metadata_dict['sha256'], metadata_dict['file_path'])
self._hash_index.add_entry(
metadata_dict['sha256'],
metadata_dict['file_path'],
metadata_dict.get('autov3') or None,
)
await self._persist_current_cache()
return True
except Exception as e:
@@ -1578,7 +1676,11 @@ class ModelScanner:
sha_value = cache_entry.get('sha256')
if sha_value:
self._hash_index.add_entry(sha_value.lower(), normalized_new_path)
self._hash_index.add_entry(
sha_value.lower(),
normalized_new_path,
cache_entry.get('autov3') or None,
)
all_folders = set(item['folder'] for item in cache.raw_data)
cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
@@ -1736,7 +1838,11 @@ class ModelScanner:
if old_sha:
self._hash_index.remove_by_path(file_path)
if new_sha:
self._hash_index.add_entry(new_sha, file_path)
self._hash_index.add_entry(
new_sha,
file_path,
desired_entry.get('autov3') or None,
)
# ---- Incremental version index update ----
new_civitai = desired_entry.get("civitai")
@@ -1787,6 +1893,74 @@ class ModelScanner:
return True
async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool:
"""Persist an AutoV3 hash for a single model (single write path used by the backfill service).
Locates the in-memory cache entry by ``file_path`` and updates only its
``autov3`` field: the in-memory hash index, the SQLite snapshot via
:meth:`PersistentModelCache.update_single_model`, and the
``.metadata.json`` sidecar. sha256, tags, and every other field are
left untouched, so the persistent delta only ever differs in autov3.
Returns:
``True`` when the entry was found and updated, ``False`` otherwise.
Never raises — failures are logged and swallowed.
"""
try:
if self._cache is None:
return False
entry = next(
(item for item in self._cache.raw_data if item.get('file_path') == file_path),
None,
)
if entry is None:
return False
# Normalize once so the memory entry, sidecar, and SQLite row agree.
autov3 = (autov3 or "").lower()
# Capture the pre-mutation state so update_single_model only sees
# an autov3 delta between old and new.
old_item = dict(entry)
entry['autov3'] = autov3 or ''
# Prefer add_entry when a sha256 is known so the sha256 and autov3
# maps stay in sync; fall back to an autov3-only registration.
sha_value = entry.get('sha256')
checked_autov3 = entry.get('autov3') or None
if sha_value:
self._hash_index.add_entry(sha_value.lower(), file_path, checked_autov3)
elif checked_autov3:
self._hash_index.add_autov3(checked_autov3, file_path)
persistent = getattr(self, '_persistent_cache', None)
if persistent is not None:
await asyncio.get_event_loop().run_in_executor(
None,
persistent.update_single_model,
model_type,
entry,
old_item,
)
# Sidecar write-back: JSON null encodes the checked-unavailable
# state. Skip silently when the sidecar does not exist.
metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json"
if os.path.exists(metadata_path):
with open(metadata_path, 'r', encoding='utf-8') as handle:
payload = json.load(handle)
if not isinstance(payload, dict):
payload = {}
payload['autov3'] = entry['autov3'] or None
await MetadataManager.save_metadata(metadata_path, payload)
return True
except Exception as exc:
logger.warning("Failed to update AutoV3 for %s: %s", file_path, exc)
return False
@staticmethod
def _cache_entries_differ(a: Dict[str, Any], b: Dict[str, Any]) -> bool:
"""Return ``True`` when two cache-entry dicts differ in any field.

View File

@@ -3,7 +3,7 @@ import logging
import os
import sqlite3
import threading
from dataclasses import dataclass
from dataclasses import dataclass, field
from typing import Dict, List, Mapping, Optional, Sequence, Tuple
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
@@ -18,6 +18,7 @@ class PersistedCacheData:
raw_data: List[Dict]
hash_rows: List[Tuple[str, str]]
excluded_models: List[str]
autov3_hash_rows: List[Tuple[str, str]] = field(default_factory=list)
DEFAULT_LICENSE_FLAGS = 127 # 127 (0b1111111) encodes default CivitAI permissions with all commercial modes enabled.
@@ -36,6 +37,7 @@ class PersistentModelCache:
"size",
"modified",
"sha256",
"autov3",
"base_model",
"preview_url",
"preview_nsfw_level",
@@ -118,6 +120,10 @@ class PersistentModelCache:
"SELECT sha256, file_path FROM hash_index WHERE model_type = ?",
(model_type,),
).fetchall()
autov3_rows = conn.execute(
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
(model_type,),
).fetchall()
excluded = conn.execute(
"SELECT file_path FROM excluded_models WHERE model_type = ?",
(model_type,),
@@ -191,6 +197,8 @@ class PersistentModelCache:
"hash_status": row["hash_status"] or "completed",
"hf_url": row["hf_url"] or "",
}
if row["autov3"] is not None:
item["autov3"] = (row["autov3"] or "").lower()
raw_data.append(item)
hash_pairs = [(entry["sha256"].lower(), entry["file_path"]) for entry in hash_rows if entry["sha256"]]
@@ -201,10 +209,21 @@ class PersistentModelCache:
if sha_value:
hash_pairs.append((sha_value.lower(), item["file_path"]))
excluded_paths = [row["file_path"] for row in excluded]
return PersistedCacheData(raw_data=raw_data, hash_rows=hash_pairs, excluded_models=excluded_paths)
autov3_pairs = [
(entry["autov3"].lower(), entry["file_path"])
for entry in autov3_rows
if entry["autov3"]
]
def save_cache(self, model_type: str, raw_data: Sequence[Dict], hash_index: Dict[str, List[str]], excluded_models: Sequence[str]) -> None:
excluded_paths = [row["file_path"] for row in excluded]
return PersistedCacheData(
raw_data=raw_data,
hash_rows=hash_pairs,
excluded_models=excluded_paths,
autov3_hash_rows=autov3_pairs,
)
def save_cache(self, model_type: str, raw_data: Sequence[Dict], hash_index: Dict[str, List[str]], excluded_models: Sequence[str], autov3_hash_index: Optional[Dict[str, List[str]]] = None) -> None:
if not self.is_enabled():
return
if not self._schema_initialized:
@@ -251,6 +270,10 @@ class PersistentModelCache:
"DELETE FROM hash_index WHERE model_type = ? AND file_path = ?",
to_remove_models,
)
conn.executemany(
"DELETE FROM autov3_index WHERE model_type = ? AND file_path = ?",
to_remove_models,
)
conn.executemany(
"DELETE FROM excluded_models WHERE model_type = ? AND file_path = ?",
to_remove_models,
@@ -373,6 +396,52 @@ class PersistentModelCache:
hash_inserts,
)
if autov3_hash_index is not None:
existing_autov3_rows = conn.execute(
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
(model_type,),
).fetchall()
existing_autov3_map: Dict[str, set] = {}
for row in existing_autov3_rows:
autov3_value = (row["autov3"] or "").lower()
if not autov3_value:
continue
existing_autov3_map.setdefault(autov3_value, set()).add(row["file_path"])
new_autov3_map: Dict[str, set] = {}
for autov3_value, paths in autov3_hash_index.items():
normalized_autov3 = (autov3_value or "").lower()
if not normalized_autov3:
continue
bucket = new_autov3_map.setdefault(normalized_autov3, set())
for path in paths:
if path:
bucket.add(path)
autov3_inserts: List[Tuple[str, str, str]] = []
autov3_deletes: List[Tuple[str, str, str]] = []
all_autov3 = set(existing_autov3_map.keys()) | set(new_autov3_map.keys())
for autov3_value in all_autov3:
existing_paths = existing_autov3_map.get(autov3_value, set())
new_paths = new_autov3_map.get(autov3_value, set())
for path in existing_paths - new_paths:
autov3_deletes.append((model_type, autov3_value, path))
for path in new_paths - existing_paths:
autov3_inserts.append((model_type, autov3_value, path))
if autov3_deletes:
conn.executemany(
"DELETE FROM autov3_index WHERE model_type = ? AND autov3 = ? AND file_path = ?",
autov3_deletes,
)
if autov3_inserts:
conn.executemany(
"INSERT OR IGNORE INTO autov3_index (model_type, autov3, file_path) VALUES (?, ?, ?)",
autov3_inserts,
)
existing_excluded_rows = conn.execute(
"SELECT file_path FROM excluded_models WHERE model_type = ?",
(model_type,),
@@ -435,6 +504,7 @@ class PersistentModelCache:
size INTEGER,
modified REAL,
sha256 TEXT,
autov3 TEXT,
base_model TEXT,
preview_url TEXT,
preview_nsfw_level INTEGER,
@@ -472,6 +542,13 @@ class PersistentModelCache:
PRIMARY KEY (model_type, sha256, file_path)
);
CREATE TABLE IF NOT EXISTS autov3_index (
model_type TEXT NOT NULL,
autov3 TEXT NOT NULL,
file_path TEXT NOT NULL,
PRIMARY KEY (model_type, autov3, file_path)
);
CREATE TABLE IF NOT EXISTS excluded_models (
model_type TEXT NOT NULL,
file_path TEXT NOT NULL,
@@ -504,6 +581,7 @@ class PersistentModelCache:
"license_flags": f"INTEGER DEFAULT {DEFAULT_LICENSE_FLAGS}",
"hash_status": "TEXT DEFAULT 'completed'",
"hf_url": "TEXT DEFAULT ''",
"autov3": "TEXT",
}
for column, definition in required_columns.items():
@@ -549,6 +627,12 @@ class PersistentModelCache:
if license_flags is None:
license_flags = DEFAULT_LICENSE_FLAGS
autov3_value = item.get("autov3")
if autov3_value is None:
autov3_column = None
else:
autov3_column = (autov3_value or "").lower()
return (
model_type,
item.get("file_path"),
@@ -558,6 +642,7 @@ class PersistentModelCache:
int(item.get("size") or 0),
float(item.get("modified") or 0.0),
(item.get("sha256") or "").lower() or None,
autov3_column,
item.get("base_model") or "",
item.get("preview_url") or "",
int(item.get("preview_nsfw_level") or 0),
@@ -663,6 +748,25 @@ class PersistentModelCache:
(model_type, new_sha, file_path),
)
# --- autov3_index ---
new_autov3: Optional[str] = new_item.get("autov3")
if new_autov3 is not None:
new_autov3 = (new_autov3 or "").lower()
old_autov3: Optional[str] = (old_item.get("autov3") if old_item else None)
if old_autov3 is not None:
old_autov3 = (old_autov3 or "").lower()
if new_autov3 != old_autov3:
if old_autov3:
conn.execute(
"DELETE FROM autov3_index WHERE model_type = ? AND autov3 = ? AND file_path = ?",
(model_type, old_autov3, file_path),
)
if new_autov3:
conn.execute(
"INSERT OR IGNORE INTO autov3_index (model_type, autov3, file_path) VALUES (?, ?, ?)",
(model_type, new_autov3, file_path),
)
conn.execute("COMMIT")
except Exception:
conn.execute("ROLLBACK")
@@ -676,6 +780,40 @@ class PersistentModelCache:
exc,
)
def get_models_missing_autov3(self, model_type: str) -> List[str]:
"""Return file paths whose models lack an AutoV3 checked state.
Only rows with a completed sha256 and a NULL autov3 column qualify —
rows with '' (checked-unavailable) or a value are never returned, so
the backfill query self-terminates.
"""
if not self.is_enabled():
return []
if not self._schema_initialized:
self._initialize_schema()
if not self._schema_initialized:
return []
try:
with self._db_lock:
conn = self._connect(readonly=True)
try:
rows = conn.execute(
"SELECT file_path FROM models "
"WHERE model_type = ? AND autov3 IS NULL "
"AND sha256 IS NOT NULL AND sha256 != ''",
(model_type,),
).fetchall()
finally:
conn.close()
return [row["file_path"] for row in rows]
except Exception as exc:
logger.warning(
"Failed to query models missing autov3 for %s: %s",
model_type,
exc,
)
return []
def _load_tags(self, conn: sqlite3.Connection, model_type: str) -> Dict[str, List[str]]:
tag_rows = conn.execute(
"SELECT file_path, tag FROM model_tags WHERE model_type = ?",

View File

@@ -80,6 +80,17 @@ CIVITAI_USER_MODEL_TYPES = [
# Default chunk size in megabytes used for hashing large files.
DEFAULT_HASH_CHUNK_SIZE_MB = 4
# Upper bound for a safetensors header block (bytes). Real headers are at most
# a few MB (tensor name/shape lists); the cap prevents a crafted file with an
# absurd 64-bit header length from forcing a multi-GB allocation during scan.
MAX_SAFETENSORS_HEADER_BYTES = 64 * 1024 * 1024
# First 12 chars of the SHA256 of an empty byte string. Some (re-packaging)
# training tools write this placeholder into safetensors metadata instead of a
# real hash; it must never be treated as a valid AutoV3 — several broken
# models sharing it would collide in the hash index and falsely match recipes.
INVALID_AUTOV3_EMPTY_HASH = "e3b0c44298fc"
# Auto-organize settings
AUTO_ORGANIZE_BATCH_SIZE = (
50 # Process models in batches to avoid overwhelming the system

View File

@@ -9,6 +9,8 @@ from typing import Any
from .constants import (
CARD_PREVIEW_WIDTH,
DEFAULT_HASH_CHUNK_SIZE_MB,
INVALID_AUTOV3_EMPTY_HASH,
MAX_SAFETENSORS_HEADER_BYTES,
PREVIEW_EXTENSIONS,
)
from .exif_utils import ExifUtils
@@ -90,6 +92,8 @@ def read_safetensors_metadata(file_path: str) -> dict[str, Any]:
if len(header_len_bytes) < 8:
return {}
header_len = struct.unpack("<Q", header_len_bytes)[0]
if header_len > MAX_SAFETENSORS_HEADER_BYTES:
return {}
header_bytes = f.read(header_len)
if len(header_bytes) < header_len:
return {}
@@ -123,8 +127,16 @@ def calculate_autov3(file_path: str) -> str | None:
return None
embedded_hash = metadata.get("sshs_model_hash") or metadata.get("modelspec.hash_sha256")
if embedded_hash and isinstance(embedded_hash, str) and len(embedded_hash) >= 12:
return embedded_hash[:12]
if embedded_hash and isinstance(embedded_hash, str):
# OneTrainer writes modelspec.hash_sha256 with a "0x" prefix.
embedded_hash = embedded_hash.strip().removeprefix("0x").removeprefix("0X")
if len(embedded_hash) >= 12:
autov3 = embedded_hash[:12].lower()
# The empty-string SHA256 placeholder written by some repackaging
# tools is not a real hash; treat it as unavailable so broken
# models never share one bogus value.
if autov3 != INVALID_AUTOV3_EMPTY_HASH:
return autov3
return None

View File

@@ -6,7 +6,7 @@ import time
from typing import Any, Dict, Optional, Type, Union
from .models import BaseModelMetadata, LoraMetadata
from .file_utils import normalize_path, find_preview_file, calculate_sha256
from .file_utils import normalize_path, find_preview_file, calculate_sha256, calculate_autov3
from .lora_metadata import extract_lora_metadata, extract_checkpoint_metadata
logger = logging.getLogger(__name__)
@@ -210,6 +210,11 @@ class MetadataManager:
hash_duration = time.perf_counter() - start_hash_time
logger.info(f"SHA256 hash calculated for {real_path} in {hash_duration:.3f}s")
# AutoV3 reads only the safetensors header, so it is cheap even for
# large files. At creation time we always know the checked state:
# store "" when no recognized hash is embedded (checked-unavailable).
autov3 = calculate_autov3(real_path)
# Create instance based on model type
if model_class.__name__ == "CheckpointMetadata":
metadata = model_class(
@@ -257,6 +262,9 @@ class MetadataManager:
usage_tips="{}"
)
# Record the AutoV3 state explicitly ("" = checked, no value).
metadata.autov3 = autov3 or ""
# Try to extract model-specific metadata
# await MetadataManager._enrich_metadata(metadata, real_path)

View File

@@ -2,9 +2,40 @@ from dataclasses import dataclass, asdict, field
from typing import Dict, Optional, List, Any
from datetime import datetime
import os
from .constants import INVALID_AUTOV3_EMPTY_HASH
from .model_utils import determine_base_model
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.
Returns ``None`` when no Civitai data, no matching file, or no usable
AutoV3 hash is available.
"""
if not civitai_data or not sha256:
return None
target_sha = sha256.lower()
for file_info in civitai_data.get("files") or []:
if not isinstance(file_info, dict):
continue
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 None
@dataclass
class BaseModelMetadata:
"""Base class for all model metadata structures"""
@@ -35,6 +66,7 @@ class BaseModelMetadata:
metadata_source: Optional[str] = None # Last provider that supplied metadata
last_checked_at: float = 0 # Last checked timestamp
hash_status: str = "completed" # Hash calculation status: pending | calculating | completed | failed
autov3: Optional[str] = None # CivitAI AutoV3 hash (12-char lowercase hex); "" = checked but unavailable, None = not checked
trainedWords: List[str] = field(
default_factory=list
) # Trigger words / activation prompts (source-agnostic)
@@ -58,6 +90,14 @@ class BaseModelMetadata:
"""Create instance from dictionary"""
data_copy = data.copy()
# autov3 three-state semantics: an explicit key means the value is known.
# JSON null in the sidecar ("checked but unavailable") is normalized to ""
# in memory; an absent key stays None ("not checked yet"). autov3 is a known
# field, so it flows through fields_to_use below and never leaks into
# _unknown_fields.
if "autov3" in data_copy:
data_copy["autov3"] = data_copy["autov3"] or ""
# Use cached fields if available, otherwise compute them
if not hasattr(cls, "_known_fields_cache"):
known_fields = set()
@@ -97,11 +137,29 @@ class BaseModelMetadata:
if hasattr(self, "_unknown_fields"):
result.update(self._unknown_fields)
# autov3 three-state semantics: emit the key only when the value is known.
# "" is serialized as JSON null ("checked but unavailable"); an absent key
# means "not checked yet". Done after unknown fields so a stale unknown
# copy can never override the typed field.
if self.autov3 is not None:
result["autov3"] = self.autov3 or None
else:
result.pop("autov3", None)
return result
def update_civitai_info(self, civitai_data: Dict) -> None:
"""Update Civitai information"""
"""Update Civitai information.
Civitai's AutoV3 is the authoritative hash for recipe matching, so
whenever the version metadata reports an AutoV3 for the file whose
SHA256 matches this model, it takes precedence over the locally
extracted header hash.
"""
self.civitai = civitai_data
autov3 = autov3_from_civitai_files(civitai_data, self.sha256)
if autov3:
self.autov3 = autov3
def update_file_info(self, file_path: str, update_timestamps: bool = False) -> None:
"""
@@ -190,13 +248,15 @@ class LoraMetadata(BaseModelMetadata):
if "description" in model_data:
description = model_data["description"]
sha256_value = (file_info.get("hashes") or {}).get("SHA256", "").lower()
return cls(
file_name=os.path.splitext(file_name)[0],
model_name=model_data.get("name", os.path.splitext(file_name)[0]),
file_path=save_path.replace(os.sep, "/"),
size=file_info.get("sizeKB", 0) * 1024,
modified=datetime.now().timestamp(),
sha256=(file_info.get("hashes") or {}).get("SHA256", "").lower(),
sha256=sha256_value,
base_model=base_model,
preview_url="", # Will be updated after preview download
preview_nsfw_level=0, # Will be updated after preview download
@@ -204,6 +264,7 @@ class LoraMetadata(BaseModelMetadata):
civitai=version_info,
tags=tags,
modelDescription=description,
autov3=autov3_from_civitai_files(version_info, sha256_value),
)
@@ -220,6 +281,7 @@ class CheckpointMetadata(BaseModelMetadata):
"""Create CheckpointMetadata instance from Civitai version info"""
file_name = file_info.get("name", "")
base_model = determine_base_model(version_info.get("baseModel", ""))
sha256_value = (file_info.get("hashes") or {}).get("SHA256", "").lower()
sub_type = version_info.get("type", "checkpoint")
# Extract tags and description if available
@@ -237,7 +299,7 @@ class CheckpointMetadata(BaseModelMetadata):
file_path=save_path.replace(os.sep, "/"),
size=file_info.get("sizeKB", 0) * 1024,
modified=datetime.now().timestamp(),
sha256=(file_info.get("hashes") or {}).get("SHA256", "").lower(),
sha256=sha256_value,
base_model=base_model,
preview_url="", # Will be updated after preview download
preview_nsfw_level=0,
@@ -246,6 +308,7 @@ class CheckpointMetadata(BaseModelMetadata):
sub_type=sub_type,
tags=tags,
modelDescription=description,
autov3=autov3_from_civitai_files(version_info, sha256_value),
)
@@ -262,6 +325,7 @@ class EmbeddingMetadata(BaseModelMetadata):
"""Create EmbeddingMetadata instance from Civitai version info"""
file_name = file_info.get("name", "")
base_model = determine_base_model(version_info.get("baseModel", ""))
sha256_value = (file_info.get("hashes") or {}).get("SHA256", "").lower()
sub_type = version_info.get("type", "embedding")
# Extract tags and description if available
@@ -279,7 +343,7 @@ class EmbeddingMetadata(BaseModelMetadata):
file_path=save_path.replace(os.sep, "/"),
size=file_info.get("sizeKB", 0) * 1024,
modified=datetime.now().timestamp(),
sha256=(file_info.get("hashes") or {}).get("SHA256", "").lower(),
sha256=sha256_value,
base_model=base_model,
preview_url="", # Will be updated after preview download
preview_nsfw_level=0,
@@ -288,4 +352,5 @@ class EmbeddingMetadata(BaseModelMetadata):
sub_type=sub_type,
tags=tags,
modelDescription=description,
autov3=autov3_from_civitai_files(version_info, sha256_value),
)

View File

@@ -0,0 +1,350 @@
"""Tests for Autov3BackfillService."""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any, Dict, List, Optional
import pytest
from py.services.autov3_backfill_service import Autov3BackfillService
from py.services.model_cache import ModelCache
from py.services.model_hash_index import ModelHashIndex
from py.services.model_scanner import ModelScanner
from py.services.persistent_model_cache import DEFAULT_LICENSE_FLAGS, PersistentModelCache
@pytest.fixture(autouse=True)
def reset_backfill_singleton() -> None:
"""Reset the service singleton so every test starts from a fresh instance."""
Autov3BackfillService._instance = None
yield
Autov3BackfillService._instance = None
def _entry(file_path: str, sha256: str, autov3: Optional[str] = None) -> Dict[str, Any]:
return {
'file_path': file_path,
'file_name': Path(file_path).stem,
'model_name': Path(file_path).stem,
'folder': '',
'size': 1,
'modified': 1.0,
'sha256': sha256,
'autov3': autov3,
'base_model': '',
'preview_url': '',
'preview_nsfw_level': 0,
'from_civitai': True,
'favorite': False,
'notes': '',
'usage_tips': '',
'metadata_source': None,
'exclude': False,
'db_checked': False,
'last_checked_at': 0.0,
'tags': [],
'civitai': None,
'civitai_deleted': False,
'skip_metadata_refresh': False,
'license_flags': DEFAULT_LICENSE_FLAGS,
'hash_status': 'completed',
'hf_url': '',
}
class RecordingScanner:
"""Duck-typed scanner double persisting updates to a real cache."""
def __init__(
self,
model_type: str,
persistent_cache: PersistentModelCache,
entries: List[Dict[str, Any]],
) -> None:
self.model_type = model_type
self._persistent_cache = persistent_cache
self.entries: Dict[str, Dict[str, Any]] = {entry['file_path']: entry for entry in entries}
self.update_calls: List[tuple] = []
async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool:
self.update_calls.append((model_type, file_path, autov3))
entry = self.entries.get(file_path)
if entry is None:
return False
old_item = dict(entry)
new_item = dict(entry)
new_item['autov3'] = autov3
self._persistent_cache.update_single_model(model_type, new_item, old_item)
entry['autov3'] = autov3
return True
def _make_store(tmp_path: Path, monkeypatch, name: str = 'cache.sqlite') -> PersistentModelCache:
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
return PersistentModelCache(db_path=str(tmp_path / name))
def _write_file(tmp_path: Path, name: str) -> str:
path = tmp_path / name
path.write_text(name, encoding='utf-8')
return path.as_posix()
async def test_backfill_updates_models_and_self_terminates(tmp_path: Path, monkeypatch) -> None:
store = _make_store(tmp_path, monkeypatch)
path_a = _write_file(tmp_path, 'a.txt')
path_b = _write_file(tmp_path, 'b.txt')
checked = (tmp_path / 'checked.txt').as_posix()
valued = (tmp_path / 'valued.txt').as_posix()
entries = [
_entry(path_a, 'hash-a'),
_entry(path_b, 'hash-b'),
_entry(checked, 'hash-checked', autov3=''),
_entry(valued, 'hash-valued', autov3='a1b2c3d4e5f6'),
]
store.save_cache(
'dummy',
entries,
{e['sha256']: [e['file_path']] for e in entries},
[],
)
scanner = RecordingScanner('dummy', store, entries)
updated = await Autov3BackfillService.get_instance().backfill(scanner)
# Non-safetensors files yield no embedded hash, so both are marked ''.
assert updated == 2
assert set(scanner.update_calls) == {('dummy', path_a, ''), ('dummy', path_b, '')}
# Self-terminating: the driving query now finds no remaining rows.
assert store.get_models_missing_autov3('dummy') == []
persisted = store.load_cache('dummy')
items = {item['file_path']: item for item in persisted.raw_data}
assert items[path_a]['autov3'] == ''
assert items[path_b]['autov3'] == ''
# Checked-unavailable and valued rows are never recomputed or touched.
assert items[checked]['autov3'] == ''
assert items[valued]['autov3'] == 'a1b2c3d4e5f6'
async def test_backfill_skips_missing_files_without_marking(tmp_path: Path, monkeypatch) -> None:
store = _make_store(tmp_path, monkeypatch)
existing = _write_file(tmp_path, 'existing.txt')
missing = (tmp_path / 'missing.txt').as_posix()
entries = [_entry(existing, 'hash-existing'), _entry(missing, 'hash-missing')]
store.save_cache(
'dummy',
entries,
{'hash-existing': [existing], 'hash-missing': [missing]},
[],
)
scanner = RecordingScanner('dummy', store, entries)
updated = await Autov3BackfillService.get_instance().backfill(scanner)
assert updated == 1
assert scanner.update_calls == [('dummy', existing, '')]
# The missing row was not marked, so it still appears in the query.
assert store.get_models_missing_autov3('dummy') == [missing]
async def test_backfill_returns_zero_when_same_type_already_running(tmp_path: Path, monkeypatch) -> None:
store = _make_store(tmp_path, monkeypatch)
scanner = RecordingScanner('dummy', store, [])
service = Autov3BackfillService.get_instance()
service._running_types = {'dummy'}
try:
assert await service.backfill(scanner) == 0
finally:
service._running_types = set()
assert scanner.update_calls == []
async def test_backfill_runs_concurrently_for_different_model_types(tmp_path: Path, monkeypatch) -> None:
"""Scanners initialize in parallel (lora_manager.py), so a backfill for one
model type must not skip another type's backfill."""
store = _make_store(tmp_path, monkeypatch)
lora_file = _write_file(tmp_path, 'lora.txt')
ckpt_file = _write_file(tmp_path, 'ckpt.txt')
store.save_cache(
'lora',
[_entry(lora_file, 'hash-lora')],
{'hash-lora': [lora_file]},
[],
)
store.save_cache(
'checkpoint',
[_entry(ckpt_file, 'hash-ckpt')],
{'hash-ckpt': [ckpt_file]},
[],
)
lora_scanner = RecordingScanner('lora', store, [_entry(lora_file, 'hash-lora')])
ckpt_scanner = RecordingScanner('checkpoint', store, [_entry(ckpt_file, 'hash-ckpt')])
service = Autov3BackfillService.get_instance()
service._running_types = {'checkpoint'} # Simulate a checkpoint backfill in flight
try:
# The lora backfill must still run while checkpoint is in progress.
assert await service.backfill(lora_scanner) == 1
assert lora_scanner.update_calls == [('lora', lora_file, '')]
finally:
service._running_types = set()
async def test_backfill_never_raises_on_failure(tmp_path: Path, monkeypatch) -> None:
store = _make_store(tmp_path, monkeypatch)
existing = _write_file(tmp_path, 'boom.txt')
class RaisingScanner(RecordingScanner):
async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool:
raise RuntimeError('boom')
entries = [_entry(existing, 'hash-boom')]
store.save_cache('dummy', entries, {'hash-boom': [existing]}, [])
scanner = RaisingScanner('dummy', store, entries)
updated = await Autov3BackfillService.get_instance().backfill(scanner)
assert updated == 0
async def test_backfill_uses_default_cache_when_scanner_has_none(tmp_path: Path, monkeypatch) -> None:
store = _make_store(tmp_path, monkeypatch)
existing = _write_file(tmp_path, 'model.txt')
entries = [_entry(existing, 'hash-x')]
store.save_cache('dummy', entries, {'hash-x': [existing]}, [])
from py.services import persistent_model_cache as pmc_module
monkeypatch.setattr(pmc_module, 'get_persistent_cache', lambda: store)
class BareScanner:
model_type = 'dummy'
async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool:
entry = next(e for e in entries if e['file_path'] == file_path)
old_item = dict(entry)
new_item = dict(entry)
new_item['autov3'] = autov3
store.update_single_model(model_type, new_item, old_item)
return True
updated = await Autov3BackfillService.get_instance().backfill(BareScanner())
assert updated == 1
assert store.get_models_missing_autov3('dummy') == []
async def test_backfill_idempotent_second_run_is_noop(tmp_path: Path, monkeypatch) -> None:
store = _make_store(tmp_path, monkeypatch)
existing = _write_file(tmp_path, 'idem.txt')
entries = [_entry(existing, 'hash-idem')]
store.save_cache('dummy', entries, {'hash-idem': [existing]}, [])
scanner = RecordingScanner('dummy', store, entries)
service = Autov3BackfillService.get_instance()
assert await service.backfill(scanner) == 1
# A re-run has nothing left to do.
assert await service.backfill(scanner) == 0
assert len(scanner.update_calls) == 1
async def test_backfill_end_to_end_through_scanner_lazy_import(tmp_path: Path, monkeypatch) -> None:
"""Drive the scanner's lazy-import trigger (`_run_autov3_backfill`) end to end."""
store = _make_store(tmp_path, monkeypatch)
path_a = _write_file(tmp_path, 'alpha.txt')
path_b = _write_file(tmp_path, 'beta.txt')
entries = [_entry(path_a, 'hash-alpha'), _entry(path_b, 'hash-beta')]
store.save_cache(
'dummy',
entries,
{'hash-alpha': [path_a], 'hash-beta': [path_b]},
[],
)
class RealScanner(ModelScanner):
def __init__(self) -> None:
self.model_type = 'dummy'
self._persistent_cache = store
self._cache = ModelCache(raw_data=[dict(e) for e in entries], folders=[])
self._hash_index = ModelHashIndex()
await RealScanner()._run_autov3_backfill()
assert store.get_models_missing_autov3('dummy') == []
persisted = store.load_cache('dummy')
items = {item['file_path']: item for item in persisted.raw_data}
assert items[path_a]['autov3'] == ''
assert items[path_b]['autov3'] == ''
async def test_backfill_prefers_civitai_autov3_from_sidecar(tmp_path: Path, monkeypatch) -> None:
"""Backfill uses the Civitai AutoV3 for the SHA256-matching file when the
sidecar carries Civitai metadata, even if the file itself has no embedded
header hash (the checkpoint case)."""
store = _make_store(tmp_path, monkeypatch)
path = _write_file(tmp_path, 'ckpt.txt') # non-safetensors: no header hash
sidecar = tmp_path / 'ckpt.metadata.json'
sidecar.write_text(
json.dumps({
"sha256": "hash-ckpt",
"civitai": {
"files": [
{"name": "other.safetensors", "hashes": {"SHA256": "zzz999"}},
{"name": "ckpt.safetensors", "hashes": {"SHA256": "HASH-CKPT", "AutoV3": "ABCDEF1234567890"}},
]
},
}),
encoding='utf-8',
)
store.save_cache('dummy', [_entry(path, 'hash-ckpt')], {'hash-ckpt': [path]}, [])
scanner = RecordingScanner('dummy', store, [_entry(path, 'hash-ckpt')])
updated = await Autov3BackfillService.get_instance().backfill(scanner)
assert updated == 1
assert scanner.update_calls == [('dummy', path, 'abcdef123456')]
persisted = store.load_cache('dummy')
items = {item['file_path']: item for item in persisted.raw_data}
assert items[path]['autov3'] == 'abcdef123456'
# Self-terminating: the row is marked and the driving query empties.
assert store.get_models_missing_autov3('dummy') == []
async def test_backfill_falls_back_to_header_when_sidecar_has_no_match(tmp_path: Path, monkeypatch) -> None:
"""When the sidecar's Civitai files do not contain a SHA256 match, the
backfill falls back to the embedded header hash ('' for non-safetensors)."""
store = _make_store(tmp_path, monkeypatch)
path = _write_file(tmp_path, 'plain.txt')
sidecar = tmp_path / 'plain.metadata.json'
sidecar.write_text(
json.dumps({
"sha256": "hash-plain",
"civitai": {"files": [{"name": "other.safetensors", "hashes": {"SHA256": "zzz999", "AutoV3": "ABCDEF123456"}}]},
}),
encoding='utf-8',
)
store.save_cache('dummy', [_entry(path, 'hash-plain')], {'hash-plain': [path]}, [])
scanner = RecordingScanner('dummy', store, [_entry(path, 'hash-plain')])
updated = await Autov3BackfillService.get_instance().backfill(scanner)
assert updated == 1
assert scanner.update_calls == [('dummy', path, '')]

View File

@@ -1318,3 +1318,70 @@ class TestHfGroupKey:
"hf_url": "https://huggingface.co/user/repo",
}
assert BaseModelService._extract_group_key(item) == "hf:user/repo"
class TestApplyHashFilters:
"""_apply_hash_filters matches items by SHA256 or non-empty AutoV3."""
def _make_service(self):
return DummyService(model_type="stub", scanner=object(), metadata_class=BaseModelMetadata)
@pytest.mark.asyncio
async def test_matches_item_by_autov3(self):
service = self._make_service()
data = [
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": "abcdef123456"},
{"file_path": "/m/two.safetensors", "sha256": "b" * 64, "autov3": ""},
]
result = await service._apply_hash_filters(data, {"single_hash": "ABCDEF123456"})
assert [item["file_path"] for item in result] == ["/m/one.safetensors"]
@pytest.mark.asyncio
async def test_matches_item_by_sha256(self):
service = self._make_service()
data = [
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": ""},
]
result = await service._apply_hash_filters(data, {"single_hash": "A" * 64})
assert [item["file_path"] for item in result] == ["/m/one.safetensors"]
@pytest.mark.asyncio
async def test_empty_or_absent_autov3_never_matches(self):
service = self._make_service()
data = [
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": ""},
{"file_path": "/m/two.safetensors", "sha256": "b" * 64},
]
result = await service._apply_hash_filters(data, {"single_hash": "cdef123456ab"})
assert result == []
@pytest.mark.asyncio
async def test_multiple_hashes_match_autov3_and_sha256(self):
service = self._make_service()
data = [
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": "abcdef123456"},
{"file_path": "/m/two.safetensors", "sha256": "b" * 64, "autov3": ""},
]
result = await service._apply_hash_filters(
data, {"multiple_hashes": ["abcdef123456", "c" * 64]}
)
assert [item["file_path"] for item in result] == ["/m/one.safetensors"]
@pytest.mark.asyncio
async def test_no_hash_filters_returns_data_unchanged(self):
service = self._make_service()
data = [
{"file_path": "/m/one.safetensors", "sha256": "a" * 64, "autov3": "abcdef123456"},
]
result = await service._apply_hash_filters(data, {})
assert result == data

View File

@@ -321,3 +321,91 @@ class TestCacheEntryValidator:
assert result.is_valid is True
assert result.repaired is False
class TestAutov3Validation:
"""AutoV3 optional-field validation semantics."""
def _entry(self, **overrides):
# Fully-populated entry so that autov3 is the only candidate repair.
entry = {
'file_path': '/models/test.safetensors',
'sha256': 'abc123',
'file_name': 'test.safetensors',
'model_name': 'Test Model',
'folder': 'test_folder',
'size': 1024,
'modified': 1234567890.0,
'tags': ['tag1'],
'preview_url': 'http://example.com/preview.jpg',
'base_model': 'SD1.5',
'from_civitai': True,
'favorite': True,
'exclude': False,
'db_checked': True,
'preview_nsfw_level': 1,
'notes': 'Test notes',
'usage_tips': 'Test tips',
'hash_status': 'completed',
}
entry.update(overrides)
return entry
def test_validate_valid_autov3_normalized_to_lowercase(self):
"""Uppercase 12-hex autov3 is normalized to lowercase under auto_repair."""
result = CacheEntryValidator.validate(
self._entry(autov3='ABCDEF123456'), auto_repair=True
)
assert result.is_valid is True
assert result.entry['autov3'] == 'abcdef123456'
assert result.repaired is True
def test_validate_autov3_empty_string_is_valid(self):
"""Empty autov3 means checked-but-unavailable and is valid."""
result = CacheEntryValidator.validate(
self._entry(autov3=''), auto_repair=False
)
assert result.is_valid is True
assert result.repaired is False
def test_validate_autov3_none_is_valid_and_not_counted_as_repair(self):
"""autov3 None (not checked) is valid and is NOT counted as a repair."""
result = CacheEntryValidator.validate(
self._entry(autov3=None), auto_repair=True
)
assert result.is_valid is True
assert result.repaired is False
assert result.entry['autov3'] is None
def test_validate_absent_autov3_is_valid_and_not_counted_as_repair(self):
"""A missing autov3 field is valid and is NOT counted as a repair."""
result = CacheEntryValidator.validate(self._entry(), auto_repair=True)
assert result.is_valid is True
assert result.repaired is False
assert 'autov3' not in result.entry
def test_validate_short_autov3_still_valid_and_repaired_to_none(self):
"""A malformed autov3 does not invalidate the entry (optional field);
with auto_repair the value is repaired to None."""
result = CacheEntryValidator.validate(
self._entry(autov3='abc'), auto_repair=True
)
assert result.is_valid is True
assert result.entry['autov3'] is None
assert result.repaired is True
def test_validate_non_string_autov3_still_valid_and_repaired_to_none(self):
"""A non-string autov3 does not invalidate the entry (optional field);
with auto_repair the value is repaired to None."""
result = CacheEntryValidator.validate(
self._entry(autov3=123), auto_repair=True
)
assert result.is_valid is True
assert result.entry['autov3'] is None
assert result.repaired is True

View File

@@ -112,6 +112,68 @@ async def test_update_model_metadata_merges_and_persists():
)
@pytest.mark.asyncio
async def test_update_model_metadata_propagates_civitai_autov3():
helpers = build_service()
local = {
"sha256": "111aabbf94dd9e59c05d842fccf57bec915b2a3c237f6b54f8d614e40858d717",
"autov3": "",
"model_name": "Local",
}
remote = {
"source": "api",
"model": {"name": "Remote Model", "description": "", "tags": []},
"images": [],
"files": [
{
"name": "other.safetensors",
"hashes": {"SHA256": "ZZZ999"},
},
{
"name": "model.safetensors",
"hashes": {
"SHA256": "111aabbf94dd9e59c05d842fccf57bec915b2a3c237f6b54f8d614e40858d717",
"AutoV3": "8A582E901D7F",
},
},
],
}
result = await helpers.service.update_model_metadata(
"path/to/model.metadata.json",
local,
remote,
helpers.default_provider,
)
# Civitai-first: the '' (checked-unavailable) state is upgraded in-session
# by the freshly fetched metadata, without any header re-read.
assert result["autov3"] == "8a582e901d7f"
@pytest.mark.asyncio
async def test_update_model_metadata_keeps_autov3_without_matching_file():
helpers = build_service()
local = {"sha256": "abc123", "autov3": "", "model_name": "Local"}
remote = {
"source": "api",
"model": {"name": "Remote Model", "description": "", "tags": []},
"images": [],
"files": [{"name": "other.safetensors", "hashes": {"SHA256": "ZZZ999", "AutoV3": "ABCDEF123456"}}],
}
result = await helpers.service.update_model_metadata(
"path/to/model.metadata.json",
local,
remote,
helpers.default_provider,
)
assert result["autov3"] == ""
@pytest.mark.asyncio
async def test_fetch_and_update_model_success_updates_cache(tmp_path):
helpers = build_service()

View File

@@ -111,3 +111,145 @@ class TestModelHashIndexGetDuplicateFilenames:
index.add_entry("abc123", "/a/lora.safetensors")
assert len(index) == 1
assert index.get_duplicate_filenames() == {}
class TestModelHashIndexAutov3:
"""AutoV3 hash index behavior."""
def test_add_entry_with_autov3_supports_lookup_by_autov3(self):
index = ModelHashIndex()
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="AbCdEf123456")
assert index.has_hash("abcdef123456") is True
assert index.get_path("abcdef123456") == "/models/lora.safetensors"
assert index.get_all_autov3() == {"abcdef123456": "/models/lora.safetensors"}
def test_add_entry_without_autov3_creates_no_autov3_lookup(self):
index = ModelHashIndex()
index.add_entry("b" * 64, "/models/lora.safetensors")
assert index.has_hash("abcdef123456") is False
assert index.get_path("abcdef123456") is None
assert index.get_all_autov3() == {}
def test_add_autov3_standalone_supports_lookup(self):
index = ModelHashIndex()
index.add_autov3("cdef123456ab", "/models/only_autov3.safetensors")
assert index.has_hash("cdef123456ab") is True
assert index.get_path("cdef123456ab") == "/models/only_autov3.safetensors"
assert index.get_all_autov3() == {"cdef123456ab": "/models/only_autov3.safetensors"}
def test_remove_by_path_removes_autov3_mapping(self):
index = ModelHashIndex()
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
index.remove_by_path("/models/lora.safetensors")
assert index.has_hash("abcdef123456") is False
assert index.get_all_autov3() == {}
def test_remove_by_hash_removes_autov3_mapping(self):
index = ModelHashIndex()
sha256 = "a" * 64
index.add_entry(sha256, "/models/lora.safetensors", autov3="abcdef123456")
index.remove_by_hash(sha256)
assert index.has_hash("abcdef123456") is False
assert index.get_all_autov3() == {}
def test_clear_empties_autov3_index(self):
index = ModelHashIndex()
index.add_entry("a" * 64, "/models/a.safetensors", autov3="aaaaabbbbbcc")
index.add_entry("b" * 64, "/models/b.safetensors", autov3="dddddeeeeeff")
index.clear()
assert index.get_all_autov3() == {}
assert index.has_hash("aaaaabbbbbcc") is False
def test_same_autov3_last_write_wins(self):
index = ModelHashIndex()
index.add_entry("a" * 64, "/models/first.safetensors", autov3="abcdef123456")
index.add_entry("b" * 64, "/models/second.safetensors", autov3="abcdef123456")
assert index.get_path("abcdef123456") == "/models/second.safetensors"
assert index.get_all_autov3() == {"abcdef123456": "/models/second.safetensors"}
def test_dispatch_len_10_hits_autov2(self):
index = ModelHashIndex()
sha256 = "a" * 64
index.add_entry(sha256, "/models/lora.safetensors")
assert index.get_path(sha256[:10]) == "/models/lora.safetensors"
assert index.has_hash(sha256[:10]) is True
def test_dispatch_len_64_hits_sha256(self):
index = ModelHashIndex()
sha256 = "b" * 64
index.add_entry(sha256, "/models/lora.safetensors")
assert index.get_path(sha256) == "/models/lora.safetensors"
assert index.has_hash(sha256) is True
def test_dispatch_len_12_hits_autov3(self):
index = ModelHashIndex()
index.add_entry("c" * 64, "/models/lora.safetensors", autov3="cdef123456ab")
assert index.get_path("cdef123456ab") == "/models/lora.safetensors"
assert index.has_hash("cdef123456ab") is True
def test_add_entry_drops_stale_autov3_for_replaced_path(self):
# A file replaced in place (new content → new sha256 and new autov3)
# must not keep the old autov3 mapping — it would survive into the
# persisted snapshot and make lookups resolve the wrong file.
index = ModelHashIndex()
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
index.add_entry("b" * 64, "/models/lora.safetensors", autov3="fedcba654321")
assert index.get_path("abcdef123456") is None
assert index.has_hash("abcdef123456") is False
assert index.get_path("fedcba654321") == "/models/lora.safetensors"
assert index.get_all_autov3() == {"fedcba654321": "/models/lora.safetensors"}
def test_add_entry_without_autov3_drops_stale_mapping_for_replaced_path(self):
# Replaced file whose new content has no embedded hash: the stale
# autov3 mapping must be dropped, not left pointing at the path.
index = ModelHashIndex()
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
index.add_entry("b" * 64, "/models/lora.safetensors")
assert index.get_path("abcdef123456") is None
assert index.get_all_autov3() == {}
def test_add_entry_re_registration_with_same_autov3_is_idempotent(self):
index = ModelHashIndex()
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
index.add_entry("a" * 64, "/models/lora.safetensors", autov3="abcdef123456")
assert index.get_path("abcdef123456") == "/models/lora.safetensors"
assert index.get_all_autov3() == {"abcdef123456": "/models/lora.safetensors"}
def test_add_entry_same_sha_without_autov3_preserves_existing_mapping(self):
# A lazy-hash completion (checkpoint_scanner) re-registers the SAME
# file with the same sha256 but omits autov3. That must never clear
# the previously registered autov3 mapping.
index = ModelHashIndex()
index.add_entry("a" * 64, "/models/ckpt.safetensors", autov3="abcdef123456")
index.add_entry("a" * 64, "/models/ckpt.safetensors")
assert index.get_path("abcdef123456") == "/models/ckpt.safetensors"
assert index.get_all_autov3() == {"abcdef123456": "/models/ckpt.safetensors"}
def test_add_entry_same_sha_with_new_autov3_drops_old_mapping(self):
# Re-registration with an explicit, different autov3 (metadata
# correction) must drop the stale mapping for that path.
index = ModelHashIndex()
index.add_entry("a" * 64, "/models/ckpt.safetensors", autov3="abcdef123456")
index.add_entry("a" * 64, "/models/ckpt.safetensors", autov3="fedcba654321")
assert index.get_path("abcdef123456") is None
assert index.has_hash("abcdef123456") is False
assert index.get_path("fedcba654321") == "/models/ckpt.safetensors"
assert index.get_all_autov3() == {"fedcba654321": "/models/ckpt.safetensors"}

View File

@@ -341,3 +341,112 @@ def test_update_single_model_update_hash(tmp_path: Path, monkeypatch):
new_hash_pairs = [p for p in persisted.hash_rows if p[0] == 'new-hash']
assert len(new_hash_pairs) == 1
assert new_hash_pairs[0][1] == file_path
# ── get_models_missing_autov3 ─────────────────────────────────────────
def _autov3_entry(file_path: str, sha256: str, autov3=None) -> dict:
"""Minimal model entry for the models table (autov3 tri-state preserved)."""
return {
'file_path': file_path,
'file_name': Path(file_path).stem,
'model_name': Path(file_path).stem,
'folder': '',
'size': 1,
'modified': 1.0,
'sha256': sha256,
'autov3': autov3,
'base_model': '',
'preview_url': '',
'preview_nsfw_level': 0,
'from_civitai': True,
'favorite': False,
'notes': '',
'usage_tips': '',
'metadata_source': None,
'exclude': False,
'db_checked': False,
'last_checked_at': 0.0,
'tags': [],
'civitai': None,
'civitai_deleted': False,
'skip_metadata_refresh': False,
'license_flags': DEFAULT_LICENSE_FLAGS,
'hash_status': 'completed',
'hf_url': '',
}
def test_get_models_missing_autov3_filters_rows(tmp_path: Path, monkeypatch) -> None:
"""Only NULL-autov3 rows with a completed sha256 qualify."""
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
db_path = tmp_path / 'cache.sqlite'
store = PersistentModelCache(db_path=str(db_path))
null_path = (tmp_path / 'null.txt').as_posix()
checked_path = (tmp_path / 'checked.txt').as_posix()
valued_path = (tmp_path / 'valued.txt').as_posix()
empty_sha_path = (tmp_path / 'empty_sha.txt').as_posix()
store.save_cache(
'dummy',
[
_autov3_entry(null_path, 'hash-null'),
_autov3_entry(checked_path, 'hash-checked', autov3=''),
_autov3_entry(valued_path, 'hash-valued', autov3='a1b2c3d4e5f6'),
_autov3_entry(empty_sha_path, ''),
],
{},
[],
)
assert store.get_models_missing_autov3('dummy') == [null_path]
def test_get_models_missing_autov3_filters_by_model_type(tmp_path: Path, monkeypatch) -> None:
"""Only rows of the requested model_type are returned."""
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
db_path = tmp_path / 'cache.sqlite'
store = PersistentModelCache(db_path=str(db_path))
lora_path = (tmp_path / 'lora.txt').as_posix()
checkpoint_path = (tmp_path / 'checkpoint.txt').as_posix()
store.save_cache('lora', [_autov3_entry(lora_path, 'hash-lora')], {}, [])
store.save_cache('checkpoint', [_autov3_entry(checkpoint_path, 'hash-checkpoint')], {}, [])
assert store.get_models_missing_autov3('lora') == [lora_path]
assert store.get_models_missing_autov3('checkpoint') == [checkpoint_path]
def test_get_models_missing_autov3_empty_on_clean_db(tmp_path: Path, monkeypatch) -> None:
"""A freshly created database has no rows to backfill."""
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
store = PersistentModelCache(db_path=str(tmp_path / 'cache.sqlite'))
assert store.get_models_missing_autov3('dummy') == []
def test_get_models_missing_autov3_disabled_cache_returns_empty(tmp_path: Path, monkeypatch) -> None:
"""When the persistent cache is disabled the query is a no-op."""
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '1')
store = PersistentModelCache(db_path=str(tmp_path / 'cache.sqlite'))
assert store.get_models_missing_autov3('dummy') == []
def test_get_models_missing_autov3_self_terminates_after_marking(tmp_path: Path, monkeypatch) -> None:
"""Once a row receives a checked state it drops out of the query."""
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
db_path = tmp_path / 'cache.sqlite'
store = PersistentModelCache(db_path=str(db_path))
file_path = (tmp_path / 'm.txt').as_posix()
store.save_cache('dummy', [_autov3_entry(file_path, 'hash-m')], {}, [])
assert store.get_models_missing_autov3('dummy') == [file_path]
# Mark the row '' (checked-unavailable) and re-query.
old_item = {'file_path': file_path, 'tags': [], 'sha256': 'hash-m'}
new_item = _autov3_entry(file_path, 'hash-m', autov3='')
store.update_single_model('dummy', new_item, old_item=old_item)
assert store.get_models_missing_autov3('dummy') == []

View File

@@ -1,15 +1,26 @@
import hashlib
import json
import os
import struct
import pytest
from py.utils.constants import MAX_SAFETENSORS_HEADER_BYTES
from py.utils.file_utils import (
calculate_autov3,
calculate_sha256,
find_preview_file,
get_preview_extension,
)
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("<Q", len(header)) + header + payload)
@pytest.mark.asyncio
async def test_calculate_sha256(tmp_path):
file_path = tmp_path / "sample.bin"
@@ -49,3 +60,126 @@ def test_find_preview_file_supports_example_extension(tmp_path):
)
def test_get_preview_extension(preview_name, expected):
assert get_preview_extension(preview_name) == expected
class TestCalculateAutov3:
def test_returns_first_12_chars_from_sshs_model_hash(self, tmp_path):
file_path = tmp_path / "lora.safetensors"
_write_safetensors(file_path, {"sshs_model_hash": "abcdef1234567890abcdef"})
assert calculate_autov3(str(file_path)) == "abcdef123456"
def test_lowercases_uppercase_embedded_hash(self, tmp_path):
file_path = tmp_path / "upper.safetensors"
_write_safetensors(file_path, {"sshs_model_hash": "ABCDEF1234567890ABCDEF"})
assert calculate_autov3(str(file_path)) == "abcdef123456"
def test_returns_first_12_chars_from_modelspec_hash_sha256(self, tmp_path):
file_path = tmp_path / "model.safetensors"
_write_safetensors(file_path, {"modelspec.hash_sha256": "00112233445566778899aabb"})
assert calculate_autov3(str(file_path)) == "001122334455"
def test_prefers_sshs_model_hash_over_modelspec_hash(self, tmp_path):
file_path = tmp_path / "model.safetensors"
_write_safetensors(
file_path,
{
"sshs_model_hash": "aaaaaaaaaaaaaaaaaaaaaaaa",
"modelspec.hash_sha256": "bbbbbbbbbbbbbbbbbbbbbbbb",
},
)
assert calculate_autov3(str(file_path)) == "aaaaaaaaaaaa"
def test_returns_none_for_non_safetensors_file(self, tmp_path):
file_path = tmp_path / "plain.bin"
file_path.write_bytes(b"just some plain bytes, not a safetensors file")
assert calculate_autov3(str(file_path)) is None
def test_returns_none_for_empty_file(self, tmp_path):
file_path = tmp_path / "empty.safetensors"
file_path.write_bytes(b"")
assert calculate_autov3(str(file_path)) is None
def test_returns_none_for_header_length_shorter_than_8_bytes(self, tmp_path):
file_path = tmp_path / "short.safetensors"
file_path.write_bytes(b"\x10\x00")
assert calculate_autov3(str(file_path)) is None
def test_returns_none_for_truncated_json_header(self, tmp_path):
file_path = tmp_path / "truncated.safetensors"
file_path.write_bytes(struct.pack("<Q", 100) + b'{"__metadata__":')
assert calculate_autov3(str(file_path)) is None
def test_returns_none_when_metadata_lacks_recognized_hash(self, tmp_path):
file_path = tmp_path / "model.safetensors"
_write_safetensors(file_path, {"ss_model_name": "something-else"})
assert calculate_autov3(str(file_path)) is None
def test_returns_none_when_embedded_hash_is_too_short(self, tmp_path):
file_path = tmp_path / "model.safetensors"
_write_safetensors(file_path, {"sshs_model_hash": "abc123"})
assert calculate_autov3(str(file_path)) is None
def test_returns_none_when_embedded_hash_is_not_a_string(self, tmp_path):
file_path = tmp_path / "model.safetensors"
_write_safetensors(file_path, {"sshs_model_hash": 123456})
assert calculate_autov3(str(file_path)) is None
def test_strips_0x_prefix_from_modelspec_hash(self, tmp_path):
# OneTrainer writes modelspec.hash_sha256 with a "0x" prefix.
file_path = tmp_path / "onetrainer.safetensors"
_write_safetensors(
file_path, {"modelspec.hash_sha256": "0x1585b50b9d7d66778899aabb"}
)
assert calculate_autov3(str(file_path)) == "1585b50b9d7d"
def test_returns_none_for_empty_string_sha256_placeholder(self, tmp_path):
# Repackaging tools sometimes write the SHA256-of-empty placeholder
# instead of a real hash; it must not be treated as a valid AutoV3.
file_path = tmp_path / "broken.safetensors"
_write_safetensors(
file_path,
{"sshs_model_hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"},
)
assert calculate_autov3(str(file_path)) is None
def test_strips_0x_prefix_from_empty_hash_placeholder(self, tmp_path):
file_path = tmp_path / "onetrainer_broken.safetensors"
_write_safetensors(
file_path,
{"modelspec.hash_sha256": "0xe3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"},
)
assert calculate_autov3(str(file_path)) is None
def test_returns_none_for_header_length_over_limit(self, tmp_path):
# A crafted file whose 64-bit header length exceeds the cap must not
# trigger a giant allocation; it is rejected before any read.
file_path = tmp_path / "huge_header.safetensors"
file_path.write_bytes(
struct.pack("<Q", MAX_SAFETENSORS_HEADER_BYTES + 1) + b'{"__metadata__": {}}'
)
assert calculate_autov3(str(file_path)) is None
def test_reads_header_at_exact_limit_boundary(self, tmp_path):
# A header length exactly at the cap is still valid.
file_path = tmp_path / "at_limit.safetensors"
file_path.write_bytes(
struct.pack("<Q", MAX_SAFETENSORS_HEADER_BYTES) + b'{"__metadata__": {"sshs_model_hash": "abcdef1234567890abcdef"}}'
)
# The read returns {} because the actual bytes are shorter than the
# claimed length (short-read guard), without allocating anything huge.
assert calculate_autov3(str(file_path)) is None

View File

@@ -1,10 +1,18 @@
import json
import struct
import pytest
from py.utils.metadata_manager import MetadataManager
from py.utils.models import BaseModelMetadata
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("<Q", len(header)) + header + payload)
@pytest.mark.asyncio
async def test_base_model_metadata_sets_empty_civitai_dict():
metadata = BaseModelMetadata(
@@ -35,3 +43,35 @@ async def test_create_default_metadata_uses_empty_civitai(tmp_path):
payload = json.loads(metadata_path.read_text(encoding="utf-8"))
assert payload.get("civitai") == {}
@pytest.mark.asyncio
async def test_create_default_metadata_writes_autov3_from_safetensors(tmp_path):
model_path = tmp_path / "example.safetensors"
_write_safetensors(model_path, {"sshs_model_hash": "abcdef1234567890abcdef"})
metadata = await MetadataManager.create_default_metadata(str(model_path))
assert metadata is not None
assert metadata.autov3 == "abcdef123456"
metadata_path = model_path.with_suffix(".metadata.json")
payload = json.loads(metadata_path.read_text(encoding="utf-8"))
assert payload.get("autov3") == "abcdef123456"
@pytest.mark.asyncio
async def test_create_default_metadata_writes_autov3_null_for_non_safetensors(tmp_path):
model_path = tmp_path / "example.safetensors"
model_path.write_bytes(b"stub")
metadata = await MetadataManager.create_default_metadata(str(model_path))
assert metadata is not None
assert metadata.autov3 == ""
metadata_path = model_path.with_suffix(".metadata.json")
payload = json.loads(metadata_path.read_text(encoding="utf-8"))
assert payload.get("autov3") is None

View File

@@ -125,3 +125,200 @@ class TestLoraMetadataConsistency:
# Type is stored in civitai dict
assert metadata.civitai.get("model", {}).get("type") == "Lora"
class TestBaseModelMetadataAutov3:
"""Three-state autov3 semantics on BaseModelMetadata (None/''/12-hex)."""
def _make_dict(self, **overrides):
data = {
"file_name": "model",
"model_name": "Model",
"file_path": "/tmp/model.safetensors",
"size": 0,
"modified": 0.0,
"sha256": "deadbeef",
"base_model": "Unknown",
"preview_url": "",
}
data.update(overrides)
return data
def test_from_dict_null_autov3_normalizes_to_empty_string(self):
metadata = BaseModelMetadata.from_dict(self._make_dict(autov3=None))
assert metadata.autov3 == ""
def test_to_dict_emits_autov3_null_when_checked_unavailable(self):
metadata = BaseModelMetadata.from_dict(self._make_dict(autov3=None))
payload = metadata.to_dict()
assert "autov3" in payload
assert payload["autov3"] is None
def test_from_dict_absent_autov3_stays_none(self):
metadata = BaseModelMetadata.from_dict(self._make_dict())
assert metadata.autov3 is None
def test_to_dict_omits_autov3_when_not_checked(self):
metadata = BaseModelMetadata.from_dict(self._make_dict())
assert "autov3" not in metadata.to_dict()
def test_from_dict_preserves_lowercase_autov3_value(self):
metadata = BaseModelMetadata.from_dict(self._make_dict(autov3="abcdef123456"))
assert metadata.autov3 == "abcdef123456"
assert metadata.to_dict()["autov3"] == "abcdef123456"
@pytest.mark.parametrize("model_cls", [LoraMetadata, CheckpointMetadata, EmbeddingMetadata])
def test_from_civitai_info_extracts_autov3(self, model_cls):
# Civitai versions can ship multiple files; AutoV3 is taken from the
# file whose SHA256 matches the downloaded file.
version_info = {
"baseModel": "SDXL",
"model": {"name": "Test", "description": "", "tags": []},
"files": [
{"name": "other.safetensors", "sizeKB": 100, "hashes": {"SHA256": "zzz999"}},
{
"name": "model.safetensors",
"sizeKB": 1000,
"hashes": {"SHA256": "abc123", "AutoV3": "ABCDEF1234567890ABCDEF"},
},
],
}
file_info = {
"name": "model.safetensors",
"sizeKB": 1000,
"hashes": {"SHA256": "abc123"},
}
metadata = model_cls.from_civitai_info(version_info, file_info, "/test/model.safetensors")
assert metadata.autov3 == "abcdef123456"
@pytest.mark.parametrize(
"hashes",
[
{},
{"SHA256": "abc123"},
{"SHA256": "abc123", "AutoV3": "abc"},
{"SHA256": "abc123", "AutoV3": 123},
{"SHA256": "abc123", "AutoV3": None},
],
)
@pytest.mark.parametrize("model_cls", [LoraMetadata, CheckpointMetadata, EmbeddingMetadata])
def test_from_civitai_info_autov3_none_when_missing_or_invalid(self, model_cls, hashes):
version_info = {
"baseModel": "SDXL",
"model": {"name": "Test", "description": "", "tags": []},
"files": [{"name": "model.safetensors", "sizeKB": 1000, "hashes": hashes}],
}
file_info = {"name": "model.safetensors", "sizeKB": 1000, "hashes": {"SHA256": "abc123"}}
metadata = model_cls.from_civitai_info(version_info, file_info, "/test/model.safetensors")
assert metadata.autov3 is None
def test_autov3_from_civitai_files_matches_sha256_case_insensitively(self):
from py.utils.models import autov3_from_civitai_files
civitai = {
"files": [
{"name": "a.safetensors", "hashes": {"SHA256": "111AAA"}},
{"name": "b.safetensors", "hashes": {"SHA256": "222BBB", "AutoV3": "ABCDEF123456"}},
]
}
assert autov3_from_civitai_files(civitai, "222bbb") == "abcdef123456"
assert autov3_from_civitai_files(civitai, "111aaa") is None
assert autov3_from_civitai_files(civitai, "999999") is None
assert autov3_from_civitai_files(None, "222bbb") is None
assert autov3_from_civitai_files(civitai, "") is None
def test_autov3_from_civitai_files_ignores_files_without_sha256(self):
from py.utils.models import autov3_from_civitai_files
civitai = {
"files": [
{"name": "a.safetensors", "hashes": {"AutoV3": "ABCDEF123456"}},
{"name": "b.safetensors", "hashes": {"SHA256": "222BBB", "AutoV3": "123456ABCDEF"}},
]
}
assert autov3_from_civitai_files(civitai, "222bbb") == "123456abcdef"
def test_autov3_from_civitai_files_rejects_empty_hash_placeholder(self):
from py.utils.models import autov3_from_civitai_files
civitai = {
"files": [
{"name": "b.safetensors", "hashes": {"SHA256": "222BBB", "AutoV3": "E3B0C44298FC1C149AFB..."}},
]
}
# The empty-string SHA256 placeholder must never be adopted as a value.
assert autov3_from_civitai_files(civitai, "222bbb") is None
def test_update_civitai_info_populates_autov3_from_matching_file(self):
metadata = CheckpointMetadata(
file_name="kreamania",
model_name="Kreamania",
file_path="/test/kreamania_variant5.safetensors",
size=1000,
modified=1234567890.0,
sha256="111AABBF94DD9E59C05D842FCCF57BEC915B2A3C237F6B54F8D614E40858D717",
base_model="FLUX.1 D",
preview_url="",
autov3="",
)
version_info = {
"files": [
{
"name": "kreamania_variant5.safetensors",
"hashes": {"SHA256": "111aabbf94dd9e59c05d842fccf57bec915b2a3c237f6b54f8d614e40858d717", "AutoV3": "8A582E901D7F"},
}
]
}
metadata.update_civitai_info(version_info)
assert metadata.autov3 == "8a582e901d7f"
def test_update_civitai_info_prefers_civitai_over_header_value(self):
metadata = BaseModelMetadata(
file_name="model",
model_name="Model",
file_path="/test/model.safetensors",
size=1000,
modified=1234567890.0,
sha256="abc123",
base_model="SDXL",
preview_url="",
autov3="e3b0c44298fc", # stale header-extracted value
)
version_info = {
"files": [
{"name": "model.safetensors", "hashes": {"SHA256": "abc123", "AutoV3": "DEF456ABC789"}}
]
}
metadata.update_civitai_info(version_info)
assert metadata.autov3 == "def456abc789"
def test_update_civitai_info_keeps_autov3_when_no_match(self):
metadata = BaseModelMetadata(
file_name="model",
model_name="Model",
file_path="/test/model.safetensors",
size=1000,
modified=1234567890.0,
sha256="abc123",
base_model="SDXL",
preview_url="",
autov3="",
)
metadata.update_civitai_info({"files": [{"name": "other.safetensors", "hashes": {"SHA256": "zzz999", "AutoV3": "DEF456ABC789"}}]})
assert metadata.autov3 == ""