add sofia models

conf/gtauav_balanced_sofia.gin

@@ -0,0 +1,39 @@
+# GTA-UAV Balanced (SOFIA-Tiny backbone): SOFIA v7.1 student trained from scratch
+# с двухуровневой text fusion:
+#   1. Mid-level: Text-FiLM в SAT и UAV heads (модулирует feature map перед GGeM/CHP).
+#   2. Late-level: GatedFusion на дескрипторах (как в DINOv3/StripNet вариантах).
+#
+# Trainable (~5-7M):
+#   - SOFIA backbone (Tiny, ~5M, from scratch — нет pretrained)
+#   - SOFIA heads (SatHead GGeM+BN+Linear, UAVHead AltitudeFiLM+CHP+BN+Linear, +Text-FiLM)
+#   - DGTRS-CLIP LoRA (rank=4, ~147K)
+#   - TextFusionMLP (3*768 -> 1024 -> 1024, ~3.4M, shared)
+#   - Gates α_q, α_g + learnable τ
+#
+# Altitude (drone_height метры) подаётся в UAVHead.AltitudeFiLM из dataset meta CSV.
+# Для sat — altitude=None → FiLM passthrough (γ=1, β=0).
+#
+# Note: SOFIA from scratch — нужно больше эпох или warmup, чем frozen DINOv3/StripNet.
+#       Mamba-2 backend (mamba_ssm) даёт 2-8x speedup vs torch fallback.
+
+include 'conf/gtauav_balanced.gin'
+
+# ---- Backbone ----
+TrainConfigGTAUAV.backbone = "sofia"
+TrainConfigGTAUAV.sofia_preset = "Tiny"
+TrainConfigGTAUAV.sofia_d_descriptor = 1024
+TrainConfigGTAUAV.sofia_use_text_film_uav = True
+TrainConfigGTAUAV.sofia_use_text_film_sat = True
+TrainConfigGTAUAV.sofia_lora_rank = 4
+# Mamba-1 used for Tiny (Mamba-2 torch fallback has a pre-existing reshape bug
+# with channels not divisible by default headdim; switch to "mamba2" for M/L
+# presets where channels % 64 == 0 OR install mamba_ssm CUDA kernels).
+TrainConfigGTAUAV.sofia_mamba_variant = "mamba1"
+TrainConfigGTAUAV.sofia_mamba_backend = "auto"  # mamba_ssm if installed else torch fallback
+
+# ---- Training overrides ----
+TrainConfigGTAUAV.gradient_checkpointing = False  # SOFIA from-scratch — keep activations live
+TrainConfigGTAUAV.shared_encoder = True  # ignored by SOFIA but kept for logging compat
+
+# ---- Output ----
+TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_sofia"

conf/gtauav_balanced_sofia_v1.gin

@@ -0,0 +1,42 @@
+# GTA-UAV Balanced (SOFIA v1 backbone): StripNet+DCNv4 hierarchical CNN
+# (~3-30M params depending on variant) trained from scratch с двухуровневой
+# text fusion:
+#   1. Mid-level: Text-FiLM в SAT и UAV heads (модулирует [B,C,8,8] перед GGeM).
+#   2. Late-level: GatedFusion на дескрипторах.
+#
+# UAV head: AltitudeFiLM(drone_height) + [TextFiLM] + GGeM + BN + Linear.
+# SAT head: [TextFiLM] + GGeM + BN + Linear.
+# Один backbone shared между sat/uav.
+#
+# Variant -> размер модели:
+#   tiny_tiny: dims [16, 32, 80, 128]   (~0.4M)
+#   tiny     : dims [32, 64, 128, 256]  (~1M)
+#   small    : dims [64, 128, 320, 512] (~3M, default)
+#   small_v2 : dims [64, 128, 256, 384] (~2M)
+#
+# Trainable (с small variant + text fusion):
+#   - SOFIA v1 backbone (~3M) + heads (~0.6M)
+#   - DGTRS-CLIP LoRA (rank=4, ~147K)
+#   - TextFusionMLP (3*768 -> 1024 -> 1024, ~3.4M, shared)
+#   - Gates α_q, α_g + learnable τ
+# Total trainable ~7M.
+#
+# Note: DCNv4 требует CUDA — обучение только на GPU. Не работает на CPU.
+
+include 'conf/gtauav_balanced.gin'
+
+# ---- Backbone ----
+TrainConfigGTAUAV.backbone = "sofia_v1"
+TrainConfigGTAUAV.sofia_v1_variant = "tiny"
+TrainConfigGTAUAV.sofia_v1_d_descriptor = 1024
+TrainConfigGTAUAV.sofia_v1_use_text_film_uav = True
+TrainConfigGTAUAV.sofia_v1_use_text_film_sat = True
+TrainConfigGTAUAV.sofia_v1_use_film_altitude = True
+TrainConfigGTAUAV.sofia_v1_lora_rank = 4
+
+# ---- Training overrides ----
+TrainConfigGTAUAV.gradient_checkpointing = False  # SOFIA v1 from-scratch
+TrainConfigGTAUAV.shared_encoder = True
+
+# ---- Output ----
+TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_sofia_v1"

conf/gtauav_baseline_sofia.gin

@@ -0,0 +1,13 @@
+# GTA-UAV Baseline (SOFIA-Tiny backbone): no text fusion. Reference R@1 для
+# computing Δ R@1 vs gtauav_balanced_sofia.gin.
+#
+# В baseline_mode=True:
+#   - Text-FiLM отключается (SOFIA heads работают только с altitude).
+#   - DGTRS-CLIP не загружается, TextFusionMLP не строится.
+#   - GatedFusion gates = 1.0 (text игнорируется).
+
+include 'conf/gtauav_balanced_sofia.gin'
+
+TrainConfigGTAUAV.baseline_mode = True
+TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_sofia"
+TrainConfigGTAUAV.use_gradcam = False

conf/gtauav_baseline_sofia_v1.gin

@@ -0,0 +1,8 @@
+# GTA-UAV Baseline (SOFIA v1 backbone): no text fusion. Reference R@1 для
+# computing Δ R@1 vs gtauav_balanced_sofia_v1.gin.
+
+include 'conf/gtauav_balanced_sofia_v1.gin'
+
+TrainConfigGTAUAV.baseline_mode = True
+TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_sofia_v1"
+TrainConfigGTAUAV.use_gradcam = False

scripts/test_dcn.py

@@ -0,0 +1,11 @@
+import torch
+from DCNv4 import DCNv4
+
+dcn = DCNv4(channels=64, group=4).cuda().eval()
+print('--- 50 forward calls, no_grad ---')
+with torch.no_grad():
+    for s in range(50):
+        x = torch.randn(8, 4096, 64, device='cuda')
+        _ = dcn(x)
+        if s % 10 == 0:
+            print(f'{s}: {torch.cuda.memory_allocated() / 1e6:.1f} MB')

src/datasets/gtauav_dataset.py

@@ -104,6 +104,10 @@ class GTAUAVDataset(Dataset):
         image_transform: Fallback single transform for both (used if drone/sat not set).
         filter_meta: Path to seg_filter.json (exclude 90%+ bg/water).
         drop_caption_prob: Probability of dropping captions (ablation).
+        meta_csvs: Optional list of CSV paths with `img_name,drone_height,...`
+            columns (e.g. `cross-area-drone2sate-{train,test}_drone_meta.csv`).
+            When provided, an `altitude` (meters) field is attached per drone
+            sample. Defaults to scanning the rgb_root for *_drone_meta.csv.
         seed: Random seed.
     """
 
@@ -117,6 +121,7 @@ class GTAUAVDataset(Dataset):
         image_transform: Callable[[Image.Image], torch.Tensor] | None = None,
         filter_meta: str | None = None,
         drop_caption_prob: float = 0.0,
+        meta_csvs: list[str] | None = None,
         seed: int = 0,
     ) -> None:
         self.rgb_root = Path(rgb_root)
@@ -131,6 +136,9 @@ class GTAUAVDataset(Dataset):
         if filter_meta is not None:
             self._load_filter(Path(filter_meta))
 
+        # Load drone altitude index (img_name -> meters). Empty dict if no CSVs.
+        self.altitude_index: dict[str, float] = self._load_altitude_index(meta_csvs)
+
         # Load caption index.
         LOGGER.info("📚 Loading caption index from %s", caption_root)
         self.caption_index = _load_caption_index(self.caption_root)
@@ -141,6 +149,43 @@ class GTAUAVDataset(Dataset):
         self._load_pairs(Path(pair_json))
         LOGGER.info("✅ Loaded %d pairs from %s", len(self.entries), pair_json)
 
+    def _load_altitude_index(self, meta_csvs: list[str] | None) -> dict[str, float]:
+        """Build {drone_img_name: altitude_meters} from drone_meta CSVs.
+
+        If `meta_csvs` is None, auto-discovers `*_drone_meta.csv` (TSV format,
+        columns img_name/drone_height/...) under `self.rgb_root`. Missing or
+        unreadable files are skipped silently — altitude defaults to 0.0
+        downstream when an entry is missing.
+        """
+        if meta_csvs is None:
+            meta_csvs = [str(p) for p in sorted(self.rgb_root.glob("*_drone_meta.csv"))]
+            # Prefer `*_drone_meta_new.csv` if present (overrides original).
+            new_csvs = [str(p) for p in sorted(self.rgb_root.glob("*_drone_meta_new.csv"))]
+            meta_csvs = new_csvs or meta_csvs
+        index: dict[str, float] = {}
+        for csv_path in meta_csvs:
+            path = Path(csv_path)
+            if not path.exists():
+                continue
+            try:
+                with path.open() as f:
+                    header = f.readline().rstrip("\n").split("\t")
+                    name_idx = header.index("img_name")
+                    height_idx = header.index("drone_height")
+                    for line in f:
+                        parts = line.rstrip("\n").split("\t")
+                        if len(parts) <= max(name_idx, height_idx):
+                            continue
+                        try:
+                            index[parts[name_idx]] = float(parts[height_idx])
+                        except ValueError:
+                            continue
+            except (OSError, ValueError) as exc:
+                LOGGER.warning("Failed to parse drone meta CSV %s: %s", path, exc)
+        if index:
+            LOGGER.info("📐 Altitude index: %d drones (from %d CSV)", len(index), len(meta_csvs))
+        return index
+
     def _load_filter(self, path: Path) -> None:
         with open(path) as f:
             meta = json.load(f)
@@ -200,6 +245,7 @@ class GTAUAVDataset(Dataset):
                 "caption_l2": l2,
                 "caption_l3": l3,
                 "sat_captions": sat_captions,
+                "altitude": self.altitude_index.get(drone_name, 0.0),
             })
 
     def _load_image(self, directory: str, filename: str, transform: Callable | None = None) -> torch.Tensor:
@@ -267,6 +313,7 @@ class GTAUAVDataset(Dataset):
             "pair_id": entry["drone_name"],
             "sat_name": sat_name,
             "positive_weight": pos_weight,
+            "altitude": float(entry["altitude"]),
         }
 
 
@@ -286,6 +333,7 @@ def collate_gtauav_batch(
         "pair_ids": [b["pair_id"] for b in batch],
         "sat_names": [b["sat_name"] for b in batch],
         "positive_weights": torch.tensor([b["positive_weight"] for b in batch], dtype=torch.float32),
+        "altitude": torch.tensor([b["altitude"] for b in batch], dtype=torch.float32),
     }
 
 
@@ -371,6 +419,7 @@ class GTAUAVDroneQuery(Dataset):
             "caption_l2": entry["caption_l2"],
             "caption_l3": entry["caption_l3"],
             "valid_sat_names": list(entry["sat_candidates"]),
+            "altitude": float(entry.get("altitude", 0.0)),
         }
 
 
@@ -392,4 +441,5 @@ def collate_drone_query(batch: list[dict[str, Any]]) -> dict[str, Any]:
         "caption_l2": [b["caption_l2"] for b in batch],
         "caption_l3": [b["caption_l3"] for b in batch],
         "valid_sat_names": [b["valid_sat_names"] for b in batch],
+        "altitude": torch.tensor([b["altitude"] for b in batch], dtype=torch.float32),
     }

src/models/asymmetric_encoder.py

@@ -462,6 +462,7 @@ class AsymmetricEncoder(nn.Module):
         caption_l1: list[str] | None = None,
         caption_l2: list[str] | None = None,
         caption_l3: list[str] | None = None,
+        altitude: torch.Tensor | None = None,    # noqa: ARG002 — accepted for API parity with SOFIAFusionEncoder
     ) -> torch.Tensor:
         """Encode drone → normalized query embedding with per-sample text mask."""
         drone_feat = self.encode_drone(drone_img)
@@ -492,6 +493,7 @@ class AsymmetricEncoder(nn.Module):
         sat_caption_l1: list[str] | None = None,
         sat_caption_l2: list[str] | None = None,
         sat_caption_l3: list[str] | None = None,
+        altitude: torch.Tensor | None = None,    # noqa: ARG002 — accepted for API parity with SOFIAFusionEncoder
     ) -> dict[str, torch.Tensor]:
         """Forward pass.
 

src/models/sofia_v1/__init__.py

@@ -0,0 +1,41 @@
+"""SOFIA v1 — StripNet + DCNv4 hierarchical CNN backbone for CVGL.
+
+Lightweight 4-stage backbone (~5–30M params depending on variant). Outputs
+features at 4 scales (last stage is 8x8 for 256x256 input).
+
+Variants (in `stripnet_model_dcn.VARIANT_MAP`):
+- `tiny_tiny`: dims [16, 32, 80, 128]
+- `tiny`     : dims [32, 64, 128, 256]
+- `small`    : dims [64, 128, 320, 512]   (default)
+- `small_v2` : dims [64, 128, 256, 384]
+
+Use `SOFIAv1FusionEncoder` from `src.models.sofia_v1_fusion_encoder` for
+end-to-end CVGL training with DGTRS-CLIP captions and altitude.
+"""
+
+from .config import SOFIAv1Config
+from .heads import SatHeadV1, UAVHeadV1
+from .model import SOFIAv1
+from .stripnet_model_dcn import (
+    StripNetDCN,
+    VARIANT_MAP,
+    build_stripnet_dcn,
+    get_stripnet_dcn_small,
+    get_stripnet_dcn_small_v2,
+    get_stripnet_dcn_tiny,
+    get_stripnet_dcn_tiny_tiny,
+)
+
+__all__ = [
+    "SOFIAv1",
+    "SOFIAv1Config",
+    "SatHeadV1",
+    "UAVHeadV1",
+    "StripNetDCN",
+    "build_stripnet_dcn",
+    "get_stripnet_dcn_small",
+    "get_stripnet_dcn_small_v2",
+    "get_stripnet_dcn_tiny",
+    "get_stripnet_dcn_tiny_tiny",
+    "VARIANT_MAP",
+]

src/models/sofia_v1/config.py

@@ -0,0 +1,44 @@
+"""SOFIA v1 configuration.
+
+Lightweight 4-stage StripNet+DCNv4 backbone variants (`tiny_tiny`/`tiny`/`small`
+/`small_v2`) plus simple GGeM-based heads with optional altitude-FiLM and
+text-FiLM modulation.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass
+class SOFIAv1Config:
+    # -------- Backbone --------
+    variant: Literal["tiny_tiny", "tiny", "small", "small_v2"] = "small"
+    in_channels: int = 3
+    input_size: int = 256
+    # DCN op variant. "v2" (default) uses torchvision DeformConv2d — stable.
+    # "v4" uses OpenGVLab DCNv4 — faster but has known C++-extension memory
+    # leak (~9 MB per forward) that OOMs in long training runs.
+    dcn_variant: Literal["v2", "v4"] = "v2"
+
+    # -------- Heads --------
+    d_descriptor: int = 1024
+    return_normalized: bool = False     # False → wrapper handles L2 after gated fusion
+
+    # Altitude-FiLM (UAV head only).
+    use_film_altitude: bool = True
+    altitude_norm: float = 500.0
+
+    # Text-FiLM (mid-level fusion). Both can be toggled independently.
+    use_text_film_uav: bool = True
+    use_text_film_sat: bool = True
+    text_film_dim: int = 1024
+    text_film_hidden: int = 256
+
+    def summary(self) -> str:
+        return (
+            f"SOFIAv1Config(variant={self.variant}, d={self.d_descriptor}, "
+            f"film_alt={self.use_film_altitude}, "
+            f"text_film(sat={self.use_text_film_sat},uav={self.use_text_film_uav}))"
+        )

src/models/sofia_v1/heads.py

@@ -0,0 +1,99 @@
+"""Heads for SOFIA v1 (StripNet+DCNv4 backbone).
+
+Designed parallel to SOFIA v7.1 heads but lighter — basic GGeM pooling
+instead of CHP, optional altitude-FiLM (UAV) and text-FiLM (both).
+
+The heads produce un-normalized D-dim descriptors when `return_normalized=False`,
+so that a fusion wrapper can blend with text via gated fusion before the final
+L2 normalization.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from src.models.sofia_v71.layers import AltitudeFiLM, GGeM, TextFiLM
+
+
+class SatHeadV1(nn.Module):
+    """Satellite head: [TextFiLM] + GGeM + BN + Linear [+ L2]."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        d_descriptor: int,
+        return_normalized: bool = False,
+        use_text_film: bool = False,
+        text_film_dim: int = 1024,
+        text_film_hidden: int = 256,
+    ) -> None:
+        super().__init__()
+        self.return_normalized = return_normalized
+        self.use_text_film = use_text_film
+        if use_text_film:
+            self.text_film = TextFiLM(in_channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
+        self.ggem = GGeM(in_channels)
+        self.bn = nn.BatchNorm1d(in_channels, affine=False)
+        self.proj = nn.Linear(in_channels, d_descriptor)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        text_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.use_text_film:
+            x = self.text_film(x, text_emb)
+        g = self.ggem(x)
+        g = self.bn(g)
+        g = self.proj(g)
+        if self.return_normalized:
+            g = F.normalize(g, p=2, dim=-1)
+        return g
+
+
+class UAVHeadV1(nn.Module):
+    """UAV head: AltitudeFiLM [+ TextFiLM] + GGeM + BN + Linear [+ L2]."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        d_descriptor: int,
+        use_film: bool = True,
+        altitude_norm: float = 500.0,
+        return_normalized: bool = False,
+        use_text_film: bool = False,
+        text_film_dim: int = 1024,
+        text_film_hidden: int = 256,
+    ) -> None:
+        super().__init__()
+        self.return_normalized = return_normalized
+        self.use_film = use_film
+        self.use_text_film = use_text_film
+        if use_film:
+            self.film = AltitudeFiLM(in_channels, altitude_norm=altitude_norm)
+        if use_text_film:
+            self.text_film = TextFiLM(in_channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
+        self.ggem = GGeM(in_channels)
+        self.bn = nn.BatchNorm1d(in_channels, affine=False)
+        self.proj = nn.Linear(in_channels, d_descriptor)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        altitude: Optional[torch.Tensor] = None,
+        text_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.use_film:
+            x = self.film(x, altitude)
+        if self.use_text_film:
+            x = self.text_film(x, text_emb)
+        g = self.ggem(x)
+        g = self.bn(g)
+        g = self.proj(g)
+        if self.return_normalized:
+            g = F.normalize(g, p=2, dim=-1)
+        return g

src/models/sofia_v1/model.py

@@ -0,0 +1,93 @@
+"""SOFIA v1 model: StripNet+DCNv4 backbone + asymmetric Sat/UAV heads.
+
+Architecture:
+    img [B,3,256,256] --> StripNetDCN (4 stages) --> [B, C_4, 8, 8]
+                                                        |
+                                  ┌─────────────────────┴────────────────┐
+                                  ▼                                      ▼
+                              SatHeadV1                            UAVHeadV1
+                              [TextFiLM]                          AltitudeFiLM
+                              GGeM                                [TextFiLM]
+                              BN + Linear                         GGeM
+                              [L2]                                BN + Linear
+                                                                  [L2]
+
+Heads return un-normalized D-dim descriptors (for downstream gated text fusion).
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+
+from .config import SOFIAv1Config
+from .heads import SatHeadV1, UAVHeadV1
+from .stripnet_model_dcn import VARIANT_MAP, StripNetDCN
+
+
+class SOFIAv1(nn.Module):
+    """SOFIA v1: shared StripNet+DCNv4 backbone with asymmetric heads."""
+
+    def __init__(self, cfg: SOFIAv1Config) -> None:
+        super().__init__()
+        if cfg.variant not in VARIANT_MAP:
+            raise ValueError(f"Unknown variant {cfg.variant!r}")
+        self.cfg = cfg
+
+        # Single shared backbone for sat + uav (saves params; v1 stays lightweight).
+        self.backbone: StripNetDCN = VARIANT_MAP[cfg.variant](dcn_variant=cfg.dcn_variant)
+        last_channels = self.backbone.embed_dims[-1]
+        self.feature_channels = last_channels
+
+        self.sat_head = SatHeadV1(
+            in_channels=last_channels,
+            d_descriptor=cfg.d_descriptor,
+            return_normalized=cfg.return_normalized,
+            use_text_film=cfg.use_text_film_sat,
+            text_film_dim=cfg.text_film_dim,
+            text_film_hidden=cfg.text_film_hidden,
+        )
+        self.uav_head = UAVHeadV1(
+            in_channels=last_channels,
+            d_descriptor=cfg.d_descriptor,
+            use_film=cfg.use_film_altitude,
+            altitude_norm=cfg.altitude_norm,
+            return_normalized=cfg.return_normalized,
+            use_text_film=cfg.use_text_film_uav,
+            text_film_dim=cfg.text_film_dim,
+            text_film_hidden=cfg.text_film_hidden,
+        )
+
+    def _extract_last(self, x: torch.Tensor) -> torch.Tensor:
+        """Run backbone, return only the deepest feature map [B, C, 8, 8]."""
+        feats: List[torch.Tensor] = self.backbone(x)
+        return feats[-1]
+
+    def forward(
+        self,
+        sat: Optional[torch.Tensor] = None,
+        uav: Optional[torch.Tensor] = None,
+        altitude: Optional[torch.Tensor] = None,
+        text_emb_sat: Optional[torch.Tensor] = None,
+        text_emb_uav: Optional[torch.Tensor] = None,
+        return_features: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        result: Dict[str, torch.Tensor] = {}
+
+        if sat is not None:
+            f_sat = self._extract_last(sat)
+            g_sat = self.sat_head(f_sat, text_emb=text_emb_sat)
+            result["g_sat"] = g_sat
+            if return_features:
+                result["features_sat"] = f_sat
+
+        if uav is not None:
+            f_uav = self._extract_last(uav)
+            g_uav = self.uav_head(f_uav, altitude=altitude, text_emb=text_emb_uav)
+            result["g_uav"] = g_uav
+            if return_features:
+                result["features_uav"] = f_uav
+
+        return result

src/models/sofia_v1/stripnet_blocks_dcn.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+from torchvision.ops import DeformConv2d
+
+class DCNBlock(nn.Module):
+    """
+    StripNet-style block but uses deformable conv instead of rigid convs.
+    """
+    def __init__(self, in_ch, out_ch, hidden_ratio=0.25, modulation=True):
+        super().__init__()
+        hidden_ch = max(1, int(out_ch * hidden_ratio))
+
+        # 1x1 reduce
+        self.reduce = nn.Sequential(
+            nn.Conv2d(in_ch, hidden_ch, kernel_size=1, bias=False),
+            nn.BatchNorm2d(hidden_ch),
+            nn.ReLU(inplace=True),
+        )
+
+        # Offset conv (predicts offsets and optional mask)
+        offset_channels = 2 * 3 * 3 if not modulation else 3 * 3 * 3
+        self.offset_conv = nn.Conv2d(hidden_ch, offset_channels,
+                                     kernel_size=3, padding=1)
+
+        # Deformable conv
+        self.deform = DeformConv2d(hidden_ch, hidden_ch,
+                                   kernel_size=3, padding=1, bias=False)
+        self.bn = nn.BatchNorm2d(hidden_ch)
+        self.act = nn.ReLU(inplace=True)
+
+        # 1x1 expand
+        self.expand = nn.Sequential(
+            nn.Conv2d(hidden_ch, out_ch, kernel_size=1, bias=False),
+            nn.BatchNorm2d(out_ch),
+        )
+
+        self.residual = (in_ch == out_ch)
+
+    def forward(self, x):
+        identity = x
+
+        x = self.reduce(x)
+        offset = self.offset_conv(x)
+
+        if offset.shape[1] == 18:  # DCNv1
+            x = self.deform(x, offset)
+        else:  # DCNv2: last 9 channels are mask
+            o, mask = offset.split([18, 9], dim=1)
+            mask = mask.sigmoid()
+            x = self.deform(x, o, mask)
+
+        x = self.act(self.bn(x))
+        x = self.expand(x)
+
+        if self.residual:
+            x = x + identity
+        return x

src/models/sofia_v1/stripnet_blocks_dcn_new.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+from DCNv4 import DCNv4  # ← your installed OpenGVLab version
+import math
+
+class DCNBlockV4(nn.Module):
+    """
+    StripNet-style block that uses OpenGVLab DCNv4 instead of torchvision DCNv2.
+    """
+
+    def __init__(self, in_ch, out_ch, hidden_ratio=0.25,
+                 kernel_size=3, stride=1, dilation=1, group=4,
+                 offset_scale=1.0, use_bias=False):
+        super().__init__()
+        assert kernel_size in (3, 5, 7)
+        pad = (kernel_size // 2) * dilation
+
+        # Hidden channels — must satisfy (hidden_ch // group) % 16 == 0
+        hidden_ch = max(16, int(out_ch * hidden_ratio))
+        hidden_ch = math.ceil(hidden_ch / 16) * 16
+        # increase until kernel constraint satisfied
+        while (hidden_ch // group) % 16 != 0:
+            hidden_ch += 16
+
+        #print(f"[DCNv4] adjusted hidden_ch={hidden_ch}, group={group}")
+
+
+        self.reduce = nn.Sequential(
+            nn.Conv2d(in_ch, hidden_ch, 1, bias=False),
+            self.make_gn(hidden_ch),
+            nn.ReLU(inplace=True),
+        )
+
+        # DCNv4 core
+        self.dcn = DCNv4(
+            channels=hidden_ch,
+            kernel_size=kernel_size,
+            stride=stride,
+            pad=pad,
+            dilation=dilation,
+            group=group,
+            offset_scale=offset_scale,
+            dw_kernel_size=None,
+            center_feature_scale=False,
+            remove_center=False,
+            output_bias=True,
+            without_pointwise=False,
+        )
+
+
+        self.expand = nn.Sequential(
+            nn.Conv2d(hidden_ch, out_ch, 1, bias=False),
+            self.make_gn(out_ch),
+
+        )
+
+        self.residual = (in_ch == out_ch and stride == 1)
+        self.stride = stride
+
+    def make_gn(self, num_channels):
+        num_groups = max(1, num_channels // 16)
+        return nn.GroupNorm(num_groups, num_channels)
+
+    def forward(self, x):
+        """
+        Input:  [B, C, H, W]
+        Output: [B, C, H', W']
+        """
+        identity = x
+
+        x = self.reduce(x)
+        B, C, H, W = x.shape
+
+        # Flatten for DCNv4
+        x_seq = x.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)
+
+        # --- clamp activations to avoid huge values ---
+        x_seq = torch.clamp(x_seq, -50.0, 50.0)
+
+        x_seq = self.dcn(x_seq)
+        x_seq = torch.nan_to_num(x_seq, nan=0.0, posinf=1e4, neginf=-1e4)
+
+        # --- back to (B, C, H, W) ---
+        x = x_seq.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
+
+        # --- expand + residual ---
+        x = self.expand(x)
+        if self.residual:
+            x = x + identity
+        return x
+
+
+def test_dcnblock_v4():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+
+    block = DCNBlockV4(
+        in_ch=128,  # must match input
+        out_ch=128,
+        hidden_ratio=0.25,
+        kernel_size=3,
+        stride=1,
+        dilation=1,
+        group=5,
+        offset_scale=1.0,
+    ).to(device)
+
+    # ✅ internal feature map (not RGB)
+    x = torch.randn(2, 128, 128, 128, device=device, requires_grad=True)
+
+    with torch.cuda.amp.autocast_mode.autocast(enabled=torch.cuda.is_available()):
+        y = block(x)
+
+    print(f"Input  shape : {x.shape}")
+    print(f"Output shape : {y.shape}")
+    assert y.shape == x.shape
+
+    loss = y.mean()
+    loss.backward()
+    print("✅ DCNBlockV4 test passed.\n")
+
+
+
+if __name__ == "__main__":
+    test_dcnblock_v4()

src/models/sofia_v1/stripnet_model_dcn.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+
+from src.models.sofia_v1.stripnet_blocks_dcn import DCNBlock
+from src.models.sofia_v1.stripnet_blocks_dcn_new import DCNBlockV4
+from src.models.stripnet.model import OverlapPatchEmbed
+
+
+class StripNetDCN(nn.Module):
+    """4-stage hierarchical CNN backbone: OverlapPatchEmbed + DCN blocks per stage.
+
+    Output: list of [B, C_i, H_i, W_i] features at each stage. For 256x256 input
+    with default downsampling (4, 2, 2, 2), final stage is 8x8.
+
+    DCN variant:
+      - "v2" (default): torchvision `DeformConv2d` — stable, no memory leaks.
+      - "v4": OpenGVLab DCNv4 — faster on CUDA but has a known C++ extension
+        memory leak (~9 MB per forward call) that causes OOM in long training
+        runs. Only use if you have a patched DCNv4 build.
+    """
+
+    def __init__(
+        self,
+        in_chans: int = 3,
+        embed_dims: list[int] = [64, 128, 256, 512],
+        depths: list[int] = [3, 4, 6, 3],
+        dcn_variant: str = "v2",
+    ) -> None:
+        super().__init__()
+        if dcn_variant not in ("v2", "v4"):
+            raise ValueError(f"dcn_variant must be 'v2' or 'v4', got {dcn_variant!r}")
+        self.num_stages = len(embed_dims)
+        self.embed_dims = embed_dims
+        self.dcn_variant = dcn_variant
+
+        for i in range(self.num_stages):
+            patch_embed = OverlapPatchEmbed(
+                patch_size=7 if i == 0 else 3,
+                stride=4 if i == 0 else 2,
+                in_chans=in_chans if i == 0 else self.embed_dims[i - 1],
+                embed_dim=self.embed_dims[i],
+            )
+
+            block_cls = DCNBlockV4 if dcn_variant == "v4" else DCNBlock
+            block = nn.ModuleList([
+                block_cls(
+                    in_ch=self.embed_dims[i],
+                    out_ch=self.embed_dims[i],
+                ) for _ in range(depths[i])
+            ])
+
+            setattr(self, f"patch_embed{i + 1}", patch_embed)
+            setattr(self, f"block{i + 1}", block)
+
+    def forward_features(self, x: torch.Tensor) -> list[torch.Tensor]:
+        outs = []
+        for i in range(self.num_stages):
+            patch_embed = getattr(self, f"patch_embed{i + 1}")
+            block = getattr(self, f"block{i + 1}")
+            x, H, W = patch_embed(x)
+            for blk in block:
+                x = blk(x)
+            outs.append(x)
+        return outs
+
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        return self.forward_features(x)
+
+def get_stripnet_dcn_small(dcn_variant: str = "v2") -> StripNetDCN:
+    return StripNetDCN(
+        in_chans=3,
+        embed_dims=[64, 128, 320, 512],
+        depths=[2, 2, 4, 2],
+        dcn_variant=dcn_variant,
+    )
+
+
+def get_stripnet_dcn_small_v2(dcn_variant: str = "v2") -> StripNetDCN:
+    return StripNetDCN(
+        in_chans=3,
+        embed_dims=[64, 128, 256, 384],
+        depths=[2, 2, 4, 2],
+        dcn_variant=dcn_variant,
+    )
+
+
+def get_stripnet_dcn_tiny(dcn_variant: str = "v2") -> StripNetDCN:
+    return StripNetDCN(
+        in_chans=3,
+        embed_dims=[32, 64, 128, 256],
+        depths=[3, 3, 5, 2],
+        dcn_variant=dcn_variant,
+    )
+
+
+def get_stripnet_dcn_tiny_tiny(dcn_variant: str = "v2") -> StripNetDCN:
+    return StripNetDCN(
+        in_chans=3,
+        embed_dims=[16, 32, 80, 128],
+        depths=[3, 3, 5, 2],
+        dcn_variant=dcn_variant,
+    )
+
+
+VARIANT_MAP = {
+    "small": get_stripnet_dcn_small,
+    "small_v2": get_stripnet_dcn_small_v2,
+    "tiny": get_stripnet_dcn_tiny,
+    "tiny_tiny": get_stripnet_dcn_tiny_tiny,
+}
+
+
+def build_stripnet_dcn(variant: str = "small", dcn_variant: str = "v2") -> StripNetDCN:
+    """Factory: variant in {tiny_tiny, tiny, small, small_v2}, dcn_variant in {v2, v4}."""
+    if variant not in VARIANT_MAP:
+        raise ValueError(f"Unknown variant {variant!r}. Available: {list(VARIANT_MAP)}")
+    return VARIANT_MAP[variant](dcn_variant=dcn_variant)

src/models/sofia_v71/README.md

@@ -0,0 +1,331 @@
+# SOFIA v7.1 — PyTorch implementation
+
+Reference implementation of **SOFIA v7.1** CVGL student model targeting Jetson Orin NX with 500 MB – 1 GB VRAM after INT8 quantization.
+
+Full design rationale: see `2_hypotesis/temp_hypotesis/HYP_SOFIA_v7_UltraDeep_дизайн.md`.
+
+## Architecture
+
+```
+Input (256×256×3, sat or UAV)
+   ↓
+Stem: dual-conv (3→C_mid→C_stem_out), ×2 down
+   ↓
+Stage 1 (shared sat/UAV): StripDCN-lite × d1, ×2 down
+   ↓
+Stage 2 (shared): StripMixConv × d2, ×2 down
+   ↓
+Stage 3 (separate): MambaVision MV5 × d3, ×2 down
+   ↓
+Stage 4 (separate): MambaVision MV1 × d4, ×2 down (→ 8×8)
+   ↓
+Ultra-lite 1×1 Neck → F̃ [B, C_n, 8, 8]
+   ↓
+┌─ Sat-Head: GGeM → BN → Linear → L2 → g_sat
+└─ UAV-Head: FiLM(altitude) → CHP (polar+FFT+mag) → BN → Linear → L2 → g_uav
+```
+
+## Presets
+
+| Preset | Params | FLOPs | INT8 weights | FP16 weights | Target latency |
+|--------|-------:|------:|-------------:|-------------:|---------------:|
+| **M** (default) | ~500 M | ~132 G | ~500 MB | ~1 GB | ~18 ms |
+| **L** | ~1 B | ~283 G | ~1 GB | ~2 GB | ~20 ms |
+| Tiny | ~5 M | ~1.4 G | ~5 MB | ~10 MB | ~4 ms |
+
+*Latency estimates на Jetson Orin NX 8GB (INT8 TRT mixed-precision для DCN/Mamba в FP16).*
+
+## Dependencies
+
+Required:
+- Python ≥ 3.9
+- PyTorch ≥ 2.0
+- torchvision ≥ 0.15 (для `deform_conv2d`)
+
+Optional (для production Mamba speed):
+- [`mamba_ssm`](https://github.com/state-spaces/mamba) — ускоренные backends:
+  - v1 `selective_scan_fn` для Mamba-1 (~5× vs Python loop)
+  - v2 `Mamba2` модуль для Mamba-2 SSD dual form (~2–8× vs Mamba-1)
+- [`causal-conv1d`](https://github.com/Dao-AILab/causal-conv1d) — ускоренный 1D conv для Mamba
+
+## Mamba variants (приоритет по качеству/скорости)
+
+| Variant | Описание | Speed | Quality | Когда использовать |
+|---------|----------|:-----:|:-------:|--------------------|
+| **`mamba2`** (default) | SSD dual form, scalar A per head | **2–8× faster** | best | Main choice если mamba_ssm v2 доступен |
+| `mamba1` | Original selective scan, diagonal A | ref | +0% vs mamba2 | Legacy / reproducibility / если v2 недоступен |
+| `efficient_vmamba` | Atrous scan (2 directions, no CUDA kernel) | 2–3× faster | ~−0.3% R@1 | Speed fallback без mamba_ssm |
+
+### Выбор через config
+
+```python
+from code_sofia_v71 import sofia_m_config, SOFIAv71
+
+cfg = sofia_m_config()
+
+# Вариант 1: Mamba-2 (default, preferred) — если mamba_ssm v2 доступен
+cfg.mamba_variant = "mamba2"
+cfg.mamba_backend = "auto"   # falls back to torch if unavailable
+
+# Вариант 2: EfficientVMamba — speed без зависимости от mamba_ssm
+cfg.mamba_variant = "efficient_vmamba"
+
+# Вариант 3: Mamba-1 — legacy / для сравнения
+cfg.mamba_variant = "mamba1"
+
+model = SOFIAv71(cfg)
+```
+
+### Проверка доступности backends
+
+```python
+from code_sofia_v71 import is_mamba_ssm_available, is_mamba2_available
+
+print(f"Mamba-1 CUDA: {is_mamba_ssm_available()}")
+print(f"Mamba-2 CUDA: {is_mamba2_available()}")
+```
+
+### Несовместимость параметров между variants
+
+Mamba-1 и Mamba-2 используют **разную параметризацию $A$** (diagonal-per-channel vs scalar-per-head) — state_dict **НЕ совместим**. EfficientVMamba имеет N независимых scanners каждый с собственными параметрами Mamba-1.
+
+**Следствие:** нельзя train в `mamba1` и load в `mamba2` checkpoint. Выбор variant нужно зафиксировать до training.
+
+## Quick start
+
+```python
+import torch
+from code_sofia_v71 import build_sofia
+
+# Default = SOFIA-M (500 MB INT8 target)
+model = build_sofia("M")
+model.eval()
+
+sat = torch.randn(2, 3, 256, 256)   # satellite image
+uav = torch.randn(2, 3, 256, 256)   # UAV image
+altitude = torch.tensor([120.0, 450.0])   # meters
+
+out = model(sat=sat, uav=uav, altitude=altitude)
+g_sat = out["g_sat"]   # [2, 512] L2-normalized
+g_uav = out["g_uav"]   # [2, 512] L2-normalized
+
+# Retrieval similarity
+similarity = (g_sat * g_uav).sum(dim=-1)   # cosine
+```
+
+## Verification
+
+```bash
+# Smoke test (Tiny preset, CPU, fast)
+python -m sofia_v71.verify
+
+# Full SOFIA-M check with latency benchmark (GPU)
+python -m sofia_v71.verify --preset M --device cuda --benchmark
+
+# Maximum scale
+python -m sofia_v71.verify --preset L --device cuda
+```
+
+Expected output for SOFIA-M:
+```
+Total params: ~500 M
+...
+FP16 weights: ~1000 MB  (~0.98 GB)
+INT8 weights: ~500 MB   (~0.49 GB)
+...
+out[g_sat]: (1, 512)
+out[g_uav]: (1, 512)
+```
+
+## Key novel components
+
+### `layers.py`
+- **`GGeM`** — per-channel learnable exponent Generalized Mean pooling (F11)
+- **`CircularHarmonicPool`** — formally SO(2)-invariant UAV pool via polar → 1D FFT → magnitude (NOVEL NH2)
+- **`AltitudeFiLM`** — telemetry-conditioned modulation (NOVEL NH4)
+- **`RoPE2D`** — 2D rotary positional embedding
+
+### `blocks.py`
+- **`StripDCNLiteBlock`** — Strip DW MBConv с DCN offset на одной оси (NOVEL)
+- **`StripMixConvBlock`** — Strip + MixConv (3/5/7 kernels) для multi-scale
+- **`MambaVisionBlock`** — Mamba ∥ MHSA [∥ Strip] + FFN (MV5 variant — NOVEL 3-way)
+- **`SimpleMambaBlock`** — reference pure-PyTorch Mamba-1 (replace with `mamba_ssm` в production)
+
+### `model.py`
+- **`SOFIAv71`** — full model with optional weight-sharing
+- **`SatHead` vs `UAVHead`** — asymmetric physics-motivated design (NOVEL NH1)
+- **`RingAuxHead`** — LPN Square-Ring training-only aux
+
+## KD taps
+
+Backbone forward returns features at stages s0/f1/f2/f3/f4. Использовать для
+hierarchical knowledge distillation:
+
+```python
+out = model(sat=sat, uav=uav, return_features=True)
+f3_sat = out["features_sat"]["f3"]   # для teacher feature alignment
+```
+
+## Production notes
+
+### Mamba backend selection
+
+Default `mamba_variant="mamba2"` с `backend="auto"`:
+- Use `Mamba2` from `mamba_ssm.modules.mamba2` если установлено (CUDA, fast)
+- Иначе fallback на pure-PyTorch simplified SSD scan (slow, работает)
+
+Для Mamba-1 legacy путь — `mamba_variant="mamba1"`, backend так же резолвится.
+
+Для EfficientVMamba — pure PyTorch, не нуждается в mamba_ssm.
+
+## Quantization (`quant.py`)
+
+Reusable utilities для INT8 PTQ/QAT в SOFIA.
+
+### `OffsetClampSTE` — DCN-M2
+
+Hard clamp DCN offsets к `[-k, +k]` со STE backward. Уже подключён в `StripDCNLiteBlock` через флаг `use_offset_clamp_ste=True` (default).
+
+```python
+from code_sofia_v71 import OffsetClampSTE, offset_clamp_ste
+
+# Module form
+clamp = OffsetClampSTE(kernel_size=7)
+clamped = clamp(offsets)
+
+# Functional form
+clamped = offset_clamp_ste(offsets, kernel_size=7)
+```
+
+Backward — identity (gradient проходит насквозь). Можно ставить с epoch 1, не только в QAT.
+
+### `KScaledFakeQuant` — k-scaled fake quantization
+
+Multi-bin fake quant для long-tail distributions (Mamba Δ, y).
+
+```python
+from code_sofia_v71 import KScaledFakeQuant
+
+fq = KScaledFakeQuant(num_bins=3)
+
+# Calibration
+fq.start_calibration()
+with torch.no_grad():
+    for batch in cal_loader:
+        _ = model(batch)
+fq.finalize_calibration(percentile=99.9)
+# fq теперь в PTQ-режиме
+```
+
+### `KScaledMamba2Block` — drop-in для Mamba2Block
+
+Подкласс `Mamba2Block` с k-scaled fake-quant узлами на критических путях
+(`x_main`, `delta`, `y`). Внутреннее состояние scan'а `h_t` остаётся в
+model dtype (R5 reparam principle).
+
+```python
+from code_sofia_v71 import KScaledMamba2Block
+
+mamba = KScaledMamba2Block(
+    channels=192,
+    d_state=64,
+    headdim=64,
+    num_bins=3,
+    targets=("x_main", "delta", "y"),    # subset
+)
+
+mamba.start_calibration()
+# ... run calibration data ...
+mamba.finalize_calibration(percentile=99.9)
+```
+
+**Constraint:** только `backend='torch'` поддерживается (mamba_ssm CUDA kernel
+не имеет hooks для k-scaled fake-quant). Для full INT8 deploy на TRT нужен
+custom plugin — отдельный deploy concern.
+
+### Model-wide helpers
+
+```python
+from code_sofia_v71 import start_calibration, finalize_calibration, set_quant_enabled
+
+# Начать калибровку всех KScaledFakeQuant и KScaledMamba2Block в модели
+start_calibration(model)
+with torch.no_grad():
+    for batch in cal_loader:
+        _ = model(batch)
+finalize_calibration(model, percentile=99.9)
+
+# Toggle on/off для FP-vs-INT8 ablation
+set_quant_enabled(model, False)    # FP forward
+y_fp = model(x)
+set_quant_enabled(model, True)
+y_q = model(x)
+```
+
+### Smoke test
+
+```bash
+python -m sofia_v71.quant
+```
+
+Прогонит unit-test на:
+- DCN-M2 clamp (gradient pass-through)
+- KScaledFakeQuant calibration + quant error
+- KScaledMamba2Block FP-vs-INT8 diff
+
+### DCN INT8 QAT
+StripDCN использует `torchvision.ops.deform_conv2d`. Для INT8 deploy на TRT:
+- Применить QAT с DCN-M1..M4 modifications (см. HYP Phase 7)
+- Offset predictor: per-channel scale, offset clamping `[-k, +k]`
+- Mask: FP16 micro-block внутри INT8 graph
+
+### TensorRT export
+```python
+import torch
+model.eval()
+model_fuse = model  # apply fuse_reparam passes separately
+torch.onnx.export(
+    model,
+    (sat, uav, altitude),
+    "sofia_v71.onnx",
+    input_names=["sat", "uav", "altitude"],
+    output_names=["g_sat", "g_uav"],
+    opset_version=17,
+)
+# Then:
+# trtexec --onnx=sofia_v71.onnx --int8 --fp16 --saveEngine=sofia.plan
+```
+
+## Training
+
+See `HYP_SOFIA_v7_UltraDeep_дизайн.md` Phase 8 for full training recipe:
+- Loss: InfoNCE (Sample4Geo mining) + Ring aux + (opt) cross-view consistency
+- Curriculum: PALW sigmoid warmup, 60 epochs, 4 phases
+- Optimizer: AdamW, LR 3e-4 cosine, wd 0.05
+- Temperature: τ 0.1 → 0.01 cosine decay
+
+## File layout
+
+```
+code_sofia_v71/
+├── __init__.py       — public exports
+├── config.py         — SOFIAConfig + presets (M/L/Tiny)
+├── layers.py         — GGeM, CHP, FiLM, RoPE2D, SE, LayerNorm2d
+├── blocks.py         — StripDCN, StripMixConv, Mamba, MambaVision, Downsample
+├── model.py          — Stem, Backbone, Neck, Heads, SOFIAv71
+├── verify.py         — parameter counter + benchmark script
+└── README.md         — this file
+```
+
+## References to design doc
+
+Все architectural decisions обоснованы в
+`2_hypotesis/temp_hypotesis/HYP_SOFIA_v7_UltraDeep_дизайн.md`:
+- Phase 1: requirements R1–R13
+- Phase 2': DCN / Strip / MambaVision operator catalog
+- Phase 3': backbone design with 4 candidates (E/F/G)
+- Phase 4'': CVGL-Aware Head v7.1-α (Asymmetric + CHP + FiLM)
+- Phase 5': ablation matrix
+- Phase 6: MambaVision MV5 operational details
+- Phase 7: StripDCN QAT strategy
+- Phase 8: training pipeline

src/models/sofia_v71/__init__.py

@@ -0,0 +1,126 @@
+"""SOFIA v7.1 — student model for cross-view geo-localization (CVGL).
+
+Architecture:
+    stem → stage1 (StripDCN-lite) → stage2 (StripMixConv)
+         → stage3 (MambaVision MV5) → stage4 (MambaVision MV1)
+         → 1×1 neck → asymmetric Sat/UAV heads
+
+Key novel components:
+- StripDCN-lite: Strip DW with adaptive offset on one axis
+- MambaVision MV5: Mamba ∥ MHSA ∥ Strip 3-way parallel
+- CircularHarmonicPool: formally SO(2)-invariant UAV head
+- AltitudeFiLM: telemetry-aware conditioning
+
+Presets (see config.py):
+- SOFIA-M (~500 M params, 500 MB INT8) — default
+- SOFIA-L (~1 B params, 1 GB INT8) — max scale
+- SOFIA-Tiny (~5 M) — reference from original v7.1 spec
+
+Quick start:
+    from sofia_v71 import build_sofia
+    model = build_sofia("M")
+    out = model(sat=sat_tensor, uav=uav_tensor, altitude=alt_tensor)
+    g_sat, g_uav = out["g_sat"], out["g_uav"]
+
+See README.md and HYP_SOFIA_v7_UltraDeep_дизайн.md for full design rationale.
+"""
+
+from .config import (
+    SOFIAConfig,
+    sofia_l_config,
+    sofia_m_config,
+    sofia_tiny_config,
+    DEFAULT_CONFIG,
+)
+from .model import (
+    SOFIAv71,
+    build_sofia,
+    Backbone,
+    Stem,
+    UltraLiteNeck,
+    SatHead,
+    UAVHead,
+    RingAuxHead,
+)
+from .layers import (
+    GGeM,
+    CircularHarmonicPool,
+    AltitudeFiLM,
+    TextFiLM,
+    RoPE2D,
+    SqueezeExcite,
+    LayerNorm2d,
+)
+from .blocks import (
+    StripDCNLiteBlock,
+    StripMixConvBlock,
+    SimpleMambaBlock,
+    Mamba2Block,
+    EfficientVMambaBlock,
+    MambaVisionBlock,
+    EVSSBridge,
+    Downsample,
+    build_mamba_block,
+    is_mamba_ssm_available,
+    is_mamba2_available,
+)
+from .quant import (
+    # DCN-M2
+    OffsetClampSTE,
+    offset_clamp_ste,
+    # k-scaled
+    KScaledFakeQuant,
+    KScaledMamba2Block,
+    # model-wide helpers
+    start_calibration,
+    finalize_calibration,
+    set_quant_enabled,
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    # Config
+    "SOFIAConfig",
+    "sofia_m_config",
+    "sofia_l_config",
+    "sofia_tiny_config",
+    "DEFAULT_CONFIG",
+    # Top-level
+    "SOFIAv71",
+    "build_sofia",
+    # Model parts
+    "Backbone",
+    "Stem",
+    "UltraLiteNeck",
+    "SatHead",
+    "UAVHead",
+    "RingAuxHead",
+    # Layers
+    "GGeM",
+    "CircularHarmonicPool",
+    "AltitudeFiLM",
+    "TextFiLM",
+    "RoPE2D",
+    "SqueezeExcite",
+    "LayerNorm2d",
+    # Blocks
+    "StripDCNLiteBlock",
+    "StripMixConvBlock",
+    "SimpleMambaBlock",
+    "Mamba2Block",
+    "EfficientVMambaBlock",
+    "MambaVisionBlock",
+    "EVSSBridge",
+    "Downsample",
+    "build_mamba_block",
+    "is_mamba_ssm_available",
+    "is_mamba2_available",
+    # Quantization
+    "OffsetClampSTE",
+    "offset_clamp_ste",
+    "KScaledFakeQuant",
+    "KScaledMamba2Block",
+    "start_calibration",
+    "finalize_calibration",
+    "set_quant_enabled",
+]

src/models/sofia_v71/config.py

@@ -0,0 +1,210 @@
+"""SOFIA v7.1 configuration system.
+
+Two scale presets targeting Jetson Orin NX INT8 deployment:
+
+| Preset   | Params   | INT8 size | FP16 size | Target latency |
+|----------|---------:|----------:|----------:|---------------:|
+| SOFIA-M  | ~500 M   | ~500 MB   | ~1 GB     | ~18 ms         |
+| SOFIA-L  | ~1 B     | ~1 GB     | ~2 GB     | ~20 ms         |
+
+Based on HYP_SOFIA_v7_UltraDeep_дизайн.md (Phases 2'–5' + Phase 4''
+CVGL-Aware Head + revision to ultra-lite 1x1 neck).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List, Literal, Optional
+
+
+@dataclass
+class SOFIAConfig:
+    """Configuration for SOFIA v7.1 student architecture.
+
+    Controls backbone width/depth, neck, and CVGL-Aware Head.
+    """
+
+    # -------- Input --------
+    input_size: int = 256
+    in_channels: int = 3
+
+    # -------- Stem (dual-conv FastViT-style) --------
+    stem_mid: int = 64          # intermediate channels
+    stem_out: int = 128         # output channels into stage 1
+
+    # -------- Backbone dimensions per stage s1..s4 --------
+    embed_dims: List[int] = field(
+        default_factory=lambda: [256, 512, 1280, 1536]
+    )
+    depths: List[int] = field(default_factory=lambda: [3, 4, 15, 3])
+
+    # -------- Stage 1–2 block parameters --------
+    mbconv_expand: int = 4
+    se_ratio: int = 16
+    strip_kernel_s1: int = 7    # strip DW kernel size stage 1
+    strip_kernel_s2: int = 5    # strip DW kernel size stage 2
+    mix_kernels: List[int] = field(
+        default_factory=lambda: [3, 5, 7]
+    )   # MixConv DW kernel sizes stage 2
+    use_dcn_strip: bool = True  # adaptive offset on horizontal strip
+
+    # -------- Stage 3–4 (MambaVision) --------
+    mamba_d_state: int = 16
+    mamba_dt_rank: Optional[int] = None     # auto = max(1, C // 16)
+    mamba_backend: Literal["auto", "torch", "mamba_ssm"] = "auto"
+    # "auto" uses mamba_ssm if importable, else torch fallback
+
+    # Mamba variant — one of:
+    #   "mamba2"            (preferred, SSD dual form, 2-8x faster)
+    #   "mamba1"            (original selective scan, mature)
+    #   "efficient_vmamba"  (speed fallback, atrous scan ~2-3x speedup)
+    mamba_variant: Literal["mamba1", "mamba2", "efficient_vmamba"] = "mamba2"
+    # Per-variant tunables passed through to factory
+    mamba_extra_kwargs: dict = field(default_factory=lambda: {
+        "d_state_mamba2": 64,   # Mamba-2 typically uses N=64 (not 16)
+        "headdim": 64,          # Mamba-2 head dim
+        "expand": 2,            # Mamba-2 inner expansion factor
+        "d_conv": 4,            # Mamba-2 local conv kernel
+        "n_directions": 2,      # EfficientVMamba: 2 or 4 atrous directions
+    })
+
+    num_heads_s3: int = 8
+    num_heads_s4: int = 8
+    use_strip_branch_s3: bool = True    # MV5 with Strip branch
+    use_strip_branch_s4: bool = False   # MV1 without Strip branch
+    ffn_expand: int = 4
+
+    # -------- EVSS-style bridge (B6-inspired, opt-in) --------
+    # When True, inserts a within-resolution dual-path refinement block
+    # right after each downsample to stage 3 and stage 4. Smooths semantic
+    # gap between heterogeneous stages (DCN → MambaVision). Adds ~165K params
+    # and ~30 MMAC per insertion at C=192, HW=256.
+    use_evss_bridge: bool = False
+    evss_bridge_locations: List[str] = field(
+        default_factory=lambda: ["pre_stage3"]    # subset of {pre_stage3, pre_stage4}
+    )
+
+    # -------- Neck (ultra-lite 1x1 projection) --------
+    neck_channels: int = 192    # C_n, output channels from neck
+
+    # -------- CVGL-Aware Head v7.1-α --------
+    d_descriptor: int = 512     # global descriptor dimensionality
+    use_asymmetric_heads: bool = True   # Sat vs UAV different heads
+    chp_rings: int = 8
+    chp_angles: int = 16
+    chp_harmonics: int = 4
+    use_film_altitude: bool = True
+    altitude_norm: float = 500.0    # divides altitude in meters
+    ring_count: int = 4             # LPN rings auxiliary
+    use_ring_aux: bool = True       # training-only ring aux branch
+
+    # -------- Text fusion (extension for caption-conditioned heads) --------
+    return_normalized: bool = True   # if False, heads return pre-L2 features (for late gated fusion)
+    use_text_film_sat: bool = False  # text-FiLM modulation in SatHead before GGeM
+    use_text_film_uav: bool = False  # text-FiLM modulation in UAVHead alongside altitude FiLM
+    text_film_dim: int = 1024        # text embedding dim feeding FiLM (matches TextFusionMLP out_dim)
+    text_film_hidden: int = 256
+
+    # -------- Weight-sharing --------
+    share_stages_1_2: bool = True   # sat ↔ UAV shared weights stages 1-2
+
+    # -------- KD taps (enable for future teacher KD) --------
+    enable_kd_taps: bool = True
+
+    # -------- Deployment hints --------
+    precision: Literal["fp32", "fp16", "int8_mixed"] = "fp16"
+
+    def validate(self) -> None:
+        assert len(self.embed_dims) == 4, (
+            f"embed_dims must have 4 entries, got {len(self.embed_dims)}"
+        )
+        assert len(self.depths) == 4, (
+            f"depths must have 4 entries, got {len(self.depths)}"
+        )
+        assert self.input_size % 32 == 0, (
+            f"input_size must be divisible by 32 (4 downsamples × stem), "
+            f"got {self.input_size}"
+        )
+
+    def summary(self) -> str:
+        return (
+            f"SOFIAConfig(stem={self.stem_mid}/{self.stem_out}, "
+            f"dims={self.embed_dims}, depths={self.depths}, "
+            f"neck={self.neck_channels}, d={self.d_descriptor}, "
+            f"precision={self.precision})"
+        )
+
+
+# ============================================================
+# Scale Presets
+# ============================================================
+
+def sofia_m_config() -> SOFIAConfig:
+    """SOFIA-M: ~500 M params target (~500 MB INT8, ~1 GB FP16).
+
+    Fits in 500 MB VRAM after INT8 quantization on Jetson Orin NX.
+    Expected latency: ~18 ms.
+    """
+    return SOFIAConfig(
+        stem_mid=64,
+        stem_out=128,
+        embed_dims=[256, 512, 1280, 1536],
+        depths=[3, 4, 15, 3],
+        neck_channels=192,
+        d_descriptor=512,
+    )
+
+
+def sofia_l_config() -> SOFIAConfig:
+    """SOFIA-L: ~1 B params target (~1 GB INT8, ~2 GB FP16).
+
+    Fits in 1 GB VRAM after INT8 quantization on Jetson Orin NX.
+    Expected latency: ~20 ms.
+    """
+    return SOFIAConfig(
+        stem_mid=64,
+        stem_out=128,
+        embed_dims=[256, 512, 1536, 2048],
+        depths=[3, 4, 20, 3],
+        neck_channels=256,
+        d_descriptor=1024,
+    )
+
+
+def sofia_tiny_config() -> SOFIAConfig:
+    """SOFIA-Tiny: ~5 M params (matches original v7.1 spec).
+
+    Reference for research comparisons. Not optimized for 500 MB INT8 target.
+    `num_heads_*` is set to 4 so `head_dim` (channels // heads) is divisible
+    by 4 — required by RoPE 2D in `MambaVisionBlock` (s3: 176/4=44, s4: 224/4=56).
+    `mamba_extra_kwargs.headdim=16` because Mamba-2 requires channels % headdim == 0;
+    176 and 224 are not divisible by the default 64.
+    """
+    return SOFIAConfig(
+        stem_mid=16,
+        stem_out=32,
+        embed_dims=[48, 96, 176, 224],
+        depths=[2, 3, 4, 2],
+        num_heads_s3=4,
+        num_heads_s4=4,
+        neck_channels=128,
+        d_descriptor=512,
+        mamba_extra_kwargs={
+            "d_state_mamba2": 64,
+            "headdim": 16,
+            "expand": 2,
+            "d_conv": 4,
+            "n_directions": 2,
+        },
+    )
+
+
+# Default preset
+DEFAULT_CONFIG = sofia_m_config
+
+
+if __name__ == "__main__":
+    for name, fn in [("M", sofia_m_config), ("L", sofia_l_config), ("Tiny", sofia_tiny_config)]:
+        cfg = fn()
+        cfg.validate()
+        print(f"SOFIA-{name}: {cfg.summary()}")

src/models/sofia_v71/layers.py

@@ -0,0 +1,398 @@
+"""SOFIA v7.1 custom layers.
+
+Includes:
+- GGeM: Generalized Mean Pooling with per-channel learnable exponent (F11)
+- CircularHarmonicPool: Formally SO(2)-invariant pooling via polar + FFT magnitude (NH2 novel)
+- AltitudeFiLM: FiLM conditioning on UAV altitude (NH4 novel)
+- RoPE2D: 2D Rotary Position Embedding for attention
+- SqueezeExcite: standard SE block
+- LayerNorm2d: channel-last LN wrapper for 2D features
+
+All rotation-invariance and FiLM modules are NOVEL contributions of SOFIA v7.1-α.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ============================================================
+# GGeM: Generalized Mean Pooling (F11)
+# ============================================================
+
+class GGeM(nn.Module):
+    """Per-channel learnable Generalized Mean pooling.
+
+    Formula:
+        GGeM(F)_c = (1/HW · Σ F_{c,h,w}^{p_c})^{1/p_c}
+        p_c = softplus(p_hat_c) ∈ (0, ∞)
+    """
+
+    def __init__(self, channels: int, init_p: float = 3.0, eps: float = 1e-6) -> None:
+        super().__init__()
+        # softplus^{-1}(init_p) = log(exp(init_p) - 1)
+        hat_init = math.log(math.exp(init_p) - 1.0)
+        self.hat_p = nn.Parameter(torch.full((channels,), hat_init))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: [B, C, H, W] -> [B, C]"""
+        p = F.softplus(self.hat_p).view(1, -1, 1, 1)    # [1, C, 1, 1]
+        x_clamped = x.clamp(min=self.eps)
+        x_pow = x_clamped.pow(p)
+        x_mean = x_pow.mean(dim=(2, 3), keepdim=True)
+        out = x_mean.pow(1.0 / p)
+        return out.flatten(1)
+
+
+# ============================================================
+# CHP: Circular Harmonic Pool (NH2 novel — formally SO(2)-invariant)
+# ============================================================
+
+class CircularHarmonicPool(nn.Module):
+    """Formally SO(2) rotation-invariant pooling.
+
+    Algorithm:
+        1. Sample input feature map at polar grid (r, θ) via bilinear grid_sample
+        2. Apply 1D real FFT along θ-axis
+        3. Keep magnitudes of first N harmonics (invariant to shift = rotation)
+        4. GGeM pool over rings r (per-channel-per-harmonic)
+        5. Flatten to descriptor [B, C * N]
+
+    Output is theoretically invariant to input rotation of any angle.
+
+    See HYP Phase 4'' Section 4''.1 NH2 and Section 4''.5 for formal proof.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        rings: int = 8,
+        angles: int = 16,
+        harmonics: int = 4,
+        r_min: float = 0.1,
+        r_max: float = 1.0,
+    ) -> None:
+        super().__init__()
+        assert harmonics <= angles // 2 + 1, (
+            f"harmonics {harmonics} cannot exceed angles//2+1 = {angles // 2 + 1}"
+        )
+
+        self.channels = channels
+        self.rings = rings
+        self.angles = angles
+        self.harmonics = harmonics
+
+        # GGeM over rings (per channel × per harmonic)
+        self.ggem = GGeM(channels * harmonics, init_p=3.0)
+
+        # Precompute polar grid in normalized [-1, 1] coords for grid_sample
+        grid = self._make_polar_grid(rings, angles, r_min, r_max)    # [R, T, 2]
+        self.register_buffer("polar_grid", grid, persistent=False)
+
+    @staticmethod
+    def _make_polar_grid(R: int, T: int, r_min: float, r_max: float) -> torch.Tensor:
+        r_values = torch.linspace(r_min, r_max, R)    # [R]
+        theta_values = torch.linspace(0.0, 2 * math.pi, T + 1)[:-1]    # [T]
+        r_grid, theta_grid = torch.meshgrid(r_values, theta_values, indexing="ij")    # [R, T]
+        x = r_grid * torch.cos(theta_grid)
+        y = r_grid * torch.sin(theta_grid)
+        grid = torch.stack([x, y], dim=-1)    # [R, T, 2]
+        return grid
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, C, H, W] feature map
+        Returns:
+            descriptor: [B, C * harmonics]
+        """
+        B, C, H, W = x.shape
+        assert C == self.channels, (
+            f"Input channels {C} != expected {self.channels}"
+        )
+
+        # 1. Polar sampling
+        # grid: [B, R, T, 2]
+        grid = self.polar_grid.unsqueeze(0).expand(B, -1, -1, -1)
+        polar = F.grid_sample(
+            x, grid,
+            mode="bilinear",
+            padding_mode="zeros",
+            align_corners=True,
+        )    # [B, C, R, T]
+
+        # 2. 1D real FFT along angular axis
+        polar_fft = torch.fft.rfft(polar, dim=-1)    # [B, C, R, T//2+1]
+        polar_fft = polar_fft[..., : self.harmonics]    # [B, C, R, N]
+
+        # 3. Magnitude (rotation invariant)
+        magnitude = polar_fft.abs()    # [B, C, R, N]
+
+        # 4. Reshape for GGeM: treat (C, N) as combined channel dim, rings as spatial
+        # Shape: [B, C*N, R, 1]
+        magnitude_reshaped = (
+            magnitude
+            .permute(0, 1, 3, 2)    # [B, C, N, R]
+            .reshape(B, C * self.harmonics, self.rings, 1)
+        )
+
+        # 5. GGeM pool over rings (H=R, W=1)
+        descriptor = self.ggem(magnitude_reshaped)    # [B, C*N]
+
+        return descriptor
+
+
+# ============================================================
+# FiLM: Altitude-conditioned modulation (NH4 novel)
+# ============================================================
+
+class AltitudeFiLM(nn.Module):
+    """FiLM modulation conditioned on scalar altitude.
+
+    F' = γ(h) · F + β(h)  where γ,β ∈ R^C are produced by MLP.
+
+    At altitude=None, produces identity (γ=1, β=0) via zero-init of final layer.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        hidden_dim: int = 64,
+        altitude_norm: float = 500.0,
+    ) -> None:
+        super().__init__()
+        self.channels = channels
+        self.altitude_norm = altitude_norm
+
+        self.mlp = nn.Sequential(
+            nn.Linear(1, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, 2 * channels),
+        )
+        # Zero-init last layer → initial γ=0 before residual, β=0
+        nn.init.zeros_(self.mlp[-1].weight)
+        nn.init.zeros_(self.mlp[-1].bias)
+
+    def forward(self, x: torch.Tensor, altitude: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            x: [B, C, H, W]
+            altitude: [B] or [B, 1] scalar altitude in meters (or None for neutral)
+        Returns:
+            [B, C, H, W]
+        """
+        B = x.shape[0]
+        if altitude is None:
+            altitude = torch.zeros(B, 1, device=x.device, dtype=x.dtype)
+        elif altitude.dim() == 1:
+            altitude = altitude.unsqueeze(-1)
+
+        h_norm = altitude.to(x.dtype) / self.altitude_norm
+        gamma_beta = self.mlp(h_norm)    # [B, 2C]
+        gamma, beta = gamma_beta.chunk(2, dim=-1)
+        # Residual form: γ = 1 + delta_γ (starts at identity)
+        gamma = gamma.view(B, self.channels, 1, 1) + 1.0
+        beta = beta.view(B, self.channels, 1, 1)
+        return gamma * x + beta
+
+
+# ============================================================
+# TextFiLM: text-conditioned modulation (extension for caption fusion)
+# ============================================================
+
+class TextFiLM(nn.Module):
+    """FiLM modulation conditioned on a text embedding.
+
+    F' = γ(z) · F + β(z) where γ,β ∈ R^C are produced by an MLP
+    from z ∈ R^{D_txt}. Identity at init via zero-init of last layer
+    (so γ=1, β=0 before training shifts the residual).
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        text_dim: int = 1024,
+        hidden_dim: int = 256,
+    ) -> None:
+        super().__init__()
+        self.channels = channels
+        self.text_dim = text_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(text_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, 2 * channels),
+        )
+        nn.init.zeros_(self.mlp[-1].weight)
+        nn.init.zeros_(self.mlp[-1].bias)
+
+    def forward(self, x: torch.Tensor, text_emb: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            x: [B, C, H, W]
+            text_emb: [B, D_txt] or None for no-op (identity)
+        Returns:
+            [B, C, H, W]
+        """
+        if text_emb is None:
+            return x
+        B = x.shape[0]
+        gamma_beta = self.mlp(text_emb.to(x.dtype))    # [B, 2C]
+        gamma, beta = gamma_beta.chunk(2, dim=-1)
+        gamma = gamma.view(B, self.channels, 1, 1) + 1.0
+        beta = beta.view(B, self.channels, 1, 1)
+        return gamma * x + beta
+
+
+# ============================================================
+# RoPE 2D
+# ============================================================
+
+class RoPE2D(nn.Module):
+    """2D Rotary Position Embedding.
+
+    Splits head_dim into two halves: first half gets x-position encoding,
+    second half gets y-position encoding. For each half, applies standard
+    1D RoPE rotation.
+
+    Reference: RoFormer (B49) adapted for 2D.
+    """
+
+    def __init__(self, head_dim: int, max_resolution: int = 64, base: float = 10000.0) -> None:
+        super().__init__()
+        assert head_dim % 2 == 0, "head_dim must be even for RoPE"
+        assert head_dim % 4 == 0, "head_dim must be divisible by 4 for 2D RoPE"
+        self.head_dim = head_dim
+        self.half_dim = head_dim // 2    # dedicated to each axis
+        self.max_resolution = max_resolution
+
+        # Frequencies for each axis (half_dim per axis, sin+cos pairs)
+        freqs = 1.0 / (base ** (torch.arange(0, self.half_dim, 2).float() / self.half_dim))
+        self.register_buffer("freqs", freqs, persistent=False)
+
+    def _make_embeds(self, H: int, W: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Produce cos/sin embeddings for HW tokens in raster order."""
+        y_pos = torch.arange(H, device=device, dtype=torch.float32)
+        x_pos = torch.arange(W, device=device, dtype=torch.float32)
+
+        freqs_y = torch.einsum("i,j->ij", y_pos, self.freqs)    # [H, half_dim/2]
+        freqs_x = torch.einsum("i,j->ij", x_pos, self.freqs)    # [W, half_dim/2]
+
+        # Expand to full grid: [H, W, half_dim/2] each
+        freqs_y = freqs_y.unsqueeze(1).expand(-1, W, -1)    # [H, W, half_dim/2]
+        freqs_x = freqs_x.unsqueeze(0).expand(H, -1, -1)    # [H, W, half_dim/2]
+
+        # Concatenate: x-axis freqs into first half, y-axis into second half
+        # Each half is [H, W, half_dim/2]; we pair-up for complex rotation
+        freqs_combined_x = torch.cat([freqs_x, freqs_x], dim=-1)    # [H, W, half_dim]
+        freqs_combined_y = torch.cat([freqs_y, freqs_y], dim=-1)    # [H, W, half_dim]
+
+        freqs_full = torch.cat([freqs_combined_x, freqs_combined_y], dim=-1)    # [H, W, head_dim]
+        cos = freqs_full.cos().reshape(H * W, -1)
+        sin = freqs_full.sin().reshape(H * W, -1)
+        return cos, sin
+
+    @staticmethod
+    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+        """Rotate: (x1, x2) -> (-x2, x1)."""
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat([-x2, x1], dim=-1)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, H: int, W: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            q, k: [B, heads, HW, head_dim]
+            H, W: spatial dims for position computation
+        Returns:
+            q_rot, k_rot with positional encoding applied
+        """
+        cos, sin = self._make_embeds(H, W, q.device)
+        cos = cos.to(q.dtype).unsqueeze(0).unsqueeze(0)    # [1, 1, HW, head_dim]
+        sin = sin.to(q.dtype).unsqueeze(0).unsqueeze(0)
+        q_rot = (q * cos) + (self._rotate_half(q) * sin)
+        k_rot = (k * cos) + (self._rotate_half(k) * sin)
+        return q_rot, k_rot
+
+
+# ============================================================
+# SE: Squeeze-Excite
+# ============================================================
+
+class SqueezeExcite(nn.Module):
+    """Standard Squeeze-Excite channel attention."""
+
+    def __init__(self, channels: int, reduction: int = 16) -> None:
+        super().__init__()
+        hidden = max(1, channels // reduction)
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Conv2d(channels, hidden, 1),
+            nn.SiLU(inplace=True),
+            nn.Conv2d(hidden, channels, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        s = self.pool(x)
+        s = self.fc(s)
+        return x * s
+
+
+# ============================================================
+# LayerNorm2d: LN over channels for (B, C, H, W) layout
+# ============================================================
+
+class LayerNorm2d(nn.Module):
+    """LayerNorm over C dimension for 4D tensors."""
+
+    def __init__(self, channels: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(channels, eps=eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # [B, C, H, W] → [B, H, W, C] → LN → [B, C, H, W]
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = x.permute(0, 3, 1, 2).contiguous()
+        return x
+
+
+if __name__ == "__main__":
+    # Smoke test
+    torch.manual_seed(0)
+
+    # GGeM test
+    g = GGeM(64)
+    x = torch.randn(2, 64, 8, 8)
+    out = g(x)
+    print(f"GGeM out: {out.shape}")    # [2, 64]
+
+    # CHP test — verify rotation invariance
+    chp = CircularHarmonicPool(32, rings=8, angles=16, harmonics=4)
+    x = torch.randn(1, 32, 16, 16)
+    out1 = chp(x)
+    # Rotate x by 90° and verify invariance (approximately)
+    x_rot = torch.rot90(x, k=1, dims=(-2, -1))
+    out2 = chp(x_rot)
+    diff = (out1 - out2).abs().max().item()
+    print(f"CHP rotation-invariance max diff: {diff:.4e} (should be small)")
+    print(f"CHP out shape: {out1.shape}")    # [1, 128]
+
+    # FiLM test
+    film = AltitudeFiLM(64)
+    x = torch.randn(2, 64, 8, 8)
+    altitudes = torch.tensor([150.0, 300.0])
+    out = film(x, altitudes)
+    print(f"FiLM out: {out.shape}")    # [2, 64, 8, 8]
+
+    # RoPE2D test
+    rope = RoPE2D(32)
+    q = torch.randn(2, 4, 64, 32)    # [B, heads, HW, head_dim]
+    k = torch.randn(2, 4, 64, 32)
+    q_r, k_r = rope(q, k, 8, 8)
+    print(f"RoPE out: q={q_r.shape}, k={k_r.shape}")

src/models/sofia_v71/model.py

@@ -0,0 +1,589 @@
+"""SOFIA v7.1 full model.
+
+Architecture:
+    stem → stage1 (StripDCN-lite) → stage2 (StripMixConv)
+         → stage3 (MambaVision MV5) → stage4 (MambaVision MV1)
+         → 1×1 neck projection → {Sat-Head, UAV-Head} + (training) Ring aux
+
+Features:
+- Siamese backbone with optional weight-sharing stages 1-2
+- Asymmetric Sat/UAV heads (CVGL-Aware Head v7.1-α)
+- Training-time Ring LPN aux for partial matching robustness
+- KD taps on F2, F3, F4 for future teacher-student distillation
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import SOFIAConfig
+from .blocks import (
+    Downsample,
+    EVSSBridge,
+    MambaVisionBlock,
+    StripDCNLiteBlock,
+    StripMixConvBlock,
+)
+from .layers import (
+    AltitudeFiLM,
+    CircularHarmonicPool,
+    GGeM,
+    LayerNorm2d,
+    TextFiLM,
+)
+
+
+# ============================================================
+# Stem
+# ============================================================
+
+class Stem(nn.Module):
+    """Dual-conv FastViT-style stem: 3 → mid (s=2) → out (s=1).
+
+    Downsampling ×2 total (input 256 → 128).
+    """
+
+    def __init__(self, in_channels: int, mid_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, mid_channels, 3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(mid_channels)
+        self.act1 = nn.SiLU(inplace=True)
+        self.conv2 = nn.Conv2d(mid_channels, out_channels, 3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.act2 = nn.SiLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.act1(self.bn1(self.conv1(x)))
+        x = self.act2(self.bn2(self.conv2(x)))
+        return x
+
+
+# ============================================================
+# Backbone (shared s1-2 optional, separate s3-4)
+# ============================================================
+
+class Backbone(nn.Module):
+    """SOFIA v7.1 backbone: 4 stages + stem.
+
+    Stage 1: StripDCN-lite blocks
+    Stage 2: StripMixConv blocks
+    Stage 3: MambaVision MV5 blocks (Mamba ∥ MHSA ∥ Strip)
+    Stage 4: MambaVision MV1 blocks (Mamba ∥ MHSA)
+    """
+
+    def __init__(self, cfg: SOFIAConfig) -> None:
+        super().__init__()
+        self.cfg = cfg
+        dims = cfg.embed_dims
+        depths = cfg.depths
+
+        # Stem: input → stem_out
+        self.stem = Stem(cfg.in_channels, cfg.stem_mid, cfg.stem_out)
+
+        # Stage 1: stem_out → dims[0], no downsample at block level
+        # (stem already did 4× total; stage 1 operates at 64×64 for 256 input)
+        self.ds1 = Downsample(cfg.stem_out, dims[0])
+        self.stage1 = nn.Sequential(*[
+            StripDCNLiteBlock(
+                in_channels=dims[0],
+                out_channels=dims[0],
+                expand=cfg.mbconv_expand,
+                kernel=cfg.strip_kernel_s1,
+                stride=1,
+                se_ratio=cfg.se_ratio,
+                use_dcn=cfg.use_dcn_strip,
+            )
+            for _ in range(depths[0])
+        ])
+
+        # Stage 2
+        self.ds2 = Downsample(dims[0], dims[1])
+        self.stage2 = nn.Sequential(*[
+            StripMixConvBlock(
+                in_channels=dims[1],
+                out_channels=dims[1],
+                expand=cfg.mbconv_expand,
+                strip_kernel=cfg.strip_kernel_s2,
+                mix_kernels=cfg.mix_kernels,
+                stride=1,
+                se_ratio=cfg.se_ratio,
+            )
+            for _ in range(depths[1])
+        ])
+
+        # Stage 3 — optional EVSS bridge after downsample (B6-inspired refinement
+        # to smooth DCN→MambaVision semantic gap)
+        self.ds3 = Downsample(dims[1], dims[2])
+        if cfg.use_evss_bridge and "pre_stage3" in cfg.evss_bridge_locations:
+            self.bridge3 = EVSSBridge(
+                channels=dims[2],
+                mamba_d_state=cfg.mamba_d_state,
+                mamba_dt_rank=cfg.mamba_dt_rank,
+                mamba_variant="mamba1",    # lightweight refinement, not primary mixer
+                mamba_backend=cfg.mamba_backend,
+                se_reduction=cfg.se_ratio,
+            )
+        else:
+            self.bridge3 = None
+        self.stage3 = nn.Sequential(*[
+            MambaVisionBlock(
+                channels=dims[2],
+                num_heads=cfg.num_heads_s3,
+                d_state=cfg.mamba_d_state,
+                dt_rank=cfg.mamba_dt_rank,
+                use_strip_branch=cfg.use_strip_branch_s3,
+                ffn_expand=cfg.ffn_expand,
+                strip_kernel=cfg.strip_kernel_s1,
+                mamba_backend=cfg.mamba_backend,
+                mamba_variant=cfg.mamba_variant,
+                mamba_extra_kwargs=cfg.mamba_extra_kwargs,
+            )
+            for _ in range(depths[2])
+        ])
+
+        # Stage 4 — optional EVSS bridge after downsample
+        self.ds4 = Downsample(dims[2], dims[3])
+        if cfg.use_evss_bridge and "pre_stage4" in cfg.evss_bridge_locations:
+            self.bridge4 = EVSSBridge(
+                channels=dims[3],
+                mamba_d_state=cfg.mamba_d_state,
+                mamba_dt_rank=cfg.mamba_dt_rank,
+                mamba_variant="mamba1",
+                mamba_backend=cfg.mamba_backend,
+                se_reduction=cfg.se_ratio,
+            )
+        else:
+            self.bridge4 = None
+        self.stage4 = nn.Sequential(*[
+            MambaVisionBlock(
+                channels=dims[3],
+                num_heads=cfg.num_heads_s4,
+                d_state=cfg.mamba_d_state,
+                dt_rank=cfg.mamba_dt_rank,
+                use_strip_branch=cfg.use_strip_branch_s4,
+                ffn_expand=cfg.ffn_expand,
+                strip_kernel=cfg.strip_kernel_s1,
+                mamba_backend=cfg.mamba_backend,
+                mamba_variant=cfg.mamba_variant,
+                mamba_extra_kwargs=cfg.mamba_extra_kwargs,
+            )
+            for _ in range(depths[3])
+        ])
+
+    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Returns dict with F1-F4 taps for KD/multi-scale use."""
+        s0 = self.stem(x)                  # stem_out, H/2
+        f1 = self.stage1(self.ds1(s0))     # dims[0], H/4
+        f2 = self.stage2(self.ds2(f1))     # dims[1], H/8
+
+        # Stage 3 with optional EVSS bridge.
+        ds3_out = self.ds3(f2)
+        if self.bridge3 is not None:
+            ds3_out = self.bridge3(ds3_out)
+        f3 = self.stage3(ds3_out)          # dims[2], H/16
+
+        # Stage 4 with optional EVSS bridge.
+        ds4_out = self.ds4(f3)
+        if self.bridge4 is not None:
+            ds4_out = self.bridge4(ds4_out)
+        f4 = self.stage4(ds4_out)          # dims[3], H/32 (= 8×8 for 256 input)
+
+        return {"s0": s0, "f1": f1, "f2": f2, "f3": f3, "f4": f4}
+
+
+# ============================================================
+# Neck: ultra-lite 1×1 projection
+# ============================================================
+
+class UltraLiteNeck(nn.Module):
+    """1×1 projection + BN + activation."""
+
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.SiLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj(x)
+
+
+# ============================================================
+# Heads
+# ============================================================
+
+class SatHead(nn.Module):
+    """Satellite view head: [TextFiLM] + GGeM + BN + Linear [+ L2].
+
+    `text_emb` is optional; when None or `text_film` is disabled, behaves
+    identically to the original head. `return_normalized=False` returns
+    pre-L2 descriptors for late gated fusion in a wrapper.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        d_descriptor: int,
+        return_normalized: bool = True,
+        use_text_film: bool = False,
+        text_film_dim: int = 1024,
+        text_film_hidden: int = 256,
+    ) -> None:
+        super().__init__()
+        self.return_normalized = return_normalized
+        self.use_text_film = use_text_film
+        if use_text_film:
+            self.text_film = TextFiLM(channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
+        self.ggem = GGeM(channels)
+        self.bn = nn.BatchNorm1d(channels, affine=False)
+        self.proj = nn.Linear(channels, d_descriptor)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        text_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """x: [B, C, H, W], text_emb: [B, D_txt] or None -> [B, d]"""
+        if self.use_text_film:
+            x = self.text_film(x, text_emb)
+        g = self.ggem(x)
+        g = self.bn(g)
+        g = self.proj(g)
+        if self.return_normalized:
+            g = F.normalize(g, p=2, dim=-1)
+        return g
+
+
+class UAVHead(nn.Module):
+    """UAV view head: FiLM(altitude) [+ TextFiLM] + CHP + BN + Linear [+ L2].
+
+    Formally SO(2)-invariant via CHP. Altitude-aware via FiLM. Optional
+    text-FiLM is applied AFTER altitude-FiLM (zero-init β so it starts as
+    identity). `return_normalized=False` returns pre-L2 descriptors for
+    late gated fusion in a wrapper.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        d_descriptor: int,
+        rings: int = 8,
+        angles: int = 16,
+        harmonics: int = 4,
+        use_film: bool = True,
+        altitude_norm: float = 500.0,
+        return_normalized: bool = True,
+        use_text_film: bool = False,
+        text_film_dim: int = 1024,
+        text_film_hidden: int = 256,
+    ) -> None:
+        super().__init__()
+        self.return_normalized = return_normalized
+        self.use_film = use_film
+        self.use_text_film = use_text_film
+        if use_film:
+            self.film = AltitudeFiLM(channels, altitude_norm=altitude_norm)
+        if use_text_film:
+            self.text_film = TextFiLM(channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
+
+        self.chp = CircularHarmonicPool(
+            channels=channels,
+            rings=rings,
+            angles=angles,
+            harmonics=harmonics,
+        )
+        chp_dim = channels * harmonics
+        self.bn = nn.BatchNorm1d(chp_dim, affine=False)
+        self.proj = nn.Linear(chp_dim, d_descriptor)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        altitude: Optional[torch.Tensor] = None,
+        text_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """x: [B, C, H, W], altitude: [B] or None, text_emb: [B, D_txt] or None -> [B, d]"""
+        if self.use_film:
+            x = self.film(x, altitude)
+        if self.use_text_film:
+            x = self.text_film(x, text_emb)
+        g = self.chp(x)
+        g = self.bn(g)
+        g = self.proj(g)
+        if self.return_normalized:
+            g = F.normalize(g, p=2, dim=-1)
+        return g
+
+
+class RingAuxHead(nn.Module):
+    """LPN Square-Ring pool + per-ring Linear + L2 (training-only auxiliary).
+
+    Used for partial-matching robustness. Drop at inference.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        rings: int = 4,
+        d_per_ring: int = 128,
+        feature_size: int = 8,
+    ) -> None:
+        super().__init__()
+        self.rings = rings
+        self.feature_size = feature_size
+        self.d_per_ring = d_per_ring
+
+        # Per-ring GGeM
+        self.ggems = nn.ModuleList([GGeM(channels) for _ in range(rings)])
+        self.projs = nn.ModuleList([
+            nn.Linear(channels, d_per_ring) for _ in range(rings)
+        ])
+
+        # Precompute ring masks
+        masks = self._make_ring_masks(rings, feature_size)    # [R, H, W]
+        self.register_buffer("ring_masks", masks, persistent=False)
+
+    @staticmethod
+    def _make_ring_masks(R: int, S: int) -> torch.Tensor:
+        """Concentric square rings on SxS feature map."""
+        masks = torch.zeros(R, S, S)
+        center = (S - 1) / 2.0
+        # Define ring by Chebyshev distance thresholds
+        for r in range(R):
+            r_min = r * (S / (2 * R))
+            r_max = (r + 1) * (S / (2 * R))
+            for i in range(S):
+                for j in range(S):
+                    dist = max(abs(i - center), abs(j - center))
+                    if r_min <= dist < r_max or (r == R - 1 and dist <= r_max):
+                        masks[r, i, j] = 1.0
+        # Normalize each ring to sum to 1 for pooling stability (not required but cleaner)
+        masks = masks / masks.sum(dim=(1, 2), keepdim=True).clamp(min=1.0)
+        return masks
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Args:
+            x: [B, C, H, W]
+        Returns:
+            List of R tensors each [B, d_per_ring], L2-normalized
+        """
+        B, C, H, W = x.shape
+        outs = []
+        for r in range(self.rings):
+            mask = self.ring_masks[r].to(x.dtype)    # [H, W]
+            # Apply mask
+            x_masked = x * mask.unsqueeze(0).unsqueeze(0)    # [B, C, H, W]
+            g = self.ggems[r](x_masked)    # [B, C]
+            g = self.projs[r](g)    # [B, d_per_ring]
+            g = F.normalize(g, p=2, dim=-1)
+            outs.append(g)
+        return outs
+
+
+# ============================================================
+# Full SOFIA v7.1 model
+# ============================================================
+
+class SOFIAv71(nn.Module):
+    """Full SOFIA v7.1 model.
+
+    Forward signature:
+        forward(sat: [B,3,H,W], uav: [B,3,H,W], altitude: [B] = None,
+                return_features: bool = False) -> dict
+    """
+
+    def __init__(self, cfg: SOFIAConfig) -> None:
+        super().__init__()
+        cfg.validate()
+        self.cfg = cfg
+
+        # Backbones (shared or separate depending on config)
+        self.backbone_shared = Backbone(cfg)
+        if cfg.share_stages_1_2:
+            # Single backbone, use one for both
+            self.backbone_sat = None
+            self.backbone_uav = None
+        else:
+            # Fully separate backbones (rare case)
+            self.backbone_sat = Backbone(cfg)
+            self.backbone_uav = Backbone(cfg)
+            # Clear shared
+            del self.backbone_shared
+            self.backbone_shared = None
+
+        # Neck (shared)
+        self.neck = UltraLiteNeck(cfg.embed_dims[-1], cfg.neck_channels)
+
+        # Heads
+        if cfg.use_asymmetric_heads:
+            self.sat_head = SatHead(
+                cfg.neck_channels, cfg.d_descriptor,
+                return_normalized=cfg.return_normalized,
+                use_text_film=cfg.use_text_film_sat,
+                text_film_dim=cfg.text_film_dim,
+                text_film_hidden=cfg.text_film_hidden,
+            )
+            self.uav_head = UAVHead(
+                channels=cfg.neck_channels,
+                d_descriptor=cfg.d_descriptor,
+                rings=cfg.chp_rings,
+                angles=cfg.chp_angles,
+                harmonics=cfg.chp_harmonics,
+                use_film=cfg.use_film_altitude,
+                altitude_norm=cfg.altitude_norm,
+                return_normalized=cfg.return_normalized,
+                use_text_film=cfg.use_text_film_uav,
+                text_film_dim=cfg.text_film_dim,
+                text_film_hidden=cfg.text_film_hidden,
+            )
+        else:
+            # Symmetric: use SatHead for both
+            self.sat_head = SatHead(
+                cfg.neck_channels, cfg.d_descriptor,
+                return_normalized=cfg.return_normalized,
+                use_text_film=cfg.use_text_film_sat,
+                text_film_dim=cfg.text_film_dim,
+                text_film_hidden=cfg.text_film_hidden,
+            )
+            self.uav_head = self.sat_head
+
+        # Ring aux (training-only)
+        # Resolution chain: input/2 (stem) → /2 (ds1) → /2 (ds2) → /2 (ds3) → /2 (ds4)
+        # Total ×32 downsampling → input 256 gives 8×8 final feature map
+        if cfg.use_ring_aux:
+            feat_size = cfg.input_size // 32
+            self.ring_head = RingAuxHead(
+                channels=cfg.neck_channels,
+                rings=cfg.ring_count,
+                d_per_ring=cfg.d_descriptor // cfg.ring_count,
+                feature_size=feat_size,
+            )
+        else:
+            self.ring_head = None
+
+    def _extract_features(self, x: torch.Tensor, view: str) -> Dict[str, torch.Tensor]:
+        """Run backbone for sat or uav view."""
+        if self.cfg.share_stages_1_2:
+            return self.backbone_shared(x)
+        else:
+            bb = self.backbone_sat if view == "sat" else self.backbone_uav
+            return bb(x)
+
+    def forward(
+        self,
+        sat: Optional[torch.Tensor] = None,
+        uav: Optional[torch.Tensor] = None,
+        altitude: Optional[torch.Tensor] = None,
+        text_emb_sat: Optional[torch.Tensor] = None,
+        text_emb_uav: Optional[torch.Tensor] = None,
+        return_features: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            sat: [B, 3, H, W] satellite image (or None to skip)
+            uav: [B, 3, H, W] UAV image (or None to skip)
+            altitude: [B] altitude in meters for UAV (or None = neutral)
+            text_emb_sat: [B, D_txt] caption embedding for SatHead text-FiLM (or None)
+            text_emb_uav: [B, D_txt] caption embedding for UAVHead text-FiLM (or None)
+            return_features: if True, also return backbone features F2-F4 for KD
+
+        Returns:
+            dict with g_sat, g_uav (global descriptors), optional features,
+            and optional rings (training only)
+        """
+        result: Dict[str, torch.Tensor] = {}
+
+        # SAT path
+        if sat is not None:
+            feats_sat = self._extract_features(sat, view="sat")
+            f4_sat = feats_sat["f4"]
+            neck_sat = self.neck(f4_sat)
+            g_sat = self.sat_head(neck_sat, text_emb=text_emb_sat)
+            result["g_sat"] = g_sat
+            if return_features:
+                result["features_sat"] = feats_sat
+                result["neck_sat"] = neck_sat
+            if self.training and self.ring_head is not None:
+                result["rings_sat"] = self.ring_head(neck_sat)
+
+        # UAV path
+        if uav is not None:
+            feats_uav = self._extract_features(uav, view="uav")
+            f4_uav = feats_uav["f4"]
+            neck_uav = self.neck(f4_uav)
+            if self.cfg.use_asymmetric_heads:
+                g_uav = self.uav_head(neck_uav, altitude=altitude, text_emb=text_emb_uav)
+            else:
+                g_uav = self.uav_head(neck_uav, text_emb=text_emb_uav)
+            result["g_uav"] = g_uav
+            if return_features:
+                result["features_uav"] = feats_uav
+                result["neck_uav"] = neck_uav
+            if self.training and self.ring_head is not None:
+                result["rings_uav"] = self.ring_head(neck_uav)
+
+        return result
+
+    def count_parameters(self, only_trainable: bool = True) -> Dict[str, int]:
+        """Count parameters per module."""
+        counts = {}
+        total = 0
+        for name, module in self.named_children():
+            if module is None:
+                continue
+            n = sum(p.numel() for p in module.parameters() if not only_trainable or p.requires_grad)
+            counts[name] = n
+            total += n
+        counts["_total"] = total
+        return counts
+
+
+# ============================================================
+# Build helper
+# ============================================================
+
+def build_sofia(preset: str = "M") -> SOFIAv71:
+    """Build SOFIA v7.1 model from preset name."""
+    from .config import sofia_m_config, sofia_l_config, sofia_tiny_config
+    preset_map = {
+        "M": sofia_m_config,
+        "L": sofia_l_config,
+        "Tiny": sofia_tiny_config,
+    }
+    if preset not in preset_map:
+        raise ValueError(f"Unknown preset '{preset}'. Available: {list(preset_map)}")
+    cfg = preset_map[preset]()
+    return SOFIAv71(cfg)
+
+
+if __name__ == "__main__":
+    # Smoke test with Tiny preset (fast)
+    print("Building SOFIA-Tiny for smoke test...")
+    model = build_sofia("Tiny")
+    model.eval()
+
+    counts = model.count_parameters()
+    total_m = counts["_total"] / 1e6
+    print(f"Total params: {total_m:.2f} M")
+    print("Per-module:")
+    for k, v in counts.items():
+        print(f"  {k}: {v / 1e6:.3f} M")
+
+    # Dry forward
+    sat = torch.randn(1, 3, 256, 256)
+    uav = torch.randn(1, 3, 256, 256)
+    alt = torch.tensor([150.0])
+    with torch.no_grad():
+        out = model(sat=sat, uav=uav, altitude=alt)
+    for k, v in out.items():
+        if isinstance(v, torch.Tensor):
+            print(f"  out[{k}]: {tuple(v.shape)}")
+        else:
+            print(f"  out[{k}]: {type(v).__name__}")

src/models/sofia_v71/quant.py

@@ -0,0 +1,565 @@
+"""SOFIA v7.1 quantization utilities.
+
+Two reusable building blocks for INT8 deployment:
+
+1. **OffsetClampSTE** (DCN-M2): hard clamp DCN offsets to physical range
+   ``[-k, +k]`` with straight-through gradient. Bounded distribution =
+   stable percentile clipping during PTQ + cleaner offset semantics.
+
+2. **KScaledFakeQuant** (R5-style): multi-bin fake quantization. Bins are
+   chosen by magnitude percentiles; each bin has its own scale. Adapts
+   quantization resolution to long-tail distributions (Mamba Δ, output y).
+
+Plus a drop-in:
+
+3. **KScaledMamba2Block**: ``Mamba2Block`` subclass with k-scaled fake-quant
+   nodes on ``x_main``, ``delta``, and per-token ``y`` outputs. Implements
+   R5 reparam principle (state ``h_t`` kept in model dtype; only observable
+   tensors quantized).
+
+All components are PTQ/QAT-compatible (STE backward, calibration mode).
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .blocks import Mamba2Block
+
+
+# ============================================================
+# DCN-M2: Offset clamp with STE
+# ============================================================
+
+class _OffsetClampSTE(torch.autograd.Function):
+    """Hard clamp [low, high] in forward, identity gradient in backward."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, low: float, high: float) -> torch.Tensor:
+        return torch.clamp(x, low, high)
+
+    @staticmethod
+    def backward(ctx, grad: torch.Tensor):
+        # STE: pass gradient unmodified
+        return grad, None, None
+
+
+def offset_clamp_ste(offsets: torch.Tensor, kernel_size: int) -> torch.Tensor:
+    """Functional DCN-M2: clamp offsets to ``[-k, +k]`` with STE backward.
+
+    Args:
+        offsets: arbitrary-shape tensor of DCN offsets.
+        kernel_size: physical receptive-field bound k. Output is clamped to
+            absolute values <= k. Standard choice equals the DW kernel size.
+
+    Returns:
+        Clamped tensor (same shape as ``offsets``). Backward propagates
+        gradient as if no clamp was applied, so training can still adjust
+        weights even when an offset hits the boundary.
+    """
+    k = float(kernel_size)
+    return _OffsetClampSTE.apply(offsets, -k, k)
+
+
+class OffsetClampSTE(nn.Module):
+    """Module wrapper for :func:`offset_clamp_ste`.
+
+    Use as a drop-in layer between ``offset_predictor`` and ``deform_conv2d``::
+
+        self.clamp = OffsetClampSTE(kernel_size=7)
+        offsets = self.clamp(self.offset_predictor(x))
+    """
+
+    def __init__(self, kernel_size: int) -> None:
+        super().__init__()
+        self.kernel_size = int(kernel_size)
+
+    def forward(self, offsets: torch.Tensor) -> torch.Tensor:
+        return offset_clamp_ste(offsets, self.kernel_size)
+
+    def extra_repr(self) -> str:
+        return f"kernel_size={self.kernel_size}, range=[-{self.kernel_size}, +{self.kernel_size}]"
+
+
+# ============================================================
+# K-Scaled Fake Quantization (R5 style, multi-bin)
+# ============================================================
+
+class _STEFakeQuant(torch.autograd.Function):
+    """Fake-quantize: ``round(x / s) * s`` clamped to [qmin, qmax], STE backward.
+
+    `scale` is a tensor broadcastable with `x` (per-tensor scalar or
+    per-element via gathered per-bin scales).
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        scale: torch.Tensor,
+        qmin: int,
+        qmax: int,
+    ) -> torch.Tensor:
+        x_int = torch.round(x / scale).clamp(qmin, qmax)
+        return x_int * scale
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        return grad_output, None, None, None
+
+
+class KScaledFakeQuant(nn.Module):
+    """k-scaled (multi-bin magnitude-bucketed) fake quantization.
+
+    Algorithm:
+        1. ``g(x) = bin index in [0..k-1] based on |x|`` using thresholds
+           ``t_0 < t_1 < ... < t_{k-2}``.
+        2. ``s(x) = scales[g(x)]`` per-element scale.
+        3. ``x_q = round(x / s) * s`` clamped to ``[qmin, qmax]``.
+        4. STE backward.
+
+    Setting ``num_bins=1`` reduces to standard single-scale fake quant.
+
+    Calibration workflow::
+
+        fq = KScaledFakeQuant(num_bins=3)
+        fq.start_calibration()
+        for batch in cal_loader:
+            _ = model(batch)        # collect stats
+        fq.finalize_calibration(percentile=99.9)
+        # fq is now active (calibrating=False) and applies fake-quant in forward
+    """
+
+    def __init__(
+        self,
+        num_bins: int = 3,
+        qmin: int = -128,
+        qmax: int = 127,
+        max_calibration_samples: int = 200_000,
+    ) -> None:
+        super().__init__()
+        if num_bins < 1:
+            raise ValueError(f"num_bins must be >= 1, got {num_bins}")
+        if qmax <= qmin:
+            raise ValueError(f"qmax ({qmax}) must be > qmin ({qmin})")
+
+        self.num_bins = num_bins
+        self.qmin = qmin
+        self.qmax = qmax
+        self.max_calibration_samples = max_calibration_samples
+
+        # Thresholds: (num_bins - 1,) — bin boundaries on |x|. For num_bins=1, empty.
+        n_thr = max(0, num_bins - 1)
+        self.register_buffer("thresholds", torch.zeros(n_thr))
+
+        # Scales: (num_bins,) — one scale per bin. Default 1.0 = no quant effect.
+        self.register_buffer("scales", torch.ones(num_bins))
+
+        # Per-instance state (not buffers — runtime-only)
+        self.calibrating: bool = False
+        self.enabled: bool = True
+        self._calib_buffer: List[torch.Tensor] = []
+        self._n_collected: int = 0
+
+    # -------------------- Calibration API --------------------
+
+    def start_calibration(self) -> None:
+        """Begin collecting input distribution stats."""
+        self.calibrating = True
+        self._calib_buffer = []
+        self._n_collected = 0
+
+    def stop_calibration(self) -> None:
+        """End calibration without finalizing scales (e.g. abort)."""
+        self.calibrating = False
+        self._calib_buffer = []
+        self._n_collected = 0
+
+    def finalize_calibration(self, percentile: float = 99.9) -> None:
+        """Compute thresholds and per-bin scales from collected stats.
+
+        Args:
+            percentile: per-bin tail percentile for scale computation
+                (typical 99.9). Robust to extreme outliers.
+        """
+        if not self._calib_buffer:
+            self.calibrating = False
+            return
+
+        all_data = torch.cat(self._calib_buffer)
+        abs_data = all_data.abs()
+
+        if self.num_bins == 1:
+            # Single scale: percentile of |x|
+            tail = torch.quantile(abs_data, percentile / 100.0).item()
+            self.scales[0] = max(tail / float(self.qmax), 1e-8)
+        else:
+            # Threshold positions: equal-mass quantile splits
+            qpts = torch.linspace(0.0, 1.0, self.num_bins + 1, device=abs_data.device)[1:-1]
+            thr_vals = torch.quantile(abs_data, qpts)
+            self.thresholds.copy_(thr_vals.to(self.thresholds.device))
+
+            # Per-bin scale: per-bin tail percentile / qmax
+            for g in range(self.num_bins):
+                if g == 0:
+                    lo = 0.0
+                    hi = thr_vals[0].item()
+                elif g == self.num_bins - 1:
+                    lo = thr_vals[-1].item()
+                    hi = float("inf")
+                else:
+                    lo = thr_vals[g - 1].item()
+                    hi = thr_vals[g].item()
+
+                mask = (abs_data >= lo) & (abs_data < hi)
+                if mask.any():
+                    bin_tail = torch.quantile(
+                        abs_data[mask], percentile / 100.0
+                    ).item()
+                    self.scales[g] = max(bin_tail / float(self.qmax), 1e-8)
+                else:
+                    self.scales[g] = 1.0
+
+        self._calib_buffer = []
+        self._n_collected = 0
+        self.calibrating = False
+
+    # -------------------- Forward --------------------
+
+    def _bin_index(self, x: torch.Tensor) -> torch.Tensor:
+        """Per-element bin index (long tensor, shape == x.shape)."""
+        if self.num_bins == 1:
+            return torch.zeros_like(x, dtype=torch.long)
+
+        abs_x = x.abs()
+        idx = torch.zeros_like(x, dtype=torch.long)
+        for i in range(self.num_bins - 1):
+            t = self.thresholds[i]
+            idx = torch.where(abs_x >= t, idx + 1, idx)
+        return idx
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.enabled:
+            return x
+
+        if self.calibrating:
+            # Subsample to bound memory
+            with torch.no_grad():
+                flat = x.detach().flatten()
+                budget = self.max_calibration_samples - self._n_collected
+                if budget > 0:
+                    take = min(flat.numel(), budget)
+                    if take < flat.numel():
+                        # Random subsample for representativeness
+                        idx = torch.randperm(flat.numel(), device=flat.device)[:take]
+                        sample = flat[idx]
+                    else:
+                        sample = flat
+                    self._calib_buffer.append(sample.cpu())
+                    self._n_collected += sample.numel()
+            return x
+
+        # Apply k-scaled fake-quant
+        bin_idx = self._bin_index(x)
+        scale = self.scales[bin_idx]    # shape == x.shape
+        return _STEFakeQuant.apply(x, scale, self.qmin, self.qmax)
+
+    def extra_repr(self) -> str:
+        return (
+            f"num_bins={self.num_bins}, qmin={self.qmin}, qmax={self.qmax}, "
+            f"calibrating={self.calibrating}, enabled={self.enabled}"
+        )
+
+
+# ============================================================
+# K-Scaled Mamba-2 drop-in
+# ============================================================
+
+class KScaledMamba2Block(Mamba2Block):
+    """Drop-in replacement for :class:`Mamba2Block` with k-scaled fake-quant.
+
+    Adds ``KScaledFakeQuant`` nodes on three observable tensors inside the
+    Mamba-2 forward pass:
+
+    - ``x_main``: scan input after conv1d + SiLU
+    - ``delta``: time-step ``Δ`` after softplus (addresses MF2 long tail)
+    - ``y``:      per-token scan output (addresses MF4 state propagation)
+
+    The recurrent state ``h_t`` is **not** quantized — kept in model dtype
+    inside the scan loop. This implements R5's reparam principle: only
+    observable I/O is quantized; internal state stays high-precision.
+
+    Backend constraint: forces ``backend='torch'``. The mamba_ssm CUDA
+    backend would require k-scaled support inside the custom kernel (out of
+    scope here). For a fast deploy path, train+QAT in torch, then export to
+    a custom TRT plugin that does the same INT8 scan.
+
+    Args:
+        channels: model dim ``C``.
+        d_state: state size ``N`` (Mamba-2 default 64).
+        headdim: head dim (must divide channels).
+        d_conv: local conv kernel.
+        expand: inner expansion factor.
+        backend: must be ``'torch'`` (or ``'auto'``, which resolves to torch).
+        num_bins: ``k`` for k-scaled (typical 2-4).
+        targets: which tensors to quantize. Subset of
+            ``('x_main', 'delta', 'y')``.
+
+    Example::
+
+        mamba = KScaledMamba2Block(channels=192, num_bins=3,
+                                   targets=('delta', 'y'))
+        mamba.start_calibration()
+        for batch in cal_loader:
+            _ = model_with_mamba(batch)
+        mamba.finalize_calibration()
+        # Now in PTQ mode — forward applies fake-quant on chosen targets.
+    """
+
+    SUPPORTED_TARGETS: Tuple[str, ...] = ("x_main", "delta", "y")
+
+    def __init__(
+        self,
+        channels: int,
+        d_state: int = 64,
+        headdim: int = 64,
+        d_conv: int = 4,
+        expand: int = 2,
+        backend: str = "torch",
+        num_bins: int = 3,
+        targets: Sequence[str] = ("x_main", "delta", "y"),
+    ) -> None:
+        # Force torch — k-scaled requires explicit forward-pass control
+        if backend == "auto":
+            backend = "torch"
+        if backend != "torch":
+            raise ValueError(
+                "KScaledMamba2Block currently supports only backend='torch'. "
+                f"Got backend='{backend}'. The mamba_ssm CUDA kernel does not "
+                "expose hooks for k-scaled fake-quant; integrating INT8 scan "
+                "via TRT plugin is a separate deploy-time concern."
+            )
+
+        super().__init__(
+            channels=channels,
+            d_state=d_state,
+            headdim=headdim,
+            d_conv=d_conv,
+            expand=expand,
+            backend="torch",
+        )
+
+        unknown = set(targets) - set(self.SUPPORTED_TARGETS)
+        if unknown:
+            raise ValueError(
+                f"Unknown k-scaled targets: {unknown}. "
+                f"Supported: {self.SUPPORTED_TARGETS}"
+            )
+        self.targets = set(targets)
+        self.num_bins = num_bins
+
+        # Build fake-quant nodes only for active targets
+        if "x_main" in self.targets:
+            self.fq_x_main = KScaledFakeQuant(num_bins=num_bins)
+        if "delta" in self.targets:
+            self.fq_delta = KScaledFakeQuant(num_bins=num_bins)
+        if "y" in self.targets:
+            self.fq_y = KScaledFakeQuant(num_bins=num_bins)
+
+    # -------------------- Calibration helpers --------------------
+
+    def _all_fqs(self) -> List[KScaledFakeQuant]:
+        out: List[KScaledFakeQuant] = []
+        for name in self.SUPPORTED_TARGETS:
+            attr = f"fq_{name}"
+            mod = getattr(self, attr, None)
+            if isinstance(mod, KScaledFakeQuant):
+                out.append(mod)
+        return out
+
+    def start_calibration(self) -> None:
+        for fq in self._all_fqs():
+            fq.start_calibration()
+
+    def finalize_calibration(self, percentile: float = 99.9) -> None:
+        for fq in self._all_fqs():
+            fq.finalize_calibration(percentile=percentile)
+
+    def set_quant_enabled(self, enabled: bool) -> None:
+        for fq in self._all_fqs():
+            fq.enabled = enabled
+
+    # -------------------- Forward --------------------
+
+    def _forward_torch(self, x: torch.Tensor) -> torch.Tensor:
+        """Mamba-2 torch forward with k-scaled fake-quant on selected paths."""
+        B, L, C = x.shape
+
+        # In projection (matches parent layout)
+        z_in = self.in_proj(x)
+        xz, BC, dt = torch.split(
+            z_in,
+            [2 * self.d_inner, 2 * self.d_state, self.nheads],
+            dim=-1,
+        )
+        x_main, z_gate = xz.chunk(2, dim=-1)
+        B_p, C_p = BC.chunk(2, dim=-1)
+
+        # Local conv + SiLU
+        x_main_t = x_main.transpose(1, 2)
+        x_main_t = self.conv1d(x_main_t)[..., :L]
+        x_main = F.silu(x_main_t).transpose(1, 2)
+
+        # K-SCALED QUANT: x_main (scan input)
+        if "x_main" in self.targets:
+            x_main = self.fq_x_main(x_main)
+
+        # Δ
+        dt = F.softplus(dt)
+        if "delta" in self.targets:
+            dt = self.fq_delta(dt)
+
+        A = -torch.exp(self.A_log)    # [nheads]
+        x_head = x_main.view(B, L, self.nheads, self.headdim)
+
+        # Sequential scan. R5 reparam: state h kept in model dtype throughout.
+        h = torch.zeros(
+            B, self.nheads, self.headdim, self.d_state,
+            device=x.device, dtype=x.dtype,
+        )
+        ys: List[torch.Tensor] = []
+        for t in range(L):
+            dt_t = dt[:, t]
+            B_t = B_p[:, t]
+            C_t = C_p[:, t]
+            x_t = x_head[:, t]
+
+            dA = torch.exp(dt_t * A.unsqueeze(0)).unsqueeze(-1).unsqueeze(-1)
+            dB = (dt_t.unsqueeze(-1) * B_t.unsqueeze(1)).unsqueeze(2)
+
+            h = dA * h + dB * x_t.unsqueeze(-1)
+            y_t = (h * C_t.unsqueeze(1).unsqueeze(1)).sum(dim=-1)
+            y_t = y_t + self.D.unsqueeze(0).unsqueeze(-1) * x_t
+
+            # K-SCALED QUANT: per-token y (after scan, before gate)
+            if "y" in self.targets:
+                y_t = self.fq_y(y_t)
+
+            ys.append(y_t)
+
+        y = torch.stack(ys, dim=1)
+        y = y.reshape(B, L, self.d_inner)
+        y = y * F.silu(z_gate)
+        y = self.out_proj(y)
+        return y
+
+
+# ============================================================
+# Model-wide calibration helpers
+# ============================================================
+
+def start_calibration(model: nn.Module) -> None:
+    """Walk model, start calibration on every quant module."""
+    seen = set()
+    # First, KScaledMamba2Block instances drive their internal FQs.
+    for m in model.modules():
+        if isinstance(m, KScaledMamba2Block):
+            m.start_calibration()
+            for fq in m._all_fqs():
+                seen.add(id(fq))
+    # Then standalone KScaledFakeQuant not owned by a KScaledMamba2Block.
+    for m in model.modules():
+        if isinstance(m, KScaledFakeQuant) and id(m) not in seen:
+            m.start_calibration()
+
+
+def finalize_calibration(model: nn.Module, percentile: float = 99.9) -> None:
+    """Walk model, finalize all quant modules with the same percentile."""
+    seen = set()
+    for m in model.modules():
+        if isinstance(m, KScaledMamba2Block):
+            m.finalize_calibration(percentile=percentile)
+            for fq in m._all_fqs():
+                seen.add(id(fq))
+    for m in model.modules():
+        if isinstance(m, KScaledFakeQuant) and id(m) not in seen:
+            m.finalize_calibration(percentile=percentile)
+
+
+def set_quant_enabled(model: nn.Module, enabled: bool) -> None:
+    """Toggle every quant module on/off. Useful for FP-vs-INT8 comparison."""
+    for m in model.modules():
+        if isinstance(m, KScaledFakeQuant):
+            m.enabled = enabled
+        elif isinstance(m, KScaledMamba2Block):
+            m.set_quant_enabled(enabled)
+
+
+# ============================================================
+# Smoke tests
+# ============================================================
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+
+    print("=== DCN-M2 OffsetClampSTE ===")
+    offsets = torch.randn(2, 14, 16, 16, requires_grad=True) * 5.0    # outliers ~ ±15
+    print(f"  in:  range [{offsets.min().item():.2f}, {offsets.max().item():.2f}]")
+    clamped = offset_clamp_ste(offsets, kernel_size=3)
+    print(f"  out: range [{clamped.min().item():.2f}, {clamped.max().item():.2f}]")
+    clamped.sum().backward()
+    grad_mean = offsets.grad.mean().item()
+    grad_std = offsets.grad.std().item()
+    print(f"  grad mean={grad_mean:.4f} std={grad_std:.4f} "
+          f"(STE → exactly 1.0 for every element)")
+
+    print("\n=== KScaledFakeQuant (num_bins=3) ===")
+    fq = KScaledFakeQuant(num_bins=3)
+    fq.start_calibration()
+    with torch.no_grad():
+        for _ in range(20):
+            x = torch.randn(64, 32) * 1.5
+            x[0, 0] = 25.0      # outlier
+            x[3, 7] = -40.0
+            _ = fq(x)
+    fq.finalize_calibration(percentile=99.9)
+    print(f"  thresholds: {fq.thresholds.tolist()}")
+    print(f"  scales:     {[round(s, 5) for s in fq.scales.tolist()]}")
+
+    x_test = torch.randn(8, 32) * 3
+    x_test[0, 0] = 30.0
+    x_q = fq(x_test)
+    err = (x_q - x_test).abs().mean().item()
+    rel_err = err / x_test.abs().mean().item()
+    print(f"  avg abs err: {err:.4e}  (relative: {rel_err:.2%})")
+
+    print("\n=== KScaledMamba2Block ===")
+    block = KScaledMamba2Block(
+        channels=128, d_state=32, headdim=32, num_bins=3,
+        targets=("delta", "y"),
+    )
+    n_params = sum(p.numel() for p in block.parameters()) / 1e3
+    print(f"  params: {n_params:.1f} K")
+
+    # Calibration on dummy data
+    block.start_calibration()
+    with torch.no_grad():
+        for _ in range(3):
+            _ = block(torch.randn(1, 64, 128))
+    block.finalize_calibration(percentile=99.9)
+
+    # FP vs k-scaled-INT8 comparison
+    x = torch.randn(2, 64, 128)
+    block.set_quant_enabled(False)
+    y_fp = block(x)
+    block.set_quant_enabled(True)
+    y_q = block(x)
+    diff = (y_fp - y_q).abs().mean().item()
+    rel = diff / y_fp.abs().mean().item()
+    print(f"  FP output  range [{y_fp.min().item():.3f}, {y_fp.max().item():.3f}]")
+    print(f"  INT8 output range [{y_q.min().item():.3f}, {y_q.max().item():.3f}]")
+    print(f"  abs diff: {diff:.4e}  (relative: {rel:.2%})")

src/models/sofia_v71/verify.py

@@ -0,0 +1,158 @@
+"""Verify SOFIA v7.1 model scales: param count, memory footprint, forward pass.
+
+Run:
+    python -m sofia_v71.verify
+    python -m sofia_v71.verify --preset L
+    python -m sofia_v71.verify --preset M --benchmark
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from typing import Optional
+
+import torch
+
+from .blocks import is_mamba_ssm_available, is_mamba2_available
+from .config import sofia_m_config, sofia_l_config, sofia_tiny_config
+from .model import SOFIAv71
+
+
+def count_parameters(model: torch.nn.Module) -> dict:
+    """Count parameters per named child."""
+    counts = {}
+    for name, param in model.named_parameters():
+        top = name.split(".")[0]
+        counts[top] = counts.get(top, 0) + param.numel()
+    counts["_total"] = sum(counts.values())
+    return counts
+
+
+def estimate_quantized_size(n_params: int, precision: str = "int8") -> float:
+    """Estimate on-disk / VRAM weight storage in MB."""
+    bytes_per_param = {"fp32": 4, "fp16": 2, "int8": 1, "int4": 0.5}[precision]
+    return n_params * bytes_per_param / (1024 ** 2)
+
+
+def test_forward(model: SOFIAv71, device: str = "cpu", batch_size: int = 1) -> dict:
+    """Dry forward pass, return output shapes."""
+    model = model.to(device).eval()
+    cfg = model.cfg
+    sat = torch.randn(batch_size, 3, cfg.input_size, cfg.input_size, device=device)
+    uav = torch.randn(batch_size, 3, cfg.input_size, cfg.input_size, device=device)
+    altitude = torch.rand(batch_size, device=device) * 300.0 + 50.0
+
+    shapes = {}
+    with torch.no_grad():
+        out = model(sat=sat, uav=uav, altitude=altitude)
+    for k, v in out.items():
+        if isinstance(v, torch.Tensor):
+            shapes[k] = tuple(v.shape)
+        elif isinstance(v, list):
+            shapes[k] = f"list[{len(v)} × {tuple(v[0].shape)}]"
+        elif isinstance(v, dict):
+            for kk, vv in v.items():
+                if isinstance(vv, torch.Tensor):
+                    shapes[f"{k}.{kk}"] = tuple(vv.shape)
+    return shapes
+
+
+def benchmark(model: SOFIAv71, device: str = "cpu", n_warmup: int = 3, n_runs: int = 20) -> dict:
+    """Measure forward pass latency."""
+    model = model.to(device).eval()
+    cfg = model.cfg
+    sat = torch.randn(1, 3, cfg.input_size, cfg.input_size, device=device)
+    uav = torch.randn(1, 3, cfg.input_size, cfg.input_size, device=device)
+    altitude = torch.tensor([150.0], device=device)
+
+    # Warmup
+    with torch.no_grad():
+        for _ in range(n_warmup):
+            _ = model(sat=sat, uav=uav, altitude=altitude)
+        if device.startswith("cuda"):
+            torch.cuda.synchronize()
+
+        # Measured runs
+        times = []
+        for _ in range(n_runs):
+            if device.startswith("cuda"):
+                torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            _ = model(sat=sat, uav=uav, altitude=altitude)
+            if device.startswith("cuda"):
+                torch.cuda.synchronize()
+            times.append(time.perf_counter() - t0)
+
+    times_ms = [t * 1000 for t in times]
+    return {
+        "mean_ms": sum(times_ms) / len(times_ms),
+        "min_ms": min(times_ms),
+        "max_ms": max(times_ms),
+        "std_ms": (sum((t - sum(times_ms) / len(times_ms)) ** 2 for t in times_ms) / len(times_ms)) ** 0.5,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Verify SOFIA v7.1 model")
+    parser.add_argument("--preset", choices=["Tiny", "M", "L"], default="Tiny",
+                        help="Model preset (default: Tiny for fast smoke test)")
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--benchmark", action="store_true")
+    parser.add_argument("--batch-size", type=int, default=1)
+    args = parser.parse_args()
+
+    cfg_fn = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config}[args.preset]
+    cfg = cfg_fn()
+
+    # Report mamba_ssm availability
+    print(f"\n=== Environment ===")
+    print(f"  mamba_ssm v1 (Mamba-1 selective scan): "
+          f"{'available' if is_mamba_ssm_available() else 'NOT available'}")
+    print(f"  mamba_ssm Mamba-2 (SSD dual form):     "
+          f"{'available' if is_mamba2_available() else 'NOT available'}")
+    print(f"  configured variant: {cfg.mamba_variant}, backend: {cfg.mamba_backend}")
+
+    print(f"\n=== SOFIA-{args.preset} configuration ===")
+    print(cfg.summary())
+
+    print("\n=== Building model ===")
+    model = SOFIAv71(cfg)
+    counts = count_parameters(model)
+    total = counts["_total"]
+
+    print(f"\nTotal params: {total / 1e6:.2f} M")
+    print("\nPer-module breakdown:")
+    for k, v in sorted(counts.items(), key=lambda x: -x[1]):
+        if k == "_total":
+            continue
+        pct = 100 * v / total
+        print(f"  {k:30s} {v / 1e6:>8.3f} M  ({pct:5.1f}%)")
+
+    print("\n=== Memory footprint estimates ===")
+    for prec in ["fp32", "fp16", "int8", "int4"]:
+        mb = estimate_quantized_size(total, prec)
+        print(f"  {prec.upper():6s} weights: {mb:>8.1f} MB  ({mb / 1024:.2f} GB)")
+
+    print("\n=== Forward pass test ===")
+    try:
+        shapes = test_forward(model, device=args.device, batch_size=args.batch_size)
+        for k, v in shapes.items():
+            print(f"  out[{k}]: {v}")
+    except Exception as e:
+        print(f"  ERROR during forward: {type(e).__name__}: {e}")
+        import traceback
+        traceback.print_exc()
+
+    if args.benchmark:
+        print("\n=== Latency benchmark ===")
+        try:
+            stats = benchmark(model, device=args.device)
+            print(f"  mean: {stats['mean_ms']:.2f} ms  min: {stats['min_ms']:.2f}  "
+                  f"max: {stats['max_ms']:.2f}  std: {stats['std_ms']:.2f}")
+        except Exception as e:
+            print(f"  ERROR during benchmark: {type(e).__name__}: {e}")
+
+
+if __name__ == "__main__":
+    main()

src/training/train_gtauav.py

@@ -54,6 +54,14 @@ from src.models.asymmetric_encoder import (
     get_drone_train_transform,
     get_satellite_train_transform,
 )
+from src.models.sofia_fusion_encoder import SOFIAFusionEncoder
+from src.models.sofia_v1 import SOFIAv1Config
+from src.models.sofia_v1_fusion_encoder import SOFIAv1FusionEncoder
+from src.models.sofia_v71 import (
+    sofia_l_config,
+    sofia_m_config,
+    sofia_tiny_config,
+)
 
 LOGGER = logging.getLogger("caption_test.train_gtauav")
 
@@ -91,11 +99,30 @@ class TrainConfigGTAUAV:
     mona_last_n_blocks: int = 12  # inject adapters only in last 12 of 24 ViT blocks
     gradient_checkpointing: bool = True  # trade compute for VRAM (allows larger batch)
     # StripNet backbone option (replaces DINOv3 when backbone="stripnet").
-    backbone: str = "dinov3"  # "dinov3" or "stripnet"
+    backbone: str = "dinov3"  # "dinov3", "stripnet", or "sofia"
     stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth"
     stripnet_mona_last_n_stages: int = 2  # Conv-MONA in last N of 4 StripNet stages (0 = disable MONA)
     stripnet_freeze: bool = True  # If False, StripNet backbone is fully trainable (full fine-tune)
     stripnet_backbone_lr_factor: float = 0.1  # Backbone LR = learning_rate * factor (only when unfrozen)
+    # SOFIA backbone options (used when backbone="sofia"). Trained from scratch — no pretrained checkpoint.
+    sofia_preset: str = "Tiny"  # "Tiny" | "M" | "L"
+    sofia_d_descriptor: int = 1024  # retrieval space (1024 = match TextFusionMLP out_dim)
+    sofia_use_text_film_uav: bool = True  # mid-level text-FiLM in UAV head
+    sofia_use_text_film_sat: bool = True  # mid-level text-FiLM in SAT head
+    sofia_lora_rank: int = 4
+    sofia_mamba_variant: str = "mamba2"  # "mamba1" | "mamba2" | "efficient_vmamba"
+    sofia_mamba_backend: str = "auto"  # "auto" | "torch" | "mamba_ssm"
+    # EVSSBridge (B6-inspired refinement between heterogeneous stages, opt-in).
+    sofia_use_evss_bridge: bool = False
+    sofia_evss_bridge_locations: list[str] = field(default_factory=lambda: ["pre_stage3"])
+    # SOFIA v1 backbone options (used when backbone="sofia_v1"). StripNet+DCN, from scratch.
+    sofia_v1_variant: str = "small"  # "tiny_tiny" | "tiny" | "small" | "small_v2"
+    sofia_v1_dcn_variant: str = "v2"  # "v2" (torchvision DeformConv2d, stable) | "v4" (OpenGVLab, leaky)
+    sofia_v1_d_descriptor: int = 1024
+    sofia_v1_use_text_film_uav: bool = True
+    sofia_v1_use_text_film_sat: bool = True
+    sofia_v1_use_film_altitude: bool = True
+    sofia_v1_lora_rank: int = 4
 
     # Training.
     resume_from: str | None = None  # path to checkpoint for resuming
@@ -167,7 +194,7 @@ def _atomic_save(obj: dict, path: Path) -> None:
 
 
 def _build_param_groups(
-    model: AsymmetricEncoder,
+    model: nn.Module,
     lr: float,
     text_lr_factor: float,
     stripnet_backbone_lr_factor: float = 0.1,
@@ -177,7 +204,8 @@ def _build_param_groups(
     Groups:
       - text_encoder.*       → lr * text_lr_factor (default 1e-5)
       - image_encoder.backbone.* (when StripNet unfrozen) → lr * stripnet_backbone_lr_factor (default 1e-5)
-      - everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv) → lr
+      - everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv,
+                         SOFIA backbone+heads when backbone="sofia") → lr
     """
     text_params = []
     backbone_params = []
@@ -249,9 +277,13 @@ def _embed_drone_queries(
     embs: list[torch.Tensor] = []
     for batch in tqdm(loader, desc="  dss-embed-queries", unit="batch", leave=False):
         drone_img = batch["drone_img"].to(device, non_blocking=True)
+        altitude = batch.get("altitude")
+        if altitude is not None:
+            altitude = altitude.to(device, non_blocking=True)
         q = model.encode_query(
             drone_img,
             batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
+            altitude=altitude,
         )
         embs.append(q.cpu())
 
@@ -336,9 +368,13 @@ def _evaluate(
         if max_batches is not None and i >= max_batches:
             break
         drone_img = batch["drone_img"].to(device, non_blocking=True)
+        altitude = batch.get("altitude")
+        if altitude is not None:
+            altitude = altitude.to(device, non_blocking=True)
         q = model.encode_query(
             drone_img,
             batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
+            altitude=altitude,
         )
         query_embs.append(q.cpu())
         query_valid_names.extend(batch["valid_sat_names"])
@@ -575,6 +611,19 @@ def train(cfg: TrainConfigGTAUAV) -> None:
 
     if cfg.resume_from is not None:
         LOGGER.info("Resuming from %s", cfg.resume_from)
+        if cfg.backbone == "sofia":
+            model, resume_ckpt = SOFIAFusionEncoder.load_checkpoint(
+                cfg.resume_from,
+                lrsclip_path=cfg.lrsclip_path,
+                device=cfg.device,
+            )
+        elif cfg.backbone == "sofia_v1":
+            model, resume_ckpt = SOFIAv1FusionEncoder.load_checkpoint(
+                cfg.resume_from,
+                lrsclip_path=cfg.lrsclip_path,
+                device=cfg.device,
+            )
+        else:
             model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
                 cfg.resume_from,
                 dino_web_path=cfg.dino_web_path,
@@ -585,11 +634,55 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         start_epoch = resume_ckpt.get("epoch", -1) + 1
     else:
         mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)"
-        if cfg.backbone == "stripnet":
+        if cfg.backbone == "sofia":
+            enc_str = f"SOFIA-{cfg.sofia_preset} (text-FiLM uav={cfg.sofia_use_text_film_uav}, sat={cfg.sofia_use_text_film_sat})"
+        elif cfg.backbone == "sofia_v1":
+            enc_str = f"SOFIAv1-{cfg.sofia_v1_variant} (StripNet+DCNv4, text-FiLM uav={cfg.sofia_v1_use_text_film_uav}, sat={cfg.sofia_v1_use_text_film_sat})"
+        elif cfg.backbone == "stripnet":
             enc_str = "StripNet-small (shared, 512→1024 proj)"
         else:
             enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)"
         LOGGER.info("Building model — %s, %s", mode_str, enc_str)
+        if cfg.backbone == "sofia":
+            preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config}
+            if cfg.sofia_preset not in preset_map:
+                raise ValueError(f"Unknown sofia_preset={cfg.sofia_preset!r}")
+            sofia_cfg = preset_map[cfg.sofia_preset]()
+            sofia_cfg.d_descriptor = cfg.sofia_d_descriptor
+            sofia_cfg.text_film_dim = cfg.sofia_d_descriptor
+            sofia_cfg.use_text_film_uav = cfg.sofia_use_text_film_uav and not cfg.baseline_mode
+            sofia_cfg.use_text_film_sat = cfg.sofia_use_text_film_sat and not cfg.baseline_mode
+            sofia_cfg.mamba_variant = cfg.sofia_mamba_variant
+            sofia_cfg.mamba_backend = cfg.sofia_mamba_backend
+            sofia_cfg.use_evss_bridge = cfg.sofia_use_evss_bridge
+            sofia_cfg.evss_bridge_locations = list(cfg.sofia_evss_bridge_locations)
+            model = SOFIAFusionEncoder(
+                sofia_cfg=sofia_cfg,
+                lrsclip_path=cfg.lrsclip_path,
+                init_gate=cfg.init_gate,
+                baseline_mode=cfg.baseline_mode,
+                lora_rank=cfg.sofia_lora_rank,
+                device=cfg.device,
+            ).to(cfg.device)
+        elif cfg.backbone == "sofia_v1":
+            sofia_v1_cfg = SOFIAv1Config(
+                variant=cfg.sofia_v1_variant,
+                dcn_variant=cfg.sofia_v1_dcn_variant,
+                d_descriptor=cfg.sofia_v1_d_descriptor,
+                text_film_dim=cfg.sofia_v1_d_descriptor,
+                use_text_film_uav=cfg.sofia_v1_use_text_film_uav and not cfg.baseline_mode,
+                use_text_film_sat=cfg.sofia_v1_use_text_film_sat and not cfg.baseline_mode,
+                use_film_altitude=cfg.sofia_v1_use_film_altitude,
+            )
+            model = SOFIAv1FusionEncoder(
+                sofia_cfg=sofia_v1_cfg,
+                lrsclip_path=cfg.lrsclip_path,
+                init_gate=cfg.init_gate,
+                baseline_mode=cfg.baseline_mode,
+                lora_rank=cfg.sofia_v1_lora_rank,
+                device=cfg.device,
+            ).to(cfg.device)
+        else:
             model = AsymmetricEncoder(
                 dino_web_path=cfg.dino_web_path,
                 dino_sat_path=cfg.dino_sat_path,
@@ -608,7 +701,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         LOGGER.info("embed_dim=%d", model.embed_dim)
 
     # --- Gradient checkpointing (trade compute for VRAM) ---
-    # StripNet doesn't expose set_gradient_checkpointing — skip silently.
+    # StripNet/SOFIA don't expose set_gradient_checkpointing — only DGTRS gets it.
     if cfg.gradient_checkpointing and cfg.backbone == "dinov3":
         if cfg.shared_encoder:
             model.image_encoder.set_gradient_checkpointing(True)
@@ -618,10 +711,10 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         if model.text_encoder is not None:
             model.text_encoder.transformer.gradient_checkpointing = True
         LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)")
-    elif cfg.gradient_checkpointing and cfg.backbone == "stripnet":
+    elif cfg.gradient_checkpointing and cfg.backbone in ("stripnet", "sofia", "sofia_v1"):
         if model.text_encoder is not None:
             model.text_encoder.transformer.gradient_checkpointing = True
-        LOGGER.info("Gradient checkpointing enabled (DGTRS only; StripNet doesn't support)")
+        LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", cfg.backbone)
 
     n_trainable = sum(p.numel() for p in model.trainable_parameters())
     n_total = sum(p.numel() for p in model.parameters())
@@ -879,11 +972,14 @@ def train(cfg: TrainConfigGTAUAV) -> None:
 
             drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
             sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
+            altitude = batch.get("altitude")
+            if altitude is not None:
+                altitude = altitude.to(cfg.device, non_blocking=True)
 
             # Model forward in AMP (fp16 for DINOv3/DGTRS encoders).
             with autocast(device_type="cuda", enabled=cfg.use_amp):
                 if cfg.baseline_mode:
-                    embeddings = model(drone_img=drone_img, sat_img=sat_img)
+                    embeddings = model(drone_img=drone_img, sat_img=sat_img, altitude=altitude)
                 else:
                     embeddings = model(
                         drone_img=drone_img,
@@ -894,6 +990,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
                         sat_caption_l1=batch["sat_caption_l1"],
                         sat_caption_l2=batch["sat_caption_l2"],
                         sat_caption_l3=batch["sat_caption_l3"],
+                        altitude=altitude,
                     )
             # Loss — InfoNCE or WeightedInfoNCE. Only the latter uses positive_weights.
             queue_neg = neg_bank.get_queue() if neg_bank is not None else None
@@ -1103,20 +1200,23 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         history.append(epoch_record)
 
         # Save checkpoint. Model architecture flags go into the ckpt so
-        # `AsymmetricEncoder.load_checkpoint` can rebuild the right shape.
+        # `AsymmetricEncoder.load_checkpoint` (or `SOFIAFusionEncoder.load_checkpoint`)
-        _atomic_save(
+        # can rebuild the right shape.
-            obj={
+        ckpt_obj = {
             "epoch": epoch,
             "model_state": model.state_dict(),
             "optimizer_state": optimizer.state_dict(),
             "loss_state": loss_fn.state_dict(),
             "baseline_mode": cfg.baseline_mode,
-                "shared_encoder": cfg.shared_encoder,
+            "backbone": cfg.backbone,
-                "mona_bottleneck": cfg.mona_bottleneck,
+        }
-                "mona_last_n_blocks": cfg.mona_last_n_blocks,
+        if cfg.backbone in ("sofia", "sofia_v1"):
-            },
+            ckpt_obj["sofia_cfg"] = model.sofia_cfg
-            path=output_dir / f"ckpt_epoch{epoch:03d}.pt",
+        else:
-        )
+            ckpt_obj["shared_encoder"] = cfg.shared_encoder
+            ckpt_obj["mona_bottleneck"] = cfg.mona_bottleneck
+            ckpt_obj["mona_last_n_blocks"] = cfg.mona_last_n_blocks
+        _atomic_save(obj=ckpt_obj, path=output_dir / f"ckpt_epoch{epoch:03d}.pt")
         LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch)
 
     # Save history.