mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-09 23:40:15 -03:00
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:
@@ -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) |
|
| `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 |
|
| `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"` |
|
| `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
|
- `preview_url` — Updated if preview file is moved/removed
|
||||||
- `sha256` — Updated during hash calculation (when `hash_status="pending"`)
|
- `sha256` — Updated during hash calculation (when `hash_status="pending"`)
|
||||||
- `hash_status` — Updated during hash calculation
|
- `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
|
- `last_checked_at` — Timestamp of scan
|
||||||
- `metadata_source` — Set based on metadata provider
|
- `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` |
|
| `metadata_source` | `null` |
|
||||||
| `last_checked_at` | `0` |
|
| `last_checked_at` | `0` |
|
||||||
| `hash_status` | `"completed"` |
|
| `hash_status` | `"completed"` |
|
||||||
|
| `autov3` | absent (not checked) or `null` (checked, no value) |
|
||||||
| `usage_tips` | `"{}"` (LoRA only) |
|
| `usage_tips` | `"{}"` (LoRA only) |
|
||||||
| `model_type` | `"checkpoint"` or `"embedding"` (not present in LoRA models) |
|
| `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 |
|
| Version | Date | Changes |
|
||||||
|---------|------|---------|
|
|---------|------|---------|
|
||||||
|
| 1.1 | 2026-08 | Added `autov3` field (CivitAI AutoV3 hash with three-state semantics) |
|
||||||
| 1.0 | 2026-03 | Initial schema documentation |
|
| 1.0 | 2026-03 | Initial schema documentation |
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|||||||
@@ -2119,20 +2119,25 @@ class RecipeManagementHandler:
|
|||||||
for item in getattr(parent_cache_data, "raw_data", []):
|
for item in getattr(parent_cache_data, "raw_data", []):
|
||||||
if item.get("sha256", "").lower() == model_hash.lower():
|
if item.get("sha256", "").lower() == model_hash.lower():
|
||||||
local_cache[model_hash.lower()] = item
|
local_cache[model_hash.lower()] = item
|
||||||
# Compute AutoV3 so the parser can also match on
|
# Register the AutoV3 hash so the parser can also
|
||||||
# that hash type (CivitAI metadata resources use
|
# match on that hash type (CivitAI metadata
|
||||||
# AutoV3).
|
# resources use AutoV3). Prefer the stored cache
|
||||||
file_path = item.get("file_path")
|
# field; only compute it when the entry has none.
|
||||||
if file_path and os.path.exists(file_path):
|
autov3 = (item.get("autov3") or "").lower()
|
||||||
try:
|
if not autov3:
|
||||||
from ...utils.file_utils import (
|
file_path = item.get("file_path")
|
||||||
calculate_autov3,
|
if file_path and os.path.exists(file_path):
|
||||||
)
|
try:
|
||||||
autov3 = calculate_autov3(file_path)
|
from ...utils.file_utils import (
|
||||||
if autov3:
|
calculate_autov3,
|
||||||
local_cache[autov3.lower()] = item
|
)
|
||||||
except Exception:
|
autov3 = (
|
||||||
pass
|
calculate_autov3(file_path) or ""
|
||||||
|
).lower()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
if autov3:
|
||||||
|
local_cache[autov3] = item
|
||||||
break
|
break
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|||||||
140
py/services/autov3_backfill_service.py
Normal file
140
py/services/autov3_backfill_service.py
Normal 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)
|
||||||
@@ -446,20 +446,31 @@ class BaseModelService(ABC):
|
|||||||
async def _apply_hash_filters(
|
async def _apply_hash_filters(
|
||||||
self, data: List[Dict], hash_filters: Dict
|
self, data: List[Dict], hash_filters: Dict
|
||||||
) -> List[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")
|
single_hash = hash_filters.get("single_hash")
|
||||||
multiple_hashes = hash_filters.get("multiple_hashes")
|
multiple_hashes = hash_filters.get("multiple_hashes")
|
||||||
|
|
||||||
if single_hash:
|
if single_hash:
|
||||||
# Filter by single hash
|
# Filter by single hash (SHA256 or AutoV3)
|
||||||
single_hash = single_hash.lower()
|
|
||||||
return [
|
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:
|
elif multiple_hashes:
|
||||||
# Filter by multiple hashes
|
# Filter by multiple hashes (SHA256 or AutoV3)
|
||||||
hash_set = set(hash.lower() for hash in multiple_hashes)
|
hash_set = {hash.lower() for hash in multiple_hashes}
|
||||||
return [item for item in data if item.get("sha256", "").lower() in hash_set]
|
return [item for item in data if matches_hash_set(item, hash_set)]
|
||||||
|
|
||||||
return data
|
return data
|
||||||
|
|
||||||
|
|||||||
@@ -59,6 +59,7 @@ class CacheEntryValidator:
|
|||||||
'notes': ('', False),
|
'notes': ('', False),
|
||||||
'usage_tips': ('', False),
|
'usage_tips': ('', False),
|
||||||
'hash_status': ('completed', False),
|
'hash_status': ('completed', False),
|
||||||
|
'autov3': (None, False),
|
||||||
}
|
}
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -119,8 +120,13 @@ class CacheEntryValidator:
|
|||||||
if is_required:
|
if is_required:
|
||||||
errors.append(f"Required field '{field_name}' is missing or None")
|
errors.append(f"Required field '{field_name}' is missing or None")
|
||||||
if auto_repair:
|
if auto_repair:
|
||||||
working_entry[field_name] = cls._get_default_copy(default_value)
|
# A missing optional field whose default is None is already
|
||||||
repaired = True
|
# 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
|
continue
|
||||||
|
|
||||||
# Validate field type and value
|
# Validate field type and value
|
||||||
@@ -175,6 +181,15 @@ class CacheEntryValidator:
|
|||||||
# that invalidates the entry, but we also don't mark it repaired.
|
# that invalidates the entry, but we also don't mark it repaired.
|
||||||
pass
|
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
|
# Determine if entry is valid
|
||||||
# Entry is valid if no critical required field errors remain after repair
|
# Entry is valid if no critical required field errors remain after repair
|
||||||
# Critical fields are file_path and sha256
|
# Critical fields are file_path and sha256
|
||||||
@@ -242,6 +257,19 @@ class CacheEntryValidator:
|
|||||||
"""
|
"""
|
||||||
expected_type = type(default_value)
|
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
|
# Special handling for numeric types
|
||||||
if expected_type == int:
|
if expected_type == int:
|
||||||
if not isinstance(value, (int, float)):
|
if not isinstance(value, (int, float)):
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ from datetime import datetime
|
|||||||
from typing import Any, Dict, List, Optional
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
from ..utils.models import CheckpointMetadata
|
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 ..utils.metadata_manager import MetadataManager
|
||||||
from ..config import config
|
from ..config import config
|
||||||
from .model_scanner import ModelScanner
|
from .model_scanner import ModelScanner
|
||||||
@@ -62,6 +62,11 @@ class CheckpointScanner(ModelScanner):
|
|||||||
# Find preview image
|
# Find preview image
|
||||||
preview_url = find_preview_file(base_name, dir_path)
|
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
|
# Create metadata WITHOUT calculating hash
|
||||||
metadata = CheckpointMetadata(
|
metadata = CheckpointMetadata(
|
||||||
file_name=base_name,
|
file_name=base_name,
|
||||||
@@ -77,6 +82,7 @@ class CheckpointScanner(ModelScanner):
|
|||||||
sub_type="checkpoint",
|
sub_type="checkpoint",
|
||||||
from_civitai=False, # Mark as local model since no hash yet
|
from_civitai=False, # Mark as local model since no hash yet
|
||||||
hash_status="pending", # Mark hash as pending
|
hash_status="pending", # Mark hash as pending
|
||||||
|
autov3=autov3 or "",
|
||||||
)
|
)
|
||||||
|
|
||||||
# Save the created metadata
|
# Save the created metadata
|
||||||
@@ -120,7 +126,11 @@ class CheckpointScanner(ModelScanner):
|
|||||||
# that queries get_hash_by_filename first) will miss on every
|
# that queries get_hash_by_filename first) will miss on every
|
||||||
# lookup and keep calling back into this method, creating a
|
# lookup and keep calling back into this method, creating a
|
||||||
# tight loop that never populates the index.
|
# 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
|
return metadata.sha256
|
||||||
|
|
||||||
async with self._hash_calculation_lock:
|
async with self._hash_calculation_lock:
|
||||||
@@ -132,7 +142,11 @@ class CheckpointScanner(ModelScanner):
|
|||||||
and metadata.hash_status == "completed"
|
and metadata.hash_status == "completed"
|
||||||
and metadata.sha256
|
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
|
return metadata.sha256
|
||||||
|
|
||||||
task = self._hash_calculation_tasks.get(real_path)
|
task = self._hash_calculation_tasks.get(real_path)
|
||||||
@@ -185,7 +199,11 @@ class CheckpointScanner(ModelScanner):
|
|||||||
if metadata.hash_status == "completed" and metadata.sha256:
|
if metadata.hash_status == "completed" and metadata.sha256:
|
||||||
# Populate the in-memory hash index even for pre-computed
|
# Populate the in-memory hash index even for pre-computed
|
||||||
# hashes, mirroring the fix in calculate_hash_for_model.
|
# 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
|
return metadata.sha256
|
||||||
|
|
||||||
# Update status to calculating
|
# Update status to calculating
|
||||||
@@ -202,7 +220,11 @@ class CheckpointScanner(ModelScanner):
|
|||||||
await MetadataManager.save_metadata(file_path, metadata)
|
await MetadataManager.save_metadata(file_path, metadata)
|
||||||
|
|
||||||
# Update hash index
|
# 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
|
# Update the in-memory cache entry so that subsequent
|
||||||
# _persist_current_cache / _save_persistent_cache calls
|
# _persist_current_cache / _save_persistent_cache calls
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ from typing import Any, Awaitable, Callable, Dict, Iterable, Optional
|
|||||||
from ..services.settings_manager import SettingsManager
|
from ..services.settings_manager import SettingsManager
|
||||||
from ..utils.civitai_utils import resolve_license_payload
|
from ..utils.civitai_utils import resolve_license_payload
|
||||||
from ..utils.model_utils import determine_base_model
|
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 .connectivity_guard import OFFLINE_FRIENDLY_MESSAGE, is_expected_offline_error
|
||||||
from .errors import RateLimitError
|
from .errors import RateLimitError
|
||||||
|
|
||||||
@@ -152,6 +153,18 @@ class MetadataSyncService:
|
|||||||
civitai_metadata.get("baseModel")
|
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(
|
await self._preview_service.ensure_preview_for_metadata(
|
||||||
metadata_path, local_metadata, civitai_metadata.get("images", [])
|
metadata_path, local_metadata, civitai_metadata.get("images", [])
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -8,11 +8,12 @@ class ModelHashIndex:
|
|||||||
self._hash_to_path: Dict[str, str] = {}
|
self._hash_to_path: Dict[str, str] = {}
|
||||||
self._filename_to_hash: Dict[str, str] = {}
|
self._filename_to_hash: Dict[str, str] = {}
|
||||||
self._autov2_to_path: Dict[str, str] = {}
|
self._autov2_to_path: Dict[str, str] = {}
|
||||||
|
self._autov3_to_path: Dict[str, str] = {}
|
||||||
# New data structures for tracking duplicates
|
# New data structures for tracking duplicates
|
||||||
self._duplicate_hashes: Dict[str, List[str]] = {} # sha256 -> list of paths
|
self._duplicate_hashes: Dict[str, List[str]] = {} # sha256 -> list of paths
|
||||||
self._duplicate_filenames: Dict[str, List[str]] = {} # filename -> 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"""
|
"""Add or update hash index entry"""
|
||||||
if not sha256 or not file_path:
|
if not sha256 or not file_path:
|
||||||
return
|
return
|
||||||
@@ -33,9 +34,13 @@ class ModelHashIndex:
|
|||||||
self._duplicate_hashes.setdefault(sha256, []).append(file_path)
|
self._duplicate_hashes.setdefault(sha256, []).append(file_path)
|
||||||
|
|
||||||
# Track duplicates by filename - FIXED LOGIC
|
# Track duplicates by filename - FIXED LOGIC
|
||||||
|
is_re_registration = False
|
||||||
if filename in self._filename_to_hash:
|
if filename in self._filename_to_hash:
|
||||||
existing_hash = self._filename_to_hash[filename]
|
existing_hash = self._filename_to_hash[filename]
|
||||||
existing_path = self._hash_to_path.get(existing_hash)
|
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 this is a different file with the same filename
|
||||||
if existing_path and existing_path != file_path:
|
if existing_path and existing_path != file_path:
|
||||||
@@ -67,6 +72,30 @@ class ModelHashIndex:
|
|||||||
# AutoV2 = first 10 chars of SHA256
|
# AutoV2 = first 10 chars of SHA256
|
||||||
if len(sha256) >= 10:
|
if len(sha256) >= 10:
|
||||||
self._autov2_to_path[sha256[:10]] = file_path
|
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:
|
def _get_filename_from_path(self, file_path: str) -> str:
|
||||||
"""Extract filename without extension from path"""
|
"""Extract filename without extension from path"""
|
||||||
@@ -167,6 +196,11 @@ class ModelHashIndex:
|
|||||||
for k in autov2_keys_to_remove:
|
for k in autov2_keys_to_remove:
|
||||||
del self._autov2_to_path[k]
|
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:
|
def remove_by_hash(self, sha256: str) -> None:
|
||||||
"""Remove entry by hash"""
|
"""Remove entry by hash"""
|
||||||
sha256 = sha256.lower()
|
sha256 = sha256.lower()
|
||||||
@@ -189,6 +223,11 @@ class ModelHashIndex:
|
|||||||
autov2_key = sha256[:10]
|
autov2_key = sha256[:10]
|
||||||
if autov2_key in self._autov2_to_path:
|
if autov2_key in self._autov2_to_path:
|
||||||
del self._autov2_to_path[autov2_key]
|
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
|
# Update filename-to-hash and duplicate filenames for all paths
|
||||||
for path_to_remove in paths_to_remove:
|
for path_to_remove in paths_to_remove:
|
||||||
@@ -209,22 +248,26 @@ class ModelHashIndex:
|
|||||||
del self._duplicate_filenames[fname]
|
del self._duplicate_filenames[fname]
|
||||||
|
|
||||||
def has_hash(self, hash_value: str) -> bool:
|
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()
|
normalized = hash_value.lower()
|
||||||
if normalized in self._hash_to_path:
|
if normalized in self._hash_to_path:
|
||||||
return True
|
return True
|
||||||
if len(normalized) == 10:
|
if len(normalized) == 10:
|
||||||
return normalized in self._autov2_to_path
|
return normalized in self._autov2_to_path
|
||||||
|
if len(normalized) == 12:
|
||||||
|
return normalized in self._autov3_to_path
|
||||||
return False
|
return False
|
||||||
|
|
||||||
def get_path(self, hash_value: str) -> Optional[str]:
|
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()
|
normalized = hash_value.lower()
|
||||||
path = self._hash_to_path.get(normalized)
|
path = self._hash_to_path.get(normalized)
|
||||||
if path is not None:
|
if path is not None:
|
||||||
return path
|
return path
|
||||||
if len(normalized) == 10:
|
if len(normalized) == 10:
|
||||||
return self._autov2_to_path.get(normalized)
|
return self._autov2_to_path.get(normalized)
|
||||||
|
if len(normalized) == 12:
|
||||||
|
return self._autov3_to_path.get(normalized)
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def get_hash(self, file_path: str) -> Optional[str]:
|
def get_hash(self, file_path: str) -> Optional[str]:
|
||||||
@@ -243,6 +286,7 @@ class ModelHashIndex:
|
|||||||
self._hash_to_path.clear()
|
self._hash_to_path.clear()
|
||||||
self._filename_to_hash.clear()
|
self._filename_to_hash.clear()
|
||||||
self._autov2_to_path.clear()
|
self._autov2_to_path.clear()
|
||||||
|
self._autov3_to_path.clear()
|
||||||
self._duplicate_hashes.clear()
|
self._duplicate_hashes.clear()
|
||||||
self._duplicate_filenames.clear()
|
self._duplicate_filenames.clear()
|
||||||
|
|
||||||
@@ -253,6 +297,10 @@ class ModelHashIndex:
|
|||||||
def get_all_filenames(self) -> Set[str]:
|
def get_all_filenames(self) -> Set[str]:
|
||||||
"""Get all filenames in the index"""
|
"""Get all filenames in the index"""
|
||||||
return set(self._filename_to_hash.keys())
|
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]]:
|
def get_duplicate_hashes(self) -> Dict[str, List[str]]:
|
||||||
"""Get dictionary of duplicate hashes and their paths"""
|
"""Get dictionary of duplicate hashes and their paths"""
|
||||||
|
|||||||
@@ -7,9 +7,9 @@ import shutil
|
|||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Set, Type, Union
|
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 ..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.metadata_manager import MetadataManager
|
||||||
from ..utils.civitai_utils import resolve_license_info
|
from ..utils.civitai_utils import resolve_license_info
|
||||||
from .model_cache import ModelCache
|
from .model_cache import ModelCache
|
||||||
@@ -86,6 +86,7 @@ class ModelScanner:
|
|||||||
self._persistent_cache = get_persistent_cache()
|
self._persistent_cache = get_persistent_cache()
|
||||||
self._name_display_mode = self._resolve_name_display_mode()
|
self._name_display_mode = self._resolve_name_display_mode()
|
||||||
self._cancel_requested = False # Flag for cancellation
|
self._cancel_requested = False # Flag for cancellation
|
||||||
|
self._autov3_backfill_scheduled = False # One-time AutoV3 backfill trigger per process
|
||||||
try:
|
try:
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
except RuntimeError:
|
except RuntimeError:
|
||||||
@@ -225,6 +226,19 @@ class ModelScanner:
|
|||||||
if not isinstance(notes, str):
|
if not isinstance(notes, str):
|
||||||
notes = str(notes)
|
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] = {
|
entry: Dict[str, Any] = {
|
||||||
'file_path': normalized_path,
|
'file_path': normalized_path,
|
||||||
# file_name is always stored WITHOUT extension (e.g. "OWSMianne_ANIMA_V1",
|
# 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),
|
'size': int(get_value('size', 0) or 0),
|
||||||
'modified': float(get_value('modified', 0.0) or 0.0),
|
'modified': float(get_value('modified', 0.0) or 0.0),
|
||||||
'sha256': (get_value('sha256', '') or '').lower(),
|
'sha256': (get_value('sha256', '') or '').lower(),
|
||||||
|
'autov3': entry_autov3,
|
||||||
'base_model': get_value('base_model', '') or '',
|
'base_model': get_value('base_model', '') or '',
|
||||||
'preview_url': preview_url,
|
'preview_url': preview_url,
|
||||||
'preview_nsfw_level': int(get_value('preview_nsfw_level', 0) or 0),
|
'preview_nsfw_level': int(get_value('preview_nsfw_level', 0) or 0),
|
||||||
@@ -473,6 +488,13 @@ class ModelScanner:
|
|||||||
if sha_value and path:
|
if sha_value and path:
|
||||||
hash_index.add_entry(sha_value.lower(), 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] = {}
|
tags_count: Dict[str, int] = {}
|
||||||
adjusted_raw_data: List[Dict[str, Any]] = []
|
adjusted_raw_data: List[Dict[str, Any]] = []
|
||||||
for item in persisted.raw_data:
|
for item in persisted.raw_data:
|
||||||
@@ -541,8 +563,30 @@ class ModelScanner:
|
|||||||
'scanner_type': self.model_type,
|
'scanner_type': self.model_type,
|
||||||
'pageType': page_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
|
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:
|
async def _save_persistent_cache(self, scan_result: CacheBuildResult) -> None:
|
||||||
if not scan_result or not getattr(self, '_persistent_cache', None):
|
if not scan_result or not getattr(self, '_persistent_cache', None):
|
||||||
return
|
return
|
||||||
@@ -555,6 +599,7 @@ class ModelScanner:
|
|||||||
return
|
return
|
||||||
|
|
||||||
hash_snapshot = self._build_hash_index_snapshot(scan_result.hash_index)
|
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()
|
loop = asyncio.get_event_loop()
|
||||||
try:
|
try:
|
||||||
await loop.run_in_executor(
|
await loop.run_in_executor(
|
||||||
@@ -563,7 +608,8 @@ class ModelScanner:
|
|||||||
self.model_type,
|
self.model_type,
|
||||||
list(scan_result.raw_data),
|
list(scan_result.raw_data),
|
||||||
hash_snapshot,
|
hash_snapshot,
|
||||||
list(scan_result.excluded_models)
|
list(scan_result.excluded_models),
|
||||||
|
autov3_snapshot,
|
||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
logger.warning("%s Scanner: Failed to persist cache: %s", self.model_type.capitalize(), exc)
|
logger.warning("%s Scanner: Failed to persist cache: %s", self.model_type.capitalize(), exc)
|
||||||
@@ -589,6 +635,20 @@ class ModelScanner:
|
|||||||
bucket.append(path)
|
bucket.append(path)
|
||||||
return snapshot
|
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:
|
async def _persist_current_cache(self) -> None:
|
||||||
if self._cache is None or not getattr(self, '_persistent_cache', None):
|
if self._cache is None or not getattr(self, '_persistent_cache', None):
|
||||||
return
|
return
|
||||||
@@ -880,7 +940,11 @@ class ModelScanner:
|
|||||||
|
|
||||||
# Update hash index if available
|
# Update hash index if available
|
||||||
if 'sha256' in model_data and 'file_path' in model_data:
|
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
|
# Update tags count
|
||||||
if 'tags' in model_data and model_data['tags']:
|
if 'tags' in model_data and model_data['tags']:
|
||||||
@@ -1130,6 +1194,36 @@ class ModelScanner:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to compute SHA256 for {file_path}: {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
|
# Skip excluded models
|
||||||
if model_data.get('exclude', False):
|
if model_data.get('exclude', False):
|
||||||
excluded_models.append(model_data['file_path'])
|
excluded_models.append(model_data['file_path'])
|
||||||
@@ -1322,7 +1416,7 @@ class ModelScanner:
|
|||||||
sha_value = result.get('sha256')
|
sha_value = result.get('sha256')
|
||||||
model_path = result.get('file_path')
|
model_path = result.get('file_path')
|
||||||
if sha_value and model_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 []:
|
for tag in result.get('tags') or []:
|
||||||
tags_count[tag] = tags_count.get(tag, 0) + 1
|
tags_count[tag] = tags_count.get(tag, 0) + 1
|
||||||
@@ -1391,7 +1485,11 @@ class ModelScanner:
|
|||||||
await self._cache.resort()
|
await self._cache.resort()
|
||||||
|
|
||||||
# Update the hash index
|
# 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()
|
await self._persist_current_cache()
|
||||||
return True
|
return True
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -1578,7 +1676,11 @@ class ModelScanner:
|
|||||||
|
|
||||||
sha_value = cache_entry.get('sha256')
|
sha_value = cache_entry.get('sha256')
|
||||||
if sha_value:
|
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)
|
all_folders = set(item['folder'] for item in cache.raw_data)
|
||||||
cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
|
cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
|
||||||
@@ -1736,7 +1838,11 @@ class ModelScanner:
|
|||||||
if old_sha:
|
if old_sha:
|
||||||
self._hash_index.remove_by_path(file_path)
|
self._hash_index.remove_by_path(file_path)
|
||||||
if new_sha:
|
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 ----
|
# ---- Incremental version index update ----
|
||||||
new_civitai = desired_entry.get("civitai")
|
new_civitai = desired_entry.get("civitai")
|
||||||
@@ -1787,6 +1893,74 @@ class ModelScanner:
|
|||||||
|
|
||||||
return True
|
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
|
@staticmethod
|
||||||
def _cache_entries_differ(a: Dict[str, Any], b: Dict[str, Any]) -> bool:
|
def _cache_entries_differ(a: Dict[str, Any], b: Dict[str, Any]) -> bool:
|
||||||
"""Return ``True`` when two cache-entry dicts differ in any field.
|
"""Return ``True`` when two cache-entry dicts differ in any field.
|
||||||
|
|||||||
@@ -3,7 +3,7 @@ import logging
|
|||||||
import os
|
import os
|
||||||
import sqlite3
|
import sqlite3
|
||||||
import threading
|
import threading
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass, field
|
||||||
from typing import Dict, List, Mapping, Optional, Sequence, Tuple
|
from typing import Dict, List, Mapping, Optional, Sequence, Tuple
|
||||||
|
|
||||||
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
|
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
|
||||||
@@ -18,6 +18,7 @@ class PersistedCacheData:
|
|||||||
raw_data: List[Dict]
|
raw_data: List[Dict]
|
||||||
hash_rows: List[Tuple[str, str]]
|
hash_rows: List[Tuple[str, str]]
|
||||||
excluded_models: List[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.
|
DEFAULT_LICENSE_FLAGS = 127 # 127 (0b1111111) encodes default CivitAI permissions with all commercial modes enabled.
|
||||||
@@ -36,6 +37,7 @@ class PersistentModelCache:
|
|||||||
"size",
|
"size",
|
||||||
"modified",
|
"modified",
|
||||||
"sha256",
|
"sha256",
|
||||||
|
"autov3",
|
||||||
"base_model",
|
"base_model",
|
||||||
"preview_url",
|
"preview_url",
|
||||||
"preview_nsfw_level",
|
"preview_nsfw_level",
|
||||||
@@ -118,6 +120,10 @@ class PersistentModelCache:
|
|||||||
"SELECT sha256, file_path FROM hash_index WHERE model_type = ?",
|
"SELECT sha256, file_path FROM hash_index WHERE model_type = ?",
|
||||||
(model_type,),
|
(model_type,),
|
||||||
).fetchall()
|
).fetchall()
|
||||||
|
autov3_rows = conn.execute(
|
||||||
|
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
|
||||||
|
(model_type,),
|
||||||
|
).fetchall()
|
||||||
excluded = conn.execute(
|
excluded = conn.execute(
|
||||||
"SELECT file_path FROM excluded_models WHERE model_type = ?",
|
"SELECT file_path FROM excluded_models WHERE model_type = ?",
|
||||||
(model_type,),
|
(model_type,),
|
||||||
@@ -191,6 +197,8 @@ class PersistentModelCache:
|
|||||||
"hash_status": row["hash_status"] or "completed",
|
"hash_status": row["hash_status"] or "completed",
|
||||||
"hf_url": row["hf_url"] or "",
|
"hf_url": row["hf_url"] or "",
|
||||||
}
|
}
|
||||||
|
if row["autov3"] is not None:
|
||||||
|
item["autov3"] = (row["autov3"] or "").lower()
|
||||||
raw_data.append(item)
|
raw_data.append(item)
|
||||||
|
|
||||||
hash_pairs = [(entry["sha256"].lower(), entry["file_path"]) for entry in hash_rows if entry["sha256"]]
|
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:
|
if sha_value:
|
||||||
hash_pairs.append((sha_value.lower(), item["file_path"]))
|
hash_pairs.append((sha_value.lower(), item["file_path"]))
|
||||||
|
|
||||||
excluded_paths = [row["file_path"] for row in excluded]
|
autov3_pairs = [
|
||||||
return PersistedCacheData(raw_data=raw_data, hash_rows=hash_pairs, excluded_models=excluded_paths)
|
(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():
|
if not self.is_enabled():
|
||||||
return
|
return
|
||||||
if not self._schema_initialized:
|
if not self._schema_initialized:
|
||||||
@@ -251,6 +270,10 @@ class PersistentModelCache:
|
|||||||
"DELETE FROM hash_index WHERE model_type = ? AND file_path = ?",
|
"DELETE FROM hash_index WHERE model_type = ? AND file_path = ?",
|
||||||
to_remove_models,
|
to_remove_models,
|
||||||
)
|
)
|
||||||
|
conn.executemany(
|
||||||
|
"DELETE FROM autov3_index WHERE model_type = ? AND file_path = ?",
|
||||||
|
to_remove_models,
|
||||||
|
)
|
||||||
conn.executemany(
|
conn.executemany(
|
||||||
"DELETE FROM excluded_models WHERE model_type = ? AND file_path = ?",
|
"DELETE FROM excluded_models WHERE model_type = ? AND file_path = ?",
|
||||||
to_remove_models,
|
to_remove_models,
|
||||||
@@ -373,6 +396,52 @@ class PersistentModelCache:
|
|||||||
hash_inserts,
|
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(
|
existing_excluded_rows = conn.execute(
|
||||||
"SELECT file_path FROM excluded_models WHERE model_type = ?",
|
"SELECT file_path FROM excluded_models WHERE model_type = ?",
|
||||||
(model_type,),
|
(model_type,),
|
||||||
@@ -435,6 +504,7 @@ class PersistentModelCache:
|
|||||||
size INTEGER,
|
size INTEGER,
|
||||||
modified REAL,
|
modified REAL,
|
||||||
sha256 TEXT,
|
sha256 TEXT,
|
||||||
|
autov3 TEXT,
|
||||||
base_model TEXT,
|
base_model TEXT,
|
||||||
preview_url TEXT,
|
preview_url TEXT,
|
||||||
preview_nsfw_level INTEGER,
|
preview_nsfw_level INTEGER,
|
||||||
@@ -472,6 +542,13 @@ class PersistentModelCache:
|
|||||||
PRIMARY KEY (model_type, sha256, file_path)
|
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 (
|
CREATE TABLE IF NOT EXISTS excluded_models (
|
||||||
model_type TEXT NOT NULL,
|
model_type TEXT NOT NULL,
|
||||||
file_path TEXT NOT NULL,
|
file_path TEXT NOT NULL,
|
||||||
@@ -504,6 +581,7 @@ class PersistentModelCache:
|
|||||||
"license_flags": f"INTEGER DEFAULT {DEFAULT_LICENSE_FLAGS}",
|
"license_flags": f"INTEGER DEFAULT {DEFAULT_LICENSE_FLAGS}",
|
||||||
"hash_status": "TEXT DEFAULT 'completed'",
|
"hash_status": "TEXT DEFAULT 'completed'",
|
||||||
"hf_url": "TEXT DEFAULT ''",
|
"hf_url": "TEXT DEFAULT ''",
|
||||||
|
"autov3": "TEXT",
|
||||||
}
|
}
|
||||||
|
|
||||||
for column, definition in required_columns.items():
|
for column, definition in required_columns.items():
|
||||||
@@ -549,6 +627,12 @@ class PersistentModelCache:
|
|||||||
if license_flags is None:
|
if license_flags is None:
|
||||||
license_flags = DEFAULT_LICENSE_FLAGS
|
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 (
|
return (
|
||||||
model_type,
|
model_type,
|
||||||
item.get("file_path"),
|
item.get("file_path"),
|
||||||
@@ -558,6 +642,7 @@ class PersistentModelCache:
|
|||||||
int(item.get("size") or 0),
|
int(item.get("size") or 0),
|
||||||
float(item.get("modified") or 0.0),
|
float(item.get("modified") or 0.0),
|
||||||
(item.get("sha256") or "").lower() or None,
|
(item.get("sha256") or "").lower() or None,
|
||||||
|
autov3_column,
|
||||||
item.get("base_model") or "",
|
item.get("base_model") or "",
|
||||||
item.get("preview_url") or "",
|
item.get("preview_url") or "",
|
||||||
int(item.get("preview_nsfw_level") or 0),
|
int(item.get("preview_nsfw_level") or 0),
|
||||||
@@ -663,6 +748,25 @@ class PersistentModelCache:
|
|||||||
(model_type, new_sha, file_path),
|
(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")
|
conn.execute("COMMIT")
|
||||||
except Exception:
|
except Exception:
|
||||||
conn.execute("ROLLBACK")
|
conn.execute("ROLLBACK")
|
||||||
@@ -676,6 +780,40 @@ class PersistentModelCache:
|
|||||||
exc,
|
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]]:
|
def _load_tags(self, conn: sqlite3.Connection, model_type: str) -> Dict[str, List[str]]:
|
||||||
tag_rows = conn.execute(
|
tag_rows = conn.execute(
|
||||||
"SELECT file_path, tag FROM model_tags WHERE model_type = ?",
|
"SELECT file_path, tag FROM model_tags WHERE model_type = ?",
|
||||||
|
|||||||
@@ -80,6 +80,17 @@ CIVITAI_USER_MODEL_TYPES = [
|
|||||||
# Default chunk size in megabytes used for hashing large files.
|
# Default chunk size in megabytes used for hashing large files.
|
||||||
DEFAULT_HASH_CHUNK_SIZE_MB = 4
|
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 settings
|
||||||
AUTO_ORGANIZE_BATCH_SIZE = (
|
AUTO_ORGANIZE_BATCH_SIZE = (
|
||||||
50 # Process models in batches to avoid overwhelming the system
|
50 # Process models in batches to avoid overwhelming the system
|
||||||
|
|||||||
@@ -9,6 +9,8 @@ from typing import Any
|
|||||||
from .constants import (
|
from .constants import (
|
||||||
CARD_PREVIEW_WIDTH,
|
CARD_PREVIEW_WIDTH,
|
||||||
DEFAULT_HASH_CHUNK_SIZE_MB,
|
DEFAULT_HASH_CHUNK_SIZE_MB,
|
||||||
|
INVALID_AUTOV3_EMPTY_HASH,
|
||||||
|
MAX_SAFETENSORS_HEADER_BYTES,
|
||||||
PREVIEW_EXTENSIONS,
|
PREVIEW_EXTENSIONS,
|
||||||
)
|
)
|
||||||
from .exif_utils import ExifUtils
|
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:
|
if len(header_len_bytes) < 8:
|
||||||
return {}
|
return {}
|
||||||
header_len = struct.unpack("<Q", header_len_bytes)[0]
|
header_len = struct.unpack("<Q", header_len_bytes)[0]
|
||||||
|
if header_len > MAX_SAFETENSORS_HEADER_BYTES:
|
||||||
|
return {}
|
||||||
header_bytes = f.read(header_len)
|
header_bytes = f.read(header_len)
|
||||||
if len(header_bytes) < header_len:
|
if len(header_bytes) < header_len:
|
||||||
return {}
|
return {}
|
||||||
@@ -123,8 +127,16 @@ def calculate_autov3(file_path: str) -> str | None:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
embedded_hash = metadata.get("sshs_model_hash") or metadata.get("modelspec.hash_sha256")
|
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:
|
if embedded_hash and isinstance(embedded_hash, str):
|
||||||
return embedded_hash[:12]
|
# 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
|
return None
|
||||||
|
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ import time
|
|||||||
from typing import Any, Dict, Optional, Type, Union
|
from typing import Any, Dict, Optional, Type, Union
|
||||||
|
|
||||||
from .models import BaseModelMetadata, LoraMetadata
|
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
|
from .lora_metadata import extract_lora_metadata, extract_checkpoint_metadata
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -210,6 +210,11 @@ class MetadataManager:
|
|||||||
hash_duration = time.perf_counter() - start_hash_time
|
hash_duration = time.perf_counter() - start_hash_time
|
||||||
logger.info(f"SHA256 hash calculated for {real_path} in {hash_duration:.3f}s")
|
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
|
# Create instance based on model type
|
||||||
if model_class.__name__ == "CheckpointMetadata":
|
if model_class.__name__ == "CheckpointMetadata":
|
||||||
metadata = model_class(
|
metadata = model_class(
|
||||||
@@ -257,6 +262,9 @@ class MetadataManager:
|
|||||||
usage_tips="{}"
|
usage_tips="{}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Record the AutoV3 state explicitly ("" = checked, no value).
|
||||||
|
metadata.autov3 = autov3 or ""
|
||||||
|
|
||||||
# Try to extract model-specific metadata
|
# Try to extract model-specific metadata
|
||||||
# await MetadataManager._enrich_metadata(metadata, real_path)
|
# await MetadataManager._enrich_metadata(metadata, real_path)
|
||||||
|
|
||||||
|
|||||||
@@ -2,9 +2,40 @@ from dataclasses import dataclass, asdict, field
|
|||||||
from typing import Dict, Optional, List, Any
|
from typing import Dict, Optional, List, Any
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
import os
|
import os
|
||||||
|
from .constants import INVALID_AUTOV3_EMPTY_HASH
|
||||||
from .model_utils import determine_base_model
|
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
|
@dataclass
|
||||||
class BaseModelMetadata:
|
class BaseModelMetadata:
|
||||||
"""Base class for all model metadata structures"""
|
"""Base class for all model metadata structures"""
|
||||||
@@ -35,6 +66,7 @@ class BaseModelMetadata:
|
|||||||
metadata_source: Optional[str] = None # Last provider that supplied metadata
|
metadata_source: Optional[str] = None # Last provider that supplied metadata
|
||||||
last_checked_at: float = 0 # Last checked timestamp
|
last_checked_at: float = 0 # Last checked timestamp
|
||||||
hash_status: str = "completed" # Hash calculation status: pending | calculating | completed | failed
|
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(
|
trainedWords: List[str] = field(
|
||||||
default_factory=list
|
default_factory=list
|
||||||
) # Trigger words / activation prompts (source-agnostic)
|
) # Trigger words / activation prompts (source-agnostic)
|
||||||
@@ -58,6 +90,14 @@ class BaseModelMetadata:
|
|||||||
"""Create instance from dictionary"""
|
"""Create instance from dictionary"""
|
||||||
data_copy = data.copy()
|
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
|
# Use cached fields if available, otherwise compute them
|
||||||
if not hasattr(cls, "_known_fields_cache"):
|
if not hasattr(cls, "_known_fields_cache"):
|
||||||
known_fields = set()
|
known_fields = set()
|
||||||
@@ -97,11 +137,29 @@ class BaseModelMetadata:
|
|||||||
if hasattr(self, "_unknown_fields"):
|
if hasattr(self, "_unknown_fields"):
|
||||||
result.update(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
|
return result
|
||||||
|
|
||||||
def update_civitai_info(self, civitai_data: Dict) -> None:
|
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
|
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:
|
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:
|
if "description" in model_data:
|
||||||
description = model_data["description"]
|
description = model_data["description"]
|
||||||
|
|
||||||
|
sha256_value = (file_info.get("hashes") or {}).get("SHA256", "").lower()
|
||||||
|
|
||||||
return cls(
|
return cls(
|
||||||
file_name=os.path.splitext(file_name)[0],
|
file_name=os.path.splitext(file_name)[0],
|
||||||
model_name=model_data.get("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, "/"),
|
file_path=save_path.replace(os.sep, "/"),
|
||||||
size=file_info.get("sizeKB", 0) * 1024,
|
size=file_info.get("sizeKB", 0) * 1024,
|
||||||
modified=datetime.now().timestamp(),
|
modified=datetime.now().timestamp(),
|
||||||
sha256=(file_info.get("hashes") or {}).get("SHA256", "").lower(),
|
sha256=sha256_value,
|
||||||
base_model=base_model,
|
base_model=base_model,
|
||||||
preview_url="", # Will be updated after preview download
|
preview_url="", # Will be updated after preview download
|
||||||
preview_nsfw_level=0, # 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,
|
civitai=version_info,
|
||||||
tags=tags,
|
tags=tags,
|
||||||
modelDescription=description,
|
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"""
|
"""Create CheckpointMetadata instance from Civitai version info"""
|
||||||
file_name = file_info.get("name", "")
|
file_name = file_info.get("name", "")
|
||||||
base_model = determine_base_model(version_info.get("baseModel", ""))
|
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")
|
sub_type = version_info.get("type", "checkpoint")
|
||||||
|
|
||||||
# Extract tags and description if available
|
# Extract tags and description if available
|
||||||
@@ -237,7 +299,7 @@ class CheckpointMetadata(BaseModelMetadata):
|
|||||||
file_path=save_path.replace(os.sep, "/"),
|
file_path=save_path.replace(os.sep, "/"),
|
||||||
size=file_info.get("sizeKB", 0) * 1024,
|
size=file_info.get("sizeKB", 0) * 1024,
|
||||||
modified=datetime.now().timestamp(),
|
modified=datetime.now().timestamp(),
|
||||||
sha256=(file_info.get("hashes") or {}).get("SHA256", "").lower(),
|
sha256=sha256_value,
|
||||||
base_model=base_model,
|
base_model=base_model,
|
||||||
preview_url="", # Will be updated after preview download
|
preview_url="", # Will be updated after preview download
|
||||||
preview_nsfw_level=0,
|
preview_nsfw_level=0,
|
||||||
@@ -246,6 +308,7 @@ class CheckpointMetadata(BaseModelMetadata):
|
|||||||
sub_type=sub_type,
|
sub_type=sub_type,
|
||||||
tags=tags,
|
tags=tags,
|
||||||
modelDescription=description,
|
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"""
|
"""Create EmbeddingMetadata instance from Civitai version info"""
|
||||||
file_name = file_info.get("name", "")
|
file_name = file_info.get("name", "")
|
||||||
base_model = determine_base_model(version_info.get("baseModel", ""))
|
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")
|
sub_type = version_info.get("type", "embedding")
|
||||||
|
|
||||||
# Extract tags and description if available
|
# Extract tags and description if available
|
||||||
@@ -279,7 +343,7 @@ class EmbeddingMetadata(BaseModelMetadata):
|
|||||||
file_path=save_path.replace(os.sep, "/"),
|
file_path=save_path.replace(os.sep, "/"),
|
||||||
size=file_info.get("sizeKB", 0) * 1024,
|
size=file_info.get("sizeKB", 0) * 1024,
|
||||||
modified=datetime.now().timestamp(),
|
modified=datetime.now().timestamp(),
|
||||||
sha256=(file_info.get("hashes") or {}).get("SHA256", "").lower(),
|
sha256=sha256_value,
|
||||||
base_model=base_model,
|
base_model=base_model,
|
||||||
preview_url="", # Will be updated after preview download
|
preview_url="", # Will be updated after preview download
|
||||||
preview_nsfw_level=0,
|
preview_nsfw_level=0,
|
||||||
@@ -288,4 +352,5 @@ class EmbeddingMetadata(BaseModelMetadata):
|
|||||||
sub_type=sub_type,
|
sub_type=sub_type,
|
||||||
tags=tags,
|
tags=tags,
|
||||||
modelDescription=description,
|
modelDescription=description,
|
||||||
|
autov3=autov3_from_civitai_files(version_info, sha256_value),
|
||||||
)
|
)
|
||||||
|
|||||||
350
tests/services/test_autov3_backfill_service.py
Normal file
350
tests/services/test_autov3_backfill_service.py
Normal 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, '')]
|
||||||
@@ -1318,3 +1318,70 @@ class TestHfGroupKey:
|
|||||||
"hf_url": "https://huggingface.co/user/repo",
|
"hf_url": "https://huggingface.co/user/repo",
|
||||||
}
|
}
|
||||||
assert BaseModelService._extract_group_key(item) == "hf: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
|
||||||
|
|||||||
@@ -321,3 +321,91 @@ class TestCacheEntryValidator:
|
|||||||
|
|
||||||
assert result.is_valid is True
|
assert result.is_valid is True
|
||||||
assert result.repaired is False
|
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
|
||||||
|
|||||||
@@ -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
|
@pytest.mark.asyncio
|
||||||
async def test_fetch_and_update_model_success_updates_cache(tmp_path):
|
async def test_fetch_and_update_model_success_updates_cache(tmp_path):
|
||||||
helpers = build_service()
|
helpers = build_service()
|
||||||
|
|||||||
@@ -111,3 +111,145 @@ class TestModelHashIndexGetDuplicateFilenames:
|
|||||||
index.add_entry("abc123", "/a/lora.safetensors")
|
index.add_entry("abc123", "/a/lora.safetensors")
|
||||||
assert len(index) == 1
|
assert len(index) == 1
|
||||||
assert index.get_duplicate_filenames() == {}
|
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"}
|
||||||
|
|||||||
@@ -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']
|
new_hash_pairs = [p for p in persisted.hash_rows if p[0] == 'new-hash']
|
||||||
assert len(new_hash_pairs) == 1
|
assert len(new_hash_pairs) == 1
|
||||||
assert new_hash_pairs[0][1] == file_path
|
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') == []
|
||||||
|
|||||||
@@ -1,15 +1,26 @@
|
|||||||
import hashlib
|
import hashlib
|
||||||
|
import json
|
||||||
import os
|
import os
|
||||||
|
import struct
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
from py.utils.constants import MAX_SAFETENSORS_HEADER_BYTES
|
||||||
from py.utils.file_utils import (
|
from py.utils.file_utils import (
|
||||||
|
calculate_autov3,
|
||||||
calculate_sha256,
|
calculate_sha256,
|
||||||
find_preview_file,
|
find_preview_file,
|
||||||
get_preview_extension,
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_calculate_sha256(tmp_path):
|
async def test_calculate_sha256(tmp_path):
|
||||||
file_path = tmp_path / "sample.bin"
|
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):
|
def test_get_preview_extension(preview_name, expected):
|
||||||
assert 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
|
||||||
|
|||||||
@@ -1,10 +1,18 @@
|
|||||||
import json
|
import json
|
||||||
|
import struct
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from py.utils.metadata_manager import MetadataManager
|
from py.utils.metadata_manager import MetadataManager
|
||||||
from py.utils.models import BaseModelMetadata
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_base_model_metadata_sets_empty_civitai_dict():
|
async def test_base_model_metadata_sets_empty_civitai_dict():
|
||||||
metadata = BaseModelMetadata(
|
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"))
|
payload = json.loads(metadata_path.read_text(encoding="utf-8"))
|
||||||
|
|
||||||
assert payload.get("civitai") == {}
|
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
|
||||||
|
|||||||
@@ -125,3 +125,200 @@ class TestLoraMetadataConsistency:
|
|||||||
|
|
||||||
# Type is stored in civitai dict
|
# Type is stored in civitai dict
|
||||||
assert metadata.civitai.get("model", {}).get("type") == "Lora"
|
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 == ""
|
||||||
|
|||||||
Reference in New Issue
Block a user