Skip to content

Using local models on filesystem #39

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
servantftransperfect wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Using local models on filesystem #39

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7f12442
Using local models on filesystem
Jun 16, 2026
b780da6
memory optimization
Jun 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions src/romav2/features.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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)
Expand All @@ -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)
Expand Down
12 changes: 9 additions & 3 deletions src/romav2/matcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand All @@ -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(
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
61 changes: 44 additions & 17 deletions src/romav2/romav2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -79,6 +87,7 @@ class Cfg:
setting: Setting = "precise"
compile: bool = False
name: str = "RoMa v2"
weights: dict = None

# settings
H_lr: int
Expand All @@ -94,11 +103,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
Expand Down Expand Up @@ -167,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")
Expand All @@ -179,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"],
Expand Down Expand Up @@ -249,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"],
Expand All @@ -260,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:
Expand Down Expand Up @@ -302,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)
Expand Down Expand Up @@ -340,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"]
Expand All @@ -358,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

Expand Down