Upload export_sam3.py

export_sam3.py

@@ -0,0 +1,675 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#   "torch>=2.7",
+#   "torchvision",
+#   "onnx",
+#   "onnxruntime",
+#   "onnxsim",
+#   "timm",
+#   "ftfy==6.1.1",
+#   "regex",
+#   "iopath",
+#   "huggingface_hub",
+#   "einops",
+#   "pycocotools",
+#   "psutil",
+# ]
+# ///
+"""Export SAM3 tracker ONNX files and package a darktable .dtmodel.
+
+This script exports the SAM3 interactive tracker path behind a SAM2-shaped ONNX
+interface for darktable compatibility. It expects a user-downloaded sam3.pt in
+the repo root, clones the SAM3 source repo if needed, pins it to the tested
+commit, writes encoder.onnx and decoder.onnx into mask-object-sam3/, then
+packages mask-object-sam3.dtmodel.
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import os
+import shutil
+import subprocess
+import sys
+import warnings
+import zipfile
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+SAM3_REPO_URL = "https://github.com/facebookresearch/sam3.git"
+SAM3_COMMIT = "8e451d5eb43c817b64ae7577fb7b9ae223db88a9"
+MODEL_ID = "mask-object-sam3"
+
+
+try:
+    import onnxruntime  # type: ignore
+
+    ONNXRUNTIME_EXISTS = True
+except ImportError:
+    ONNXRUNTIME_EXISTS = False
+
+
+try:
+    import onnxsim  # type: ignore
+
+    ONNXSIM_EXISTS = True
+except ImportError:
+    ONNXSIM_EXISTS = False
+
+
+class SAM3TrackerEncoderOnnxModel(nn.Module):
+    """Wrap SAM3 tracker image encoder with a SAM2-shaped ONNX interface."""
+
+    SAM3_IMAGE_SIZE = 1008
+
+    def __init__(self, image_model: nn.Module, tracker_model: nn.Module):
+        super().__init__()
+        self.image_model = image_model
+        self.tracker_model = tracker_model
+        self.feat_sizes = [(288, 288), (144, 144), (72, 72)]
+
+    @torch.no_grad()
+    def forward(self, image: torch.Tensor):
+        image = F.interpolate(
+            image,
+            size=(self.SAM3_IMAGE_SIZE, self.SAM3_IMAGE_SIZE),
+            mode="bilinear",
+            align_corners=False,
+        )
+        backbone_out = self.image_model.backbone.forward_image(image)[
+            "sam2_backbone_out"
+        ]
+        backbone_out["backbone_fpn"][0] = self.tracker_model.sam_mask_decoder.conv_s0(
+            backbone_out["backbone_fpn"][0]
+        )
+        backbone_out["backbone_fpn"][1] = self.tracker_model.sam_mask_decoder.conv_s1(
+            backbone_out["backbone_fpn"][1]
+        )
+        _, vision_feats, _, _ = self.tracker_model._prepare_backbone_features(
+            backbone_out
+        )
+
+        # Match SAM3InteractiveImagePredictor image embedding path.
+        vision_feats[-1] = vision_feats[-1] + self.tracker_model.no_mem_embed
+
+        feats = [
+            feat.permute(1, 2, 0).view(1, -1, *feat_size)
+            for feat, feat_size in zip(vision_feats[::-1], self.feat_sizes[::-1])
+        ][::-1]
+
+        high_res_feats_0 = F.interpolate(
+            feats[0], size=(256, 256), mode="bilinear", align_corners=False
+        )
+        high_res_feats_1 = F.interpolate(
+            feats[1], size=(128, 128), mode="bilinear", align_corners=False
+        )
+        image_embed = F.interpolate(
+            feats[2], size=(64, 64), mode="bilinear", align_corners=False
+        )
+
+        return high_res_feats_0, high_res_feats_1, image_embed
+
+
+class SAM3TrackerDecoderOnnxModel(nn.Module):
+    """Wrap SAM3 tracker decoder with a SAM2-shaped ONNX interface."""
+
+    SAM2_IMAGE_SIZE = 1024.0
+    SAM3_IMAGE_SIZE = 1008.0
+
+    def __init__(self, tracker_model: nn.Module, multimask_output: bool):
+        super().__init__()
+        self.prompt_encoder = tracker_model.sam_prompt_encoder
+        self.mask_decoder = tracker_model.sam_mask_decoder
+        self.image_size = tracker_model.image_size
+        self.multimask_output = multimask_output
+
+    @torch.no_grad()
+    def forward(
+        self,
+        image_embed: torch.Tensor,
+        high_res_feats_0: torch.Tensor,
+        high_res_feats_1: torch.Tensor,
+        point_coords: torch.Tensor,
+        point_labels: torch.Tensor,
+        mask_input: torch.Tensor,
+        has_mask_input: torch.Tensor,
+    ):
+        image_embed = F.interpolate(
+            image_embed, size=(72, 72), mode="bilinear", align_corners=False
+        )
+        high_res_feats_0 = F.interpolate(
+            high_res_feats_0, size=(288, 288), mode="bilinear", align_corners=False
+        )
+        high_res_feats_1 = F.interpolate(
+            high_res_feats_1, size=(144, 144), mode="bilinear", align_corners=False
+        )
+        mask_input = F.interpolate(
+            mask_input, size=(288, 288), mode="bilinear", align_corners=False
+        )
+        point_coords = point_coords * (self.SAM3_IMAGE_SIZE / self.SAM2_IMAGE_SIZE)
+
+        sparse_embedding = self._embed_points(point_coords, point_labels)
+        dense_embedding = self._embed_masks(mask_input, has_mask_input)
+
+        masks, iou_predictions, _, _ = self.mask_decoder(
+            image_embeddings=image_embed,
+            image_pe=self.prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_embedding,
+            dense_prompt_embeddings=dense_embedding,
+            multimask_output=self.multimask_output,
+            repeat_image=False,
+            high_res_features=[high_res_feats_0, high_res_feats_1],
+        )
+
+        masks = torch.clamp(masks, -32.0, 32.0)
+        masks = F.interpolate(
+            masks, size=(256, 256), mode="bilinear", align_corners=False
+        )
+        return masks, iou_predictions
+
+    def _embed_points(
+        self, point_coords: torch.Tensor, point_labels: torch.Tensor
+    ) -> torch.Tensor:
+        point_coords = point_coords + 0.5
+
+        padding_point = torch.zeros(
+            (point_coords.shape[0], 1, 2), device=point_coords.device
+        )
+        padding_label = -torch.ones(
+            (point_labels.shape[0], 1),
+            dtype=point_labels.dtype,
+            device=point_labels.device,
+        )
+        point_coords = torch.cat([point_coords, padding_point], dim=1)
+        point_labels = torch.cat([point_labels, padding_label], dim=1)
+
+        point_coords[:, :, 0] = point_coords[:, :, 0] / self.image_size
+        point_coords[:, :, 1] = point_coords[:, :, 1] / self.image_size
+
+        point_embedding = self.prompt_encoder.pe_layer._pe_encoding(point_coords)
+        point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding)
+
+        point_embedding = point_embedding * (point_labels != -1)
+        point_embedding = (
+            point_embedding
+            + self.prompt_encoder.not_a_point_embed.weight
+            * (point_labels == -1)
+        )
+
+        for i in range(self.prompt_encoder.num_point_embeddings):
+            point_embedding = (
+                point_embedding
+                + self.prompt_encoder.point_embeddings[i].weight
+                * (point_labels == i)
+            )
+
+        return point_embedding
+
+    def _embed_masks(
+        self, mask_input: torch.Tensor, has_mask_input: torch.Tensor
+    ) -> torch.Tensor:
+        target_size = self.prompt_encoder.mask_input_size
+        if mask_input.shape[-2:] != target_size:
+            mask_input = F.interpolate(
+                mask_input.float(),
+                size=target_size,
+                mode="bilinear",
+                align_corners=False,
+                antialias=True,
+            )
+
+        mask_embedding = self.prompt_encoder.mask_downscaling(mask_input)
+        no_mask_embedding = self.prompt_encoder.no_mask_embed.weight.reshape(
+            1, -1, 1, 1
+        )
+
+        return has_mask_input * mask_embedding + (1 - has_mask_input) * no_mask_embedding
+
+
+def to_numpy(tensor: torch.Tensor):
+    return tensor.detach().cpu().numpy()
+
+
+def run(command: list[str], cwd: Path | None = None) -> None:
+    print(" ".join(command))
+    subprocess.run(command, cwd=cwd, check=True)
+
+
+def ensure_sam3_repo(sam3_root: Path) -> None:
+    if shutil.which("git") is None:
+        raise RuntimeError("git was not found. Install git and run this script again.")
+
+    if not sam3_root.exists():
+        print(f"Cloning SAM3 into {sam3_root}...")
+        run(["git", "clone", SAM3_REPO_URL, str(sam3_root)])
+    elif not (sam3_root / ".git").is_dir():
+        raise RuntimeError(f"{sam3_root} exists but is not a git checkout")
+
+    print(f"Pinning SAM3 to {SAM3_COMMIT}...")
+    try:
+        run(["git", "checkout", SAM3_COMMIT], cwd=sam3_root)
+    except subprocess.CalledProcessError:
+        run(["git", "fetch", "origin", SAM3_COMMIT], cwd=sam3_root)
+        run(["git", "checkout", SAM3_COMMIT], cwd=sam3_root)
+
+    if not (sam3_root / "sam3").is_dir():
+        raise RuntimeError(f"SAM3 package folder not found under {sam3_root}")
+
+
+def add_sam3_to_path(sam3_root: Path) -> None:
+    sam3_root = sam3_root.resolve()
+    sam3_root_text = str(sam3_root)
+    if sam3_root_text not in sys.path:
+        sys.path.insert(0, sam3_root_text)
+
+
+@contextlib.contextmanager
+def redirect_cuda_allocations_to_cpu(enabled: bool):
+    """Redirect SAM3 construction-time CUDA cache tensors to CPU."""
+
+    if not enabled:
+        yield
+        return
+
+    original_zeros = torch.zeros
+    original_arange = torch.arange
+
+    def rewrite_cuda_device(kwargs: dict) -> dict:
+        if kwargs.get("device") == "cuda":
+            kwargs = dict(kwargs)
+            kwargs["device"] = "cpu"
+        return kwargs
+
+    def zeros_wrapper(*args, **kwargs):
+        return original_zeros(*args, **rewrite_cuda_device(kwargs))
+
+    def arange_wrapper(*args, **kwargs):
+        return original_arange(*args, **rewrite_cuda_device(kwargs))
+
+    torch.zeros = zeros_wrapper
+    torch.arange = arange_wrapper
+    try:
+        yield
+    finally:
+        torch.zeros = original_zeros
+        torch.arange = original_arange
+
+
+def build_models(sam3_root: Path, checkpoint: Path, device: str):
+    add_sam3_to_path(sam3_root)
+
+    import sam3.model_builder as model_builder
+
+    force_real_rope(model_builder)
+    disable_fused_addmm_activation()
+
+    print("Loading SAM3 image model with interactive tracker...")
+    with redirect_cuda_allocations_to_cpu(enabled=device == "cpu"):
+        image_model = model_builder.build_sam3_image_model(
+            checkpoint_path=str(checkpoint),
+            load_from_HF=False,
+            enable_inst_interactivity=True,
+            device=device,
+            eval_mode=True,
+        )
+
+    if image_model.inst_interactive_predictor is None:
+        raise RuntimeError("SAM3 image model did not create inst_interactive_predictor")
+
+    tracker_model = image_model.inst_interactive_predictor.model
+    image_model.eval()
+    tracker_model.eval()
+    image_model.to(device)
+    tracker_model.to(device)
+    image_model.float()
+    tracker_model.float()
+    if device == "cpu":
+        add_parameter_dtype_input_hooks(image_model)
+        add_parameter_dtype_input_hooks(tracker_model)
+    return image_model, tracker_model
+
+
+def force_real_rope(model_builder_module) -> None:
+    """Force SAM3 ViT construction to avoid complex-valued RoPE ops."""
+
+    original_create_vit_backbone = model_builder_module._create_vit_backbone
+
+    def create_vit_backbone_real_rope(
+        compile_mode=None,
+        use_fa3=False,
+        use_rope_real=False,
+    ):
+        return original_create_vit_backbone(
+            compile_mode=compile_mode,
+            use_fa3=use_fa3,
+            use_rope_real=True,
+        )
+
+    model_builder_module._create_vit_backbone = create_vit_backbone_real_rope
+
+
+def disable_fused_addmm_activation() -> None:
+    """Replace SAM3's fused addmm+activation with ONNX-friendly ops."""
+
+    import torch.nn.functional as functional
+    import sam3.model.vitdet as vitdet
+    import sam3.perflib.fused as fused
+
+    def export_friendly_addmm_act(activation, linear, mat1):
+        output = functional.linear(mat1.float(), linear.weight.float(), linear.bias.float())
+        if activation in [functional.relu, torch.nn.ReLU]:
+            return functional.relu(output)
+        if activation in [functional.gelu, torch.nn.GELU]:
+            return functional.gelu(output)
+        raise ValueError(f"Unexpected activation {activation}")
+
+    fused.addmm_act = export_friendly_addmm_act
+    vitdet.addmm_act = export_friendly_addmm_act
+
+
+def add_parameter_dtype_input_hooks(model: nn.Module) -> None:
+    """Cast floating inputs to module parameter dtype during CPU export."""
+
+    def first_parameter_dtype(module: nn.Module):
+        for parameter in module.parameters(recurse=False):
+            return parameter.dtype
+        return None
+
+    def cast_value(value, dtype):
+        if torch.is_tensor(value) and value.is_floating_point() and value.dtype != dtype:
+            return value.to(dtype=dtype)
+        if isinstance(value, tuple):
+            return tuple(cast_value(item, dtype) for item in value)
+        if isinstance(value, list):
+            return [cast_value(item, dtype) for item in value]
+        if isinstance(value, dict):
+            return {key: cast_value(item, dtype) for key, item in value.items()}
+        return value
+
+    def hook(module, args):
+        dtype = first_parameter_dtype(module)
+        if dtype is None:
+            return args
+        return tuple(cast_value(arg, dtype) for arg in args)
+
+    hook_types = (
+        nn.Conv2d,
+        nn.ConvTranspose2d,
+        nn.LayerNorm,
+        nn.Linear,
+        nn.MultiheadAttention,
+    )
+    for module in model.modules():
+        if isinstance(module, hook_types):
+            module.register_forward_pre_hook(hook)
+
+
+def export_encoder(
+    image_model: nn.Module, tracker_model: nn.Module, output: Path, opset: int
+):
+    encoder = SAM3TrackerEncoderOnnxModel(image_model, tracker_model)
+    encoder.eval()
+
+    device = tracker_model.device
+    dummy_image = torch.randn(1, 3, 1024, 1024, dtype=torch.float32, device=device)
+    high_res_feats_0, high_res_feats_1, image_embed = encoder(dummy_image)
+
+    print("Encoder output shapes:")
+    print(f"  high_res_feats_0: {tuple(high_res_feats_0.shape)}")
+    print(f"  high_res_feats_1: {tuple(high_res_feats_1.shape)}")
+    print(f"  image_embed:      {tuple(image_embed.shape)}")
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=torch.jit.TracerWarning)
+        warnings.filterwarnings("ignore", category=UserWarning)
+        print(f"Exporting encoder to {output}...")
+        torch.onnx.export(
+            encoder,
+            dummy_image,
+            str(output),
+            export_params=True,
+            verbose=False,
+            opset_version=opset,
+            do_constant_folding=True,
+            input_names=["image"],
+            output_names=["high_res_feats_0", "high_res_feats_1", "image_embed"],
+            dynamo=False,
+        )
+
+    simplify_model(output)
+    verify_model(output, {"image": to_numpy(dummy_image)}, "Encoder")
+    return high_res_feats_0, high_res_feats_1, image_embed
+
+
+def export_decoder(
+    tracker_model: nn.Module,
+    output: Path,
+    opset: int,
+    high_res_feats_0: torch.Tensor,
+    high_res_feats_1: torch.Tensor,
+    image_embed: torch.Tensor,
+    multimask_output: bool,
+):
+    decoder = SAM3TrackerDecoderOnnxModel(
+        tracker_model, multimask_output=multimask_output
+    )
+    decoder.eval()
+
+    device = tracker_model.device
+    dummy_inputs = {
+        "image_embed": image_embed,
+        "high_res_feats_0": high_res_feats_0,
+        "high_res_feats_1": high_res_feats_1,
+        "point_coords": torch.randint(
+            low=0, high=1024, size=(1, 5, 2), dtype=torch.float32, device=device
+        ),
+        "point_labels": torch.randint(
+            low=0, high=2, size=(1, 5), dtype=torch.float32, device=device
+        ),
+        "mask_input": torch.randn(1, 1, 256, 256, dtype=torch.float32, device=device),
+        "has_mask_input": torch.tensor([1], dtype=torch.float32, device=device),
+    }
+
+    masks, iou_predictions = decoder(**dummy_inputs)
+    print("Decoder output shapes:")
+    print(f"  masks:           {tuple(masks.shape)}")
+    print(f"  iou_predictions: {tuple(iou_predictions.shape)}")
+
+    dynamic_axes = {
+        "point_coords": {1: "num_points"},
+        "point_labels": {1: "num_points"},
+    }
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=torch.jit.TracerWarning)
+        warnings.filterwarnings("ignore", category=UserWarning)
+        print(f"Exporting decoder to {output}...")
+        torch.onnx.export(
+            decoder,
+            tuple(dummy_inputs.values()),
+            str(output),
+            export_params=True,
+            verbose=False,
+            opset_version=opset,
+            do_constant_folding=True,
+            input_names=list(dummy_inputs.keys()),
+            output_names=["masks", "iou_predictions"],
+            dynamic_axes=dynamic_axes,
+            dynamo=False,
+        )
+
+    simplify_model(output)
+    verify_model(output, {k: to_numpy(v) for k, v in dummy_inputs.items()}, "Decoder")
+
+
+def simplify_model(model_path: Path) -> None:
+    if not ONNXSIM_EXISTS:
+        print("Warning: onnx-simplifier not installed, skipping simplification.")
+        return
+
+    import onnx
+
+    print(f"Simplifying {model_path.name}...")
+    onnx_model = onnx.load(str(model_path))
+    model_simp, check = onnxsim.simplify(onnx_model)
+    if not check:
+        print("Warning: simplified model failed validation, keeping original.")
+        return
+
+    to_remove = [
+        vi
+        for vi in model_simp.graph.value_info
+        if len(vi.type.tensor_type.shape.dim) == 0
+    ]
+    for vi in to_remove:
+        model_simp.graph.value_info.remove(vi)
+    if to_remove:
+        print(f"Removed {len(to_remove)} empty shape annotations.")
+
+    onnx.save(model_simp, str(model_path))
+    print("Graph simplification passed.")
+
+
+def verify_model(model_path: Path, dummy_inputs: dict, name: str) -> None:
+    if not ONNXRUNTIME_EXISTS:
+        print(f"Warning: onnxruntime not installed, skipping {name} validation.")
+        return
+
+    providers = ["CPUExecutionProvider"]
+    session = onnxruntime.InferenceSession(str(model_path), providers=providers)
+    outputs = session.run(None, dummy_inputs)
+    print(f"{name} ran successfully with ONNXRuntime.")
+    for output_meta, output_value in zip(session.get_outputs(), outputs):
+        print(f"  {output_meta.name}: {tuple(output_value.shape)}")
+
+
+def package_dtmodel(output_dir: Path, package_path: Path) -> None:
+    required_files = ["config.json", "encoder.onnx", "decoder.onnx"]
+    missing = [name for name in required_files if not (output_dir / name).is_file()]
+    if missing:
+        raise FileNotFoundError(
+            f"Missing required file(s) in {output_dir}: {', '.join(missing)}"
+        )
+
+    if package_path.exists():
+        package_path.unlink()
+
+    print(f"Packaging {package_path}...")
+    with zipfile.ZipFile(package_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for name in required_files:
+            file_path = output_dir / name
+            zf.write(file_path, f"{MODEL_ID}/{name}")
+    print(f"Created: {package_path}")
+
+
+def export(
+    sam3_root: Path,
+    checkpoint: Path,
+    output_dir: Path,
+    package_path: Path,
+    opset: int,
+    device: str,
+    multimask_output: bool,
+) -> None:
+    if device == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA requested but torch.cuda.is_available() is false")
+
+    ensure_sam3_repo(sam3_root)
+
+    if not checkpoint.is_file():
+        raise FileNotFoundError(
+            f"sam3.pt not found: {checkpoint}\n"
+            "Download it from https://huggingface.co/facebook/sam3 after "
+            "requesting access, then place it next to export_sam3.py."
+        )
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    image_model, tracker_model = build_models(sam3_root, checkpoint, device)
+
+    encoder_path = output_dir / "encoder.onnx"
+    decoder_path = output_dir / "decoder.onnx"
+
+    high_res_feats_0, high_res_feats_1, image_embed = export_encoder(
+        image_model,
+        tracker_model,
+        encoder_path,
+        opset,
+    )
+    export_decoder(
+        tracker_model,
+        decoder_path,
+        opset,
+        high_res_feats_0,
+        high_res_feats_1,
+        image_embed,
+        multimask_output,
+    )
+    package_dtmodel(output_dir, package_path)
+
+    print("Done.")
+    print(f"Encoder: {encoder_path}")
+    print(f"Decoder: {decoder_path}")
+    print(f"Package: {package_path}")
+
+
+def main() -> None:
+    repo_root = Path(__file__).resolve().parent
+
+    parser = argparse.ArgumentParser(
+        description="Export SAM3 tracker encoder.onnx/decoder.onnx and package a darktable .dtmodel."
+    )
+    parser.add_argument(
+        "--sam3-root",
+        type=Path,
+        default=repo_root / "sam3",
+        help="Path where the facebookresearch/sam3 repo will be cloned or reused.",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=Path,
+        default=repo_root / "sam3.pt",
+        help="Path to the user-downloaded sam3.pt checkpoint.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=repo_root / MODEL_ID,
+        help="Output directory containing config.json and generated ONNX files.",
+    )
+    parser.add_argument(
+        "--package",
+        type=Path,
+        default=repo_root / f"{MODEL_ID}.dtmodel",
+        help="Output .dtmodel package path.",
+    )
+    parser.add_argument("--opset", type=int, default=20)
+    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
+    parser.add_argument(
+        "--single-mask",
+        action="store_true",
+        help="Export decoder with a single mask output instead of three masks.",
+    )
+    args = parser.parse_args()
+
+    os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+
+    export(
+        sam3_root=args.sam3_root,
+        checkpoint=args.checkpoint,
+        output_dir=args.output_dir,
+        package_path=args.package,
+        opset=args.opset,
+        device=args.device,
+        multimask_output=not args.single_mask,
+    )
+
+
+if __name__ == "__main__":
+    main()