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)