Spaces:
Running on Zero
Running on Zero
| """VideoMaMa++ demo: SAM3 mask track -> VideoMaMa++ alpha matte. | |
| The VideoMaMa++ weights and model code live in a private Hugging Face repo. | |
| They are downloaded at startup with the HF_TOKEN Space secret into the Hub | |
| cache (outside the app folder) and are never served to visitors. | |
| """ | |
| import os | |
| import sys | |
| import time | |
| import uuid | |
| from pathlib import Path | |
| try: | |
| import spaces # ZeroGPU; must be imported before torch touches CUDA | |
| GPU = spaces.GPU | |
| except ImportError: # local run | |
| def GPU(fn=None, **_): | |
| return fn if fn is not None else (lambda f: f) | |
| import cv2 | |
| import gradio as gr | |
| import imageio_ffmpeg | |
| import numpy as np | |
| import subprocess | |
| import torch | |
| from huggingface_hub import snapshot_download | |
| from PIL import Image | |
| from painter import paint_mask, paint_points | |
| PRIVATE_REPO = os.environ.get("VIDEOMAMAPP_REPO", "SammyLim/VideoMaMa-PlusPlus") | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| APP_DIR = Path(__file__).resolve().parent | |
| OUT_DIR = Path(os.environ.get("VIDEOMAMAPP_OUT", "/tmp/videomamapp_outputs")) | |
| # Wan2.1 base files (public) are cached here, outside the app folder. | |
| os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", str(Path.home() / ".cache" / "videomamapp_models")) | |
| os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface") | |
| MAX_FRAMES = 48 # hard cap per run | |
| DEFAULT_FRAMES = 24 | |
| MAX_LONG_SIDE = 1536 # the paper's processing size; larger videos are capped here | |
| MIN_LONG_SIDE = 256 | |
| SCALES = {"100% (original)": 1.0, "75%": 0.75, "50%": 0.5, "25%": 0.25} | |
| DEFAULT_SCALE = "50%" | |
| GREEN = np.array([120, 255, 155], np.float32) | |
| # ---------------------------------------------------------------- models | |
| print("[startup] downloading VideoMaMa++ (private repo)...", flush=True) | |
| _local = os.environ.get("VIDEOMAMAPP_LOCAL_DIR") # local testing only | |
| if not _local and not HF_TOKEN: | |
| raise RuntimeError( | |
| "HF_TOKEN is not set. Add a read token with access to the private model repo " | |
| f"{PRIVATE_REPO} and to facebook/sam3 as a Space secret named HF_TOKEN.") | |
| _private = Path(_local) if _local else Path(snapshot_download(PRIVATE_REPO, token=HF_TOKEN)) | |
| sys.path.insert(0, str(_private / "code")) | |
| import videomamapp_inference as vmpp # noqa: E402 | |
| from sam3_tracker import Sam3Tracker # noqa: E402 | |
| print("[startup] loading VideoMaMa++ ...", flush=True) | |
| PIPE = vmpp.load_pipeline(_private / "videomamapp_dit.safetensors", _private / "prompt_context.pt", device="cpu") | |
| PIPE.to("cuda") | |
| print("[startup] loading SAM3 ...", flush=True) | |
| SAM3 = Sam3Tracker(device="cuda", token=HF_TOKEN) | |
| print("[startup] ready", flush=True) | |
| # ---------------------------------------------------------------- video io | |
| def read_video(path, max_frames, stride): | |
| cap = cv2.VideoCapture(str(path)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) or 24.0 | |
| orig = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))) | |
| frames, i = [], 0 | |
| while len(frames) < max_frames: | |
| ok, f = cap.read() | |
| if not ok: | |
| break | |
| if i % stride == 0: | |
| f = cv2.cvtColor(f, cv2.COLOR_BGR2RGB) | |
| h, w = f.shape[:2] | |
| s = MAX_LONG_SIDE / max(h, w) | |
| if s < 1: | |
| f = cv2.resize(f, (round(w * s / 2) * 2, round(h * s / 2) * 2), interpolation=cv2.INTER_AREA) | |
| frames.append(f) | |
| i += 1 | |
| cap.release() | |
| return frames, fps / stride, orig | |
| def processing_size(orig, scale): | |
| """(W, H) the model runs at: `scale` of the original size, where the | |
| original's long side is first capped at MAX_LONG_SIDE; floored at | |
| MIN_LONG_SIDE, both sides multiples of 16.""" | |
| w, h = orig | |
| long_side = max(MIN_LONG_SIDE, min(max(w, h), MAX_LONG_SIDE) * SCALES[scale]) | |
| f = long_side / max(w, h) | |
| return max(16, round(w * f / 16) * 16), max(16, round(h * f / 16) * 16) | |
| def gpu_seconds(n, size): | |
| """Estimated GPU time. VideoMaMa++ per frame (6-frame chunks, overlap 2) fitted | |
| to A100 runs at 512-1536 px (0.15-2.78 s/frame), plus SAM3 at ~0.1 s/frame.""" | |
| mp = size[0] * size[1] / 1e6 | |
| return n * (0.1 + 1.135 * mp + 0.718 * mp * mp) | |
| def write_video(frames, path, fps): | |
| """H.264 mp4 that browsers can play.""" | |
| h, w = frames[0].shape[:2] | |
| W, H = w - w % 2, h - h % 2 | |
| cmd = [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", | |
| "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{W}x{H}", "-r", f"{fps:.4f}", "-i", "-", | |
| "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "18", "-movflags", "+faststart", str(path)] | |
| p = subprocess.Popen(cmd, stdin=subprocess.PIPE) | |
| for f in frames: | |
| if f.ndim == 2: | |
| f = np.repeat(f[..., None], 3, 2) | |
| p.stdin.write(np.ascontiguousarray(f[:H, :W]).tobytes()) | |
| p.stdin.close() | |
| p.wait() | |
| return str(path) | |
| # ---------------------------------------------------------------- GPU work | |
| def gpu_first_frame_mask(frame, points, labels): | |
| return SAM3.first_frame_mask(frame, points, labels) | |
| def _run_duration(frames, points, labels, text, max_size): | |
| # GPU seconds to reserve. Visitors are refused when this exceeds their | |
| # remaining ZeroGPU quota, so keep it close to the real cost (model | |
| # loading + the estimate, with 50% headroom). | |
| h, w = frames[0].shape[:2] | |
| return int(min(240, 20 + 1.5 * gpu_seconds(len(frames), (w, h)))) | |
| def gpu_track_and_matte(frames, points, labels, text, max_size): | |
| t0 = time.time() | |
| if text: | |
| masks = SAM3.track_text(frames, text) | |
| else: | |
| masks = SAM3.track_points(frames, points, labels) | |
| torch.cuda.synchronize() | |
| t1 = time.time() | |
| mattes = vmpp.matte(PIPE, frames, masks, max_size=int(max_size)) | |
| t2 = time.time() | |
| timing = {"gpu": torch.cuda.get_device_name(), "sam3": t1 - t0, "matte": t2 - t1} | |
| return masks, mattes, timing | |
| # ---------------------------------------------------------------- handlers | |
| def on_load(video, n_frames, stride, scale, state): | |
| if video is None: | |
| raise gr.Error("Upload a video or pick an example first.") | |
| frames, fps, orig = read_video(video, int(n_frames), int(stride)) | |
| if not frames: | |
| raise gr.Error("Could not read any frames from this video.") | |
| state = {"frames": frames, "fps": fps, "orig": orig, "points": [], "labels": [], "mask": None} | |
| msg = f"Loaded {len(frames)} frames ({orig[0]}脳{orig[1]}). Click on the object, or use a text prompt." | |
| return state, Image.fromarray(frames[0]), msg, on_scale(state, scale) | |
| def on_scale(state, scale): | |
| if not state or "frames" not in state: | |
| return "" | |
| size = processing_size(state["orig"], scale) | |
| sec = gpu_seconds(len(state["frames"]), size) | |
| return (f"Runs at **{size[0]}脳{size[1]}** on {len(state['frames'])} frames: " | |
| f"about **{sec:.0f} s** of GPU time (A100 estimate).") | |
| def on_click(state, kind, evt: gr.SelectData): | |
| if not state or "frames" not in state: | |
| raise gr.Error("Load a video first.") | |
| x, y = evt.index[0], evt.index[1] | |
| state["points"].append([x, y]) | |
| state["labels"].append(1 if kind == "Positive" else 0) | |
| frame = state["frames"][0] | |
| state["mask"] = gpu_first_frame_mask(frame, state["points"], state["labels"]) | |
| view = paint_points(paint_mask(frame, state["mask"]), state["points"], state["labels"]) | |
| return state, Image.fromarray(view) | |
| def on_clear(state): | |
| if not state or "frames" not in state: | |
| return state, None | |
| state.update(points=[], labels=[], mask=None) | |
| return state, Image.fromarray(state["frames"][0]) | |
| def on_run(state, text, scale, progress=gr.Progress()): | |
| if not state or "frames" not in state: | |
| raise gr.Error("Load a video first.") | |
| text = (text or "").strip() | |
| if not text and not state["points"]: | |
| raise gr.Error("Click on the object in the first frame, or type a text prompt.") | |
| frames, points = state["frames"], state["points"] | |
| # Work at the processing resolution: SAM3, VideoMaMa++ and the output | |
| # videos all use the downscaled frames. | |
| h, w = frames[0].shape[:2] | |
| size = processing_size(state["orig"], scale) | |
| max_size = max(size) | |
| if size != (w, h): | |
| interp = cv2.INTER_AREA if size[0] < w else cv2.INTER_CUBIC | |
| frames = [cv2.resize(f, size, interpolation=interp) for f in frames] | |
| points = [[x * size[0] / w, y * size[1] / h] for x, y in points] | |
| progress(0.1, desc=f"SAM3 + VideoMaMa++ on {len(frames)} frames ...") | |
| t0 = time.time() | |
| masks, mattes, timing = gpu_track_and_matte(frames, points, state["labels"], text, max_size) | |
| took = time.time() - t0 | |
| run = OUT_DIR / uuid.uuid4().hex[:10] | |
| run.mkdir(parents=True, exist_ok=True) | |
| fps = state["fps"] | |
| comps = [] | |
| for f, a in zip(frames, mattes): | |
| al = (a.astype(np.float32) / 255.0)[..., None] | |
| comps.append((f.astype(np.float32) * al + GREEN * (1 - al)).astype(np.uint8)) | |
| alpha_v = write_video(mattes, run / "alpha.mp4", fps) | |
| green_v = write_video(comps, run / "greenscreen.mp4", fps) | |
| mask_v = write_video(masks, run / "sam3_mask.mp4", fps) | |
| write = time.time() - t0 - took | |
| wait = took - timing["sam3"] - timing["matte"] | |
| print(f"[run] {len(frames)} frames {frames[0].shape[1]}x{frames[0].shape[0]} ({scale}) on {timing['gpu']}: " | |
| f"SAM3 {timing['sam3']:.1f}s, VideoMaMa++ {timing['matte']:.1f}s, GPU wait/transfer {wait:.1f}s, " | |
| f"encode {write:.1f}s", flush=True) | |
| return alpha_v, green_v, mask_v, ( | |
| f"Done: {len(frames)} frames at {size[0]}脳{size[1]} in {took + write:.0f} s " | |
| f"(SAM3 {timing['sam3']:.0f} s, VideoMaMa++ {timing['matte']:.0f} s).") | |
| def on_new_video(): | |
| return None, None, "", None, None, None, "" | |
| # ---------------------------------------------------------------- UI | |
| CSS = """ | |
| #title {text-align: center; margin-bottom: 0;} | |
| #links {text-align: center;} | |
| #run_btn {font-size: 1.1rem;} | |
| #note {background: var(--color-accent-soft); border-left: 4px solid var(--color-accent); | |
| border-radius: 6px; padding: 10px 14px;} | |
| """ | |
| with gr.Blocks(title="VideoMaMa++ demo") as demo: | |
| gr.Markdown("# VideoMaMa++: Temporally Consistent Video Matting via Preserve-and-Refine Tokens", elem_id="title") | |
| gr.Markdown( | |
| "**NeurIPS 2026** 路 [Project page](https://cvlab-kaist.github.io/VideoMaMa-PlusPlus) 路 " | |
| "[Paper](https://cvlab-kaist.github.io/VideoMaMa-PlusPlus/source/_NeurIPS_2026__LongMaMa.pdf) 路 " | |
| "[GitHub](https://github.com/cvlab-kaist/VideoMaMa-PlusPlus)", | |
| elem_id="links", | |
| ) | |
| gr.Markdown( | |
| "1. Upload a video (or pick an example) and press **Load video**. " | |
| "2. Click the object on the first frame (or type what it is, e.g. `person`). " | |
| "3. Press **Run matting**: SAM3 tracks the object and VideoMaMa++ turns the mask track into an alpha matte." | |
| ) | |
| gr.Markdown( | |
| f"**Speed vs. quality:** to keep runs short, the demo starts with a reduced setting of " | |
| f"**{DEFAULT_FRAMES} frames at {DEFAULT_SCALE} resolution**, so these results are lower in quality than " | |
| "VideoMaMa++ can produce. For the best mattes, set **Processing resolution** to **100% (original)**. " | |
| "That takes longer and uses more of your GPU quota.", | |
| elem_id="note", | |
| ) | |
| state = gr.State(None) | |
| with gr.Row(): | |
| with gr.Column(): | |
| video_in = gr.Video(label="Input video", sources=["upload"]) | |
| with gr.Row(): | |
| n_frames = gr.Slider(6, MAX_FRAMES, value=DEFAULT_FRAMES, step=1, label="Frames to process", | |
| info="Taken from the start of the video. Press Load video again after changing.") | |
| stride = gr.Slider(1, 4, value=1, step=1, label="Frame stride", | |
| info="Use every k-th frame to cover a longer span at the same cost.") | |
| load_btn = gr.Button("Load video", variant="secondary") | |
| status = gr.Markdown("") | |
| with gr.Column(): | |
| first_frame = gr.Image(label="First frame: click the object", type="pil", interactive=False) | |
| with gr.Row(): | |
| click_kind = gr.Radio(["Positive", "Negative"], value="Positive", label="Click type") | |
| clear_btn = gr.Button("Clear clicks") | |
| text_prompt = gr.Textbox(label="Or a text prompt (overrides clicks)", placeholder="e.g. person, dog, hair") | |
| with gr.Accordion("Settings", open=True): | |
| scale = gr.Radio(list(SCALES), value=DEFAULT_SCALE, label="Processing resolution", | |
| info=f"Share of the original size; 100% is capped at {MAX_LONG_SIDE} px on the long side, " | |
| "the paper's setting. Lower is much faster; fewer frames also helps.") | |
| estimate = gr.Markdown("") | |
| run_btn = gr.Button("Run matting", variant="primary", elem_id="run_btn") | |
| with gr.Row(): | |
| alpha_out = gr.Video(label="Alpha matte (VideoMaMa++)", autoplay=True, loop=True) | |
| green_out = gr.Video(label="Greenscreen composite", autoplay=True, loop=True) | |
| mask_out = gr.Video(label="SAM3 mask track (input guide)", autoplay=True, loop=True) | |
| examples = sorted(str(p) for p in (APP_DIR / "samples").glob("*.mp4")) | |
| if examples: | |
| gr.Examples(examples=examples, inputs=[video_in], label="Examples") | |
| load_btn.click(on_load, [video_in, n_frames, stride, scale, state], [state, first_frame, status, estimate]) | |
| scale.change(on_scale, [state, scale], [estimate]) | |
| first_frame.select(on_click, [state, click_kind], [state, first_frame]) | |
| clear_btn.click(on_clear, [state], [state, first_frame]) | |
| run_btn.click(on_run, [state, text_prompt, scale], [alpha_out, green_out, mask_out, status]) | |
| video_in.change(on_new_video, None, [state, first_frame, status, alpha_out, green_out, mask_out, estimate]) | |
| if __name__ == "__main__": | |
| OUT_DIR.mkdir(parents=True, exist_ok=True) | |
| demo.queue(max_size=20).launch( | |
| css=CSS, | |
| allowed_paths=[str(OUT_DIR), str(APP_DIR / "samples")], | |
| # Never serve the private weights or the model cache. | |
| blocked_paths=[str(_private), os.environ["DIFFSYNTH_MODEL_BASE_PATH"], str(Path.home() / ".cache")], | |
| ) | |