diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 6b8da433b..d04331a74 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -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( diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index 3120b6ddd..dc577dadd 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -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_' # ============================================================