Fold larryvrh/MiniMax-H3-Turbo-Lora v4-600-EMA instead of the lightx2v 8-step file

h3_turbo_lora.py

@@ -1,21 +1,31 @@
-"""The lightx2v 8-step Turbo LoRA, folded into the same bf16 weights H3-World is merged into.
-
-This mirrors the `h3_lora.py` of the Spaces that already run this adapter on MiniMax-H3 —
-[`MiniMaxAI/MiniMax-H3-Turbo-Lora`](https://huggingface.co/spaces/MiniMaxAI/MiniMax-H3-Turbo-Lora)
-(its `lightx` set) and
-[`hugging-apps/minimax-h3-turbo-sla-demo`](https://huggingface.co/spaces/hugging-apps/minimax-h3-turbo-sla-demo):
-
-* `lightx2v/Minimax-h3-Turbo` is a **PEFT checkpoint against the diffusers module tree itself** —
-  `transformer_blocks.N.attn.to_q.lora_A.default.weight` and friends — so unlike H3-World it needs
-  no key conversion at all: 312 targets (50 transformer blocks + 2 token-refiner blocks x
-  `to_q/to_k/to_v/to_out.0/ff.net.0.proj/ff.net.2`) map name-for-name.
-* Rank is 128 and the file's own safetensors metadata records `alpha: 8`, so the fold scale is
-  `alpha / rank = 0.0625` — exactly what `set_adapters(weights=1.0)` applies in lightx2v's
-  reference inference script, and what the Spaces above use for this repo.
-* `minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors` is an 8-NFE distillation, so the step count
-  is overridden to **8** when it is active. No scheduler swap and no CFG change: MiniMax-H3 is
-  already guidance-distilled and its own `MiniMaxH3Scheduler` is what every one of those Spaces
-  keeps using with the turbo LoRA folded.
+"""The `larryvrh/MiniMax-H3-Turbo-Lora` few-step LoRA, folded into the same bf16 weights H3-World is merged into.
+
+`minimax_h3_turbo_v4_step600_ema.safetensors` is the repo's recommended checkpoint (its `v4`
+line, step 600, EMA). Everything below comes out of the file itself and the model card:
+
+* It is a **ComfyUI-side checkpoint in the original MiniMax-H3 module naming** —
+  `blocks.N.attn.qkv_proj.lora_A.weight`, `token_refiner.blocks.N.mlp.fc1`,
+  `blocks.N.adaln_proj.linear`, `final_layer.adaln_proj.linear` — so, exactly like H3-World, the
+  keys have to be replayed through `convert_minimax_h3_to_diffusers.py`'s renames before the
+  deltas mean anything against the diffusers port: 259 LoRA pairs become 363 weight deltas.
+
+  One transform differs from `app.py`'s H3-World merge and it matters: the LoRA was trained
+  against `comfy.ldm.minimax.model`, whose `Attention.forward` reads
+  `self.qkv_proj(x).split(heads * head_dim, dim=-1)` — i.e. `Comfy-Org/MiniMax-H3` stores the
+  fused QKV as contiguous `[q_all; k_all; v_all]`, the reference's *in-memory* layout, not the raw
+  release shards' per-head interleave. So these rows are split into contiguous thirds and **not**
+  de-interleaved. The `mlp.fc1` `[gate; value]` -> `[value; gate]` swap diffusers' `SwiGLU` needs
+  still applies (comfy's swiglu also reads `[gate; value]`), and `adaln_proj.linear` /
+  `final_layer.adaln_proj.linear` are pure renames — the conversion script permutes no rows there.
+* The file's own safetensors metadata records `application: "W_eff = W + lora_B @ lora_A"` and the
+  card states *"applied as a plain low-rank update, alpha = rank, so no extra scaling"* — hence a
+  fold scale of exactly **1.0**, with `H3_TURBO_STRENGTH` left as the card's strength dial (nudge
+  to ~1.05-1.2 for smear, ~0.8-0.95 for over-sharp grain). Ranks are mixed on purpose: 64 for the
+  attention/FFN projections, 16 for the AdaLN projections.
+* The card's useful step range is **4-8**, with 6-8 recommended and no benefit past 8, so the step
+  count is overridden to **8 NFE** when the adapter is active. No scheduler swap and no CFG
+  change: MiniMax-H3 is already guidance-distilled and its own `MiniMaxH3Scheduler` is what this
+  Space keeps using with the turbo LoRA folded.
 
 The adapter is *folded* into the weights rather than wrapped as a runtime module, for the same
 reason H3-World is: this Space patches `MiniMaxH3AttnProcessor` and drives the transformer's live
@@ -29,29 +39,59 @@ import os
 
 import torch
 
-TURBO_REPO = os.environ.get("H3_TURBO_REPO", "lightx2v/Minimax-h3-Turbo")
-TURBO_FILE = os.environ.get("H3_TURBO_FILE", "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors")
-# The distillation's own NFE count; the card offers 4 as the faster, softer alternative.
+TURBO_REPO = os.environ.get("H3_TURBO_REPO", "larryvrh/MiniMax-H3-Turbo-Lora")
+TURBO_FILE = os.environ.get("H3_TURBO_FILE", "minimax_h3_turbo_v4_step600_ema.safetensors")
+# The card's step range is 4-8 (6-8 recommended, nothing gained past 8).
 TURBO_STEPS = int(os.environ.get("H3_TURBO_STEPS", "8"))
-# 0 = read `alpha` out of the file's safetensors metadata (it is `8` for every lightx2v file).
-TURBO_ALPHA = float(os.environ.get("H3_TURBO_ALPHA", "0"))
+# alpha == rank in this checkpoint, so the update is applied as-is; this is the card's strength dial.
 TURBO_STRENGTH = float(os.environ.get("H3_TURBO_STRENGTH", "1.0"))
 
-SUFFIX_A, SUFFIX_B = ".lora_A.default.weight", ".lora_B.default.weight"
+SUFFIX_A, SUFFIX_B = ".lora_A.weight", ".lora_B.weight"
 
 
-def load() -> dict:
-    """Download the adapter and return ``{"label", "scale", "entries"}``; entries are ``(param_key, A, B)``."""
+def _target_name(source_name: str) -> str:
+    """Rename one original-layout module path onto its diffusers module path."""
+    if source_name.startswith("token_refiner.blocks."):
+        target = source_name.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
+    elif source_name.startswith("blocks."):
+        target = source_name.replace("blocks.", "transformer_blocks.", 1)
+    else:
+        target = source_name
+    return target.replace("final_layer.adaln_proj.linear", "norm_out.linear")
+
+
+def _targets(name: str, b_weight, inner_dim: int):
+    """Yield ``(diffusers_param_key, row_transformed_B)`` for one original LoRA base name."""
+    target = _target_name(name)
+
+    if target.endswith(".attn.qkv_proj"):
+        prefix = target.removesuffix("qkv_proj")
+        if b_weight.shape[0] != 3 * inner_dim:
+            raise ValueError(
+                f"{name} lora_B has {b_weight.shape[0]} rows, expected 3 x {inner_dim} for a fused QKV"
+            )
+        # Contiguous thirds — the comfy layout this adapter was trained against already holds
+        # `[q_all; k_all; v_all]`, so there is no per-head de-interleave here (see the module docstring).
+        for kind, part in zip(("q", "k", "v"), b_weight.split(inner_dim, dim=0)):
+            yield f"{prefix}to_{kind}.weight", part.contiguous()
+    elif target.endswith(".mlp.fc1"):
+        # SwiGLU gate/value swap: the checkpoint stores [gate, value]; diffusers stores [value, gate].
+        gate, value = b_weight.chunk(2, dim=0)
+        yield target.replace(".mlp.fc1", ".ff.net.0.proj") + ".weight", torch.cat([value, gate]).contiguous()
+    elif target.endswith(".mlp.fc2"):
+        yield target.replace(".mlp.fc2", ".ff.net.2") + ".weight", b_weight
+    elif target.endswith(".attn.out_proj"):
+        yield target.replace(".attn.out_proj", ".attn.to_out.0") + ".weight", b_weight
+    else:  # `adaln_proj.linear`, `norm_out.linear` — pure renames
+        yield target + ".weight", b_weight
+
+
+def load(inner_dim: int) -> dict:
+    """Download the adapter and return ``{"label", "scale", "entries", ...}``; entries are ``(param_key, A, B)``."""
     from huggingface_hub import hf_hub_download
-    from safetensors import safe_open
     from safetensors.torch import load_file
 
-    path = hf_hub_download(TURBO_REPO, TURBO_FILE)
-    with safe_open(path, framework="pt") as handle:
-        metadata = handle.metadata() or {}
-    alpha = TURBO_ALPHA or float(metadata.get("alpha", 8))
-
-    lora = load_file(path)
+    lora = load_file(hf_hub_download(TURBO_REPO, TURBO_FILE))
     bases = sorted({key[: -len(SUFFIX_A)] for key in lora if key.endswith(SUFFIX_A)})
     if not bases:
         raise ValueError(f"No lora_A/lora_B pairs found in {TURBO_FILE}")
@@ -62,25 +102,29 @@ def load() -> dict:
         if f"{name}{SUFFIX_B}" not in lora:
             raise ValueError(f"LoRA is missing the lora_B twin of {name}{SUFFIX_A}")
 
-    ranks = {lora[f"{name}{SUFFIX_A}"].shape[0] for name in bases}
-    if len(ranks) != 1:
-        raise ValueError(f"LoRA mixes ranks {sorted(ranks)}")
-    rank = ranks.pop()
+    # Ranks are mixed by design (64 on attention/FFN, 16 on AdaLN) and alpha == rank for every one of
+    # them, so the fold scale is rank-independent.
+    ranks = sorted({lora[f"{name}{SUFFIX_A}"].shape[0] for name in bases})
+
+    entries = []
+    for name in bases:
+        a = lora[f"{name}{SUFFIX_A}"]
+        for key, b_part in _targets(name, lora[f"{name}{SUFFIX_B}"], inner_dim):
+            entries.append((key, a, b_part))
 
-    entries = [(f"{name}.weight", lora[f"{name}{SUFFIX_A}"], lora[f"{name}{SUFFIX_B}"]) for name in bases]
     return {
         "label": f"{TURBO_REPO}/{TURBO_FILE}",
-        "scale": alpha / rank * TURBO_STRENGTH,
+        "scale": TURBO_STRENGTH,  # alpha == rank -> the update is applied as-is
         "entries": entries,
-        "rank": rank,
-        "alpha": alpha,
+        "targets": len(bases),
+        "ranks": ranks,
     }
 
 
 def _fold(transformer, spec: dict, sign: float) -> None:
     """Add ``sign * scale * (B @ A)`` to every target weight, in place.
 
-    The factors ride to the weight's own device before the matmul — the deltas are 5 TFLOP in
+    The factors ride to the weight's own device before the matmul — the deltas are a few TFLOP in
     total, which is seconds on the card and minutes on a Space's two vCPUs.
     """
     params = dict(transformer.named_parameters())
@@ -100,7 +144,7 @@ def _fold(transformer, spec: dict, sign: float) -> None:
 
 def prepare(transformer) -> str:
     """Load the adapter and validate it against ``transformer``, without folding it yet."""
-    spec = load()
+    spec = load(transformer.config.num_attention_heads * transformer.config.attention_head_dim)
     params = dict(transformer.named_parameters())
     missed = [key for key, _, _ in spec["entries"] if key not in params]
     if missed:
@@ -116,9 +160,9 @@ def prepare(transformer) -> str:
             )
     transformer._turbo_state = {"active": False, "spec": spec}
     return (
-        f"Turbo LoRA ready · {spec['label']} · {len(spec['entries'])} targets, "
-        f"rank {spec['rank']}, alpha {spec['alpha']:g} -> scale {spec['scale']:.4f} · "
-        f"{TURBO_STEPS} steps when active"
+        f"Turbo LoRA ready · {spec['label']} · {spec['targets']} targets -> "
+        f"{len(spec['entries'])} weight deltas, rank {'/'.join(str(r) for r in spec['ranks'])}, "
+        f"scale {spec['scale']:.4g} · {TURBO_STEPS} steps when active"
     )