Update app.py

app.py

@@ -67,10 +67,14 @@ DEFAULT_LORA_D_STRENGTH = 0.0
 DEFAULT_LORA_E_STRENGTH = 0.0
 DEFAULT_LORA_F_STRENGTH = 0.0
 DEFAULT_LORA_G_STRENGTH = 0.0
-# Some `diffusion_model.blocks.*` checkpoints store SwiGLU's fc1 gate/value halves in the opposite order
-# diffusers expects. Leave off first; if the LoRA's effect looks inverted/broken rather than just weak or
-# strong, set H3_LORA_SWAP_FC1=1 and compare.
-SWAP_FC1_HALVES = os.environ.get("H3_LORA_SWAP_FC1", "0") == "1"
+# Per-LoRA, not global: different training pipelines can store SwiGLU's fc1 gate/value halves in either order,
+# and one flag can only be right for however many of the 8 files happen to agree. `lora1` (the Distilled/Turbo
+# LoRA) is confirmed needing the swap by InstantX's official conversion of the same lineage
+# (MiniMax-H3-Turbo-Lora-Diffusers/convert.py: "SwiGLU fc1 halves are swapped to match Diffusers' [value; gate]
+# layout"); the rest default off until tested individually — set H3_LORA_SWAP_FC1_NAMES to a comma-separated
+# list of LORA_FILES keys (e.g. "lora1,lorac") to override. Replaces H3_LORA_SWAP_FC1, which no longer does
+# anything.
+SWAP_FC1_NAMES = {name for name in os.environ.get("H3_LORA_SWAP_FC1_NAMES", "lora1").split(",") if name}
 
 # Must stay identical to the conditioner's table: the *label* goes over the wire, so a canvas that half does not
 # know is rejected there and surfaces as a failure here. This is the workflow's "Target Dimension" node.
@@ -155,39 +159,71 @@ def lower_duration_floor(seconds: float = MIN_UI_DURATION) -> None:
 
     MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(seconds))
 
-def _convert_diffusion_model_lora(raw: dict, base_shapes: dict) -> dict:
+def _convert_diffusion_model_lora(raw: dict, base_shapes: dict, swap_fc1: bool) -> tuple[dict, dict]:
     """Rename a `diffusion_model.blocks.*` (original-checkpoint) LoRA state dict onto
     `MiniMaxH3Transformer3DModel`'s (`transformer_blocks.*`) naming, so `load_lora_adapter` can attach it.
     `raw` maps original key -> tensor. `base_shapes` maps the *unwrapped* base model's parameter names to their
     shapes — captured once before any adapter is attached, since `load_lora_adapter` wraps each target Linear in
     a PEFT layer and renames its weight to `<name>.base_layer.weight`, so a live `transformer.state_dict()` call
     after the first adapter attaches would no longer have `to_q.weight` etc. under their original names.
+
+    Returns `(converted_weights, network_alphas)` — `network_alphas` is `load_lora_adapter`'s per-module `alpha`
+    map. Built for every converted module, not just ones whose raw file carries an explicit `.alpha` key: PEFT's
+    default scaling isn't guaranteed to land on `alpha == rank` when a LoRA mixes ranks across target types —
+    InstantX's own Turbo-LoRA conversion needs `network_alphas` for exactly this reason (attn/mlp modules rank
+    64, AdaLN modules rank 16), even though that file carries no `.alpha` keys at all. So `alpha = rank` is
+    synthesized for every module first, then overridden wherever the raw file specifies something else.
     """
     import re
 
-    out = {}
+    out: dict = {}
+    raw_alphas: dict[str, float] = {}  # raw ComfyUI base name -> alpha, from real `.alpha` keys only
 
-    # Family A: `[diffusion_model.]blocks.N.(attn|mlp|adaln_proj).LEAF.(lora_A|lora_B).weight` — covers Mylo,
-    # VBVR, AIO_V2, moawxx, the Furry Realism LoRA, and (minus its `diffusion_model.` prefix) the Turbo LoRA.
-    standard = re.compile(
-        r"^(?:diffusion_model\.)?blocks\.(\d+)\.(attn|mlp|adaln_proj)\.([\w.]+)\.(lora_[AB])\.weight$"
-    )
-    # Family B (Kohya-style): `lora_unet_blocks_N_TARGET.(lora_down|lora_up).weight` — covers SB and Fluid
+    # Family A: standard (non-Kohya) naming — covers Mylo, VBVR, AIO_V2, moawxx, Anthro Realism, and (once
+    # `diffusion_model.` is stripped) the Distilled/Turbo LoRA. Matched by module base name rather than one
+    # fixed pattern, and renamed by substitution — this is what lets `token_refiner.blocks.*` and the top-level
+    # `final_layer.adaln_proj` resolve onto real targets instead of falling through unmatched. Ported from
+    # InstantX's official `MiniMax-H3-Turbo-Lora-Diffusers/convert.py`, written for this exact LoRA family.
+    standard_ab = re.compile(r"^(?:diffusion_model\.)?(.+)\.(lora_[AB])\.weight$")
+    standard_alpha = re.compile(r"^(?:diffusion_model\.)?(.+)\.alpha$")
+
+    # Family B (Kohya-style): `lora_unet_blocks_N_TARGET.(lora_down|lora_up|alpha)` — covers SB and Fluid
     # Enhancer. `lora_down`/`lora_up` are the same A/B convention under a different name.
-    kohya = re.compile(
+    kohya_ab = re.compile(
         r"^lora_unet_blocks_(\d+)_(attn_out_proj|attn_qkv_proj|mlp_fc1|mlp_fc2)\.(lora_down|lora_up)\.weight$"
     )
+    kohya_alpha = re.compile(r"^lora_unet_blocks_(\d+)_(attn_out_proj|attn_qkv_proj|mlp_fc1|mlp_fc2)\.alpha$")
     kohya_targets = {
         "attn_out_proj": ("attn", "out_proj"),
         "attn_qkv_proj": ("attn", "qkv_proj"),
         "mlp_fc1": ("mlp", "fc1"),
         "mlp_fc2": ("mlp", "fc2"),
     }
-    kohya_ab = {"lora_down": "lora_A", "lora_up": "lora_B"}
-
-    def emit(block: str, kind: str, leaf: str, ab: str, tensor) -> None:
-        prefix = f"transformer_blocks.{block}."
-        if kind == "attn" and leaf == "qkv_proj":
+    kohya_ab_name = {"lora_down": "lora_A", "lora_up": "lora_B"}
+
+    def rename_base(name: str) -> str:
+        """ComfyUI module path (before `.lora_*`/`.alpha`) -> Diffusers module path."""
+        if name.startswith("token_refiner.blocks."):
+            name = "token_refiner.refiner_blocks." + name[len("token_refiner.blocks."):]
+        elif name.startswith("blocks."):
+            name = "transformer_blocks." + name[len("blocks."):]
+        name = name.replace("final_layer.adaln_proj.linear", "norm_out.linear")
+        name = name.replace(".attn.out_proj", ".attn.to_out.0")
+        name = name.replace(".mlp.fc2", ".ff.net.2")
+        name = name.replace(".mlp.fc1", ".ff.net.0.proj")
+        return name
+
+    def target_bases(raw_base: str) -> list[str]:
+        """Diffusers-side base name(s) for one pre-rename module path — one, except `attn.qkv_proj`, which fans
+        out to `to_q`/`to_k`/`to_v` (same rank, so the same alpha applies to all three)."""
+        if raw_base.endswith(".attn.qkv_proj"):
+            prefix = rename_base(raw_base[: -len("attn.qkv_proj")])
+            return [f"{prefix}attn.to_q", f"{prefix}attn.to_k", f"{prefix}attn.to_v"]
+        return [rename_base(raw_base)]
+
+    def emit(raw_base: str, ab: str, tensor) -> None:
+        if raw_base.endswith(".attn.qkv_proj"):
+            prefix = rename_base(raw_base[: -len("attn.qkv_proj")])
             if ab == "lora_A":
                 # Shared low-rank input side — identical for q, k, v.
                 out[f"{prefix}attn.to_q.{ab}.weight"] = tensor
@@ -198,57 +234,70 @@ def _convert_diffusion_model_lora(raw: dict, base_shapes: dict) -> dict:
                 k_out = base_shapes[f"{prefix}attn.to_k.weight"][0]
                 v_out = base_shapes[f"{prefix}attn.to_v.weight"][0]
                 assert tensor.shape[0] == q_out + k_out + v_out, (
-                    f"blocks.{block}.attn.qkv_proj.{ab}: expected {q_out + k_out + v_out} rows "
+                    f"{raw_base}.{ab}: expected {q_out + k_out + v_out} rows "
                     f"(q{q_out}+k{k_out}+v{v_out}), got {tensor.shape[0]}"
                 )
                 out[f"{prefix}attn.to_q.{ab}.weight"] = tensor[:q_out].clone()
                 out[f"{prefix}attn.to_k.{ab}.weight"] = tensor[q_out:q_out + k_out].clone()
                 out[f"{prefix}attn.to_v.{ab}.weight"] = tensor[q_out + k_out:].clone()
-        elif kind == "attn" and leaf == "out_proj":
-            out[f"{prefix}attn.to_out.0.{ab}.weight"] = tensor
-        elif kind == "mlp" and leaf == "fc1":
-            if ab == "lora_B" and SWAP_FC1_HALVES:
-                half = tensor.shape[0] // 2
-                tensor = torch.cat([tensor[half:], tensor[:half]], dim=0)
-            out[f"{prefix}ff.net.0.proj.{ab}.weight"] = tensor
-        elif kind == "mlp" and leaf == "fc2":
-            out[f"{prefix}ff.net.2.{ab}.weight"] = tensor
-        elif kind == "adaln_proj" and leaf == "linear":
-            out[f"{prefix}adaln_proj.linear.{ab}.weight"] = tensor
-        else:
-            print(f"[lora-convert] no mapping for blocks.{block}.{kind}.{leaf}.{ab}, skipping", flush=True)
+            return
+
+        if raw_base.endswith(".mlp.fc1") and ab == "lora_B" and swap_fc1:
+            half = tensor.shape[0] // 2
+            tensor = torch.cat([tensor[half:], tensor[:half]], dim=0)
+
+        key_base = rename_base(raw_base)
+        if f"{key_base}.weight" not in base_shapes:
+            # The more permissive substitution-based rename can produce a name that isn't an actual target on
+            # the live model — validated here rather than trusting the rename blindly, since it no longer
+            # checks against a fixed whitelist of known `kind`s the way the old anchored regex did.
+            print(f"[lora-convert] '{raw_base}' renamed to '{key_base}', which isn't a real target — skipping", flush=True)
+            return
+        out[f"{key_base}.{ab}.weight"] = tensor
 
     for key, raw_tensor in raw.items():
         # Some files (fp16-labeled ones especially) don't match the bf16 transformer's dtype; PEFT expects the
         # adapter's dtype to match the wrapped base layer's.
         tensor = raw_tensor.to(torch.bfloat16)
 
-        match = standard.match(key)
+        match = standard_ab.match(key)
         if match:
-            block, kind, leaf, ab = match.groups()
-            emit(block, kind, leaf, ab, tensor)
+            raw_base, ab = match.groups()
+            emit(raw_base, ab, tensor)
             continue
 
-        match = kohya.match(key)
+        match = kohya_ab.match(key)
         if match:
             block, target, direction = match.groups()
             kind, leaf = kohya_targets[target]
-            emit(block, kind, leaf, kohya_ab[direction], tensor)
+            emit(f"blocks.{block}.{kind}.{leaf}", kohya_ab_name[direction], tensor)
             continue
 
-        if key.endswith(".alpha"):
-            # Per-module rank/alpha scaling isn't threaded through — matched modules get PEFT's default scaling
-            # (scale 1.0), and the UI slider is what actually controls each LoRA's visible strength here. This
-            # is a known simplification: a file's built-in alpha may have scaled it up or down from its raw
-            # rank, so its slider range that "feels right" may not match what the file's author intended or
-            # tested at. It isn't a bug — the Furry Realism LoRA's `.alpha` keys were already dropped the same
-            # way and it loads and works fine — just worth knowing if a LoRA's effect seems unexpectedly
-            # weak/strong across its whole slider range rather than at a specific value.
+        match = standard_alpha.match(key)
+        if match:
+            (raw_base,) = match.groups()
+            raw_alphas[raw_base] = float(raw_tensor)
+            continue
+
+        match = kohya_alpha.match(key)
+        if match:
+            block, target = match.groups()
+            kind, leaf = kohya_targets[target]
+            raw_alphas[f"blocks.{block}.{kind}.{leaf}"] = float(raw_tensor)
             continue
 
         print(f"[lora-convert] skipping unrecognized key: {key}", flush=True)
 
-    return out
+    network_alphas: dict[str, float] = {}
+    for out_key, out_tensor in out.items():
+        if out_key.endswith(".lora_B.weight"):
+            base = out_key[: -len(".lora_B.weight")]
+            network_alphas[f"{base}.alpha"] = float(out_tensor.shape[1])
+    for raw_base, alpha in raw_alphas.items():
+        for base in target_bases(raw_base):
+            network_alphas[f"{base}.alpha"] = alpha
+
+    return out, network_alphas
 
 PIPE = None
 FILM = None
@@ -354,8 +403,12 @@ def load_models() -> str | None:
                     path = hf_hub_download(LORA_REPO, filename)
                     with safe_open(path, framework="pt") as handle:
                         raw = {k: handle.get_tensor(k) for k in handle.keys()}
-                    converted = _convert_diffusion_model_lora(raw, base_shapes)
-                    pipe.transformer.load_lora_adapter(converted, adapter_name=name, prefix=None)
+                    converted, network_alphas = _convert_diffusion_model_lora(
+                        raw, base_shapes, swap_fc1=name in SWAP_FC1_NAMES
+                    )
+                    pipe.transformer.load_lora_adapter(
+                        converted, adapter_name=name, prefix=None, network_alphas=network_alphas
+                    )
                     # `load_lora_adapter` warns-and-continues on a zero-key match instead of raising, so count
                     # matched layers ourselves and fail loudly if a file attached nothing.
                     matched = sum(