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)