diff --git a/docs/metadata-json-schema.md b/docs/metadata-json-schema.md index 4fc7f037..235733c6 100644 --- a/docs/metadata-json-schema.md +++ b/docs/metadata-json-schema.md @@ -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 | --- diff --git a/py/routes/handlers/recipe_handlers.py b/py/routes/handlers/recipe_handlers.py index 68480ce7..367baa7f 100644 --- a/py/routes/handlers/recipe_handlers.py +++ b/py/routes/handlers/recipe_handlers.py @@ -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 diff --git a/py/services/autov3_backfill_service.py b/py/services/autov3_backfill_service.py new file mode 100644 index 00000000..07b3d98f --- /dev/null +++ b/py/services/autov3_backfill_service.py @@ -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) diff --git a/py/services/base_model_service.py b/py/services/base_model_service.py index 2cd61f11..524e43b6 100644 --- a/py/services/base_model_service.py +++ b/py/services/base_model_service.py @@ -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 diff --git a/py/services/cache_entry_validator.py b/py/services/cache_entry_validator.py index 753ca743..176ede09 100644 --- a/py/services/cache_entry_validator.py +++ b/py/services/cache_entry_validator.py @@ -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)): diff --git a/py/services/checkpoint_scanner.py b/py/services/checkpoint_scanner.py index 23cebc10..ae4dc8e9 100644 --- a/py/services/checkpoint_scanner.py +++ b/py/services/checkpoint_scanner.py @@ -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 diff --git a/py/services/metadata_sync_service.py b/py/services/metadata_sync_service.py index 37c547ef..d5b4af36 100644 --- a/py/services/metadata_sync_service.py +++ b/py/services/metadata_sync_service.py @@ -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", []) ) diff --git a/py/services/model_hash_index.py b/py/services/model_hash_index.py index 2b890e29..fb796fa1 100644 --- a/py/services/model_hash_index.py +++ b/py/services/model_hash_index.py @@ -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""" diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index a3425413..b06c1027 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -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. diff --git a/py/services/persistent_model_cache.py b/py/services/persistent_model_cache.py index 9d11ec6c..af6032aa 100644 --- a/py/services/persistent_model_cache.py +++ b/py/services/persistent_model_cache.py @@ -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 = ?", diff --git a/py/utils/constants.py b/py/utils/constants.py index fcb203bc..134cf875 100644 --- a/py/utils/constants.py +++ b/py/utils/constants.py @@ -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 diff --git a/py/utils/file_utils.py b/py/utils/file_utils.py index 720eb4aa..910c6c86 100644 --- a/py/utils/file_utils.py +++ b/py/utils/file_utils.py @@ -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(" 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 diff --git a/py/utils/metadata_manager.py b/py/utils/metadata_manager.py index 122de8fe..7dbe1b84 100644 --- a/py/utils/metadata_manager.py +++ b/py/utils/metadata_manager.py @@ -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) diff --git a/py/utils/models.py b/py/utils/models.py index 3b95a35c..9e029661 100644 --- a/py/utils/models.py +++ b/py/utils/models.py @@ -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), ) diff --git a/tests/services/test_autov3_backfill_service.py b/tests/services/test_autov3_backfill_service.py new file mode 100644 index 00000000..9b1325a3 --- /dev/null +++ b/tests/services/test_autov3_backfill_service.py @@ -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, '')] diff --git a/tests/services/test_base_model_service.py b/tests/services/test_base_model_service.py index 25ac6d2d..9309170e 100644 --- a/tests/services/test_base_model_service.py +++ b/tests/services/test_base_model_service.py @@ -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 diff --git a/tests/services/test_cache_entry_validator.py b/tests/services/test_cache_entry_validator.py index 58fc75ff..a1b7580e 100644 --- a/tests/services/test_cache_entry_validator.py +++ b/tests/services/test_cache_entry_validator.py @@ -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 diff --git a/tests/services/test_metadata_sync_service.py b/tests/services/test_metadata_sync_service.py index 0997932c..8fb87618 100644 --- a/tests/services/test_metadata_sync_service.py +++ b/tests/services/test_metadata_sync_service.py @@ -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() diff --git a/tests/services/test_model_hash_index.py b/tests/services/test_model_hash_index.py index 927df442..fdb592e9 100644 --- a/tests/services/test_model_hash_index.py +++ b/tests/services/test_model_hash_index.py @@ -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"} diff --git a/tests/services/test_persistent_model_cache.py b/tests/services/test_persistent_model_cache.py index c898fd78..edded38f 100644 --- a/tests/services/test_persistent_model_cache.py +++ b/tests/services/test_persistent_model_cache.py @@ -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') == [] diff --git a/tests/utils/test_file_utils.py b/tests/utils/test_file_utils.py index df75c764..250b0c4c 100644 --- a/tests/utils/test_file_utils.py +++ b/tests/utils/test_file_utils.py @@ -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("