Upload app.py with huggingface_hub

app.py

@@ -19,30 +19,33 @@ 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
-MASK_FILL = 128  # masked region painted neutral grey in the reference the model fills
 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 pipeline once at module scope (ZeroGPU registers it) ---------------
+# --- 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):
@@ -58,15 +61,29 @@ def _pick_resolution(first_frame: Image.Image, preset: str):
 
 
 def first_frame(video):
-    """Populate the mask editor with the uploaded video's first frame."""
     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):
-    """Extract a binary mask (H,W) from a gr.ImageEditor value — union of painted layers."""
     if not editor_value:
         return None
     layers = editor_value.get("layers") or []
@@ -87,16 +104,41 @@ def _mask_from_editor(editor_value, width, height):
     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 = args[3] if len(args) > 3 else "Fast"
-    num_frames = args[4] if len(args) > 4 else 73
+    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(50 + int(num_frames) * per_frame)
+    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_editor, prompt, preset, num_frames, mask_mode, seed, randomize,
+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.")
@@ -114,29 +156,27 @@ def inpaint(video, mask_editor, prompt, preset, num_frames, mask_mode, seed, ran
     width, height = _pick_resolution(frames[0], preset)
     num_frames = int(num_frames)
 
-    mask = _mask_from_editor(mask_editor, width, height)
-    if mask is None:
-        raise gr.Error("Draw a mask over the region to inpaint (use the brush on the frame).")
+    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)]
 
-    fill = 0 if mask_mode == "black" else MASK_FILL
-    if mask_mode == "attn":
-        # Full original video as reference; attention mask tells the model to ignore the
-        # reference (and regenerate from prompt) inside the painted 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)
-    else:
-        # Masked region painted a flat color in the reference; the model fills it.
-        ref = []
-        for fr in orig:
-            m = fr.copy()
-            m[mask] = fill
-            ref.append(Image.fromarray(m))
-        attn_mask = None
+    # 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(
@@ -160,7 +200,7 @@ def inpaint(video, mask_editor, prompt, preset, num_frames, mask_mode, seed, ran
         return_dict=False,
     )
 
-    # Composite: keep original pixels outside the mask, generated pixels inside (feathered).
+    # 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]
@@ -178,30 +218,33 @@ 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. "
-        "Upload a clip, **brush over the area to replace** on the first frame, describe what should appear there. "
-        "The mask is applied across all frames. "
+        "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) · "
-        "base: distilled LTX-2.3."
+        "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_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,
-            )
+            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)")
-                mask_mode = gr.Radio(["attn", "grey", "black"], value="attn",
-                                     label="Mask mechanism (debug)")
                 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")
@@ -209,10 +252,15 @@ with gr.Blocks(title="LTX-2.3 Video Inpainting") as demo:
             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_editor, prompt, preset, num_frames, mask_mode, seed, randomize],
+        inputs=[video_in, mask_source, mask_text, mask_editor, prompt, preset, num_frames, seed, randomize],
         outputs=[video_out, used_seed],
     )