Update app.py

app.py

@@ -370,17 +370,47 @@ pipeline = LTX23DistilledA2VPipeline(
     distilled_checkpoint_path=checkpoint_path,
     spatial_upsampler_path=spatial_upsampler_path,
     gemma_root=gemma_root,
-    loras=[
-        LoraPathStrengthAndSDOps(
-            distilled_lora_path,
-            1.0,
-            LTXV_LORA_COMFY_RENAMING_MAP,
-        )
-    ],
+    loras=[],
     quantization=QuantizationPolicy.fp8_cast(),  # keep FP8 quantization unchanged
 )
 # ----------------------------------------------------------------
 
+DISTILLED_DEFAULT_STATE: dict[str, torch.Tensor] | None = None
+
+def prepare_distilled_default_state():
+    global DISTILLED_DEFAULT_STATE
+
+    if DISTILLED_DEFAULT_STATE is not None:
+        return
+
+    print("Preparing distilled default LoRA state on CPU...")
+    tmp_ledger = pipeline.model_ledger.__class__(
+        dtype=pipeline.model_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=(
+            LoraPathStrengthAndSDOps(
+                distilled_lora_path,
+                1.0,
+                LTXV_LORA_COMFY_RENAMING_MAP,
+            ),
+        ),
+        quantization=None,
+    )
+
+    distilled_transformer = tmp_ledger.transformer()
+    DISTILLED_DEFAULT_STATE = {
+        k: v.detach().cpu().contiguous()
+        for k, v in distilled_transformer.state_dict().items()
+    }
+
+    del distilled_transformer
+    del tmp_ledger
+    gc.collect()
+    print("Distilled default LoRA state prepared.")
+
 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) -> tuple[str, str]:
     rp = round(float(pose_strength), 2)
     rg = round(float(general_strength), 2)
@@ -552,18 +582,15 @@ _orig_spatial_upsampler_factory = ledger.spatial_upsampler
 _orig_text_encoder_factory = ledger.text_encoder
 _orig_gemma_embeddings_factory = ledger.gemma_embeddings_processor
 
-# Keep everything else cached as before.
-_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()
-
 # Call the original factories once to create the cached instances we will serve by default.
 _transformer = _orig_transformer_factory()
+_transformer = _orig_transformer_factory()
+if DISTILLED_DEFAULT_STATE is not None:
+    with torch.no_grad():
+        missing, unexpected = _transformer.load_state_dict(DISTILLED_DEFAULT_STATE, strict=False)
+        if missing or unexpected:
+            print(f"[Distilled default] load_state_dict mismatch: missing={len(missing)}, unexpected={len(unexpected)}")
+    print("[Distilled default] applied to transformer.")
 _video_encoder = _orig_video_encoder_factory()
 _video_decoder = _orig_video_decoder_factory()
 _audio_encoder = _orig_audio_encoder_factory()