From 2e4e8afc72e4a3aa1ff9ae70a96f0f4c2b671c5f Mon Sep 17 00:00:00 2001 From: Bryon Lewis Date: Tue, 26 May 2026 11:18:59 -0400 Subject: [PATCH] basic sam3 upgrade --- .../example-docker-containers/.env.example | 10 +- .../example-docker-containers/Dockerfile | 2 +- .../example-docker-containers/README.md | 55 +- .../SAM3Demo/SAM3Demo.py | 692 ++++++++++++------ .../SAM3Demo/SAM3Demo.xml | 4 +- .../SAM3Demo/test_sam31_loader.py | 81 ++ .../download_sam3_ckpts.sh | 14 +- 7 files changed, 621 insertions(+), 237 deletions(-) create mode 100644 dive-dsa-slicer/example-docker-containers/SAM3Demo/test_sam31_loader.py diff --git a/dive-dsa-slicer/example-docker-containers/.env.example b/dive-dsa-slicer/example-docker-containers/.env.example index d5300036..107778ab 100644 --- a/dive-dsa-slicer/example-docker-containers/.env.example +++ b/dive-dsa-slicer/example-docker-containers/.env.example @@ -1,6 +1,14 @@ # Copy to .env and set your Hugging Face token (read access). -# Request access: https://huggingface.co/facebook/sam3 +# Request access: https://huggingface.co/facebook/sam3.1 HF_TOKEN=hf_your_token_here # Optional image tag # TAG=latest + +# Optional SAM3Demo GPU memory tuning (see README.md) +# SAM3_MAX_NUM_OBJECTS=2 +# SAM3_GROUNDING_BATCH_SIZE=1 +# SAM3_POSTPROCESS_BATCH_SIZE=2 +# SAM3_PROPAGATION_CHUNK_SIZE=2 +# SAM3_OFFLOAD_VIDEO_TO_CPU=true +# SAM3_USE_FA3=true diff --git a/dive-dsa-slicer/example-docker-containers/Dockerfile b/dive-dsa-slicer/example-docker-containers/Dockerfile index 7fa78163..7af0111e 100644 --- a/dive-dsa-slicer/example-docker-containers/Dockerfile +++ b/dive-dsa-slicer/example-docker-containers/Dockerfile @@ -51,7 +51,7 @@ RUN \ # install step \ pip install --no-cache-dir opencv-python-headless && \ # SAM dependencies - pip install --no-cache-dir sam2 sam3 huggingface_hub psutil pycocotools torch pillow && \ + pip install --no-cache-dir sam2 'sam3 @ git+https://github.com/facebookresearch/sam3.git' huggingface_hub psutil pycocotools torch pillow einops && \ # clean up \ rm -rf /root/.cache/pip/* diff --git a/dive-dsa-slicer/example-docker-containers/README.md b/dive-dsa-slicer/example-docker-containers/README.md index e91623f9..ac692abe 100644 --- a/dive-dsa-slicer/example-docker-containers/README.md +++ b/dive-dsa-slicer/example-docker-containers/README.md @@ -38,7 +38,7 @@ Uses SAM 2 video propagation from an existing track bounding box or mask. Requir ## SAM3Demo -Uses [SAM 3](https://github.com/facebookresearch/sam3) open-vocabulary text prompts to segment all matching instances on the current DIVE frame and ingest them as new tracks with masks. +Uses [SAM 3.1](https://github.com/facebookresearch/sam3) (Object Multiplex) open-vocabulary text prompts to segment all matching instances on the current DIVE frame and ingest them as new tracks with masks. ### Parameters @@ -46,20 +46,55 @@ Uses [SAM 3](https://github.com/facebookresearch/sam3) open-vocabulary text prom - **DIVEFrameId**: frame index (auto-filled from the DIVE UI when launched from the app). - **DIVETrackType**: classification label for new tracks. - **ConfidenceThreshold**: minimum detection score (default `0.5`). -- **TrackingFrames**: number of frames to track forward from the prompt frame (`1` = single frame only; values greater than `1` use SAM3 video propagation, similar to SAM2Demo). +- **TrackingFrames**: number of frames to track forward from the prompt frame (`1` = single frame only; values greater than `1` propagate masks with SAM 3.1 video tracking, similar to SAM2Demo). ### Device selection -Single-frame inference (`TrackingFrames` = 1) uses CUDA when available, otherwise CPU. Video propagation (`TrackingFrames` > 1) requires CUDA because the SAM3 video predictor loads the model on GPU. Apple MPS is detected but SAM3 runs on CPU for single-frame mode only. - -If you see `RuntimeError: The NVIDIA driver on your system is too old` while the log says `Using inference device: cpu`, PyTorch cannot use the GPU (driver/CUDA mismatch). Single-frame mode should still work on CPU after rebuilding the image (the container patches SAM3’s hardcoded CUDA precompute paths). For GPU inference, update the host NVIDIA driver to match the PyTorch CUDA version in the image, or rebuild with a CPU-only PyTorch wheel. +SAM 3.1 multiplex inference requires CUDA. If you see `RuntimeError: The NVIDIA driver on your system is too old`, update the host NVIDIA driver to match the PyTorch CUDA version in the image. ### Hugging Face checkpoints -SAM 3 weights are gated on Hugging Face. Request access to [facebook/sam3](https://huggingface.co/facebook/sam3), then either: +SAM 3.1 weights are gated on Hugging Face. Request access to [facebook/sam3.1](https://huggingface.co/facebook/sam3.1), then either: + +1. Copy `.env.example` to `.env`, set `HF_TOKEN`, and run `docker compose build --no-cache` to prefetch `sam3.1_multiplex.pt` into `/opt/SAM3/models`, or +2. Set `HF_TOKEN` (or run `hf auth login`) at runtime so `build_sam3_predictor` can download on first use, or +3. Set `SAM3_CHECKPOINT` to a local `sam3.1_multiplex.pt` path inside the container. + +The image installs `sam3` from GitHub (SAM 3.1 APIs are not yet on PyPI) and downloads the text tokenizer vocab (`bpe_simple_vocab_16e6.txt.gz`) into `/opt/SAM3/assets`. Override with `SAM3_BPE_PATH` if needed. + +### Reducing GPU memory + +SAM 3.1 multiplex loads a large detector + tracker stack on GPU. The checkpoint size is similar to SAM 3, but peak VRAM is usually higher because every job uses the full multiplex path (including single-frame jobs). + +**Job parameters (no code changes):** + +- Set **TrackingFrames** to `1` when you only need one frame (still uses multiplex, but avoids propagation memory). +- Use a higher **ConfidenceThreshold** to drop weak detections (fewer masks/objects held in memory). + +**Container environment variables** (on GPUs ≤20 GiB, conservative defaults apply automatically): + +| Variable | Typical default (≤20 GiB) | Effect | +|----------|---------------------------|--------| +| `SAM3_MAX_NUM_OBJECTS` | `4` | Cap parallel object slots (was 16 in upstream sam3). | +| `SAM3_GROUNDING_BATCH_SIZE` | `2` | Detector grounding chunk size (upstream default `16` often OOMs on 16 GiB). | +| `SAM3_POSTPROCESS_BATCH_SIZE` | `2` | Postprocess batch size. | +| `SAM3_PROPAGATION_CHUNK_SIZE` | `3` | Propagate a few frames per stream, then `empty_cache` (set `0` for one shot). | +| `SAM3_OFFLOAD_VIDEO_TO_CPU` | `true` | Keep decoded frames on CPU. | +| `SAM3_USE_FA3` | `true` on Ampere+ only | Flash Attention 3 when supported. | + +Example if propagation still OOMs on a 16 GiB GPU: + +```bash +SAM3_MAX_NUM_OBJECTS=2 \ +SAM3_GROUNDING_BATCH_SIZE=1 \ +SAM3_PROPAGATION_CHUNK_SIZE=2 \ +SAM3_OFFLOAD_VIDEO_TO_CPU=true +``` + +Also set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` (enabled by default in `SAM3Demo.py`). + +### Pre-Ampere GPUs (Turing / Volta / Pascal) -1. Copy `.env.example` to `.env`, set `HF_TOKEN`, and run `docker compose build --no-cache` to prefetch weights into `/opt/SAM3/models`, or -2. Set `HF_TOKEN` (or run `hf auth login`) at runtime so `build_sam3_image_model` can download on first use, or -3. Set `SAM3_CHECKPOINT` to a local `sam3.pt` path inside the container. +If you see `RuntimeError: No available kernel` during propagation, the GPU is below CUDA capability 8.0 (no Flash Attention). `SAM3Demo.py` patches sam3 to use math SDPA and runs inference in fp32 on those GPUs. Ensure you are on the latest `SAM3Demo.py` from this repo. -The image build also downloads the text tokenizer vocab (`bpe_simple_vocab_16e6.txt.gz`) into `/opt/SAM3/assets` (the PyPI `sam3` package does not include it). Override with `SAM3_BPE_PATH` if needed. +**Hard limits:** Model weights still occupy several GB on GPU; there is no CPU fallback for SAM 3.1 multiplex in this demo. For the lowest VRAM on single-frame text prompts only, the older SAM 3 image-model path (`sam3.pt`) used less memory but is not SAM 3.1. diff --git a/dive-dsa-slicer/example-docker-containers/SAM3Demo/SAM3Demo.py b/dive-dsa-slicer/example-docker-containers/SAM3Demo/SAM3Demo.py index 70d78688..77eda8b0 100644 --- a/dive-dsa-slicer/example-docker-containers/SAM3Demo/SAM3Demo.py +++ b/dive-dsa-slicer/example-docker-containers/SAM3Demo/SAM3Demo.py @@ -1,39 +1,253 @@ """ -SAM3 text-prompt segmentation task for DIVE. +SAM 3.1 text-prompt segmentation task for DIVE. Segments all instances matching a text prompt on a video frame and optionally -propagates masks forward through subsequent frames. Creates one track per -instance and uploads masks/annotations to Girder. +propagates masks forward through subsequent frames using the SAM 3.1 Object +Multiplex video predictor. Creates one track per instance and uploads +masks/annotations to Girder. -Uses GPU (CUDA) for video propagation when available; single-frame inference -also runs on CPU. +Requires CUDA (SAM 3.1 multiplex inference). """ +import gc import importlib.util import json import logging import os import pprint -import subprocess import tempfile from contextlib import nullcontext from pathlib import Path -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional import cv2 import girder_client import numpy as np import torch -from PIL import Image from pycocotools import mask as mask_utils -from ctk_cli import CLIArgumentParser # noqa: I004 -from slicer_cli_web import ctk_cli_adjustment # noqa +try: + from ctk_cli import CLIArgumentParser # noqa: I004 + from slicer_cli_web import ctk_cli_adjustment # noqa +except ImportError: # pragma: no cover + # Allows importing this module outside the slicer runtime (useful for local sanity tests). + CLIArgumentParser = None # type: ignore[assignment] + ctk_cli_adjustment = None # type: ignore[assignment] logging.basicConfig(level=logging.CRITICAL) +# Reduce CUDA fragmentation on long propagation runs (safe no-op if unsupported). +os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True') + _EXAMPLE_ROOT = Path(__file__).resolve().parent.parent +def _patch_sam3_offload_state_to_cpu_kwarg() -> None: + """ + Compatibility patch for SAM 3.1 predictor/model init_state mismatch. + + Some SAM3 releases/pins call `init_state(offload_state_to_cpu=...)`, while + `Sam3MultiplexTrackingWithInteractivity.init_state` may not accept that kwarg. + We ignore it to avoid a hard TypeError during `start_session`. + """ + + try: + import inspect + from sam3.model.sam3_multiplex_tracking import ( + Sam3MultiplexTrackingWithInteractivity, + ) + except Exception: + # If sam3 isn't installed or module paths differ, just skip. + return + + try: + if "offload_state_to_cpu" in inspect.signature( + Sam3MultiplexTrackingWithInteractivity.init_state + ).parameters: + return + except Exception: + # If we can't inspect the signature, still apply the safe wrapper. + pass + + orig_init_state = Sam3MultiplexTrackingWithInteractivity.init_state + + def wrapped_init_state(self, *args, **kwargs): + kwargs.pop("offload_state_to_cpu", None) + return orig_init_state(self, *args, **kwargs) + + Sam3MultiplexTrackingWithInteractivity.init_state = wrapped_init_state + + +def _patch_sam3_multiplex_checkpoint_preload() -> None: + """ + SAM 3.1 HF checkpoints use `detector.*` and `tracker.model.*` keys for the + combined demo model. `build_sam3_multiplex_video_predictor` also builds a + tracker-only sub-model and incorrectly preloads the same checkpoint there, + which produces thousands of missing/unexpected key warnings and leaves + weights uninitialized. Skip that preload; load once into the full model. + """ + + try: + import sam3.model_builder as model_builder + except Exception: + return + + if getattr(model_builder, '_sam3demo_skip_tracker_preload', False): + return + + original_build_tracker = model_builder.build_sam3_multiplex_video_model + + def build_tracker_without_checkpoint_preload( + checkpoint_path=None, + load_from_HF=True, + **kwargs, + ): + return original_build_tracker( + checkpoint_path=None, + load_from_HF=False, + **kwargs, + ) + + model_builder.build_sam3_multiplex_video_model = build_tracker_without_checkpoint_preload + model_builder._sam3demo_skip_tracker_preload = True + + +def _patch_sam3_functional_attention_sdpa() -> None: + """ + On pre-Ampere GPUs, sam3's functional_attention forces FLASH_ATTENTION SDPA only + when use_fa3=False, which raises "No available kernel". Use MATH/efficient SDPA. + """ + + try: + import sam3.model.decoder as decoder_mod + import torch.nn.functional as torchF + from sam3.sam.rope import apply_rotary_enc, apply_rotary_enc_real + from torch import Tensor + from torch.nn.attention import sdpa_kernel, SDPBackend + except Exception: + return + + if getattr(decoder_mod, '_sam3demo_sdpa_patched', False): + return + + def functional_attention( + q: Tensor, + k: Tensor, + v: Tensor, + *, + dropout: float, + num_heads: int, + num_k_exclude_rope: int = 0, + freqs_cis: Optional[Tensor] = None, + freqs_cis_real: Optional[Tensor] = None, + freqs_cis_imag: Optional[Tensor] = None, + use_fa3: bool = False, + use_rope_real: bool = False, + rope_k_repeat: bool, + ): + b, n, cq = q.shape + _, m, ck = k.shape + _, _, cv = v.shape + if b > 1: + assert k.shape[0] == v.shape[0] == b + else: + assert k.shape[0] == b == 1, f'{q.shape=} {k.shape=} {v.shape=}' + assert v.shape[1] == m + + q = q.reshape(b, n, num_heads, cq // num_heads).transpose(1, 2) + k = k.reshape(b, m, num_heads, ck // num_heads).transpose(1, 2) + v = v.reshape(v.shape[0], m, num_heads, cv // num_heads).transpose(1, 2) + + if freqs_cis is not None: + num_k_rope = k.size(-2) - num_k_exclude_rope + if use_rope_real: + q, k[:, :, :num_k_rope] = apply_rotary_enc_real( + q, + k[:, :, :num_k_rope], + freqs_cis_real=freqs_cis_real, + freqs_cis_imag=freqs_cis_imag, + repeat_freqs_k=rope_k_repeat, + ) + else: + q, k[:, :, :num_k_rope] = apply_rotary_enc( + q, + k[:, :, :num_k_rope], + freqs_cis, + repeat_freqs_k=rope_k_repeat, + ) + + if use_fa3 and _cuda_supports_flash_attention_3(): + from sam3.perflib.fa3 import flash_attn_func + + assert dropout == 0.0 + out = flash_attn_func( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ) + elif _cuda_supports_flash_attention_3(): + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + out = torchF.scaled_dot_product_attention( + q, k, v, dropout_p=dropout + ) + else: + with sdpa_kernel( + [SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION] + ): + out = torchF.scaled_dot_product_attention( + q, k, v, dropout_p=dropout + ) + out = out.transpose(1, 2) + out = out.reshape(b, n, cv) + return out + + decoder_mod.functional_attention = functional_attention + decoder_mod._sam3demo_sdpa_patched = True + + +def _configure_pytorch_sdpa_for_gpu() -> None: + """Prefer math/mem-efficient attention on GPUs without Flash Attention.""" + if not torch.cuda.is_available() or _cuda_supports_flash_attention_3(): + return + torch.backends.cuda.enable_math_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(True) + try: + torch.backends.cuda.enable_flash_sdp(False) + except Exception: + pass + + +def _cuda_inference_context(): + """bf16 autocast on Ampere+; fp32 inference on older GPUs for SDPA compatibility.""" + if _cuda_supports_flash_attention_3(): + return torch.autocast(device_type='cuda', dtype=torch.bfloat16) + return nullcontext() + + +def _prepare_sam31_multiplex_checkpoint(checkpoint_path: Path) -> dict: + """Load a SAM 3.1 multiplex checkpoint and apply OSS key remaps when needed.""" + ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=True) + if isinstance(ckpt, dict) and isinstance(ckpt.get('model'), dict): + ckpt = ckpt['model'] + + if not isinstance(ckpt, dict): + raise RuntimeError(f'Unexpected checkpoint format in {checkpoint_path}') + + needs_remap = any( + key.startswith('sam3_model.') or key.startswith('sam2_predictor.') + for key in ckpt + ) + if not needs_remap: + return ckpt + + remapped: dict = {} + for key, value in ckpt.items(): + new_key = key + if key.startswith('sam3_model.'): + new_key = 'detector.' + key[len('sam3_model.') :] + elif key.startswith('sam2_predictor.'): + new_key = 'tracker.' + key[len('sam2_predictor.') :] + remapped[new_key] = value + return remapped + + def _load_sam2_helpers(): """Reuse frame extraction and mask upload helpers from SAM2Demo.""" module_path = _EXAMPLE_ROOT / 'SAM2Demo' / 'SAM2Demo.py' @@ -43,7 +257,12 @@ def _load_sam2_helpers(): return module.process_masks_folder, module.extract_frames -process_masks_folder, extract_frames = _load_sam2_helpers() +try: + process_masks_folder, extract_frames = _load_sam2_helpers() +except Exception: + # Allows local import for sanity tests without full slicer dependencies (e.g. hydra). + process_masks_folder = None + extract_frames = None def resolve_device() -> str: @@ -55,14 +274,99 @@ def resolve_device() -> str: return 'cpu' -def model_device_string(device: str) -> str: - """SAM3 builders only accept cuda or cpu device strings.""" - if device == 'mps': - print('MPS detected but SAM3 uses CPU for this device string; continuing on CPU.') - return 'cpu' - if device in ('cuda', 'cpu'): - return device - return 'cpu' +def _env_bool(name: str, default: bool) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in ('1', 'true', 'yes', 'on') + + +def _env_int(name: str, default: int, minimum: int = 1) -> int: + raw = os.environ.get(name) + if raw is None: + return default + try: + return max(minimum, int(raw)) + except ValueError: + return default + + +def _cuda_supports_flash_attention_3() -> bool: + if not torch.cuda.is_available(): + return False + try: + major, _minor = torch.cuda.get_device_capability() + return major >= 8 + except Exception: + return False + + +def _gpu_total_gib() -> Optional[float]: + if not torch.cuda.is_available(): + return None + try: + return torch.cuda.get_device_properties(0).total_memory / (1024**3) + except Exception: + return None + + +def sam3_memory_settings(tracking_frames: int) -> dict: + """ + Runtime memory knobs (override via environment variables). + + SAM3_MAX_NUM_OBJECTS: cap parallel object slots. + SAM3_GROUNDING_BATCH_SIZE / SAM3_POSTPROCESS_BATCH_SIZE: detector batching (defaults are lower on <=20 GiB GPUs). + SAM3_PROPAGATION_CHUNK_SIZE: propagate N frames per stream call, then empty CUDA cache (0 = one shot). + SAM3_OFFLOAD_VIDEO_TO_CPU: keep decoded frames on CPU. + SAM3_USE_FA3: Flash Attention 3 when the GPU supports it (Ampere+). + """ + gpu_gib = _gpu_total_gib() + conservative = gpu_gib is not None and gpu_gib <= 20.0 + + if conservative: + default_max_objects = 4 + default_grounding_batch = 2 + default_postprocess_batch = 2 + default_prop_chunk = 3 if tracking_frames > 1 else 0 + else: + default_max_objects = 8 if tracking_frames <= 1 else 12 + default_grounding_batch = 4 if tracking_frames > 1 else 8 + default_postprocess_batch = 4 if tracking_frames > 1 else 8 + default_prop_chunk = 0 + + max_num_objects = _env_int('SAM3_MAX_NUM_OBJECTS', default_max_objects, minimum=1) + grounding_batch_size = _env_int( + 'SAM3_GROUNDING_BATCH_SIZE', default_grounding_batch, minimum=1 + ) + postprocess_batch_size = _env_int( + 'SAM3_POSTPROCESS_BATCH_SIZE', default_postprocess_batch, minimum=1 + ) + propagation_chunk_size = _env_int( + 'SAM3_PROPAGATION_CHUNK_SIZE', default_prop_chunk, minimum=0 + ) + offload_video = _env_bool('SAM3_OFFLOAD_VIDEO_TO_CPU', default=True) + use_fa3 = _env_bool('SAM3_USE_FA3', default=_cuda_supports_flash_attention_3()) + return { + 'max_num_objects': max_num_objects, + 'grounding_batch_size': grounding_batch_size, + 'postprocess_batch_size': postprocess_batch_size, + 'propagation_chunk_size': propagation_chunk_size, + 'offload_video_to_cpu': offload_video, + 'use_fa3': use_fa3, + 'gpu_gib': gpu_gib, + } + + +def tune_multiplex_predictor_memory(predictor, mem: dict) -> None: + """Lower peak VRAM on the built multiplex model (sam3 defaults batch size to 16).""" + model = getattr(predictor, 'model', None) + if model is None: + return + model.max_num_objects = mem['max_num_objects'] + model.batched_grounding_batch_size = mem['grounding_batch_size'] + model.postprocess_batch_size = mem['postprocess_batch_size'] + if mem['grounding_batch_size'] <= 1: + model.use_batched_grounding = False def allocate_track_ids(existing_tracks: list, count: int) -> List[str]: @@ -148,39 +452,6 @@ def save_and_record_mask( track['end'] = max(track['end'], frame_idx) -def run_sam3_text_inference( - image_path: Path, - text_prompt: str, - confidence: float, - device: str, - checkpoint_path: Optional[Path], -) -> dict: - """Run SAM3 image inference with a text prompt on a single frame.""" - from sam3.model_builder import build_sam3_image_model - from sam3.model.sam3_image_processor import Sam3Processor - - load_from_hf = checkpoint_path is None - ckpt = str(checkpoint_path) if checkpoint_path else None - model_device = model_device_string(device) - print(f'Loading SAM3 image model on {model_device} with checkpoint: {ckpt or "HF Hub"}.') - - model = build_sam3_image_model( - device=model_device, - checkpoint_path=ckpt, - load_from_HF=load_from_hf, - bpe_path=resolve_bpe_path(), - ) - processor = Sam3Processor( - model, - device=model_device, - confidence_threshold=confidence, - ) - - image = Image.open(image_path).convert('RGB') - state = processor.set_image(image) - return processor.set_text_prompt(prompt=text_prompt, state=state) - - def _masks_from_outputs(outputs: dict) -> Dict[int, np.ndarray]: """Parse SAM3 video predictor outputs into {obj_id: binary_mask}.""" obj_ids = outputs.get('out_obj_ids', []) @@ -202,6 +473,22 @@ def _masks_from_outputs(outputs: dict) -> Dict[int, np.ndarray]: return masks +def _consume_propagation_stream( + predictor, + stream_request: dict, +) -> Dict[int, Dict[int, np.ndarray]]: + mask_dict: Dict[int, Dict[int, np.ndarray]] = {} + for response in predictor.handle_stream_request(stream_request): + frame_idx = response.get('frame_index') + if frame_idx is None: + continue + outputs = response.get('outputs', {}) + masks = _masks_from_outputs(outputs) + if masks: + mask_dict[int(frame_idx)] = masks + return mask_dict + + def collect_video_propagation( predictor, session_id: str, @@ -209,37 +496,59 @@ def collect_video_propagation( start_frame_index: int = 0, max_frame_num_to_track: int, confidence: float, + propagation_chunk_size: int = 0, ) -> Dict[int, Dict[int, np.ndarray]]: """ Propagate a text prompt forward through a frame directory session. + When propagation_chunk_size > 0, propagation is split into smaller stream + calls with torch.cuda.empty_cache() between chunks to reduce peak VRAM. + Returns: Mapping of clip-relative frame index -> {sam3_obj_id: mask}. """ mask_dict: Dict[int, Dict[int, np.ndarray]] = {} - stream_request = { - 'type': 'propagate_in_video', - 'session_id': session_id, - 'propagation_direction': 'forward', - 'start_frame_index': start_frame_index, - 'max_frame_num_to_track': max_frame_num_to_track, - 'output_prob_thresh': confidence, - } - for response in predictor.handle_stream_request(stream_request): - frame_idx = response.get('frame_index') - if frame_idx is None: - continue - outputs = response.get('outputs', {}) - masks = _masks_from_outputs(outputs) - if masks: - mask_dict[int(frame_idx)] = masks + end_limit = start_frame_index + max_frame_num_to_track + + if propagation_chunk_size <= 0: + stream_request = { + 'type': 'propagate_in_video', + 'session_id': session_id, + 'propagation_direction': 'forward', + 'start_frame_index': start_frame_index, + 'max_frame_num_to_track': max_frame_num_to_track, + 'output_prob_thresh': confidence, + } + mask_dict.update(_consume_propagation_stream(predictor, stream_request)) + else: + current = start_frame_index + while current <= end_limit: + frames_left = end_limit - current + 1 + chunk_frames = min(propagation_chunk_size, frames_left) + stream_request = { + 'type': 'propagate_in_video', + 'session_id': session_id, + 'propagation_direction': 'forward', + 'start_frame_index': current, + 'max_frame_num_to_track': chunk_frames - 1, + 'output_prob_thresh': confidence, + } + print( + f'Propagating frames {current}–{current + chunk_frames - 1} ' + f'(chunk size {chunk_frames})...' + ) + mask_dict.update(_consume_propagation_stream(predictor, stream_request)) + current += chunk_frames + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.empty_cache() if torch.cuda.is_available(): torch.cuda.synchronize() return mask_dict -def run_sam3_video_text_inference( +def run_sam3_text_inference( frame_dir: Path, prompt_frame_index: int, text_prompt: str, @@ -247,34 +556,56 @@ def run_sam3_video_text_inference( confidence: float, checkpoint_path: Optional[Path], ) -> Dict[int, Dict[int, np.ndarray]]: - """Run SAM3 video text prompt + forward propagation on an extracted frame clip.""" - from sam3.model_builder import build_sam3_video_predictor - + """Run SAM 3.1 text prompt on a frame clip; propagate when tracking_frames > 1.""" + _configure_pytorch_sdpa_for_gpu() + _patch_sam3_offload_state_to_cpu_kwarg() + _patch_sam3_multiplex_checkpoint_preload() + _patch_sam3_functional_attention_sdpa() + from sam3.model_builder import build_sam3_multiplex_video_predictor + + mem = sam3_memory_settings(tracking_frames) ckpt = str(checkpoint_path) if checkpoint_path else None - predictor = build_sam3_video_predictor( + print( + f'Loading SAM 3.1 multiplex predictor with checkpoint: {ckpt or "HF Hub (facebook/sam3.1)"}.' + ) + gpu_note = f', gpu={mem["gpu_gib"]:.1f}GiB' if mem.get('gpu_gib') else '' + print( + 'SAM 3.1 memory settings: ' + f'max_num_objects={mem["max_num_objects"]}, ' + f'grounding_batch_size={mem["grounding_batch_size"]}, ' + f'postprocess_batch_size={mem["postprocess_batch_size"]}, ' + f'propagation_chunk_size={mem["propagation_chunk_size"]}, ' + f'offload_video_to_cpu={mem["offload_video_to_cpu"]}, ' + f'use_fa3={mem["use_fa3"]}' + f'{gpu_note}' + ) + predictor = build_sam3_multiplex_video_predictor( checkpoint_path=ckpt, bpe_path=resolve_bpe_path(), + max_num_objects=mem['max_num_objects'], + use_fa3=mem['use_fa3'], + use_rope_real=mem['use_fa3'], async_loading_frames=False, + default_output_prob_thresh=confidence, ) + tune_multiplex_predictor_memory(predictor, mem) - autocast_ctx = ( - torch.autocast(device_type='cuda', dtype=torch.bfloat16) - if torch.cuda.is_available() - else nullcontext() - ) - + mask_dict: Dict[int, Dict[int, np.ndarray]] = {} session_id = None try: - with autocast_ctx: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + with _cuda_inference_context(): response = predictor.handle_request( { 'type': 'start_session', 'resource_path': str(frame_dir), + 'offload_video_to_cpu': mem['offload_video_to_cpu'], } ) session_id = response['session_id'] - predictor.handle_request( + prompt_response = predictor.handle_request( { 'type': 'add_prompt', 'session_id': session_id, @@ -283,14 +614,22 @@ def run_sam3_video_text_inference( 'output_prob_thresh': confidence, } ) - - mask_dict = collect_video_propagation( - predictor, - session_id, - start_frame_index=prompt_frame_index, - max_frame_num_to_track=tracking_frames, - confidence=confidence, - ) + prompt_masks = _masks_from_outputs(prompt_response.get('outputs', {})) + if prompt_masks: + mask_dict[prompt_frame_index] = prompt_masks + + if tracking_frames > 1: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + propagated = collect_video_propagation( + predictor, + session_id, + start_frame_index=prompt_frame_index, + max_frame_num_to_track=tracking_frames - 1, + confidence=confidence, + propagation_chunk_size=mem['propagation_chunk_size'], + ) + mask_dict.update(propagated) finally: if session_id is not None: try: @@ -303,6 +642,10 @@ def run_sam3_video_text_inference( ) except Exception: pass + del predictor + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() return mask_dict @@ -323,50 +666,6 @@ def map_sam3_obj_ids_to_track_ids( return {obj_id: track_id for obj_id, track_id in zip(sam3_obj_ids, track_ids)} -def build_single_frame_annotation_output( - state: dict, - track_ids: List[str], - frame_idx: int, - track_type: str, - working_directory: Path, -) -> Path: - """Convert SAM3 image processor state into DIVE mask/track artifacts.""" - masks = state.get('masks') - if masks is None or len(masks) == 0: - raise RuntimeError( - 'No instances matched text prompt with confidence >= threshold. ' - f'Found {0 if masks is None else len(masks)} mask(s).' - ) - - if len(track_ids) < len(masks): - raise RuntimeError('Not enough track ids allocated for detected instances.') - - output_dir = working_directory / 'output/masks' - output_dir.mkdir(parents=True, exist_ok=True) - rle_masks: dict = {} - track_data = {'tracks': {}, 'groups': {}, 'version': 2} - - for index, track_id in enumerate(track_ids[: len(masks)]): - mask_tensor = masks[index] - mask_bin = mask_tensor.detach().cpu().numpy().astype(np.uint8) - save_and_record_mask( - mask_bin, - output_dir, - track_id, - frame_idx, - track_type, - rle_masks, - track_data, - ) - - with open(output_dir / 'RLE_MASKS.json', 'w') as f: - json.dump(rle_masks, f, indent=2) - with open(output_dir / 'TrackJSON.json', 'w') as f: - json.dump(track_data, f, indent=2) - - return output_dir - - def build_video_annotation_output( mask_dict: Dict[int, Dict[int, np.ndarray]], obj_id_to_track_id: Dict[int, str], @@ -407,7 +706,7 @@ def resolve_checkpoint_path() -> Optional[Path]: env_path = os.environ.get('SAM3_CHECKPOINT') if env_path and Path(env_path).exists(): return Path(env_path) - default = Path('/opt/SAM3/models/sam3.pt') + default = Path('/opt/SAM3/models/sam3.1_multiplex.pt') if default.exists(): return default return None @@ -465,65 +764,49 @@ def process_input_args(args, gc: girder_client.GirderClient) -> None: existing_tracks = gc.get('dive_annotation/track', {'folderId': dataset_id}) checkpoint_path = resolve_checkpoint_path() - if tracking_frames <= 1: - frame_path = _extract_single_frame( - args.DIVEVideo, start_frame, working_dir_path - ) - state = run_sam3_text_inference( - frame_path, - text_prompt, - confidence, - device, - checkpoint_path, + if process_masks_folder is None or extract_frames is None: + raise RuntimeError( + 'SAM2 helper functions are unavailable. This should only happen outside the slicer container environment.' ) - num_instances = len(state.get('masks', [])) - track_ids = allocate_track_ids(existing_tracks, num_instances) - print(f'Detected {num_instances} instance(s); assigning track ids: {track_ids}') - output_dir = build_single_frame_annotation_output( - state, - track_ids, - start_frame, - track_type, - working_dir_path, - ) - else: - if device != 'cuda': - raise RuntimeError( - f'Video propagation over {tracking_frames} frames requires CUDA. ' - 'Set TrackingFrames to 1 for CPU inference, or run on a GPU host.' - ) - frame_dir = extract_frames( - args.DIVEVideo, - start_frame, - tracking_frames, - working_dir_path, - ) - print( - f'Propagating text prompt from frame {start_frame} ' - f'for {tracking_frames} frame(s)...' - ) - mask_dict = run_sam3_video_text_inference( - frame_dir, - prompt_frame_index=0, - text_prompt=text_prompt, - tracking_frames=tracking_frames, - confidence=confidence, - checkpoint_path=checkpoint_path, - ) - obj_id_to_track_id = map_sam3_obj_ids_to_track_ids(mask_dict, existing_tracks) - print( - f'Propagated {len(mask_dict)} frame(s); ' - f'track mapping: {obj_id_to_track_id}' - ) - output_dir = build_video_annotation_output( - mask_dict, - obj_id_to_track_id, - start_frame, - track_type, - working_dir_path, + if device != 'cuda': + raise RuntimeError( + 'SAM 3.1 requires CUDA. Run on a GPU host with a compatible NVIDIA driver.' ) + frame_dir = extract_frames( + args.DIVEVideo, + start_frame, + tracking_frames, + working_dir_path, + ) + action = ( + f'Propagating text prompt from frame {start_frame} ' + f'for {tracking_frames} frame(s)...' + if tracking_frames > 1 + else f'Running text segmentation on frame {start_frame}...' + ) + print(action) + mask_dict = run_sam3_text_inference( + frame_dir, + prompt_frame_index=0, + text_prompt=text_prompt, + tracking_frames=tracking_frames, + confidence=confidence, + checkpoint_path=checkpoint_path, + ) + obj_id_to_track_id = map_sam3_obj_ids_to_track_ids(mask_dict, existing_tracks) + print( + f'Processed {len(mask_dict)} frame(s); track mapping: {obj_id_to_track_id}' + ) + output_dir = build_video_annotation_output( + mask_dict, + obj_id_to_track_id, + start_frame, + track_type, + working_dir_path, + ) + process_masks_folder( gc, dataset_id, @@ -533,33 +816,6 @@ def process_input_args(args, gc: girder_client.GirderClient) -> None: ) -def _extract_single_frame( - video_path: Union[str, Path], - frame_number: int, - working_directory: Path, -) -> Path: - """Extract one frame from a video as a JPEG.""" - frame_dir = working_directory / 'frames' - frame_dir.mkdir(parents=True, exist_ok=True) - output_path = frame_dir / '00000.jpg' - select_filter = f"select=eq(n\\,{frame_number})" - subprocess.run( - [ - 'ffmpeg', - '-i', str(video_path), - '-vf', select_filter, - '-vsync', '0', - '-frames:v', '1', - '-q:v', '2', - str(output_path), - ], - check=True, - ) - if not output_path.exists(): - raise RuntimeError(f'Failed to extract frame {frame_number} from {video_path}') - return output_path - - def main(args) -> None: gc = girder_client.GirderClient(apiUrl=args.girderApiUrl) gc.setToken(args.girderToken) @@ -569,4 +825,8 @@ def main(args) -> None: if __name__ == '__main__': + if CLIArgumentParser is None: + raise RuntimeError( + 'CLIArgumentParser is unavailable. This script is intended to run in the slicer container environment.' + ) main(CLIArgumentParser().parse_args()) diff --git a/dive-dsa-slicer/example-docker-containers/SAM3Demo/SAM3Demo.xml b/dive-dsa-slicer/example-docker-containers/SAM3Demo/SAM3Demo.xml index fe616fc8..0da68b62 100644 --- a/dive-dsa-slicer/example-docker-containers/SAM3Demo/SAM3Demo.xml +++ b/dive-dsa-slicer/example-docker-containers/SAM3Demo/SAM3Demo.xml @@ -1,8 +1,8 @@ DIVE - SAM3 Text Segmentation - Segment all instances matching a text prompt on a DIVE video frame using SAM3, optionally propagate masks forward through subsequent frames, and ingest annotations + SAM 3.1 Text Segmentation + Segment all instances matching a text prompt on a DIVE video frame using SAM 3.1, optionally propagate masks forward through subsequent frames, and ingest annotations 0.1.0 Apache 2.0 diff --git a/dive-dsa-slicer/example-docker-containers/SAM3Demo/test_sam31_loader.py b/dive-dsa-slicer/example-docker-containers/SAM3Demo/test_sam31_loader.py new file mode 100644 index 00000000..8c1602ad --- /dev/null +++ b/dive-dsa-slicer/example-docker-containers/SAM3Demo/test_sam31_loader.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +"""Local sanity checks for SAM 3.1 multiplex checkpoint loading patches.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +# Allow importing SAM3Demo helpers when run from repo root or SAM3Demo/ +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from SAM3Demo import ( # noqa: E402 + _patch_sam3_multiplex_checkpoint_preload, + _patch_sam3_offload_state_to_cpu_kwarg, + _prepare_sam31_multiplex_checkpoint, +) + + +def test_patches_apply() -> None: + _patch_sam3_offload_state_to_cpu_kwarg() + _patch_sam3_multiplex_checkpoint_preload() + import sam3.model_builder as model_builder + + assert getattr(model_builder, '_sam3demo_skip_tracker_preload', False) + print('OK: multiplex checkpoint preload patch applied') + + +def test_checkpoint_key_layout(checkpoint_path: Path) -> None: + ckpt = _prepare_sam31_multiplex_checkpoint(checkpoint_path) + keys = list(ckpt.keys()) + has_detector = any(k.startswith('detector.') for k in keys) + has_tracker = any(k.startswith('tracker.model.') for k in keys) + print(f'Checkpoint keys: {len(keys)} total') + print(f' detector.* present: {has_detector}') + print(f' tracker.model.* present: {has_tracker}') + if not (has_detector and has_tracker): + raise SystemExit( + 'Checkpoint does not look like SAM 3.1 multiplex (expected detector.* and tracker.model.*)' + ) + print('OK: checkpoint key layout') + + +def test_build_predictor(checkpoint_path: Path | None) -> None: + import torch + + if not torch.cuda.is_available(): + print('SKIP: CUDA not available for predictor build') + return + + _patch_sam3_offload_state_to_cpu_kwarg() + _patch_sam3_multiplex_checkpoint_preload() + from sam3.model_builder import build_sam3_multiplex_video_predictor + + ckpt = str(checkpoint_path) if checkpoint_path else None + print(f'Building predictor (checkpoint={ckpt or "HF Hub"})...') + predictor = build_sam3_multiplex_video_predictor( + checkpoint_path=ckpt, + use_fa3=False, + use_rope_real=False, + async_loading_frames=False, + ) + print(f'OK: predictor built: {type(predictor).__name__}') + + +def main() -> None: + test_patches_apply() + ckpt_env = os.environ.get('SAM3_CHECKPOINT', '/opt/SAM3/models/sam3.1_multiplex.pt') + ckpt_path = Path(ckpt_env) + if ckpt_path.is_file(): + test_checkpoint_key_layout(ckpt_path) + if '--build' in sys.argv: + test_build_predictor(ckpt_path) + else: + print(f'SKIP: checkpoint not found at {ckpt_path} (set SAM3_CHECKPOINT to test layout/build)') + if '--build' in sys.argv: + test_build_predictor(None) + + +if __name__ == '__main__': + main() diff --git a/dive-dsa-slicer/example-docker-containers/download_sam3_ckpts.sh b/dive-dsa-slicer/example-docker-containers/download_sam3_ckpts.sh index 1327beed..40a951fc 100755 --- a/dive-dsa-slicer/example-docker-containers/download_sam3_ckpts.sh +++ b/dive-dsa-slicer/example-docker-containers/download_sam3_ckpts.sh @@ -1,6 +1,6 @@ #!/bin/bash -# Download SAM3 checkpoints from Hugging Face when HF_TOKEN is available. -# Request access at https://huggingface.co/facebook/sam3 before building or running. +# Download SAM 3.1 checkpoints from Hugging Face when HF_TOKEN is available. +# Request access at https://huggingface.co/facebook/sam3.1 before building or running. set -euo pipefail @@ -18,17 +18,17 @@ if [ ! -f "${ASSETS_DIR}/${BPE_NAME}" ]; then fi if [ -z "${HF_TOKEN:-}" ]; then - echo "HF_TOKEN not set; skipping SAM3 checkpoint download." - echo "SAM3 will download checkpoints at runtime if Hugging Face credentials are configured." + echo "HF_TOKEN not set; skipping SAM 3.1 checkpoint download." + echo "SAM 3.1 will download checkpoints at runtime if Hugging Face credentials are configured." exit 0 fi if ! command -v hf &> /dev/null; then - echo "hf CLI not found; install huggingface_hub to prefetch SAM3 weights." + echo "hf CLI not found; install huggingface_hub to prefetch SAM 3.1 weights." exit 0 fi -hf download facebook/sam3 sam3.pt config.json \ +hf download facebook/sam3.1 sam3.1_multiplex.pt config.json \ --local-dir "$MODEL_DIR" \ --token "$HF_TOKEN" -echo "SAM3 checkpoints downloaded to $MODEL_DIR" +echo "SAM 3.1 checkpoints downloaded to $MODEL_DIR"