Running VideoMaMa on NVIDIA DGX Spark (ARM64 / GB10)

Accelerated Computing DGX Spark / GB10 User Forum DGX Spark / GB10
1

Running VideoMaMa on NVIDIA DGX Spark (ARM64 / GB10)

Platform: Lenovo ThinkStation PGX — NVIDIA GB10 Blackwell Superchip, 128GB unified memory, DGX OS, CUDA 13.0, Driver 580.142

This guide documents how to get VideoMaMa — a video matting model combining SAM2 mask tracking with Stable Video Diffusion — running on the DGX Spark. Several ARM64-specific issues needed to be resolved that don’t appear on x86 systems.

What is VideoMaMa?

VideoMaMa is a video matting pipeline that:

  1. Uses SAM2 to track an object through video frames based on user click prompts
  2. Uses a fine-tuned Stable Video Diffusion UNet to generate high-quality alpha mattes
  3. Provides a Gradio web UI for interactive use

Checkpoint Structure

Download the following checkpoints and place them as shown:

checkpoints/
├── VideoMaMa/                          # HuggingFace: cvlab-kaist/VideoMaMa
│   ├── dino_projection_mlp.pth
│   └── unet/
│       ├── config.json
│       └── diffusion_pytorch_model.safetensors
├── sam2/                               # Meta: facebook/sam2
│   └── sam2.1_hiera_large.pt
└── stable-video-diffusion-img2vid-xt/  # HuggingFace: stabilityai/stable-video-diffusion-img2vid-xt
    ├── model_index.json
    ├── feature_extractor/
    ├── image_encoder/
    └── vae/

ARM64 Issues and Fixes

Fix 1 — Gradio version incompatibility

Gradio 4.x introduced a streaming WebSocket path in gradio-client that throws Method not implemented on ARM64, silently breaking all UI events including .select() clicks. No JS errors surface until you open DevTools:

stream.ts:185 Method not implemented.
api_info.ts:401 Too many arguments provided for the endpoint.

Fix: Pin to Gradio 3.50.2 + gradio-client 0.6.1 which uses a polling approach that works reliably on ARM64.

RUN uv pip install --no-cache "gradio==3.50.2" "gradio-client==0.6.1" --system

Fix 2 — Hardcoded localhost binding

The upstream app.py binds Gradio to 127.0.0.1, making the UI unreachable from outside the container even with port mapping. Patched to 0.0.0.0.

Fix 3 — gr.SelectData events silently fail on ARM64

Even with Gradio 3.50.2, gr.SelectData click events on gr.Image don’t reach Python on ARM64. Additionally, gr.HTML() strips <script> tags in Gradio 3.50.2 as a security measure, so the JS fix can’t be injected that way.

Fix: Replace .select() with a hidden textbox + hidden button wired to demo.load(_js=...):

  • JS onclick captures pixel coordinates from the image element
  • Coordinates are written to a hidden gr.Textbox
  • A hidden gr.Button is programmatically clicked to trigger the Python backend
  • The JS handler is injected via demo.load(_js=...) which Gradio 3.x executes reliably

Fix 4 — OOM on long videos

VideoMaMa is built on Stable Video Diffusion which has a native 25-frame context window. Passing all frames at once (e.g. 121 frames) causes an out-of-memory crash that can take down the entire system — GB10 uses unified memory shared between CPU and GPU.

Fix: Chunked inference wrapper that processes 25-frame overlapping windows and frees memory between chunks with torch.cuda.empty_cache().

Memory budget on GB10 (128GB unified):

Component Memory
SAM2 + SVD UNet + DINO (idle) ~20GB
Per 25-frame chunk @ 1024×576 fp16 ~51GB
Available headroom ~57GB ✅

Docker Setup

docker-compose.yml

services:
  video-mama:
    build:
      context: .
      dockerfile: Dockerfile
    container_name: video-mama
    volumes:
      - ./checkpoints:/workspace/VideoMaMa/checkpoints
      - pip-cache:/root/.cache
    ports:
      - "7860:7860"
    environment:
      - CUDA_VISIBLE_DEVICES=0
    ipc: host
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

volumes:
  pip-cache:

Dockerfile

# Use existing local SM121-capable torch — no download needed
FROM vllm-node-tf5:latest

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

RUN apt-get update && apt-get install -y --no-install-recommends \
    ffmpeg cmake \
    libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev \
    libswscale-dev libswresample-dev libavfilter-dev pkg-config \
    libx11-dev \
    && rm -rf /var/lib/apt/lists/*

RUN pip install uv

WORKDIR /workspace

RUN git clone https://github.com/cvlab-kaist/VideoMaMa.git /workspace/VideoMaMa

WORKDIR /workspace/VideoMaMa

RUN uv pip install --no-cache git+https://github.com/facebookresearch/sam2.git --system

RUN uv pip install --no-cache --no-deps \
    "git+https://github.com/EasternJournalist/utils3d.git#egg=utils3d" --system

RUN uv pip install --no-cache --no-deps -e . --system

RUN DS_BUILD_OPS=0 uv pip install --no-cache deepspeed==0.17.1 --system || \
    echo "deepspeed skipped, continuing..."

RUN uv pip install --no-cache --system \
    "accelerate>=1.9.0" \
    "diffusers>=0.31.0" \
    "opencv-python-headless>=4.9.0.80" \
    "tokenizers>=0.20.3" \
    "transformers==4.57.0" \
    "numpy>=1.23.5,<2" \
    psutil \
    matplotlib \
    scipy \
    setuptools \
    omegaconf \
    tabulate \
    pandas \
    wandb \
    datasets \
    "peft>0.15" \
    easydict \
    "boto3>=1.37.22" \
    ftfy \
    scikit-image \
    moviepy \
    pyarrow \
    imageio imageio-ffmpeg \
    pycocotools \
    einops safetensors \
    timm

# Install gradio 3.50.2 + gradio-client 0.6.1 — matches the working WSL install.
# Gradio 4.x introduced a streaming WebSocket path in gradio-client that throws
# "Method not implemented" on ARM64, breaking all events including .select().
RUN uv pip install --no-cache "gradio==3.50.2" "gradio-client==0.6.1" --system


# Patch app.py:
#   1. Bind to 0.0.0.0 so the port is reachable from the host
#   2. Disable frpc share link (port is exposed directly via Docker)
#   3. Fix gr.State(None) → gr.State({}) to avoid JS schema serialization
#      errors that silently prevent .select() click events on ARM64
#   4. Replace broken .select() with JS onclick → hidden textbox → gr.Button
#      workaround (Gradio 4.44 .select() silently fails on ARM64 browsers)
RUN python3 - << 'EOF'
path = "/workspace/VideoMaMa/demo/app.py"
src = open(path).read()

src = src.replace('server_name="127.0.0.1"', 'server_name="0.0.0.0"')

# Fix image display size to match video player width
src = src.replace(
    'button {border-radius: 8px !important;}',
    'button {border-radius: 8px !important;}\n'
    '#first_frame_display img { width: 100% !important; height: 100% !important; object-fit: contain !important; }\n'
    '#first_frame_display { min-height: 440px !important; }'
)
src = src.replace('share=True', 'share=False')
# Hide the frames slider — we now process all frames always
src = src.replace(
    'num_frames_slider = gr.Slider(',
    'num_frames_slider = gr.Slider(visible=False,'
)
src = src.replace('video_state = gr.State(None)', 'video_state = gr.State({})')
src = src.replace('if video_state is None or', 'if not video_state or')
src = src.replace('if video_state is None:', 'if not video_state:')

# Patch videomama_wrapper.py to process frames in 25-frame chunks
# (VideoMaMa is based on SVD which has a fixed 25-frame context window;
#  passing all frames at once OOMs the system)
wrapper_patch = '''
import torch
import numpy as np
from PIL import Image

_original_videomama = videomama

def videomama(pipeline, frames_np, mask_frames_np, chunk_size=25, overlap=4):
    """
    Chunked VideoMaMa inference — processes video in overlapping windows
    to avoid OOM on long videos. chunk_size=25 matches SVD context window.
    """
    n = len(frames_np)
    if n <= chunk_size:
        return _original_videomama(pipeline, frames_np, mask_frames_np)

    print(f"Processing {n} frames in chunks of {chunk_size} (overlap={overlap})...")
    output_frames = []
    start = 0
    while start < n:
        end = min(start + chunk_size, n)
        chunk_frames = frames_np[start:end]
        chunk_masks = mask_frames_np[start:end]
        print(f"  Chunk {start}-{end} ({len(chunk_frames)} frames)...")
        chunk_out = _original_videomama(pipeline, chunk_frames, chunk_masks)
        # Skip overlap frames from previous chunk to avoid seam artifacts
        skip = overlap if start > 0 else 0
        output_frames.extend(chunk_out[skip:])
        # Advance by chunk_size minus overlap for smooth transitions
        start += chunk_size - overlap
        # Free memory between chunks
        torch.cuda.empty_cache()

    return output_frames[:n]
'''

# Append the chunked wrapper to videomama_wrapper.py
vw_path = "/workspace/VideoMaMa/demo/videomama_wrapper.py"
vw_src = open(vw_path).read()
vw_src += "\n" + wrapper_patch
open(vw_path, "w").write(vw_src)
print("[DONE] Patched videomama_wrapper.py with chunked inference")

# Always process all frames at original FPS — no sampling, no FPS adjustment
old_extract = '''    # If video has more frames than max_frames, randomly sample
    if len(all_frames) > max_frames:
        print(f"Video has {len(all_frames)} frames, randomly sampling {max_frames} frames...")
        # Sort indices to maintain temporal order
        sampled_indices = sorted(np.random.choice(len(all_frames), max_frames, replace=False))
        frames = [all_frames[i] for i in sampled_indices]
        print(f"Sampled frame indices: {sampled_indices}")
        
        # Adjust FPS to maintain normal playback speed
        # If we sampled N frames from M total frames, adjust FPS proportionally
        adjusted_fps = original_fps * (len(frames) / len(all_frames))
    else:
        frames = all_frames
        adjusted_fps = original_fps
        print(f"Video has {len(frames)} frames (≤ {max_frames}), using all frames")
    
    print(f"Using {len(frames)} frames from video (Original FPS: {original_fps:.2f}, Adjusted FPS: {adjusted_fps:.2f})")
    
    return frames, adjusted_fps'''

new_extract = '''    frames = all_frames
    print(f"Using all {len(frames)} frames at original FPS: {original_fps:.2f}")
    return frames, original_fps'''

src = src.replace(old_extract, new_extract)

# Patch sam_refine to accept coords from hidden textbox instead of gr.SelectData
old_fn = '''def sam_refine(video_state, point_prompt, click_state, evt: gr.SelectData):
    """
    Add click and update mask on first frame
    
    Args:
        video_state: Dictionary with video data
        point_prompt: "Positive" or "Negative"
        click_state: [[points], [labels]]
        evt: Gradio SelectData event with click coordinates
    """
    if not video_state or "frames" not in video_state:
        return None, video_state, click_state
    
    # Add new click
    x, y = evt.index[0], evt.index[1]'''

new_fn = '''def sam_refine(video_state, point_prompt, click_state, click_coords):
    """
    Add click and update mask on first frame.
    click_coords: JSON string "[x, y]" from hidden textbox via JS onclick.
    """
    if not video_state or "frames" not in video_state:
        return None, video_state, click_state
    if not click_coords:
        return None, video_state, click_state
    import json as _json
    try:
        coords = _json.loads(click_coords)
        x, y = int(coords[0]), int(coords[1])
    except Exception:
        return None, video_state, click_state'''

src = src.replace(old_fn, new_fn)

# Replace the gr.Image block and .select() wiring with JS onclick version
old_image = '''            first_frame_display = gr.Image(
                label="First Frame",
                type="pil",
                interactive=True
            )'''

new_image = '''            first_frame_display = gr.Image(
                label="First Frame",
                type="pil",
                interactive=True,
                elem_id="first_frame_display"
            )
            click_coords = gr.Textbox(value="", visible=False, elem_id="click_coords")
            click_trigger = gr.Button("__click_trigger__", visible=False, elem_id="click_trigger")'''

src = src.replace(old_image, new_image)

# Replace .select() wiring with .click() on the hidden trigger button
old_select = '''    first_frame_display.select(
        fn=sam_refine,
        inputs=[video_state, point_prompt, click_state],
        outputs=[first_frame_display, video_state, click_state]
    )'''

new_select = '''    # Use _js to inject click handler via Gradio 3.x's built-in JS support.
    # gr.HTML() strips <script> tags in Gradio 3.50.2, so we attach the handler
    # via the load event's _js parameter which executes reliably.
    click_trigger.click(
        fn=sam_refine,
        inputs=[video_state, point_prompt, click_state, click_coords],
        outputs=[first_frame_display, video_state, click_state]
    )

    # Inject JS via demo.load which runs after the page is ready
    def noop():
        return None
    demo.load(
        fn=noop,
        inputs=[],
        outputs=[],
        _js="""
        () => {
            let attachedImg = null;
            function attachClickHandler() {
                const img = document.querySelector('#first_frame_display img');
                const coordBox = document.querySelector('#click_coords textarea');
                const btn = document.querySelector('#click_trigger');
                if (!img || !coordBox || !btn) return;
                if (img === attachedImg) return;
                attachedImg = img;
                img.style.cursor = 'crosshair';
                img.addEventListener('click', function(e) {
                    const rect = img.getBoundingClientRect();
                    const scaleX = img.naturalWidth / rect.width;
                    const scaleY = img.naturalHeight / rect.height;
                    const x = Math.round((e.clientX - rect.left) * scaleX);
                    const y = Math.round((e.clientY - rect.top) * scaleY);
                    coordBox.value = JSON.stringify([x, y]);
                    coordBox.dispatchEvent(new Event('input', {bubbles: true}));
                    setTimeout(() => btn.click(), 50);
                });
                console.log('VideoMaMa: click handler attached');
            }
            setInterval(attachClickHandler, 300);
            attachClickHandler();
        }
        """
    )'''

src = src.replace(old_select, new_select)

open(path, "w").write(src)
print("[DONE] Patched app.py with JS onclick workaround")
EOF

ENV PYTHONPATH=/workspace/VideoMaMa

CMD ["python3", "/workspace/VideoMaMa/demo/app.py"]

Usage

# Build and start
cd VideoMaMa
docker compose build
docker compose up -d video-mama

# Access the UI
http://<your-dgx-ip>:7860

# Monitor inference
docker logs -f video-mama

Workflow:

  1. Upload a video
  2. Click Load Video
  3. Click on the object in the first frame (crosshair cursor)
  4. Add Positive/Negative clicks to refine the mask
  5. Click 🚀 Run Matting

Performance (DGX Spark, GB10)

Stage Details Time
SAM2 tracking 121 frames ~40s
VideoMaMa inference per 25-frame chunk ~varies by video resolution
Total memory at inference models + one chunk ~71GB / 128GB

Notes

  • The base image vllm-node-tf5 is a community SM121-optimized image for the GB10 Blackwell. Any ARM64 CUDA 13 base with PyTorch should work as a substitute — just ensure SM121 support for the GB10.
  • Chunked inference with overlap=4 produces smooth transitions. Increase overlap if you see seam artifacts between chunks.
  • Output videos (matting result, mask track, greenscreen composite) are saved to ./outputs/ inside the container and served via the Gradio UI.
  • The frpc ARM64 binary for Gradio’s share link feature is unavailable — expose the port directly via Docker instead.
3 Likes