Spaces:
Sleeping
Sleeping
| 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 ---------------------------------------------------- | |
| 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) --- | |
| 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) | |