From 7f12442efa5aad2ba6aedaefcee87718c78e96bd Mon Sep 17 00:00:00 2001 From: Fabien Servant Date: Tue, 16 Jun 2026 14:52:27 +0200 Subject: [PATCH 1/2] Using local models on filesystem --- src/romav2/features.py | 8 ++++++-- src/romav2/romav2.py | 12 ++++++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/src/romav2/features.py b/src/romav2/features.py index 6836eaa..92107bf 100644 --- a/src/romav2/features.py +++ b/src/romav2/features.py @@ -88,6 +88,7 @@ class Cfg: default_factory=lambda: [11, 17] ) # [4, 11, 17, 23] for dinov3 style weights_path: str | None = None + module_path: str | None = None # Path to local directory containing module def __new__(cls, cfg: Cfg) -> nn.Module: partial_wrap = partial( @@ -101,9 +102,10 @@ def __new__(cls, cfg: Cfg) -> nn.Module: normalizer = imagenet # TODO: this will break in distributed if not available locally dinov3_vitl16: nn.Module = torch.hub.load( - repo_or_dir="facebookresearch/dinov3:adc254450203739c8149213a7a69d8d905b4fcfa", + repo_or_dir="facebookresearch/dinov3:adc254450203739c8149213a7a69d8d905b4fcfa" if cfg.module_path is None else cfg.module_path, model="dinov3_vitl16", pretrained=cfg.weights_path is not None, + source="github" if cfg.module_path is None else "local", weights=cfg.weights_path, skip_validation=True, ).to(device) @@ -118,7 +120,9 @@ def __new__(cls, cfg: Cfg) -> nn.Module: normalizer = imagenet dinov2_vit14: nn.Module = torch.hub.load( - "facebookresearch/dinov2", "dinov2_vitl14" + repo_or_dir="facebookresearch/dinov2" if cfg.module_path is None else cfg.module_path, + model="dinov2_vitl14", + source="github" if cfg.module_path is None else "local" ).to(device) dinov2_vit14.mask_token = None layers = _get_layers(cfg.layer_idx, dinov2_vit14) diff --git a/src/romav2/romav2.py b/src/romav2/romav2.py index 65a1a3c..cb60641 100644 --- a/src/romav2/romav2.py +++ b/src/romav2/romav2.py @@ -79,6 +79,7 @@ class Cfg: setting: Setting = "precise" compile: bool = False name: str = "RoMa v2" + weights: dict = None # settings H_lr: int @@ -94,11 +95,14 @@ def __init__(self, cfg: Cfg | None = None): if cfg is None: # default cfg = RoMaV2.Cfg() + + if cfg.weights is None: + weights = torch.hub.load_state_dict_from_url( + "https://github.com/Parskatt/RoMaV2/releases/download/weights/romav2.pt" + ) + else: + weights = cfg.weights - weights = torch.hub.load_state_dict_from_url( - "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) self.cfg = cfg From b780da6118de6f6f0a83a586e529102c212fefc9 Mon Sep 17 00:00:00 2001 From: Fabien Servant Date: Fri, 19 Jun 2026 16:40:18 +0200 Subject: [PATCH 2/2] memory optimization --- src/romav2/matcher.py | 12 ++++++++--- src/romav2/romav2.py | 49 +++++++++++++++++++++++++++++++------------ 2 files changed, 45 insertions(+), 16 deletions(-) diff --git a/src/romav2/matcher.py b/src/romav2/matcher.py index 1c1ef6d..a75d7da 100644 --- a/src/romav2/matcher.py +++ b/src/romav2/matcher.py @@ -32,7 +32,8 @@ def _compute_match_embeddings( W_A: int, H_B: int, W_B: int, -) -> torch.Tensor: + return_attention: bool, +) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]: attn_AB_logits = (1 / temp * cosine_similarity(f_A, f_B)).reshape( B, H_A * W_A, H_B * W_B ) @@ -42,8 +43,10 @@ def _compute_match_embeddings( match_emb = einsum( attn_AB, pos_emb_grid, "B H_A W_A H_B W_B, B H_B W_B D -> B H_A W_A D" ) - attn_AB_logits = attn_AB_logits.reshape(B, H_A, W_A, H_B, W_B) - return attn_AB_logits, attn_AB, match_emb + if return_attention: + attn_AB_logits = attn_AB_logits.reshape(B, H_A, W_A, H_B, W_B) + return attn_AB_logits, attn_AB, match_emb + return None, None, match_emb def _compute_head_preds( @@ -121,6 +124,7 @@ def forward( img_A: torch.Tensor, img_B: torch.Tensor, bidirectional: bool, + return_attention: bool = True, ): preds = {} f_A = torch.cat(f_list_A, dim=-1) @@ -158,6 +162,7 @@ def forward( W_A=W_A, H_B=H_B, W_B=W_B, + return_attention=return_attention, ) warp_AB, confidence_AB = _compute_head_preds( f_list_A=f_list_A, @@ -179,6 +184,7 @@ def forward( W_A=W_B, H_B=H_A, W_B=W_A, + return_attention=return_attention, ) warp_BA, confidence_BA = _compute_head_preds( f_list_A=f_list_B, diff --git a/src/romav2/romav2.py b/src/romav2/romav2.py index cb60641..e5906b7 100644 --- a/src/romav2/romav2.py +++ b/src/romav2/romav2.py @@ -67,6 +67,14 @@ def _map_confidence(*, confidence: torch.Tensor, threshold: float | None): return overlap, precision +def _move_optional_tensor( + tensor: torch.Tensor | None, output_device: torch.device | str | None +) -> torch.Tensor | None: + if tensor is None or output_device is None: + return tensor + return tensor.to(output_device) + + class RoMaV2(nn.Module): @dataclass(frozen=True) class Cfg: @@ -171,6 +179,7 @@ def forward( img_B_lr: torch.Tensor, img_A_hr: torch.Tensor | None = None, img_B_hr: torch.Tensor | None = None, + return_intermediates: bool = True, ) -> dict[str, tuple[torch.Tensor, torch.Tensor] | torch.Tensor]: if torch.get_float32_matmul_precision() != "highest": raise RuntimeError("Float32 matmul precision must be set to highest") @@ -183,10 +192,16 @@ def forward( f_B = self.f(img_B_lr) # match feats matcher_output = self.matcher( - f_A, f_B, img_A=img_A_lr, img_B=img_B_lr, bidirectional=self.bidirectional + f_A, + f_B, + img_A=img_A_lr, + img_B=img_B_lr, + bidirectional=self.bidirectional, + return_attention=return_intermediates, ) # return matcher_output - predictions["matcher"] = matcher_output + if return_intermediates: + predictions["matcher"] = matcher_output warp_AB, confidence_AB = ( matcher_output["warp_AB"], matcher_output["confidence_AB"], @@ -253,8 +268,6 @@ def forward( ) else: refiner_output_BA = None - predictions[f"refiner_{patch_size}_AB"] = refiner_output_AB - predictions[f"refiner_{patch_size}_BA"] = refiner_output_BA warp_AB, confidence_AB = ( refiner_output_AB["warp"], refiner_output_AB["confidence"], @@ -264,6 +277,9 @@ def forward( refiner_output_BA["warp"], refiner_output_BA["confidence"], ) + if return_intermediates: + predictions[f"refiner_{patch_size}_AB"] = refiner_output_AB + predictions[f"refiner_{patch_size}_BA"] = refiner_output_BA predictions["warp_AB"] = warp_AB predictions["confidence_AB"] = confidence_AB if self.bidirectional: @@ -306,6 +322,7 @@ def match( self, img_like_A: ImageLike, img_like_B: ImageLike, + output_device: torch.device | str | None = None, ) -> dict[str, torch.Tensor]: self.eval() img_A = self._load_image(img_like_A) @@ -344,7 +361,13 @@ def match( img_A_hr = None img_B_hr = None - preds = self(img_A_lr, img_B_lr, img_A_hr=img_A_hr, img_B_hr=img_B_hr) + preds = self( + img_A_lr, + img_B_lr, + img_A_hr=img_A_hr, + img_B_hr=img_B_hr, + return_intermediates=False, + ) warp_AB = preds["warp_AB"] confidence_AB = preds["confidence_AB"] @@ -362,14 +385,14 @@ def match( precision_BA = None preds = { - "warp_AB": warp_AB.clone(), - "confidence_AB": confidence_AB.clone(), - "overlap_AB": overlap_AB.clone(), - "precision_AB": precision_AB.clone(), - "warp_BA": warp_BA.clone() if warp_BA is not None else None, - "confidence_BA": confidence_BA.clone() if confidence_BA is not None else None, - "overlap_BA": overlap_BA.clone() if overlap_BA is not None else None, - "precision_BA": precision_BA.clone() if precision_BA is not None else None, + "warp_AB": _move_optional_tensor(warp_AB, output_device), + "confidence_AB": _move_optional_tensor(confidence_AB, output_device), + "overlap_AB": _move_optional_tensor(overlap_AB, output_device), + "precision_AB": _move_optional_tensor(precision_AB, output_device), + "warp_BA": _move_optional_tensor(warp_BA, output_device), + "confidence_BA": _move_optional_tensor(confidence_BA, output_device), + "overlap_BA": _move_optional_tensor(overlap_BA, output_device), + "precision_BA": _move_optional_tensor(precision_BA, output_device), } return preds