"""
LW-DETR (Lightweight DETR), adapted from
https://github.com/Atten4Vis/LW-DETR

Paper "LW-DETR: A Transformer Replacement to YOLO for Real-Time Detection", https://arxiv.org/abs/2406.03459

Changes from original:
* Uses lite_refpoint_refine mode only (fixed reference points across decoder layers)
* Uses bbox_reparam mode only (anchor-delta parameterization)

Generated by gpt-5.2-codex xhigh.
"""

# Reference license: Apache-2.0

import copy
import math
from typing import Any
from typing import Optional

import torch
import torch.nn.functional as F
from torch import nn
from torchvision.ops import MLP
from torchvision.ops import Conv2dNormActivation
from torchvision.ops import boxes as box_ops

from birder.common import training_utils
from birder.layers import LayerNorm2d
from birder.model_registry import registry
from birder.net.base import DetectorBackbone
from birder.net.detection.base import DetectionBaseNet
from birder.net.detection.deformable_detr import HungarianMatcher
from birder.net.detection.deformable_detr import MultiScaleDeformableAttention
from birder.ops.soft_nms import SoftNMS


def _get_clones(module: nn.Module, N: int) -> nn.ModuleList:
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


def _apply_reference_delta(reference_points: torch.Tensor, bbox_delta: torch.Tensor) -> torch.Tensor:
    ref_xy = bbox_delta[..., :2] * reference_points[..., 2:] + reference_points[..., :2]
    ref_wh = bbox_delta[..., 2:].exp() * reference_points[..., 2:]
    return torch.concat((ref_xy, ref_wh), dim=-1)


def _gen_sine_embed_for_position(pos_tensor: torch.Tensor, dim: int, temperature: float = 10000.0) -> torch.Tensor:
    scale = 2 * math.pi
    dim_t = torch.arange(dim, dtype=torch.float32, device=pos_tensor.device)
    dim_t = temperature ** (2 * (dim_t // 2) / dim)

    x_embed = pos_tensor[..., 0] * scale
    y_embed = pos_tensor[..., 1] * scale
    pos_x = x_embed[..., None] / dim_t
    pos_y = y_embed[..., None] / dim_t
    pos_x = torch.stack((pos_x[..., 0::2].sin(), pos_x[..., 1::2].cos()), dim=-1).flatten(-2)
    pos_y = torch.stack((pos_y[..., 0::2].sin(), pos_y[..., 1::2].cos()), dim=-1).flatten(-2)

    if pos_tensor.size(-1) == 2:
        pos = torch.concat((pos_y, pos_x), dim=-1)
    elif pos_tensor.size(-1) == 4:
        w_embed = pos_tensor[..., 2] * scale
        h_embed = pos_tensor[..., 3] * scale
        pos_w = w_embed[..., None] / dim_t
        pos_h = h_embed[..., None] / dim_t
        pos_w = torch.stack((pos_w[..., 0::2].sin(), pos_w[..., 1::2].cos()), dim=-1).flatten(-2)
        pos_h = torch.stack((pos_h[..., 0::2].sin(), pos_h[..., 1::2].cos()), dim=-1).flatten(-2)
        pos = torch.concat((pos_y, pos_x, pos_w, pos_h), dim=-1)
    else:
        raise ValueError(f"Unknown pos_tensor shape(-1): {pos_tensor.size(-1)}")

    return pos


class MultiheadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, attn_drop: float = 0.0, proj_drop: float = 0.0) -> None:
        super().__init__()
        assert d_model % num_heads == 0, "d_model should be divisible by num_heads"

        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.scale = self.head_dim**-0.5

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(d_model, d_model)
        self.proj_drop = nn.Dropout(proj_drop)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.xavier_uniform_(self.q_proj.weight)
        nn.init.xavier_uniform_(self.k_proj.weight)
        nn.init.xavier_uniform_(self.v_proj.weight)
        nn.init.xavier_uniform_(self.proj.weight)
        if self.q_proj.bias is not None:
            nn.init.zeros_(self.q_proj.bias)
            nn.init.zeros_(self.k_proj.bias)
            nn.init.zeros_(self.v_proj.bias)
            nn.init.zeros_(self.proj.bias)

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        B, l_q, C = query.size()
        q = self.q_proj(query).reshape(B, l_q, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(key).reshape(B, key.size(1), self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(value).reshape(B, value.size(1), self.num_heads, self.head_dim).transpose(1, 2)

        if key_padding_mask is not None:
            # key_padding_mask is expected to be boolean (True = masked)
            # SDPA expects True = attend, so we invert
            attn_mask = ~key_padding_mask[:, None, None, :]
        else:
            attn_mask = None

        attn = F.scaled_dot_product_attention(  # pylint: disable=not-callable
            q, k, v, attn_mask=attn_mask, dropout_p=self.attn_drop.p if self.training else 0.0, scale=self.scale
        )

        attn = attn.transpose(1, 2).reshape(B, l_q, C)
        x = self.proj(attn)
        x = self.proj_drop(x)

        return x


class Bottleneck(nn.Module):
    def __init__(self, channels: int, expansion: float = 1.0) -> None:
        super().__init__()
        hidden = int(channels * expansion)
        self.conv1 = Conv2dNormActivation(
            channels, hidden, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), activation_layer=nn.SiLU
        )
        self.conv2 = Conv2dNormActivation(
            hidden, channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), activation_layer=nn.SiLU
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.conv2(self.conv1(x))


class C2f(nn.Module):
    """
    CSP Bottleneck with 2 convolutions
    """

    def __init__(self, in_channels: int, out_channels: int, num_blocks: int, expansion: float = 0.5) -> None:
        super().__init__()
        hidden = int(out_channels * expansion)
        self.conv1 = Conv2dNormActivation(
            in_channels, hidden * 2, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), activation_layer=nn.SiLU
        )
        self.blocks = nn.ModuleList([Bottleneck(hidden, expansion=1.0) for _ in range(num_blocks)])
        self.conv2 = Conv2dNormActivation(
            hidden * (2 + num_blocks),
            out_channels,
            kernel_size=(1, 1),
            stride=(1, 1),
            padding=(0, 0),
            activation_layer=nn.SiLU,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = self.conv1(x)
        c = y.size(1) // 2
        y = list(y.split([c, c], dim=1))
        y.extend(block(y[-1]) for block in self.blocks)

        return self.conv2(torch.concat(y, dim=1))


class MultiScaleProjector(nn.Module):
    """
    LW-DETR Multi-Scale Projector that creates pyramid features from backbone feature maps

    For each output scale, applies per-level sampling (upsample/downsample) to each input
    feature map, resizes to match, concatenates them, and passes through C2f + LayerNorm.
    """

    def __init__(self, in_channels: list[int], out_channels: int, scale_factors: list[float], num_blocks: int) -> None:
        super().__init__()
        self.use_extra_pool = False

        self.stage_scales: list[float] = []
        self.stages = nn.ModuleList()
        for scale in scale_factors:
            scale_samplers: list[nn.Module] = []
            fused_channels = 0

            for in_dim in in_channels:
                out_dim = in_dim
                if scale == 4.0:
                    # 4x upsample via two ConvTranspose2d
                    layers: nn.Module = nn.Sequential(
                        nn.ConvTranspose2d(in_dim, in_dim // 2, kernel_size=(2, 2), stride=(2, 2), padding=(0, 0)),
                        LayerNorm2d(in_dim // 2, eps=1e-6),
                        nn.GELU(),
                        nn.ConvTranspose2d(in_dim // 2, in_dim // 4, kernel_size=(2, 2), stride=(2, 2), padding=(0, 0)),
                    )
                    out_dim = in_dim // 4
                elif scale == 2.0:
                    # 2x upsample
                    if in_dim > 512:
                        layers = nn.Sequential(
                            Conv2dNormActivation(in_dim, in_dim // 2, kernel_size=1, activation_layer=nn.ReLU),
                            nn.ConvTranspose2d(
                                in_dim // 2, in_dim // 4, kernel_size=(2, 2), stride=(2, 2), padding=(0, 0)
                            ),
                        )
                        out_dim = in_dim // 4
                    else:
                        layers = nn.ConvTranspose2d(
                            in_dim, in_dim // 2, kernel_size=(2, 2), stride=(2, 2), padding=(0, 0)
                        )
                        out_dim = in_dim // 2
                elif scale == 1.0:
                    # Identity (will resize to match in forward)
                    layers = nn.Identity()
                    out_dim = in_dim
                elif scale == 0.5:
                    # 2x downsample
                    layers = Conv2dNormActivation(
                        in_dim, in_dim, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), activation_layer=nn.SiLU
                    )
                    out_dim = in_dim
                elif scale == 0.25:
                    # Will use max pooling on previous scale output
                    self.use_extra_pool = True
                    continue
                else:
                    raise ValueError(f"Unsupported scale_factor: {scale}")

                scale_samplers.append(layers)
                fused_channels += out_dim

            if len(scale_samplers) > 0:
                fuser = nn.Sequential(
                    C2f(fused_channels, out_channels, num_blocks),
                    LayerNorm2d(out_channels, eps=1e-6),
                )
                stage = nn.ModuleDict(
                    {
                        "samplers": nn.ModuleList(scale_samplers),
                        "fuser": fuser,
                    }
                )
                self.stage_scales.append(float(scale))
                self.stages.append(stage)

    def forward(self, features: list[torch.Tensor]) -> list[torch.Tensor]:
        results: list[torch.Tensor] = []

        # Use the middle feature level as reference for target size computation
        ref_idx = len(features) // 2
        ref_h, ref_w = features[ref_idx].shape[-2:]

        for stage_idx, stage in enumerate(self.stages):
            scale = self.stage_scales[stage_idx]
            target_h = int(ref_h * scale)
            target_w = int(ref_w * scale)

            samplers = stage["samplers"]
            fuser = stage["fuser"]

            fused_feats: list[torch.Tensor] = []
            for level_idx, sampler in enumerate(samplers):
                feat = sampler(features[level_idx])
                if feat.shape[-2:] != (target_h, target_w):
                    feat = F.interpolate(feat, size=(target_h, target_w), mode="nearest")

                fused_feats.append(feat)

            fused = torch.concat(fused_feats, dim=1)
            results.append(fuser(fused))

        if self.use_extra_pool is True:
            results.append(F.max_pool2d(results[-1], kernel_size=(1, 1), stride=(2, 2), padding=(0, 0)))

        return results


class LWDeformableTransformerDecoderLayer(nn.Module):
    def __init__(
        self,
        d_model: int,
        d_ffn: int,
        dropout: float,
        n_levels: int,
        self_heads: int,
        cross_heads: int,
        n_points: int,
        num_queries: int,
        group_detr: int,
    ) -> None:
        super().__init__()
        self.num_queries = num_queries
        self.group_detr = group_detr

        self.self_attn = MultiheadAttention(d_model, self_heads, attn_drop=dropout)
        self.norm1 = nn.LayerNorm(d_model)

        self.cross_attn = MultiScaleDeformableAttention(d_model, n_levels, cross_heads, n_points)
        self.norm2 = nn.LayerNorm(d_model)

        self.linear1 = nn.Linear(d_model, d_ffn)
        self.linear2 = nn.Linear(d_ffn, d_model)
        self.norm3 = nn.LayerNorm(d_model)

        self.activation = nn.ReLU()
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        tgt: torch.Tensor,
        query_pos: torch.Tensor,
        reference_points: torch.Tensor,
        src: torch.Tensor,
        src_spatial_shapes: torch.Tensor,
        src_shapes: list[list[int]],
        level_start_index: torch.Tensor,
        src_padding_mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        q_k = tgt + query_pos

        # Group-DETR: during training, split queries by groups so each group attends only within itself
        if self.training is True and self.group_detr > 1:
            B, N, C = q_k.size()
            num_queries_per_group = N // self.group_detr
            qk_grouped = q_k.view(B, self.group_detr, num_queries_per_group, C)
            qk_grouped = qk_grouped.reshape(B * self.group_detr, num_queries_per_group, C)

            tgt_grouped = tgt.view(B, self.group_detr, num_queries_per_group, C)
            tgt_grouped = tgt_grouped.reshape(B * self.group_detr, num_queries_per_group, C)

            tgt2 = self.self_attn(qk_grouped, qk_grouped, value=tgt_grouped)
            tgt2 = tgt2.view(B, self.group_detr, num_queries_per_group, C)
            tgt2 = tgt2.reshape(B, N, C)
        else:
            tgt2 = self.self_attn(q_k, q_k, value=tgt)

        tgt = tgt + self.dropout(tgt2)
        tgt = self.norm1(tgt)

        tgt2 = self.cross_attn(
            tgt + query_pos, reference_points, src, src_spatial_shapes, level_start_index, src_padding_mask, src_shapes
        )
        tgt = tgt + self.dropout(tgt2)
        tgt = self.norm2(tgt)

        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout(tgt2)
        tgt = self.norm3(tgt)

        return tgt


# pylint: disable=invalid-name
class LW_DETRDecoder(nn.Module):
    """
    LW-DETR Decoder with learned queries and fixed reference points

    Uses learned content queries and reference point embeddings. Supports Group-DETR
    training where queries are replicated into multiple groups with independent matching.
    In two-stage mode, reference points come from encoder proposals instead of learned embeddings.
    """

    def __init__(
        self,
        hidden_dim: int,
        num_classes: int,
        num_queries: int,
        num_decoder_layers: int,
        num_levels: int,
        num_self_heads: int,
        num_cross_heads: int,
        dim_feedforward: int,
        dropout: float,
        num_decoder_points: int,
        group_detr: int,
        two_stage: bool,
    ) -> None:
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_queries = num_queries
        self.num_layers = num_decoder_layers
        self.num_levels = num_levels
        self.group_detr = group_detr
        self.two_stage = two_stage

        # Learnable content queries and reference points (per Group-DETR)
        self.content_queries = nn.Parameter(torch.empty(num_queries * group_detr, hidden_dim))
        self.reference_point_embed = nn.Parameter(torch.empty(num_queries * group_detr, 4))

        decoder_layer = LWDeformableTransformerDecoderLayer(
            hidden_dim,
            dim_feedforward,
            dropout,
            num_levels,
            num_self_heads,
            num_cross_heads,
            num_decoder_points,
            num_queries,
            group_detr,
        )
        self.layers = _get_clones(decoder_layer, num_decoder_layers)

        # Sinusoidal embeddings: 4D coords -> 4 * (hidden_dim/2) = 2 * hidden_dim
        self.query_pos_head = MLP(2 * hidden_dim, [hidden_dim, hidden_dim], activation_layer=nn.ReLU)

        self.class_embed = nn.Linear(hidden_dim, num_classes)
        self.bbox_embed = MLP(hidden_dim, [hidden_dim, hidden_dim, 4], activation_layer=nn.ReLU)

        # Weights initialization
        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        nn.init.xavier_uniform_(self.content_queries)
        nn.init.zeros_(self.reference_point_embed)
        nn.init.constant_(self.class_embed.bias, bias_value)
        nn.init.zeros_(self.bbox_embed[-2].weight)
        nn.init.zeros_(self.bbox_embed[-2].bias)

    def reset_classifier(self, num_classes: int) -> None:
        self.class_embed = nn.Linear(self.hidden_dim, num_classes)
        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        nn.init.constant_(self.class_embed.bias, bias_value)

    # pylint: disable=too-many-locals,too-many-branches
    def forward(
        self,
        feats: list[torch.Tensor],
        spatial_shapes: list[list[int]],
        level_start_index: list[int],
        valid_ratios: torch.Tensor,
        padding_mask: Optional[list[torch.Tensor]] = None,
        enc_reference_points: Optional[torch.Tensor] = None,
        return_intermediates: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        memory = []
        mask_flatten = []
        for idx, feat in enumerate(feats):
            feat = feat.flatten(2).permute(0, 2, 1)
            memory.append(feat)
            if padding_mask is not None:
                mask_flatten.append(padding_mask[idx].flatten(1))

        memory = torch.concat(memory, dim=1)
        memory_padding_mask = torch.concat(mask_flatten, dim=1) if len(mask_flatten) > 0 else None

        # Two-stage mode: combine learned reference point deltas with encoder proposals
        if self.two_stage is True:
            assert enc_reference_points is not None
            # Slice content queries and reference point embeddings to match encoder proposals
            num_enc_queries = enc_reference_points.size(1)
            if self.training is True:
                if self.group_detr > 1:
                    if num_enc_queries % self.group_detr != 0:
                        raise ValueError("num_enc_queries must be divisible by group_detr in training mode")
                    num_queries_per_group = num_enc_queries // self.group_detr
                    content_queries_weight = self.content_queries.reshape(
                        self.group_detr, self.num_queries, self.hidden_dim
                    )[:, :num_queries_per_group].reshape(-1, self.hidden_dim)
                    refpoint_embed_weight = self.reference_point_embed.reshape(self.group_detr, self.num_queries, 4)[
                        :, :num_queries_per_group
                    ].reshape(-1, 4)
                else:
                    content_queries_weight = self.content_queries[:num_enc_queries]
                    refpoint_embed_weight = self.reference_point_embed[:num_enc_queries]
            else:
                num_to_use = min(self.num_queries, num_enc_queries)
                content_queries_weight = self.content_queries[:num_to_use]
                refpoint_embed_weight = self.reference_point_embed[:num_to_use]
                enc_reference_points = enc_reference_points[:, :num_to_use]

            # Learned embeddings act as deltas on top of encoder proposals (bbox_reparam formula)
            refpoint_embed = refpoint_embed_weight.unsqueeze(0).expand(memory.size(0), -1, -1)
            refpoint_cxcy = refpoint_embed[..., :2] * enc_reference_points[..., 2:] + enc_reference_points[..., :2]
            refpoint_wh = refpoint_embed[..., 2:].exp() * enc_reference_points[..., 2:]
            reference_points = torch.concat([refpoint_cxcy, refpoint_wh], dim=-1)
        else:
            if self.training is True:
                content_queries_weight = self.content_queries
                refpoint_embed_weight = self.reference_point_embed
            else:
                content_queries_weight = self.content_queries[: self.num_queries]
                refpoint_embed_weight = self.reference_point_embed[: self.num_queries]

            init_ref_points_unact = refpoint_embed_weight.unsqueeze(0).expand(memory.size(0), -1, -1)

            # One-stage mode: reference points are un-activated, need sigmoid
            reference_points = init_ref_points_unact.sigmoid()

        target = content_queries_weight.unsqueeze(0).expand(memory.size(0), -1, -1)

        spatial_shapes_tensor = torch.tensor(spatial_shapes, dtype=torch.long, device=memory.device)
        level_start_index_tensor = torch.tensor(level_start_index, dtype=torch.long, device=memory.device)

        # lite_refpoint_refine mode: compute query_pos ONCE from initial reference points
        # Reference points stay fixed throughout all decoder layers
        if reference_points.shape[-1] == 4:
            reference_points_input = (
                reference_points[:, :, None] * torch.concat([valid_ratios, valid_ratios], dim=-1)[:, None]
            )
        else:
            reference_points_input = reference_points[:, :, None] * valid_ratios[:, None]

        query_sine_embed = _gen_sine_embed_for_position(reference_points_input[:, :, 0], self.hidden_dim // 2)
        query_pos = self.query_pos_head(query_sine_embed)

        bboxes_list: list[torch.Tensor] = []
        logits_list: list[torch.Tensor] = []
        layer_bboxes = torch.empty(0, device=target.device, dtype=target.dtype)
        class_logits = torch.empty(0, device=target.device, dtype=target.dtype)
        last_idx = len(self.layers) - 1
        for layer_idx, decoder_layer in enumerate(self.layers):
            target = decoder_layer(
                target,
                query_pos,
                reference_points_input,
                memory,
                spatial_shapes_tensor,
                spatial_shapes,
                level_start_index_tensor,
                memory_padding_mask,
            )

            if return_intermediates is True or layer_idx == last_idx:
                # Each layer predicts boxes from the same initial reference points
                bbox_delta = self.bbox_embed(target)
                layer_bboxes = _apply_reference_delta(reference_points, bbox_delta)
                class_logits = self.class_embed(target)
                if return_intermediates is True:
                    bboxes_list.append(layer_bboxes)
                    logits_list.append(class_logits)

        if return_intermediates is True:
            out_bboxes = torch.stack(bboxes_list)
            out_logits = torch.stack(logits_list)
        else:
            out_bboxes = layer_bboxes
            out_logits = class_logits

        return (out_bboxes, out_logits)


# pylint: disable=invalid-name
class LW_DETR(DetectionBaseNet):
    default_size = (640, 640)

    def __init__(
        self,
        num_classes: int,
        backbone: DetectorBackbone,
        *,
        config: Optional[dict[str, Any]] = None,
        size: Optional[tuple[int, int]] = None,
        export_mode: bool = False,
    ) -> None:
        super().__init__(num_classes, backbone, config=config, size=size, export_mode=export_mode)
        assert self.config is not None, "must set config"

        # Sigmoid based classification (like multi-label networks)
        self.num_classes = self.num_classes - 1

        dropout = 0.0
        num_decoder_layers: int = self.config.get("num_decoder_layers", 3)
        projector_dim: int = self.config.get("projector_dim", 256)
        projector_scale_factors: list[float] = self.config["projector_scale_factors"]
        projector_num_blocks: int = self.config.get("projector_num_blocks", 3)
        decoder_dim: int = self.config.get("decoder_dim", 256)
        decoder_self_heads: int = self.config.get("decoder_self_heads", 8)
        decoder_cross_heads: int = self.config.get("decoder_cross_heads", decoder_self_heads)
        decoder_points: int = self.config.get("decoder_points", 2)
        num_queries: int = self.config.get("num_queries", 300)
        two_stage: bool = self.config.get("two_stage", False)
        soft_nms: bool = self.config.get("soft_nms", False)

        self.soft_nms = None
        if soft_nms is True:
            self.soft_nms = SoftNMS()

        self.ia_alpha = 0.25
        self.loss_class = 1.0
        self.loss_bbox = 5.0
        self.loss_giou = 2.0
        self.aux_loss = True
        self.num_queries = num_queries
        self.two_stage = two_stage
        self.group_detr = 13

        self.backbone.return_channels = self.backbone.return_channels[-3:]
        self.backbone.return_stages = self.backbone.return_stages[-3:]

        num_decoder_levels = len(projector_scale_factors)

        self.projector = MultiScaleProjector(
            in_channels=list(self.backbone.return_channels),
            out_channels=projector_dim,
            scale_factors=projector_scale_factors,
            num_blocks=projector_num_blocks,
        )
        self.decoder_input_proj = nn.ModuleList(
            [
                (
                    nn.Conv2d(projector_dim, decoder_dim, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
                    if projector_dim != decoder_dim
                    else nn.Identity()
                )
                for _ in range(num_decoder_levels)
            ]
        )
        self.decoder = LW_DETRDecoder(
            hidden_dim=decoder_dim,
            num_classes=self.num_classes,
            num_queries=self.num_queries,
            num_decoder_layers=num_decoder_layers,
            num_levels=num_decoder_levels,
            num_self_heads=decoder_self_heads,
            num_cross_heads=decoder_cross_heads,
            dim_feedforward=decoder_dim * 4,
            dropout=dropout,
            num_decoder_points=decoder_points,
            group_detr=self.group_detr,
            two_stage=self.two_stage,
        )

        # Two-stage encoder output modules
        if self.two_stage is True:
            self.enc_output = nn.ModuleList([nn.Linear(decoder_dim, decoder_dim) for _ in range(self.group_detr)])
            self.enc_output_norm = nn.ModuleList([nn.LayerNorm(decoder_dim) for _ in range(self.group_detr)])
            # Encoder output class and bbox embeddings
            self.enc_out_class_embed = nn.ModuleList(
                [copy.deepcopy(self.decoder.class_embed) for _ in range(self.group_detr)]
            )
            self.enc_out_bbox_embed = nn.ModuleList(
                [copy.deepcopy(self.decoder.bbox_embed) for _ in range(self.group_detr)]
            )

            # Weights initialization
            for enc_output in self.enc_output:
                nn.init.xavier_uniform_(enc_output.weight)
                if enc_output.bias is not None:
                    nn.init.zeros_(enc_output.bias)

            for enc_output_norm in self.enc_output_norm:
                nn.init.ones_(enc_output_norm.weight)
                nn.init.zeros_(enc_output_norm.bias)

        else:
            self.enc_output = nn.ModuleList([nn.Identity()])
            self.enc_output_norm = nn.ModuleList([nn.Identity()])
            self.enc_out_class_embed = nn.ModuleList([nn.Identity()])
            self.enc_out_bbox_embed = nn.ModuleList([nn.Identity()])

        self.decoder_dim = decoder_dim
        self.matcher = HungarianMatcher(cost_class=self.loss_class, cost_bbox=self.loss_bbox, cost_giou=self.loss_giou)

        if self.export_mode is False:
            self.forward = torch.compiler.disable(recursive=False)(self.forward)  # type: ignore[method-assign]

    def reset_classifier(self, num_classes: int) -> None:
        self.num_classes = num_classes
        self.decoder.reset_classifier(num_classes)
        if self.two_stage is True:
            self.enc_out_class_embed = nn.ModuleList(
                [copy.deepcopy(self.decoder.class_embed) for _ in range(self.group_detr)]
            )

    def freeze(self, freeze_classifier: bool = True) -> None:
        for param in self.parameters():
            param.requires_grad_(False)
        if freeze_classifier is False:
            for param in self.decoder.class_embed.parameters():
                param.requires_grad_(True)

            if self.two_stage is True:
                for enc_class_embed in self.enc_out_class_embed:
                    for param in enc_class_embed.parameters():
                        param.requires_grad_(True)

    @staticmethod
    def _get_src_permutation_idx(indices: list[tuple[torch.Tensor, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]:
        batch_idx = torch.concat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
        src_idx = torch.concat([src for (src, _) in indices])
        return (batch_idx, src_idx)

    def _class_loss(
        self,
        cls_logits: torch.Tensor,
        box_output: torch.Tensor,
        targets: list[dict[str, torch.Tensor]],
        indices: list[tuple[torch.Tensor, torch.Tensor]],
        num_boxes: float,
    ) -> torch.Tensor:
        idx = self._get_src_permutation_idx(indices)
        pred_scores = cls_logits.sigmoid()
        target_scores = torch.zeros_like(cls_logits)
        pos_mask = torch.zeros_like(cls_logits, dtype=torch.bool)
        if idx[0].numel() > 0:
            target_classes = torch.concat([t["labels"][J] for t, (_, J) in zip(targets, indices)], dim=0)
            src_boxes = box_output[idx]
            target_boxes = torch.concat([t["boxes"][J] for t, (_, J) in zip(targets, indices)], dim=0)
            iou = box_ops.box_iou(
                box_ops.box_convert(src_boxes, in_fmt="cxcywh", out_fmt="xyxy"),
                box_ops.box_convert(target_boxes, in_fmt="cxcywh", out_fmt="xyxy"),
            ).diag()
            iou = iou.clamp(min=0).detach()
            mixed_target = pred_scores.detach()[idx[0], idx[1], target_classes] ** self.ia_alpha
            mixed_target = mixed_target * (iou ** (1.0 - self.ia_alpha))
            target_scores[idx[0], idx[1], target_classes] = mixed_target.to(target_scores.dtype)
            pos_mask[idx[0], idx[1], target_classes] = True

        neg_mask = ~pos_mask

        pos_loss = F.binary_cross_entropy_with_logits(cls_logits, target_scores, reduction="none")
        neg_loss = F.binary_cross_entropy_with_logits(cls_logits, torch.zeros_like(cls_logits), reduction="none")
        neg_loss = neg_loss * pred_scores.pow(2)

        loss_pos = pos_loss[pos_mask].sum()
        loss_neg = neg_loss[neg_mask].sum()
        return (loss_pos + loss_neg) / num_boxes

    def _box_loss(
        self,
        box_output: torch.Tensor,
        targets: list[dict[str, torch.Tensor]],
        indices: list[tuple[torch.Tensor, torch.Tensor]],
        num_boxes: float,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        idx = self._get_src_permutation_idx(indices)
        src_boxes = box_output[idx]
        target_boxes = torch.concat([t["boxes"][i] for t, (_, i) in zip(targets, indices)], dim=0)

        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction="none")
        loss_bbox = loss_bbox.sum() / num_boxes

        loss_giou = 1 - torch.diag(
            box_ops.generalized_box_iou(
                box_ops.box_convert(src_boxes, in_fmt="cxcywh", out_fmt="xyxy"),
                box_ops.box_convert(target_boxes, in_fmt="cxcywh", out_fmt="xyxy"),
            )
        )
        loss_giou = loss_giou.sum() / num_boxes

        return (loss_bbox, loss_giou)

    def _match_with_groups(
        self, cls_logits: torch.Tensor, box_output: torch.Tensor, targets: list[dict[str, torch.Tensor]]
    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
        """
        Perform Hungarian matching with Group-DETR support

        Reshapes queries into a larger batch (B * G) so a single matcher call
        can be used, then combines indices with proper group offsets.
        """

        if self.group_detr <= 1:
            return self.matcher(cls_logits, box_output, targets)  # type: ignore[no-any-return]

        batch_size = cls_logits.size(0)
        group_detr = self.group_detr
        total_queries = cls_logits.size(1)
        num_queries_per_group = total_queries // group_detr

        # Reshape from (B, G*Q, C) to (B*G, Q, C) to batch all groups together
        cls_logits_grouped = cls_logits.reshape(batch_size * group_detr, num_queries_per_group, -1)
        box_output_grouped = box_output.reshape(batch_size * group_detr, num_queries_per_group, -1)

        # Duplicate targets for each group: each group within a batch matches against the same targets
        targets_grouped = [targets[batch_idx] for batch_idx in range(batch_size) for _ in range(group_detr)]

        grouped_indices = self.matcher(cls_logits_grouped, box_output_grouped, targets_grouped)

        all_indices: list[tuple[torch.Tensor, torch.Tensor]] = []
        for batch_idx in range(batch_size):
            src_indices_list: list[torch.Tensor] = []
            tgt_indices_list: list[torch.Tensor] = []

            base = batch_idx * group_detr
            for group_idx in range(group_detr):
                src_idx, tgt_idx = grouped_indices[base + group_idx]
                # Adjust source indices by group offset to map back to original query positions
                src_indices_list.append(src_idx + group_idx * num_queries_per_group)
                tgt_indices_list.append(tgt_idx)

            # Combine indices from all groups for this batch element
            all_indices.append((torch.concat(src_indices_list), torch.concat(tgt_indices_list)))

        return all_indices

    def _compute_loss_from_outputs(
        self, targets: list[dict[str, torch.Tensor]], cls_logits: torch.Tensor, box_output: torch.Tensor
    ) -> dict[str, torch.Tensor]:
        num_boxes = sum(len(t["labels"]) for t in targets)
        num_boxes = torch.as_tensor(num_boxes, dtype=torch.float, device=cls_logits.device)
        if training_utils.is_dist_available_and_initialized() is True:
            torch.distributed.all_reduce(num_boxes)

        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1)

        # Normalize by group_detr for Group-DETR training
        num_boxes_normalized = num_boxes * self.group_detr

        loss_ce_list = []
        loss_bbox_list = []
        loss_giou_list = []
        for idx in range(cls_logits.size(0)):
            indices = self._match_with_groups(cls_logits[idx], box_output[idx], targets)

            loss_ce_list.append(
                self._class_loss(cls_logits[idx], box_output[idx], targets, indices, num_boxes_normalized)
            )
            loss_bbox_i, loss_giou_i = self._box_loss(box_output[idx], targets, indices, num_boxes_normalized)
            loss_bbox_list.append(loss_bbox_i)
            loss_giou_list.append(loss_giou_i)

        loss_ce = torch.stack(loss_ce_list).sum() * self.loss_class
        loss_bbox = torch.stack(loss_bbox_list).sum() * self.loss_bbox
        loss_giou = torch.stack(loss_giou_list).sum() * self.loss_giou

        losses = {"labels": loss_ce, "boxes": loss_bbox, "giou": loss_giou}

        return losses

    def _compute_enc_output_loss(
        self, targets: list[dict[str, torch.Tensor]], enc_cls_logits: torch.Tensor, enc_box_output: torch.Tensor
    ) -> dict[str, torch.Tensor]:
        num_boxes = sum(len(t["labels"]) for t in targets)
        num_boxes = torch.as_tensor(num_boxes, dtype=torch.float, device=enc_cls_logits.device)
        if training_utils.is_dist_available_and_initialized() is True:
            torch.distributed.all_reduce(num_boxes)

        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1)
        num_boxes_normalized = num_boxes * self.group_detr

        indices = self._match_with_groups(enc_cls_logits, enc_box_output, targets)
        loss_ce = self._class_loss(enc_cls_logits, enc_box_output, targets, indices, num_boxes_normalized)
        loss_bbox, loss_giou = self._box_loss(enc_box_output, targets, indices, num_boxes_normalized)

        return {
            "labels_enc": loss_ce * self.loss_class,
            "boxes_enc": loss_bbox * self.loss_bbox,
            "giou_enc": loss_giou * self.loss_giou,
        }

    @torch.jit.unused  # type: ignore[untyped-decorator]
    @torch.compiler.disable()  # type: ignore[untyped-decorator]
    def compute_loss(
        self,
        targets: list[dict[str, torch.Tensor]],
        out_logits: torch.Tensor,
        out_bboxes: torch.Tensor,
        image_sizes: torch.Tensor,
        enc_cls_logits: Optional[torch.Tensor] = None,
        enc_box_output: Optional[torch.Tensor] = None,
    ) -> dict[str, torch.Tensor]:
        losses: dict[str, torch.Tensor] = {}

        # Compute encoder output loss first (before modifying targets in-place)
        if self.two_stage is True and enc_cls_logits is not None and enc_box_output is not None:
            enc_targets = copy.deepcopy(targets)
            for idx, target in enumerate(enc_targets):
                boxes = target["boxes"]
                boxes = box_ops.box_convert(boxes, in_fmt="xyxy", out_fmt="cxcywh")
                scale = image_sizes[idx].flip(0).repeat(2).float()  # flip to [W, H], repeat to [W, H, W, H]
                boxes = boxes / scale
                enc_targets[idx]["boxes"] = boxes
                enc_targets[idx]["labels"] = target["labels"] - 1

            enc_losses = self._compute_enc_output_loss(enc_targets, enc_cls_logits, enc_box_output)
            losses.update(enc_losses)

        # Transform targets for decoder loss
        for idx, target in enumerate(targets):
            boxes = target["boxes"]
            boxes = box_ops.box_convert(boxes, in_fmt="xyxy", out_fmt="cxcywh")
            scale = image_sizes[idx].flip(0).repeat(2).float()  # flip to [W, H], repeat to [W, H, W, H]
            boxes = boxes / scale
            targets[idx]["boxes"] = boxes
            targets[idx]["labels"] = target["labels"] - 1

        if self.aux_loss is False:
            out_bboxes = out_bboxes[-1:].contiguous()
            out_logits = out_logits[-1:].contiguous()

        decoder_losses = self._compute_loss_from_outputs(targets, out_logits, out_bboxes)
        losses.update(decoder_losses)

        return losses

    @staticmethod
    def _get_valid_ratio(mask: torch.Tensor) -> torch.Tensor:
        _, H, W = mask.size()
        valid_h = torch.sum(~mask[:, :, 0], 1)
        valid_w = torch.sum(~mask[:, 0, :], 1)
        valid_ratio_h = valid_h.float() / H
        valid_ratio_w = valid_w.float() / W
        return torch.stack([valid_ratio_w, valid_ratio_h], dim=-1)

    def gen_encoder_output_proposals(
        self, memory: torch.Tensor, memory_padding_mask: Optional[torch.Tensor], spatial_shapes: list[list[int]]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Generate encoder output proposals for two-stage detection

        Creates bounding box proposals from encoder memory by generating a grid of anchors
        across all spatial positions and feature levels.
        """

        N = memory.size(0)
        proposals = []
        cur = 0
        for lvl, (H, W) in enumerate(spatial_shapes):
            if memory_padding_mask is not None:
                mask_flatten = memory_padding_mask[:, cur : cur + H * W].view(N, H, W, 1)
                valid_h = torch.sum(~mask_flatten[:, :, 0, 0], 1)
                valid_w = torch.sum(~mask_flatten[:, 0, :, 0], 1)
            else:
                valid_h = torch.full((N,), H, dtype=torch.float32, device=memory.device)
                valid_w = torch.full((N,), W, dtype=torch.float32, device=memory.device)

            grid_y, grid_x = torch.meshgrid(
                torch.linspace(0, H - 1, H, dtype=torch.float32, device=memory.device),
                torch.linspace(0, W - 1, W, dtype=torch.float32, device=memory.device),
                indexing="ij",
            )
            grid = torch.stack([grid_x, grid_y], dim=-1)  # (H, W, 2)

            scale = torch.stack([valid_w, valid_h], dim=1).view(N, 1, 1, 2)
            grid = (grid.unsqueeze(0).expand(N, -1, -1, -1) + 0.5) / scale
            wh = torch.ones_like(grid) * 0.05 * (2.0**lvl)
            proposal = torch.concat([grid, wh], dim=-1).view(N, -1, 4)
            proposals.append(proposal)
            cur += H * W

        output_proposals = torch.concat(proposals, dim=1)
        output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(dim=-1, keepdim=True)

        # bbox_reparam mode: keep proposals in normalized [0,1] space (no un-sigmoid transform)
        if memory_padding_mask is not None:
            output_proposals = output_proposals.masked_fill(memory_padding_mask.unsqueeze(-1), 0.0)

        output_proposals = output_proposals.masked_fill(~output_proposals_valid, 0.0)

        output_memory = memory
        if memory_padding_mask is not None:
            output_memory = output_memory.masked_fill(memory_padding_mask.unsqueeze(-1), 0.0)

        output_memory = output_memory.masked_fill(~output_proposals_valid, 0.0)

        return (output_memory, output_proposals.to(memory.dtype))

    def postprocess_detections(
        self, class_logits: torch.Tensor, box_regression: torch.Tensor, image_sizes: torch.Tensor
    ) -> list[dict[str, torch.Tensor]]:
        prob = class_logits.sigmoid()
        topk = min(self.num_queries, prob.shape[1] * prob.shape[2])
        topk_values, topk_indexes = torch.topk(prob.view(prob.shape[0], -1), k=topk, dim=1)
        scores = topk_values
        topk_boxes = topk_indexes // prob.shape[2]
        labels = topk_indexes % prob.shape[2]
        labels += 1  # background offset

        boxes = box_ops.box_convert(box_regression, in_fmt="cxcywh", out_fmt="xyxy")
        boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).expand(-1, -1, 4))

        img_h, img_w = image_sizes.unbind(1)
        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
        boxes = boxes * scale_fct[:, None, :]

        detections: list[dict[str, torch.Tensor]] = []
        for s, lbl, b in zip(scores, labels, boxes):
            if self.soft_nms is not None:
                soft_scores, keep = self.soft_nms(b, s, lbl, score_threshold=0.001)
                s[keep] = soft_scores
                b = b[keep]
                s = s[keep]
                lbl = lbl[keep]

            detections.append({"boxes": b, "scores": s, "labels": lbl})

        return detections

    # pylint: disable=too-many-locals
    def forward_net(
        self, x: torch.Tensor, masks: Optional[torch.Tensor]
    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor, torch.Tensor]:
        features: dict[str, torch.Tensor] = self.backbone.detection_features(x)
        stages = self.backbone.return_stages
        feature_list = [features[stage] for stage in stages]

        # Align multi-scale backbone features to a common resolution (LW-DETR expects same spatial size here)
        if len(feature_list) > 1:
            ref_idx = len(feature_list) // 2
            ref_h, ref_w = feature_list[ref_idx].shape[-2:]
            feature_list = [
                (
                    F.interpolate(feat, size=(ref_h, ref_w), mode="nearest")
                    if feat.shape[-2:] != (ref_h, ref_w)
                    else feat
                )
                for feat in feature_list
            ]

        proj_feats = self.projector(feature_list)

        mask_list: Optional[list[torch.Tensor]] = None
        if masks is not None:
            mask_list = []
            for feat in proj_feats:
                mask = F.interpolate(masks[None].float(), size=feat.shape[-2:], mode="nearest").to(torch.bool)[0]
                mask_list.append(mask)

        decoder_feats: list[torch.Tensor] = []
        for idx, proj in enumerate(self.decoder_input_proj):
            decoder_feats.append(proj(proj_feats[idx]))

        spatial_shapes: list[list[int]] = []
        for feat in decoder_feats:
            spatial_shapes.append([feat.shape[-2], feat.shape[-1]])

        level_start_index = [0]
        for shape in spatial_shapes[:-1]:
            level_start_index.append(level_start_index[-1] + shape[0] * shape[1])

        if mask_list is not None:
            valid_ratios = torch.stack([self._get_valid_ratio(m) for m in mask_list], dim=1)
        else:
            B = x.size(0)
            num_levels = len(decoder_feats)
            valid_ratios = torch.ones((B, num_levels, 2), device=proj_feats[0].device, dtype=torch.float32)

        # Two-stage: generate encoder proposals and select top-k per group
        enc_cls_logits: Optional[torch.Tensor] = None
        enc_box_output: Optional[torch.Tensor] = None
        enc_reference_points: Optional[torch.Tensor] = None

        if self.two_stage is True:
            # Flatten decoder features to create memory
            memory_list = []
            mask_flatten_list = []
            for idx, feat in enumerate(decoder_feats):
                memory_list.append(feat.flatten(2).permute(0, 2, 1))
                if mask_list is not None:
                    mask_flatten_list.append(mask_list[idx].flatten(1))

            memory = torch.concat(memory_list, dim=1)
            memory_padding_mask = torch.concat(mask_flatten_list, dim=1) if len(mask_flatten_list) > 0 else None

            # Generate encoder output proposals
            output_memory, output_proposals = self.gen_encoder_output_proposals(
                memory, memory_padding_mask, spatial_shapes
            )

            # Handle case where spatial positions < num_queries
            num_proposals = output_memory.size(1)
            topk = min(self.num_queries, num_proposals)

            if self.training is True:
                # Apply per-group encoder heads and select top-k per group
                enc_cls_list: list[torch.Tensor] = []
                enc_box_list: list[torch.Tensor] = []
                refpoint_list: list[torch.Tensor] = []

                for enc_output, enc_output_norm, enc_out_class, enc_out_bbox in zip(
                    self.enc_output, self.enc_output_norm, self.enc_out_class_embed, self.enc_out_bbox_embed
                ):
                    # Apply per-group encoder output projection
                    output_memory_g = enc_output_norm(enc_output(output_memory))

                    # Compute class and bbox predictions using bbox_reparam formula
                    enc_cls_g = enc_out_class(output_memory_g)
                    enc_box_delta_g = enc_out_bbox(output_memory_g)
                    # bbox_reparam: delta * proposal_wh + proposal_xy for center, exp(delta) * proposal_wh for size
                    enc_box_cxcy_g = enc_box_delta_g[..., :2] * output_proposals[..., 2:] + output_proposals[..., :2]
                    enc_box_wh_g = enc_box_delta_g[..., 2:].exp() * output_proposals[..., 2:]
                    enc_box_g = torch.concat([enc_box_cxcy_g, enc_box_wh_g], dim=-1)

                    # Select top-k proposals for this group
                    _, topk_indices_g = torch.topk(enc_cls_g.max(dim=-1)[0], topk, dim=1)

                    # Gather selected proposals for this group
                    topk_cls_g = torch.gather(
                        enc_cls_g, 1, topk_indices_g.unsqueeze(-1).expand(-1, -1, enc_cls_g.size(-1))
                    )
                    topk_box_g = torch.gather(enc_box_g, 1, topk_indices_g.unsqueeze(-1).expand(-1, -1, 4))

                    enc_cls_list.append(topk_cls_g)
                    enc_box_list.append(topk_box_g)
                    # Reference points (boxes are already in normalized space with bbox_reparam)
                    refpoint_list.append(topk_box_g.detach())

                # Concatenate all groups
                enc_cls_logits = torch.concat(enc_cls_list, dim=1)
                enc_box_output = torch.concat(enc_box_list, dim=1)
                enc_reference_points = torch.concat(refpoint_list, dim=1)
            else:
                # In inference, use only first group (TorchScript-friendly, avoids extra compute)
                enc_output = self.enc_output[0]
                enc_output_norm = self.enc_output_norm[0]
                enc_out_class = self.enc_out_class_embed[0]
                enc_out_bbox = self.enc_out_bbox_embed[0]

                output_memory_g = enc_output_norm(enc_output(output_memory))
                enc_cls_g = enc_out_class(output_memory_g)
                enc_box_delta_g = enc_out_bbox(output_memory_g)
                enc_box_cxcy_g = enc_box_delta_g[..., :2] * output_proposals[..., 2:] + output_proposals[..., :2]
                enc_box_wh_g = enc_box_delta_g[..., 2:].exp() * output_proposals[..., 2:]
                enc_box_g = torch.concat([enc_box_cxcy_g, enc_box_wh_g], dim=-1)

                _, topk_indices_g = torch.topk(enc_cls_g.max(dim=-1)[0], topk, dim=1)
                enc_cls_logits = torch.gather(
                    enc_cls_g, 1, topk_indices_g.unsqueeze(-1).expand(-1, -1, enc_cls_g.size(-1))
                )
                enc_box_output = torch.gather(enc_box_g, 1, topk_indices_g.unsqueeze(-1).expand(-1, -1, 4))
                enc_reference_points = enc_box_output.detach()

        out_bboxes, out_logits = self.decoder(
            decoder_feats,
            spatial_shapes,
            level_start_index,
            valid_ratios,
            padding_mask=mask_list,
            enc_reference_points=enc_reference_points,
            return_intermediates=self.training is True,
        )

        return (enc_cls_logits, enc_box_output, out_bboxes, out_logits)

    def forward(
        self,
        x: torch.Tensor,
        targets: Optional[list[dict[str, torch.Tensor]]] = None,
        masks: Optional[torch.Tensor] = None,
        image_sizes: Optional[list[tuple[int, int]]] = None,
    ) -> tuple[list[dict[str, torch.Tensor]], dict[str, torch.Tensor]]:
        self._input_check(targets)
        image_sizes_tensor = self._to_img_list(x, image_sizes).image_sizes

        enc_cls_logits, enc_box_output, out_bboxes, out_logits = self.forward_net(x, masks)

        losses: dict[str, torch.Tensor] = {}
        detections: list[dict[str, torch.Tensor]] = []
        if self.training is True:
            assert targets is not None, "targets should not be none when in training mode"

            # Compute all losses (encoder + decoder) in compute_loss
            losses = self.compute_loss(
                targets,
                out_logits,
                out_bboxes,
                image_sizes_tensor,
                enc_cls_logits=enc_cls_logits,
                enc_box_output=enc_box_output,
            )
        else:
            detections = self.postprocess_detections(out_logits, out_bboxes, image_sizes_tensor)

        return (detections, losses)


registry.register_model_config(
    "lw_detr",
    LW_DETR,
    config={"projector_scale_factors": [1.0], "decoder_cross_heads": 16},
)
registry.register_model_config(
    "lw_detr_2stg",
    LW_DETR,
    config={"projector_scale_factors": [1.0], "decoder_cross_heads": 16, "two_stage": True},
)
registry.register_model_config(
    "lw_detr_l",
    LW_DETR,
    config={
        "projector_dim": 384,
        "projector_scale_factors": [2.0, 0.5],
        "decoder_dim": 384,
        "decoder_self_heads": 12,
        "decoder_cross_heads": 24,
        "decoder_points": 4,
    },
)
registry.register_model_config(
    "lw_detr_l_2stg",
    LW_DETR,
    config={
        "projector_dim": 384,
        "projector_scale_factors": [2.0, 0.5],
        "decoder_dim": 384,
        "decoder_self_heads": 12,
        "decoder_cross_heads": 24,
        "decoder_points": 4,
        "two_stage": True,
    },
)
