From 48997fcd3d42a133198340b62bf7817003bcbc37 Mon Sep 17 00:00:00 2001
From: Kristian Myram <kristian.myram@gmail.com>
Date: Thu, 15 Oct 2026 01:10:28 +0200
Subject: [PATCH] S2.2: make the reference script run without torchvision

Preprocess by hand, tap the neck's ModelOutput, report forward time.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 tools/sam2_reference.py | 36 ++++++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/tools/sam2_reference.py b/tools/sam2_reference.py
index e40019134..2cfd03fde 100644
--- a/tools/sam2_reference.py
+++ b/tools/sam2_reference.py
@@ -11,7 +11,8 @@ import sys
 import numpy as np
 import torch
 from safetensors.torch import save_file
-from transformers import Sam2Model, Sam2Processor
+import torch.nn.functional as F
+from transformers import Sam2Model
 
 
 def synthetic_image(size=256):
@@ -30,17 +31,33 @@ def main(model_dir, out):
     torch.manual_seed(0)
     torch.set_num_threads(1)
     model = Sam2Model.from_pretrained(model_dir).eval()
-    proc = Sam2Processor.from_pretrained(model_dir)
     image = synthetic_image()
-    inputs = proc(images=image, input_points=[[[[115.0, 128.0]]]], input_labels=[[[1]]], return_tensors="pt")
+    # Preprocess by hand (no torchvision needed): resize to 1024x1024, ImageNet mean/std.
+    x = torch.from_numpy(image).permute(2, 0, 1)[None].float() / 255.0
+    x = F.interpolate(x, size=(1024, 1024), mode="bilinear", align_corners=False)
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+    scale = 1024.0 / image.shape[0]
+    inputs = {
+        "pixel_values": (x - mean) / std,
+        "input_points": torch.tensor([[[[115.0 * scale, 128.0 * scale]]]]),
+        "input_labels": torch.tensor([[[1]]], dtype=torch.int32),
+    }
     taps = {}
 
     def tap(name):
         def hook(_m, _i, o):
-            outs = o if isinstance(o, (tuple, list)) else [o]
-            for k, t in enumerate(outs):
-                if isinstance(t, torch.Tensor):
-                    taps[f"{name}.{k}"] = t.detach().float().contiguous()
+            def flat(x):
+                if isinstance(x, torch.Tensor):
+                    yield x
+                elif isinstance(x, (tuple, list)):
+                    for y in x:
+                        yield from flat(y)
+                elif hasattr(x, "to_tuple"):  # a ModelOutput (the neck)
+                    yield from flat(x.to_tuple())
+
+            for k, t in enumerate(flat(o)):
+                taps[f"{name}.{k}"] = t.detach().float().clone().contiguous()
         return hook
 
     # Stage taps by module path: patch embed, each Hiera block, the neck, prompt encoder, decoder.
@@ -50,12 +67,15 @@ def main(model_dir, out):
             len(parts) == 4 and parts[:3] == ["vision_encoder", "backbone", "blocks"]
         ):
             mod.register_forward_hook(tap(name))
+    import time
     with torch.no_grad():
+        t0 = time.perf_counter()
         res = model(**inputs, multimask_output=True)
+        print(f"forward (1 thread, first run): {time.perf_counter() - t0:.2f} s")
     taps["input.pixel_values"] = inputs["pixel_values"].float()
     taps["out.iou_scores"] = res.iou_scores.float()
     taps["out.pred_masks"] = res.pred_masks.float()
-    save_file({k: v for k, v in taps.items()}, out)
+    save_file({k: v.clone().contiguous() for k, v in taps.items()}, out)
     print(f"wrote {len(taps)} tensors to {out}")