diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index 002c57d82..fd0ff8730 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -303,19 +303,19 @@ def finalize_network(net, name, family, lora_scale, t0, unmapped=0, mismatch=0, return net -def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool: - """LoRA-style rank-and-dim sanity check against the live module weight. - - Honors SDNQ-quantized modules by reading the original shape from the - dequantizer rather than the packed weight tensor. - """ +def module_shape(sd_module): + """The live weight shape of a module, read from the dequantizer for SDNQ-quantized layers; None without a weight.""" if not hasattr(sd_module, "weight"): - return False + return None if hasattr(sd_module, "sdnq_dequantizer"): - mod_shape = sd_module.sdnq_dequantizer.original_shape - else: - mod_shape = sd_module.weight.shape - if len(mod_shape) < 2 or len(down_w.shape) < 2 or len(up_w.shape) < 2: + return tuple(sd_module.sdnq_dequantizer.original_shape) + return tuple(sd_module.weight.shape) + + +def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool: + """LoRA-style rank-and-dim sanity check against the live module weight.""" + mod_shape = module_shape(sd_module) + if mod_shape is None or len(mod_shape) < 2 or len(down_w.shape) < 2 or len(up_w.shape) < 2: return False return down_w.shape[1] == mod_shape[1] and up_w.shape[0] == mod_shape[0] @@ -576,6 +576,7 @@ def try_load_lora(name, network_on_disk, lora_scale, *, network_prefix=NETWORK_PREFIX_DEFAULT, group_by_suffixes_fn=group_by_suffixes, network_alpha=None, + adapt_weights=None, arch_name="generic"): """Generic LoRA loader (handles DoRA via the universal ``finalize_updown`` hook). @@ -584,6 +585,10 @@ def try_load_lora(name, network_on_disk, lora_scale, *, ``network_alpha`` is a file-level alpha for files without alpha tensors; a file carrying any alpha of its own keeps those and ignores it. + + ``adapt_weights(sd_module, network_key, w)`` lets an arch refit a delta onto + a module whose live layout differs from the trained one (a pruned AdaLN + basis, for instance) before the shape check; returning None keeps ``w``. """ t0 = time.time() state_dict = read_state_dict(network_on_disk.filename, what="network") @@ -647,6 +652,9 @@ def try_load_lora(name, network_on_disk, lora_scale, *, skipped += 1 continue + if adapt_weights is not None: + target_w = adapt_weights(sd_module, network_key, target_w) or target_w + if not shapes_match(sd_module, target_w["lora_down.weight"], target_w["lora_up.weight"]): if l.debug: _module = f'{getattr(sd_module, "weight", None).shape if hasattr(sd_module, "weight") else "?"}' diff --git a/pipelines/minimax/minimax_lora.py b/pipelines/minimax/minimax_lora.py index 17079f74d..26961b684 100644 --- a/pipelines/minimax/minimax_lora.py +++ b/pipelines/minimax/minimax_lora.py @@ -12,6 +12,9 @@ lands on the fused SwiGLU projection with its two output halves swapped. import re +import torch + +from modules.logger import log from modules.lora import native_adapter @@ -215,8 +218,38 @@ _BIND_KWARGS = dict( ) +def pruned_basis(sd_module, rank, width): + """The AdaLN curve basis of the loaded transformer that owns ``sd_module``, or None on an unpruned model.""" + from modules import shared + pipe = getattr(shared.sd_model, "pipe", shared.sd_model) + for component in ("transformer", "transformer_ref"): + transformer = getattr(pipe, component, None) + basis = getattr(getattr(transformer, "time_embedder", None), "basis", None) + if basis is None or tuple(basis.shape) != (rank, width): + continue + if any(module is sd_module for module in transformer.modules()): + return basis + return None + + +def project_pruned_adaln(sd_module, network_key, w): + """Refit an AdaLN delta trained on the released time embedding onto the pruned curve basis: the pruned class + stores ``W @ P``, so ``up @ down`` lands exactly as ``up @ (down @ P)``.""" + down = w.get("lora_down.weight") + shape = native_adapter.module_shape(sd_module) + if down is None or down.ndim != 2 or shape is None or len(shape) != 2 or down.shape[1] == shape[1]: + return None + basis = pruned_basis(sd_module, shape[1], down.shape[1]) + if basis is None: + return None + projected = dict(w) + projected["lora_down.weight"] = (down.to(dtype=torch.float32, device=basis.device) @ basis.to(dtype=torch.float32).T).to(dtype=down.dtype, device=down.device) + log.debug(f'Network load: type=LoRA arch=minimaxh3 key={network_key} adaln projected {down.shape[1]}->{shape[1]}') + return projected + + def try_load_lora(name, network_on_disk, lora_scale): - return native_adapter.try_load_lora(name, network_on_disk, lora_scale, network_alpha=file_alpha(network_on_disk), **_BIND_KWARGS) + return native_adapter.try_load_lora(name, network_on_disk, lora_scale, network_alpha=file_alpha(network_on_disk), adapt_weights=project_pruned_adaln, **_BIND_KWARGS) def try_load_lokr(name, network_on_disk, lora_scale): diff --git a/test/test-minimax-native-adapters.py b/test/test-minimax-native-adapters.py index 9d8e7ff30..8b375f239 100644 --- a/test/test-minimax-native-adapters.py +++ b/test/test-minimax-native-adapters.py @@ -24,6 +24,8 @@ Save formats exercised, each seen in a published LoRA: - diffusers names with kohya suffixes: the alibaba-pai Acc LoRAs. - peft wrapper around a ``dit`` attribute (``base_model.model.dit.blocks.0...``): the mvp-lab RAVEN LoRA. +- pruned checkpoints (``MiniMaxH3PrunedTransformer3DModel``): AdaLN deltas + trained at the released width refit onto the rank-k curve basis. The reference module tree is the real ``MiniMaxH3Transformer3DModel`` at tiny dims, so module names and target shapes are authoritative. Every LoRA file @@ -201,14 +203,36 @@ class _MockSdModel: self.__class__.__name__ = 'MiniMaxH3ModularPipeline' -def install_mock_pipe(): +def install_mock_pipe(transformer=REF): """Point shared.sd_model at a mock exposing the reference transformer; re-installed per load so stamps do not leak.""" - sd_model = _MockSdModel(_MockPipeline(REF)) + sd_model = _MockSdModel(_MockPipeline(transformer)) from modules.modeldata import model_data model_data.sd_model = sd_model return sd_model +CURVE_RANK = 2 + + +class _CurveTimeEmbedder(nn.Module): + """The pruned class's table lookup, reduced to the basis buffer the loader reads.""" + + def __init__(self, rank, width): + super().__init__() + self.register_buffer('basis', torch.randn(rank, width)) + + +def pruned_reference(): + """REF with every AdaLN projection folded onto a rank-2 curve basis, the layout of MiniMaxH3PrunedTransformer3DModel.""" + import copy + pruned = copy.deepcopy(REF) + pruned.time_embedder = _CurveTimeEmbedder(CURVE_RANK, pruned.transformer_blocks[0].adaln_proj.linear.in_features) + for block in pruned.transformer_blocks: + block.adaln_proj.linear = nn.Linear(CURVE_RANK, block.adaln_proj.linear.out_features) + pruned.norm_out.linear = nn.Linear(CURVE_RANK, pruned.norm_out.linear.out_features) + return pruned + + # ============================================================ # Synthesizers and helpers # ============================================================ @@ -280,8 +304,8 @@ class _MockNetworkOnDisk: self.metadata = metadata or {} -def load_native(state_dict, name='test', metadata=None): - install_mock_pipe() +def load_native(state_dict, name='test', metadata=None, transformer=REF): + install_mock_pipe(transformer) with TempLora(state_dict, name=name, metadata=metadata) as nod: return M.try_load(name, nod, lora_scale=1.0) @@ -605,6 +629,22 @@ def test_non_numeric_metadata_alpha_is_ignored(): return True +def test_adaln_deltas_project_onto_the_pruned_basis(): + """On a pruned transformer an AdaLN delta trained at the released width binds as up @ (down @ basis.T); every other layer binds verbatim.""" + pruned = pruned_reference() + basis = pruned.time_embedder.basis + sd = {k: v for k, v in synth_diffusers(suffix=KOHYA).items() if not k.startswith('time_embedder.')} # the pruned class drops the time embedder MLP + net = load_native(sd, name='pruned', transformer=pruned) + assert net is not None and net.mismatch == 0, f'mismatch={None if net is None else net.mismatch}' + expected = identity_deltas(sd) + curve_keys = [k for k in expected if k.endswith('_adaln_proj_linear') and '_refiner_' not in k] + ['lora_transformer_norm_out_linear'] + for key in curve_keys: + expected[key] = expected[key] @ basis.T + assert_same_deltas(native_deltas(net), expected) + assert tuple(net.modules['lora_transformer_norm_out_linear'].down_model.weight.shape) == (RANK, CURVE_RANK) + return True + + # ============================================================ # Tests - real files # ============================================================ @@ -665,6 +705,7 @@ def run_tests(): test_metadata_alpha_scales_an_alphaless_file, test_metadata_alpha_yields_to_alpha_tensors, test_non_numeric_metadata_alpha_is_ignored, + test_adaln_deltas_project_onto_the_pruned_basis, ]: run_test(CAT_LOADER, fn)