Deploy learned baseline + hybrid multi-view tracking ensemble

7df6a88 6 months ago

19.7 kB

"""
Perceiver-based transformer for 3D roof wireframe prediction.
Architecture overview:
    Input tokens [B, T, D]
        |
        v
    input_proj: Linear -> GELU -> Linear -> LayerNorm   =>  [B, T, hidden]
        |
        v
    Perceiver latent bottleneck (N PerceiverLatentLayers):
        Learnable latent embeddings [L, hidden] are broadcast to batch.
        Each layer: cross-attn(latents <- tokens) -> self-attn(latents) -> FFN
        Output: latents [B, L, hidden]
        |
        v
    Segment decoder (M SegmentDecoderLayers):
        Learnable query embeddings [S, hidden] are broadcast to batch.
        Each layer: cross-attn(queries <- latents) -> self-attn(queries) -> FFN
        Output: queries [B, S, hidden]
        |
        v
    segment_head: Linear -> 6D -> (midpoint, half_vector)
        + query_offsets (learnable per-query bias)
        endpoints = midpoint +/- half_vector  ->  [B, S, 2, 3]
"""

import torch
import torch.nn as nn

from .attention import MultiHeadSDPA, FeedForward


# ---------------------------------------------------------------------------
# Building blocks
# ---------------------------------------------------------------------------

class AttnResidual(nn.Module):
    """Pre-norm attention + residual + dropout."""

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        dropout: float = 0.0,
        kv_heads: int | None = None,
        norm_class=None,
        qk_norm: bool = False,
        qk_norm_type: str = "l2",
    ):
        super().__init__()
        norm_class = norm_class or nn.LayerNorm
        self.norm = norm_class(d_model)
        self.attn = MultiHeadSDPA(d_model, num_heads, kv_heads=kv_heads, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
        self.drop = nn.Dropout(dropout)

    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        memory_key_padding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        res = x
        x = self.norm(x)
        x = self.attn(x, memory, key_padding_mask=memory_key_padding_mask)
        return res + self.drop(x)


class FFNResidual(nn.Module):
    """Pre-norm feed-forward + residual + dropout."""

    def __init__(
        self,
        d_model: int,
        dim_ff: int,
        dropout: float = 0.0,
        activation: str = "gelu",
        norm_class=None,
    ):
        super().__init__()
        norm_class = norm_class or nn.LayerNorm
        self.norm = norm_class(d_model)
        self.ffn = FeedForward(d_model, dim_ff, activation=activation)
        self.drop = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        res = x
        x = self.norm(x)
        x = self.ffn(x)
        return res + self.drop(x)


# ---------------------------------------------------------------------------
# Perceiver encoder layer
# ---------------------------------------------------------------------------

class PerceiverLatentLayer(nn.Module):
    """Single Perceiver latent layer.
    If use_cross=True:  cross-attn(latents <- points) -> self-attn -> FFN
    If use_cross=False: self-attn -> FFN  (saves compute in deep stacks)
    """

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        dim_ff: int,
        dropout: float = 0.0,
        activation: str = "gelu",
        kv_heads_cross: int | None = None,
        kv_heads_self: int | None = None,
        use_cross: bool = True,
        norm_class=None,
        qk_norm: bool = False,
        qk_norm_type: str = "l2",
    ):
        super().__init__()
        self.use_cross = use_cross
        if use_cross:
            self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)

    def forward(
        self,
        latents: torch.Tensor,
        points: torch.Tensor,
        points_key_padding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if self.use_cross:
            latents = self.cross(latents, points, memory_key_padding_mask=points_key_padding_mask)
        latents = self.self_attn(latents, latents)
        latents = self.ffn(latents)
        return latents


# ---------------------------------------------------------------------------
# Segment decoder layer
# ---------------------------------------------------------------------------

class SegmentDecoderLayer(nn.Module):
    """Single segment decoder layer.
    cross-attn(queries <- latents) -> [cross-attn(queries <- inputs)] -> self-attn(queries) -> FFN
    If input_xattn=True, adds a second cross-attention that attends directly
    to the projected input tokens (bypassing the latent bottleneck). This gives
    queries access to fine-grained point-level detail for vertex precision.
    """

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        dim_ff: int,
        dropout: float = 0.0,
        activation: str = "gelu",
        kv_heads_cross: int | None = None,
        kv_heads_self: int | None = None,
        norm_class=None,
        input_xattn: bool = False,
        qk_norm: bool = False,
        qk_norm_type: str = "l2",
    ):
        super().__init__()
        self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
        self.input_xattn = input_xattn
        if input_xattn:
            self.cross_input = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)

    def forward(
        self,
        queries: torch.Tensor,
        latents: torch.Tensor,
        src: torch.Tensor | None = None,
        src_key_padding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        queries = self.cross(queries, latents)
        if self.input_xattn and src is not None:
            queries = self.cross_input(queries, src, memory_key_padding_mask=src_key_padding_mask)
        queries = self.self_attn(queries, queries)
        queries = self.ffn(queries)
        return queries


# ---------------------------------------------------------------------------
# Full model
# ---------------------------------------------------------------------------

class TokenTransformerSegments(nn.Module):
    """Perceiver transformer that predicts 3D roof wireframe segments.
    Takes point-cloud tokens and outputs segment endpoints as [B, S, 2, 3]
    where S is the number of segments and each segment has two 3D endpoints.
    Args:
        segments:       Number of predicted segments (S).
        in_dim:         Dimensionality of input tokens.
        hidden:         Internal hidden dimension throughout the model.
        num_heads:      Number of attention heads.
        kv_heads_cross: Grouped-query heads for cross-attention (None = standard MHA).
        kv_heads_self:  Grouped-query heads for self-attention (None = standard MHA).
        dim_feedforward: FFN intermediate dimension.
        dropout:        Dropout rate applied after attention and FFN.
        latent_tokens:  Number of learnable latent embeddings (L) in the bottleneck.
        latent_layers:  Number of PerceiverLatentLayers (N).
        decoder_layers: Number of SegmentDecoderLayers (M).
    """

    def __init__(
        self,
        segments: int = 32,
        in_dim: int = 128,
        hidden: int = 128,
        num_heads: int = 4,
        kv_heads_cross: int | None = 2,
        kv_heads_self: int | None = 0,
        dim_feedforward: int = 256,
        dropout: float = 0.01,
        latent_tokens: int = 64,
        latent_layers: int = 2,
        decoder_layers: int = 2,
        cross_attn_interval: int = 1,
        norm_class=None,
        activation: str = "gelu",
        segment_conf: bool = False,
        pre_encoder_layers: int = 0,
        segment_param: str = "midpoint_halfvec",
        length_floor: float = 0.0,
        decoder_input_xattn: bool = False,
        qk_norm: bool = False,
        qk_norm_type: str = "l2",
    ):
        super().__init__()
        self.segments = segments
        self.out_vertices = segments * 2
        self.segment_param = segment_param
        self.decoder_input_xattn = decoder_input_xattn
        norm_class = norm_class or nn.LayerNorm

        # Treat 0 as "use standard MHA"
        if kv_heads_cross is not None and kv_heads_cross <= 0:
            kv_heads_cross = None
        if kv_heads_self is not None and kv_heads_self <= 0:
            kv_heads_self = None

        # -- Input projection --
        self.input_proj = nn.Sequential(
            nn.Linear(in_dim, dim_feedforward),
            nn.GELU(),
            nn.Linear(dim_feedforward, hidden),
            norm_class(hidden),
        )

        # -- Optional pre-encoder: self-attention on full token sequence --
        if pre_encoder_layers > 0:
            self.pre_encoder = nn.ModuleList([
                SelfAttentionEncoderLayer(
                    d_model=hidden,
                    num_heads=num_heads,
                    dim_ff=dim_feedforward,
                    dropout=dropout,
                    activation=activation,
                    kv_heads=kv_heads_self,
                    norm_class=norm_class,
                    qk_norm=qk_norm, qk_norm_type=qk_norm_type,
                )
                for _ in range(pre_encoder_layers)
            ])
        else:
            self.pre_encoder = None

        # -- Perceiver latent bottleneck --
        self.latent_embed = nn.Embedding(latent_tokens, hidden)
        N = latent_layers
        self.latent_layers = nn.ModuleList([
            PerceiverLatentLayer(
                d_model=hidden,
                num_heads=num_heads,
                dim_ff=dim_feedforward,
                dropout=dropout,
                activation=activation,
                kv_heads_cross=kv_heads_cross,
                kv_heads_self=kv_heads_self,
                use_cross=(i == 0) or (i == N - 1) or (i % cross_attn_interval == 0),
                norm_class=norm_class,
                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
            )
            for i in range(N)
        ])

        # -- Segment decoder --
        self.query_embed = nn.Embedding(segments, hidden)
        self.decoder_layers = nn.ModuleList([
            SegmentDecoderLayer(
                d_model=hidden,
                num_heads=num_heads,
                dim_ff=dim_feedforward,
                dropout=dropout,
                activation=activation,
                kv_heads_cross=kv_heads_cross,
                kv_heads_self=kv_heads_self,
                norm_class=norm_class,
                input_xattn=decoder_input_xattn,
                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
            )
            for _ in range(decoder_layers)
        ])

        # -- Output head --
        if segment_param == "midpoint_dir_len":
            self.segment_head = nn.Linear(hidden, 7)  # mid(3) + dir(3) + len(1)
        else:
            self.segment_head = nn.Linear(hidden, 6)  # mid(3) + half(3)
        self.query_offsets = nn.Parameter(torch.zeros(segments, 2, 3))

        nn.init.trunc_normal_(self.segment_head.weight, mean=0.0, std=1e-3)
        if self.segment_head.bias is not None:
            nn.init.zeros_(self.segment_head.bias)
            if segment_param == "midpoint_dir_len":
                # softplus(0.5) * 0.1 ≈ 0.097 default length in normalized space
                self.segment_head.bias.data[6] = 0.5
        nn.init.normal_(self.query_offsets, mean=0.0, std=0.05)

        # -- Optional confidence head --
        self.segment_conf = segment_conf
        if segment_conf:
            self.conf_head = nn.Linear(hidden, 1)
            nn.init.zeros_(self.conf_head.bias)

    def forward(
        self,
        tokens: torch.Tensor,
        mask: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor | list]:
        """
        Args:
            tokens: Input point-cloud tokens [B, T, in_dim].
            mask:   Boolean validity mask [B, T]. True = valid token.
        Returns:
            Dict with keys:
                "vertices": [B, S*2, 3] flattened endpoints.
                "segments": [B, S, 2, 3] segment endpoints.
                "edges":    Per-batch list of (start, end) index pairs into vertices.
                "conf":     [B, S] logits (only if segment_conf=True).
        """
        B = tokens.shape[0]

        # Project input tokens
        src = self.input_proj(tokens)  # [B, T, hidden]

        # Padding mask (True where padded) for cross-attention
        pad_mask = ~mask.bool() if mask is not None else None

        # Optional pre-encoder: self-attention on full token sequence
        if self.pre_encoder is not None:
            for layer in self.pre_encoder:
                src = layer(src, key_padding_mask=pad_mask)

        # Perceiver latent bottleneck
        latents = self.latent_embed.weight.unsqueeze(0).expand(B, -1, -1)
        for layer in self.latent_layers:
            latents = layer(latents, src, points_key_padding_mask=pad_mask)

        # Segment decoder
        queries = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1)
        for layer in self.decoder_layers:
            queries = layer(queries, latents,
                            src=src if self.decoder_input_xattn else None,
                            src_key_padding_mask=pad_mask if self.decoder_input_xattn else None)

        # Predict segments -> endpoints
        if self.segment_param == "midpoint_dir_len":
            raw = self.segment_head(queries)  # [B, S, 7]
            mid = raw[:, :, :3] + self.query_offsets[:, 0, :].unsqueeze(0)
            direction = torch.nn.functional.normalize(raw[:, :, 3:6], dim=-1)
            length = torch.nn.functional.softplus(raw[:, :, 6:7]) * 0.1
            half = direction * length * 0.5
        else:
            raw = self.segment_head(queries).view(B, self.segments, 2, 3)
            raw = raw + self.query_offsets.unsqueeze(0)
            mid, half = raw[:, :, 0], raw[:, :, 1]
        seg_params = torch.stack([mid - half, mid + half], dim=2)

        vertices = seg_params.reshape(B, self.out_vertices, 3)
        edges = [[(2 * i, 2 * i + 1) for i in range(self.segments)] for _ in range(B)]

        out = {"vertices": vertices, "segments": seg_params, "edges": edges,
               "src": src, "pad_mask": pad_mask, "queries": queries}
        if self.segment_conf:
            out["conf"] = self.conf_head(queries).squeeze(-1)  # [B, S]
        return out


# ---------------------------------------------------------------------------
# Encoder-only layer (self-attention on full token sequence)
# ---------------------------------------------------------------------------

class SelfAttentionEncoderLayer(nn.Module):
    """Single self-attention layer: self-attn(tokens) -> FFN."""

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        dim_ff: int,
        dropout: float = 0.0,
        activation: str = "gelu",
        kv_heads: int | None = None,
        norm_class=None,
        qk_norm: bool = False,
        qk_norm_type: str = "l2",
    ):
        super().__init__()
        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)

    def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor | None = None) -> torch.Tensor:
        x = self.self_attn(x, x, memory_key_padding_mask=key_padding_mask)
        x = self.ffn(x)
        return x


# ---------------------------------------------------------------------------
# End-to-end model: tokenizer embeddings + perceiver
# ---------------------------------------------------------------------------

class EdgeDepthSegmentsModel(nn.Module):
    """Tokenizer embeddings + transformer for 3D roof wireframes.
    Supports two architectures via the `arch` parameter:
    - "perceiver": Perceiver latent bottleneck (default, O(L*T) attention)
    - "transformer": Standard self-attention encoder (O(T^2) attention)
    Both share the same decoder, output head, and tokenizer.
    """

    def __init__(
        self,
        seq_cfg,
        segments: int = 32,
        hidden: int = 128,
        num_heads: int = 4,
        kv_heads_cross: int | None = 2,
        kv_heads_self: int | None = 0,
        dim_feedforward: int = 256,
        dropout: float = 0.1,
        latent_tokens: int = 64,
        latent_layers: int = 1,
        decoder_layers: int = 2,
        label_emb_dim: int = 16,
        src_emb_dim: int = 2,
        behind_emb_dim: int = 8,
        fourier_seed: int = 0,
        cross_attn_interval: int = 1,
        norm_class=None,
        activation: str = "gelu",
        segment_conf: bool = False,
        use_vote_features: bool = False,
        arch: str = "perceiver",
        encoder_layers: int = 4,
        pre_encoder_layers: int = 0,
        segment_param: str = "midpoint_halfvec",
        length_floor: float = 0.0,
        decoder_input_xattn: bool = False,
        qk_norm: bool = False,
        qk_norm_type: str = "l2",
        learnable_fourier: bool = False,
    ):
        super().__init__()
        self.seq_cfg = seq_cfg

        from .tokenizer import EdgeDepthSequenceBuilder
        self.tokenizer = EdgeDepthSequenceBuilder(
            seq_cfg,
            label_emb_dim=label_emb_dim,
            src_emb_dim=src_emb_dim,
            behind_emb_dim=behind_emb_dim,
            fourier_seed=fourier_seed,
            use_vote_features=use_vote_features,
            learnable_fourier=learnable_fourier,
        )

        if arch == "transformer":
            raise ValueError(
                "arch='transformer' is no longer supported. "
                "TransformerSegments has been removed; use arch='perceiver'.")
        else:
            self.segmenter = TokenTransformerSegments(
                segments=segments,
                in_dim=self.tokenizer.out_dim,
                hidden=hidden,
                num_heads=num_heads,
                kv_heads_cross=kv_heads_cross,
                kv_heads_self=kv_heads_self,
                dim_feedforward=dim_feedforward,
                dropout=dropout,
                latent_tokens=latent_tokens,
                latent_layers=latent_layers,
                decoder_layers=decoder_layers,
                cross_attn_interval=cross_attn_interval,
                norm_class=norm_class,
                activation=activation,
                segment_conf=segment_conf,
                pre_encoder_layers=pre_encoder_layers,
                segment_param=segment_param,
                length_floor=length_floor,
                decoder_input_xattn=decoder_input_xattn,
                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
            )

    def forward_tokens(self, tokens: torch.Tensor, mask: torch.Tensor):
        """Run the segmenter on pre-built token tensors."""
        return self.segmenter(tokens, mask)

	"""
	Perceiver-based transformer for 3D roof wireframe prediction.

	Architecture overview:

	Input tokens [B, T, D]
	\|
	v
	input_proj: Linear -> GELU -> Linear -> LayerNorm => [B, T, hidden]
	\|
	v
	Perceiver latent bottleneck (N PerceiverLatentLayers):
	Learnable latent embeddings [L, hidden] are broadcast to batch.
	Each layer: cross-attn(latents <- tokens) -> self-attn(latents) -> FFN
	Output: latents [B, L, hidden]
	\|
	v
	Segment decoder (M SegmentDecoderLayers):
	Learnable query embeddings [S, hidden] are broadcast to batch.
	Each layer: cross-attn(queries <- latents) -> self-attn(queries) -> FFN
	Output: queries [B, S, hidden]
	\|
	v
	segment_head: Linear -> 6D -> (midpoint, half_vector)
	+ query_offsets (learnable per-query bias)
	endpoints = midpoint +/- half_vector -> [B, S, 2, 3]
	"""

	import torch
	import torch.nn as nn

	from .attention import MultiHeadSDPA, FeedForward


	# ---------------------------------------------------------------------------
	# Building blocks
	# ---------------------------------------------------------------------------

	class AttnResidual(nn.Module):
	"""Pre-norm attention + residual + dropout."""

	def __init__(
	self,
	d_model: int,
	num_heads: int,
	dropout: float = 0.0,
	kv_heads: int \| None = None,
	norm_class=None,
	qk_norm: bool = False,
	qk_norm_type: str = "l2",
	):
	super().__init__()
	norm_class = norm_class or nn.LayerNorm
	self.norm = norm_class(d_model)
	self.attn = MultiHeadSDPA(d_model, num_heads, kv_heads=kv_heads, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
	self.drop = nn.Dropout(dropout)

	def forward(
	self,
	x: torch.Tensor,
	memory: torch.Tensor,
	memory_key_padding_mask: torch.Tensor \| None = None,
	) -> torch.Tensor:
	res = x
	x = self.norm(x)
	x = self.attn(x, memory, key_padding_mask=memory_key_padding_mask)
	return res + self.drop(x)


	class FFNResidual(nn.Module):
	"""Pre-norm feed-forward + residual + dropout."""

	def __init__(
	self,
	d_model: int,
	dim_ff: int,
	dropout: float = 0.0,
	activation: str = "gelu",
	norm_class=None,
	):
	super().__init__()
	norm_class = norm_class or nn.LayerNorm
	self.norm = norm_class(d_model)
	self.ffn = FeedForward(d_model, dim_ff, activation=activation)
	self.drop = nn.Dropout(dropout)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	res = x
	x = self.norm(x)
	x = self.ffn(x)
	return res + self.drop(x)


	# ---------------------------------------------------------------------------
	# Perceiver encoder layer
	# ---------------------------------------------------------------------------

	class PerceiverLatentLayer(nn.Module):
	"""Single Perceiver latent layer.

	If use_cross=True: cross-attn(latents <- points) -> self-attn -> FFN
	If use_cross=False: self-attn -> FFN (saves compute in deep stacks)
	"""

	def __init__(
	self,
	d_model: int,
	num_heads: int,
	dim_ff: int,
	dropout: float = 0.0,
	activation: str = "gelu",
	kv_heads_cross: int \| None = None,
	kv_heads_self: int \| None = None,
	use_cross: bool = True,
	norm_class=None,
	qk_norm: bool = False,
	qk_norm_type: str = "l2",
	):
	super().__init__()
	self.use_cross = use_cross
	if use_cross:
	self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
	self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
	self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)

	def forward(
	self,
	latents: torch.Tensor,
	points: torch.Tensor,
	points_key_padding_mask: torch.Tensor \| None = None,
	) -> torch.Tensor:
	if self.use_cross:
	latents = self.cross(latents, points, memory_key_padding_mask=points_key_padding_mask)
	latents = self.self_attn(latents, latents)
	latents = self.ffn(latents)
	return latents


	# ---------------------------------------------------------------------------
	# Segment decoder layer
	# ---------------------------------------------------------------------------

	class SegmentDecoderLayer(nn.Module):
	"""Single segment decoder layer.

	cross-attn(queries <- latents) -> [cross-attn(queries <- inputs)] -> self-attn(queries) -> FFN

	If input_xattn=True, adds a second cross-attention that attends directly
	to the projected input tokens (bypassing the latent bottleneck). This gives
	queries access to fine-grained point-level detail for vertex precision.
	"""

	def __init__(
	self,
	d_model: int,
	num_heads: int,
	dim_ff: int,
	dropout: float = 0.0,
	activation: str = "gelu",
	kv_heads_cross: int \| None = None,
	kv_heads_self: int \| None = None,
	norm_class=None,
	input_xattn: bool = False,
	qk_norm: bool = False,
	qk_norm_type: str = "l2",
	):
	super().__init__()
	self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
	self.input_xattn = input_xattn
	if input_xattn:
	self.cross_input = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
	self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
	self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)

	def forward(
	self,
	queries: torch.Tensor,
	latents: torch.Tensor,
	src: torch.Tensor \| None = None,
	src_key_padding_mask: torch.Tensor \| None = None,
	) -> torch.Tensor:
	queries = self.cross(queries, latents)
	if self.input_xattn and src is not None:
	queries = self.cross_input(queries, src, memory_key_padding_mask=src_key_padding_mask)
	queries = self.self_attn(queries, queries)
	queries = self.ffn(queries)
	return queries


	# ---------------------------------------------------------------------------
	# Full model
	# ---------------------------------------------------------------------------

	class TokenTransformerSegments(nn.Module):
	"""Perceiver transformer that predicts 3D roof wireframe segments.

	Takes point-cloud tokens and outputs segment endpoints as [B, S, 2, 3]
	where S is the number of segments and each segment has two 3D endpoints.

	Args:
	segments: Number of predicted segments (S).
	in_dim: Dimensionality of input tokens.
	hidden: Internal hidden dimension throughout the model.
	num_heads: Number of attention heads.
	kv_heads_cross: Grouped-query heads for cross-attention (None = standard MHA).
	kv_heads_self: Grouped-query heads for self-attention (None = standard MHA).
	dim_feedforward: FFN intermediate dimension.
	dropout: Dropout rate applied after attention and FFN.
	latent_tokens: Number of learnable latent embeddings (L) in the bottleneck.
	latent_layers: Number of PerceiverLatentLayers (N).
	decoder_layers: Number of SegmentDecoderLayers (M).
	"""

	def __init__(
	self,
	segments: int = 32,
	in_dim: int = 128,
	hidden: int = 128,
	num_heads: int = 4,
	kv_heads_cross: int \| None = 2,
	kv_heads_self: int \| None = 0,
	dim_feedforward: int = 256,
	dropout: float = 0.01,
	latent_tokens: int = 64,
	latent_layers: int = 2,
	decoder_layers: int = 2,
	cross_attn_interval: int = 1,
	norm_class=None,
	activation: str = "gelu",
	segment_conf: bool = False,
	pre_encoder_layers: int = 0,
	segment_param: str = "midpoint_halfvec",
	length_floor: float = 0.0,
	decoder_input_xattn: bool = False,
	qk_norm: bool = False,
	qk_norm_type: str = "l2",
	):
	super().__init__()
	self.segments = segments
	self.out_vertices = segments * 2
	self.segment_param = segment_param
	self.decoder_input_xattn = decoder_input_xattn
	norm_class = norm_class or nn.LayerNorm

	# Treat 0 as "use standard MHA"
	if kv_heads_cross is not None and kv_heads_cross <= 0:
	kv_heads_cross = None
	if kv_heads_self is not None and kv_heads_self <= 0:
	kv_heads_self = None

	# -- Input projection --
	self.input_proj = nn.Sequential(
	nn.Linear(in_dim, dim_feedforward),
	nn.GELU(),
	nn.Linear(dim_feedforward, hidden),
	norm_class(hidden),
	)

	# -- Optional pre-encoder: self-attention on full token sequence --
	if pre_encoder_layers > 0:
	self.pre_encoder = nn.ModuleList([
	SelfAttentionEncoderLayer(
	d_model=hidden,
	num_heads=num_heads,
	dim_ff=dim_feedforward,
	dropout=dropout,
	activation=activation,
	kv_heads=kv_heads_self,
	norm_class=norm_class,
	qk_norm=qk_norm, qk_norm_type=qk_norm_type,
	)
	for _ in range(pre_encoder_layers)
	])
	else:
	self.pre_encoder = None

	# -- Perceiver latent bottleneck --
	self.latent_embed = nn.Embedding(latent_tokens, hidden)
	N = latent_layers
	self.latent_layers = nn.ModuleList([
	PerceiverLatentLayer(
	d_model=hidden,
	num_heads=num_heads,
	dim_ff=dim_feedforward,
	dropout=dropout,
	activation=activation,
	kv_heads_cross=kv_heads_cross,
	kv_heads_self=kv_heads_self,
	use_cross=(i == 0) or (i == N - 1) or (i % cross_attn_interval == 0),
	norm_class=norm_class,
	qk_norm=qk_norm, qk_norm_type=qk_norm_type,
	)
	for i in range(N)
	])

	# -- Segment decoder --
	self.query_embed = nn.Embedding(segments, hidden)
	self.decoder_layers = nn.ModuleList([
	SegmentDecoderLayer(
	d_model=hidden,
	num_heads=num_heads,
	dim_ff=dim_feedforward,
	dropout=dropout,
	activation=activation,
	kv_heads_cross=kv_heads_cross,
	kv_heads_self=kv_heads_self,
	norm_class=norm_class,
	input_xattn=decoder_input_xattn,
	qk_norm=qk_norm, qk_norm_type=qk_norm_type,
	)
	for _ in range(decoder_layers)
	])

	# -- Output head --
	if segment_param == "midpoint_dir_len":
	self.segment_head = nn.Linear(hidden, 7) # mid(3) + dir(3) + len(1)
	else:
	self.segment_head = nn.Linear(hidden, 6) # mid(3) + half(3)
	self.query_offsets = nn.Parameter(torch.zeros(segments, 2, 3))

	nn.init.trunc_normal_(self.segment_head.weight, mean=0.0, std=1e-3)
	if self.segment_head.bias is not None:
	nn.init.zeros_(self.segment_head.bias)
	if segment_param == "midpoint_dir_len":
	# softplus(0.5) * 0.1 ≈ 0.097 default length in normalized space
	self.segment_head.bias.data[6] = 0.5
	nn.init.normal_(self.query_offsets, mean=0.0, std=0.05)

	# -- Optional confidence head --
	self.segment_conf = segment_conf
	if segment_conf:
	self.conf_head = nn.Linear(hidden, 1)
	nn.init.zeros_(self.conf_head.bias)

	def forward(
	self,
	tokens: torch.Tensor,
	mask: torch.Tensor \| None = None,
	) -> dict[str, torch.Tensor \| list]:
	"""
	Args:
	tokens: Input point-cloud tokens [B, T, in_dim].
	mask: Boolean validity mask [B, T]. True = valid token.

	Returns:
	Dict with keys:
	"vertices": [B, S*2, 3] flattened endpoints.
	"segments": [B, S, 2, 3] segment endpoints.
	"edges": Per-batch list of (start, end) index pairs into vertices.
	"conf": [B, S] logits (only if segment_conf=True).
	"""
	B = tokens.shape[0]

	# Project input tokens
	src = self.input_proj(tokens) # [B, T, hidden]

	# Padding mask (True where padded) for cross-attention
	pad_mask = ~mask.bool() if mask is not None else None

	# Optional pre-encoder: self-attention on full token sequence
	if self.pre_encoder is not None:
	for layer in self.pre_encoder:
	src = layer(src, key_padding_mask=pad_mask)

	# Perceiver latent bottleneck
	latents = self.latent_embed.weight.unsqueeze(0).expand(B, -1, -1)
	for layer in self.latent_layers:
	latents = layer(latents, src, points_key_padding_mask=pad_mask)

	# Segment decoder
	queries = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1)
	for layer in self.decoder_layers:
	queries = layer(queries, latents,
	src=src if self.decoder_input_xattn else None,
	src_key_padding_mask=pad_mask if self.decoder_input_xattn else None)

	# Predict segments -> endpoints
	if self.segment_param == "midpoint_dir_len":
	raw = self.segment_head(queries) # [B, S, 7]
	mid = raw[:, :, :3] + self.query_offsets[:, 0, :].unsqueeze(0)
	direction = torch.nn.functional.normalize(raw[:, :, 3:6], dim=-1)
	length = torch.nn.functional.softplus(raw[:, :, 6:7]) * 0.1
	half = direction * length * 0.5
	else:
	raw = self.segment_head(queries).view(B, self.segments, 2, 3)
	raw = raw + self.query_offsets.unsqueeze(0)
	mid, half = raw[:, :, 0], raw[:, :, 1]
	seg_params = torch.stack([mid - half, mid + half], dim=2)

	vertices = seg_params.reshape(B, self.out_vertices, 3)
	edges = [[(2 * i, 2 * i + 1) for i in range(self.segments)] for _ in range(B)]

	out = {"vertices": vertices, "segments": seg_params, "edges": edges,
	"src": src, "pad_mask": pad_mask, "queries": queries}
	if self.segment_conf:
	out["conf"] = self.conf_head(queries).squeeze(-1) # [B, S]
	return out


	# ---------------------------------------------------------------------------
	# Encoder-only layer (self-attention on full token sequence)
	# ---------------------------------------------------------------------------

	class SelfAttentionEncoderLayer(nn.Module):
	"""Single self-attention layer: self-attn(tokens) -> FFN."""

	def __init__(
	self,
	d_model: int,
	num_heads: int,
	dim_ff: int,
	dropout: float = 0.0,
	activation: str = "gelu",
	kv_heads: int \| None = None,
	norm_class=None,
	qk_norm: bool = False,
	qk_norm_type: str = "l2",
	):
	super().__init__()
	self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
	self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)

	def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor \| None = None) -> torch.Tensor:
	x = self.self_attn(x, x, memory_key_padding_mask=key_padding_mask)
	x = self.ffn(x)
	return x


	# ---------------------------------------------------------------------------
	# End-to-end model: tokenizer embeddings + perceiver
	# ---------------------------------------------------------------------------

	class EdgeDepthSegmentsModel(nn.Module):
	"""Tokenizer embeddings + transformer for 3D roof wireframes.

	Supports two architectures via the `arch` parameter:
	- "perceiver": Perceiver latent bottleneck (default, O(L*T) attention)
	- "transformer": Standard self-attention encoder (O(T^2) attention)

	Both share the same decoder, output head, and tokenizer.
	"""

	def __init__(
	self,
	seq_cfg,
	segments: int = 32,
	hidden: int = 128,
	num_heads: int = 4,
	kv_heads_cross: int \| None = 2,
	kv_heads_self: int \| None = 0,
	dim_feedforward: int = 256,
	dropout: float = 0.1,
	latent_tokens: int = 64,
	latent_layers: int = 1,
	decoder_layers: int = 2,
	label_emb_dim: int = 16,
	src_emb_dim: int = 2,
	behind_emb_dim: int = 8,
	fourier_seed: int = 0,
	cross_attn_interval: int = 1,
	norm_class=None,
	activation: str = "gelu",
	segment_conf: bool = False,
	use_vote_features: bool = False,
	arch: str = "perceiver",
	encoder_layers: int = 4,
	pre_encoder_layers: int = 0,
	segment_param: str = "midpoint_halfvec",
	length_floor: float = 0.0,
	decoder_input_xattn: bool = False,
	qk_norm: bool = False,
	qk_norm_type: str = "l2",
	learnable_fourier: bool = False,
	):
	super().__init__()
	self.seq_cfg = seq_cfg

	from .tokenizer import EdgeDepthSequenceBuilder
	self.tokenizer = EdgeDepthSequenceBuilder(
	seq_cfg,
	label_emb_dim=label_emb_dim,
	src_emb_dim=src_emb_dim,
	behind_emb_dim=behind_emb_dim,
	fourier_seed=fourier_seed,
	use_vote_features=use_vote_features,
	learnable_fourier=learnable_fourier,
	)

	if arch == "transformer":
	raise ValueError(
	"arch='transformer' is no longer supported. "
	"TransformerSegments has been removed; use arch='perceiver'.")
	else:
	self.segmenter = TokenTransformerSegments(
	segments=segments,
	in_dim=self.tokenizer.out_dim,
	hidden=hidden,
	num_heads=num_heads,
	kv_heads_cross=kv_heads_cross,
	kv_heads_self=kv_heads_self,
	dim_feedforward=dim_feedforward,
	dropout=dropout,
	latent_tokens=latent_tokens,
	latent_layers=latent_layers,
	decoder_layers=decoder_layers,
	cross_attn_interval=cross_attn_interval,
	norm_class=norm_class,
	activation=activation,
	segment_conf=segment_conf,
	pre_encoder_layers=pre_encoder_layers,
	segment_param=segment_param,
	length_floor=length_floor,
	decoder_input_xattn=decoder_input_xattn,
	qk_norm=qk_norm, qk_norm_type=qk_norm_type,
	)

	def forward_tokens(self, tokens: torch.Tensor, mask: torch.Tensor):
	"""Run the segmenter on pre-built token tensors."""
	return self.segmenter(tokens, mask)