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
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from diffusers.pipelines.ltx2.utils import 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 = "linoyts/ltx2.3-inpainting-lora"
LORA_FILE = "ltx-2.3-22b-ic-lora-inpainting.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)
pipe.load_lora_weights(load_file(_lora_path), adapter_name="inpaint")
pipe.set_adapters("inpaint", LORA_SCALE)

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):
    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)
    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
    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(80 + int(num_frames) * per_frame)


# --- 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(track_tqdm=True)):
    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())
    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 (generator: yields live step status across the ZeroGPU boundary) ---
@spaces.GPU(duration=_duration)
def inpaint(video, mask_video, prompt, preset, num_frames, seed, randomize):
    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]

    state = {"step": 0}
    holder = {}

    def _cb(p, i, t, kw):
        state["step"] = i + 1
        return {}

    def _run():
        try:
            holder["out"] = 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=torch.Generator(device="cuda").manual_seed(seed),
                output_type="np", return_dict=False, callback_on_step_end=_cb,
            )
        except Exception as e:  # surface in main thread
            holder["err"] = e

    th = threading.Thread(target=_run)
    th.start()
    yield None, gr.update(), "### ⏳ Preparing…"
    while th.is_alive():
        s = state["step"]
        msg = f"### 🪄 Denoising — step {s}/{NUM_STEPS}" if s else "### ⏳ Loading model / encoding…"
        yield gr.update(), gr.update(), msg
        time.sleep(0.4)
    th.join()
    if "err" in holder:
        raise holder["err"]
    yield gr.update(), gr.update(), "### 🎬 Decoding & encoding video…"
    video_out, audio_out = holder["out"]

    gen = (np.clip(video_out[0], 0, 1) * 255).astype(np.uint8)
    orig_arr = np.stack(orig).astype(np.float32)
    n = min(len(gen), len(orig_arr), len(masks))
    out = np.empty((n, height, width, 3), dtype=np.uint8)
    for i in range(n):
        soft = np.array(Image.fromarray((masks[i] * 255).astype(np.uint8)).filter(ImageFilter.GaussianBlur(3))) / 255.0
        out[i] = (gen[i].astype(np.float32) * soft[:, :, None] + orig_arr[i] * (1 - soft[:, :, None])).astype(np.uint8)

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    _export(out, audio_out, out_path)
    yield out_path, seed, "### ✅ Done"


# --- 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/ltx-community/ltx2.3-inpainting-lora) 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")
            status = gr.Markdown("")
            used_seed = gr.Number(label="Seed used", interactive=False)

    # 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, used_seed, status])

    gr.Examples(
        examples=[
            ["examples/cat.mp4", "examples/cat_mask.mp4",
             "a fluffy white dog with floppy ears in the same spot, soft panting and a gentle bark, quiet indoor room tone",
             "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, glinting metallic reflections and smooth joints; upbeat electronic music and soft mechanical whirring",
             "Fast (768×448)", 49, 42, False],
        ],
        inputs=[video_in, mask_video, prompt, preset, num_frames, seed, randomize],
        outputs=[video_out, used_seed, status], fn=inpaint, cache_examples=True, cache_mode="lazy",
    )

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