Update app.py

app.py

@@ -44,8 +44,13 @@ import gradio as gr
 import numpy as np
 from huggingface_hub import hf_hub_download, snapshot_download
 from safetensors.torch import load_file
-from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
+from ltx_core.loader.primitives import (
+    StateDict,
+    LoraPathStrengthAndSDOps,
+    LoraStateDictWithStrength,
+)
 from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
+from safetensors import safe_open
 
 from ltx_core.loader.fuse_loras import apply_loras
 
@@ -344,6 +349,21 @@ pipeline = LTX23DistilledA2VPipeline(
 )
 # ----------------------------------------------------------------
 
+LORA_STATE_DICT_CACHE: dict[str, StateDict] = {}
+
+def _load_lora_state_dict(path: str) -> StateDict:
+    if path in LORA_STATE_DICT_CACHE:
+        return LORA_STATE_DICT_CACHE[path]
+
+    with safe_open(path, framework="pt", device="cpu") as f:
+        tensors = {k: f.get_tensor(k).contiguous() for k in f.keys()}
+
+    size = sum(t.numel() * t.element_size() for t in tensors.values())
+    dtypes = {t.dtype for t in tensors.values()}
+    sd = StateDict(sd=tensors, device=torch.device("cpu"), size=size, dtype=dtypes)
+    LORA_STATE_DICT_CACHE[path] = sd
+    return sd
+
 def _make_lora_key(pose_strength: float, general_strength: float, motion_strength: float, dreamlay_strength: float, mself_strength: float, dramatic_strength: float, fluid_strength: float, liquid_strength: float, demopose_strength: float, voice_strength: float, realism_strength: float, transition_strength: float, physics_strength: float, reasoning_strength: float, twostep_strength: float) -> tuple[str, str]:
     rp = round(float(pose_strength), 2)
     rg = round(float(general_strength), 2)
@@ -434,15 +454,30 @@ def apply_current_loras_to_transformer(
         fused_state = LORA_STATE_CACHE[key]
     else:
         loras = [
-            LoraPathStrengthAndSDOps(path, strength, LTXV_LORA_COMFY_RENAMING_MAP)
+            LoraStateDictWithStrength(
+                state_dict=_load_lora_state_dict(path),
+                strength=strength,
+            )
             for path, strength in _collect_lora_specs(
-                pose_strength, general_strength, motion_strength, dreamlay_strength,
-                mself_strength, dramatic_strength, fluid_strength, liquid_strength,
-                demopose_strength, voice_strength, realism_strength, transition_strength,
-                physics_strength, reasoning_strength, twostep_strength,
+                pose_strength,
+                general_strength,
+                motion_strength,
+                dreamlay_strength,
+                mself_strength,
+                dramatic_strength,
+                fluid_strength,
+                liquid_strength,
+                demopose_strength,
+                voice_strength,
+                realism_strength,
+                transition_strength,
+                physics_strength,
+                reasoning_strength,
+                twostep_strength,
             )
+            if strength != 0.0
         ]
-
+        
         base_model_sd = _StateDictModel(
             {k: v.clone() for k, v in BASE_TRANSFORMER_STATE.items()}
         )
@@ -453,6 +488,12 @@ def apply_current_loras_to_transformer(
             dtype=pipeline.model_ledger.dtype,
         )
         
+        fused_state = (
+            fused_model_sd.sd
+            if hasattr(fused_model_sd, "sd")
+            else fused_model_sd
+        )
+        
         fused_state = fused_model_sd.sd if hasattr(fused_model_sd, "sd") else fused_model_sd
         LORA_STATE_CACHE[key] = fused_state