Port encode to new LTX-2 GemmaTextEncoder API

app.py

@@ -61,17 +61,19 @@ model_ledger = ModelLedger(
     gemma_root_path=gemma_local_path,
 )
 
-# Load text encoder once and keep it in memory
+# Load text encoder and embeddings processor once and keep them in memory
 text_encoder = model_ledger.text_encoder().to("cuda")
+embeddings_processor = model_ledger.gemma_embeddings_processor().to("cuda")
 
 print("=" * 80)
 print("Text encoder loaded and ready!")
 print("=" * 80)
 
-def encode_text_simple(text_encoder, prompt: str):
-    """Simple text encoding without using pipeline_utils."""
-    v_context, a_context, _ = text_encoder(prompt)
-    return v_context, a_context
+def encode_text_simple(prompt: str):
+    """Encode a prompt: Gemma hidden states -> embeddings processor."""
+    hidden_states, attention_mask = text_encoder.encode(prompt)
+    out = embeddings_processor.process_hidden_states(hidden_states, attention_mask)
+    return out.video_encoding, out.audio_encoding
 
 @spaces.GPU()
 def encode_prompt(
@@ -125,13 +127,13 @@ def encode_prompt(
                 )
 
         # Encode the positive prompt using the pre-loaded text encoder
-        video_context, audio_context = encode_text_simple(text_encoder, final_prompt)
+        video_context, audio_context = encode_text_simple(final_prompt)
 
         # Encode negative prompt if provided
         video_context_negative = None
         audio_context_negative = None
         if negative_prompt:
-            video_context_negative, audio_context_negative = encode_text_simple(text_encoder, negative_prompt)
+            video_context_negative, audio_context_negative = encode_text_simple(negative_prompt)
 
         # Save embeddings to file
         output_dir = Path("embeddings")
@@ -141,7 +143,7 @@ def encode_prompt(
         # Save embeddings (with negative contexts if provided)
         embedding_data = {
             'video_context': video_context.cpu(),
-            'audio_context': audio_context.cpu(),
+            'audio_context': audio_context.cpu() if audio_context is not None else None,
             'prompt': final_prompt,
             'original_prompt': prompt if enhance_prompt else final_prompt,
         }
@@ -149,7 +151,9 @@ def encode_prompt(
         # Add negative contexts if they were encoded
         if video_context_negative is not None:
             embedding_data['video_context_negative'] = video_context_negative.cpu()
-            embedding_data['audio_context_negative'] = audio_context_negative.cpu()
+            embedding_data['audio_context_negative'] = (
+                audio_context_negative.cpu() if audio_context_negative is not None else None
+            )
             embedding_data['negative_prompt'] = negative_prompt
 
         torch.save(embedding_data, output_path)