""" Model definition + loading for Img2LocJakarta 2.0 (DINOv3 + conditional flow matching). Architecture (must match the training checkpoint exactly): DinoV3CFMRegressor ├── encoder_net : DinoFeatureEncoder │ └── backbone : DINOv3 ViT-L/16 (fine-tuned), CLS token -> (B, 1024) ├── velocity_head : VelocityMLP v(y_t, t, cond) -> R^2 (AdaLN residual blocks) └── target_mean/var : buffers, standardiser of the local (x, y) metre targets The model is generative. Gaussian noise "particles" are integrated through the learned velocity field (Heun / RK2) onto the geographic target; the cloud of arrival points is the predictive distribution. `sample_particles` returns it in local metres, `sample_lonlat` in WGS-84 degrees. The checkpoint bundles the fine-tuned backbone weights together with the head, so only the DINOv3 *architecture* source (facebookresearch/dinov3) is needed, with ``pretrained=False`` -- no separate LVD-1689M download. """ from __future__ import annotations import math import os import subprocess import sys from pathlib import Path from typing import Optional import torch import torch.nn as nn # --- constants that are part of the trained model ----------------------------- DINOV3_GITHUB = "https://github.com/facebookresearch/dinov3" BACKBONE_NAME = "dinov3_vitl16" COND_DIM = 1024 HIDDEN_DIM = 512 NUM_BLOCKS = 4 TARGET_EPS = 1e-6 # Local flat-earth projection origin (Jakarta city centre) used during training. ORIGIN_LON = 106.828320 ORIGIN_LAT = -6.227468 EARTH_RADIUS_M = 6371000.0 # --- coordinate conversion ----------------------------------------------------- def local_xy_to_lonlat(xy_meters: torch.Tensor) -> torch.Tensor: """(..., 2) [x east, y north] metres -> (..., 2) [lon, lat] degrees.""" lat0_rad = math.radians(ORIGIN_LAT) dlon = xy_meters[..., 0] / (EARTH_RADIUS_M * math.cos(lat0_rad)) dlat = xy_meters[..., 1] / EARTH_RADIUS_M return torch.stack( [ORIGIN_LON + torch.rad2deg(dlon), ORIGIN_LAT + torch.rad2deg(dlat)], dim=-1 ) # --- modules ------------------------------------------------------------------- class SinusoidalTimeEmbedding(nn.Module): def __init__(self, dim: int = 128): super().__init__() assert dim % 2 == 0 self.dim = dim def forward(self, t: torch.Tensor) -> torch.Tensor: if t.dim() == 1: t = t[:, None] half = self.dim // 2 freqs = torch.exp( torch.linspace(0, math.log(10_000.0), half, device=t.device, dtype=t.dtype) ) args = t * (1.0 / freqs)[None, :] return torch.cat([torch.sin(args), torch.cos(args)], dim=-1) class DinoFeatureEncoder(nn.Module): def __init__(self, backbone: nn.Module, pool: str = "cls"): super().__init__() self.backbone = backbone self.pool = pool if hasattr(self.backbone, "head"): self.backbone.head = nn.Identity() def forward(self, images: torch.Tensor) -> torch.Tensor: out = self.backbone(images) if isinstance(out, (tuple, list)): out = out[0] if out.dim() == 2: return out if out.dim() == 3: return out[:, 0] if self.pool == "cls" else out.mean(dim=1) raise ValueError(f"Unexpected backbone output shape: {tuple(out.shape)}") class _VelocityResBlock(nn.Module): def __init__(self, hidden_dim: int, t_dim: int, dropout: float): super().__init__() self.norm = nn.LayerNorm(hidden_dim) self.adaLN = nn.Linear(t_dim, 2 * hidden_dim) self.ff = nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.SiLU(), nn.Dropout(dropout), nn.Linear(hidden_dim, hidden_dim), ) def forward(self, x, t_emb): scale, shift = self.adaLN(t_emb).chunk(2, dim=-1) h = self.norm(x) * (1.0 + scale) + shift return x + self.ff(h) class VelocityMLP(nn.Module): def __init__(self, cond_dim, y_dim=2, t_dim=128, hidden_dim=512, num_blocks=4, dropout=0.0): super().__init__() self.t_embed = SinusoidalTimeEmbedding(t_dim) self.input_proj = nn.Linear(cond_dim + y_dim, hidden_dim) self.blocks = nn.ModuleList( [_VelocityResBlock(hidden_dim, t_dim, dropout) for _ in range(num_blocks)] ) self.out_norm = nn.LayerNorm(hidden_dim) self.out_proj = nn.Linear(hidden_dim, y_dim) def forward(self, y_t, t, cond): t_emb = self.t_embed(t) x = self.input_proj(torch.cat([cond, y_t], dim=-1)) for block in self.blocks: x = block(x, t_emb) return self.out_proj(self.out_norm(x)) class DinoV3CFMRegressor(nn.Module): """DINOv3 encoder + flow-matching velocity head. `pixel_values` must already be resized to 256x256 and ImageNet-normalised (see ``GeoTagPredictor`` / the transform in ``inference.py``). """ def __init__( self, backbone: nn.Module, cond_dim: int = COND_DIM, hidden_dim: int = HIDDEN_DIM, num_blocks: int = NUM_BLOCKS, ): super().__init__() self.encoder_net = DinoFeatureEncoder(backbone) self.velocity_head = VelocityMLP(cond_dim, 2, 128, hidden_dim, num_blocks, 0.0) self.register_buffer("target_mean", torch.zeros(2)) self.register_buffer("target_var", torch.ones(2)) def encode(self, pixel_values: torch.Tensor) -> torch.Tensor: return self.encoder_net(pixel_values) def velocity(self, y_t, t, cond): return self.velocity_head(y_t, t, cond) @torch.no_grad() def sample_particles( self, pixel_values: torch.Tensor, num_particles: int = 256, steps: int = 8, generator: Optional[torch.Generator] = None, ) -> torch.Tensor: """Heun (RK2) flow of Gaussian noise onto the target. Returns (B, N, 2) metres.""" B, dt = pixel_values.shape[0], 1.0 / steps cond = self.encode(pixel_values).repeat_interleave(num_particles, dim=0) y = torch.randn( B * num_particles, 2, device=cond.device, dtype=cond.dtype, generator=generator ) for s in range(steps): t0 = torch.full((B * num_particles, 1), s / steps, device=cond.device, dtype=cond.dtype) t1 = torch.full((B * num_particles, 1), (s + 1) / steps, device=cond.device, dtype=cond.dtype) k1 = self.velocity(y, t0, cond) k2 = self.velocity(y + dt * k1, t1, cond) y = y + 0.5 * dt * (k1 + k2) std = torch.sqrt(self.target_var + TARGET_EPS) return (y * std + self.target_mean).view(B, num_particles, 2) @torch.no_grad() def sample_lonlat(self, pixel_values, num_particles: int = 256, steps: int = 8, generator=None): """(B, N, 2) particle cloud as [lon, lat] degrees.""" return local_xy_to_lonlat(self.sample_particles(pixel_values, num_particles, steps, generator)) # -- construction helpers -------------------------------------------------- @staticmethod def build_backbone(dinov3_dir: Optional[str] = None) -> nn.Module: """Instantiate the DINOv3 ViT-L/16 architecture (no pretrained download). Source lookup: `dinov3_dir`, $GEOTAG_DINOV3_DIR, an importable `dinov3` package, the torch-hub cache, else a one-off `git clone` into the hub cache. Imports the backbone factory directly (skips hubconf.py's heavy eval deps). """ try: import dinov3 # noqa: F401 except ImportError: src = dinov3_dir or os.environ.get("GEOTAG_DINOV3_DIR") if not src: src = str(Path(torch.hub.get_dir()) / "facebookresearch_dinov3_main") if not Path(src).is_dir(): subprocess.run(["git", "clone", "--depth", "1", DINOV3_GITHUB, src], check=True) sys.path.insert(0, src) from dinov3.hub import backbones return getattr(backbones, BACKBONE_NAME)(pretrained=False) @classmethod def from_pretrained( cls, model_id_or_path: str, *, filename: str = "model.safetensors", device: str | torch.device = "cpu", backbone: Optional[nn.Module] = None, dinov3_dir: Optional[str] = None, ) -> "DinoV3CFMRegressor": """Load weights from a local ``.safetensors`` file or a HuggingFace repo id. The file is a full state_dict: ``encoder_net.backbone.*``, ``velocity_head.*`` and the ``target_mean`` / ``target_var`` standardiser buffers. """ if os.path.isfile(model_id_or_path): weights_path = model_id_or_path elif os.path.isdir(model_id_or_path): weights_path = os.path.join(model_id_or_path, filename) else: from huggingface_hub import hf_hub_download weights_path = hf_hub_download(repo_id=model_id_or_path, filename=filename) if backbone is None: backbone = cls.build_backbone(dinov3_dir) model = cls(backbone) if weights_path.endswith(".safetensors"): from safetensors.torch import load_file state_dict = load_file(weights_path) else: state_dict = torch.load(weights_path, map_location="cpu") missing, unexpected = model.load_state_dict(state_dict, strict=False) bad = [k for k in missing if k.startswith(("velocity_head.", "target_"))] if bad: raise RuntimeError(f"checkpoint is missing required weights: {bad[:5]}") # extra head weights mean the checkpoint was trained with a different head # (e.g. more blocks): loading it would silently drop part of the model extra = [k for k in unexpected if k.startswith(("velocity_head.", "target_"))] if extra: raise RuntimeError(f"checkpoint has head weights this model does not: {extra[:5]}") if missing or unexpected: print(f"[modeling_geotag] missing={missing[:5]} unexpected={unexpected[:5]}") model.to(device).eval() return model