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, DISTILLED_SIGMA_VALUES
from diffusers.utils import load_video, encode_video
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

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

# Distilled speed LoRA, stacked on the dev transformer. Mirrors the validated
# ComfyUI workflow (UNETLoader = dev + LoraLoaderModelOnly = distill @ 0.6),
# which samples at cfg 1 over DISTILLED_SIGMA_VALUES.
DISTILL_LORA_REPO = "Kijai/LTX2.3_comfy"
DISTILL_LORA_WEIGHT_NAME = (
    "loras/ltx-2.3-22b-distilled-1.1_lora-dynamic_fro09_avg_rank_111_bf16.safetensors"
)
DISTILL_LORA_SCALE = 0.7

# 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 = len(DISTILLED_SIGMA_VALUES)  # 8
DEFAULT_GUIDANCE = 1.0
DEFAULT_LORA_SCALE = 1.5
DEFAULT_DURATION = 5.0


def _normalize_dynamic_rank_lora(state_dict):
    """Pre-compensate for the diffusers LTX-2 loader dropping per-layer alphas.

    `LTX2LoraLoaderMixin.load_lora_into_transformer` passes `network_alphas=None`,
    so `get_peft_kwargs` falls back to `lora_alpha = <rank of the first lora_B>`
    with an empty `alpha_pattern`. peft then scales every layer by
    `lora_alpha / rank`. That is harmless for a fixed-rank LoRA, but this one is
    rank-dynamic (3..384), so only ~1% of its 1660 layers would land on the
    intended scale.

    Every `.alpha` in the file equals its own layer's rank, i.e. the intended
    scale is exactly 1.0 everywhere. Multiplying `lora_B` by `rank / lora_alpha`
    makes peft's division cancel out. Computed from the same dict object that is
    handed to diffusers, so the `lora_alpha` pick matches whatever ordering it sees.
    """
    state_dict = {k: v for k, v in state_dict.items() if not k.endswith(".alpha")}
    ranks = {k: v.shape[1] for k, v in state_dict.items() if k.endswith("lora_B.weight")}
    if not ranks:
        return state_dict
    lora_alpha = next(iter(ranks.values()))
    return {
        k: (v.float() * (ranks[k] / lora_alpha)).to(v.dtype) if k in ranks else v
        for k, v in state_dict.items()
    }


# ─── 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")
print("[MODEL] Loading distill LoRA from", DISTILL_LORA_REPO)
_distill_sd = _normalize_dynamic_rank_lora(
    load_file(hf_hub_download(DISTILL_LORA_REPO, DISTILL_LORA_WEIGHT_NAME))
)
pipe.load_lora_weights(_distill_sd, adapter_name="distill")
del _distill_sd
pipe.set_adapters(["crossview", "distill"], [DEFAULT_LORA_SCALE, DISTILL_LORA_SCALE])
print("[MODEL] Model loaded and LoRAs 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,
    distill_strength=DISTILL_LORA_SCALE,
    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)

    # Distilled sampling: the distill LoRA stacked on the dev transformer, as in
    # the validated ComfyUI workflow. DISTILLED_SIGMA_VALUES is an 8-value
    # schedule, so it only applies at its native step count; any other step count
    # falls back to the scheduler's own spacing.
    pipe.set_adapters(["crossview", "distill"], [lora_scale, float(distill_strength)])
    steps = int(num_steps)
    sigmas = DISTILLED_SIGMA_VALUES if steps == len(DISTILLED_SIGMA_VALUES) else None
    cfg = guidance_scale
    stg = 0.0

    # 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=steps,
        sigmas=sigmas,
        guidance_scale=cfg,
        stg_scale=stg,
        audio_guidance_scale=cfg,
        audio_stg_scale=stg,
        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-Prompt - Video Camera Control by Prompt", css=CUSTOM_CSS) as demo:

    with gr.Column(elem_classes=["fillable"]):
        gr.HTML("""
        <div id="header">
            <h1>🎥 LTX CrossView-Prompt - Video Camera Control by Prompt</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):
                    distill_strength = gr.Slider(
                        minimum=0.0, maximum=1.0, value=DISTILL_LORA_SCALE, step=0.05,
                        label="Distill LoRA strength",
                        info="0.6 matches the ComfyUI workflow. Practitioners often "
                             "run 0.25-0.5 — a stronger distill trades away motion, "
                             "and a weaker one usually wants more steps.",
                    )
                    duration_seconds = gr.Slider(
                        minimum=1, maximum=5, value=DEFAULT_DURATION, 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",
                        info="8 uses the distilled sigma schedule. Any other value "
                             "falls back to the scheduler's default spacing.",
                    )
                    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_car_snow.mp4", "far to the right", "higher", "further", 5, 42, False, 1.5, 1.0, 8, "", 0.7],
            ["assets/ref_armor_dusk.mp4", "far to the left", "higher", "further", 5, 42, False, 1.5, 1.0, 8, "", 0.7],
            ["assets/ref_blacksmith.mp4", "slightly to the right", "higher", "same distance", 5, 7517777301363702000, False, 1.5, 1.0, 8, "", 0.7],
        ]
        gr.Examples(
            examples=examples,
            inputs=[reference_video, azimuth, elevation, distance, duration_seconds, seed, randomize_seed, lora_scale, guidance_scale, num_steps, negative_prompt, distill_strength],
            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, distill_strength],
        outputs=[output_video, used_seed, used_prompt],
    )

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