mirror of
https://github.com/vladmandic/automatic
synced 2026-08-26 06:30:44 +02:00
CalamitousFelicitousness
5e1a0719ab
detect_quant now reads _quantization_metadata as the authoritative quant source when present, and resolves marker-file formats from the marked layers' stored weight dtypes instead of file-wide dtype voting, which mislabeled fp8 files carrying extra uint8 tensors and had no nvfp4 mapping at all. Schema bump so cached probe entries refresh.
430 lines
20 KiB
Python
430 lines
20 KiB
Python
"""Header-only model analysis: architecture fingerprinting, precision/quant
|
|
detection, and embedded metadata extraction from safetensors/gguf files
|
|
without reading tensor data. Consumed by the civitai remote peek and the
|
|
local model audit; analyze_header is pure so both paths share one analyzer.
|
|
"""
|
|
import os
|
|
import re
|
|
import json
|
|
from collections import Counter
|
|
from dataclasses import dataclass
|
|
|
|
|
|
SCHEMA_VERSION = 5
|
|
MAX_HEADER_BYTES = 16 * 1024 * 1024
|
|
STRIP_PREFIXES = ('model.diffusion_model.', 'diffusion_model.', 'net.')
|
|
# companion families bundled alongside the diffusion core in full checkpoints
|
|
COMPANION_PREFIXES = ('text_encoders.', 'cond_stage_model.', 'conditioner.', 'first_stage_model.', 'vae.', 'vocoder.', 'audio_vae.')
|
|
LORA_SUFFIXES = ('.lora_down.weight', '.lora_up.weight', '.lora_A.weight', '.lora_B.weight', '.hada_w1_a', '.lokr_w1', '.lokr_w2', '.dora_scale', '.diff', '.diff_b')
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArchFingerprint:
|
|
family: str
|
|
display: str
|
|
required: tuple # regex fragments, all must match at least one inner key
|
|
forbidden: tuple = ()
|
|
weight: int = 1
|
|
|
|
|
|
# Markers verified against real headers (local library + civitai ranged peeks).
|
|
# Families not listed report 'unknown' rather than guessing.
|
|
FINGERPRINTS = (
|
|
ArchFingerprint('sdxl', 'Stable Diffusion XL', (r'^input_blocks\.', r'^middle_block\.', r'^label_emb\.')),
|
|
ArchFingerprint('sd', 'Stable Diffusion 1.x/2.x', (r'^input_blocks\.', r'^middle_block\.'), forbidden=(r'^label_emb\.',)),
|
|
ArchFingerprint('sd3', 'Stable Diffusion 3', (r'^joint_blocks\.', r'^context_embedder\.')),
|
|
ArchFingerprint('f1', 'FLUX.1', (r'^double_blocks\.\d+\.img_attn', r'^vector_in\.'), forbidden=(r'^distilled_guidance_layer\.', r'stream_modulation')),
|
|
ArchFingerprint('chroma', 'Chroma', (r'^double_blocks\.\d+\.img_attn', r'^distilled_guidance_layer\.')),
|
|
ArchFingerprint('f2', 'FLUX.2', (r'^double_blocks\.\d+\.img_attn', r'stream_modulation'), forbidden=(r'^vector_in\.',)),
|
|
ArchFingerprint('qwen', 'Qwen Image', (r'^transformer_blocks\.\d+\.img_mlp', r'^time_text_embed\.'), forbidden=(r'audio_attn',)),
|
|
ArchFingerprint('krea2', 'Krea 2', (r'^blocks\.\d+\.attn', r'^txtfusion\.')),
|
|
ArchFingerprint('wanai-animate', 'Wan 2.2 Animate', (r'^blocks\.\d+\.cross_attn', r'^patch_embedding\.', r'^face_adapter\.|^motion_encoder\.|^pose_patch_embedding\.')),
|
|
ArchFingerprint('wanai', 'Wan DiT', (r'^blocks\.\d+\.cross_attn', r'^blocks\.\d+\.self_attn', r'^patch_embedding\.', r'^time_projection\.'), forbidden=(r'^face_adapter\.', r'^motion_encoder\.', r'^pose_patch_embedding\.', r'^llm_adapter\.')),
|
|
ArchFingerprint('anima', 'Anima', (r'^blocks\.\d+\.cross_attn', r'^llm_adapter\.')),
|
|
ArchFingerprint('ernieimage', 'Ernie Image', (r'^layers\.\d+\.self_attention', r'^adaLN_sa_ln|^final_norm\.')),
|
|
ArchFingerprint('zimage', 'Z-Image', (r'^layers\.\d+\.attention', r'^noise_refiner\.', r'^cap_embedder\.')),
|
|
ArchFingerprint('ltxvideo', 'LTX Video', (r'^transformer_blocks\.\d+\.attn1', r'^adaln_single\.|^vocoder\.')),
|
|
ArchFingerprint('ideogram4', 'Ideogram 4', (r'^layers\.\d+\.attention\.(qkv|o)\.', r'^adaln_proj\.|^t_embedding\.')),
|
|
)
|
|
|
|
# name/text to family hints, ordered specific-first; matched against lowercased
|
|
# text with spaces removed. Used for lora trainer metadata and for names that
|
|
# imply a base (folder names, base tags).
|
|
FAMILY_NAME_HINTS = (
|
|
('stable-diffusion-xl', 'sdxl'), ('sdxl', 'sdxl'), ('sd_xl', 'sdxl'),
|
|
('animagine', 'sdxl'), ('pony', 'sdxl'), ('illustrious', 'sdxl'), ('noobai', 'sdxl'),
|
|
('stable-diffusion-3', 'sd3'), ('sd3', 'sd3'),
|
|
('stable-diffusion-v1', 'sd'), ('sd-v1', 'sd'), ('sd_v1', 'sd'), ('sd15', 'sd'), ('sd1.5', 'sd'), ('sd2', 'sd'),
|
|
('flux.2', 'f2'), ('flux2', 'f2'), ('klein', 'f2'), ('chroma', 'chroma'), ('flux', 'f1'),
|
|
('krea', 'krea2'), ('qwen', 'qwen'), ('wan', 'wanai'), ('anima', 'anima'),
|
|
('ideogram', 'ideogram4'), ('z-image', 'zimage'), ('zimage', 'zimage'),
|
|
('ernie', 'ernieimage'), ('ltx', 'ltxvideo'),
|
|
)
|
|
|
|
|
|
def family_from_text(text: str) -> str | None:
|
|
compact = (text or '').lower().replace(' ', '')
|
|
for hint, family in FAMILY_NAME_HINTS:
|
|
if hint in compact:
|
|
return family
|
|
return None
|
|
LORA_KEY_HINTS = (
|
|
('txtfusion', 'krea2'),
|
|
('distilled_guidance', 'chroma'),
|
|
('stream_modulation', 'f2'),
|
|
('double_blocks', 'f1'),
|
|
('joint_blocks', 'sd3'),
|
|
('llm_adapter', 'anima'),
|
|
('cross_attn', 'wanai'),
|
|
('img_mlp', 'qwen'),
|
|
('noise_refiner', 'zimage'),
|
|
('context_refiner', 'zimage'),
|
|
# diffusers flux layout; f1 vs f2 not separable from lora keys
|
|
('single_transformer_blocks', 'f1'),
|
|
('transformer_blocks', 'qwen'),
|
|
)
|
|
# checked in order; trainers sometimes write a bogus modelspec.architecture
|
|
# while ss_base_model_version names the real base
|
|
LORA_META_KEYS = ('ss_base_model_version', 'modelspec.architecture', 'ss_sd_model_name')
|
|
FAMILY_DISPLAY = {fp.family: fp.display for fp in FINGERPRINTS}
|
|
|
|
# expert-role vocabulary for dual-transformer archs, matched against
|
|
# __metadata__ values; wan uses phrases because a bare 'high' appears in
|
|
# unrelated metadata (e.g. 'high quality')
|
|
ROLE_WORDS = {
|
|
'wanai': (('high-noise', ('high noise',)), ('low-noise', ('low noise',))),
|
|
'ideogram4': (('uncond', ('uncond', 'unconditional')),),
|
|
}
|
|
|
|
|
|
def strip_key(key: str) -> str:
|
|
for prefix in STRIP_PREFIXES:
|
|
if key.startswith(prefix):
|
|
return key[len(prefix):]
|
|
return key
|
|
|
|
|
|
def split_words(text: str) -> str:
|
|
text = re.sub(r'([a-z0-9])([A-Z])', r'\1 \2', text or '')
|
|
return re.sub(r'[^a-zA-Z0-9]+', ' ', text).lower()
|
|
|
|
|
|
def detect_kind(inner_keys: list) -> str:
|
|
"""Container-level classification before arch matching."""
|
|
if any(k.endswith(LORA_SUFFIXES) or '.lora_down.' in k or '.lora_up.' in k or '.lora.down.' in k or '.lora.up.' in k or '.lora_A.' in k or '.lora_B.' in k for k in inner_keys):
|
|
return 'lora'
|
|
top = Counter(k.split('.')[0] for k in inner_keys)
|
|
total = sum(top.values())
|
|
vae_keys = top.get('encoder', 0) + top.get('decoder', 0) + top.get('quant_conv', 0) + top.get('post_quant_conv', 0)
|
|
if total > 0 and top.get('decoder', 0) > 0 and vae_keys > total * 0.9:
|
|
return 'vae'
|
|
diffusion_core = any(k.startswith(('input_blocks', 'double_blocks', 'blocks.', 'layers.', 'transformer_blocks', 'joint_blocks')) for k in inner_keys)
|
|
llm_markers = any(('.mlp.gate_proj' in k or 'embed_tokens' in k or '.DenseReluDense.' in k) for k in inner_keys) or top.get('shared', 0) > 0
|
|
if not diffusion_core and llm_markers:
|
|
return 'text-encoder'
|
|
return 'model'
|
|
|
|
|
|
def lora_context_dim(shapes: dict | None) -> int | None:
|
|
"""Cross-attention context width from a lora down/A tensor: 768 for sd1.x,
|
|
1024 for sd2.x, 2048 for sdxl. Arbitrates when trainer metadata lies."""
|
|
for k, shape in (shapes or {}).items():
|
|
if 'attn2' in k and 'to_k' in k and ('lora_down' in k or '.lora.down.' in k or 'lora_A' in k):
|
|
if isinstance(shape, (list, tuple)) and len(shape) == 2:
|
|
return int(shape[1])
|
|
return None
|
|
|
|
|
|
def match_lora_base(core_keys: list, metadata: dict | None, shapes: dict | None = None):
|
|
"""Loras train a module subset, so full fingerprints cannot apply. Trainer
|
|
metadata names the base; keyless PEFT loras resolve via distinctive
|
|
substrings. Returns (family, confidence, marker) or None."""
|
|
def unet_split(family, confidence, marker):
|
|
# sd-vs-sdxl from metadata or key heuristics is unreliable; shapes win
|
|
dim = lora_context_dim(shapes)
|
|
if dim is not None:
|
|
return ('sdxl' if dim >= 2048 else 'sd'), 0.8, f'shape:context-dim={dim}'
|
|
return family, confidence, marker
|
|
|
|
for key in LORA_META_KEYS:
|
|
family = family_from_text(str((metadata or {}).get(key, '')))
|
|
if family is not None:
|
|
if family in ('sd', 'sdxl'):
|
|
return unet_split(family, 0.9, f'metadata:{key}')
|
|
return family, 0.9, f'metadata:{key}'
|
|
joined = '\n'.join(core_keys)
|
|
for hint, family in LORA_KEY_HINTS:
|
|
if hint in joined:
|
|
return family, 0.5, f'key:{hint}'
|
|
# bare blocks.N.attn without self/cross variants is the krea2 DiT layout
|
|
if re.search(r'blocks\.\d+\.attn\.', joined) and 'self_attn' not in joined:
|
|
return 'krea2', 0.5, 'key:blocks.attn'
|
|
if 'input_blocks' in joined or 'down_blocks' in joined:
|
|
return unet_split('sdxl' if 'te2' in joined else 'sd', 0.5, 'key:unet-blocks')
|
|
return None
|
|
|
|
|
|
def match_arch(keys: list, metadata: dict | None, shapes: dict | None = None) -> dict:
|
|
inner = [strip_key(k) for k in keys]
|
|
core = [k for k in inner if not k.startswith(COMPANION_PREFIXES)]
|
|
prefixes = {k[: len(k) - len(s)] for k, s in ((k, strip_key(k)) for k in keys) if k != s}
|
|
kind = detect_kind(core)
|
|
matched_family = None
|
|
matched_markers = []
|
|
candidates = []
|
|
confidence = 0.0
|
|
if kind == 'model':
|
|
for fp in FINGERPRINTS:
|
|
hits = [pat for pat in fp.required if any(re.search(pat, k) for k in core)]
|
|
if len(hits) < len(fp.required):
|
|
continue
|
|
if any(any(re.search(pat, k) for k in core) for pat in fp.forbidden):
|
|
continue
|
|
candidates.append(fp)
|
|
if matched_family is None:
|
|
matched_family = fp.family
|
|
matched_markers = hits
|
|
confidence = 1.0 if matched_family and len(candidates) == 1 else (0.7 if matched_family else 0.0)
|
|
elif kind == 'lora':
|
|
resolved = match_lora_base(core, metadata, shapes)
|
|
if resolved is not None:
|
|
matched_family, confidence, marker = resolved
|
|
matched_markers = [marker]
|
|
variant = None
|
|
if matched_family is not None and metadata:
|
|
meta_text = split_words(' '.join(str(v) for v in metadata.values() if isinstance(v, str)))
|
|
for suffix, words in ROLE_WORDS.get(matched_family, ()):
|
|
if any(re.search(rf'\b{w}\b', meta_text) for w in words):
|
|
variant = suffix
|
|
break
|
|
return {
|
|
'kind': kind,
|
|
'family': matched_family or 'unknown',
|
|
'display': FAMILY_DISPLAY.get(matched_family, 'Unknown') if matched_family else 'Unknown',
|
|
'confidence': confidence,
|
|
'variant': variant,
|
|
'detected_prefix': sorted(prefixes)[0] if prefixes else '',
|
|
'matched_markers': matched_markers,
|
|
'candidates': [{'family': fp.family, 'display': fp.display} for fp in candidates],
|
|
}
|
|
|
|
|
|
# storage dtype of a marked layer's weight identifies its comfy_quant format
|
|
COMFY_QUANT_DTYPE_FORMATS = {'I8': 'int8_tensorwise', 'F8_E4M3': 'float8_e4m3fn', 'U8': 'nvfp4'}
|
|
|
|
|
|
def detect_quant(keys: list, dtypes: Counter, container: str, metadata: dict | None = None, key_dtypes: dict | None = None) -> dict:
|
|
if container == 'gguf':
|
|
quantized = {d: n for d, n in dtypes.items() if d not in ('F32', 'F16', 'BF16')}
|
|
dominant = max(quantized, key=quantized.get) if quantized else None
|
|
return {'scheme': 'gguf' if dominant else None, 'format': dominant, 'marked_layers': None, 'source': 'gguf-qtype'}
|
|
quant_metadata = (metadata or {}).get('_quantization_metadata')
|
|
if quant_metadata:
|
|
try:
|
|
layers = (json.loads(quant_metadata) if isinstance(quant_metadata, str) else quant_metadata).get('layers') or {}
|
|
except Exception:
|
|
layers = {}
|
|
formats = Counter(entry.get('format') for entry in layers.values() if isinstance(entry, dict) and entry.get('format'))
|
|
if formats:
|
|
return {'scheme': 'comfy_quant', 'format': formats.most_common(1)[0][0], 'marked_layers': len(layers), 'source': 'header'}
|
|
comfy = [k for k in keys if k.endswith('.comfy_quant')]
|
|
if comfy:
|
|
weight_dtypes = Counter()
|
|
for marker in comfy:
|
|
dtype = (key_dtypes or {}).get(f"{marker[: -len('.comfy_quant')]}.weight")
|
|
if dtype:
|
|
weight_dtypes[dtype] += 1
|
|
dominant = weight_dtypes.most_common(1)[0][0] if weight_dtypes else None
|
|
return {'scheme': 'comfy_quant', 'format': COMFY_QUANT_DTYPE_FORMATS.get(dominant), 'marked_layers': len(comfy), 'source': 'weight-dtype'}
|
|
has_fp8 = dtypes.get('F8_E4M3', 0) + dtypes.get('F8_E5M2', 0) > 0
|
|
has_scales = any(k.endswith(('scaled_fp8', '.scale_weight', '.scale_input', '.weight_scale')) for k in keys)
|
|
if has_fp8 and has_scales:
|
|
return {'scheme': 'scaled_fp8', 'format': 'float8_e4m3fn' if dtypes.get('F8_E4M3') else 'float8_e5m2', 'marked_layers': None, 'source': 'marker'}
|
|
return {'scheme': None, 'format': None, 'marked_layers': None, 'source': None}
|
|
|
|
|
|
def analyze_header(header: dict, container: str = 'safetensors', arch_metadata: dict | None = None) -> dict:
|
|
metadata = header.get('__metadata__') or {}
|
|
if arch_metadata:
|
|
metadata = {**metadata, **arch_metadata}
|
|
entries = {k: v for k, v in header.items() if k != '__metadata__' and isinstance(v, dict)}
|
|
keys = list(entries)
|
|
dtypes = Counter(v.get('dtype') for v in entries.values() if v.get('dtype'))
|
|
params = 0
|
|
# dominant dtype is element-weighted over core tensors only: bundled
|
|
# companions and scale scalars must not outvote the diffusion weights
|
|
core_elements = Counter()
|
|
for k, v in entries.items():
|
|
shape = v.get('shape')
|
|
if not isinstance(shape, (list, tuple)) or k.endswith('.comfy_quant'):
|
|
continue
|
|
n = 1
|
|
for dim in shape:
|
|
n *= int(dim)
|
|
params += n
|
|
if not strip_key(k).startswith(COMPANION_PREFIXES) and v.get('dtype'):
|
|
core_elements[v['dtype']] += n
|
|
inner_keys = [strip_key(k) for k in keys]
|
|
shapes = {k: v.get('shape') for k, v in entries.items()}
|
|
arch = match_arch(keys, metadata, shapes)
|
|
key_dtypes = {strip_key(k): v.get('dtype') for k, v in entries.items()}
|
|
quant = detect_quant(inner_keys, dtypes, container, metadata=metadata, key_dtypes=key_dtypes)
|
|
flags = []
|
|
if not metadata:
|
|
flags.append('no_metadata')
|
|
if any(k.startswith(COMPANION_PREFIXES) for k in inner_keys):
|
|
flags.append('companion_bundled')
|
|
if quant['scheme'] == 'comfy_quant':
|
|
flags.append('comfy_marker')
|
|
if arch['family'] == 'wanai' and arch['variant'] is None:
|
|
flags.append('role_ambiguous')
|
|
if arch['kind'] == 'lora' and 0 < arch['confidence'] < 0.9:
|
|
flags.append('lora_base_inferred')
|
|
return {
|
|
'schema': SCHEMA_VERSION,
|
|
'ok': True,
|
|
'error': None,
|
|
'container': container,
|
|
'tensors': len(keys),
|
|
'params': params,
|
|
'dtypes': dict(dtypes),
|
|
'dominant_dtype': core_elements.most_common(1)[0][0] if core_elements else (dtypes.most_common(1)[0][0] if dtypes else None),
|
|
'arch': arch,
|
|
'quant': quant,
|
|
'metadata': metadata,
|
|
'metadata_present': bool(metadata),
|
|
'flags': flags,
|
|
}
|
|
|
|
|
|
def error_result(container: str, error: str, flag: str = 'unreadable') -> dict:
|
|
return {
|
|
'schema': SCHEMA_VERSION, 'ok': False, 'error': error, 'container': container,
|
|
'tensors': 0, 'params': 0, 'dtypes': {}, 'dominant_dtype': None,
|
|
'arch': {'kind': 'unknown', 'family': 'unknown', 'display': 'Unknown', 'confidence': 0.0, 'variant': None, 'detected_prefix': '', 'matched_markers': [], 'candidates': []},
|
|
'quant': {'scheme': None, 'format': None, 'marked_layers': None, 'source': None},
|
|
'metadata': {}, 'metadata_present': False, 'flags': [flag],
|
|
}
|
|
|
|
|
|
def read_safetensors_header(path: str) -> dict:
|
|
"""Raw header read: 8-byte length prefix + JSON. Does not touch the global
|
|
sd_metadata cache and keeps ss_tag_frequency at full fidelity."""
|
|
with open(path, 'rb') as f:
|
|
header_len = int.from_bytes(f.read(8), 'little')
|
|
if header_len <= 0 or header_len > MAX_HEADER_BYTES:
|
|
raise ValueError(f'implausible header length: {header_len}')
|
|
return json.loads(f.read(header_len).decode('utf-8'))
|
|
|
|
|
|
def comfy_marker_format(path: str) -> str | None:
|
|
"""Exact comfy_quant format string from the first marker tensor's bytes;
|
|
only possible for local files, remote peeks stay dtype-inferred."""
|
|
try:
|
|
with open(path, 'rb') as f:
|
|
header_len = int.from_bytes(f.read(8), 'little')
|
|
if header_len <= 0 or header_len > MAX_HEADER_BYTES:
|
|
return None
|
|
header = json.loads(f.read(header_len).decode('utf-8'))
|
|
for k, v in header.items():
|
|
if k == '__metadata__' or not k.endswith('.comfy_quant') or not isinstance(v, dict):
|
|
continue
|
|
start, end = v.get('data_offsets', (0, 0))
|
|
if end <= start or end - start > 4096:
|
|
return None
|
|
f.seek(8 + header_len + start)
|
|
fmt = json.loads(f.read(end - start).decode('utf-8')).get('format')
|
|
return str(fmt) if fmt else None
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
# quant format / dtype to the precision token used in filenames; fp8 stays
|
|
# variant-specific because e4m3fn and e5m2 differ in kernel support
|
|
QUANT_PRECISION_TOKENS = {'int8_tensorwise': 'int8', 'float8_e4m3fn': 'fp8_e4m3fn', 'float8_e5m2': 'fp8_e5m2', 'nvfp4': 'nvfp4', 'mxfp8': 'mxfp8'}
|
|
DTYPE_PRECISION_TOKENS = {'F32': 'fp32', 'F16': 'fp16', 'BF16': 'bf16', 'F8_E4M3': 'fp8_e4m3fn', 'F8_E5M2': 'fp8_e5m2'}
|
|
|
|
|
|
def precision_token(probe: dict) -> str | None:
|
|
"""Filename token for a probe's true precision; None when the container
|
|
encodes it already (gguf) or nothing is known."""
|
|
quant = probe.get('quant') or {}
|
|
scheme = quant.get('scheme')
|
|
if scheme == 'gguf':
|
|
return None
|
|
if scheme in ('comfy_quant', 'scaled_fp8'):
|
|
fmt = quant.get('format') or ''
|
|
return QUANT_PRECISION_TOKENS.get(fmt, re.sub(r'[^a-z0-9]', '', fmt.lower()) or None)
|
|
return DTYPE_PRECISION_TOKENS.get(probe.get('dominant_dtype') or '')
|
|
|
|
|
|
def probe_safetensors_file(path: str) -> dict:
|
|
try:
|
|
header = read_safetensors_header(path)
|
|
except Exception as e:
|
|
return error_result('safetensors', str(e), 'corrupt_header')
|
|
result = analyze_header(header)
|
|
if result['quant']['scheme'] == 'comfy_quant':
|
|
fmt = comfy_marker_format(path)
|
|
if fmt:
|
|
result['quant']['format'] = fmt
|
|
result['quant']['source'] = 'marker'
|
|
return result
|
|
|
|
|
|
def probe_gguf_file(path: str) -> dict:
|
|
try:
|
|
from modules.ggml import install_gguf
|
|
install_gguf()
|
|
import gguf
|
|
reader = gguf.GGUFReader(path)
|
|
header = {}
|
|
for tensor in reader.tensors:
|
|
header[str(tensor.name)] = {'dtype': str(tensor.tensor_type.name), 'shape': [int(d) for d in reversed(tensor.shape)]}
|
|
arch_metadata = {}
|
|
for key in ('general.architecture', 'general.name'):
|
|
fld = reader.fields.get(key)
|
|
if fld is not None:
|
|
try:
|
|
arch_metadata[key] = str(bytes(fld.parts[fld.data[0]]).decode('utf-8'))
|
|
except Exception:
|
|
pass
|
|
return analyze_header(header, container='gguf', arch_metadata=arch_metadata)
|
|
except Exception as e:
|
|
return error_result('gguf', str(e))
|
|
|
|
|
|
probe_cache = None
|
|
|
|
|
|
def probe_file(path: str, use_cache: bool = True) -> dict:
|
|
"""Dispatch by extension with an mtime-validated persistent cache."""
|
|
global probe_cache # pylint: disable=global-statement
|
|
ext = os.path.splitext(path)[1].lower()
|
|
if ext not in ('.safetensors', '.gguf'):
|
|
return error_result('unknown', f'unsupported extension: {ext}', 'unsupported')
|
|
try:
|
|
stat = os.stat(path)
|
|
except OSError as e:
|
|
return error_result('unknown', str(e))
|
|
if probe_cache is None:
|
|
from modules import paths
|
|
from modules.json_helpers import readfile
|
|
probe_cache = readfile(paths.probe_cache_file, silent=True, lock=True, as_type='dict')
|
|
entry = probe_cache.get(path) if use_cache else None
|
|
if entry and entry.get('mtime') == stat.st_mtime and entry.get('size') == stat.st_size and entry.get('schema') == SCHEMA_VERSION:
|
|
return entry['probe']
|
|
probe = probe_safetensors_file(path) if ext == '.safetensors' else probe_gguf_file(path)
|
|
probe_cache[path] = {'mtime': stat.st_mtime, 'size': stat.st_size, 'schema': SCHEMA_VERSION, 'probe': probe}
|
|
return probe
|
|
|
|
|
|
def save_probe_cache():
|
|
"""Flush the probe cache to disk; call once per scan, not per file."""
|
|
if probe_cache is None:
|
|
return
|
|
from modules import paths
|
|
from modules.json_helpers import writefile
|
|
writefile(probe_cache, paths.probe_cache_file, silent=True, atomic=True)
|