diff --git a/pipelines/minimax/minimax_lora.py b/pipelines/minimax/minimax_lora.py index 171f47d42..d6fda7c13 100644 --- a/pipelines/minimax/minimax_lora.py +++ b/pipelines/minimax/minimax_lora.py @@ -49,6 +49,7 @@ STANDALONE_RENAMES = { "time_embedder.proj_in": "time_embedder.linear_1", "time_embedder.proj_out": "time_embedder.linear_2", "final_layer.adaln_proj.linear": "norm_out.linear", + "final_layer.norm": "norm_out.norm", "final_layer.video_out": "proj_out", "final_layer.audio_out": "audio_proj_out", } diff --git a/test/test-minimax-native-adapters.py b/test/test-minimax-native-adapters.py index 094946fab..1c03bb605 100644 --- a/test/test-minimax-native-adapters.py +++ b/test/test-minimax-native-adapters.py @@ -389,7 +389,8 @@ def oracle_mapping(state_dict, network_alpha=None): a per-key alpha or else the file-level one applied as ``alpha / rank``. """ if any(k.startswith(REFERENCE_PREFIXES) for k in state_dict): - sd = {k.replace('base_model.model.dit.', 'diffusion_model.', 1) if k.startswith('base_model.model.dit.') else k: v for k, v in state_dict.items()} + ignored = [k for k in state_dict if k.endswith(('.diff', '.diff_b'))] # weight and bias residuals of an extraction: not pairs, the converter rejects them + sd = {k.replace('base_model.model.dit.', 'diffusion_model.', 1) if k.startswith('base_model.model.dit.') else k: v for k, v in state_dict.items() if k not in ignored} sd = {k[:-len('.lora_a')] + ('.lora_A.weight' if k.endswith('.lora_a') else '.lora_B.weight') if k.endswith(('.lora_a', '.lora_b')) else k: v for k, v in sd.items()} # lowercase peft names the converter does not read converted = convert_diffusers(sd) out = {} @@ -397,7 +398,7 @@ def oracle_mapping(state_dict, network_alpha=None): if key.endswith('.lora_A.weight'): path = key[len('transformer.'):-len('.lora_A.weight')] out[path] = (value, converted[f'transformer.{path}.lora_B.weight'], 1.0) - return out, [] + return out, ignored has_alpha = any(k.endswith('.alpha') for k in state_dict) out, ignored = {}, [] for key, value in state_dict.items(): @@ -646,6 +647,30 @@ def test_adaln_deltas_project_onto_the_pruned_basis(): return True +def test_extraction_residuals_bind(): + """Bias residuals ride their pair as ex_bias, norm weight residuals bind through the full family, and the final norm resolves.""" + out_proj = REF.transformer_blocks[0].attn.to_out[0] + linear_1 = REF.time_embedder.linear_1 + sd = { + 'blocks.0.attn.out_proj.lora_A.weight': torch.randn(RANK, out_proj.in_features), + 'blocks.0.attn.out_proj.lora_B.weight': torch.randn(out_proj.out_features, RANK), + 'blocks.0.attn.out_proj.diff_b': torch.randn(out_proj.out_features), + 'time_embedder.proj_in.lora_A.weight': torch.randn(RANK, linear_1.in_features), + 'time_embedder.proj_in.lora_B.weight': torch.randn(linear_1.out_features, RANK), + 'time_embedder.proj_in.diff_b': torch.randn(linear_1.out_features), + 'blocks.0.norm1.diff': torch.randn(REF.transformer_blocks[0].norm1.weight.shape[0]), + 'final_layer.norm.diff': torch.randn(REF.norm_out.norm.weight.shape[0]), + } + net = load_native(sd, name='residuals') + assert net is not None, 'chain returned nothing' + mods = net.modules + assert torch.equal(mods['lora_transformer_transformer_blocks_0_attn_to_out_0'].ex_bias, sd['blocks.0.attn.out_proj.diff_b']) + assert torch.equal(mods['lora_transformer_time_embedder_linear_1'].ex_bias, sd['time_embedder.proj_in.diff_b']) + assert torch.equal(mods['lora_transformer_transformer_blocks_0_norm1'].weight, sd['blocks.0.norm1.diff']) + assert torch.equal(mods['lora_transformer_norm_out_norm'].weight, sd['final_layer.norm.diff']) + return True + + # ============================================================ # Tests - real files # ============================================================ @@ -660,10 +685,14 @@ def test_real_files_match_reference_loaders(): return True for path in REAL_FILES: name = os.path.basename(path) - with safetensors.safe_open(path, framework='pt') as f: - metadata = f.metadata() or {} + try: + with safetensors.safe_open(path, framework='pt') as f: + metadata = f.metadata() or {} + sd = safetensors.torch.load_file(path) + except safetensors.SafetensorError as e: # a malformed file the loader cannot open either + log.warning(f' {name}: unreadable, skipped: {e}') + continue network_alpha = M.file_alpha(_MockNetworkOnDisk(path, name, metadata)) - sd = safetensors.torch.load_file(path) native = native_mapping(sd, network_alpha) oracle, ignored = oracle_mapping(sd, network_alpha) probes = assert_same_factors(native, oracle, name) @@ -707,6 +736,7 @@ def run_tests(): test_metadata_alpha_yields_to_alpha_tensors, test_non_numeric_metadata_alpha_is_ignored, test_adaln_deltas_project_onto_the_pruned_basis, + test_extraction_residuals_bind, ]: run_test(CAT_LOADER, fn)