feat(lora): project adaln deltas onto the pruned minimax basis

The pruned transformer stores W @ P for each AdaLN projection against
the rank-8 basis in time_embedder.basis, so deltas trained at the
released width failed the shape check. try_load_lora gains an
adapt_weights hook and the MiniMax loader refits lora_down as down @ P,
which is exact.

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 "?"}'

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):

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)