Update app.py

app.py

@@ -37,7 +37,7 @@ GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
 
 LORA_REPO = os.environ.get("H3_LORA_REPO", "dagloop5/LoRA")
 LORA_FILES = {
-    "lora1": os.environ.get("H3_LORA_A_FILE", "minimax_h3_turbo_v4_step600_ema.safetensors"),
+    "lora1": os.environ.get("H3_LORA_1_FILE", "minimax_h3_turbo_v4_step600_ema.safetensors"),
     "loraa": os.environ.get("H3_LORA_A_FILE", "Mylo_lora_epoch31.safetensors"),
     "lorab": os.environ.get("H3_LORA_B_FILE", "VBVR_H3_attn_only.safetensors"),
     "lorac": os.environ.get("H3_LORA_C_FILE", "AIO_V2.safetensors"),
@@ -144,19 +144,27 @@ def _convert_diffusion_model_lora(raw: dict, base_shapes: dict) -> dict:
     import re
 
     out = {}
-    pattern = re.compile(r"^diffusion_model\.blocks\.(\d+)\.(attn|mlp)\.(\w+)\.(lora_[AB])\.weight$")
 
-    for key, raw_tensor in raw.items():
-        match = pattern.match(key)
-        if not match:
-            print(f"[lora-convert] skipping unrecognized key: {key}", flush=True)
-            continue
-        # 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)
-        block, kind, leaf, ab = match.groups()
+    # 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
+    # Enhancer. `lora_down`/`lora_up` are the same A/B convention under a different name.
+    kohya = re.compile(
+        r"^lora_unet_blocks_(\d+)_(attn_out_proj|attn_qkv_proj|mlp_fc1|mlp_fc2)\.(lora_down|lora_up)\.weight$"
+    )
+    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":
             if ab == "lora_A":
                 # Shared low-rank input side — identical for q, k, v.
@@ -168,28 +176,55 @@ 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"{key}: expected {q_out + k_out + v_out} rows (q{q_out}+k{k_out}+v{v_out}), "
-                    f"got {tensor.shape[0]}"
+                    f"blocks.{block}.attn.qkv_proj.{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:
-                # Some checkpoints store the SwiGLU gate/value halves in the opposite order diffusers expects.
                 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 {key}, skipping", flush=True)
+            print(f"[lora-convert] no mapping for blocks.{block}.{kind}.{leaf}.{ab}, skipping", flush=True)
+
+    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)
+        if match:
+            block, kind, leaf, ab = match.groups()
+            emit(block, kind, leaf, ab, tensor)
+            continue
+
+        match = kohya.match(key)
+        if match:
+            block, target, direction = match.groups()
+            kind, leaf = kohya_targets[target]
+            emit(block, kind, leaf, kohya_ab[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.
+            continue
+
+        print(f"[lora-convert] skipping unrecognized key: {key}", flush=True)
 
     return out