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)