Update app.py

app.py

@@ -14,7 +14,6 @@ if not os.path.exists(LTX_REPO_DIR):
     print(f"Cloning {LTX_REPO_URL}...")
     subprocess.run(["git", "clone", "--depth", "1", LTX_REPO_URL, LTX_REPO_DIR], check=True)
 
-# Install ltx-core and ltx-pipelines if not already installed
 try:
     import ltx_pipelines  # noqa: F401
 except ImportError:
@@ -32,6 +31,7 @@ sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src"))
 import logging
 import random
 import tempfile
+from pathlib import Path
 
 import torch
 torch._dynamo.config.suppress_errors = True
@@ -40,7 +40,8 @@ torch._dynamo.config.disable = True
 import spaces
 import gradio as gr
 import numpy as np
-from huggingface_hub import hf_hub_download, snapshot_download
+from gradio_client import Client, handle_file
+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
@@ -53,118 +54,178 @@ logging.getLogger().setLevel(logging.INFO)
 MAX_SEED = np.iinfo(np.int32).max
 DEFAULT_PROMPT = (
     "An astronaut hatches from a fragile egg on the surface of the Moon, "
-    "the shell cracking and peeling apart in gentle low-gravity motion."
+    "the shell cracking and peeling apart in gentle low-gravity motion. "
+    "Fine lunar dust lifts and drifts outward with each movement, floating "
+    "in slow arcs before settling back onto the ground. The astronaut pushes "
+    "free in a deliberate, weightless motion, small fragments of the egg "
+    "tumbling and spinning through the air."
 )
 DEFAULT_HEIGHT = 1024
 DEFAULT_WIDTH = 1536
 DEFAULT_FRAME_RATE = 24.0
 
-# Download models from Hugging Face
+# Model repo
 LTX_MODEL_REPO = "diffusers-internal-dev/ltx-23"
-GEMMA_MODEL_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"
 
+# Text encoder space URL - must be a 2.3-compatible text encoder
+TEXT_ENCODER_SPACE = "multimodalart/gemma-text-encoder-ltx23"
+
+# Download model checkpoints
 print("=" * 80)
-print("Downloading models from Hugging Face...")
+print("Downloading LTX-2.3 distilled model...")
 print("=" * 80)
 
-DISTILLED_CHECKPOINT = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-22b-distilled.safetensors")
-SPATIAL_UPSAMPLER = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-spatial-upscaler-x2-1.0.safetensors")
-GEMMA_ROOT = snapshot_download(repo_id=GEMMA_MODEL_REPO)
-
-print(f"Distilled checkpoint: {DISTILLED_CHECKPOINT}")
-print(f"Spatial upsampler: {SPATIAL_UPSAMPLER}")
-print(f"Gemma root: {GEMMA_ROOT}")
+checkpoint_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-22b-distilled.safetensors")
+spatial_upsampler_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-spatial-upscaler-x2-1.0.safetensors")
 
-# Initialize pipeline
-print("=" * 80)
-print("Loading LTX-2.3 Distilled pipeline...")
-print("=" * 80)
+print(f"Checkpoint: {checkpoint_path}")
+print(f"Spatial upsampler: {spatial_upsampler_path}")
 
+# Initialize pipeline WITHOUT text encoder (gemma_root=None)
+# Text encoding will be done by external space
 pipeline = DistilledPipeline(
-    distilled_checkpoint_path=DISTILLED_CHECKPOINT,
-    spatial_upsampler_path=SPATIAL_UPSAMPLER,
-    gemma_root=GEMMA_ROOT,
+    distilled_checkpoint_path=checkpoint_path,
+    spatial_upsampler_path=spatial_upsampler_path,
+    gemma_root=None,
     loras=[],
     quantization=QuantizationPolicy.fp8_cast(),
 )
 
-# Preload all models so first request is fast.
-# On ZeroGPU, .to('cuda') is intercepted and actual GPU allocation
-# happens inside the @spaces.GPU decorated function.
-print("Preloading models...")
-ledger = pipeline.model_ledger
-_text_encoder = ledger.text_encoder()
-_transformer = ledger.transformer()
-_video_encoder = ledger.video_encoder()
-_video_decoder = ledger.video_decoder()
-_audio_decoder = ledger.audio_decoder()
-_vocoder = ledger.vocoder()
-_spatial_upsampler = ledger.spatial_upsampler()
-
-ledger.text_encoder = lambda: _text_encoder
-ledger.transformer = lambda: _transformer
-ledger.video_encoder = lambda: _video_encoder
-ledger.video_decoder = lambda: _video_decoder
-ledger.audio_decoder = lambda: _audio_decoder
-ledger.vocoder = lambda: _vocoder
-ledger.spatial_upsampler = lambda: _spatial_upsampler
-
-print("All models preloaded!")
+# Connect to text encoder space
+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
+
+print("=" * 80)
+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')
-@torch.inference_mode()
 def generate_video(
     input_image,
     prompt: str,
     duration: float,
-    enhance_prompt: bool,
-    seed: int,
-    randomize_seed: bool,
-    height: int,
-    width: int,
+    enhance_prompt: bool = True,
+    seed: int = 42,
+    randomize_seed: bool = True,
+    height: int = DEFAULT_HEIGHT,
+    width: int = DEFAULT_WIDTH,
     progress=gr.Progress(track_tqdm=True),
 ):
-    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
-    num_frames = int(duration * DEFAULT_FRAME_RATE) + 1
-    num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1
-
-    images = []
-    if input_image is not None:
-        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
-            temp_path = f.name
+    """Generate a video based on the given parameters."""
+    try:
+        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+
+        frame_rate = DEFAULT_FRAME_RATE
+        num_frames = int(duration * frame_rate) + 1
+        # 8k+1 format
+        num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1
+
+        # Handle image input
+        images = []
+        temp_image_path = None
+        if input_image is not None:
+            output_dir = Path("outputs")
+            output_dir.mkdir(exist_ok=True)
+            temp_image_path = output_dir / f"temp_input_{current_seed}.jpg"
             if hasattr(input_image, "save"):
-                input_image.save(temp_path)
+                input_image.save(temp_image_path)
             else:
-                from shutil import copy2
-                copy2(str(input_image), temp_path)
-            images = [ImageConditioningInput(path=temp_path, frame_idx=0, strength=1.0)]
-
-    tiling_config = TilingConfig.default()
-    video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
-
-    video, audio = pipeline(
-        prompt=prompt,
-        seed=current_seed,
-        height=int(height),
-        width=int(width),
-        num_frames=num_frames,
-        frame_rate=DEFAULT_FRAME_RATE,
-        images=images,
-        tiling_config=tiling_config,
-        enhance_prompt=enhance_prompt,
-    )
+                temp_image_path = Path(input_image)
+            images = [ImageConditioningInput(path=str(temp_image_path), frame_idx=0, strength=1.0)]
 
-    output_path = tempfile.mktemp(suffix=".mp4")
-    encode_video(
-        video=video,
-        fps=DEFAULT_FRAME_RATE,
-        audio=audio,
-        output_path=output_path,
-        video_chunks_number=video_chunks_number,
-    )
+        # 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:
+            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="",
+                api_name="/encode_prompt",
+            )
+            embedding_path = result[0]
+            print(f"Embeddings received from: {embedding_path}")
+
+            embeddings = torch.load(embedding_path)
+            video_context = embeddings["video_context"].to("cuda")
+            audio_context = embeddings["audio_context"].to("cuda")
+            print("Embeddings loaded successfully")
+        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."
+            )
 
-    return output_path, current_seed
+        # 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
+
+        try:
+            tiling_config = TilingConfig.default()
+            video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
+
+            video, audio = pipeline(
+                prompt=prompt,
+                seed=current_seed,
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                frame_rate=frame_rate,
+                images=images,
+                tiling_config=tiling_config,
+                enhance_prompt=False,  # Already enhanced by text encoder space
+            )
+
+            output_path = tempfile.mktemp(suffix=".mp4")
+            encode_video(
+                video=video,
+                fps=frame_rate,
+                audio=audio,
+                output_path=output_path,
+                video_chunks_number=video_chunks_number,
+            )
+
+            return str(output_path), current_seed
+        finally:
+            # Restore original text encoder method
+            pipeline.model_ledger.text_encoder = original_text_encoder_fn
+
+    except Exception as e:
+        import traceback
+        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
+        return None, current_seed
 
 
 with gr.Blocks(title="LTX-2.3 Distilled") as demo:
@@ -180,15 +241,16 @@ with gr.Blocks(title="LTX-2.3 Distilled") as demo:
             input_image = gr.Image(label="Input Image (Optional)", type="pil")
             prompt = gr.Textbox(
                 label="Prompt",
-                value=DEFAULT_PROMPT,
+                info="for best results - make it as elaborate as possible",
+                value="Make this image come alive with cinematic motion, smooth animation",
                 lines=3,
-                placeholder="Describe the video you want to generate...",
+                placeholder="Describe the motion and animation you want...",
             )
             with gr.Row():
-                duration = gr.Slider(label="Duration (seconds)", minimum=1.0, maximum=10.0, value=5.0, step=0.5)
+                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")
+            generate_btn = gr.Button("Generate Video", variant="primary", size="lg")
 
             with gr.Accordion("Advanced Settings", open=False):
                 seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=10, step=1)
@@ -210,5 +272,9 @@ with gr.Blocks(title="LTX-2.3 Distilled") as demo:
     )
 
 
+css = """
+.gradio-container .contain{max-width: 1200px !important; margin: 0 auto !important}
+"""
+
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(theme=gr.themes.Citrus(), css=css)