app.py

import sys
from pathlib import Path

# Add packages to Python path
current_dir = Path(__file__).parent
sys.path.insert(0, str(current_dir / "packages" / "ltx-pipelines" / "src"))
sys.path.insert(0, str(current_dir / "packages" / "ltx-core" / "src"))

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_core.tiling import TilingConfig
from ltx_pipelines.constants import (
    DEFAULT_SEED,
    DEFAULT_HEIGHT,
    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_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):
    """Download from HuggingFace Hub or use local checkpoint."""
    if repo_id is None and filename is None:
        raise ValueError("Please supply at least one of `repo_id` or `filename`")

    if repo_id is not None:
        if filename is None:
            raise ValueError("If repo_id is specified, filename must also be specified.")
        print(f"Downloading {filename} from {repo_id}...")
        ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename)
        print(f"Downloaded to {ckpt_path}")
    else:
        ckpt_path = filename

    return ckpt_path


# Initialize pipeline at startup
print("=" * 80)
print("Loading LTX-2 2-stage 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(
    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)

@spaces.GPU(duration=300)
def generate_video(
    input_image,
    prompt: str,
    duration: float,
    negative_prompt: str,
    seed: int,
    randomize_seed: bool,
    num_inference_steps: int,
    cfg_guidance_scale: float,
    height: int,
    width: int,
    progress=gr.Progress()
):
    """Generate a video based on the given parameters."""
    try:
        # Randomize seed if checkbox is enabled
        if randomize_seed:
            import random
            seed = random.randint(0, 1000000)

        # Calculate num_frames from duration (using fixed 24 fps)
        frame_rate = 24.0
        num_frames = int(duration * frame_rate) + 1  # +1 to ensure we meet the duration

        # 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"

        # Handle image input
        images = []
        if input_image is not None:
            # Save uploaded image temporarily
            temp_image_path = output_dir / f"temp_input_{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
            # Format: (image_path, frame_idx, strength)
            images = [(str(temp_image_path), 0, 1.0)]

        # Run inference
        progress(0, desc="Generating video (2-stage)...")
        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(),
        )

        progress(1.0, desc="Done!")
        return str(output_path)

    except Exception as e:
        import traceback
        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
        print(error_msg)
        return None


# Create Gradio interface
with gr.Blocks(title="LTX-2 Image-to-Video") as demo:
    gr.Markdown("# LTX-2 Image-to-Video Generation")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Input Image",
                type="pil",
                sources=["upload"]
            )

            prompt = gr.Textbox(
                label="Prompt",
                value="Make this image come alive with cinematic motion, smooth animation",
                lines=3,
                placeholder="Describe the motion and animation you want..."
            )

            duration = gr.Slider(
                label="Duration (seconds)",
                minimum=1.0,
                maximum=10.0,
                value=5.0,
                step=0.1
            )

            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,
                    maximum=1000000,
                    value=DEFAULT_SEED,
                    step=1
                )

                randomize_seed = gr.Checkbox(
                    label="Randomize Seed",
                    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",
                        value=DEFAULT_WIDTH,
                        precision=0
                    )
                    height = gr.Number(
                        label="Height",
                        value=DEFAULT_HEIGHT,
                        precision=0
                    )

        with gr.Column():
            output_video = gr.Video(label="Generated Video", autoplay=True)

    generate_btn.click(
        fn=generate_video,
        inputs=[
            input_image,
            prompt,
            duration,
            negative_prompt,
            seed,
            randomize_seed,
            num_inference_steps,
            cfg_guidance_scale,
            height,
            width,
        ],
        outputs=output_video
    )

    # Add example
    gr.Examples(
        examples=[
            [
                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/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,
            ]
        ],
        inputs=[input_image, prompt, duration],
        label="Example"
    )


if __name__ == "__main__":
    demo.launch()