| """WiSER CIR end-to-end CIR set-prediction model. |
| |
| Pipeline (AC-4 + AC-5 split): |
| |
| voxel_level --[scene_encode]--> scene_tokens (no TX, once per scene) |
| scene_tokens --[modulate_scene_tx(tx_emb)]--> tx_mem (per (scene, TX)) |
| tx_mem + rx --[head]--> (exists, delay_ns, peak_db) per query |
| |
| `SparseCsiDetrModel.encode_tx(...)` remains a thin helper for Fourier + MLP. |
| `SparseCsiDetrModel.scene_encode(...)` is the TX-independent backbone |
| forward used by the cache wrapper; `modulate_scene_tx(...)` is the |
| TX-conditioned pass. The trainer's `CachedSceneTxForward` wraps these so |
| `C_enc ≤ U` counts real scene-encoder forwards and `C_txmod ≤ P` counts |
| real TX-modulation forwards. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
|
|
| import torch |
| from torch import nn |
|
|
| from ..config import ModelConfig as SharedModelConfig |
| from .backbone import BackboneConfig, build_backbone |
| from .detr_head import CIRPathSetDETRHead, DetrHeadConfig, ContinuousFourierEmbed |
|
|
|
|
| @dataclass(slots=True) |
| class ModelConfig: |
| """Runtime model config. Derive from `config.ModelConfig` via `.from_shared()`.""" |
|
|
| backbone: BackboneConfig = field(default_factory=BackboneConfig) |
| head: DetrHeadConfig = field(default_factory=DetrHeadConfig) |
|
|
| @classmethod |
| def from_shared(cls, shared: SharedModelConfig) -> "ModelConfig": |
| return cls( |
| backbone=BackboneConfig( |
| kind=shared.backbone_kind, |
| channels=shared.backbone_channels, |
| tx_emb_channels=shared.backbone_tx_emb_channels, |
| num_downsample_stages=shared.backbone_downsample_stages, |
| blocks_per_stage=shared.backbone_blocks_per_stage, |
| num_heads=shared.backbone_num_heads, |
| mlp_ratio=shared.backbone_mlp_ratio, |
| ), |
| head=DetrHeadConfig( |
| channels=shared.backbone_channels, |
| num_queries=shared.query_budget, |
| num_decoder_layers=shared.head_num_decoder_layers, |
| num_heads=shared.head_num_heads, |
| dropout=shared.head_dropout, |
| db_low=shared.db_low, |
| db_high=shared.db_high, |
| peak_db_head_arch=getattr(shared, "peak_db_head_arch", "flat"), |
| peak_db_hidden=int(getattr(shared, "peak_db_hidden", 512)), |
| peak_db_init_mean_db=float(getattr(shared, "peak_db_init_mean_db", -55.0)), |
| delay_head_arch=getattr(shared, "delay_head_arch", "flat"), |
| delay_hidden=int(getattr(shared, "delay_hidden", 512)), |
| delay_init_mean_ns=float(getattr(shared, "delay_init_mean_ns", 4.0)), |
| exists_head_arch=getattr(shared, "exists_head_arch", "flat"), |
| exists_hidden=int(getattr(shared, "exists_hidden", 512)), |
| exists_init_prob=float(getattr(shared, "exists_init_prob", 0.2)), |
| ), |
| ) |
|
|
|
|
| class SparseCsiDetrModel(nn.Module): |
| """The importable model class referenced by AC-1's smoke test. |
| |
| Exposes `scene_encode` and `modulate_scene_tx` as two separable stages so |
| the cache wrapper can enforce AC-5 invariants. The backbone itself |
| implements both: for the Round-2 `dense_fallback` kind, `scene_encode` |
| runs every sparse block WITHOUT AdaLN-zero modulation (zero condition |
| vector), and `modulate_scene_tx` adds the per-TX modulation pass. The |
| `trellis2` kind delegates to the vendored `ModulatedSparseTransformerBlock`. |
| """ |
|
|
| def __init__(self, config: ModelConfig | None = None) -> None: |
| super().__init__() |
| self.config = config or ModelConfig() |
| self.backbone = build_backbone(self.config.backbone) |
| self.tx_embed = ContinuousFourierEmbed(in_dim=3, num_bands=8) |
| self.tx_proj = nn.Sequential( |
| nn.Linear(self.tx_embed.out_dim, self.config.backbone.tx_emb_channels), |
| nn.GELU(), |
| nn.Linear(self.config.backbone.tx_emb_channels, self.config.backbone.tx_emb_channels), |
| ) |
| self.head = CIRPathSetDETRHead(self.config.head) |
|
|
| @classmethod |
| def from_shared(cls, shared: SharedModelConfig) -> "SparseCsiDetrModel": |
| return cls(ModelConfig.from_shared(shared)) |
|
|
| def encode_tx(self, tx_xyz_norm: torch.Tensor) -> torch.Tensor: |
| return self.tx_proj(self.tx_embed(tx_xyz_norm)) |
|
|
| def scene_encode(self, voxel_level: dict) -> dict: |
| """TX-independent backbone pass; produces reusable scene tokens. |
| |
| Delegates to `backbone.scene_encode_forward` which runs the TX-FREE |
| block stack (`DensePlainBlock` × S for dense_fallback, or plain |
| `SparseTransformerBlock` × S for trellis2) + the downsample stack. |
| The returned dict is safe to cache per unique scene in the batch. |
| """ |
| return self.backbone.scene_encode_forward(voxel_level) |
|
|
| def modulate_scene_tx(self, scene_tokens: dict, tx_emb: torch.Tensor) -> dict: |
| """Per-(scene, TX) modulation pass on the cached scene tokens. |
| |
| Delegates to `backbone.modulation_forward`: |
| * dense_fallback → `DenseAdaLNBlock × blocks_per_stage` with TX as |
| AdaLN-zero condition (`shift, scale, gate` all derived from TX). |
| * trellis2 → `ModulatedSparseTransformerBlock × blocks_per_stage` |
| with TX projected through `tx_mod_proj` into the vendored |
| AdaLN-zero path. |
| |
| This is the ONLY code path where TX enters the scene tensor; RX never |
| enters here. |
| """ |
| if tx_emb.dim() == 1: |
| tx_emb = tx_emb.unsqueeze(0) |
| return self.backbone.modulation_forward(scene_tokens, tx_emb) |
|
|
| def modulate_scene(self, voxel_level: dict, tx_emb: torch.Tensor) -> dict: |
| """Legacy one-shot: scene_encode + modulate_scene_tx in one call.""" |
| scene = self.scene_encode(voxel_level) |
| return self.modulate_scene_tx(scene, tx_emb) |
|
|
| def head_forward( |
| self, |
| tx_mem: torch.Tensor, |
| rx_xyz_norm: torch.Tensor, |
| tx_mem_key_padding_mask: torch.Tensor | None = None, |
| ) -> dict[str, torch.Tensor]: |
| return self.head(tx_mem, rx_xyz_norm, tx_mem_key_padding_mask) |
|
|
|
|
| __all__ = ["ModelConfig", "SparseCsiDetrModel"] |
|
|
