Update app.py

app.py

@@ -61,6 +61,8 @@ from ltx_pipelines.utils.helpers import (
     simple_denoising_func,
 )
 from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video
+from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
+from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
 
 # Force-patch xformers attention into the LTX attention module.
 from ltx_core.model.transformer import attention as _attn_mod
@@ -298,122 +300,86 @@ pipeline = LTX23DistilledA2VPipeline(
 )
 # ----------------------------------------------------------------
 
-# ---- REPLACE apply_loras_to_pipeline WITH THIS IMPLEMENTATION ----
 def apply_loras_to_pipeline(pose_strength: float, general_strength: float, motion_strength: float):
     """
-    Rebuild the transformer with the requested LoRA strengths and hot-swap into the ledger.
-    Strategy (in order):
-      1. If the ledger/builder object exists and exposes 'lora' + 'loras' + the original transformer factory,
-         mutate builder.loras (clear and append .lora(...)) and call the saved original factory to rebuild.
-      2. If a dedicated ledger.apply_loras or pipeline.set_loras exists, try calling them.
-      3. Fallback: set pipeline.loras attribute (best-effort, may be ignored).
+    Build a temporary ModelLedger configured with the requested LoRAs, build the transformer,
+    and hot-swap it into the existing ledger without recreating the pipeline object.
+
+    Strategy:
+      1. Construct LoraPathStrengthAndSDOps entries for any non-zero strengths.
+      2. Use ledger.with_loras(...) to get a temporary ledger configured with those loras.
+      3. Optionally clear the existing cached transformer to reduce peak VRAM, then build the
+         transformer from the temporary ledger and hot-swap it into the live ledger.
+      4. If anything fails, print diagnostics and leave the existing pipeline in place.
     """
     ledger = pipeline.model_ledger
 
-    # Build convenience list (paths may be None if download failed)
-    lora_entries = [
+    # Build convenience list and convert to the LTX primitive type (with sd_ops)
+    entries = [
         (pose_lora_path, float(pose_strength)),
         (general_lora_path, float(general_strength)),
         (motion_lora_path, float(motion_strength)),
     ]
 
-    # 1) Preferred: mutate the builder used by the ledger, then call the original transformer factory.
-    builder = getattr(ledger, "builder", None) or getattr(ledger, "_builder", None)
+    # Keep only nonzero strengths and valid paths (zero == disabled)
+    loras_for_builder = [
+        LoraPathStrengthAndSDOps(path, strength, LTXV_LORA_COMFY_RENAMING_MAP)
+        for path, strength in entries
+        if path is not None and float(strength) != 0.0
+    ]
+
+    if len(loras_for_builder) == 0:
+        print("[LoRA] No nonzero LoRA strengths — skipping rebuild.")
+        return
+
+    try:
+        # Create a temporary ledger configured with the extra LoRAs.
+        # with_loras accepts an iterable of LoraPathStrengthAndSDOps.
+        tmp_ledger = ledger.with_loras(tuple(loras_for_builder))
+        print(f"[LoRA] Built temporary ledger with {len(loras_for_builder)} LoRA(s).")
 
-    if builder is not None:
+        # Attempt to free previously cached transformer instance to reduce peak VRAM.
+        # (We cached instances earlier and replaced ledger.<component> with lambdas returning them.)
         try:
-            # Clear existing builder.loras if present.
-            if hasattr(builder, "loras") and isinstance(getattr(builder, "loras"), list):
-                try:
-                    builder.loras.clear()
-                except Exception:
-                    # fallback: reassign
-                    try:
-                        setattr(builder, "loras", [])
-                    except Exception:
-                        pass
-
-            # If there is an explicit builder.lora(...) helper, use it to register LoRAs.
-            if hasattr(builder, "lora"):
-                for path, strength in lora_entries:
-                    try:
-                        if path is None:
-                            continue
-                        # Always register (builder.lora should create the proper LoRA object expected by build).
-                        # The builder may accept strength==0.0; if not, register only >0.0 to avoid extra work.
-                        if float(strength) != 0.0:
-                            builder.lora(path, float(strength))
-                    except Exception as e:
-                        print(f"[LoRA] builder.lora(...) failed for {path}: {type(e).__name__}: {e}")
-            else:
-                # If there's no builder.lora helper, try to append sensible objects into builder.loras.
-                # We don't know the exact LoRA object type, so we skip if we can't.
-                print("[LoRA] builder has no 'lora' helper; attempting to set builder.loras to an empty list.")
-                try:
-                    setattr(builder, "loras", [])
-                except Exception:
-                    pass
-
-            # Use the saved original transformer factory to (re)create a transformer using the mutated builder.
-            # The original factory should call builder.build(...) internally.
-            if "_orig_transformer_factory" in globals():
-                print("[LoRA] Rebuilding transformer from builder and hot-swapping into ledger...")
-                try:
-                    new_transformer = _orig_transformer_factory()
-                    # free previous cached transformer to reduce peak memory
-                    global _transformer
-                    try:
-                        # delete previous Python ref and empty cache; underlying CUDA memory should be released when
-                        # the module is garbage-collected. We proactively call torch.cuda.empty_cache().
-                        del _transformer
-                    except Exception:
-                        pass
-                    torch.cuda.empty_cache()
-
-                    # install new transformer as the one returned by ledger.transformer()
-                    _transformer = new_transformer
-                    ledger.transformer = lambda: _transformer
-                    print("[LoRA] Transformer rebuilt and hot-swapped successfully.")
-                    return
-                except Exception as e:
-                    print(f"[LoRA] Error while rebuilding transformer via original factory: {type(e).__name__}: {e}")
-                    # fallthrough to other attempts
-            else:
-                print("[LoRA] _orig_transformer_factory is not available; cannot rebuild transformer via builder.")
+            # If ModelLedger implements clear_vram, call it to release cached GPU tensors.
+            if hasattr(ledger, "clear_vram"):
+                ledger.clear_vram()
+                print("[LoRA] Cleared old ledger VRAM cache before building new transformer.")
         except Exception as e:
-            print(f"[LoRA] Unexpected error while manipulating builder: {type(e).__name__}: {e}")
-            # proceed to fallbacks below
+            print(f"[LoRA] Warning: ledger.clear_vram() failed: {type(e).__name__}: {e}")
+
+        # Build the new transformer from the temporary ledger (this will load & fuse LoRAs).
+        print("[LoRA] Building transformer from temporary ledger (this may take time / spike VRAM)...")
+        new_transformer = tmp_ledger.transformer()  # returns an X0Model moved to device
+        print("[LoRA] New transformer built successfully.")
+
+        # Replace cached transformer instance and hot-swap ledger.transformer to return the new one.
+        global _transformer
+        try:
+            # Remove old Python ref if present to allow GC.
+            del _transformer
+        except Exception:
+            pass
+        torch.cuda.empty_cache()
+        _transformer = new_transformer
+        ledger.transformer = lambda: _transformer
+        print("[LoRA] Hot-swapped new transformer into ledger successfully.")
+
+        # Done
+        return
 
-    # 2) Try high-level APIs if present (some pipeline/ledger versions expose helpers)
-    try:
-        if hasattr(ledger, "apply_loras"):
-            print("[LoRA] Calling ledger.apply_loras(...)")
-            # Best-effort call (signature may vary across versions)
-            try:
-                ledger.apply_loras([{"path": p, "strength": s} for p, s in lora_entries if p is not None])
-            except Exception:
-                # also try a simpler form
-                ledger.apply_loras(lora_entries)
-            return
-        if hasattr(pipeline, "set_loras"):
-            print("[LoRA] Calling pipeline.set_loras(...)")
-            try:
-                pipeline.set_loras(lora_entries)
-                return
-            except Exception as e:
-                print(f"[LoRA] pipeline.set_loras failed: {type(e).__name__}: {e}")
     except Exception as e:
-        print(f"[LoRA] Exception when trying high-level apply APIs: {type(e).__name__}: {e}")
+        import traceback
+        print(f"[LoRA] Error during builder-based LoRA application: {type(e).__name__}: {e}")
+        print(traceback.format_exc())
 
-    # 3) Fallback: set pipeline.loras attribute (best-effort, may not be used)
+    # Final fallback (should rarely hit if above path works):
     try:
         print("[LoRA] Falling back to pipeline.loras attribute assignment (best-effort).")
-        pipeline.loras = [(p, float(s)) for p, s in lora_entries if p is not None]
+        pipeline.loras = [(p, float(s)) for p, s in entries if p is not None]
     except Exception as e:
         print(f"[LoRA] Fallback pipeline.loras assignment failed: {type(e).__name__}: {e}")
-
     print("[LoRA] apply_loras_to_pipeline finished (some approaches may not have taken effect).")
-# ---- END replacement ----
 
 # ---- REPLACE PRELOAD BLOCK START ----
 # Preload all models for ZeroGPU tensor packing.