some improvements

app.py

@@ -9,6 +9,9 @@ import numpy as np
 import random
 import spaces
 import gradio as gr
+from gradio_client import Client, handle_file
+import torch
+from pathlib import Path
 from typing import Optional
 from huggingface_hub import hf_hub_download
 from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
@@ -33,11 +36,13 @@ DEFAULT_PROMPT = "An astronaut hatches from a fragile egg on the surface of the
 
 # HuggingFace Hub defaults
 DEFAULT_REPO_ID = "Lightricks/LTX-2"
-DEFAULT_GEMMA_REPO_ID = "google/gemma-3-12b-it-qat-q4_0-unquantized"
 DEFAULT_CHECKPOINT_FILENAME = "ltx-2-19b-dev-fp8.safetensors"
 DEFAULT_DISTILLED_LORA_FILENAME = "ltx-2-19b-distilled-lora-384.safetensors"
 DEFAULT_SPATIAL_UPSAMPLER_FILENAME = "ltx-2-spatial-upscaler-x2-1.0.safetensors"
 
+# Text encoder space URL
+TEXT_ENCODER_SPACE = "linoyts/gemma-text-encoder"
+
 def get_hub_or_local_checkpoint(repo_id: Optional[str] = None, filename: Optional[str] = None):
     """Download from HuggingFace Hub or use local checkpoint."""
     if repo_id is None and filename is None:
@@ -68,28 +73,40 @@ print(f"Initializing pipeline with:")
 print(f"  checkpoint_path={checkpoint_path}")
 print(f"  distilled_lora_path={distilled_lora_path}")
 print(f"  spatial_upsampler_path={spatial_upsampler_path}")
-print(f"  gemma_root={DEFAULT_GEMMA_REPO_ID}")
+print(f"  text_encoder_space={TEXT_ENCODER_SPACE}")
 
+# Initialize pipeline WITHOUT text encoder (gemma_root=None)
+# Text encoding will be done by external space
 pipeline = TI2VidTwoStagesPipeline(
     checkpoint_path=checkpoint_path,
     distilled_lora_path=distilled_lora_path,
     distilled_lora_strength=DEFAULT_LORA_STRENGTH,
     spatial_upsampler_path=spatial_upsampler_path,
-    gemma_root=DEFAULT_GEMMA_REPO_ID,
+    gemma_root=None,
     loras=[],
     fp8transformer=False,
     local_files_only=False
 )
 
+# Initialize text encoder client
+print(f"Connecting to text encoder space: {TEXT_ENCODER_SPACE}")
+try:
+    text_encoder_client = Client(TEXT_ENCODER_SPACE)
+    print("✓ Text encoder client connected!")
+except Exception as e:
+    print(f"⚠ Warning: Could not connect to text encoder space: {e}")
+    text_encoder_client = None
+
 @spaces.GPU(duration=300)
 def generate_video(
     input_image,
     prompt: str,
     duration: float,
+    enhance_prompt: bool = True,
     negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
     seed: int = 42,
     randomize_seed: bool = True,
-    num_inference_steps: int = DEFAULT_NUM_INFERENCE_STEPS,
+    num_inference_steps: int = 25,
     cfg_guidance_scale: float = DEFAULT_CFG_GUIDANCE_SCALE,
     height: int = DEFAULT_HEIGHT,
     width: int = DEFAULT_WIDTH,
@@ -107,20 +124,64 @@ def generate_video(
         # Create output directory if it doesn't exist
         output_dir = Path("outputs")
         output_dir.mkdir(exist_ok=True)
-        output_path = output_dir / f"video_{seed}.mp4"
+        output_path = output_dir / f"video_{current_seed}.mp4"
 
         # Handle image input
         images = []
+        temp_image_path = None  # Initialize to None
         if input_image is not None:
             # Save uploaded image temporarily
-            temp_image_path = output_dir / f"temp_input_{seed}.jpg"
+            temp_image_path = output_dir / f"temp_input_{current_seed}.jpg"
             if hasattr(input_image, 'save'):
                 input_image.save(temp_image_path)
             else:
                 # If it's a file path already
-                temp_image_path = input_image
+                temp_image_path = Path(input_image)
             # Format: (image_path, frame_idx, strength)
             images = [(str(temp_image_path), 0, 1.0)]
+        # Get embeddings from text encoder space
+        print(f"Encoding prompt: {prompt}")
+        
+        if text_encoder_client is None:
+            raise RuntimeError(
+                f"Text encoder client not connected. Please ensure the text encoder space "
+                f"({TEXT_ENCODER_SPACE}) is running and accessible."
+            )
+        
+        try:
+            # Prepare image for upload if it exists
+            image_input = None
+            if temp_image_path is not None:
+                image_input = handle_file(str(temp_image_path))
+            
+            result = text_encoder_client.predict(
+                prompt=prompt,
+                enhance_prompt=enhance_prompt,
+                input_image=image_input,
+                seed=current_seed,
+                negative_prompt=negative_prompt,
+                api_name="/encode_prompt"
+            )
+            embedding_path = result[0]  # Path to .pt file
+            print(f"Embeddings received from: {embedding_path}")
+
+            # Load embeddings
+            embeddings = torch.load(embedding_path)
+            video_context_positive = embeddings['video_context']
+            audio_context_positive = embeddings['audio_context']
+
+            # Load negative contexts if available
+            video_context_negative = embeddings.get('video_context_negative', None)
+            audio_context_negative = embeddings.get('audio_context_negative', None)
+
+            print("✓ Embeddings loaded successfully")
+            if video_context_negative is not None:
+                print("  ✓ Negative prompt embeddings also loaded")
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to get embeddings from text encoder space: {e}\n"
+                f"Please ensure {TEXT_ENCODER_SPACE} is running properly."
+            )
 
         # Run inference - progress automatically tracks tqdm from pipeline
         pipeline(
@@ -136,6 +197,10 @@ def generate_video(
             cfg_guidance_scale=cfg_guidance_scale,
             images=images,
             tiling_config=TilingConfig.default(),
+            video_context_positive=video_context_positive,
+            audio_context_positive=audio_context_positive,
+            video_context_negative=video_context_negative,
+            audio_context_negative=audio_context_negative,
         )
 
         return str(output_path), current_seed
@@ -166,13 +231,18 @@ with gr.Blocks(title="LTX-2 Video 🎥🔈") as demo:
                 placeholder="Describe the motion and animation you want..."
             )
 
-            duration = gr.Slider(
-                label="Duration (seconds)",
-                minimum=1.0,
-                maximum=10.0,
-                value=3.0,
-                step=0.1
-            )
+            with gr.Row():
+                duration = gr.Slider(
+                    label="Duration (seconds)",
+                    minimum=1.0,
+                    maximum=10.0,
+                    value=3.0,
+                    step=0.1
+                )
+                enhance_prompt = gr.Checkbox(
+                        label="Enhance Prompt",
+                        value=True
+                    )
 
             generate_btn = gr.Button("Generate Video", variant="primary")
 
@@ -200,7 +270,7 @@ with gr.Blocks(title="LTX-2 Video 🎥🔈") as demo:
                     label="Inference Steps",
                     minimum=1,
                     maximum=100,
-                    value=DEFAULT_NUM_INFERENCE_STEPS,
+                    value=25,
                     step=1
                 )
 
@@ -233,6 +303,7 @@ with gr.Blocks(title="LTX-2 Video 🎥🔈") as demo:
             input_image,
             prompt,
             duration,
+            enhance_prompt,
             negative_prompt,
             seed,
             randomize_seed,
@@ -254,8 +325,8 @@ with gr.Blocks(title="LTX-2 Video 🎥🔈") as demo:
             ],
             [
                 "kill_bill.jpeg",
-                "A low, subsonic drone pulses as Uma Thurman's character, Beatrix Kiddo, holds her razor-sharp katana blade steady in the cinematic lighting. A faint electrical hum fills the silence. Suddenly, accompanied by a deep metallic groan, the polished steel begins to soften and distort, like heated metal starting to lose its structural integrity. Discordant strings swell as the blade's perfect edge slowly warps and droops, molten steel beginning to flow downward in silvery rivulets while maintaining its metallic sheen—each drip producing a wet, viscous stretching sound. The transformation starts subtly at first—a slight bend in the blade—then accelerates as the metal becomes increasingly fluid, the groaning intensifying. The camera holds steady on her face as her piercing eyes gradually narrow, not with lethal focus, but with confusion and growing alarm as she watches her weapon dissolve before her eyes. She whispers under her breath, voice flat with disbelief: 'Wait, what?' Her heartbeat rises in the mix—thump... thump-thump—as her breathing quickens slightly while she witnesses this impossible transformation. Sharp violin stabs punctuate each breath. The melting intensifies, the katana's perfect form becoming increasingly abstract, dripping like liquid mercury from her grip. Molten droplets fall to the ground with soft, bell-like pings. Unintelligible whispers fade in and out as her expression shifts from calm readiness to bewilderment and concern, her heartbeat now pounding like a war drum, as her legendary instrument of vengeance literally liquefies in her hands, leaving her defenseless and disoriented. All sound cuts to silence—then a single devastating bass drop as the final droplet falls, leaving only her unsteady breathing in the dark.",
-                5.0,
+                "A low, subsonic drone pulses as Uma Thurman's character, Beatrix Kiddo, holds her razor-sharp katana blade steady in the cinematic lighting. A faint electrical hum fills the silence. Suddenly, accompanied by a deep metallic groan, the polished steel begins to soften and distort, like heated metal starting to lose its structural integrity. Discordant strings swell as the blade's perfect edge slowly warps and droops, molten steel beginning to flow downward in silvery rivulets while maintaining its metallic sheen—each drip producing a wet, viscous stretching sound. The transformation starts subtly at first—a slight bend in the blade—then accelerates as the metal becomes increasingly fluid, the groaning intensifying. The camera holds steady on her face as her piercing eyes gradually narrow, not with lethal focus, but with confusion and growing alarm as she watches her weapon dissolve before her eyes. Her heartbeat rises in the mix—thump... thump-thump—as her breathing quickens slightly while she witnesses this impossible transformation. Sharp violin stabs punctuate each breath. The melting intensifies, the katana's perfect form becoming increasingly abstract, dripping like liquid mercury from her grip. Molten droplets fall to the ground with soft, bell-like pings. Unintelligible whispers fade in and out as her expression shifts from calm readiness to bewilderment and concern, her heartbeat now pounding like a war drum, as her legendary instrument of vengeance literally liquefies in her hands, leaving her defenseless and disoriented. All sound cuts to silence—then a single devastating bass drop as the final droplet falls, leaving only her unsteady breathing in the dark.",
+                3.0,
             ]
         ],
         fn=generate_video,