Fold the 4-step Turbo LoRA into the transformer, default to 4 steps

Downloads larryvrh/MiniMax-H3-Turbo-Lora at startup and merges W + B@A
into the bf16 weights, mirroring the diffusers key conversion (fused-QKV
thirds, SwiGLU gate/value swap, shared AdaLN row layout). Folded before
the AoTI package is patched in so compiled blocks carry the update.
H3_LORA / H3_LORA_REPO configure it, H3_LORA=off disables.

README.md

@@ -35,6 +35,15 @@ Space is therefore impossible, which is why quantized demos of it run NVFP4 or f
 Besides the quality argument, unquantized weights are the ones AoTI can export; an NVFP4 checkpoint cannot be
 exported at all.
 
+## 4-step Turbo LoRA
+
+The transformer runs with [`larryvrh/MiniMax-H3-Turbo-Lora`](https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora)
+folded into its bf16 weights at startup (`h3_lora.py`), so the default is **4 sampling steps** instead of 28 — the
+LoRA's design point. The checkpoint is an early preview: sharper than the base model at 4 steps, but not a finished
+run. The fold mirrors the diffusers key conversion exactly (fused-QKV thirds, the `SwiGLU` gate/value swap, the
+shared AdaLN row layout) and happens before the AoTI package is patched in, so compiled blocks carry the update too.
+`H3_LORA=minimax_h3_turbo_4step_ema.safetensors` picks the smoother time-averaged variant, `H3_LORA=off` disables.
+
 ## AoTI-compiled blocks
 
 With `H3_AOTI=1` the 50 repeated transformer blocks run from a compiled package,
@@ -157,6 +166,8 @@ one-time `PIPE.to("cuda")` is inside the first row's 339 s and does not reappear
 | `H3_ATTENTION` | `_native_cudnn` | cuDNN's fused kernel, 10–20% faster than the SDPA default and needs nothing installed. flash-attention 3 is sm90-only and this pool is sm120. |
 | `H3_GPU_DURATION` | `900` | Seconds per request; the pool applies a 1.5 duration factor. |
 | `H3_GPU_SIZE` | `xlarge` | ZeroGPU allocation size. `large` does not fit. |
+| `H3_LORA` | `minimax_h3_turbo_4step.safetensors` | Turbo LoRA file folded into the transformer at startup. `off` disables. |
+| `H3_LORA_REPO` | `larryvrh/MiniMax-H3-Turbo-Lora` | Hub repo the LoRA is fetched from. |
 
 ## Whose GPU quota pays
 

app.py

@@ -73,6 +73,7 @@ PIPE = None
 MANAGER = None
 LOAD_ERROR: str | None = None
 LOADED_IN: float | None = None
+LORA_STATUS: str | None = None
 
 
 def status() -> str:
@@ -84,7 +85,8 @@ def status() -> str:
 
     return (
         f"Ready · transformer + VAEs **bfloat16, unquantized** · placement `{PLACEMENT}` · attention `{ATTENTION}` · "
-        f"{h3_aoti.status()} · loaded in {LOADED_IN:.0f}s · conditioner `{CONDITIONER_SPACE}`"
+        f"{h3_aoti.status()} · {LORA_STATUS or 'no LoRA'} · loaded in {LOADED_IN:.0f}s · "
+        f"conditioner `{CONDITIONER_SPACE}`"
     )
 
 
@@ -96,7 +98,7 @@ def load_models() -> str | None:
     Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a bfloat16 audio VAE decodes
     the soundtrack roughly 20 dB too quiet.
     """
-    global PIPE, MANAGER, LOAD_ERROR, LOADED_IN
+    global PIPE, MANAGER, LOAD_ERROR, LOADED_IN, LORA_STATUS
 
     if PIPE is not None or LOAD_ERROR is not None:
         return LOAD_ERROR
@@ -114,6 +116,15 @@ def load_models() -> str | None:
         print(f"[gen] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
         pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
         pipe.load_components(dtype=torch.bfloat16)
+
+        # Fold the 4-step Turbo LoRA into the bf16 weights before AoTI packages the blocks, so the compiled forward
+        # reads weights that already carry the update. `H3_LORA=off` disables.
+        import h3_lora
+
+        LORA_STATUS = h3_lora.apply_lora(pipe.transformer)
+        if LORA_STATUS:
+            print(f"[gen] {LORA_STATUS}", flush=True)
+
         pipe.transformer.set_attention_backend(ATTENTION)
 
         # Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU
@@ -237,7 +248,7 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
     return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")
 
 
-def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=28, seed=42, upsample=False, progress=gr.Progress(track_tqdm=True)):
+def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=4, seed=42, upsample=False, progress=gr.Progress(track_tqdm=True)):
     """One request. `upsample` is last and defaults off, so a positional API client that predates it is unaffected."""
     if LOAD_ERROR:
         raise gr.Error(LOAD_ERROR)
@@ -370,7 +381,7 @@ with gr.Blocks(title="MiniMax-H3") as demo:
             with gr.Accordion("Advanced options", open=False):
                 canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                 duration = gr.Slider(label="Duration (s)", minimum=MIN_UI_DURATION, maximum=MAX_UI_DURATION, step=1, value=5)
-                steps = gr.Slider(label="Steps", minimum=10, maximum=40, step=1, value=28)
+                steps = gr.Slider(label="Steps", minimum=2, maximum=40, step=1, value=4)
                 seed = gr.Number(label="Seed", value=42, precision=0)
 
         with gr.Column():

h3_lora.py

@@ -0,0 +1,88 @@
+"""Fold the MiniMax-H3 Turbo LoRA (`larryvrh/MiniMax-H3-Turbo-Lora`) into the diffusers transformer.
+
+The LoRA ships against the *reference* (ComfyUI) module tree — `blocks.N.attn.qkv_proj`, `blocks.N.mlp.fc1`,
+`token_refiner.blocks.N`, `final_layer.adaln_proj.linear` — with `alpha == rank`, so the update is exactly
+`W + lora_B @ lora_A`. The diffusers checkpoint is the same weights under different names and two layout transforms
+(see `scripts/convert_minimax_h3_to_diffusers.py` in huggingface/diffusers#14371), so each delta gets the same
+transform the base weight got:
+
+* fused `attn.qkv_proj` rows are `[q_all; k_all; v_all]` in both in-memory layouts -> split into contiguous thirds
+  onto `attn.to_q` / `to_k` / `to_v`;
+* fused `mlp.fc1` is `[gate; value]` while diffusers' `SwiGLU` fuses `[value; gate]` -> swap the halves onto
+  `ff.net.0.proj`;
+* `mlp.fc2` -> `ff.net.2`, `attn.out_proj` -> `attn.to_out.0`, `blocks.` -> `transformer_blocks.`,
+  `token_refiner.blocks.` -> `token_refiner.refiner_blocks.`, `final_layer.adaln_proj.linear` -> `norm_out.linear`;
+* the `adaln_proj.linear` modulation tables share the `[timestep][modality][param]` row layout in both trees, so
+  they map name-for-name with no reordering.
+
+The delta is folded into the bf16 weights rather than applied as a runtime wrapper for one reason: the AoTI block
+package (`h3_aoti`) reads each block's live weights, and a wrapper module would be invisible to it. The fold computes
+`lora_B @ lora_A` in float32 and rounds once on the way back into bf16.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+
+LORA_REPO = os.environ.get("H3_LORA_REPO", "larryvrh/MiniMax-H3-Turbo-Lora")
+# The trained weights; the `_ema` variant is an immature time-averaged snapshot at this checkpoint. `off` disables.
+LORA_FILE = os.environ.get("H3_LORA", "minimax_h3_turbo_4step.safetensors")
+
+
+def _delta_targets(name: str, delta: torch.Tensor, inner_dim: int) -> list[tuple[str, torch.Tensor]]:
+    """Map one reference-tree LoRA base name and its `lora_B @ lora_A` delta onto diffusers parameter key(s)."""
+    if name.startswith("token_refiner.blocks."):
+        target = name.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
+    elif name.startswith("blocks."):
+        target = name.replace("blocks.", "transformer_blocks.", 1)
+    else:
+        target = name
+    target = target.replace("final_layer.adaln_proj.linear", "norm_out.linear")
+
+    if target.endswith(".attn.qkv_proj"):
+        prefix = target.removesuffix("qkv_proj")
+        return [
+            (f"{prefix}to_{kind}.weight", part.contiguous())
+            for kind, part in zip(("q", "k", "v"), delta.split(inner_dim, dim=0))
+        ]
+    if target.endswith(".mlp.fc1"):
+        gate, value = delta.chunk(2, dim=0)
+        return [(target.replace(".mlp.fc1", ".ff.net.0.proj") + ".weight", torch.cat([value, gate]).contiguous())]
+    if target.endswith(".mlp.fc2"):
+        return [(target.replace(".mlp.fc2", ".ff.net.2") + ".weight", delta)]
+    if target.endswith(".attn.out_proj"):
+        return [(target.replace(".attn.out_proj", ".attn.to_out.0") + ".weight", delta)]
+    # `adaln_proj.linear` (block-level and the final `norm_out.linear`): identical row layout on both sides.
+    return [(target + ".weight", delta)]
+
+
+def apply_lora(transformer) -> str | None:
+    """Fold the configured Turbo LoRA into `transformer` in place. Returns a status line, or `None` when disabled."""
+    if LORA_FILE.lower() in ("", "off", "none"):
+        return None
+
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    path = hf_hub_download(LORA_REPO, LORA_FILE)
+    lora = load_file(path)
+    bases = sorted({key.rsplit(".lora_", 1)[0] for key in lora})
+
+    config = transformer.config
+    inner_dim = config.num_attention_heads * config.attention_head_dim
+    params = dict(transformer.named_parameters())
+
+    folded = 0
+    for name in bases:
+        a = lora[f"{name}.lora_A.weight"].float()
+        b = lora[f"{name}.lora_B.weight"].float()
+        delta = b @ a  # alpha == rank, so the scale is 1
+        for key, converted in _delta_targets(name, delta, inner_dim):
+            param = params.get(key)
+            if param is None:
+                raise KeyError(f"LoRA target `{key}` (from `{name}`) not found in the transformer")
+            param.data = (param.data.float() + converted).to(param.dtype)
+            folded += 1
+    return f"LoRA `{LORA_REPO}/{LORA_FILE}` folded into {folded} weights ({len(bases)} modules)"