Update app.py

app.py

@@ -45,7 +45,9 @@ from huggingface_hub import hf_hub_download
 
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
+from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
 from ltx_pipelines.distilled import DistilledPipeline
+from ltx_pipelines.utils import helpers as pipeline_helpers
 from ltx_pipelines.utils.args import ImageConditioningInput
 from ltx_pipelines.utils.media_io import encode_video
 
@@ -105,18 +107,6 @@ print("Pipeline ready!")
 print("=" * 80)
 
 
-class PrecomputedTextEncoder(torch.nn.Module):
-    """Fake text encoder that returns pre-computed embeddings."""
-
-    def __init__(self, video_context, audio_context):
-        super().__init__()
-        self.video_context = video_context
-        self.audio_context = audio_context
-
-    def forward(self, text, padding_side="left"):
-        return self.video_context, self.audio_context, None
-
-
 @spaces.GPU(duration=120, size='xlarge')
 def generate_video(
     input_image,
@@ -178,7 +168,9 @@ def generate_video(
 
             embeddings = torch.load(embedding_path)
             video_context = embeddings["video_context"].to("cuda")
-            audio_context = embeddings["audio_context"].to("cuda")
+            audio_context = embeddings["audio_context"]
+            if audio_context is not None:
+                audio_context = audio_context.to("cuda")
             print("Embeddings loaded successfully")
         except Exception as e:
             raise RuntimeError(
@@ -186,10 +178,15 @@ def generate_video(
                 f"Please ensure {TEXT_ENCODER_SPACE} is running properly."
             )
 
-        # Patch the model_ledger to return a fake text encoder with pre-computed embeddings
-        fake_encoder = PrecomputedTextEncoder(video_context, audio_context)
-        original_text_encoder_fn = pipeline.model_ledger.text_encoder
-        pipeline.model_ledger.text_encoder = lambda: fake_encoder
+        # Monkey-patch encode_prompts to return pre-computed embeddings
+        # instead of loading the text encoder + embeddings processor
+        precomputed = EmbeddingsProcessorOutput(
+            video_encoding=video_context,
+            audio_encoding=audio_context,
+            attention_mask=torch.ones(1, device="cuda"),  # dummy mask
+        )
+        original_encode_prompts = pipeline_helpers.encode_prompts
+        pipeline_helpers.encode_prompts = lambda *args, **kwargs: [precomputed]
 
         try:
             tiling_config = TilingConfig.default()
@@ -218,8 +215,8 @@ def generate_video(
 
             return str(output_path), current_seed
         finally:
-            # Restore original text encoder method
-            pipeline.model_ledger.text_encoder = original_text_encoder_fn
+            # Restore original encode_prompts
+            pipeline_helpers.encode_prompts = original_encode_prompts
 
     except Exception as e:
         import traceback