Update app.py

app.py

@@ -40,6 +40,12 @@ LORA_FILES = {
     "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"),
 }
+# Display names, keyed the same as LORA_FILES — used in the UI slider labels, the per-request report line, and
+# the status line's failure list. Keep these two dicts' keys in sync when adding a LoRA.
+LORA_LABELS = {
+    "loraa": "Anthro Enhancer",
+    "lorab": "Reasoning Enhancer",
+}
 DEFAULT_LORA_A_STRENGTH = 0.0
 DEFAULT_LORA_B_STRENGTH = 0.0
 # Some `diffusion_model.blocks.*` checkpoints store SwiGLU's fc1 gate/value halves in the opposite order
@@ -172,6 +178,7 @@ FILM_ERROR: str | None = None
 LOAD_ERROR: str | None = None
 LOADED_IN: float | None = None
 LORA_STATUS: str | None = None
+LOADED_LORAS: set[str] = set()
 
 
 def status() -> str:
@@ -260,12 +267,12 @@ def load_models() -> str | None:
             from peft.tuners.tuners_utils import BaseTunerLayer
             from safetensors import safe_open
 
-            try:
-                # Snapshot once, before either adapter attaches and wraps the target Linears — see the
-                # docstring on `_convert_diffusion_model_lora` for why this can't be read fresh per-file.
-                base_shapes = {k: tuple(v.shape) for k, v in pipe.transformer.state_dict().items()}
-                counts = {}
-                for name, filename in LORA_FILES.items():
+            # Snapshot once, before any adapter attaches and wraps the target Linears — see the docstring on
+            # `_convert_diffusion_model_lora` for why this can't be read fresh per-file.
+            base_shapes = {k: tuple(v.shape) for k, v in pipe.transformer.state_dict().items()}
+            failures = []
+            for name, filename in LORA_FILES.items():
+                try:
                     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()}
@@ -273,24 +280,25 @@ def load_models() -> str | None:
                     pipe.transformer.load_lora_adapter(converted, adapter_name=name, prefix=None)
                     # `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.
-                    counts[name] = sum(
+                    matched = sum(
                         1
                         for module in pipe.transformer.modules()
                         if isinstance(module, BaseTunerLayer) and name in module.lora_A
                     )
-                    if counts[name] == 0:
-                        raise RuntimeError(
-                            f"'{filename}' converted but matched 0 target modules on "
-                            f"MiniMaxH3Transformer3DModel — the rename map needs another look."
-                        )
-                pipe.transformer.set_adapters(list(LORA_FILES), weights=[0.0] * len(LORA_FILES))
-                LORA_STATUS = "LoRAs loaded: " + ", ".join(
-                    f"`{name}` ({LORA_REPO}/{filename}, {counts[name]} layers)"
-                    for name, filename in LORA_FILES.items()
-                )
-            except Exception as error:
-                LORA_STATUS = f"LoRA load failed: {type(error).__name__}: {error}"
-                print(f"[gen] {LORA_STATUS}", flush=True)
+                    if matched == 0:
+                        raise RuntimeError(f"'{filename}' converted but matched 0 target modules")
+                    LOADED_LORAS.add(name)
+                except Exception as error:
+                    failures.append(f"`{LORA_LABELS.get(name, name)}` ({type(error).__name__}: {error})")
+                    print(
+                        f"[gen] LoRA '{name}' ({filename}) failed to load: {type(error).__name__}: {error}",
+                        flush=True,
+                    )
+
+            if LOADED_LORAS:
+                pipe.transformer.set_adapters(list(LOADED_LORAS), weights=[0.0] * len(LOADED_LORAS))
+            LORA_STATUS = "All LoRAs loaded" if not failures else "LoRA issues: " + "; ".join(failures)
+            print(f"[gen] {LORA_STATUS}", flush=True)
 
         # Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU
         # worker.
@@ -425,8 +433,7 @@ def _generate(
     sharpen,
     multiplier,
     seed,
-    lora_a_strength,
-    lora_b_strength,
+    lora_strengths,
 ):
     """The only thing on GPU time: the denoise loop, the two decoders and the workflow's post chain.
     The mp4 is muxed here rather than in the caller: a `@spaces.GPU` return crosses a process boundary by pickling,
@@ -440,11 +447,14 @@ def _generate(
 
     booked = time.time()
 
-    # Approach B: blend the two resident LoRA adapters for this request. Cheap — `set_adapters` only updates each
-    # PEFT layer's active-adapter list and scale, no weight math — so it's safe to call on every request. Guarded
-    # in case loading failed or was disabled at startup, when the transformer carries no adapters at all.
-    if hasattr(PIPE.transformer, "peft_config"):
-        PIPE.transformer.set_adapters(list(LORA_FILES), weights=[float(lora_a_strength), float(lora_b_strength)])
+# Approach B: blend whichever resident LoRA adapters actually loaded, for this request. Cheap —
+    # `set_adapters` only updates each PEFT layer's active-adapter list and scale, no weight math — so it's safe
+    # to call on every request. Filtered to `LOADED_LORAS`: a slider for a LoRA that failed at startup has no
+    # adapter behind it, and `set_adapters` would raise if asked to activate a name that was never attached.
+    if LOADED_LORAS:
+        active = {name: strength for name, strength in lora_strengths.items() if name in LOADED_LORAS}
+        if active:
+            PIPE.transformer.set_adapters(list(active), weights=list(active.values()))
 
     if PLACEMENT == "lazy":
         PIPE.to("cuda")
@@ -558,6 +568,10 @@ def generate(
         # prepares exactly this way.
         return ImageOps.exif_transpose(Image.open(path)).convert("RGB") if path else None
 
+    # Every UI LoRA slider gets packed into one dict here — this is the only place a new LoRA's slider value
+    # needs wiring in; `_generate`, `set_adapters`, and the report line below are all keyed off this dict.
+    lora_strengths = {"loraa": float(lora_a_strength), "lorab": float(lora_b_strength)}
+
     progress(0.1, desc=f"Denoising {int(steps)} steps at {width}x{height}, {num_frames} frames ...")
     call = (
         prompt_embeds,
@@ -572,8 +586,7 @@ def generate(
         float(sharpen),
         multiplier,
         int(seed),
-        float(lora_a_strength),
-        float(lora_b_strength),
+        lora_strengths,
     )
     # The same call `spaces` will book the worker with, so the report can show the fit against the measurement.
     booked_seconds = get_duration(*call)
@@ -587,10 +600,12 @@ def generate(
 
     post = [f"RCAS {float(sharpen):.2f}" if float(sharpen) > 0 else "no sharpening"]
     post.append(f"FILM {multiplier}x -> {fps} fps" if multiplier > 1 else f"{fps} fps")
+    lora_text = " / ".join(
+        f"{LORA_LABELS.get(name, name)} {strength:.2f}" for name, strength in lora_strengths.items()
+    )
     report = (
         f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s) -> {out_frames} frames at {fps} fps · "
-        f"{int(steps)} steps of `{schedule_key}` · {' · '.join(post)} · seed {int(seed)} · "
-        f"LoRA A {float(lora_a_strength):.2f} / B {float(lora_b_strength):.2f}\n\n"
+        f"{int(steps)} steps of `{schedule_key}` · {' · '.join(post)} · seed {int(seed)} · {lora_text}\n\n"
         f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
         f"{', upsampled' if refined else ''}) · denoise + decode {denoise_seconds:.0f}s "
         f"({denoise_seconds / max(1, int(steps)):.1f} s/step) · post {post_seconds:.0f}s · "