refiners: bugfix local corr

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "romav2"
-version = "2.0.0"
+version = "2.0.1"
 description = "RoMa v2: Harder Better Faster Denser Feature Matching"
 readme = "README.md"
 authors = [
@@ -15,16 +15,14 @@ dependencies = [
     "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",

src/romav2/refiner.py

@@ -169,18 +169,17 @@ 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_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,
-                f_B_bdhw,
+                f_B.permute(0, 3, 1, 2),
                 local_radius=self.cfg.local_corr_radius,
                 warp=prev_warp,
                 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

@@ -77,7 +77,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
@@ -96,7 +96,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)