Update app.py

app.py

@@ -123,6 +123,9 @@ class LTX23DistilledA2VPipeline:
     def __init__(
         self,
         distilled_checkpoint_path: str,
+        istilled_lora_path: str,
+        distilled_lora_strength_stage_1: float,
+        distilled_lora_strength_stage_2: float,
         spatial_upsampler_path: str,
         gemma_root: str,
         loras: tuple,
@@ -134,16 +137,31 @@ class LTX23DistilledA2VPipeline:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.dtype = torch.bfloat16
 
-        self.model_ledger = ModelLedger(
+        distilled_lora_stage_1 = LoraPathStrengthAndSDOps(
+            path=distilled_lora_path,
+            strength=distilled_lora_strength_stage_1,
+            sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
+        )
+        distilled_lora_stage_2 = LoraPathStrengthAndSDOps(
+            path=distilled_lora_path,
+            strength=distilled_lora_strength_stage_2,
+            sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
+        )
+
+        self.stage_1_model_ledger = ModelLedger(
             dtype=self.dtype,
-            device=self.device,
-            checkpoint_path=distilled_checkpoint_path,
+            device=device,
+            checkpoint_path=checkpoint_path,
             gemma_root_path=gemma_root,
             spatial_upsampler_path=spatial_upsampler_path,
-            loras=loras,
+            loras=(*loras, distilled_lora_stage_1),
             quantization=quantization,
         )
 
+        self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
+            loras=(*loras, distilled_lora_stage_2),
+        )
+
         self.pipeline_components = PipelineComponents(
             dtype=self.dtype,
             device=self.device,
@@ -209,8 +227,8 @@ class LTX23DistilledA2VPipeline:
                 )
                 encoded_audio_latent = torch.cat([encoded_audio_latent, pad], dim=2)
 
-        video_encoder = self.model_ledger.video_encoder()
-        transformer = self.model_ledger.transformer()
+        video_encoder = self.stage_1_model_ledger.video_encoder()
+        transformer = self.stage_1_model_ledger.transformer()
 
         # Stage 1: Generate sigmas using LTX2Scheduler with user-specified steps
         empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(
@@ -292,10 +310,11 @@ class LTX23DistilledA2VPipeline:
         cleanup_memory()
 
         # ── Upscaling ──
+        video_encoder = self.stage_1_model_ledger.video_encoder()
         upscaled_video_latent = upsample_video(
             latent=video_state.latent[:1],
             video_encoder=video_encoder,
-            upsampler=self.model_ledger.spatial_upsampler(),
+            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
         )
 
         # ── Stage 2: Full resolution ──
@@ -309,6 +328,9 @@ class LTX23DistilledA2VPipeline:
             dtype=dtype,
             device=self.device,
         )
+
+        transformer = self.stage_2_model_ledger.transformer()
+        
         video_state, audio_state = denoise_audio_video(
             output_shape=stage_2_output_shape,
             conditionings=stage_2_conditionings,
@@ -330,21 +352,22 @@ class LTX23DistilledA2VPipeline:
         # ── Decode both video and audio ──
         decoded_video = vae_decode_video(
             video_state.latent,
-            self.model_ledger.video_decoder(),
+            self.stage_2_model_ledger.video_decoder(),
             tiling_config,
             generator,
         )
         decoded_audio_output = vae_decode_audio(
             audio_state.latent,
-            self.model_ledger.audio_decoder(),
-            self.model_ledger.vocoder(),
+            self.stage_2_model_ledger.audio_decoder(),
+            self.stage_2_model_ledger.vocoder(),
         )
 
         return decoded_video, decoded_audio_output
 
 # Model repos
-LTX_MODEL_REPO = "SulphurAI/Sulphur-2-base"
+LTX_MODEL_REPO = "TenStrip/LTX2.3-10Eros"
 GEMMA_REPO ="Lightricks/gemma-3-12b-it-qat-q4_0-unquantized"
+DISTILLED_LORA_REPO = "TenStrip/LTX2.3_Distilled_Lora_1.1_Experiments"
 
 # Download model checkpoints
 print("=" * 80)
@@ -364,10 +387,11 @@ weights_dir = Path("weights")
 weights_dir.mkdir(exist_ok=True)
 checkpoint_path = hf_hub_download(
     repo_id=LTX_MODEL_REPO,
-    filename="sulphur_distil_bf16.safetensors",
+    filename="10Eros_v1_bf16.safetensors",
     local_dir=str(weights_dir),
     local_dir_use_symlinks=False,
 )
+distilled_lora_path = hf_hub_download(repo_id=DISTILLED_LORA_REPO, filename="ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors")
 spatial_upsampler_path = hf_hub_download(repo_id="Lightricks/LTX-2.3", filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors")
 gemma_root = snapshot_download(repo_id=GEMMA_REPO)
 
@@ -415,6 +439,9 @@ print(f"[Gemma] Root ready: {gemma_root}")
 
 pipeline = LTX23DistilledA2VPipeline(
     distilled_checkpoint_path=checkpoint_path,
+    distilled_lora=distilled_lora_path,
+    distilled_lora_strength_stage_1=1.0,
+    distilled_lora_strength_stage_2=0.5,
     spatial_upsampler_path=spatial_upsampler_path,
     gemma_root=gemma_root,
     loras=[],
@@ -589,22 +616,31 @@ def apply_prepared_lora_state_to_pipeline():
 
 # Preload all models for ZeroGPU tensor packing.
 print("Preloading all models (including Gemma and audio components)...")
-ledger = pipeline.model_ledger
-
-# 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()
+
+# We now have two ledgers — stage 1 (distilled LoRA @ 1.0) and stage 2 (distilled LoRA @ 0.5).
+# Both share the same dev checkpoint and spatial upsampler; only the transformer differs.
+ledger_s1 = pipeline.stage_1_model_ledger
+ledger_s2 = pipeline.stage_2_model_ledger
+
+# Save the original factory methods from BOTH ledgers so we can rebuild individual components later.
+_orig_transformer_factory_s1 = ledger_s1.transformer
+_orig_transformer_factory_s2 = ledger_s2.transformer
+_orig_video_encoder_factory = ledger_s1.video_encoder
+_orig_video_decoder_factory = ledger_s1.video_decoder
+_orig_audio_encoder_factory = ledger_s1.audio_encoder
+_orig_audio_decoder_factory = ledger_s1.audio_decoder
+_orig_vocoder_factory = ledger_s1.vocoder
+_orig_spatial_upsampler_factory = ledger_s1.spatial_upsampler
+_orig_text_encoder_factory = ledger_s1.text_encoder
+_orig_gemma_embeddings_factory = ledger_s1.gemma_embeddings_processor
+
+# Call the factories to create cached instances.
+# Stage 1 transformer: dev checkpoint + distilled LoRA @ strength 1.0 (baked in at build time)
+_transformer_s1 = _orig_transformer_factory_s1()
+# Stage 2 transformer: dev checkpoint + distilled LoRA @ strength 0.5 (baked in at build time)
+_transformer_s2 = _orig_transformer_factory_s2()
+
+# Shared components — only need one copy since both ledgers use the same VAE/encoder paths.
 _video_encoder = _orig_video_encoder_factory()
 _video_decoder = _orig_video_decoder_factory()
 _audio_encoder = _orig_audio_encoder_factory()
@@ -615,18 +651,30 @@ _text_encoder = _orig_text_encoder_factory()
 _embeddings_processor = _orig_gemma_embeddings_factory()
 
 # Replace ledger methods with lightweight lambdas that return the cached instances.
+# Both ledgers point to the same shared model instances (except transformer).
 # 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
-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
+ledger_s1.transformer = lambda: _transformer_s1
+ledger_s2.transformer = lambda: _transformer_s2
+ledger_s1.video_encoder = lambda: _video_encoder
+ledger_s2.video_encoder = lambda: _video_encoder
+ledger_s1.video_decoder = lambda: _video_decoder
+ledger_s2.video_decoder = lambda: _video_decoder
+ledger_s1.audio_encoder = lambda: _audio_encoder
+ledger_s2.audio_encoder = lambda: _audio_encoder
+ledger_s1.audio_decoder = lambda: _audio_decoder
+ledger_s2.audio_decoder = lambda: _audio_decoder
+ledger_s1.vocoder = lambda: _vocoder
+ledger_s2.vocoder = lambda: _vocoder
+ledger_s1.spatial_upsampler = lambda: _spatial_upsampler
+ledger_s2.spatial_upsampler = lambda: _spatial_upsampler
+ledger_s1.text_encoder = lambda: _text_encoder
+ledger_s2.text_encoder = lambda: _text_encoder
+ledger_s1.gemma_embeddings_processor = lambda: _embeddings_processor
+ledger_s2.gemma_embeddings_processor = lambda: _embeddings_processor
 
 print("All models preloaded (including Gemma text encoder and audio encoder)!")
+print(f"  Stage 1 transformer: {_transformer_s1.__class__.__name__} (distilled LoRA @ 1.0 baked)")
+print(f"  Stage 2 transformer: {_transformer_s2.__class__.__name__} (distilled LoRA @ 0.5 baked)")
 # ---- REPLACE PRELOAD BLOCK END ----
 
 print("=" * 80)