app.py

import os

os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import random
import tempfile
import threading
import time

import numpy as np
import imageio.v3 as iio
import spaces
import torch
import gradio as gr
from PIL import Image, ImageFilter, ImageOps
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from diffusers import LTX2InContextPipeline, LTX2LatentUpsamplePipeline
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from diffusers.utils import load_video, encode_video

# --- Config -----------------------------------------------------------------
# Outpainting reuses the inpainting IC-LoRA: pad the input to a target aspect ratio and
# regenerate the empty margins via the conditioning_attention_mask (0 in the margins).
BASE_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
LORA_REPO = "Lightricks/LTX-2.3-22b-IC-LoRA-In-Outpainting"
LORA_FILE = "ltx-2.3-22b-ic-lora-in-outpainting-0.9.safetensors"
LORA_SCALE = 1.0
FPS = 24
NUM_STEPS = len(DISTILLED_SIGMA_VALUES)
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")

CANVASES = {
    "Landscape 16:9 (768×448)": (768, 448),
    "Portrait 9:16 (448×768)": (448, 768),
    "Square 1:1 (640×640)": (640, 640),
    "Standard 4:3 (768×576)": (768, 576),
    "Cinemascope 2.39:1 (768×320)": (768, 320),
}
FRAME_CHOICES = [49, 73, 97, 121]

pipe = LTX2InContextPipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe.vae.enable_tiling()
_lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=HF_TOKEN)
# Kept as a togglable adapter (NOT fused) so Stage 2 runs on the bare distilled model.
pipe.load_lora_weights(load_file(_lora_path), adapter_name="inpaint")
pipe.set_adapters("inpaint", LORA_SCALE)
# NOTE: AOTI temporarily disabled while validating 2-stage inference; re-enable once confirmed.
# spaces.aoti_load(module=pipe.transformer, repo_id="ltx-community/LTX-2.3-Transformer-GroupB-sm120-cu130-r0e")

# Stage-2 latent upsampler (spatial x2) — two-stage diffusers inference.
_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
    "dg845/LTX-2.3-Spatial-Upsampler-Diffusers", subfolder="latent_upsampler", torch_dtype=torch.bfloat16)
_upsampler.to("cuda")
upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=_upsampler)


def _src_fps(path, default=FPS):
    try:
        return float(iio.immeta(path, plugin="pyav").get("fps", default)) or default
    except Exception:
        return default


def _natural_pil(path, num_frames, max_side=768):
    frames = load_video(path)
    if not frames:
        return []
    fps = _src_fps(path)
    out = []
    for i in range(num_frames):
        idx = min(int(round(i / FPS * fps)), len(frames) - 1)
        f = frames[idx].convert("RGB")
        f.thumbnail((max_side, max_side), Image.LANCZOS)
        out.append(f)
    return out


def _layout(inner_size, canvas):
    cw, ch = canvas
    iw, ih = inner_size
    s = min(cw / iw, ch / ih)
    nw, nh = max(1, round(iw * s)), max(1, round(ih * s))
    ox, oy = (cw - nw) // 2, (ch - nh) // 2
    return cw, ch, nw, nh, ox, oy


def _export(video_np, audio, path):
    kw = {}
    if audio is not None:
        kw = dict(audio=audio[0].float().cpu(), audio_sample_rate=pipe.vocoder.config.output_sampling_rate)
    encode_video(video_np, fps=FPS, output_path=path, **kw)


def _duration(*args, **kwargs):
    num_frames = next((a for a in args if a in FRAME_CHOICES), 73)
    return int(100 + int(num_frames) * 2.0)  # two-stage (x2 upsample + refine), AOTI off


@spaces.GPU(duration=_duration)
def outpaint(video, canvas_key, prompt, num_frames, seed, randomize,
             progress=gr.Progress(track_tqdm=True)):
    if video is None:
        raise gr.Error("Please upload a video.")
    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    num_frames = int(num_frames)
    canvas = CANVASES[canvas_key]

    frames = _natural_pil(video, num_frames)
    if not frames:
        raise gr.Error("Could not read any frames from that video.")

    cw, ch, nw, nh, ox, oy = _layout(frames[0].size, canvas)
    if ox == 0 and oy == 0:
        raise gr.Error("That target aspect ratio matches the input — nothing to outpaint. Pick a different one.")

    # mask = the margins to fill (1 in margins -> attention 0 there)
    mask = np.ones((ch, cw), dtype=bool)
    mask[oy:oy + nh, ox:ox + nw] = False

    ref, inners = [], []
    for f in frames:
        inner = f.resize((nw, nh), Image.LANCZOS)
        inners.append(np.array(inner))
        canvas_img = Image.new("RGB", (cw, ch), (128, 128, 128))
        canvas_img.paste(inner, (ox, oy))
        ref.append(canvas_img)

    am = np.ones((num_frames, ch, cw), dtype=np.float32)
    am[:, mask] = 0.0
    attn_mask = torch.from_numpy(am)[None, None]

    desc = prompt.strip() or "the scene continues naturally beyond the original frame, consistent style and lighting"
    full_prompt = f"{desc}; seamlessly extend the scene into the empty margins, matching the existing content."

    gen_ = torch.Generator(device="cuda").manual_seed(seed)
    # --- Stage 1: outpaint generate (distilled 8-step, IC-LoRA + margin attention mask) ---
    pipe.set_adapters("inpaint", LORA_SCALE)
    video_latent, audio_latent = pipe(
        prompt=full_prompt, negative_prompt="",
        reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
        conditioning_attention_mask=attn_mask, reference_downscale_factor=1,
        width=cw, height=ch, num_frames=num_frames, frame_rate=FPS,
        num_inference_steps=NUM_STEPS, sigmas=DISTILLED_SIGMA_VALUES,
        guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
        generator=gen_, output_type="latent", return_dict=False,
    )
    # --- Stage 2a: spatial x2 latent upsample ---
    up_latent = upsample_pipe(latents=video_latent, output_type="latent", return_dict=False)[0]
    # --- Stage 2b: refine at 2x res on bare distilled (drop IC-LoRA, reference, mask) ---
    pipe.disable_lora()
    try:
        video_out, audio_out = pipe(
            prompt=full_prompt, negative_prompt="",
            latents=up_latent, audio_latents=audio_latent,
            width=cw * 2, height=ch * 2, num_frames=num_frames, frame_rate=FPS,
            num_inference_steps=len(STAGE_2_DISTILLED_SIGMA_VALUES),
            sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
            guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
            generator=gen_, output_type="np", return_dict=False,
        )
    finally:
        pipe.enable_lora()

    # keep the original pixels exactly in the center (2x-upscaled); generated pixels in the margins (feathered).
    ox2, oy2, nw2, nh2 = ox * 2, oy * 2, nw * 2, nh * 2
    H2, W2 = ch * 2, cw * 2
    gen = (np.clip(video_out[0], 0, 1) * 255).astype(np.uint8)
    mask2 = np.ones((H2, W2), dtype=bool)
    mask2[oy2:oy2 + nh2, ox2:ox2 + nw2] = False
    soft = np.array(Image.fromarray((mask2 * 255).astype(np.uint8)).filter(ImageFilter.GaussianBlur(8))) / 255.0
    soft = soft[None, :, :, None]
    n = min(len(gen), num_frames)
    base = gen[:n].astype(np.float32).copy()
    for i in range(n):
        inner2 = np.array(frames[i].resize((nw2, nh2), Image.LANCZOS)).astype(np.float32)
        base[i, oy2:oy2 + nh2, ox2:ox2 + nw2] = inner2
    out = (gen[:n].astype(np.float32) * soft + base * (1 - soft)).astype(np.uint8)

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    _export(out, audio_out, out_path)
    return out_path, seed


with gr.Blocks(title="LTX-2.3 Video Outpaint") as demo:
    gr.Markdown(
        "# 🖼️ LTX-2.3 Video Outpainting\n"
        "Extend a video to a new aspect ratio — the original stays centered and the model fills the empty "
        "margins, matching the scene. Using "
        "[LTX 2.3 Distilled](https://huggingface.co/diffusers/LTX-2.3-Distilled-Diffusers) with the "
        "[Inpainting IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-In-Outpainting), via diffusers 🧨."
    )
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Input video")
            canvas_key = gr.Dropdown(list(CANVASES), value="Landscape 16:9 (768×448)", label="Target frame / aspect")
            prompt = gr.Textbox(label="What's beyond the edges, plus any sounds (optional)", lines=3,
                                placeholder="more of the same forest extending left and right, soft dappled light; forest ambience and birdsong")
            with gr.Accordion("Settings", open=False):
                num_frames = gr.Dropdown(FRAME_CHOICES, value=73, label="Frames (24fps)")
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            run = gr.Button("Outpaint", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Outpainted result")

    run.click(outpaint, inputs=[video_in, canvas_key, prompt, num_frames, seed, randomize],
              outputs=[video_out, seed])

    gr.Examples(
        examples=[
            ["examples/portrait_clip.mp4", "Landscape 16:9 (768×448)",
             "the ocean scene continues naturally to the left and right beyond the original frame — the same rolling turquoise water and white-capped waves stretching toward a wide horizon, sky, light and motion matching seamlessly; immersive ocean ambience with rolling waves, sea spray and gusting wind", 73, 42, False],
            ["examples/landscape_clip.mp4", "Cinemascope 2.39:1 (768×320)",
             "the misty mountain vista extends much wider on both sides — the same layered ridgelines and drifting fog continuing to the horizon, the calm reflective water broadening below, colour, haze and light matching seamlessly; a gentle wind moving over the water and faint distant birdsong", 73, 42, False],
        ],
        inputs=[video_in, canvas_key, prompt, num_frames, seed, randomize],
        outputs=[video_out, seed], fn=outpaint, cache_examples=True, cache_mode="lazy",
    )

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