alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-08-08 23:10:15 -03:00
Code Issues Packages Projects Releases Wiki Activity

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

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

11
py/utils/constants.py
View File

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

16
py/utils/file_utils.py
View File

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

10
py/utils/metadata_manager.py
View File

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

73
py/utils/models.py
View File

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