Update app.py

app.py

@@ -293,66 +293,158 @@ pipeline = LTX23DistilledA2VPipeline(
     distilled_checkpoint_path=checkpoint_path,
     spatial_upsampler_path=spatial_upsampler_path,
     gemma_root=gemma_root,
-    # initial LoRAs (strengths set to zero by default; will be adjusted at runtime)
-    loras=[
-        (pose_lora_path, 0.0),
-        (general_lora_path, 0.0),
-        (motion_lora_path, 0.0),
-    ],
+    loras=[],
     quantization=QuantizationPolicy.fp8_cast(),  # keep FP8 quantization unchanged
 )
 # ----------------------------------------------------------------
 
-# ---- Insert block: helper to apply LoRA strengths at runtime (between lines 281 and 283) ----
+# ---- REPLACE apply_loras_to_pipeline WITH THIS IMPLEMENTATION ----
 def apply_loras_to_pipeline(pose_strength: float, general_strength: float, motion_strength: float):
     """
-    Best-effort attempt to set LoRA strengths on the existing pipeline without creating a new pipeline object.
-    Tries common APIs that model_ledger / pipeline implementations may provide; falls back to setting
-    pipeline.loras (for lazy fusion) and logs a note.
+    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 the spec list (keep ordering consistent with the pipeline loras list created above)
-    lora_specs = [
+    ledger = pipeline.model_ledger
+
+    # Build convenience list (paths may be None if download failed)
+    lora_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)
+
+    if builder is not None:
+        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.")
+        except Exception as e:
+            print(f"[LoRA] Unexpected error while manipulating builder: {type(e).__name__}: {e}")
+            # proceed to fallbacks below
+
+    # 2) Try high-level APIs if present (some pipeline/ledger versions expose helpers)
     try:
-        # Preferred: model_ledger exposes an apply_loras API that fuses adapters in-place.
-        if hasattr(pipeline, "model_ledger") and hasattr(pipeline.model_ledger, "apply_loras"):
-            print("[LoRA] Applying LoRAs via pipeline.model_ledger.apply_loras(...)")
-            pipeline.model_ledger.apply_loras(lora_specs)
+        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
-
-        # Alternative: pipeline may have helper 'set_loras' or similar
         if hasattr(pipeline, "set_loras"):
-            print("[LoRA] Applying LoRAs via pipeline.set_loras(...)")
-            pipeline.set_loras(lora_specs)
-            return
-
-        # Fallback: set an attribute that some pipeline implementations inspect lazily.
-        print("[LoRA] Setting pipeline.loras attribute (fallback). If pipeline supports lazy fusion, new strengths will be used.")
-        pipeline.loras = lora_specs
-        return
+            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}")
 
+    # 3) Fallback: set pipeline.loras attribute (best-effort, may not be used)
+    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]
     except Exception as e:
-        # Non-fatal: advise user and keep running with whatever pipeline state exists.
-        print(f"[LoRA] Warning: failed to apply LoRAs at runtime: {type(e).__name__}: {e}")
-        print("[LoRA] If you see this, the pipeline may require a restart or an explicit fusion API to apply LoRAs.")
-# ----------------------------------------------------------------
+        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.
 print("Preloading all models (including Gemma and audio components)...")
 ledger = pipeline.model_ledger
-_transformer = ledger.transformer()
-_video_encoder = ledger.video_encoder()
-_video_decoder = ledger.video_decoder()
-_audio_encoder = ledger.audio_encoder()
-_audio_decoder = ledger.audio_decoder()
-_vocoder = ledger.vocoder()
-_spatial_upsampler = ledger.spatial_upsampler()
-_text_encoder = ledger.text_encoder()
-_embeddings_processor = ledger.gemma_embeddings_processor()
 
+# Save the original factory methods so we can rebuild individual components later.
+# These are bound callables on ledger that will call the builder when invoked.
+_orig_transformer_factory = ledger.transformer
+_orig_video_encoder_factory = ledger.video_encoder
+_orig_video_decoder_factory = ledger.video_decoder
+_orig_audio_encoder_factory = ledger.audio_encoder
+_orig_audio_decoder_factory = ledger.audio_decoder
+_orig_vocoder_factory = ledger.vocoder
+_orig_spatial_upsampler_factory = ledger.spatial_upsampler
+_orig_text_encoder_factory = ledger.text_encoder
+_orig_gemma_embeddings_factory = ledger.gemma_embeddings_processor
+
+# Call the original factories once to create the cached instances we will serve by default.
+_transformer = _orig_transformer_factory()
+_video_encoder = _orig_video_encoder_factory()
+_video_decoder = _orig_video_decoder_factory()
+_audio_encoder = _orig_audio_encoder_factory()
+_audio_decoder = _orig_audio_decoder_factory()
+_vocoder = _orig_vocoder_factory()
+_spatial_upsampler = _orig_spatial_upsampler_factory()
+_text_encoder = _orig_text_encoder_factory()
+_embeddings_processor = _orig_gemma_embeddings_factory()
+
+# Replace ledger methods with lightweight lambdas that return the cached instances.
+# We keep the original factories above so we can call them later to rebuild components.
 ledger.transformer = lambda: _transformer
 ledger.video_encoder = lambda: _video_encoder
 ledger.video_decoder = lambda: _video_decoder
@@ -362,7 +454,9 @@ ledger.vocoder = lambda: _vocoder
 ledger.spatial_upsampler = lambda: _spatial_upsampler
 ledger.text_encoder = lambda: _text_encoder
 ledger.gemma_embeddings_processor = lambda: _embeddings_processor
+
 print("All models preloaded (including Gemma text encoder and audio encoder)!")
+# ---- REPLACE PRELOAD BLOCK END ----
 
 print("=" * 80)
 print("Pipeline ready!")