mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
refactor(loader): collapse native-load logging to one summary line
strip_prefix returns the detected prefix instead of logging it, so the load summary reports it alongside the key count and the redundant reading-state marker is gone. Four near-identical cls/file lines become one.
This commit is contained in:
@@ -245,17 +245,18 @@ def load(
|
||||
)
|
||||
quant_type = model_quant.get_quant_type(quant_args)
|
||||
|
||||
log.debug(f' cls={spec.cls.__name__} file="{os.path.basename(local_file)}" reading state_dict')
|
||||
state_dict = sd_models.read_state_dict(local_file, what="transformer")
|
||||
state_dict = strip_prefix(state_dict, spec.prefixes, spec.cls.__name__)
|
||||
state_dict, detected_prefix = strip_prefix(state_dict, spec.prefixes, spec.cls.__name__)
|
||||
check_forbidden_markers(state_dict, spec.forbidden_markers, spec.cls.__name__, local_file)
|
||||
transformer_sd, sibling_sds = partition_siblings(state_dict, spec.siblings)
|
||||
del state_dict
|
||||
|
||||
sibling_counts = {name: len(sd) for name, sd in sibling_sds.items() if sd}
|
||||
prefix_disp = f'"{detected_prefix}"' if detected_prefix else "bare"
|
||||
log.info(
|
||||
f' cls={spec.cls.__name__} custom="{os.path.basename(local_file)}" '
|
||||
f"transformer_keys={len(transformer_sd)} siblings={sibling_counts or '{}'}"
|
||||
f'cls={spec.cls.__name__} custom="{os.path.basename(local_file)}" '
|
||||
f"prefix={prefix_disp} transformer_keys={len(transformer_sd)} "
|
||||
f"siblings={sibling_counts or '{}'}"
|
||||
)
|
||||
|
||||
effective_dtype = dtype if dtype is not None else devices.dtype
|
||||
@@ -308,7 +309,7 @@ def load(
|
||||
return transformer, loaded_siblings
|
||||
|
||||
|
||||
def strip_prefix(state_dict: dict, prefixes: tuple[str, ...], type_name: str) -> dict:
|
||||
def strip_prefix(state_dict: dict, prefixes: tuple[str, ...], type_name: str) -> tuple[dict, str]:
|
||||
"""Detect and uniformly strip the most common known prefix from every key.
|
||||
|
||||
Order matters: longer prefixes win over shorter ones with the same suffix
|
||||
@@ -316,6 +317,9 @@ def strip_prefix(state_dict: dict, prefixes: tuple[str, ...], type_name: str) ->
|
||||
match the dominant prefix and others do not, raises ValueError because
|
||||
mixed prefixes indicate a malformed file rather than a recoverable export
|
||||
quirk.
|
||||
|
||||
Returns the stripped state dict and the detected prefix (empty string when
|
||||
the keys are already bare) so the caller can report it in one line.
|
||||
"""
|
||||
sorted_prefixes = sorted(prefixes, key=len, reverse=True)
|
||||
counts: dict[str, int] = {}
|
||||
@@ -328,17 +332,15 @@ def strip_prefix(state_dict: dict, prefixes: tuple[str, ...], type_name: str) ->
|
||||
break
|
||||
total = len(state_dict)
|
||||
if seen == 0:
|
||||
log.debug(f"Load model: type={type_name} native_transformer prefix=bare")
|
||||
return state_dict
|
||||
return state_dict, ""
|
||||
dominant = max(counts, key=counts.get)
|
||||
if counts[dominant] != total:
|
||||
raise ValueError(
|
||||
f"Load model: type={type_name} native_transformer has mixed prefixes "
|
||||
f"(total={total} {dominant}={counts[dominant]})"
|
||||
)
|
||||
log.debug(f'Load model: type={type_name} native_transformer prefix="{dominant}"')
|
||||
offset = len(dominant)
|
||||
return {key[offset:]: value for key, value in state_dict.items()}
|
||||
return {key[offset:]: value for key, value in state_dict.items()}, dominant
|
||||
|
||||
|
||||
def check_forbidden_markers(
|
||||
|
||||
@@ -100,15 +100,17 @@ def run_test(cat: str, fn):
|
||||
|
||||
def test_strip_prefix_bare_keys_pass_through():
|
||||
sd = {'layers.0.weight': 1, 'layers.0.bias': 2}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
assert out == sd, 'bare keys must pass through unchanged'
|
||||
assert prefix == '', 'bare keys must report an empty prefix'
|
||||
|
||||
|
||||
def test_strip_prefix_dominant_single_variant():
|
||||
sd = {f'model.diffusion_model.layers.{i}.weight': i for i in range(10)}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
assert all(k.startswith('layers.') for k in out)
|
||||
assert len(out) == 10
|
||||
assert prefix == 'model.diffusion_model.'
|
||||
|
||||
|
||||
def test_strip_prefix_picks_longest_match_first():
|
||||
@@ -117,11 +119,12 @@ def test_strip_prefix_picks_longest_match_first():
|
||||
'model.diffusion_model.layers.0.weight': 1,
|
||||
'model.diffusion_model.layers.1.weight': 2,
|
||||
}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
# If shorter prefix matched, keys would start with 'model.'
|
||||
assert 'layers.0.weight' in out
|
||||
assert 'layers.1.weight' in out
|
||||
assert not any(k.startswith('model.') for k in out)
|
||||
assert prefix == 'model.diffusion_model.'
|
||||
|
||||
|
||||
def test_strip_prefix_mixed_prefixes_raises():
|
||||
@@ -138,20 +141,23 @@ def test_strip_prefix_mixed_prefixes_raises():
|
||||
|
||||
def test_strip_prefix_net_variant():
|
||||
sd = {'net.layers.0.weight': 1, 'net.layers.1.bias': 2}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'}
|
||||
assert prefix == 'net.'
|
||||
|
||||
|
||||
def test_strip_prefix_diffusion_model_variant():
|
||||
sd = {'diffusion_model.layers.0.weight': 1, 'diffusion_model.layers.1.bias': 2}
|
||||
out = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
|
||||
assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'}
|
||||
assert prefix == 'diffusion_model.'
|
||||
|
||||
|
||||
def test_strip_prefix_custom_prefix_set():
|
||||
sd = {'lora_unet_blocks_0.weight': 1, 'lora_unet_blocks_1.weight': 2}
|
||||
out = nt.strip_prefix(sd, ('lora_unet_',), 'Test')
|
||||
out, prefix = nt.strip_prefix(sd, ('lora_unet_',), 'Test')
|
||||
assert set(out.keys()) == {'blocks_0.weight', 'blocks_1.weight'}
|
||||
assert prefix == 'lora_unet_'
|
||||
|
||||
|
||||
# ============================================================
|
||||
|
||||
Reference in New Issue
Block a user