Update app.py

app.py

@@ -33,6 +33,14 @@ PLACEMENT = os.environ.get("H3_PLACEMENT", "pack").lower()
 ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
 GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
 
+LORA_REPO = os.environ.get("H3_LORA_REPO", "dagloop5/LoRA")
+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"),
+}
+DEFAULT_LORA_A_STRENGTH = 0.0
+DEFAULT_LORA_B_STRENGTH = 0.0
+
 # 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.
 CANVASES = {
@@ -101,6 +109,7 @@ FILM = None
 FILM_ERROR: str | None = None
 LOAD_ERROR: str | None = None
 LOADED_IN: float | None = None
+LORA_STATUS: str | None = None
 
 
 def status() -> str:
@@ -113,8 +122,8 @@ def status() -> str:
     film = "FILM **ready**" if FILM is not None else f"FILM **off** ({FILM_ERROR})"
     return (
         f"Ready · transformer + VAEs **bfloat16, unquantized** · placement `{PLACEMENT}` · attention "
-        f"`{ATTENTION}` · {h3_aoti.status()} · {film} · loaded in {LOADED_IN:.0f}s · conditioner "
-        f"`{CONDITIONER_SPACE}`"
+        f"`{ATTENTION}` · {h3_aoti.status()} · {film} · {LORA_STATUS or 'no LoRA'} · loaded in {LOADED_IN:.0f}s · "
+        f"conditioner `{CONDITIONER_SPACE}`"
     )
 
 
@@ -125,7 +134,7 @@ def load_models() -> str | None:
     touched. Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a bfloat16 audio
     VAE decodes the soundtrack roughly 20 dB too quiet.
     """
-    global PIPE, FILM, FILM_ERROR, LOAD_ERROR, LOADED_IN
+    global PIPE, FILM, FILM_ERROR, LOAD_ERROR, LOADED_IN, LORA_STATUS
 
     if PIPE is not None or LOAD_ERROR is not None:
         return LOAD_ERROR
@@ -145,6 +154,21 @@ def load_models() -> str | None:
         pipe.load_components(dtype=torch.bfloat16)
         pipe.transformer.set_attention_backend(ATTENTION)
 
+        # Approach B: attach both LoRA adapters as PEFT layers on the transformer, inactive (weight 0) until a
+        # request asks for them. `load_lora_adapter` is the model-level loader (`PeftAdapterMixin`), used because
+        # `MiniMaxH3ModularPipeline` has no pipeline-level `load_lora_weights` of its own.
+        if LORA_REPO.lower() not in ("", "off", "none"):
+            try:
+                for name, filename in LORA_FILES.items():
+                    pipe.transformer.load_lora_adapter(LORA_REPO, weight_name=filename, adapter_name=name)
+                pipe.transformer.set_adapters(list(LORA_FILES), weights=[0.0] * len(LORA_FILES))
+                LORA_STATUS = "LoRAs loaded: " + ", ".join(
+                    f"`{name}` ({LORA_REPO}/{filename})" 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)
+
         # Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU
         # worker.
         import h3_aoti
@@ -278,6 +302,8 @@ def _generate(
     sharpen,
     multiplier,
     seed,
+    lora_a_strength,
+    lora_b_strength,
 ):
     """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,
@@ -290,6 +316,13 @@ def _generate(
     global FILM
 
     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)])
+
     if PLACEMENT == "lazy":
         PIPE.to("cuda")
     elif PLACEMENT == "pack":
@@ -361,6 +394,8 @@ def generate(
     interpolation=DEFAULT_INTERPOLATION,
     seed=42,
     upsample=False,
+    lora_a_strength=DEFAULT_LORA_A_STRENGTH,
+    lora_b_strength=DEFAULT_LORA_B_STRENGTH,
     progress=gr.Progress(track_tqdm=True),
 ):
     """One request through the PlagueKind graph. Every parameter but the prompt carries the default its UI
@@ -414,6 +449,8 @@ def generate(
         float(sharpen),
         multiplier,
         int(seed),
+        float(lora_a_strength),
+        float(lora_b_strength),
     )
     # The same call `spaces` will book the worker with, so the report can show the fit against the measurement.
     booked_seconds = get_duration(*call)
@@ -429,7 +466,8 @@ def generate(
     post.append(f"FILM {multiplier}x -> {fps} fps" if multiplier > 1 else f"{fps} fps")
     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)}\n\n"
+        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"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 · "
@@ -537,6 +575,20 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
                     value=DEFAULT_STEPS,
                     info="PlagueKind: 15-20 on the linear_quadratic grid.",
                 )
+                lora_a_strength = gr.Slider(
+                    label="Anthro Enhancer LoRA",
+                    minimum=0.0,
+                    maximum=2.0,
+                    step=0.05,
+                    value=DEFAULT_LORA_A_STRENGTH,
+                )
+                lora_b_strength = gr.Slider(
+                    label="Reasoning Enhancer LoRA",
+                    minimum=0.0,
+                    maximum=2.0,
+                    step=0.05,
+                    value=DEFAULT_LORA_B_STRENGTH,
+                )
                 schedule = gr.Dropdown(
                     label="Sigma schedule",
                     choices=list(SCHEDULES),
@@ -583,6 +635,8 @@ with gr.Blocks(title="PlagueKind · MiniMax-H3") as demo:
         interpolation,
         seed,
         upsample,
+        lora_a_strength,
+        lora_b_strength,
     ]
 
     # Two blocks rather than one with empty keyframe cells: the positional order below matches `generate`'s