modeling_citywalker.py

"""CityWalker waypoint-prediction model, ported to a HuggingFace `PreTrainedModel`.

Port of `model/citywalker_feat.py` + supporting modules from
https://github.com/ai4ce/CityWalker, stripped of Lightning/OmegaConf.

Architecture (inference-only):

    images (B,T,3,H,W)  ──► DINOv2 ──► obs tokens  (B,T,D)
    coords (B,T+1,2)    ──► PolarEmbedding + Linear ──► goal token  (B,1,D)
                                                    ──► concat ──► (B,T+2,D)
                            ──► TransformerEncoder (self-attention decoder)
                            ──► MLP head  ──► (waypoints_pred, arrive_pred)

Outputs:
    waypoints_pred  : (B, len_traj_pred, 2)   cumulative XY deltas in body frame
    arrive_pred     : (B, 1)                  logits
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn
import torchvision.transforms.functional as TF
from transformers import Dinov2Config, Dinov2Model, PreTrainedModel
from transformers.modeling_outputs import ModelOutput

from .configuration_citywalker import CityWalkerConfig


def _build_obs_encoder(name: str) -> Dinov2Model:
    """Build the DINOv2 backbone, working under both fresh-init and
    `from_pretrained` (which wraps __init__ in a `with torch.device("meta")`
    context starting in transformers 5.x).

    Inside the meta context, calling ``Dinov2Model.from_pretrained`` raises
    because nested `from_pretrained` calls are an anti-pattern: the outer
    loader is responsible for materializing weights. So when we detect the
    meta context, we just build the empty `Dinov2Model(config)` shell — the
    outer `from_pretrained` will populate the encoder weights from the
    bundled safetensors blob (which contains the encoder's weights via
    Phase 2's full-state-dict save).

    Outside the meta context (direct `CityWalkerModel(cfg)` construction),
    we still pull the real DINOv2 weights from `facebook/dinov2-*` so users
    instantiating from scratch get a useful backbone.
    """
    in_meta = (
        torch.device("meta") == _peek_default_device()
    )
    if in_meta:
        return Dinov2Model(Dinov2Config.from_pretrained(name))
    return Dinov2Model.from_pretrained(name)


def _peek_default_device() -> Optional[torch.device]:
    """Return the device set by the outermost `with torch.device(...)` /
    `torch.set_default_device(...)` context, or None if neither is active."""
    try:
        from transformers.modeling_utils import (
            get_torch_context_manager_or_global_device,
        )

        return get_torch_context_manager_or_global_device()
    except Exception:
        return None


# Map our `obs_encoder_type` strings (matching upstream torch.hub names) to
# the corresponding facebook/dinov2-* HF repo. We mirror only the four LVD142M
# no-register variants — same backbones, same weights, just shipped via HF
# instead of torch.hub. This is what lets us drop torch.hub entirely while
# keeping the legacy CityWalker `obs_encoder_type` strings working.
_DINOV2_HF_REPOS = {
    "dinov2_vits14": "facebook/dinov2-small",
    "dinov2_vitb14": "facebook/dinov2-base",
    "dinov2_vitl14": "facebook/dinov2-large",
    "dinov2_vitg14": "facebook/dinov2-giant",
}


@dataclass
class CityWalkerOutput(ModelOutput):
    waypoints: torch.FloatTensor = None
    arrive_logits: torch.FloatTensor = None
    token_features: Optional[torch.FloatTensor] = None
    future_features: Optional[torch.FloatTensor] = None


class PolarEmbedding(nn.Module):
    """Fourier-feature encoding of 2D body-frame coordinates in polar form."""

    def __init__(self, num_freqs: int, include_input: bool):
        super().__init__()
        self.num_freqs = num_freqs
        self.include_input = include_input
        freq_bands = 2.0 ** torch.linspace(0, num_freqs - 1, num_freqs)
        self.register_buffer("freq_bands", freq_bands)
        self.out_dim = (2 if include_input else 0) + 4 * num_freqs

    def forward(self, coords: torch.Tensor) -> torch.Tensor:
        x, y = coords[..., 0], coords[..., 1]
        r = torch.sqrt(x * x + y * y).unsqueeze(-1)
        theta = torch.atan2(y, x).unsqueeze(-1)

        parts = [r, theta] if self.include_input else []
        fb = self.freq_bands.view(1, 1, -1)
        parts.append(torch.sin(theta * fb))
        parts.append(torch.cos(theta * fb))
        parts.append(torch.sin(r * fb))
        parts.append(torch.cos(r * fb))
        return torch.cat(parts, dim=-1)


class _PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding (upstream naming preserved for weight-key parity)."""

    def __init__(self, d_model: int, max_seq_len: int):
        super().__init__()
        pos_enc = torch.zeros(max_seq_len, d_model)
        pos = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pos_enc[:, 0::2] = torch.sin(pos * div_term)
        pos_enc[:, 1::2] = torch.cos(pos * div_term)
        self.register_buffer("pos_enc", pos_enc.unsqueeze(0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pos_enc[:, : x.size(1), :]


class _FeatPredictor(nn.Module):
    """Transformer self-attention stack over (context_size + 2) tokens."""

    def __init__(self, embed_dim: int, seq_len: int, nhead: int, num_layers: int, ff_dim_factor: int):
        super().__init__()
        self.positional_encoding = _PositionalEncoding(embed_dim, max_seq_len=seq_len)
        layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=nhead,
            dim_feedforward=ff_dim_factor * embed_dim,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.sa_layer = layer
        self.sa_decoder = nn.TransformerEncoder(layer, num_layers=num_layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.sa_decoder(self.positional_encoding(x))


class CityWalkerModel(PreTrainedModel):
    """HF-compatible CityWalker model. Inference path only; training stays upstream."""

    config_class = CityWalkerConfig
    base_model_prefix = "citywalker"
    supports_gradient_checkpointing = False
    main_input_name = "images"

    def __init__(self, config: CityWalkerConfig):
        super().__init__(config)
        self.config = config

        if config.do_rgb_normalize:
            self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
            self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))

        if config.obs_encoder_type not in _DINOV2_HF_REPOS:
            raise ValueError(
                f"Unsupported obs_encoder_type: {config.obs_encoder_type!r}. "
                f"Expected one of {sorted(_DINOV2_HF_REPOS)}."
            )
        # DINOv2 backbone. See `_build_obs_encoder` — handles the case where
        # we're inside the outer `from_pretrained`'s meta-device context
        # (transformers 5.x) by building an empty shell that the outer
        # loader will fill from our safetensors blob.
        self.obs_encoder = _build_obs_encoder(
            _DINOV2_HF_REPOS[config.obs_encoder_type]
        )
        if config.freeze_obs_encoder:
            for p in self.obs_encoder.parameters():
                p.requires_grad = False
            self.obs_encoder.eval()
        self._feature_dim = config.feature_dim

        self.cord_embedding = PolarEmbedding(
            num_freqs=config.cord_num_freqs,
            include_input=config.cord_include_input,
        )
        cord_enc_dim = self.cord_embedding.out_dim * (config.context_size + 1)
        self.compress_goal_enc = nn.Linear(cord_enc_dim, self._feature_dim)

        self.predictor = _FeatPredictor(
            embed_dim=self._feature_dim,
            seq_len=config.context_size + 1,
            nhead=config.decoder_num_heads,
            num_layers=config.decoder_num_layers,
            ff_dim_factor=config.decoder_ff_dim_factor,
        )
        self.predictor_mlp = nn.Sequential(
            nn.Linear((config.context_size + 1) * self._feature_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
        )
        self.wp_predictor = nn.Linear(32, config.len_traj_pred * 2)
        self.arrive_predictor = nn.Linear(32, 1)

        self.post_init()

    def _encode_obs(self, x: torch.Tensor) -> torch.Tensor:
        """Run a batch through the DINOv2 backbone and return the CLS token.

        Upstream's torch.hub backbone returns ``head(x_norm_clstoken)`` (head
        is Identity for the pretrained variants), giving (B, feature_dim).
        HF's ``Dinov2Model`` returns ``BaseModelOutputWithPooling`` with
        ``last_hidden_state`` of shape (B, num_patches+1, feature_dim); the
        CLS token is at index 0 along the sequence dim. Using ``[:, 0]`` here
        matches upstream byte-for-byte at inference (same weights, same
        layernorm, same tokenization).
        """
        out = self.obs_encoder(pixel_values=x)
        return out.last_hidden_state[:, 0]

    def _preprocess(self, x: torch.Tensor) -> torch.Tensor:
        if self.config.do_rgb_normalize:
            x = (x - self.mean) / self.std
        if self.config.do_resize:
            x = TF.center_crop(x, list(self.config.crop))
            x = TF.resize(x, list(self.config.resize))
        return x

    def forward(
        self,
        images: torch.Tensor,
        coords: torch.Tensor,
        future_images: Optional[torch.Tensor] = None,
        return_dict: bool = True,
    ):
        """
        Args:
            images:        (B, context_size, 3, H, W) float tensor in [0, 1].
            coords:        (B, context_size + 1, 2) recent body-frame XY positions.
            future_images: optional (B, context_size, 3, H, W) for the
                feature-prediction head (unused at inference).
        """
        B, T, _, H, W = images.shape
        x = self._preprocess(images.view(B * T, 3, H, W))
        obs_enc = self._encode_obs(x).view(B, T, -1)

        future_enc: Optional[torch.Tensor] = None
        if future_images is not None:
            fx = self._preprocess(future_images.view(B * T, 3, H, W))
            future_enc = self._encode_obs(fx).view(B, T, -1)

        cord_enc = self.cord_embedding(coords).view(B, -1)
        cord_enc = self.compress_goal_enc(cord_enc).view(B, 1, -1)

        tokens = torch.cat([obs_enc, cord_enc], dim=1)
        features = self.predictor(tokens)
        dec_out = self.predictor_mlp(features.view(B, -1))

        wp = self.wp_predictor(dec_out).view(B, self.config.len_traj_pred, 2)
        wp = torch.cumsum(wp, dim=1)
        arrive = self.arrive_predictor(dec_out).view(B, 1)

        if not return_dict:
            return wp, arrive, features[:, :-1], future_enc
        return CityWalkerOutput(
            waypoints=wp,
            arrive_logits=arrive,
            token_features=features[:, :-1],
            future_features=future_enc,
        )