Update app.py

app.py

@@ -10,7 +10,7 @@ import spaces
 import gradio as gr
 from typing import Optional
 from huggingface_hub import hf_hub_download
-from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
+from ltx_pipelines.distilled import DistilledPipeline
 from ltx_core.tiling import TilingConfig
 from ltx_pipelines.constants import (
     DEFAULT_SEED,
@@ -18,22 +18,15 @@ from ltx_pipelines.constants import (
     DEFAULT_WIDTH,
     DEFAULT_NUM_FRAMES,
     DEFAULT_FRAME_RATE,
-    DEFAULT_NUM_INFERENCE_STEPS,
-    DEFAULT_CFG_GUIDANCE_SCALE,
-    DEFAULT_LORA_STRENGTH,
 )
 
-# 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."
 
 # HuggingFace Hub defaults
 DEFAULT_REPO_ID = "LTX-Colab/LTX-Video-Preview"
 DEFAULT_GEMMA_REPO_ID = "google/gemma-3-12b-it-qat-q4_0-unquantized"
-DEFAULT_CHECKPOINT_FILENAME = "ltx-2-19b-dev-rc1.safetensors"
-DEFAULT_DISTILLED_LORA_FILENAME = "ltx-2-19b-distilled-lora-384-rc1.safetensors"
+DEFAULT_CHECKPOINT_FILENAME = "ltx-2-19b-distilled-rc1.safetensors"
 DEFAULT_SPATIAL_UPSAMPLER_FILENAME = "ltx-2-spatial-upscaler-x2-1.0-rc1.safetensors"
 
 def get_hub_or_local_checkpoint(repo_id: Optional[str] = None, filename: Optional[str] = None):
@@ -55,73 +48,36 @@ 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 Distilled pipeline...")
 print("=" * 80)
 
 checkpoint_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_CHECKPOINT_FILENAME)
-distilled_lora_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_DISTILLED_LORA_FILENAME)
 spatial_upsampler_path = get_hub_or_local_checkpoint(DEFAULT_REPO_ID, DEFAULT_SPATIAL_UPSAMPLER_FILENAME)
 
 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}")
 
-pipeline = TI2VidTwoStagesPipeline(
+pipeline = DistilledPipeline(
     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,
     loras=[],
     fp8transformer=False,
-    local_files_only=False
 )
 
-# print("=" * 80)
-# print("Warming up pipeline (loading Gemma text encoder)...")
-# print("=" * 80)
-
-# # Do a dummy warmup to load all models including Gemma
-# import tempfile
-# import os
-# warmup_output = tempfile.mktemp(suffix=".mp4")
-# try:
-#     pipeline(
-#         prompt="warmup",
-#         negative_prompt="",
-#         output_path=warmup_output,
-#         seed=42,
-#         height=256,
-#         width=256,
-#         num_frames=9,
-#         frame_rate=8,
-#         num_inference_steps=1,
-#         cfg_guidance_scale=1.0,
-#         images=[],
-#         tiling_config=TilingConfig.default(),
-#     )
-#     # Clean up warmup output
-#     if os.path.exists(warmup_output):
-#         os.remove(warmup_output)
-# except Exception as e:
-#     print(f"Warmup completed with note: {e}")
-
-# print("=" * 80)
-# print("Pipeline fully loaded and ready!")
-# print("=" * 80)
+print("=" * 80)
+print("Pipeline fully loaded and ready!")
+print("=" * 80)
 
 @spaces.GPU(duration=300)
 def generate_video(
     input_image,
     prompt: str,
     duration: float,
-    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
     seed: int = 42,
     randomize_seed: bool = True,
-    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)
@@ -158,15 +114,12 @@ def generate_video(
         # Run inference - progress automatically tracks tqdm from pipeline
         pipeline(
             prompt=prompt,
-            negative_prompt=negative_prompt,
             output_path=str(output_path),
             seed=seed,
             height=height,
             width=width,
             num_frames=num_frames,
             frame_rate=frame_rate,
-            num_inference_steps=num_inference_steps,
-            cfg_guidance_scale=cfg_guidance_scale,
             images=images,
             tiling_config=TilingConfig.default(),
         )
@@ -181,8 +134,8 @@ def generate_video(
 
 
 # Create Gradio interface
-with gr.Blocks(title="LTX-2 Image-to-Video") as demo:
-    gr.Markdown("# LTX-2 Image-to-Video Generation")
+with gr.Blocks(title="LTX-2 Distilled Image-to-Video") as demo:
+    gr.Markdown("# LTX-2 Distilled Image-to-Video Generation")
 
     with gr.Row():
         with gr.Column():
@@ -210,12 +163,6 @@ with gr.Blocks(title="LTX-2 Image-to-Video") as demo:
             generate_btn = gr.Button("Generate Video", variant="primary", size="lg")
 
             with gr.Accordion("Advanced Settings", open=False):
-                negative_prompt = gr.Textbox(
-                    label="Negative Prompt",
-                    value=DEFAULT_NEGATIVE_PROMPT,
-                    lines=2
-                )
-
                 seed = gr.Slider(
                     label="Seed",
                     minimum=0,
@@ -229,22 +176,6 @@ with gr.Blocks(title="LTX-2 Image-to-Video") as demo:
                     value=True
                 )
 
-                num_inference_steps = gr.Slider(
-                    label="Inference Steps",
-                    minimum=1,
-                    maximum=100,
-                    value=DEFAULT_NUM_INFERENCE_STEPS,
-                    step=1
-                )
-
-                cfg_guidance_scale = gr.Slider(
-                    label="CFG Guidance Scale",
-                    minimum=1.0,
-                    maximum=10.0,
-                    value=DEFAULT_CFG_GUIDANCE_SCALE,
-                    step=0.1
-                )
-
                 with gr.Row():
                     width = gr.Number(
                         label="Width",
@@ -266,11 +197,8 @@ with gr.Blocks(title="LTX-2 Image-to-Video") as demo:
             input_image,
             prompt,
             duration,
-            negative_prompt,
             seed,
             randomize_seed,
-            num_inference_steps,
-            cfg_guidance_scale,
             height,
             width,
         ],
@@ -290,10 +218,9 @@ with gr.Blocks(title="LTX-2 Image-to-Video") as demo:
         inputs=[input_image, prompt, duration],
         outputs = [output_video],
         label="Example",
-        cache_examples=True,
-        cache_mode="lazy",
+        cache_examples=False,
     )
 
 
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Citrus())
+    demo.launch(theme=gr.themes.Citrus())