Update app.py

app.py

@@ -349,40 +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]
-
+    # Note: Per-request LoRA loading (no caching). 
+    # If performance becomes an issue, add caching back with correct StateDict handling.
     with safe_open(path, framework="pt", device="cpu") as f:
-        tensors = {k: f.get_tensor(k).contiguous() for k in f.keys()}
+        tensors = {}
+        for key in f.keys():
+            # Apply ComfyUI→base-model key renaming so LoRA weights match transformer keys
+            renamed_key = LTXV_LORA_COMFY_RENAMING_MAP.apply_to_key(key)
+            if renamed_key is None:
+                renamed_key = key  # Keep original if no renaming match
+            tensors[renamed_key] = f.get_tensor(key).contiguous()
 
     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)
-    rm = round(float(motion_strength), 2)
-    rd = round(float(dreamlay_strength), 2)
-    rs = round(float(mself_strength), 2)
-    rr = round(float(dramatic_strength), 2)
-    rf = round(float(fluid_strength), 2)
-    rl = round(float(liquid_strength), 2)
-    ro = round(float(demopose_strength), 2)
-    rv = round(float(voice_strength), 2)
-    re = round(float(realism_strength), 2)
-    rt = round(float(transition_strength), 2)
-    ry = round(float(physics_strength), 2)
-    ri = round(float(reasoning_strength), 2)
-    rw = round(float(twostep_strength), 2)
-    key_str = f"{pose_lora_path}:{rp}|{general_lora_path}:{rg}|{motion_lora_path}:{rm}|{dreamlay_lora_path}:{rd}|{mself_lora_path}:{rs}|{dramatic_lora_path}:{rr}|{fluid_lora_path}:{rf}|{liquid_lora_path}:{rl}|{demopose_lora_path}:{ro}|{voice_lora_path}:{rv}|{realism_lora_path}:{re}|{transition_lora_path}:{rt}|{physics_lora_path}:{ry}|{reasoning_lora_path}:{ri}|{twostep_lora_path}:{rw}"
-    key = hashlib.sha256(key_str.encode("utf-8")).hexdigest()
-    return key, key_str
+    return StateDict(sd=tensors, device=torch.device("cpu"), size=size, dtype=dtypes)
 
 def _collect_lora_specs(
     pose_strength: float,
@@ -401,8 +382,8 @@ def _collect_lora_specs(
     reasoning_strength: float,
     twostep_strength: float,
 ):
-    # Keep all 14 adapters in the active list; zero strength means no effect.
-    return [
+    """Collect (path, strength) pairs for all LoRAs with non-zero strength."""
+    specs = [
         (pose_lora_path, round(float(pose_strength), 2)),
         (general_lora_path, round(float(general_strength), 2)),
         (motion_lora_path, round(float(motion_strength), 2)),
@@ -419,6 +400,8 @@ def _collect_lora_specs(
         (reasoning_lora_path, round(float(reasoning_strength), 2)),
         (twostep_lora_path, round(float(twostep_strength), 2)),
     ]
+    # Filter out zero-strength LoRAs
+    return [(path, strength) for path, strength in specs if strength != 0.0]
 
 
 def apply_current_loras_to_transformer(
@@ -438,74 +421,60 @@ def apply_current_loras_to_transformer(
     reasoning_strength: float,
     twostep_strength: float,
 ):
-    global ACTIVE_LORA_KEY
+    global _transformer
 
-    key, _ = _make_lora_key(
+    # Collect non-zero strength LoRAs
+    lora_specs = _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
+        physics_strength, reasoning_strength, twostep_strength,
     )
 
-    if key == ACTIVE_LORA_KEY:
-        return "LoRAs already active."
+    # No LoRAs to apply — skip
+    if not lora_specs:
+        return "No LoRAs (all zero strength)."
 
-    if key in LORA_STATE_CACHE:
-        fused_state = LORA_STATE_CACHE[key]
-    else:
-        loras = [
-            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,
-            )
-            if strength != 0.0
-        ]
-        
-        base_model_sd = _StateDictModel(
-            {k: v.clone() for k, v in BASE_TRANSFORMER_STATE.items()}
-        )
-        
-        fused_model_sd = apply_loras(
-            base_model_sd,
-            loras,
-            dtype=pipeline.model_ledger.dtype,
-        )
-        
-        fused_state = (
-            fused_model_sd.sd
-            if hasattr(fused_model_sd, "sd")
-            else fused_model_sd
+    # Build base model StateDict (proper type for apply_loras)
+    base_model_sd = StateDict(
+        sd={k: v.clone() for k, v in BASE_TRANSFORMER_STATE.items()},
+        device=torch.device("cpu"),
+        size=sum(v.numel() * v.element_size() for v in BASE_TRANSFORMER_STATE.values()),
+        dtype={v.dtype for v in BASE_TRANSFORMER_STATE.values()},
+    )
+
+    # Build LoraStateDictWithStrength objects
+    loras = [
+        LoraStateDictWithStrength(
+            state_dict=_load_lora_state_dict(path),
+            strength=strength,
         )
-        
-        fused_state = fused_model_sd.sd if hasattr(fused_model_sd, "sd") else fused_model_sd
-        LORA_STATE_CACHE[key] = fused_state
+        for path, strength in lora_specs
+    ]
 
+    # Fuse LoRAs into base model
+    fused_model_sd = apply_loras(
+        base_model_sd,
+        loras,
+        dtype=pipeline.model_ledger.dtype,
+    )
+
+    # Extract plain dict from StateDict for load_state_dict
+    fused_state = fused_model_sd.sd
+
+    # Load fused state dict into transformer
     with torch.no_grad():
-        missing, unexpected = _transformer.load_state_dict(fused_state, strict=False)
+        fused_state_cuda = {
+            k: (v.to(_transformer.device) if v.device == torch.device("cpu") else v)
+            for k, v in fused_state.items()
+        }
+        missing, unexpected = _transformer.load_state_dict(fused_state_cuda, strict=False)
         if missing or unexpected:
-            print(
-                f"[LoRA] state_dict mismatch: missing={len(missing)}, unexpected={len(unexpected)}"
-            )
+            print(f"[LoRA] state_dict load: missing={len(missing)}, unexpected={len(unexpected)}")
+            if missing:
+                print(f"  Missing keys (first 5): {missing[:5]}")
 
-    ACTIVE_LORA_KEY = key
-    return f"Applied LoRAs: {key[:12]}"
+    return f"Applied {len(lora_specs)} LoRA(s)."
 
 # ---- REPLACE PRELOAD BLOCK START ----
 # Preload all models for ZeroGPU tensor packing.
@@ -533,8 +502,6 @@ BASE_TRANSFORMER_STATE = {
 class _StateDictModel:
     def __init__(self, sd: dict[str, torch.Tensor]):
         self.sd = sd
-ACTIVE_LORA_KEY: str | None = None
-LORA_STATE_CACHE: dict[str, dict[str, torch.Tensor]] = {}
 _video_encoder = _orig_video_encoder_factory()
 _video_decoder = _orig_video_decoder_factory()
 _audio_encoder = _orig_audio_encoder_factory()