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