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:
CalamitousFelicitousness
2026-07-05 18:18:21 +01:00
parent f9d2bbe080
commit 0e6158438f
2 changed files with 23 additions and 15 deletions
+11 -9
View File
@@ -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(
+12 -6
View File
@@ -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_'
# ============================================================