app.py

import spaces

import gradio as gr
import torch
import numpy as np
import os
import tempfile
from pathlib import Path

from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2InContextPipeline, LTX2ReferenceCondition
from diffusers.pipelines.ltx2.utils import DEFAULT_NEGATIVE_PROMPT
from diffusers.utils import load_video, encode_video

# ─── Constants ───────────────────────────────────────────────────────────────

BASE_MODEL = "dg845/LTX-2.3-Diffusers"
LORA_REPO = "Cseti/LTX2.3-22B_IC-LoRA-CrossView-Prompt"
LORA_WEIGHT_NAME = "LTX2.3-22B_IC-LoRA-CrossView-Prompt_v0.9_13700.safetensors"
TRIGGER_WORD = "crossview."

# Camera vocabulary (from captions_all_63.txt)
AZIMUTH_CHOICES = [
    "same angle",
    "slightly to the left",
    "slightly to the right",
    "to the left",
    "to the right",
    "far to the left",
    "far to the right",
]
ELEVATION_CHOICES = ["lower", "same height", "higher"]
DISTANCE_CHOICES = ["closer", "same distance", "further"]

# Training resolution: 768x768x81 @ 15fps
DEFAULT_WIDTH = 768
DEFAULT_HEIGHT = 512
DEFAULT_NUM_FRAMES = 81
DEFAULT_FPS = 24
DEFAULT_STEPS = 30
DEFAULT_GUIDANCE = 3.0
DEFAULT_LORA_SCALE = 1.0

# ─── Model loading (module scope) ────────────────────────────────────────────

print("[MODEL] Loading LTX2InContextPipeline from", BASE_MODEL)
pipe = LTX2InContextPipeline.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")
print("[MODEL] Loading LoRA weights from", LORA_REPO)
pipe.load_lora_weights(LORA_REPO, weight_name=LORA_WEIGHT_NAME, adapter_name="crossview")
pipe.set_adapters("crossview", DEFAULT_LORA_SCALE)
print("[MODEL] Model loaded and LoRA applied.")


# ─── Helpers ──────────────────────────────────────────────────────────────────

def build_prompt(azimuth, elevation, distance):
    """Build the camera-angle prompt from the discrete vocabulary."""
    return f"{TRIGGER_WORD} new camera angle: {azimuth}, {elevation}, {distance}."


def num_frames_for_duration(seconds, fps=DEFAULT_FPS, base=8):
    raw = seconds * fps
    return ((int(raw) - 1) // base) * base + 1


# ─── Inference ─────────────────────────────────────────────────────────────────

@spaces.GPU(duration=300, size="xlarge")
def generate(
    reference_video,
    azimuth,
    elevation,
    distance,
    duration_seconds,
    seed,
    randomize_seed,
    lora_scale,
    guidance_scale,
    num_steps,
    negative_prompt,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a new camera-angle view from a reference video."""
    if reference_video is None:
        raise gr.Error("Please upload a reference video first.")

    # Seed handling
    if randomize_seed:
        seed = torch.randint(0, 2**63 - 1, (1,)).item()
    generator = torch.Generator("cuda").manual_seed(seed)

    # Build prompt from vocabulary
    prompt = build_prompt(azimuth, elevation, distance)

    # Compute frames from duration
    num_frames = num_frames_for_duration(duration_seconds, DEFAULT_FPS)

    # Load the reference video
    ref_frames = load_video(reference_video)
    ref_cond = LTX2ReferenceCondition(frames=ref_frames, strength=1.0)

    # Update LoRA scale
    pipe.set_adapters("crossview", lora_scale)

    # Run inference (return_dict=False gives (video, audio) tuple)
    video, audio = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt if negative_prompt else DEFAULT_NEGATIVE_PROMPT,
        reference_conditions=ref_cond,
        width=DEFAULT_WIDTH,
        height=DEFAULT_HEIGHT,
        num_frames=num_frames,
        frame_rate=DEFAULT_FPS,
        num_inference_steps=num_steps,
        guidance_scale=guidance_scale,
        stg_scale=1.0,
        audio_guidance_scale=guidance_scale,
        audio_stg_scale=1.0,
        generator=generator,
        output_type="np",
        return_dict=False,
    )

    # Export to video (video[0] is the first batch element)
    tmp_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False, dir="/tmp")
    tmp_path.close()

    if audio is not None and len(audio) > 0 and audio[0] is not None:
        encode_video(
            video[0],
            fps=DEFAULT_FPS,
            output_path=tmp_path.name,
            audio=audio[0].float().cpu(),
            audio_sample_rate=pipe.vocoder.config.output_sampling_rate,
        )
    else:
        encode_video(
            video[0],
            fps=DEFAULT_FPS,
            output_path=tmp_path.name,
        )

    return tmp_path.name, seed, prompt


# ─── UI ───────────────────────────────────────────────────────────────────────

CUSTOM_CSS = """
#header { text-align: center; margin-bottom: 1rem; }
#header h1 { font-size: 2rem; margin-bottom: 0.25rem; }
#header p { color: var(--body-text-color-subdued); font-size: 0.95rem; }
.fillable { max-width: 1200px !important; margin: auto; }
"""

with gr.Blocks(title="LTX CrossView Camera Control", css=CUSTOM_CSS) as demo:

    with gr.Column(elem_classes=["fillable"]):
        gr.HTML("""
        <div id="header">
            <h1>🎥 LTX-Video 2.3 CrossView Camera Control</h1>
            <p>Upload a reference video and pick a new camera angle. The IC-LoRA re-renders the same scene from a different viewpoint.</p>
        </div>
        """)

        with gr.Row(equal_height=True):
            # ─── Left: Inputs ───
            with gr.Column(scale=1):
                gr.Markdown("### 📹 Reference Video")
                reference_video = gr.Video(
                    label="Reference video",
                    sources=["upload"],
                    format="mp4",
                )

                gr.Markdown("### 🎬 Camera Angle")
                with gr.Row():
                    azimuth = gr.Dropdown(
                        choices=AZIMUTH_CHOICES,
                        value="to the right",
                        label="Azimuth (orbit)",
                        info="Horizontal camera position around the subject",
                    )
                    elevation = gr.Dropdown(
                        choices=ELEVATION_CHOICES,
                        value="lower",
                        label="Elevation (height)",
                        info="Camera height relative to subject",
                    )
                    distance = gr.Dropdown(
                        choices=DISTANCE_CHOICES,
                        value="closer",
                        label="Distance",
                        info="Camera distance to subject",
                    )

                with gr.Accordion("Advanced", open=False):
                    duration_seconds = gr.Slider(
                        minimum=1, maximum=5, value=3, step=0.5,
                        label="Duration (seconds)",
                    )
                    lora_scale = gr.Slider(
                        minimum=0.5, maximum=2.0, value=DEFAULT_LORA_SCALE, step=0.05,
                        label="LoRA strength",
                        info="Higher = stronger viewpoint shift. Model card recommends 1.2–1.5 on distilled.",
                    )
                    guidance_scale = gr.Slider(
                        minimum=1.0, maximum=10.0, value=DEFAULT_GUIDANCE, step=0.5,
                        label="Guidance scale",
                    )
                    num_steps = gr.Slider(
                        minimum=8, maximum=50, value=DEFAULT_STEPS, step=1,
                        label="Inference steps",
                    )
                    negative_prompt = gr.Textbox(
                        value="",
                        label="Negative prompt (empty = use default)",
                        lines=2,
                        placeholder="Leave empty to use the default negative prompt",
                    )
                    seed = gr.Number(value=42, label="Seed", precision=0)
                    randomize_seed = gr.Checkbox(value=True, label="Randomize seed")

                generate_btn = gr.Button("Generate New View", variant="primary", size="lg")

            # ─── Right: Output ───
            with gr.Column(scale=1):
                gr.Markdown("### 🎥 Generated New Camera View")
                output_video = gr.Video(label="Generated video", autoplay=True, format="mp4")
                used_seed = gr.Number(label="Seed used", precision=0, interactive=False)
                used_prompt = gr.Textbox(
                    label="Prompt sent to model",
                    interactive=False,
                    lines=2,
                )

        # ─── Examples ───
        gr.Markdown("---\n### 📋 Examples")
        gr.Markdown("Click an example to populate the inputs, then click **Generate New View**.")
        examples = [
            ["assets/ref_dining.mp4", "to the right", "lower", "closer", 3, 42, False, 1.0, 3.0, 30, ""],
            ["assets/ref_dining.mp4", "to the left", "higher", "further", 3, 123, False, 1.0, 3.0, 30, ""],
            ["assets/ref_scene01.mp4", "slightly to the left", "higher", "closer", 3, 42, False, 1.0, 3.0, 30, ""],
        ]
        gr.Examples(
            examples=examples,
            inputs=[reference_video, azimuth, elevation, distance, duration_seconds, seed, randomize_seed, lora_scale, guidance_scale, num_steps, negative_prompt],
            outputs=[output_video, used_seed, used_prompt],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    # Wire up
    generate_btn.click(
        fn=generate,
        inputs=[reference_video, azimuth, elevation, distance, duration_seconds, seed, randomize_seed, lora_scale, guidance_scale, num_steps, negative_prompt],
        outputs=[output_video, used_seed, used_prompt],
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), show_error=True)