Img2LocJakarta 2.0: DINOv3 + flow-matching geolocation (code, weights, model card)

2242b6a verified 1 day ago

10.1 kB

"""
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

	"""
	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