app.py

import os

# ZeroGPU: torch.compile / dynamo are unsupported — disable before torch import.
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import random
import tempfile

import numpy as np
import spaces
import torch
import gradio as gr
from PIL import Image, ImageFilter
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 Sam3Model, Sam3Processor

# --- 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)  # 8
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]

# --- Load models once at module scope (ZeroGPU registers them) ---------------
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 = Sam3Model.from_pretrained(SAM3_REPO, token=HF_TOKEN).to("cuda")
sam3_processor = Sam3Processor.from_pretrained(SAM3_REPO, token=HF_TOKEN)


# --- Helpers ----------------------------------------------------------------
def _resample(frames, n):
    idx = np.linspace(0, len(frames) - 1, n).round().astype(int)
    return [frames[i] for i in idx]


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


def first_frame(video):
    if video is None:
        return None
    frames = load_video(video)
    return np.array(frames[0].convert("RGB")) if frames else None


def _sam3_mask(image: Image.Image, text: str, score_thr: float = 0.5):
    """Run SAM3 text-prompted segmentation; return a union boolean mask (H,W) at image size."""
    inputs = sam3_processor(images=image, text=text, return_tensors="pt").to("cuda")
    with torch.no_grad():
        outputs = sam3(**inputs)
    res = sam3_processor.post_process_instance_segmentation(
        outputs, threshold=score_thr, mask_threshold=0.5,
        target_sizes=inputs.get("original_sizes").tolist(),
    )[0]
    masks = res["masks"]
    if masks is None or len(masks) == 0:
        return None
    m = masks.cpu().numpy().astype(bool)
    return np.any(m, axis=0)  # union of all matching instances


def _mask_from_editor(editor_value, width, height):
    if not editor_value:
        return None
    layers = editor_value.get("layers") or []
    bg = editor_value.get("background")
    if bg is None:
        return None
    H0, W0 = np.asarray(bg).shape[:2]
    acc = np.zeros((H0, W0), dtype=bool)
    for layer in layers:
        arr = np.asarray(layer)
        if arr.ndim == 3 and arr.shape[2] == 4:
            acc |= arr[..., 3] > 10
        elif arr.ndim == 3:
            acc |= arr.sum(axis=2) > 10
    if not acc.any():
        return None
    m = Image.fromarray((acc * 255).astype(np.uint8)).resize((width, height), Image.NEAREST)
    return np.array(m) > 127


def _overlay(image: Image.Image, mask: np.ndarray):
    img = image.convert("RGBA")
    m = Image.fromarray((mask * 255).astype(np.uint8)).resize(image.size, Image.NEAREST)
    ov = Image.new("RGBA", image.size, (255, 45, 85, 0))
    ov.putalpha(m.point(lambda v: int(v * 0.5)))
    return Image.alpha_composite(img, ov).convert("RGB")


def _duration(*args, **kwargs):
    preset = next((a for a in args if a in RES_PRESETS), "Fast")
    num_frames = next((a for a in args if a in FRAME_CHOICES), 73)
    per_frame = 1.6 if "Quality" in str(preset) else 1.0
    return int(60 + int(num_frames) * per_frame)


# --- SAM3 mask preview (cheap GPU call) -------------------------------------
@spaces.GPU(duration=40)
def preview_mask(video, mask_text, 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'.")
    frames = load_video(video)
    if not frames:
        raise gr.Error("Could not read the video.")
    f0 = frames[0].convert("RGB")
    mask = _sam3_mask(f0, mask_text.strip())
    if mask is None:
        raise gr.Error(f"SAM3 found no '{mask_text}' in the first frame. Try a different phrase.")
    return _overlay(f0, mask)


# --- Inference --------------------------------------------------------------
@spaces.GPU(duration=_duration)
def inpaint(video, mask_source, mask_text, mask_editor, prompt, preset, num_frames, seed, randomize,
            progress=gr.Progress(track_tqdm=True)):
    if video is None:
        raise gr.Error("Please upload a video.")
    if not prompt.strip():
        raise gr.Error("Describe what should fill the masked region.")

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

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

    width, height = _pick_resolution(frames[0], preset)
    num_frames = int(num_frames)

    if mask_source == "Text (SAM3)":
        if not mask_text.strip():
            raise gr.Error("Type what to mask, e.g. 'the cat'.")
        sam_mask = _sam3_mask(frames[0].convert("RGB"), mask_text.strip())
        if sam_mask is None:
            raise gr.Error(f"SAM3 found no '{mask_text}' in the first frame.")
        mask = np.array(Image.fromarray((sam_mask * 255).astype(np.uint8)).resize((width, height), Image.NEAREST)) > 127
    else:
        mask = _mask_from_editor(mask_editor, width, height)
        if mask is None:
            raise gr.Error("Draw a mask with the brush, or switch to Text (SAM3) masking.")

    orig = [np.array(f.convert("RGB").resize((width, height), Image.LANCZOS))
            for f in _resample(frames, num_frames)]

    # Full original video as reference; the attention mask makes the model ignore the
    # reference (and regenerate from the prompt) inside the masked region.
    ref = [Image.fromarray(fr) for fr in orig]
    am = np.ones((num_frames, height, width), dtype=np.float32)
    am[:, mask] = 0.0
    attn_mask = torch.from_numpy(am)[None, None]  # (1,1,F,H,W)

    ref_cond = LTX2ReferenceCondition(frames=ref, strength=1.0)
    video_out, _audio = pipe(
        prompt=prompt,
        negative_prompt="",
        reference_conditions=[ref_cond],
        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,
    )

    # Composite generated pixels inside the mask over the original (feathered edges).
    gen = (np.clip(video_out[0], 0, 1) * 255).astype(np.uint8)
    soft = np.array(Image.fromarray((mask * 255).astype(np.uint8)).filter(ImageFilter.GaussianBlur(3))) / 255.0
    soft = soft[None, :, :, None]
    orig_arr = np.stack(orig).astype(np.float32)
    n = min(len(gen), len(orig_arr))
    out = (gen[:n].astype(np.float32) * soft + orig_arr[:n] * (1 - soft)).astype(np.uint8)

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    encode_video(out, fps=FPS, output_path=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 of the frame intact. "
        "Pick the area with a **text prompt (SAM3 auto-mask)** or by **brushing** on the first frame; "
        "the mask applies across all frames. "
        "IC-LoRA: [`linoyts/ltx2.3-inpainting-lora`](https://huggingface.co/linoyts/ltx2.3-inpainting-lora) · "
        "auto-mask: [SAM3](https://huggingface.co/facebook/sam3) · base: distilled LTX-2.3."
    )
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Input video")
            mask_source = gr.Radio(["Text (SAM3)", "Brush"], value="Text (SAM3)", label="How to mask")

            with gr.Group(visible=True) as text_group:
                mask_text = gr.Textbox(label="Object(s) to mask", placeholder="the cat")
                preview_btn = gr.Button("Preview mask")
                mask_preview = gr.Image(label="SAM3 mask preview", type="pil", interactive=False)

            with gr.Group(visible=False) as brush_group:
                mask_editor = gr.ImageEditor(
                    label="Brush the region to inpaint (loads from the video's first frame)",
                    type="numpy", layers=False,
                    brush=gr.Brush(colors=["#ff2d55"], color_mode="fixed"),
                )

            prompt = gr.Textbox(label="What should fill the masked region",
                                placeholder="a lush green bush with small white flowers", lines=2)
            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")
            used_seed = gr.Number(label="Seed used", interactive=False)

    def _toggle(src):
        return gr.update(visible=src == "Text (SAM3)"), gr.update(visible=src == "Brush")

    mask_source.change(_toggle, inputs=mask_source, outputs=[text_group, brush_group])
    video_in.change(first_frame, inputs=video_in, outputs=mask_editor)
    preview_btn.click(preview_mask, inputs=[video_in, mask_text], outputs=mask_preview)
    run.click(
        inpaint,
        inputs=[video_in, mask_source, mask_text, mask_editor, prompt, preset, num_frames, seed, randomize],
        outputs=[video_out, used_seed],
    )

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