akhaliq HF Staff commited on
Commit
456aa27
·
1 Parent(s): 25c3db5

Add the lightx2v Minimax-h3-Turbo LoRA as a per-request alternative

Browse files

h3_lora.py now loads both turbo LoRAs and switches the folded set in
place per request (unfold old / fold new, one bf16 rounding, AoTI-safe).
The lightx checkpoint is PEFT-style against the diffusers tree (rank
128, alpha 8 -> fold scale 0.0625); the larry one keeps the reference
key conversion. The UI checkbox becomes a LoRA dropdown (larry / lightx
/ off) that suggests each design point in steps (6 / 4 / 28), and the
API gains a 'lora' string parameter — the legacy use_lora bool still
works.

Files changed (4) hide show
  1. README.md +9 -6
  2. app.py +40 -13
  3. h3_lora.py +130 -83
  4. index.html +20 -2
README.md CHANGED
@@ -48,11 +48,12 @@ callers get the same treatment the old upload event gave.
48
 
49
  The transformer runs with [`larryvrh/MiniMax-H3-Turbo-Lora`](https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora)
50
  folded into its bf16 weights at startup (`h3_lora.py`), so the default is **6 sampling steps** instead of 28
51
- card's comfort zone at the current checkpoint; 4 is the design point but softer). The fold mirrors the diffusers key
52
- conversion exactly (fused-QKV thirds, the `SwiGLU` gate/value swap, the shared AdaLN row layout) and happens before
53
- the AoTI package is patched in, so compiled blocks carry the update too. The low-rank factors stay resident, so the
54
- **Turbo LoRA checkbox** (or the `use_lora` API parameter) folds/unfolds `lora_B @ lora_A` in place per request
55
- off gives the base model at whatever step count you pick. `H3_LORA` selects the file (`off` disables entirely), `H3_LORA_STRENGTH` scales the update (the card's sharpness/artifact dial: >1 against smear, <1 against over-sharp grain).
 
56
 
57
  ## AoTI-compiled blocks
58
 
@@ -178,7 +179,9 @@ one-time `PIPE.to("cuda")` is inside the first row's 339 s and does not reappear
178
  | `H3_GPU_SIZE` | `xlarge` | ZeroGPU allocation size. `large` does not fit. |
179
  | `H3_LORA` | `minimax_h3_turbo_4step_ema_ckpt850.safetensors` | Turbo LoRA file folded into the transformer at startup. `off` disables. |
180
  | `H3_LORA_REPO` | `larryvrh/MiniMax-H3-Turbo-Lora` | Hub repo the LoRA is fetched from. |
181
- | `H3_LORA_STRENGTH` | `1.0` | Scales the folded LoRA delta (sharpness/artifact trade-off). |
 
 
182
 
183
  ## Whose GPU quota pays
184
 
 
48
 
49
  The transformer runs with [`larryvrh/MiniMax-H3-Turbo-Lora`](https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora)
50
  folded into its bf16 weights at startup (`h3_lora.py`), so the default is **6 sampling steps** instead of 28
51
+ card's comfort zone at the current checkpoint; 4 is the design point but softer). The larry fold mirrors the diffusers key
52
+ conversion exactly (fused-QKV thirds, the `SwiGLU` gate/value swap, the shared AdaLN row layout); lightx keys are
53
+ already diffusers-native. Both happen before the AoTI package is patched in, so compiled blocks carry the update too,
54
+ and the low-rank factors stay resident so switching is an in-place unfold/fold through one bf16 rounding.
55
+ `H3_LORA` selects the larry file (`off` skips it), `H3_LIGHTX=off` skips lightx, `H3_LORA_DEFAULT` picks which set
56
+ starts folded, and `H3_LORA_STRENGTH` scales the larry update (the card's sharpness/artifact dial).
57
 
58
  ## AoTI-compiled blocks
59
 
 
179
  | `H3_GPU_SIZE` | `xlarge` | ZeroGPU allocation size. `large` does not fit. |
180
  | `H3_LORA` | `minimax_h3_turbo_4step_ema_ckpt850.safetensors` | Turbo LoRA file folded into the transformer at startup. `off` disables. |
181
  | `H3_LORA_REPO` | `larryvrh/MiniMax-H3-Turbo-Lora` | Hub repo the LoRA is fetched from. |
182
+ | `H3_LORA_STRENGTH` | `1.0` | Scales the larry LoRA delta (sharpness/artifact trade-off). |
183
+ | `H3_LIGHTX` | `on` | Set to `off` to skip loading the lightx2v LoRA set. |
184
+ | `H3_LORA_DEFAULT` | `larry` | Which loaded LoRA set starts folded (`larry` / `lightx`). |
185
 
186
  ## Whose GPU quota pays
187
 
app.py CHANGED
@@ -224,7 +224,7 @@ _DECODE_BASE, _DECODE_PER_DEFAULT_CANVAS, _DEFAULT_CANVAS_PIXELS = 15, 15, 960 *
224
  _PLACEMENT_ALLOWANCE, _PAD = 12, 10
225
 
226
 
227
- def get_duration(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, use_lora=True, *a, **k):
228
  height, width, num_frames, steps = int(height), int(width), int(num_frames), int(steps)
229
  latent_frames = (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
230
  patches = (height // 32) * (width // 32)
@@ -235,7 +235,7 @@ def get_duration(prompt_embeds, text_token_tags, image, last_image, height, widt
235
 
236
 
237
  @spaces.GPU(duration=get_duration, size=GPU_SIZE)
238
- def _generate(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, use_lora=True):
239
  """The only thing on GPU time: the packed-sequence denoise loop and the two decoders.
240
 
241
  Only the three generated outputs come back — a `@spaces.GPU` return crosses a process boundary by pickling, and
@@ -245,8 +245,9 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
245
 
246
  import h3_lora
247
 
248
- # Fold/unfold the turbo LoRA in place to match the request (a no-op when the state already matches).
249
- h3_lora.set_enabled(PIPE.transformer, bool(use_lora))
 
250
 
251
  if PLACEMENT == "lazy":
252
  PIPE.to("cuda")
@@ -265,7 +266,7 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
265
  num_inference_steps=int(steps),
266
  generator=torch.Generator("cpu").manual_seed(int(seed)),
267
  )
268
- return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")
269
 
270
 
271
  def _fit_keyframe(image_path, current_canvas):
@@ -302,8 +303,15 @@ def _fit_keyframe(image_path, current_canvas):
302
  return image_path, label
303
 
304
 
305
- 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):
306
- """One request. `upsample` is last and defaults off, so a positional API client that predates it is unaffected."""
 
 
 
 
 
 
 
307
  if LOAD_ERROR:
308
  raise Exception(LOAD_ERROR)
309
  if PIPE is None:
@@ -315,6 +323,8 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
315
 
316
  from diffusers.utils import encode_video
317
 
 
 
318
  # Server mode: keyframes arrive as FileData dicts, and the cover-crop / canvas-fit that used to be an upload
319
  # event in the Blocks UI runs here instead, so API callers get the same treatment.
320
  first = image_path["path"] if isinstance(image_path, dict) else image_path
@@ -340,7 +350,7 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
340
  return ImageOps.exif_transpose(Image.open(path)).convert("RGB") if path else None
341
 
342
  started = time.time()
343
- frames, audio, sampling_rate = _generate(
344
  prompt_embeds,
345
  text_token_tags,
346
  keyframe(first),
@@ -350,7 +360,7 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
350
  num_frames,
351
  steps,
352
  seed,
353
- use_lora,
354
  )
355
  generate_seconds = time.time() - started
356
 
@@ -363,7 +373,7 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
363
  f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
364
  f"{', upsampled' if refined else ''}) · "
365
  f"denoise + decode {generate_seconds:.0f}s ({generate_seconds / int(steps):.1f} s/step) · "
366
- f"turbo LoRA {'on' if use_lora else 'off'} · seed {int(seed)}"
367
  )
368
  print(f"[gen] {report}", flush=True)
369
  return FileData(path=path), report, refined
@@ -380,11 +390,15 @@ app = Server(title="MiniMax-H3 Studio")
380
  @app.api(name="generate")
381
  def _generate_api(prompt: str, image_path: FileData | None = None, last_image_path: FileData | None = None,
382
  canvas: str = DEFAULT_CANVAS, duration: float = 5, steps: int = 6, seed: float = 42,
383
- upsample: bool = False, use_lora: bool = True, request: Request = None) -> tuple[FileData, str, str]:
384
- """Generate a video with a synchronized soundtrack. Returns (video, report, refined prompt)."""
 
 
 
 
385
  # `request` is injected by the event system, not an API input; its x-ip-token bills the conditioner to the caller.
386
  ip_token = request.headers.get("x-ip-token") if request is not None else None
387
- return generate(prompt, image_path, last_image_path, canvas, duration, steps, seed, upsample, use_lora, ip_token=ip_token)
388
 
389
 
390
  @app.get("/status")
@@ -397,11 +411,24 @@ def studio_status():
397
  @app.get("/studio-config")
398
  def studio_config():
399
  """The canvas table and slider ranges, so the frontend never hardcodes a label the backend would reject."""
 
 
 
 
400
  return {
401
  "canvases": list(CANVASES),
402
  "default_canvas": DEFAULT_CANVAS,
403
  "min_duration": MIN_UI_DURATION,
404
  "max_duration": MAX_UI_DURATION,
 
 
 
 
 
 
 
 
 
405
  }
406
 
407
 
 
224
  _PLACEMENT_ALLOWANCE, _PAD = 12, 10
225
 
226
 
227
+ def get_duration(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, lora="larry", *a, **k):
228
  height, width, num_frames, steps = int(height), int(width), int(num_frames), int(steps)
229
  latent_frames = (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
230
  patches = (height // 32) * (width // 32)
 
235
 
236
 
237
  @spaces.GPU(duration=get_duration, size=GPU_SIZE)
238
+ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, lora="larry"):
239
  """The only thing on GPU time: the packed-sequence denoise loop and the two decoders.
240
 
241
  Only the three generated outputs come back — a `@spaces.GPU` return crosses a process boundary by pickling, and
 
245
 
246
  import h3_lora
247
 
248
+ # Fold the requested LoRA in place (a no-op when the state already matches). AoTI blocks read the same
249
+ # live storage, so the compiled forward carries the switch too.
250
+ active_lora = h3_lora.set_active(PIPE.transformer, lora)
251
 
252
  if PLACEMENT == "lazy":
253
  PIPE.to("cuda")
 
266
  num_inference_steps=int(steps),
267
  generator=torch.Generator("cpu").manual_seed(int(seed)),
268
  )
269
+ return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate"), active_lora
270
 
271
 
272
  def _fit_keyframe(image_path, current_canvas):
 
303
  return image_path, label
304
 
305
 
306
+ def _resolve_lora(lora, use_lora) -> str:
307
+ """`lora` (`larry` / `lightx` / `off`) wins; the legacy `use_lora` bool maps onto `larry` / `off`."""
308
+ if isinstance(lora, str) and lora in ("larry", "lightx", "off"):
309
+ return lora
310
+ return "larry" if use_lora else "off"
311
+
312
+
313
+ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=6, seed=42, upsample=False, use_lora=True, lora="", ip_token=None):
314
+ """One request. `upsample`/`use_lora` keep their defaults so a positional API client that predates them is unaffected."""
315
  if LOAD_ERROR:
316
  raise Exception(LOAD_ERROR)
317
  if PIPE is None:
 
323
 
324
  from diffusers.utils import encode_video
325
 
326
+ lora = _resolve_lora(lora, use_lora)
327
+
328
  # Server mode: keyframes arrive as FileData dicts, and the cover-crop / canvas-fit that used to be an upload
329
  # event in the Blocks UI runs here instead, so API callers get the same treatment.
330
  first = image_path["path"] if isinstance(image_path, dict) else image_path
 
350
  return ImageOps.exif_transpose(Image.open(path)).convert("RGB") if path else None
351
 
352
  started = time.time()
353
+ frames, audio, sampling_rate, active_lora = _generate(
354
  prompt_embeds,
355
  text_token_tags,
356
  keyframe(first),
 
360
  num_frames,
361
  steps,
362
  seed,
363
+ lora,
364
  )
365
  generate_seconds = time.time() - started
366
 
 
373
  f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
374
  f"{', upsampled' if refined else ''}) · "
375
  f"denoise + decode {generate_seconds:.0f}s ({generate_seconds / int(steps):.1f} s/step) · "
376
+ f"turbo LoRA {active_lora} · seed {int(seed)}"
377
  )
378
  print(f"[gen] {report}", flush=True)
379
  return FileData(path=path), report, refined
 
390
  @app.api(name="generate")
391
  def _generate_api(prompt: str, image_path: FileData | None = None, last_image_path: FileData | None = None,
392
  canvas: str = DEFAULT_CANVAS, duration: float = 5, steps: int = 6, seed: float = 42,
393
+ upsample: bool = False, use_lora: bool = True, lora: str = "", request: Request = None) -> tuple[FileData, str, str]:
394
+ """Generate a video with a synchronized soundtrack. Returns (video, report, refined prompt).
395
+
396
+ `lora` selects the turbo LoRA: `larry` (default), `lightx`, or `off`. The legacy `use_lora` bool still works
397
+ when `lora` is empty.
398
+ """
399
  # `request` is injected by the event system, not an API input; its x-ip-token bills the conditioner to the caller.
400
  ip_token = request.headers.get("x-ip-token") if request is not None else None
401
+ return generate(prompt, image_path, last_image_path, canvas, duration, steps, seed, upsample, use_lora, lora, ip_token=ip_token)
402
 
403
 
404
  @app.get("/status")
 
411
  @app.get("/studio-config")
412
  def studio_config():
413
  """The canvas table and slider ranges, so the frontend never hardcodes a label the backend would reject."""
414
+ import h3_lora
415
+
416
+ state = getattr(PIPE.transformer, "_lora_state", None) if PIPE is not None else None
417
+ sets = state["sets"] if state else {}
418
  return {
419
  "canvases": list(CANVASES),
420
  "default_canvas": DEFAULT_CANVAS,
421
  "min_duration": MIN_UI_DURATION,
422
  "max_duration": MAX_UI_DURATION,
423
+ # The LoRA dropdown: value -> {label, suggested steps}.
424
+ "loras": {
425
+ **{
426
+ name: {"label": spec["label"], "steps": {"larry": 6, "lightx": 4}.get(name, 6)}
427
+ for name, spec in sets.items()
428
+ },
429
+ "off": {"label": "off (base model)", "steps": 28},
430
+ },
431
+ "default_lora": state["active"] if state else "off",
432
  }
433
 
434
 
h3_lora.py CHANGED
@@ -1,23 +1,28 @@
1
- """Fold the MiniMax-H3 Turbo LoRA (`larryvrh/MiniMax-H3-Turbo-Lora`) into the diffusers transformer.
2
-
3
- The LoRA ships against the *reference* (ComfyUI) module tree `blocks.N.attn.qkv_proj`, `blocks.N.mlp.fc1`,
4
- `token_refiner.blocks.N`, `final_layer.adaln_proj.linear` with `alpha == rank`, so the update is exactly
5
- `W + lora_B @ lora_A`. The diffusers checkpoint is the same weights under different names and two layout transforms
6
- (see `scripts/convert_minimax_h3_to_diffusers.py` in huggingface/diffusers#14371), so each delta gets the same
7
- transform the base weight got:
8
-
9
- * fused `attn.qkv_proj` rows are `[q_all; k_all; v_all]` in both in-memory layouts -> split into contiguous thirds
10
- onto `attn.to_q` / `to_k` / `to_v`;
11
- * fused `mlp.fc1` is `[gate; value]` while diffusers' `SwiGLU` fuses `[value; gate]` -> swap the halves onto
12
- `ff.net.0.proj`;
13
- * `mlp.fc2` -> `ff.net.2`, `attn.out_proj` -> `attn.to_out.0`, `blocks.` -> `transformer_blocks.`,
14
- `token_refiner.blocks.` -> `token_refiner.refiner_blocks.`, `final_layer.adaln_proj.linear` -> `norm_out.linear`;
15
- * the `adaln_proj.linear` modulation tables share the `[timestep][modality][param]` row layout in both trees, so
16
- they map name-for-name with no reordering.
17
-
18
- The delta is folded into the bf16 weights rather than applied as a runtime wrapper for one reason: the AoTI block
19
- package (`h3_aoti`) reads each block's live weights, and a wrapper module would be invisible to it. The fold computes
20
- `lora_B @ lora_A` in float32 and rounds once on the way back into bf16.
 
 
 
 
 
21
  """
22
 
23
  from __future__ import annotations
@@ -26,16 +31,19 @@ import os
26
 
27
  import torch
28
 
29
- LORA_REPO = os.environ.get("H3_LORA_REPO", "larryvrh/MiniMax-H3-Turbo-Lora")
30
  # The recommended default per the model card: ckpt850 EMA (final checkpoint of the round, sharp at 4 steps).
31
- # `off` disables.
32
- LORA_FILE = os.environ.get("H3_LORA", "minimax_h3_turbo_4step_ema_ckpt850.safetensors")
33
- # The card's sharpness/artifact dial: >1 against blurry ghosting/smear, <1 against over-sharp grain.
34
- LORA_STRENGTH = float(os.environ.get("H3_LORA_STRENGTH", "1.0"))
 
 
 
35
 
36
 
37
- def _delta_targets(name: str, delta: torch.Tensor, inner_dim: int) -> list[tuple[str, torch.Tensor]]:
38
- """Map one reference-tree LoRA base name and its `lora_B @ lora_A` delta onto diffusers parameter key(s)."""
39
  if name.startswith("token_refiner.blocks."):
40
  target = name.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
41
  elif name.startswith("blocks."):
@@ -48,78 +56,117 @@ def _delta_targets(name: str, delta: torch.Tensor, inner_dim: int) -> list[tuple
48
  prefix = target.removesuffix("qkv_proj")
49
  return [
50
  (f"{prefix}to_{kind}.weight", part.contiguous())
51
- for kind, part in zip(("q", "k", "v"), delta.split(inner_dim, dim=0))
52
  ]
53
  if target.endswith(".mlp.fc1"):
54
- gate, value = delta.chunk(2, dim=0)
55
  return [(target.replace(".mlp.fc1", ".ff.net.0.proj") + ".weight", torch.cat([value, gate]).contiguous())]
56
  if target.endswith(".mlp.fc2"):
57
- return [(target.replace(".mlp.fc2", ".ff.net.2") + ".weight", delta)]
58
  if target.endswith(".attn.out_proj"):
59
- return [(target.replace(".attn.out_proj", ".attn.to_out.0") + ".weight", delta)]
60
  # `adaln_proj.linear` (block-level and the final `norm_out.linear`): identical row layout on both sides.
61
- return [(target + ".weight", delta)]
62
 
63
 
64
- def apply_lora(transformer) -> str | None:
65
- """Fold the configured Turbo LoRA into `transformer` in place. Returns a status line, or `None` when disabled."""
66
- if LORA_FILE.lower() in ("", "off", "none"):
67
- return None
68
-
69
  from huggingface_hub import hf_hub_download
70
  from safetensors.torch import load_file
71
 
72
- path = hf_hub_download(LORA_REPO, LORA_FILE)
73
- lora = load_file(path)
74
  bases = sorted({key.rsplit(".lora_", 1)[0] for key in lora})
75
-
76
- config = transformer.config
77
- inner_dim = config.num_attention_heads * config.attention_head_dim
78
- params = dict(transformer.named_parameters())
79
-
80
- folded = 0
81
  entries = []
82
  for name in bases:
83
- a = lora[f"{name}.lora_A.weight"].float()
84
- b = lora[f"{name}.lora_B.weight"].float()
85
- delta = (b @ a) * LORA_STRENGTH # alpha == rank, so the base scale is 1
86
- targets = _delta_targets(name, delta, inner_dim)
87
- for key, converted in targets:
88
- param = params.get(key)
89
- if param is None:
90
- raise KeyError(f"LoRA target `{key}` (from `{name}`) not found in the transformer")
91
- param.data = (param.data.float() + converted).to(param.dtype)
92
- folded += 1
93
- entries.append(name)
94
-
95
- # Keep the low-rank factors (~744 MB in bf16) so the fold can be toggled per request: adding/subtracting
96
- # B @ A in place round-trips through one bf16 rounding, and the AoTI blocks read the same live storage.
97
- transformer._lora_state = {
98
- "enabled": True,
99
- "inner_dim": inner_dim,
100
- "strength": LORA_STRENGTH,
101
- "entries": [
102
- (name, lora[f"{name}.lora_A.weight"], lora[f"{name}.lora_B.weight"])
103
- for name in entries
104
- ],
105
  }
106
- strength = "" if LORA_STRENGTH == 1.0 else f" @ strength {LORA_STRENGTH}"
107
- return f"LoRA `{LORA_REPO}/{LORA_FILE}`{strength} folded into {folded} weights ({len(bases)} modules)"
108
 
109
 
110
- def set_enabled(transformer, enabled: bool) -> bool:
111
- """Fold/unfold the LoRA in place. No-op when no LoRA was loaded or the state already matches."""
112
- state = getattr(transformer, "_lora_state", None)
113
- if state is None or state["enabled"] == enabled:
114
- return state["enabled"] if state else False
115
- import torch
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
116
 
117
- sign = (1.0 if enabled else -1.0) * state.get("strength", 1.0)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
118
  params = dict(transformer.named_parameters())
119
- for name, a, b in state["entries"]:
120
- delta = b.to(torch.float32) @ a.to(torch.float32)
121
- for key, converted in _delta_targets(name, delta, state["inner_dim"]):
122
- param = params[key]
123
- param.data = (param.data.float() + sign * converted.to(param.device)).to(param.dtype)
124
- state["enabled"] = enabled
125
- return enabled
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Turbo LoRA support for the diffusers MiniMax-H3 transformer: two 4-step LoRAs, one fold mechanism.
2
+
3
+ Both LoRAs are applied by folding `scale * (lora_B @ lora_A)` into the bf16 weights rather than as runtime wrappers,
4
+ because the AoTI block package (`h3_aoti`) reads each block's live weights and a wrapper module would be invisible
5
+ to it. Deltas are computed in float32 and round once on the way back into bf16. The low-rank factors of every loaded
6
+ LoRA stay resident, so the active one can be switched per request (`set_active`) — unfold the old, fold the new, in
7
+ place, through one bf16 rounding.
8
+
9
+ The two supported LoRAs ship in different layouts:
10
+
11
+ * `larry` (`larryvrh/MiniMax-H3-Turbo-Lora`) targets the *reference* (ComfyUI) module tree `blocks.N.attn.qkv_proj`,
12
+ `blocks.N.mlp.fc1`, `token_refiner.blocks.N`, `final_layer.adaln_proj.linear` — with `alpha == rank` (scale 1).
13
+ Each delta gets the same transform the base weights got in the diffusers conversion
14
+ (`scripts/convert_minimax_h3_to_diffusers.py`, huggingface/diffusers#14371): fused-QKV row thirds onto
15
+ `attn.to_q/k/v`, the `SwiGLU` gate/value swap onto `ff.net.0.proj`, `fc2` -> `ff.net.2`,
16
+ `blocks.` -> `transformer_blocks.`, `token_refiner.blocks.` -> `token_refiner.refiner_blocks.`,
17
+ `final_layer.adaln_proj.linear` -> `norm_out.linear`. The row transforms are applied to `lora_B` directly
18
+ (rows of `B @ A` are rows of `B`), so no full delta is ever materialized at load.
19
+
20
+ * `lightx` (`lightx2v/Minimax-h3-Turbo`) is a PEFT checkpoint against the diffusers tree itself
21
+ `transformer_blocks.N.attn.to_q.lora_A.default.weight` and friends — rank 128, `alpha == 8`, so the fold scale is
22
+ `8 / 128 = 0.0625` (matching `set_adapters(weights=1.0)` in their inference script). Keys map name-for-name.
23
+
24
+ `H3_LORA` selects the larry file (`off` skips loading it), `H3_LIGHTX=off` skips lightx, `H3_LORA_DEFAULT` picks
25
+ which set starts folded, and `H3_LORA_STRENGTH` is the larry card's sharpness/artifact dial.
26
  """
27
 
28
  from __future__ import annotations
 
31
 
32
  import torch
33
 
34
+ LARRY_REPO = os.environ.get("H3_LORA_REPO", "larryvrh/MiniMax-H3-Turbo-Lora")
35
  # The recommended default per the model card: ckpt850 EMA (final checkpoint of the round, sharp at 4 steps).
36
+ LARRY_FILE = os.environ.get("H3_LORA", "minimax_h3_turbo_4step_ema_ckpt850.safetensors")
37
+ LIGHTX_REPO = os.environ.get("H3_LIGHTX_REPO", "lightx2v/Minimax-h3-Turbo")
38
+ LIGHTX_FILE = os.environ.get("H3_LIGHTX_FILE", "minimax_h3_fl2v_turbo_4step_v0.1.safetensors")
39
+ LIGHTX_ALPHA = 8
40
+ # The card's sharpness/artifact dial for the larry LoRA: >1 against blurry ghosting/smear, <1 against grain.
41
+ LARRY_STRENGTH = float(os.environ.get("H3_LORA_STRENGTH", "1.0"))
42
+ DEFAULT_LORA = os.environ.get("H3_LORA_DEFAULT", "larry")
43
 
44
 
45
+ def _larry_targets(name: str, b: torch.Tensor, inner_dim: int) -> list[tuple[str, torch.Tensor]]:
46
+ """Map one reference-tree base name and its `lora_B` onto diffusers parameter key + row-transformed B."""
47
  if name.startswith("token_refiner.blocks."):
48
  target = name.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
49
  elif name.startswith("blocks."):
 
56
  prefix = target.removesuffix("qkv_proj")
57
  return [
58
  (f"{prefix}to_{kind}.weight", part.contiguous())
59
+ for kind, part in zip(("q", "k", "v"), b.split(inner_dim, dim=0))
60
  ]
61
  if target.endswith(".mlp.fc1"):
62
+ gate, value = b.chunk(2, dim=0)
63
  return [(target.replace(".mlp.fc1", ".ff.net.0.proj") + ".weight", torch.cat([value, gate]).contiguous())]
64
  if target.endswith(".mlp.fc2"):
65
+ return [(target.replace(".mlp.fc2", ".ff.net.2") + ".weight", b)]
66
  if target.endswith(".attn.out_proj"):
67
+ return [(target.replace(".attn.out_proj", ".attn.to_out.0") + ".weight", b)]
68
  # `adaln_proj.linear` (block-level and the final `norm_out.linear`): identical row layout on both sides.
69
+ return [(target + ".weight", b)]
70
 
71
 
72
+ def _load_larry(inner_dim: int) -> dict:
 
 
 
 
73
  from huggingface_hub import hf_hub_download
74
  from safetensors.torch import load_file
75
 
76
+ lora = load_file(hf_hub_download(LARRY_REPO, LARRY_FILE))
 
77
  bases = sorted({key.rsplit(".lora_", 1)[0] for key in lora})
 
 
 
 
 
 
78
  entries = []
79
  for name in bases:
80
+ a = lora[f"{name}.lora_A.weight"]
81
+ b = lora[f"{name}.lora_B.weight"]
82
+ entries.extend((key, a, b_part) for key, b_part in _larry_targets(name, b, inner_dim))
83
+ return {
84
+ "label": f"{LARRY_REPO}/{LARRY_FILE}",
85
+ "scale": LARRY_STRENGTH, # alpha == rank, so the base scale is 1
86
+ "entries": entries,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87
  }
 
 
88
 
89
 
90
+ def _load_lightx() -> dict:
91
+ from huggingface_hub import hf_hub_download
92
+ from safetensors.torch import load_file
93
+
94
+ lora = load_file(hf_hub_download(LIGHTX_REPO, LIGHTX_FILE))
95
+ suffix_a, suffix_b = ".lora_A.default.weight", ".lora_B.default.weight"
96
+ bases = sorted({key[: -len(suffix_a)] for key in lora if key.endswith(suffix_a)})
97
+ ranks = {lora[f"{name}{suffix_a}"].shape[0] for name in bases}
98
+ if len(ranks) != 1:
99
+ raise ValueError(f"Mixed LoRA ranks in {LIGHTX_FILE}: {sorted(ranks)}")
100
+ entries = [(f"{name}.weight", lora[f"{name}{suffix_a}"], lora[f"{name}{suffix_b}"]) for name in bases]
101
+ return {
102
+ "label": f"{LIGHTX_REPO}/{LIGHTX_FILE}",
103
+ "scale": LIGHTX_ALPHA / ranks.pop(),
104
+ "entries": entries,
105
+ }
106
+
107
+
108
+ def _apply(entries, params, sign: float) -> None:
109
+ for key, a, b in entries:
110
+ param = params.get(key)
111
+ if param is None:
112
+ raise KeyError(f"LoRA target `{key}` not found in the transformer")
113
+ delta = sign * (b.to(torch.float32) @ a.to(torch.float32))
114
+ param.data = (param.data.float() + delta.to(param.device)).to(param.dtype)
115
+
116
 
117
+ def available() -> list[str]:
118
+ """The LoRA sets that were loaded at startup, plus `off`."""
119
+ state = getattr(_PIPE_TRANSFORMER, "_lora_state", None) if _PIPE_TRANSFORMER is not None else None
120
+ return sorted(state["sets"]) + ["off"] if state else ["off"]
121
+
122
+
123
+ _PIPE_TRANSFORMER = None
124
+
125
+
126
+ def apply_lora(transformer) -> str | None:
127
+ """Load every enabled LoRA set, fold the default one into `transformer`, and stash the factors for per-request
128
+ switching. Returns a status line, or `None` when everything is disabled."""
129
+ global _PIPE_TRANSFORMER
130
+ _PIPE_TRANSFORMER = transformer
131
+
132
+ inner_dim = transformer.config.num_attention_heads * transformer.config.attention_head_dim
133
+ sets = {}
134
+ if LARRY_FILE.lower() not in ("", "off", "none"):
135
+ sets["larry"] = _load_larry(inner_dim)
136
+ if os.environ.get("H3_LIGHTX", "on").lower() not in ("", "off", "none"):
137
+ sets["lightx"] = _load_lightx()
138
+ if not sets:
139
+ return None
140
+
141
+ active = DEFAULT_LORA if DEFAULT_LORA in sets else sorted(sets)[0]
142
  params = dict(transformer.named_parameters())
143
+ _apply(sets[active]["entries"], params, sets[active]["scale"])
144
+ transformer._lora_state = {"active": active, "sets": sets}
145
+ return (
146
+ f"LoRAs loaded: "
147
+ + ", ".join(f"`{name}` ({spec['label']}, {len(spec['entries'])} weights)" for name, spec in sets.items())
148
+ + f" · active `{active}`"
149
+ )
150
+
151
+
152
+ def set_active(transformer, name: str) -> str:
153
+ """Switch the folded LoRA in place. No-op when the state already matches. Returns the active set."""
154
+ state = getattr(transformer, "_lora_state", None)
155
+ if state is None:
156
+ return "off"
157
+ name = name if name in state["sets"] else "off"
158
+ if state["active"] == name:
159
+ return name
160
+ params = dict(transformer.named_parameters())
161
+ if state["active"] != "off":
162
+ old = state["sets"][state["active"]]
163
+ _apply(old["entries"], params, -old["scale"])
164
+ if name != "off":
165
+ _apply(state["sets"][name]["entries"], params, state["sets"][name]["scale"])
166
+ state["active"] = name
167
+ return name
168
+
169
+
170
+ def set_enabled(transformer, enabled: bool) -> bool:
171
+ """Backwards-compatible boolean toggle over the default set."""
172
+ return set_active(transformer, DEFAULT_LORA if enabled else "off") != "off"
index.html CHANGED
@@ -225,7 +225,10 @@
225
  <div class="deck-label">Prompt</div>
226
  <textarea id="prompt" spellcheck="false">A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot</textarea>
227
  </div>
228
- <label class="check"><input type="checkbox" id="use_lora" checked> Turbo LoRA (fast steps)</label>
 
 
 
229
  <label class="check"><input type="checkbox" id="upsample"> Upsample prompt</label>
230
 
231
  <div>
@@ -307,6 +310,21 @@ fetch("/studio-config").then(r => r.json()).then(cfg => {
307
  sel.appendChild(o);
308
  }
309
  $("duration").min = cfg.min_duration; $("duration").max = cfg.max_duration;
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
310
  });
311
 
312
  async function pollStatus() {
@@ -405,7 +423,7 @@ $("run").addEventListener("click", async () => {
405
  steps: Number($("steps").value),
406
  seed: Number($("seed").value),
407
  upsample: $("upsample").checked,
408
- use_lora: $("use_lora").checked,
409
  });
410
  let data = null;
411
  for await (const msg of submission) {
 
225
  <div class="deck-label">Prompt</div>
226
  <textarea id="prompt" spellcheck="false">A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot</textarea>
227
  </div>
228
+ <div>
229
+ <div class="deck-label">Turbo LoRA</div>
230
+ <select id="lora"></select>
231
+ </div>
232
  <label class="check"><input type="checkbox" id="upsample"> Upsample prompt</label>
233
 
234
  <div>
 
310
  sel.appendChild(o);
311
  }
312
  $("duration").min = cfg.min_duration; $("duration").max = cfg.max_duration;
313
+
314
+ const loraSel = $("lora");
315
+ for (const [value, spec] of Object.entries(cfg.loras || { larry: { label: "larry", steps: 6 } })) {
316
+ const o = document.createElement("option");
317
+ o.value = value;
318
+ o.textContent = value === "off" ? "off (base model)" : value;
319
+ o.title = spec.label;
320
+ if (value === cfg.default_lora) o.selected = true;
321
+ loraSel.appendChild(o);
322
+ }
323
+ // Each LoRA has a design point: suggest it on switch (the slider stays free).
324
+ loraSel.addEventListener("change", () => {
325
+ const spec = (cfg.loras || {})[loraSel.value];
326
+ if (spec && spec.steps) { $("steps").value = spec.steps; $("steps-out").textContent = spec.steps; }
327
+ });
328
  });
329
 
330
  async function pollStatus() {
 
423
  steps: Number($("steps").value),
424
  seed: Number($("seed").value),
425
  upsample: $("upsample").checked,
426
+ lora: $("lora").value,
427
  });
428
  let data = null;
429
  for await (const msg of submission) {