app.py

import spaces

import gradio as gr
import torch
import numpy as np
import os
import math
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"

# CrossView IC-LoRA (step 13,700 pilot) — the actual camera-control adapter.
LORA_REPO = "Cseti/LTX2.3-22B_IC-LoRA-CrossView-Prompt"
LORA_WEIGHT_NAME = "LTX2.3-22B_IC-LoRA-CrossView-Prompt_v0.9_13700.safetensors"

# Distilled speed LoRA — matches the canonical ComfyUI workflow, which pairs the
# IC-LoRA with the LTX-2.3 distilled speed LoRA for an 8-step base pass.
# The canonical graph uses Kijai's rank-111 SVD-compressed export, but that file
# carries per-layer `.alpha` tensors that diffusers' LTX2 LoRA converter passes
# through unchanged, tripping the `all("lora" in key)` format gate. The official
# Lightricks rank-384 distillation (same weights, uncompressed) is the clean,
# diffusers-loadable equivalent.
DISTILLED_REPO = "Lightricks/LTX-2.3"
DISTILLED_WEIGHT_NAME = "ltx-2.3-22b-distilled-lora-384-1.1.safetensors"
DISTILLED_SCALE = 0.6  # canonical strength

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. The output resolution is derived from
# the reference video's aspect ratio (below), keeping ~this pixel budget.
BASE_AREA = 768 * 512          # target pixel budget for the derived resolution
DIM_MULTIPLE = 32              # LTX-2 requires spatial dims divisible by 32
DEFAULT_NUM_FRAMES = 81        # native training length; default duration maps here
DEFAULT_FPS = 24
DEFAULT_DURATION = 3.5         # 3.5s @ 24fps -> 81 frames (see num_frames_for_duration)
DEFAULT_GUIDANCE = 1.0         # distilled regime is CFG-free; 1.0 = no extra CFG pass
DEFAULT_LORA_SCALE = 1.5       # IC-LoRA strength (canonical; card recommends 1.2–1.5)

# Canonical ComfyUI base-pass sigma schedule (8 steps). The distilled LoRA was
# trained for this hand-authored schedule; the terminal 0.0 is required so the last
# step fully denoises. NOTE: we intentionally keep the shipped FlowMatchEulerDiscrete
# scheduler. ComfyUI's base sampler is `euler_ancestral`, and diffusers ships a
# matching LTXEulerAncestralRFScheduler — but swapping it in BREAKS the IC-LoRA
# reference conditioning (the reference video is ignored and the output becomes an
# unrelated text-to-video clip). FlowMatch + these explicit sigmas keeps the
# reference intact and is the closest faithful config; the only sacrifice is the
# ancestral noise re-injection, which is cosmetic here.
DISTILLED_SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.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 CrossView IC-LoRA from", LORA_REPO)
pipe.load_lora_weights(LORA_REPO, weight_name=LORA_WEIGHT_NAME, adapter_name="crossview")
print("[MODEL] Loading distilled speed LoRA from", DISTILLED_REPO)
pipe.load_lora_weights(DISTILLED_REPO, weight_name=DISTILLED_WEIGHT_NAME, adapter_name="distilled")
pipe.set_adapters(["crossview", "distilled"], [DEFAULT_LORA_SCALE, DISTILLED_SCALE])
# Fuse both adapters into the base weights, then drop the adapter tensors. This is
# required on ZeroGPU: keeping the 7.6GB distilled adapter live pushes the packed
# model past the offload disk (78GB pack -> "No space left on device"). Fusing bakes
# the deltas in (numerically identical to the live adapters at these scales) so only
# the base weights get packed. Consequence: IC-LoRA strength is fixed at
# DEFAULT_LORA_SCALE — edit the constant and rebuild to change it.
pipe.fuse_lora(components=["transformer"])
pipe.unload_lora_weights()
print(f"[MODEL] Model loaded; LoRAs fused (ic={DEFAULT_LORA_SCALE}, distilled={DISTILLED_SCALE}).")

# ─── 2x spatial latent upsampler (second pass — matches the canonical two-pass) ──
# Encodes the base video to latents, 2x-upsamples them with a learned upsampler,
# and re-decodes at double resolution (no extra diffusion refine, unlike the
# ComfyUI 3-step stage-2, but the same learned latent upscaler weights).
from diffusers import LTX2LatentUpsamplePipeline
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel

UPSAMPLER_REPO = "dg845/LTX-2.3-Spatial-Upsampler-Diffusers"
print("[MODEL] Loading spatial latent upsampler from", UPSAMPLER_REPO)
latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
    UPSAMPLER_REPO, subfolder="latent_upsampler", torch_dtype=torch.bfloat16
)
upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler)
upsample_pipe.vae.enable_tiling()
upsample_pipe.to("cuda")
print("[MODEL] Upsampler ready.")


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


def compute_dimensions(ref_w, ref_h, target_area=BASE_AREA, multiple=DIM_MULTIPLE):
    """Preserve the reference video's aspect ratio at ~target_area pixels,
    rounded to the nearest valid multiple for the LTX-2 VAE."""
    ar = ref_w / ref_h
    w = math.sqrt(target_area * ar)
    h = math.sqrt(target_area / ar)
    w = max(multiple, int(round(w / multiple)) * multiple)
    h = max(multiple, int(round(h / multiple)) * multiple)
    return int(w), int(h)


def _frame_size(frame):
    """Return (width, height) for a PIL image or an HxWxC array."""
    if hasattr(frame, "size"):          # PIL.Image
        return frame.size
    arr = np.asarray(frame)
    return arr.shape[1], arr.shape[0]   # (W, H)


def _resize_frame(frame, w, h):
    if hasattr(frame, "resize"):        # PIL.Image
        return frame.resize((w, h))
    from PIL import Image
    return Image.fromarray(np.asarray(frame)).resize((w, h))


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

@spaces.GPU(duration=420, size="xlarge")
def generate(
    reference_video,
    azimuth,
    elevation,
    distance,
    duration_seconds,
    seed,
    randomize_seed,
    guidance_scale,
    negative_prompt,
    upscale,
    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 and derive output resolution from its aspect ratio
    ref_frames = load_video(reference_video)
    src_w, src_h = _frame_size(ref_frames[0])
    width, height = compute_dimensions(src_w, src_h)
    ref_frames = [_resize_frame(f, width, height) for f in ref_frames]
    ref_cond = LTX2ReferenceCondition(frames=ref_frames, strength=1.0)
    print(f"[GEN] ref {src_w}x{src_h} -> out {width}x{height}, {num_frames} frames, "
          f"{len(DISTILLED_SIGMAS) - 1}-step distilled sigmas, guidance {guidance_scale} "
          f"(LoRAs fused: ic={DEFAULT_LORA_SCALE}, distilled={DISTILLED_SCALE})")

    # Run inference (return_dict=False gives (video, audio) tuple).
    # Guidance is set to match the canonical ComfyUI workflow (CFGGuider=1, no STG,
    # no modality-isolation guidance): at guidance_scale=1.0 every guidance flag is
    # False, so it's a single conditional forward per step — the regime the distilled
    # LoRA was trained for. Raising guidance_scale above 1 re-enables CFG (and makes
    # the negative prompt take effect). STG / modality / audio guidance stay neutral.
    video, audio = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt if negative_prompt else DEFAULT_NEGATIVE_PROMPT,
        reference_conditions=[ref_cond],
        width=width,
        height=height,
        num_frames=num_frames,
        frame_rate=DEFAULT_FPS,
        sigmas=DISTILLED_SIGMAS,
        num_inference_steps=len(DISTILLED_SIGMAS) - 1,
        guidance_scale=guidance_scale,
        stg_scale=0.0,
        modality_scale=1.0,
        audio_guidance_scale=1.0,
        audio_stg_scale=0.0,
        audio_modality_scale=1.0,
        generator=generator,
        output_type="pil" if upscale else "np",
        return_dict=False,
    )

    # Optional 2x spatial upsample pass (canonical second stage). Re-encodes the base
    # frames, upsamples the latents 2x, decodes at double resolution.
    if upscale:
        progress(0.9, desc="2x upscaling")
        print(f"[GEN] 2x upsample {width}x{height} -> {2 * width}x{2 * height}")
        out_frames = upsample_pipe(
            video=video,
            width=width,
            height=height,
            output_type="np",
            return_dict=False,
        )[0][0]
    else:
        out_frames = video[0]

    # Export to video
    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(
            out_frames,
            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(
            out_frames,
            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=DEFAULT_DURATION, step=0.5,
                        label="Duration (seconds)",
                        info="3.5s ≈ 81 frames (native training length)",
                    )
                    gr.Markdown(
                        f"**IC-LoRA strength:** fused at `{DEFAULT_LORA_SCALE}` "
                        f"(distilled speed LoRA at `{DISTILLED_SCALE}`). "
                        "Fused into the base weights to fit ZeroGPU disk — edit the constant and rebuild to change."
                    )
                    guidance_scale = gr.Slider(
                        minimum=1.0, maximum=10.0, value=DEFAULT_GUIDANCE, step=0.5,
                        label="Guidance scale (CFG)",
                        info="Canonical is 1.0 (CFG off — the distilled regime). Above 1.0 re-enables CFG and the negative prompt.",
                    )
                    gr.Markdown(
                        f"**Sampler:** FlowMatch, fixed {len(DISTILLED_SIGMAS) - 1}-step distilled sigma "
                        "schedule. STG / modality / audio guidance disabled to match ComfyUI's CFG=1."
                    )
                    negative_prompt = gr.Textbox(
                        value="",
                        label="Negative prompt (only active if guidance > 1)",
                        lines=2,
                        placeholder="Leave empty to use the default negative prompt",
                    )
                    upscale = gr.Checkbox(
                        value=True, label="2× spatial upscale",
                        info="Canonical second pass — doubles output resolution (slower).",
                    )
                    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/crossview-source-sushi.mp4", "to the right", "lower", "closer", DEFAULT_DURATION, 42, False, DEFAULT_GUIDANCE, "", True],
        ]
        gr.Examples(
            examples=examples,
            inputs=[reference_video, azimuth, elevation, distance, duration_seconds, seed, randomize_seed, guidance_scale, negative_prompt, upscale],
            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, guidance_scale, negative_prompt, upscale],
        outputs=[output_video, used_seed, used_prompt],
    )

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