Project released-checkpoint AdaLN LoRAs onto the pruned timestep coordinates

README.md

@@ -21,7 +21,7 @@ tags:
 
 [MiniMaxAI/MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) in diffusers format with the AdaLN input
 projections reduced to their reachable rank. **52 GB smaller** across the two DiT partitions, no visible change in
-output, and pruned-trained LoRAs load natively.
+output, and every published H3 LoRA loads - pruned-trained ones natively, released-trained ones projected.
 
 | | released | this repo |
 |---|---|---|
@@ -131,9 +131,13 @@ be paged in and out per step. On a card with room for all of it, expect the two
 
 ## LoRAs
 
-**Pruned-trained LoRAs load natively here.** ai-toolkit trains against a pruned checkpoint by default, so its H3
-LoRAs carry `adaln_proj.linear.lora_A` of shape `[rank, 8]`. That is a size mismatch against the released
-`Linear(2688, ...)` and an exact fit here.
+**Every published H3 LoRA loads here** - the ones trained against a pruned checkpoint natively, the ones trained
+against the released one by projection.
+
+### Pruned-trained LoRAs, natively
+
+ai-toolkit trains against a pruned checkpoint by default, so its H3 LoRAs carry `adaln_proj.linear.lora_A` of shape
+`[rank, 8]`. That is a size mismatch against the released `Linear(2688, ...)` and an exact fit here.
 
 This works because the coordinates are **not ours**. A rank-8 basis is unique only up to sign and rotation, so a
 pruned-trained LoRA is portable only between checkpoints that share the coordinate convention. Rather than derive a
@@ -154,11 +158,71 @@ The last row is why the table is copied and not recomputed: a re-derived basis s
 accurate, and would still turn every existing LoRA into noise. Both files load with all 50 AdaLN projections wrapped
 by PEFT at scale exactly 1.0, 362 modules in total.
 
-The reverse does not hold: LoRAs trained against the **released** checkpoint's 2688-wide AdaLN projections, such as
-the official 4-step turbo LoRA, do **not** load here. Their `adaln_proj` updates live in the full timestep-embedding
-space and cannot be mapped onto the 8-wide one without the training-time activations. Use those against
-`MiniMaxAI/MiniMax-H3`. Everything outside `adaln_proj` - attention, feed-forward, the token refiner - is untouched
-by the pruning and loads either way.
+### Released-trained LoRAs, by projection
+
+LoRAs trained against the **released** checkpoint's 2688-wide AdaLN projections - the 4-step turbo LoRA and its
+conversions - are mapped onto these coordinates when they load. Nothing about them is approximated beyond what the
+pruning already approximates: the identity that folded the weights folds the adapter.
+
+A LoRA on the released projection contributes `lora_B @ (lora_A @ x)`, and `x = mean + c @ basis`, so
+
+```
+lora_B @ (lora_A @ x)  =  lora_B @ ((lora_A @ basis.T) @ c)  +  lora_B @ (lora_A @ mean)
+```
+
+`lora_A @ basis.T` is an `[rank, 8]` `lora_A` over the pruned coordinates, at the same rank; `lora_B` is unchanged;
+and what is left over is a **constant**, which no `Linear(8 -> ...)` can produce. That constant is carried as a
+float32 offset on the modulation, alongside the fold's own `folded_bias` and for the same reason - it is ~98% of
+what these adapters do to the AdaLN path, so rounding it into bfloat16 would spend a full rounding step of the
+modulation on almost the whole update. The factors are computed in float64.
+
+`adaln_basis` and `adaln_mean` ship as `adaln_affine.safetensors` (97 KB) beside each partition's weights. They are
+the map the projections were folded with, read back, not re-derived - for the same reason the table is copied and
+not recomputed.
+
+**Why not PEFT's `lora_bias`.** That is the obvious place to put a constant, and it is the wrong one here.
+`lora_bias` is a `LoraConfig` flag, not a per-module one, so switching it on for the 51 AdaLN modules also switches
+it on for the other 312: it warns on every bias-free target, creates zero biases on the attention and feed-forward
+modules that then read as missing keys, and makes `fuse_lora` raise `Impossible to merge LoRA with lora_bias=True
+because the base layer has no bias` on all 51 - `adaln_proj.linear` is deliberately bias-free, which is what makes
+PEFT wrap it exactly as it wraps the released projection. It would also put the constant in bfloat16, on the wrong
+side of the argument this repo already made for `folded_bias`. The float32 buffer beside `folded_bias` avoids all
+four, and PEFT still owns the state that matters: the offsets are read through `active_adapters`, `scaling` and
+`disable_adapters`, so they follow `set_adapters`, `disable_lora` and `delete_adapters` on their own.
+
+Measured per AdaLN module over 200 off-grid timesteps, projected against full-space, both in float64:
+
+| | worst of 51 | mean of 51 |
+|---|---|---|
+| rank-8 subspace residual on the timestep curve | 1.47e-5 | - |
+| [`larryvrh/MiniMax-H3-Turbo-Lora`](https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora) v4 step 600 EMA | **2.05e-5** | 1.26e-5 |
+| [`InstantX/MiniMax-H3-Turbo-Lora-Diffusers`](https://huggingface.co/InstantX/MiniMax-H3-Turbo-Lora-Diffusers) | **1.53e-5** | 0.88e-5 |
+
+which is the subspace residual and nothing else. Both files attach to **363 modules** - 312 attention/feed-forward
+and token-refiner ones untouched, 51 AdaLN ones projected - all at scale exactly 1.0. Loaded on top of the released
+weights and on top of these, the LoRA'd AdaLN function was compared against an exact float64 evaluation at three
+probe modules: the pruned side is the closer of the two at all three, the same way the fold itself is.
+
+End to end, 960x544, 124 frames, 7 steps, seed 42, `larryvrh` v4, against the identical request on the released
+weights: **cosine 0.9958** with the LoRA on and **0.9992** with it off, and the on/off difference - the turbo effect
+itself - is 0.9609 here against 0.9641 there. Frame for frame the two runs are the same video.
+
+How much of that is the AdaLN half? On these particular files, not much - their AdaLN update is about 2e-4 of the
+modulation, and simply *discarding* those 51 modules lands at 0.9960 against the released run where projecting
+lands at 0.9958, which is the same number twice. What is not the same number is projecting them and dropping the
+constant term: **0.9947**. The half-applied version is the only one of the three that is measurably wrong, which
+is both the argument for where the constant is kept and the reason to do this exactly rather than approximately -
+an adapter that leans harder on AdaLN than a turbo LoRA does has no other way to arrive intact.
+
+A file whose AdaLN modules are not uniformly 8-wide or uniformly 2688-wide raises rather than loading some of them:
+an adapter that applies to 50 of its 51 AdaLN projections is not the adapter anyone trained.
+
+One caveat: the projection happens at load, and the constant terms are not PEFT parameters, so
+`save_lora_weights` on a projected adapter writes the projected factors *without* them. Distribute the original
+file, not a re-save of it.
+
+Everything outside `adaln_proj` and `norm_out` - attention, feed-forward, the token refiner - is untouched by the
+pruning and was always loading either way.
 
 LoRA loading for MiniMax-H3 landed in diffusers via
 [huggingface/diffusers#14408](https://github.com/huggingface/diffusers/pull/14408). This repo declares the stock
@@ -288,7 +352,12 @@ ComfyUI's `*_int8_convrot` files quantize, arrived at from their tensor inventor
 `MiniMaxH3ConvRotLinear` is a plain `nn.Linear` subclass, so PEFT wraps it as a `base_layer` and a LoRA's own
 branch reads the **unrotated** input in bfloat16. A LoRA trained against the ordinary checkpoint is therefore
 correct on top of a rotated, quantized base - and `adaln_proj`, which is where pruned-trained LoRAs put most of
-their weight, is never rotated or quantized at all.
+their weight and where a released-trained one is projected, is never rotated or quantized at all.
+
+The order below is unchanged by the AdaLN projection: the projection happens at load, before anything is rotated,
+and neither `enable_convrot` nor `quantize_` touches the modules it wrote to. Verified on `larryvrh` v4 - the
+AdaLN update at three probe modules is bit-identical before and after quantizing, with all 363 modules and all 51
+constant terms still attached.
 
 The order matters, and there is one thing not to do:
 

transformer/adaln_affine.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34f285e7aeae741666868bf5912506ab4793098357a01a9d8358f6ce7704d532
+size 96960

transformer/diffusion_pytorch_model.safetensors.index.json

@@ -1,6 +1,6 @@
 {
   "metadata": {
-    "total_size": 40235366432
+    "total_size": 40235463200
   },
   "weight_map": {
     "audio_proj_in.bias": "diffusion_pytorch_model-00001-of-00014.safetensors",
@@ -637,6 +637,8 @@
     "transformer_blocks.9.ff.net.0.proj.weight": "diffusion_pytorch_model-00003-of-00014.safetensors",
     "transformer_blocks.9.ff.net.2.weight": "diffusion_pytorch_model-00003-of-00014.safetensors",
     "transformer_blocks.9.norm1.weight": "diffusion_pytorch_model-00003-of-00014.safetensors",
-    "transformer_blocks.9.norm2.weight": "diffusion_pytorch_model-00003-of-00014.safetensors"
+    "transformer_blocks.9.norm2.weight": "diffusion_pytorch_model-00003-of-00014.safetensors",
+    "adaln_basis": "adaln_affine.safetensors",
+    "adaln_mean": "adaln_affine.safetensors"
   }
-}
+}

transformer/modeling_minimax_h3_pruned.py

@@ -17,12 +17,15 @@ Everything outside the timestep path is inherited from `MiniMaxH3Transformer3DMo
 token refiner, the output heads and `forward` itself are the released implementation, unmodified. Only what feeds the
 AdaLN projections changes.
 
-The one other thing this file adds is `enable_convrot` / `quantize_8bit`: an opt-in Hadamard conditioning of the
+Two other things this file adds. `enable_convrot` / `quantize_8bit`: an opt-in Hadamard conditioning of the
 attention and feed-forward linears that makes 8-bit *compute* (int8 or fp8 dynamic activations, via torchao) land
-within a rounding step of bfloat16. It is inert unless called, so the plain pruned path is byte for byte what it was.
+within a rounding step of bfloat16. And `load_lora_adapter`, overridden to project a LoRA trained against the
+*released* 2688-wide AdaLN projections onto these 8-wide ones. Both are inert unless used, so the plain pruned path
+is byte for byte what it was.
 """
 
 import math
+import re
 
 import torch
 import torch.nn as nn
@@ -36,10 +39,17 @@ from diffusers.models.transformers.transformer_minimax_h3 import (
     MiniMaxH3Transformer3DModel,
     MiniMaxH3TransformerBlock,
 )
+from diffusers.utils import logging
 
 
+logger = logging.get_logger(__name__)
+
 _HADAMARD_CACHE: dict = {}
 
+# The 51 AdaLN projections, as a LoRA state dict names them, with or without a `transformer.` / `transformer_ref.`
+# component prefix. Group 1 is the module path the model itself knows the projection by.
+ADALN_LORA_A_KEY = re.compile(r"(?:^|\.)((?:transformer_blocks\.\d+\.adaln_proj|norm_out)\.linear)\.lora_A\.weight$")
+
 # The linears ConvRot conditions: the block stack's attention and feed-forward projections, and nothing else.
 # This is the set ComfyUI's `*_int8_convrot` checkpoints quantize (they carry one fused `qkv_proj`; the three
 # split projections here share an input, so rotating each is the same transform). The AdaLN path, the patch
@@ -129,7 +139,58 @@ class MiniMaxH3PrunedTimeProj(nn.Module):
         return timestep
 
 
-class MiniMaxH3PrunedAdaLayerNormModulation(nn.Module):
+class MiniMaxH3PrunedAdaLN(nn.Module):
+    r"""Shared by both pruned AdaLN modules: the folded float32 bias, plus the LoRA offsets that ride alongside it.
+
+    A LoRA trained against the released 2688-wide projection contributes `lora_B @ (lora_A @ x)` to the modulation.
+    With `x = mean + c @ basis`, that splits into a coordinate term the projected factors reproduce and a *constant*
+    term, `lora_B @ (lora_A @ mean)`, which no `Linear(8 -> out)` can express. That constant is held here, as a
+    per-adapter float32 buffer added to `folded_bias` - the same place, and the same precision, the fold's own
+    constant term lives in. It is deliberately not the projection's `bias`: rounding it into bfloat16 would spend a
+    full rounding step of the modulation on a term that is most of what the adapter does to the AdaLN path.
+    """
+
+    def __init__(self, out_features: int) -> None:
+        super().__init__()
+        self.register_buffer("folded_bias", torch.zeros(out_features), persistent=True)
+        # `{adapter name: buffer attribute}`. Plain state, not a submodule: the buffers themselves are what move
+        # with the module, and being non-persistent they stay out of the checkpoint, as an adapter should.
+        self._lora_adaln_offsets: dict[str, str] = {}
+
+    def register_lora_adaln_offset(self, adapter_name: str, offset: torch.Tensor) -> None:
+        r"""Attach one adapter's constant term, on the device and in the precision `folded_bias` is kept in."""
+        attribute = self._lora_adaln_offsets.get(adapter_name)
+        if attribute is None:
+            attribute = f"lora_adaln_offset_{len(self._lora_adaln_offsets)}"
+            self._lora_adaln_offsets[adapter_name] = attribute
+        value = offset.to(device=self.folded_bias.device, dtype=torch.float32)
+        self.register_buffer(attribute, value, persistent=False)
+
+    def lora_adaln_bias(self) -> torch.Tensor:
+        r"""`folded_bias` plus every active adapter's constant term at its current scaling.
+
+        PEFT owns everything this reads - `active_adapters`, `scaling`, `disable_adapters` - so the offsets follow
+        `set_adapters`, `disable_lora` and `delete_adapters` with no bookkeeping of their own. They apply whether or
+        not an adapter is merged: `fuse_lora` folds `lora_B @ lora_A` into the projection's weight, and there is
+        nowhere in a bias-free `Linear` for this term to be folded to.
+        """
+        bias = self.folded_bias
+        offsets = self._lora_adaln_offsets
+        if not offsets:
+            return bias
+        layer = self.linear
+        scaling = getattr(layer, "scaling", None)
+        if not isinstance(scaling, dict) or getattr(layer, "disable_adapters", False):
+            return bias
+        for adapter_name in getattr(layer, "active_adapters", ()):
+            attribute = offsets.get(adapter_name)
+            if attribute is None or adapter_name not in scaling:
+                continue
+            bias = bias + getattr(self, attribute) * float(scaling[adapter_name])
+        return bias
+
+
+class MiniMaxH3PrunedAdaLayerNormModulation(MiniMaxH3PrunedAdaLN):
     r"""`MiniMaxH3AdaLayerNormModulation` over the pruned timestep coordinates.
 
     Two differences from the released module. It applies no `silu` - the table already holds the coordinates of the
@@ -143,33 +204,31 @@ class MiniMaxH3PrunedAdaLayerNormModulation(nn.Module):
     """
 
     def __init__(self, adaln_rank: int, hidden_size: int) -> None:
-        super().__init__()
-        self.hidden_size = hidden_size
         out_features = 6 * hidden_size * MINIMAX_H3_MODALITY_NUM
+        super().__init__(out_features)
+        self.hidden_size = hidden_size
         self.linear = nn.Linear(adaln_rank, out_features, bias=False)
-        self.register_buffer("folded_bias", torch.zeros(out_features), persistent=True)
 
     def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]:
         dtype = get_parameter_dtype(self.linear)
         temb = self.linear(temb.to(dtype))
-        temb = (temb.float() + self.folded_bias).to(dtype)
+        temb = (temb.float() + self.lora_adaln_bias()).to(dtype)
         temb = temb.view(-1, 6 * self.hidden_size)
         return temb.chunk(6, dim=-1)
 
 
-class MiniMaxH3PrunedAdaLayerNormOut(nn.Module):
+class MiniMaxH3PrunedAdaLayerNormOut(MiniMaxH3PrunedAdaLN):
     r"""`MiniMaxH3AdaLayerNormOut` over the pruned timestep coordinates; see the modulation module above."""
 
     def __init__(self, hidden_size: int, adaln_rank: int, eps: float) -> None:
-        super().__init__()
+        super().__init__(2 * hidden_size)
         self.norm = nn.RMSNorm(hidden_size, eps=eps)
         self.linear = nn.Linear(adaln_rank, 2 * hidden_size, bias=False)
-        self.register_buffer("folded_bias", torch.zeros(2 * hidden_size), persistent=True)
 
     def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, timestep_indices: torch.Tensor) -> torch.Tensor:
         dtype = get_parameter_dtype(self.linear)
         temb = self.linear(temb.to(dtype))
-        shift, scale = (temb.float() + self.folded_bias).to(dtype).chunk(2, dim=-1)
+        shift, scale = (temb.float() + self.lora_adaln_bias()).to(dtype).chunk(2, dim=-1)
         hidden_states = self.norm(hidden_states)
         return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select(
             0, timestep_indices
@@ -188,7 +247,8 @@ class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
     Only what builds the timestep path differs from [`MiniMaxH3Transformer3DModel`]: `time_proj` becomes an identity,
     `time_embedder` becomes [`MiniMaxH3PrunedTimeEmbedder`], and the AdaLN projections take `adaln_rank` inputs.
     `forward` is inherited unchanged. The module names are the released ones, so LoRAs trained against a pruned
-    checkpoint - what the common trainers use by default - load natively.
+    checkpoint - what the common trainers use by default - load natively, and `load_lora_adapter` projects LoRAs
+    trained against the released checkpoint's `time_embed_dim`-wide projections onto the same coordinates.
 
     Args:
         adaln_rank (`int`, defaults to `8`):
@@ -217,6 +277,8 @@ class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
         "audio_proj_out",
         "rope",
         "folded_bias",
+        "adaln_basis",
+        "adaln_mean",
     ]
 
     @register_to_config
@@ -259,6 +321,13 @@ class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
         self.time_proj = MiniMaxH3PrunedTimeProj()
         self.time_embedder = MiniMaxH3PrunedTimeEmbedder(table_size=time_table_size, adaln_rank=adaln_rank)
 
+        # 2b. The affine map the fold was performed with: `silu(time_embedder(t)) ~= adaln_mean + c(t) @ adaln_basis`.
+        #     Nothing in `forward` reads these - the folded projections already carry them. They are stored so that a
+        #     LoRA trained on the released 2688-wide projections can be mapped onto these coordinates at load time;
+        #     see `load_lora_adapter`. 97 KB per partition.
+        self.register_buffer("adaln_basis", torch.zeros(adaln_rank, time_embed_dim), persistent=True)
+        self.register_buffer("adaln_mean", torch.zeros(time_embed_dim), persistent=True)
+
         # 3. Rotary embedding over the packed (t, h, w) grid
         self.rope = MiniMaxH3RotaryPosEmbed(rope_freq_dim=rope_freq_dim, rope_theta=rope_theta)
 
@@ -303,6 +372,140 @@ class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
 
         self.gradient_checkpointing = False
 
+    # -- LoRA ---------------------------------------------------------------------------------------------------
+
+    def project_adaln_lora(self, state_dict: dict, prefix: str | None = None) -> tuple[dict, dict]:
+        r"""Map a LoRA's AdaLN factors from the released timestep embedding onto the pruned coordinates.
+
+        The released projection reads `x = silu(time_embedder(t))`, so a LoRA on it contributes
+
+            lora_B @ (lora_A @ x)  =  lora_B @ (lora_A @ (mean + basis.T @ c))
+                                   =  (lora_B @ (lora_A @ basis.T)) @ c  +  lora_B @ (lora_A @ mean)
+
+        which is a rank-preserving `[rank, adaln_rank]` `lora_A` over the pruned coordinates plus a constant output
+        offset. Both are computed in float64 from the file's own factors; the only error is the rank-8 subspace's
+        own residual on the timestep curve, 1.5e-5 relative, which is ~250x below one bfloat16 step of the weights
+        being adapted.
+
+        Returns `(state_dict, {module path: offset})`, the state dict unchanged and the offsets empty when the AdaLN
+        factors are already `adaln_rank`-wide (a LoRA trained against a pruned checkpoint - most of them).
+
+        Every AdaLN module in a file has to be one or the other. A file that mixes widths is not something this can
+        half-apply, so it raises.
+
+        `prefix` scopes this to one component's keys, the same way `load_lora_adapter` scopes the load - a file that
+        names both partitions holds two different adapters under one roof, and only one of them is going into this
+        module.
+        """
+        adaln_rank = self.config.adaln_rank
+        time_embed_dim = self.config.time_embed_dim
+
+        modules = {}
+        for key in state_dict:
+            if prefix is not None and not key.startswith(f"{prefix}."):
+                continue
+            match = ADALN_LORA_A_KEY.search(key)
+            if match is not None:
+                modules[key] = match.group(1)
+        if not modules:
+            return state_dict, {}
+
+        widths = sorted({int(state_dict[key].shape[1]) for key in modules})
+        if widths == [adaln_rank]:
+            return state_dict, {}
+        if widths != [time_embed_dim]:
+            odd = sorted(
+                {modules[key] for key in modules if int(state_dict[key].shape[1]) != max(widths)},
+                key=lambda name: (name != "norm_out.linear", name),
+            )
+            raise ValueError(
+                f"This LoRA's {len(modules)} AdaLN projections read inputs of width {widths}. On this checkpoint they "
+                f"have to be uniformly {adaln_rank} wide (trained against a pruned checkpoint, loaded as they are) or "
+                f"uniformly {time_embed_dim} wide (trained against the released checkpoint, projected onto the pruned "
+                "coordinates at load). An adapter cannot be applied to some of its AdaLN modules and not others, so "
+                f"nothing was loaded. The minority width is on: {odd[:4]}{' ...' if len(odd) > 4 else ''}."
+            )
+
+        basis = self.adaln_basis
+        mean = self.adaln_mean
+        if basis.abs().sum() == 0:
+            raise ValueError(
+                "Projecting a released-checkpoint LoRA needs `adaln_basis` and `adaln_mean`, the affine map this "
+                "checkpoint's AdaLN projections were folded with, and this model was loaded without them. Re-download "
+                "the repository: they ship as `adaln_affine.safetensors` next to the weights."
+            )
+        basis = basis.double()
+        mean = mean.double()
+
+        projected = dict(state_dict)
+        offsets = {}
+        for key, module in modules.items():
+            lora_b_key = key[: -len("lora_A.weight")] + "lora_B.weight"
+            if lora_b_key not in state_dict:
+                raise ValueError(
+                    f"{key} has no matching {lora_b_key}. The constant term of the projection is "
+                    "`lora_B @ (lora_A @ mean)`, so both factors have to be present; nothing was loaded."
+                )
+            lora_a = state_dict[key].to(device=basis.device, dtype=torch.float64)
+            lora_b = state_dict[lora_b_key].to(device=basis.device, dtype=torch.float64)
+            # float32 rather than the file's bfloat16: the projection is exact arithmetic on the file's factors and
+            # there is no reason to round it twice. PEFT casts to the adapter's dtype when it loads them.
+            projected[key] = (lora_a @ basis.T).to(torch.float32).cpu().contiguous()
+            offsets[module] = (lora_b @ (lora_a @ mean)).to(torch.float32).cpu().contiguous()
+        return projected, offsets
+
+    def load_lora_adapter(self, pretrained_model_name_or_path_or_dict, prefix="transformer", hotswap=False, **kwargs):
+        r"""`PeftAdapterMixin.load_lora_adapter`, with the AdaLN projection of [`project_adaln_lora`] in front of it.
+
+        LoRAs trained against a pruned checkpoint pass through untouched. LoRAs trained against the released
+        checkpoint's 2688-wide AdaLN projections - the official turbo LoRA and its conversions - are mapped onto the
+        pruned coordinates here, which is the only thing that ever stopped them loading. Everything outside the AdaLN
+        path is identical between the two checkpoints and is neither inspected nor changed.
+        """
+        state_dict = pretrained_model_name_or_path_or_dict
+        if not isinstance(state_dict, dict):
+            from diffusers.loaders.lora_base import _fetch_state_dict
+
+            state_dict, _ = _fetch_state_dict(
+                pretrained_model_name_or_path_or_dict=state_dict,
+                weight_name=kwargs.get("weight_name"),
+                use_safetensors=kwargs.get("use_safetensors", True),
+                local_files_only=kwargs.get("local_files_only"),
+                cache_dir=kwargs.get("cache_dir"),
+                force_download=kwargs.get("force_download", False),
+                proxies=kwargs.get("proxies"),
+                token=kwargs.get("token"),
+                revision=kwargs.get("revision"),
+                subfolder=kwargs.get("subfolder"),
+                user_agent={"file_type": "attn_procs_weights", "framework": "pytorch"},
+                allow_pickle=False,
+                metadata=kwargs.get("metadata"),
+            )
+
+        state_dict, offsets = self.project_adaln_lora(state_dict, prefix=prefix)
+        if offsets:
+            logger.info(
+                f"Projecting {len(offsets)} AdaLN LoRA modules from the released {self.config.time_embed_dim}-wide "
+                f"timestep embedding onto this checkpoint's {self.config.adaln_rank} pruned coordinates; each one's "
+                "constant term is carried as a float32 offset on the modulation."
+            )
+
+        before = set(getattr(self, "peft_config", None) or ())
+        super().load_lora_adapter(state_dict, prefix=prefix, hotswap=hotswap, **kwargs)
+        if not offsets:
+            return
+
+        added = set(getattr(self, "peft_config", None) or ()) - before
+        adapter_name = added.pop() if len(added) == 1 else kwargs.get("adapter_name")
+        if adapter_name is None:
+            raise RuntimeError(
+                "The AdaLN projection could not tell which adapter was just loaded, so its constant terms were not "
+                f"attached and the adapter is incomplete. Adapters before: {sorted(before)}, after: "
+                f"{sorted(getattr(self, 'peft_config', None) or ())}. Pass `adapter_name` explicitly."
+            )
+        for module, offset in offsets.items():
+            self.get_submodule(module.rsplit(".", 1)[0]).register_lora_adaln_offset(adapter_name, offset)
+
     # -- 8-bit compute -----------------------------------------------------------------------------------------
     #
     # `enable_convrot` and `quantize_8bit` are opt-in and change nothing until called.

transformer_ref/adaln_affine.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98fbeba868fd1d08b8232dce62abb3f5427dd49fb08462869633a3d53533b50f
+size 96960

transformer_ref/diffusion_pytorch_model.safetensors.index.json

@@ -1,6 +1,6 @@
 {
   "metadata": {
-    "total_size": 40235366432
+    "total_size": 40235463200
   },
   "weight_map": {
     "audio_proj_in.bias": "diffusion_pytorch_model-00001-of-00014.safetensors",
@@ -637,6 +637,8 @@
     "transformer_blocks.9.ff.net.0.proj.weight": "diffusion_pytorch_model-00003-of-00014.safetensors",
     "transformer_blocks.9.ff.net.2.weight": "diffusion_pytorch_model-00003-of-00014.safetensors",
     "transformer_blocks.9.norm1.weight": "diffusion_pytorch_model-00003-of-00014.safetensors",
-    "transformer_blocks.9.norm2.weight": "diffusion_pytorch_model-00003-of-00014.safetensors"
+    "transformer_blocks.9.norm2.weight": "diffusion_pytorch_model-00003-of-00014.safetensors",
+    "adaln_basis": "adaln_affine.safetensors",
+    "adaln_mean": "adaln_affine.safetensors"
   }
-}
+}

transformer_ref/modeling_minimax_h3_pruned.py

@@ -17,12 +17,15 @@ Everything outside the timestep path is inherited from `MiniMaxH3Transformer3DMo
 token refiner, the output heads and `forward` itself are the released implementation, unmodified. Only what feeds the
 AdaLN projections changes.
 
-The one other thing this file adds is `enable_convrot` / `quantize_8bit`: an opt-in Hadamard conditioning of the
+Two other things this file adds. `enable_convrot` / `quantize_8bit`: an opt-in Hadamard conditioning of the
 attention and feed-forward linears that makes 8-bit *compute* (int8 or fp8 dynamic activations, via torchao) land
-within a rounding step of bfloat16. It is inert unless called, so the plain pruned path is byte for byte what it was.
+within a rounding step of bfloat16. And `load_lora_adapter`, overridden to project a LoRA trained against the
+*released* 2688-wide AdaLN projections onto these 8-wide ones. Both are inert unless used, so the plain pruned path
+is byte for byte what it was.
 """
 
 import math
+import re
 
 import torch
 import torch.nn as nn
@@ -36,10 +39,17 @@ from diffusers.models.transformers.transformer_minimax_h3 import (
     MiniMaxH3Transformer3DModel,
     MiniMaxH3TransformerBlock,
 )
+from diffusers.utils import logging
 
 
+logger = logging.get_logger(__name__)
+
 _HADAMARD_CACHE: dict = {}
 
+# The 51 AdaLN projections, as a LoRA state dict names them, with or without a `transformer.` / `transformer_ref.`
+# component prefix. Group 1 is the module path the model itself knows the projection by.
+ADALN_LORA_A_KEY = re.compile(r"(?:^|\.)((?:transformer_blocks\.\d+\.adaln_proj|norm_out)\.linear)\.lora_A\.weight$")
+
 # The linears ConvRot conditions: the block stack's attention and feed-forward projections, and nothing else.
 # This is the set ComfyUI's `*_int8_convrot` checkpoints quantize (they carry one fused `qkv_proj`; the three
 # split projections here share an input, so rotating each is the same transform). The AdaLN path, the patch
@@ -129,7 +139,58 @@ class MiniMaxH3PrunedTimeProj(nn.Module):
         return timestep
 
 
-class MiniMaxH3PrunedAdaLayerNormModulation(nn.Module):
+class MiniMaxH3PrunedAdaLN(nn.Module):
+    r"""Shared by both pruned AdaLN modules: the folded float32 bias, plus the LoRA offsets that ride alongside it.
+
+    A LoRA trained against the released 2688-wide projection contributes `lora_B @ (lora_A @ x)` to the modulation.
+    With `x = mean + c @ basis`, that splits into a coordinate term the projected factors reproduce and a *constant*
+    term, `lora_B @ (lora_A @ mean)`, which no `Linear(8 -> out)` can express. That constant is held here, as a
+    per-adapter float32 buffer added to `folded_bias` - the same place, and the same precision, the fold's own
+    constant term lives in. It is deliberately not the projection's `bias`: rounding it into bfloat16 would spend a
+    full rounding step of the modulation on a term that is most of what the adapter does to the AdaLN path.
+    """
+
+    def __init__(self, out_features: int) -> None:
+        super().__init__()
+        self.register_buffer("folded_bias", torch.zeros(out_features), persistent=True)
+        # `{adapter name: buffer attribute}`. Plain state, not a submodule: the buffers themselves are what move
+        # with the module, and being non-persistent they stay out of the checkpoint, as an adapter should.
+        self._lora_adaln_offsets: dict[str, str] = {}
+
+    def register_lora_adaln_offset(self, adapter_name: str, offset: torch.Tensor) -> None:
+        r"""Attach one adapter's constant term, on the device and in the precision `folded_bias` is kept in."""
+        attribute = self._lora_adaln_offsets.get(adapter_name)
+        if attribute is None:
+            attribute = f"lora_adaln_offset_{len(self._lora_adaln_offsets)}"
+            self._lora_adaln_offsets[adapter_name] = attribute
+        value = offset.to(device=self.folded_bias.device, dtype=torch.float32)
+        self.register_buffer(attribute, value, persistent=False)
+
+    def lora_adaln_bias(self) -> torch.Tensor:
+        r"""`folded_bias` plus every active adapter's constant term at its current scaling.
+
+        PEFT owns everything this reads - `active_adapters`, `scaling`, `disable_adapters` - so the offsets follow
+        `set_adapters`, `disable_lora` and `delete_adapters` with no bookkeeping of their own. They apply whether or
+        not an adapter is merged: `fuse_lora` folds `lora_B @ lora_A` into the projection's weight, and there is
+        nowhere in a bias-free `Linear` for this term to be folded to.
+        """
+        bias = self.folded_bias
+        offsets = self._lora_adaln_offsets
+        if not offsets:
+            return bias
+        layer = self.linear
+        scaling = getattr(layer, "scaling", None)
+        if not isinstance(scaling, dict) or getattr(layer, "disable_adapters", False):
+            return bias
+        for adapter_name in getattr(layer, "active_adapters", ()):
+            attribute = offsets.get(adapter_name)
+            if attribute is None or adapter_name not in scaling:
+                continue
+            bias = bias + getattr(self, attribute) * float(scaling[adapter_name])
+        return bias
+
+
+class MiniMaxH3PrunedAdaLayerNormModulation(MiniMaxH3PrunedAdaLN):
     r"""`MiniMaxH3AdaLayerNormModulation` over the pruned timestep coordinates.
 
     Two differences from the released module. It applies no `silu` - the table already holds the coordinates of the
@@ -143,33 +204,31 @@ class MiniMaxH3PrunedAdaLayerNormModulation(nn.Module):
     """
 
     def __init__(self, adaln_rank: int, hidden_size: int) -> None:
-        super().__init__()
-        self.hidden_size = hidden_size
         out_features = 6 * hidden_size * MINIMAX_H3_MODALITY_NUM
+        super().__init__(out_features)
+        self.hidden_size = hidden_size
         self.linear = nn.Linear(adaln_rank, out_features, bias=False)
-        self.register_buffer("folded_bias", torch.zeros(out_features), persistent=True)
 
     def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]:
         dtype = get_parameter_dtype(self.linear)
         temb = self.linear(temb.to(dtype))
-        temb = (temb.float() + self.folded_bias).to(dtype)
+        temb = (temb.float() + self.lora_adaln_bias()).to(dtype)
         temb = temb.view(-1, 6 * self.hidden_size)
         return temb.chunk(6, dim=-1)
 
 
-class MiniMaxH3PrunedAdaLayerNormOut(nn.Module):
+class MiniMaxH3PrunedAdaLayerNormOut(MiniMaxH3PrunedAdaLN):
     r"""`MiniMaxH3AdaLayerNormOut` over the pruned timestep coordinates; see the modulation module above."""
 
     def __init__(self, hidden_size: int, adaln_rank: int, eps: float) -> None:
-        super().__init__()
+        super().__init__(2 * hidden_size)
         self.norm = nn.RMSNorm(hidden_size, eps=eps)
         self.linear = nn.Linear(adaln_rank, 2 * hidden_size, bias=False)
-        self.register_buffer("folded_bias", torch.zeros(2 * hidden_size), persistent=True)
 
     def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, timestep_indices: torch.Tensor) -> torch.Tensor:
         dtype = get_parameter_dtype(self.linear)
         temb = self.linear(temb.to(dtype))
-        shift, scale = (temb.float() + self.folded_bias).to(dtype).chunk(2, dim=-1)
+        shift, scale = (temb.float() + self.lora_adaln_bias()).to(dtype).chunk(2, dim=-1)
         hidden_states = self.norm(hidden_states)
         return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select(
             0, timestep_indices
@@ -188,7 +247,8 @@ class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
     Only what builds the timestep path differs from [`MiniMaxH3Transformer3DModel`]: `time_proj` becomes an identity,
     `time_embedder` becomes [`MiniMaxH3PrunedTimeEmbedder`], and the AdaLN projections take `adaln_rank` inputs.
     `forward` is inherited unchanged. The module names are the released ones, so LoRAs trained against a pruned
-    checkpoint - what the common trainers use by default - load natively.
+    checkpoint - what the common trainers use by default - load natively, and `load_lora_adapter` projects LoRAs
+    trained against the released checkpoint's `time_embed_dim`-wide projections onto the same coordinates.
 
     Args:
         adaln_rank (`int`, defaults to `8`):
@@ -217,6 +277,8 @@ class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
         "audio_proj_out",
         "rope",
         "folded_bias",
+        "adaln_basis",
+        "adaln_mean",
     ]
 
     @register_to_config
@@ -259,6 +321,13 @@ class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
         self.time_proj = MiniMaxH3PrunedTimeProj()
         self.time_embedder = MiniMaxH3PrunedTimeEmbedder(table_size=time_table_size, adaln_rank=adaln_rank)
 
+        # 2b. The affine map the fold was performed with: `silu(time_embedder(t)) ~= adaln_mean + c(t) @ adaln_basis`.
+        #     Nothing in `forward` reads these - the folded projections already carry them. They are stored so that a
+        #     LoRA trained on the released 2688-wide projections can be mapped onto these coordinates at load time;
+        #     see `load_lora_adapter`. 97 KB per partition.
+        self.register_buffer("adaln_basis", torch.zeros(adaln_rank, time_embed_dim), persistent=True)
+        self.register_buffer("adaln_mean", torch.zeros(time_embed_dim), persistent=True)
+
         # 3. Rotary embedding over the packed (t, h, w) grid
         self.rope = MiniMaxH3RotaryPosEmbed(rope_freq_dim=rope_freq_dim, rope_theta=rope_theta)
 
@@ -303,6 +372,140 @@ class MiniMaxH3PrunedTransformer3DModel(MiniMaxH3Transformer3DModel):
 
         self.gradient_checkpointing = False
 
+    # -- LoRA ---------------------------------------------------------------------------------------------------
+
+    def project_adaln_lora(self, state_dict: dict, prefix: str | None = None) -> tuple[dict, dict]:
+        r"""Map a LoRA's AdaLN factors from the released timestep embedding onto the pruned coordinates.
+
+        The released projection reads `x = silu(time_embedder(t))`, so a LoRA on it contributes
+
+            lora_B @ (lora_A @ x)  =  lora_B @ (lora_A @ (mean + basis.T @ c))
+                                   =  (lora_B @ (lora_A @ basis.T)) @ c  +  lora_B @ (lora_A @ mean)
+
+        which is a rank-preserving `[rank, adaln_rank]` `lora_A` over the pruned coordinates plus a constant output
+        offset. Both are computed in float64 from the file's own factors; the only error is the rank-8 subspace's
+        own residual on the timestep curve, 1.5e-5 relative, which is ~250x below one bfloat16 step of the weights
+        being adapted.
+
+        Returns `(state_dict, {module path: offset})`, the state dict unchanged and the offsets empty when the AdaLN
+        factors are already `adaln_rank`-wide (a LoRA trained against a pruned checkpoint - most of them).
+
+        Every AdaLN module in a file has to be one or the other. A file that mixes widths is not something this can
+        half-apply, so it raises.
+
+        `prefix` scopes this to one component's keys, the same way `load_lora_adapter` scopes the load - a file that
+        names both partitions holds two different adapters under one roof, and only one of them is going into this
+        module.
+        """
+        adaln_rank = self.config.adaln_rank
+        time_embed_dim = self.config.time_embed_dim
+
+        modules = {}
+        for key in state_dict:
+            if prefix is not None and not key.startswith(f"{prefix}."):
+                continue
+            match = ADALN_LORA_A_KEY.search(key)
+            if match is not None:
+                modules[key] = match.group(1)
+        if not modules:
+            return state_dict, {}
+
+        widths = sorted({int(state_dict[key].shape[1]) for key in modules})
+        if widths == [adaln_rank]:
+            return state_dict, {}
+        if widths != [time_embed_dim]:
+            odd = sorted(
+                {modules[key] for key in modules if int(state_dict[key].shape[1]) != max(widths)},
+                key=lambda name: (name != "norm_out.linear", name),
+            )
+            raise ValueError(
+                f"This LoRA's {len(modules)} AdaLN projections read inputs of width {widths}. On this checkpoint they "
+                f"have to be uniformly {adaln_rank} wide (trained against a pruned checkpoint, loaded as they are) or "
+                f"uniformly {time_embed_dim} wide (trained against the released checkpoint, projected onto the pruned "
+                "coordinates at load). An adapter cannot be applied to some of its AdaLN modules and not others, so "
+                f"nothing was loaded. The minority width is on: {odd[:4]}{' ...' if len(odd) > 4 else ''}."
+            )
+
+        basis = self.adaln_basis
+        mean = self.adaln_mean
+        if basis.abs().sum() == 0:
+            raise ValueError(
+                "Projecting a released-checkpoint LoRA needs `adaln_basis` and `adaln_mean`, the affine map this "
+                "checkpoint's AdaLN projections were folded with, and this model was loaded without them. Re-download "
+                "the repository: they ship as `adaln_affine.safetensors` next to the weights."
+            )
+        basis = basis.double()
+        mean = mean.double()
+
+        projected = dict(state_dict)
+        offsets = {}
+        for key, module in modules.items():
+            lora_b_key = key[: -len("lora_A.weight")] + "lora_B.weight"
+            if lora_b_key not in state_dict:
+                raise ValueError(
+                    f"{key} has no matching {lora_b_key}. The constant term of the projection is "
+                    "`lora_B @ (lora_A @ mean)`, so both factors have to be present; nothing was loaded."
+                )
+            lora_a = state_dict[key].to(device=basis.device, dtype=torch.float64)
+            lora_b = state_dict[lora_b_key].to(device=basis.device, dtype=torch.float64)
+            # float32 rather than the file's bfloat16: the projection is exact arithmetic on the file's factors and
+            # there is no reason to round it twice. PEFT casts to the adapter's dtype when it loads them.
+            projected[key] = (lora_a @ basis.T).to(torch.float32).cpu().contiguous()
+            offsets[module] = (lora_b @ (lora_a @ mean)).to(torch.float32).cpu().contiguous()
+        return projected, offsets
+
+    def load_lora_adapter(self, pretrained_model_name_or_path_or_dict, prefix="transformer", hotswap=False, **kwargs):
+        r"""`PeftAdapterMixin.load_lora_adapter`, with the AdaLN projection of [`project_adaln_lora`] in front of it.
+
+        LoRAs trained against a pruned checkpoint pass through untouched. LoRAs trained against the released
+        checkpoint's 2688-wide AdaLN projections - the official turbo LoRA and its conversions - are mapped onto the
+        pruned coordinates here, which is the only thing that ever stopped them loading. Everything outside the AdaLN
+        path is identical between the two checkpoints and is neither inspected nor changed.
+        """
+        state_dict = pretrained_model_name_or_path_or_dict
+        if not isinstance(state_dict, dict):
+            from diffusers.loaders.lora_base import _fetch_state_dict
+
+            state_dict, _ = _fetch_state_dict(
+                pretrained_model_name_or_path_or_dict=state_dict,
+                weight_name=kwargs.get("weight_name"),
+                use_safetensors=kwargs.get("use_safetensors", True),
+                local_files_only=kwargs.get("local_files_only"),
+                cache_dir=kwargs.get("cache_dir"),
+                force_download=kwargs.get("force_download", False),
+                proxies=kwargs.get("proxies"),
+                token=kwargs.get("token"),
+                revision=kwargs.get("revision"),
+                subfolder=kwargs.get("subfolder"),
+                user_agent={"file_type": "attn_procs_weights", "framework": "pytorch"},
+                allow_pickle=False,
+                metadata=kwargs.get("metadata"),
+            )
+
+        state_dict, offsets = self.project_adaln_lora(state_dict, prefix=prefix)
+        if offsets:
+            logger.info(
+                f"Projecting {len(offsets)} AdaLN LoRA modules from the released {self.config.time_embed_dim}-wide "
+                f"timestep embedding onto this checkpoint's {self.config.adaln_rank} pruned coordinates; each one's "
+                "constant term is carried as a float32 offset on the modulation."
+            )
+
+        before = set(getattr(self, "peft_config", None) or ())
+        super().load_lora_adapter(state_dict, prefix=prefix, hotswap=hotswap, **kwargs)
+        if not offsets:
+            return
+
+        added = set(getattr(self, "peft_config", None) or ()) - before
+        adapter_name = added.pop() if len(added) == 1 else kwargs.get("adapter_name")
+        if adapter_name is None:
+            raise RuntimeError(
+                "The AdaLN projection could not tell which adapter was just loaded, so its constant terms were not "
+                f"attached and the adapter is incomplete. Adapters before: {sorted(before)}, after: "
+                f"{sorted(getattr(self, 'peft_config', None) or ())}. Pass `adapter_name` explicitly."
+            )
+        for module, offset in offsets.items():
+            self.get_submodule(module.rsplit(".", 1)[0]).register_lora_adaln_offset(adapter_name, offset)
+
     # -- 8-bit compute -----------------------------------------------------------------------------------------
     #
     # `enable_convrot` and `quantize_8bit` are opt-in and change nothing until called.