Merge upstream RoMaV2 v2.0.1 into main (+ fork adaptations)

README.md

@@ -108,18 +108,10 @@ Running these gave me `ScanNet-1500: [34.0, 56.5, 73.9]`, and `Mega-1500: [62.8,
 
 
 ## Fused local correlation kernel
-Include the `--extra fused-local-corr` flag as:
-```bash
-uv sync --extra fused-local-corr
-```
-or 
-```bash
-uv pip install romav2[fused-local-corr]
-```
-or
-```bash
-uv add romav2[fused-local-corr]
-```
+On Linux the `fused-local-corr` kernel is a core dependency and is installed
+automatically by the Setup/Install step above (`uv sync` / `uv pip install -e .`)
+— no extra flag is required. On non-Linux platforms it is skipped and RoMa falls
+back to the pure-PyTorch local correlation.
 
 ## Settings
 By twiddling with some different settings you may reach better results on your task of interest.

pyproject.toml

@@ -1,30 +1,27 @@
 [project]
 name = "romav2"
-version = "2.0.0"
+version = "2.0.1"
 description = "RoMa v2: Harder Better Faster Denser Feature Matching"
 readme = "README.md"
 authors = [
     { name = "Johan Edstedt", email = "johan.edstedt@liu.se" }
 ]
 requires-python = ">=3.10"
 dependencies = [
-    "dataclasses>=0.8",
     "einops>=0.8.1",
     "pillow>=12.0.0",
     "rich>=14.2.0",
     "torch",
     "torchvision>=0.23.0",
     "tqdm>=4.67.1",
+    "fused-local-corr ; sys_platform == 'linux'",
 ]
 
 [build-system]
 requires = ["uv_build>=0.8.15,<0.9.0"]
 build-backend = "uv_build"
 
 [project.optional-dependencies]
-fused-local-corr = [
-    "fused-local-corr ; sys_platform == 'linux'",
-]
 eval = [
     "kornia>=0.8.2",
     "matplotlib>=3.10.7",

scripts/batch_match.py

@@ -744,6 +744,17 @@ def main():
         default=True,
         help='Disable bidirectional matching (only A->B; default is bidirectional)'
     )
+    parser.add_argument(
+        '--no-compile',
+        dest='compile',
+        action='store_false',
+        default=True,
+        help='Disable torch.compile. Compiling is enabled by default for '
+             'throughput (match() runs at fixed resolution, so the graph is '
+             'compiled once and reused across all pairs). Disable it when '
+             'profiling/tracing, since torch.compile graph capture hides the '
+             'torch.profiler.record_function annotations.'
+    )
     parser.add_argument(
         '--calibration-pairs',
         type=int,
@@ -781,10 +792,10 @@ def main():
         pairs = optimize_pair_order(pairs)
 
     # Initialize model
-    logger.info(f"Initializing RoMaV2 with setting: {args.setting}")
-    #Use compile=False for full tracing
-    #model = RoMaV2(RoMaV2.Cfg(compile=False))
-    model = RoMaV2()
+    logger.info(f"Initializing RoMaV2 with setting: {args.setting} (compile={args.compile})")
+    # compile=True is the throughput path; pass --no-compile for full tracing
+    # (torch.compile hides the profiler record_function annotations).
+    model = RoMaV2(RoMaV2.Cfg(compile=args.compile))
     model.apply_setting(args.setting)
     model.bidirectional = args.bidirectional
     model.eval()

src/romav2/refiner.py

@@ -169,8 +169,9 @@ def forward(
         )
         # Corr in other means take a kxk grid around the predicted coordinate in other image
         f_A_bdhw = f_A.permute(0, 3, 1, 2)
-        f_B_bdhw = f_BA.permute(0, 3, 1, 2)
-        d = torch.cat((f_A_bdhw, f_B_bdhw, emb_in_displacement), dim=1)
+        f_B_bdhw = f_B.permute(0, 3, 1, 2)
+        f_BA_bdhw = f_BA.permute(0, 3, 1, 2)
+        d = torch.cat((f_A_bdhw, f_BA_bdhw, emb_in_displacement), dim=1)
         if self.cfg.local_corr_radius is not None:
             local_corr = local_correlation(
                 f_A_bdhw,
@@ -180,7 +181,6 @@ def forward(
                 scale_factor=scale_factor,
             )
             d = torch.cat((d, local_corr), dim=1)
-        # d = torch.cat((f_A_bdhw, f_B_bdhw, emb_in_displacement, local_corr), dim=1)
         if self.cfg.channels_last:
             d = d.to(memory_format=torch.channels_last)
         z = self.block1(d)

src/romav2/romav2.py

@@ -78,7 +78,7 @@ class Cfg:
         anchor_width: int = 512
         anchor_height: int = 512
         setting: Setting = "precise"
-        compile: bool = True
+        compile: bool = False
         name: str = "RoMa v2"
 
     # settings
@@ -97,7 +97,8 @@ def __init__(self, cfg: Cfg | None = None):
             cfg = RoMaV2.Cfg()
 
         weights = torch.hub.load_state_dict_from_url(
-            "https://github.com/Parskatt/RoMaV2/releases/download/weights/romav2.pt"
+            "https://github.com/Parskatt/RoMaV2/releases/download/v2.0.1/romav2.0.1.pt",
+            map_location=device
         )
         self.f = Descriptor(cfg.descriptor)
         self.matcher = Matcher(cfg.matcher)