Update app.py

app.py

@@ -14,7 +14,7 @@ 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
+from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
 from ltx_core.tiling import TilingConfig
 from ltx_pipelines.constants import (
     DEFAULT_SEED,
@@ -31,8 +31,8 @@ MAX_SEED = np.iinfo(np.int32).max
 # Custom negative prompt
 DEFAULT_NEGATIVE_PROMPT = "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static"
 
-# Default prompt from docstring example
-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. 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. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot."
+# Default prompt for keyframe interpolation
+DEFAULT_PROMPT = "Smooth cinematic transition between keyframes with natural motion and consistent lighting"
 
 # HuggingFace Hub defaults
 DEFAULT_REPO_ID = "Lightricks/LTX-2"
@@ -42,6 +42,8 @@ DEFAULT_SPATIAL_UPSAMPLER_FILENAME = "ltx-2-spatial-upscaler-x2-1.0.safetensors"
 
 # Text encoder space URL
 TEXT_ENCODER_SPACE = "linoyts/gemma-text-encoder"
+# Image edit space URL
+IMAGE_EDIT_SPACE = "linoyts/Qwen-Image-Edit-2509-Fast"
 
 def get_hub_or_local_checkpoint(repo_id: Optional[str] = None, filename: Optional[str] = None):
     """Download from HuggingFace Hub or use local checkpoint."""
@@ -62,7 +64,7 @@ def get_hub_or_local_checkpoint(repo_id: Optional[str] = None, filename: Optiona
 
 # Initialize pipeline at startup
 print("=" * 80)
-print("Loading LTX-2 2-stage pipeline...")
+print("Loading LTX-2 Keyframe Interpolation pipeline...")
 print("=" * 80)
 
 checkpoint_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_CHECKPOINT_FILENAME)
@@ -77,7 +79,7 @@ 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(
+pipeline = KeyframeInterpolationPipeline(
     checkpoint_path=checkpoint_path,
     distilled_lora_path=distilled_lora_path,
     distilled_lora_strength=DEFAULT_LORA_STRENGTH,
@@ -85,7 +87,6 @@ pipeline = TI2VidTwoStagesPipeline(
     gemma_root=None,
     loras=[],
     fp8transformer=False,
-    local_files_only=False
 )
 
 # Initialize text encoder client
@@ -97,22 +98,72 @@ except Exception as e:
     print(f"⚠ Warning: Could not connect to text encoder space: {e}")
     text_encoder_client = None
 
+# Initialize image edit client
+print(f"Connecting to image edit space: {IMAGE_EDIT_SPACE}")
+try:
+    image_edit_client = Client(IMAGE_EDIT_SPACE)
+    print("✓ Image edit client connected!")
+except Exception as e:
+    print(f"⚠ Warning: Could not connect to image edit space: {e}")
+    image_edit_client = None
+
+def generate_end_frame(start_frame, edit_prompt: str):
+    """Generate an end frame from the start frame using Qwen Image Edit."""
+    try:
+        if start_frame is None:
+            raise gr.Error("Please provide a start frame first")
+
+        if image_edit_client is None:
+            raise gr.Error(
+                f"Image edit client not connected. Please ensure the image edit space "
+                f"({IMAGE_EDIT_SPACE}) is running and accessible."
+            )
+
+        # Save start frame temporarily if needed
+        output_dir = Path("outputs")
+        output_dir.mkdir(exist_ok=True)
+        temp_path = output_dir / f"temp_start_for_edit.jpg"
+
+        if hasattr(start_frame, 'save'):
+            start_frame.save(temp_path)
+            image_input = handle_file(str(temp_path))
+        else:
+            image_input = handle_file(str(start_frame))
+
+        # Call Qwen Image Edit
+        result = image_edit_client.predict(
+            images=image_input,
+            prompt=edit_prompt,
+            api_name="/infer"
+        )
+
+        # Result should be the path to the edited image
+        return result
+
+    except Exception as e:
+        import traceback
+        error_msg = f"Error generating end frame: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
+        raise gr.Error(error_msg)
+
 @spaces.GPU(duration=300)
 def generate_video(
-    input_image,
+    start_frame,
+    end_frame_upload,
+    end_frame_generated,
     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 = 25,
+    num_inference_steps: int = DEFAULT_NUM_INFERENCE_STEPS,
     cfg_guidance_scale: float = DEFAULT_CFG_GUIDANCE_SCALE,
     height: int = DEFAULT_HEIGHT,
     width: int = DEFAULT_WIDTH,
     progress=gr.Progress(track_tqdm=True)
 ):
-    """Generate a video based on the given parameters."""
+    """Generate a video with keyframe interpolation between start and end frames."""
     try:
         # Randomize seed if checkbox is enabled
         current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
@@ -124,36 +175,52 @@ 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_{current_seed}.mp4"
+        output_path = output_dir / f"keyframe_video_{current_seed}.mp4"
 
-        # Handle image input
+        # Handle keyframe inputs - build list of (image_path, frame_idx, strength)
         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_{current_seed}.jpg"
-            if hasattr(input_image, 'save'):
-                input_image.save(temp_image_path)
+        temp_paths = []
+
+        # Determine which end frame to use (uploaded or generated)
+        end_frame = end_frame_generated if end_frame_generated is not None else end_frame_upload
+
+        if start_frame is None and end_frame is None:
+            raise ValueError("Please provide at least one keyframe (start or end frame)")
+
+        # Save start frame (frame index 0) if provided
+        if start_frame is not None:
+            temp_start_path = output_dir / f"temp_start_{current_seed}.jpg"
+            if hasattr(start_frame, 'save'):
+                start_frame.save(temp_start_path)
+            else:
+                temp_start_path = Path(start_frame)
+            temp_paths.append(temp_start_path)
+            images.append((str(temp_start_path), 0, 1.0))
+
+        # Save end frame (last frame index) if provided
+        if end_frame is not None:
+            temp_end_path = output_dir / f"temp_end_{current_seed}.jpg"
+            if hasattr(end_frame, 'save'):
+                end_frame.save(temp_end_path)
             else:
-                # If it's a file path already
-                temp_image_path = Path(input_image)
-            # Format: (image_path, frame_idx, strength)
-            images = [(str(temp_image_path), 0, 1.0)]
+                temp_end_path = Path(end_frame)
+            temp_paths.append(temp_end_path)
+            images.append((str(temp_end_path), num_frames - 1, 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))
-            
+            # Use first available frame for prompt enhancement
+            first_frame_path = temp_paths[0] if temp_paths else None
+            image_input = handle_file(str(first_frame_path)) if first_frame_path else None
+
             result = text_encoder_client.predict(
                 prompt=prompt,
                 enhance_prompt=enhance_prompt,
@@ -170,6 +237,9 @@ def generate_video(
             video_context_positive = embeddings['video_context']
             audio_context_positive = embeddings['audio_context']
 
+            # Get the final prompt that was used (enhanced or original)
+            final_prompt = embeddings.get('prompt', prompt)
+
             # Load negative contexts if available
             video_context_negative = embeddings.get('video_context_negative', None)
             audio_context_negative = embeddings.get('audio_context_negative', None)
@@ -203,50 +273,74 @@ def generate_video(
             audio_context_negative=audio_context_negative,
         )
 
-        return str(output_path), current_seed
+        return str(output_path), final_prompt, current_seed
 
     except Exception as e:
         import traceback
         error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
         print(error_msg)
-        return None
+        return None, f"Error: {str(e)}", current_seed
 
 
 # Create Gradio interface
-with gr.Blocks(title="LTX-2 Video 🎥🔈") as demo:
-    gr.Markdown("# LTX-2 🎥🔈: The First Open Source Audio-Video Model")
-    gr.Markdown("State-of-the-art video & audio generation with Lightricks LTX-2 TI2V. Read more: [[model]](https://huggingface.co/Lightricks/LTX-2), [[code]](https://github.com/Lightricks/LTX-2)")
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(
-                label="Input Image (Optional)",
-                type="pil",
-            )
-
-            prompt = gr.Textbox(
-                label="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 motion and animation you want..."
-            )
+with gr.Blocks(title="LTX-2 Keyframe Interpolation 🎥🔈") as demo:
+    gr.Markdown("# LTX-2 Keyframe Interpolation 🎥🔈")
+    gr.Markdown("Generate smooth video transitions between keyframes with Lightricks LTX-2. Read more: [[model]](https://huggingface.co/Lightricks/LTX-2), [[code]](https://github.com/Lightricks/LTX-2)")
 
-            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
+    with gr.Row(elem_id="general_items"):
+        with gr.Column():
+            with gr.Group(elem_id="group_all"):
+                with gr.Row():
+                    start_frame = gr.Image(
+                        label="Start Frame (Optional)",
+                        type="pil",
                     )
+                    with gr.Tabs():
+                        with gr.Tab("Upload"):
+                            end_frame_upload = gr.Image(
+                                label="End Frame",
+                                type="pil",
+                            )
+
+                        with gr.Tab("Generate"):
+                            gr.Markdown("Generate an end frame from your start frame using AI")
+                            edit_prompt = gr.Textbox(
+                                label="Edit Prompt",
+                                placeholder="Describe the transformation (e.g., '5 seconds later, sunset lighting')",
+                                lines=2,
+                                value="5 seconds in the future"
+                            )
+                            generate_end_btn = gr.Button("Generate End Frame", variant="secondary")
+                            end_frame_generated = gr.Image(
+                                label="Generated End Frame",
+                                type="pil",
+                            )
+
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    info="Describe the motion/transition between frames",
+                    value=DEFAULT_PROMPT,
+                    lines=3,
+                    placeholder="Describe the animation style and motion..."
+                )
+
 
             generate_btn = gr.Button("Generate Video", variant="primary")
 
             with gr.Accordion("Advanced Settings", open=False):
+                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
+                    )
+                
                 negative_prompt = gr.Textbox(
                     label="Negative Prompt",
                     value=DEFAULT_NEGATIVE_PROMPT,
@@ -270,7 +364,7 @@ with gr.Blocks(title="LTX-2 Video 🎥🔈") as demo:
                     label="Inference Steps",
                     minimum=1,
                     maximum=100,
-                    value=25,
+                    value=DEFAULT_NUM_INFERENCE_STEPS,
                     step=1
                 )
 
@@ -296,11 +390,26 @@ with gr.Blocks(title="LTX-2 Video 🎥🔈") as demo:
 
         with gr.Column():
             output_video = gr.Video(label="Generated Video", autoplay=True)
+            final_prompt_output = gr.Textbox(
+                label="Final Prompt Used",
+                lines=5,
+                info="This is the prompt that was used for generation (enhanced if enabled)"
+            )
 
+    # Wire up generate end frame button
+    generate_end_btn.click(
+        fn=generate_end_frame,
+        inputs=[start_frame, edit_prompt],
+        outputs=[end_frame_generated]
+    )
+
+    # Wire up generate video button
     generate_btn.click(
         fn=generate_video,
         inputs=[
-            input_image,
+            start_frame,
+            end_frame_upload,
+            end_frame_generated,
             prompt,
             duration,
             enhance_prompt,
@@ -312,33 +421,22 @@ with gr.Blocks(title="LTX-2 Video 🎥🔈") as demo:
             height,
             width,
         ],
-        outputs=[output_video,seed]
-    )
-
-    # Add example
-    gr.Examples(
-        examples=[
-            [
-                "astronaut.jpg",
-                "An astronaut hatches from a fragile egg on the surface of the Moon, 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. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot.",
-                5.0,
-            ],
-            [
-                "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. 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,
-        inputs=[input_image, prompt, duration],
-        outputs = [output_video,seed],
-        label="Example",
-        cache_examples=True,
-        cache_mode="lazy",
+        outputs=[output_video, final_prompt_output, seed]
     )
 
 css = '''
-.gradio-container .contain{max-width: 1200px !important; margin: 0 auto !important}
+.fillable{max-width: 1100px !important}
+.dark .progress-text {color: white}
+#general_items{margin-top: 2em}
+#group_all{overflow:visible}
+#group_all .styler{overflow:visible}
+#group_tabs .tabitem{padding: 0}
+.tab-wrapper{margin-top: 0px;z-index: 999;position: absolute;width: 100%;background-color: var(--block-background-fill);padding: 0;}
+#component-9-button{width: 50%;justify-content: center}
+#component-11-button{width: 50%;justify-content: center}
+#or_item{text-align: center; padding-top: 1em; padding-bottom: 1em; font-size: 1.1em;margin-left: .5em;margin-right: .5em;width: calc(100% - 1em)}
+#fivesec{margin-top: 5em;margin-left: .5em;margin-right: .5em;width: calc(100% - 1em)}
 '''
+
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Citrus())
+    demo.launch(theme=gr.themes.Citrus(), css=css, share=True)