Update app.py

app.py

@@ -252,10 +252,6 @@ LORA_CACHE_DIR = Path("lora_cache")
 LORA_CACHE_DIR.mkdir(exist_ok=True)
 current_lora_key: str | None = None
 
-PENDING_LORA_KEY: str | None = None
-PENDING_LORA_STATE: dict[str, torch.Tensor] | None = None
-PENDING_LORA_STATUS: str = "No LoRA state prepared yet."
-
 weights_dir = Path("weights")
 weights_dir.mkdir(exist_ok=True)
 checkpoint_path = hf_hub_download(
@@ -264,7 +260,6 @@ checkpoint_path = hf_hub_download(
     local_dir=str(weights_dir),
     local_dir_use_symlinks=False,
 )
-spatial_upsampler_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors")
 
 print("[Gemma] Setting up abliterated Gemma text encoder...")
 MERGED_WEIGHTS = "/tmp/abliterated_gemma_merged.safetensors"
@@ -367,7 +362,6 @@ print(f"Transition LoRA: {transition_lora_path}")
 # ----------------------------------------------------------------
 
 print(f"Checkpoint: {checkpoint_path}")
-print(f"Spatial upsampler: {spatial_upsampler_path}")
 
 # Initialize pipeline WITH text encoder and optional audio support
 # ---- Replace block (pipeline init) lines 275-281 ----
@@ -380,162 +374,141 @@ pipeline = LTX23DistilledA2VPipeline(
 )
 # ----------------------------------------------------------------
 
-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) -> 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)
-    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}"
-    key = hashlib.sha256(key_str.encode("utf-8")).hexdigest()
-    return key, key_str
-
-
-def prepare_lora_cache(
-    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,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """
-    CPU-only step:
-    - checks cache
-    - loads cached fused transformer state_dict, or
-    - builds fused transformer on CPU and saves it
-    The resulting state_dict is stored in memory and can be applied later.
-    """
-    global PENDING_LORA_KEY, PENDING_LORA_STATE, PENDING_LORA_STATUS
-
-    ledger = pipeline.model_ledger
-    key, _ = _make_lora_key(pose_strength, general_strength, motion_strength, dreamlay_strength, mself_strength, dramatic_strength, fluid_strength, liquid_strength, demopose_strength, voice_strength, realism_strength, transition_strength)
-    cache_path = LORA_CACHE_DIR / f"{key}.safetensors"
-
-    progress(0.05, desc="Preparing LoRA state")
-    if cache_path.exists():
-        try:
-            progress(0.20, desc="Loading cached fused state")
-            state = load_file(str(cache_path))
-            PENDING_LORA_KEY = key
-            PENDING_LORA_STATE = state
-            PENDING_LORA_STATUS = f"Loaded cached LoRA state: {cache_path.name}"
-            return PENDING_LORA_STATUS
-        except Exception as e:
-            print(f"[LoRA] Cache load failed: {type(e).__name__}: {e}")
-
-    entries = [
-        (pose_lora_path, round(float(pose_strength), 2)),
-        (general_lora_path, round(float(general_strength), 2)),
-        (motion_lora_path, round(float(motion_strength), 2)),
-        (dreamlay_lora_path, round(float(dreamlay_strength), 2)),
-        (mself_lora_path, round(float(mself_strength), 2)),
-        (dramatic_lora_path, round(float(dramatic_strength), 2)),
-        (fluid_lora_path, round(float(fluid_strength), 2)),
-        (liquid_lora_path, round(float(liquid_strength), 2)),
-        (demopose_lora_path, round(float(demopose_strength), 2)),
-        (voice_lora_path, round(float(voice_strength), 2)),
-        (realism_lora_path, round(float(realism_strength), 2)),
-        (transition_lora_path, round(float(transition_strength), 2)),
-    ]
-    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
-    ]
+# Currently applied LoRA deltas — stored so they can be undone before re-applying
+_applied_lora_deltas: dict[str, torch.Tensor] = {}
+_applied_lora_config: list[tuple[str, float]] = []
 
-    if not loras_for_builder:
-        PENDING_LORA_KEY = None
-        PENDING_LORA_STATE = None
-        PENDING_LORA_STATUS = "No non-zero LoRA strengths selected; nothing to prepare."
-        return PENDING_LORA_STATUS
 
-    tmp_ledger = None
-    new_transformer_cpu = None
-    try:
-        progress(0.35, desc="Building fused CPU transformer")
-        tmp_ledger = pipeline.model_ledger.__class__(
-            dtype=ledger.dtype,
-            device=torch.device("cpu"),
-            checkpoint_path=str(checkpoint_path),
-            spatial_upsampler_path=str(spatial_upsampler_path),
-            gemma_root_path=str(gemma_root),
-            loras=tuple(loras_for_builder),
-            quantization=getattr(ledger, "quantization", None),
-        )
-        new_transformer_cpu = tmp_ledger.transformer()
+def _load_and_rename_lora_tensors(lora_path: str) -> dict[str, torch.Tensor]:
+    """Load LoRA tensors from disk and apply ComfyUI→LTX key renaming."""
+    tensors = {}
+    with safe_open(lora_path, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            tensors[key] = f.get_tensor(key)
 
-        progress(0.70, desc="Extracting fused state_dict")
-        state = {
-            k: v.detach().cpu().contiguous()
-            for k, v in new_transformer_cpu.state_dict().items()
-        }
-        save_file(state, str(cache_path))
+    renamed = {}
+    for key, tensor in tensors.items():
+        new_key = key
+        for old_substr, new_substr in LTXV_LORA_COMFY_RENAMING_MAP.items():
+            new_key = new_key.replace(old_substr, new_substr)
+        renamed[new_key] = tensor
 
-        PENDING_LORA_KEY = key
-        PENDING_LORA_STATE = state
-        PENDING_LORA_STATUS = f"Built and cached LoRA state: {cache_path.name}"
-        return PENDING_LORA_STATUS
+    return renamed
 
-    except Exception as e:
-        import traceback
-        print(f"[LoRA] Prepare failed: {type(e).__name__}: {e}")
-        print(traceback.format_exc())
-        PENDING_LORA_KEY = None
-        PENDING_LORA_STATE = None
-        PENDING_LORA_STATUS = f"LoRA prepare failed: {type(e).__name__}: {e}"
-        return PENDING_LORA_STATUS
-
-    finally:
-        try:
-            del new_transformer_cpu
-        except Exception:
-            pass
-        try:
-            del tmp_ledger
-        except Exception:
-            pass
-        gc.collect()
 
+def _compute_lora_deltas(lora_path: str, strength: float) -> dict[str, torch.Tensor]:
+    """Compute weight delta tensors for a single LoRA at given strength."""
+    tensors = _load_and_rename_lora_tensors(lora_path)
+    deltas = {}
 
-def apply_prepared_lora_state_to_pipeline():
-    """
-    Fast step: copy the already prepared CPU state into the live transformer.
-    This is the only part that should remain near generation time.
-    """
-    global current_lora_key, PENDING_LORA_KEY, PENDING_LORA_STATE
+    # Collect all base keys that have a down component
+    base_keys = set()
+    for key in tensors:
+        for suffix in [".lora_down.weight", ".lora_A.weight"]:
+            if key.endswith(suffix):
+                base_keys.add(key[: -len(suffix)])
 
-    if PENDING_LORA_STATE is None or PENDING_LORA_KEY is None:
-        print("[LoRA] No prepared LoRA state available; skipping.")
-        return False
+    for base in base_keys:
+        down = tensors.get(base + ".lora_down.weight") or tensors.get(base + ".lora_A.weight")
+        up = tensors.get(base + ".lora_up.weight") or tensors.get(base + ".lora_B.weight")
 
-    if current_lora_key == PENDING_LORA_KEY:
-        print("[LoRA] Prepared LoRA state already active; skipping.")
-        return True
+        if down is None or up is None:
+            continue
+
+        alpha_val = tensors.get(base + ".alpha")
+        scale = (alpha_val.item() / down.shape[0]) if alpha_val is not None else 1.0
+
+        down_f = down.float()
+        up_f = up.float()
+
+        if down_f.dim() == 2 and up_f.dim() == 2:
+            delta = up_f @ down_f
+        elif down_f.dim() == 4:
+            delta = (up_f.flatten(1) @ down_f.flatten(1)).view(
+                up_f.shape[0], down_f.shape[1], *up_f.shape[2:]
+            )
+        else:
+            print(f"[LoRA] Skipping {base}: unexpected dims down={down_f.dim()} up={up_f.dim()}")
+            continue
 
-    existing_transformer = _transformer
+        deltas[base + ".weight"] = (delta * strength * scale).to(torch.bfloat16)
+
+    return deltas
+
+
+def apply_loras_to_transformer(
+    pose_strength, general_strength, motion_strength, dreamlay_strength,
+    mself_strength, dramatic_strength, fluid_strength, liquid_strength,
+    demopose_strength, voice_strength, realism_strength, transition_strength,
+):
+    global _applied_lora_deltas, _applied_lora_config
+
+    lora_configs = [
+        (pose_lora_path,       round(float(pose_strength), 2)),
+        (general_lora_path,    round(float(general_strength), 2)),
+        (motion_lora_path,     round(float(motion_strength), 2)),
+        (dreamlay_lora_path,   round(float(dreamlay_strength), 2)),
+        (mself_lora_path,      round(float(mself_strength), 2)),
+        (dramatic_lora_path,   round(float(dramatic_strength), 2)),
+        (fluid_lora_path,      round(float(fluid_strength), 2)),
+        (liquid_lora_path,     round(float(liquid_strength), 2)),
+        (demopose_lora_path,   round(float(demopose_strength), 2)),
+        (voice_lora_path,      round(float(voice_strength), 2)),
+        (realism_lora_path,    round(float(realism_strength), 2)),
+        (transition_lora_path, round(float(transition_strength), 2)),
+    ]
+
+    # Skip if config hasn't changed since last application
+    if lora_configs == _applied_lora_config:
+        print("[LoRA] Config unchanged, skipping re-application.")
+        return
+
+    # Undo previously applied deltas
+    if _applied_lora_deltas:
+        print(f"[LoRA] Undoing {len(_applied_lora_deltas)} previously applied delta(s)...")
+        with torch.no_grad():
+            for name, param in _transformer.named_parameters():
+                if name in _applied_lora_deltas:
+                    param.data -= _applied_lora_deltas[name].to(
+                        device=param.device, dtype=param.dtype
+                    )
+        _applied_lora_deltas = {}
+        gc.collect()
+
+    active = [(p, s) for p, s in lora_configs if p is not None and s != 0.0]
+    if not active:
+        print("[LoRA] No active LoRAs.")
+        _applied_lora_config = lora_configs
+        return
+
+    print(f"[LoRA] Computing deltas for {len(active)} active LoRA(s)...")
+    combined_deltas: dict[str, torch.Tensor] = {}
+    for lora_path, strength in active:
+        try:
+            deltas = _compute_lora_deltas(lora_path, strength)
+            for key, delta in deltas.items():
+                if key in combined_deltas:
+                    combined_deltas[key] = combined_deltas[key] + delta
+                else:
+                    combined_deltas[key] = delta
+            print(f"[LoRA] {Path(lora_path).name}: {len(deltas)} delta(s) at strength {strength}")
+        except Exception as e:
+            import traceback
+            print(f"[LoRA] Failed on {lora_path}: {e}\n{traceback.format_exc()}")
+
+    applied_count = 0
     with torch.no_grad():
-        missing, unexpected = existing_transformer.load_state_dict(PENDING_LORA_STATE, strict=False)
-        if missing or unexpected:
-            print(f"[LoRA] load_state_dict mismatch: missing={len(missing)}, unexpected={len(unexpected)}")
+        for name, param in _transformer.named_parameters():
+            if name in combined_deltas:
+                param.data += combined_deltas[name].to(
+                    device=param.device, dtype=param.dtype
+                )
+                applied_count += 1
 
-    current_lora_key = PENDING_LORA_KEY
-    print("[LoRA] Prepared LoRA state applied to the pipeline.")
-    return True
+    _applied_lora_deltas = combined_deltas
+    _applied_lora_config = lora_configs
+    print(f"[LoRA] Applied {applied_count} weight delta(s) to live transformer.")
+    gc.collect()
 
 # ---- REPLACE PRELOAD BLOCK START ----
 # Preload all models for ZeroGPU tensor packing.
@@ -550,7 +523,6 @@ _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
 
@@ -561,7 +533,6 @@ _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()
 
@@ -573,7 +544,6 @@ ledger.video_decoder = lambda: _video_decoder
 ledger.audio_encoder = lambda: _audio_encoder
 ledger.audio_decoder = lambda: _audio_decoder
 ledger.vocoder = lambda: _vocoder
-ledger.spatial_upsampler = lambda: _spatial_upsampler
 ledger.text_encoder = lambda: _text_encoder
 ledger.gemma_embeddings_processor = lambda: _embeddings_processor
 
@@ -716,7 +686,11 @@ def generate_video(
 
         log_memory("before pipeline call")
 
-        apply_prepared_lora_state_to_pipeline()
+        apply_loras_to_transformer(
+            pose_strength, general_strength, motion_strength, dreamlay_strength,
+            mself_strength, dramatic_strength, fluid_strength, liquid_strength,
+            demopose_strength, voice_strength, realism_strength, transition_strength,
+        )
 
         video, audio = pipeline(
             prompt=prompt,
@@ -833,11 +807,6 @@ with gr.Blocks(title="LTX-2.3 Distilled") as demo:
                         label="Transition strength",
                         minimum=0.0, maximum=2.0, value=0.0, step=0.01
                     )
-                prepare_lora_btn = gr.Button("Prepare / Load LoRA Cache", variant="secondary")
-                lora_status = gr.Textbox(
-                    label="LoRA Cache Status",
-                    value="No LoRA state prepared yet.",
-                    interactive=False,
                 )
 
         with gr.Column():
@@ -907,12 +876,6 @@ with gr.Blocks(title="LTX-2.3 Distilled") as demo:
         inputs=[first_image, last_image, high_res],
         outputs=[width, height],
     )
-
-    prepare_lora_btn.click(
-        fn=prepare_lora_cache,
-        inputs=[pose_strength, general_strength, motion_strength, dreamlay_strength, mself_strength, dramatic_strength, fluid_strength, liquid_strength, demopose_strength, voice_strength, realism_strength, transition_strength],
-        outputs=[lora_status],
-    )
     
     generate_btn.click(
         fn=generate_video,