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 cv2
import numpy as np
import imageio.v3 as iio
from tqdm import tqdm
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
from transformers import Sam3VideoModel, Sam3VideoProcessor

# --- Config -----------------------------------------------------------------
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
SAM3_REPO = "facebook/sam3"
FPS = 24
NUM_STEPS = len(DISTILLED_SIGMA_VALUES)
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")

RES_PRESETS = {"Fast (768×448)": (768, 448), "Quality (960×544)": (960, 544)}
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)

sam3 = Sam3VideoModel.from_pretrained(SAM3_REPO, token=HF_TOKEN, dtype=torch.bfloat16)
sam3_processor = Sam3VideoProcessor.from_pretrained(SAM3_REPO, token=HF_TOKEN)


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


def _load_frames(path, num_frames, width, height):
    frames = load_video(path)
    if not frames:
        return []
    fps = _src_fps(path)
    return [ImageOps.fit(frames[min(int(round(i / FPS * fps)), len(frames) - 1)].convert("RGB"),
                         (width, height), Image.LANCZOS) for i in range(num_frames)]


def _pick_resolution(first_frame, preset):
    w, h = RES_PRESETS[preset]
    if first_frame.height > first_frame.width:
        w, h = h, w
    return w, h


def _sam3_video_masks(frames_pil, text, progress=None):
    sam3.to("cuda")
    W, H = frames_pil[0].size
    video = np.stack([np.array(f) for f in frames_pil])
    session = sam3_processor.init_video_session(
        video=video, inference_device="cuda", processing_device="cpu",
        video_storage_device="cpu", dtype=torch.bfloat16,
    )
    session = sam3_processor.add_text_prompt(inference_session=session, text=text)
    n = len(frames_pil)
    masks = np.zeros((n, H, W), dtype=bool)
    # Manual per-frame progress (NOT a tqdm-wrap of the SAM3 generator — that crashes
    # postprocess_outputs with track_tqdm active). progress() forwards over ZeroGPU.
    for mo in sam3.propagate_in_video_iterator(inference_session=session, max_frame_num_to_track=n):
        proc = sam3_processor.postprocess_outputs(session, mo)
        m = proc.get("masks")
        if m is not None and len(m):
            arr = m.float().cpu().numpy()
            arr = (arr.reshape(-1, arr.shape[-2], arr.shape[-1]) > 0.5).any(axis=0)
            if arr.shape != (H, W):
                arr = np.array(Image.fromarray((arr * 255).astype(np.uint8)).resize((W, H), Image.NEAREST)) > 127
            if 0 <= mo.frame_idx < n:
                masks[mo.frame_idx] = arr
        if progress is not None and 0 <= mo.frame_idx < n:
            progress((mo.frame_idx + 1) / n, desc=f"Segmenting with SAM3 — frame {mo.frame_idx + 1}/{n}")
    return masks if masks.any() else None


def _write_mask_video(masks_bool):
    vid = (masks_bool[..., None] * np.array([255, 255, 255], np.uint8)).astype(np.uint8)
    path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    iio.imwrite(path, vid, fps=FPS, plugin="pyav", codec="libx264")
    return path


def _read_mask(mask_video, num_frames, width, height):
    mframes = load_video(mask_video)
    if not mframes:
        return None
    idx = np.linspace(0, len(mframes) - 1, num_frames).round().astype(int)
    masks = np.stack([np.array(mframes[i].convert("L").resize((width, height), Image.NEAREST)) > 127 for i in idx])
    return masks if masks.any() else None


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):
    preset = next((a for a in args if isinstance(a, str) and a in RES_PRESETS), "Fast")
    num_frames = next((a for a in args if isinstance(a, int) and a in FRAME_CHOICES), 73)
    per_frame = 1.6 if "Quality" in str(preset) else 1.0
    return int(100 + int(num_frames) * per_frame * 1.5)  # two-stage (x2 upsample + refine), AOTI off


# --- Mask editing (CPU, no GPU) ---------------------------------------------
def apply_dilation(base, mask_video, dilate_px):
    """Re-dilate the UNDILATED base mask by the slider amount (every frame; no compounding).
    Falls back to the currently-shown mask as the base (e.g. after loading an example).
    Returns (mask_to_show, base_to_keep)."""
    src = base or mask_video
    if not src:
        return gr.update(), base
    if int(dilate_px) <= 0:
        return src, src
    mframes = load_video(src)
    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * int(dilate_px) + 1, 2 * int(dilate_px) + 1))
    out = [np.stack([cv2.dilate((np.array(f.convert("L")) > 127).astype(np.uint8), k) * 255] * 3, axis=-1).astype(np.uint8)
           for f in mframes]
    path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    iio.imwrite(path, np.stack(out), fps=FPS, plugin="pyav", codec="libx264")
    return path, src


# --- SAM3 mask generation ----------------------------------------------------
@spaces.GPU(duration=70)
def gen_mask(video, mask_text, preset, progress=gr.Progress()):
    if video is None:
        raise gr.Error("Upload a video first.")
    if not mask_text.strip():
        raise gr.Error("Type what to mask, e.g. 'the cat'.")
    probe = load_video(video)
    if not probe:
        raise gr.Error("Could not read the video.")
    width, height = _pick_resolution(probe[0], preset)
    frames = _load_frames(video, 25, width, height)
    masks = _sam3_video_masks(frames, mask_text.strip(), progress)
    if masks is None:
        raise gr.Error(f"SAM3 found no '{mask_text}' in the video. Try a simpler phrase (e.g. 'person', 'dog').")
    return _write_mask_video(masks)


# --- Inference -------------------------------------------------------------
@spaces.GPU(duration=_duration)
def inpaint(video, mask_video, prompt, preset, num_frames, seed, randomize, progress=gr.Progress(track_tqdm=True)):
    if video is None:
        raise gr.Error("Please upload a video.")
    if mask_video is None:
        raise gr.Error("Add a mask: type an object and click 'Generate mask', or upload a mask video.")
    if not prompt.strip():
        raise gr.Error("Describe what should fill the masked region (and any sounds).")

    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    num_frames = int(num_frames)

    probe = load_video(video)
    if not probe:
        raise gr.Error("Could not read any frames from that video.")
    width, height = _pick_resolution(probe[0], preset)

    orig_pil = _load_frames(video, num_frames, width, height)
    orig = [np.array(f) for f in orig_pil]

    masks = _read_mask(mask_video, num_frames, width, height)
    if masks is None:
        raise gr.Error("The mask is empty. White pixels mark the region to inpaint.")

    ref = [Image.fromarray(fr) for fr in orig]
    am = np.ones((num_frames, height, width), dtype=np.float32)
    am[masks] = 0.0
    attn_mask = torch.from_numpy(am)[None, None]

    gen_ = torch.Generator(device="cuda").manual_seed(seed)
    # --- Stage 1: masked inpaint generate (distilled 8-step, IC-LoRA + attention mask) ---
    pipe.set_adapters("inpaint", LORA_SCALE)
    video_latent, audio_latent = pipe(
        prompt=prompt.strip(), negative_prompt="",
        reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
        conditioning_attention_mask=attn_mask, reference_downscale_factor=1,
        width=width, height=height, 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=prompt.strip(), negative_prompt="",
            latents=up_latent, audio_latents=audio_latent,
            width=width * 2, height=height * 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()

    # composite the refined masked region over the (2x-upscaled) original to keep the rest intact
    H2, W2 = height * 2, width * 2
    gen = (np.clip(video_out[0], 0, 1) * 255).astype(np.uint8)
    orig2 = np.stack([np.array(f.resize((W2, H2), Image.LANCZOS)) for f in orig_pil]).astype(np.float32)
    n = min(len(gen), len(orig2), len(masks))
    out = np.empty((n, H2, W2, 3), dtype=np.uint8)
    for i in range(n):
        m2 = Image.fromarray((masks[i] * 255).astype(np.uint8)).resize((W2, H2), Image.NEAREST)
        soft = np.array(m2.filter(ImageFilter.GaussianBlur(6))) / 255.0
        out[i] = (gen[i].astype(np.float32) * soft[:, :, None] + orig2[i] * (1 - soft[:, :, None])).astype(np.uint8)

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


# --- UI ---------------------------------------------------------------------
with gr.Blocks(title="LTX-2.3 Video Inpainting") as demo:
    gr.Markdown(
        "# 🪄 LTX-2.3 Video Inpainting\n"
        "Mask a region of a video and regenerate it from a prompt, keeping the rest intact. Type an object and "
        "**generate a tracked mask with SAM3**, or **upload your own mask video** (white = region). "
        "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) and "
        "[SAM3](https://huggingface.co/facebook/sam3), via diffusers 🧨."
    )
    base_mask = gr.State()
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Input video")
            mask_video = gr.Video(label="Mask video (white = region to inpaint) — generated or uploaded")
            with gr.Row():
                mask_text = gr.Textbox(label="Object to mask (SAM3)", placeholder="the cat", scale=3)
                gen_btn = gr.Button("Generate mask", scale=1)
            prompt = gr.Textbox(label="Prompt — what should fill the region, plus any sounds", lines=2,
                                placeholder="a fluffy white dog with floppy ears in the same spot, soft panting and a gentle bark")
            dilate_px = gr.Slider(0, 48, value=0, step=4,
                                  label="Expand masked region (px) — grow the mask outward if it doesn't cover enough")
            with gr.Accordion("Settings", open=False):
                preset = gr.Dropdown(list(RES_PRESETS), value="Fast (768×448)", label="Resolution")
                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("Inpaint", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Inpainted result")

    # a new input video invalidates any stored base mask
    video_in.change(lambda: None, outputs=base_mask)
    # SAM3 mask -> base, then show (dilated by current slider, all frames)
    gen_btn.click(gen_mask, inputs=[video_in, mask_text, preset], outputs=base_mask).then(
        apply_dilation, inputs=[base_mask, mask_video, dilate_px], outputs=[mask_video, base_mask])
    # uploaded mask -> base, then show
    mask_video.upload(lambda v: v, inputs=mask_video, outputs=base_mask).then(
        apply_dilation, inputs=[base_mask, mask_video, dilate_px], outputs=[mask_video, base_mask])
    # slider -> re-dilate base across ALL frames, live update
    dilate_px.release(apply_dilation, inputs=[base_mask, mask_video, dilate_px], outputs=[mask_video, base_mask])

    run.click(inpaint, inputs=[video_in, mask_video, prompt, preset, num_frames, seed, randomize],
              outputs=[video_out, seed])

    gr.Examples(
        examples=[
            ["examples/cat.mp4", "examples/cat_mask.mp4",
             "a fluffy white dog with soft floppy ears, a small black nose and bright dark eyes, sitting in the same spot and looking around alertly with a happy little tail-wag, its coat ruffling gently in the breeze under warm natural daylight; immersive outdoor street ambience — distant passing traffic and car hum, footsteps on the pavement, faint chatter and city birdsong, with the dog's soft panting and a single cheerful bark",
             "Fast (768×448)", 49, 42, False],
            ["examples/man_dancing.mp4", "examples/man_dancing_mask.mp4",
             "a sleek chrome humanoid robot dancing in the same spot with smooth articulated joints and glinting metallic reflections, its polished silver plating catching the light as it pops and locks with crisp, confident moves in perfect rhythm; a funky upbeat dance track with a punchy bassline and a snappy beat — energetic TikTok-style groove — layered with subtle mechanical servo whirs and clicks timed to every move",
             "Fast (768×448)", 49, 42, False],
        ],
        inputs=[video_in, mask_video, prompt, preset, num_frames, seed, randomize],
        outputs=[video_out, seed], fn=inpaint, cache_examples=True, cache_mode="lazy",
    )

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