Default to the ckpt500 turbo weights, add a per-request LoRA toggle

- H3_LORA default is now minimax_h3_turbo_4step_ckpt500.safetensors, the
model card's recommended default (newest, sharpest, non-EMA)
- the low-rank factors stay resident after the startup fold, so
h3_lora.set_enabled() folds/unfolds B@A in place per request; the AoTI
blocks read the same live storage. UI gets a 'Turbo LoRA' checkbox
(default on), the API a use_lora parameter (default true)
- default steps 4 -> 6 per the card's updated guidance (comfort zone is
6-8 at this checkpoint; 4 is softer)

README.md

@@ -47,11 +47,12 @@ callers get the same treatment the old upload event gave.
 ## 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.
+folded into its bf16 weights at startup (`h3_lora.py`), so the default is **6 sampling steps** instead of 28 (the
+card's comfort zone at the current checkpoint; 4 is the design point but softer). 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. The low-rank factors stay resident, so the
+**Turbo LoRA checkbox** (or the `use_lora` API parameter) folds/unfolds `lora_B @ lora_A` in place per request —
+off gives the base model at whatever step count you pick. `H3_LORA` selects the file (`off` disables entirely).
 
 ## AoTI-compiled blocks
 
@@ -175,7 +176,7 @@ 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` | `minimax_h3_turbo_4step_ckpt500.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

@@ -224,7 +224,7 @@ _DECODE_BASE, _DECODE_PER_DEFAULT_CANVAS, _DEFAULT_CANVAS_PIXELS = 15, 15, 960 *
 _PLACEMENT_ALLOWANCE, _PAD = 12, 10
 
 
-def get_duration(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, *a, **k):
+def get_duration(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, use_lora=True, *a, **k):
     height, width, num_frames, steps = int(height), int(width), int(num_frames), int(steps)
     latent_frames = (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
     patches = (height // 32) * (width // 32)
@@ -235,7 +235,7 @@ def get_duration(prompt_embeds, text_token_tags, image, last_image, height, widt
 
 
 @spaces.GPU(duration=get_duration, size=GPU_SIZE)
-def _generate(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed):
+def _generate(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, use_lora=True):
     """The only thing on GPU time: the packed-sequence denoise loop and the two decoders.
 
     Only the three generated outputs come back — a `@spaces.GPU` return crosses a process boundary by pickling, and
@@ -243,6 +243,11 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
     """
     import torch
 
+    import h3_lora
+
+    # Fold/unfold the turbo LoRA in place to match the request (a no-op when the state already matches).
+    h3_lora.set_enabled(PIPE.transformer, bool(use_lora))
+
     if PLACEMENT == "lazy":
         PIPE.to("cuda")
     elif PLACEMENT == "pack":
@@ -297,7 +302,7 @@ def _fit_keyframe(image_path, current_canvas):
     return image_path, label
 
 
-def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=4, seed=42, upsample=False, ip_token=None):
+def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=6, seed=42, upsample=False, use_lora=True, ip_token=None):
     """One request. `upsample` is last and defaults off, so a positional API client that predates it is unaffected."""
     if LOAD_ERROR:
         raise Exception(LOAD_ERROR)
@@ -345,6 +350,7 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
         num_frames,
         steps,
         seed,
+        use_lora,
     )
     generate_seconds = time.time() - started
 
@@ -356,7 +362,8 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
         f"{width}x{height} · {num_frames} frames ({num_frames / FPS:.3f} s) · {int(steps)} steps · "
         f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
         f"{', upsampled' if refined else ''}) · "
-        f"denoise + decode {generate_seconds:.0f}s ({generate_seconds / int(steps):.1f} s/step) · seed {int(seed)}"
+        f"denoise + decode {generate_seconds:.0f}s ({generate_seconds / int(steps):.1f} s/step) · "
+        f"turbo LoRA {'on' if use_lora else 'off'} · seed {int(seed)}"
     )
     print(f"[gen] {report}", flush=True)
     return FileData(path=path), report, refined
@@ -372,12 +379,12 @@ app = Server(title="MiniMax-H3 Studio")
 
 @app.api(name="generate")
 def _generate_api(prompt: str, image_path: FileData | None = None, last_image_path: FileData | None = None,
-                  canvas: str = DEFAULT_CANVAS, duration: float = 5, steps: int = 4, seed: float = 42,
-                  upsample: bool = False, request: Request = None) -> tuple[FileData, str, str]:
+                  canvas: str = DEFAULT_CANVAS, duration: float = 5, steps: int = 6, seed: float = 42,
+                  upsample: bool = False, use_lora: bool = True, request: Request = None) -> tuple[FileData, str, str]:
     """Generate a video with a synchronized soundtrack. Returns (video, report, refined prompt)."""
     # `request` is injected by the event system, not an API input; its x-ip-token bills the conditioner to the caller.
     ip_token = request.headers.get("x-ip-token") if request is not None else None
-    return generate(prompt, image_path, last_image_path, canvas, duration, steps, seed, upsample, ip_token=ip_token)
+    return generate(prompt, image_path, last_image_path, canvas, duration, steps, seed, upsample, use_lora, ip_token=ip_token)
 
 
 @app.get("/status")

h3_lora.py

@@ -27,8 +27,8 @@ 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")
+# The recommended default per the model card: ckpt500, non-EMA (newest, sharpest). `off` disables.
+LORA_FILE = os.environ.get("H3_LORA", "minimax_h3_turbo_4step_ckpt500.safetensors")
 
 
 def _delta_targets(name: str, delta: torch.Tensor, inner_dim: int) -> list[tuple[str, torch.Tensor]]:
@@ -75,14 +75,46 @@ def apply_lora(transformer) -> str | None:
     params = dict(transformer.named_parameters())
 
     folded = 0
+    entries = []
     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):
+        targets = _delta_targets(name, delta, inner_dim)
+        for key, converted in targets:
             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
+        entries.append((name, [key for key, _ in targets]))
+
+    # Keep the low-rank factors (~744 MB in bf16) so the fold can be toggled per request: adding/subtracting
+    # B @ A in place round-trips through one bf16 rounding, and the AoTI blocks read the same live storage.
+    transformer._lora_state = {
+        "enabled": True,
+        "inner_dim": inner_dim,
+        "entries": [
+            (name, lora[f"{name}.lora_A.weight"], lora[f"{name}.lora_B.weight"], keys)
+            for name, keys in entries
+        ],
+    }
     return f"LoRA `{LORA_REPO}/{LORA_FILE}` folded into {folded} weights ({len(bases)} modules)"
+
+
+def set_enabled(transformer, enabled: bool) -> bool:
+    """Fold/unfold the LoRA in place. No-op when no LoRA was loaded or the state already matches."""
+    state = getattr(transformer, "_lora_state", None)
+    if state is None or state["enabled"] == enabled:
+        return state["enabled"] if state else False
+    import torch
+
+    sign = 1.0 if enabled else -1.0
+    params = dict(transformer.named_parameters())
+    for name, a, b, keys in state["entries"]:
+        delta = b.to(torch.float32) @ a.to(torch.float32)
+        for key, converted in zip(keys, _delta_targets(name, delta, state["inner_dim"])):
+            param = params[key]
+            param.data = (param.data.float() + sign * converted.to(param.device)).to(param.dtype)
+    state["enabled"] = enabled
+    return enabled

index.html

@@ -180,6 +180,7 @@
       <div class="deck-label">Prompt</div>
       <textarea id="prompt" spellcheck="false">A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot</textarea>
     </div>
+    <label class="check"><input type="checkbox" id="use_lora" checked> Turbo LoRA (fast steps)</label>
     <label class="check"><input type="checkbox" id="upsample"> Upsample prompt</label>
 
     <div>
@@ -202,8 +203,8 @@
           <input type="range" id="duration" min="2" max="14" step="1" value="5">
         </div>
         <div>
-          <div class="slider-row"><span>Steps</span><output id="steps-out">4</output></div>
-          <input type="range" id="steps" min="2" max="40" step="1" value="4">
+          <div class="slider-row"><span>Steps</span><output id="steps-out">6</output></div>
+          <input type="range" id="steps" min="2" max="40" step="1" value="6">
         </div>
         <div>
           <div class="deck-label">Seed</div>
@@ -351,6 +352,7 @@ $("run").addEventListener("click", async () => {
       steps: Number($("steps").value),
       seed: Number($("seed").value),
       upsample: $("upsample").checked,
+      use_lora: $("use_lora").checked,
     });
     let data = result.data;
     // Server mode returns the tuple as ONE Api output: unwrap [[video, report, refined]] too.