cage_extractor.py

"""
Coarse-to-Fine Gated Extractor (Cage) from VocBulwark (Appendix B.2).

Architecture from paper:
  - Three branches (fine/mid/coarse), each with 4 GSCMs
  - Channel progression: 1 → 32 → 64 → 128 → 128
  - Fine: k=3, s=2, p=1; Mid: k=5, s=2, p=2; Coarse: k=7, s=2, p=3
  - Each GSCM has InstanceNorm + LeakyReLU
  - Each branch output → 1D adaptive average pooling
  - Fusion: GSCM(384→256, k=1, s=1, p=0) → GSCM(256, k=3, s=1, p=1)
  - Dual-path pooling (AAP + AMP) → arithmetic mean
  - Decoder: Linear(256, 512) → Linear(512, watermark_bits)
"""

import torch
import torch.nn as nn
from transformers import PreTrainedModel

from .cage_config import CageExtractorConfig


class GatedSeparableConvModule(nn.Module):
    """GSCM: Gated Separable Convolution Module (Sec 4.3 + Appendix B.2).

    Depthwise separable conv with gating mechanism.
    Uses InstanceNorm and LeakyReLU as per paper.
    Supports stride and different input/output channels.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: int = 1,
    ):
        super().__init__()

        # Content branch (DSC): depthwise → pointwise → InstanceNorm → LeakyReLU
        self.content = nn.Sequential(
            nn.Conv1d(in_channels, in_channels, kernel_size,
                      stride=stride, padding=padding, groups=in_channels),
            nn.Conv1d(in_channels, out_channels, 1),
            nn.InstanceNorm1d(out_channels),
            nn.LeakyReLU(0.2),
        )

        # Gating branch: same structure but Sigmoid output
        self.gate = nn.Sequential(
            nn.Conv1d(in_channels, in_channels, kernel_size,
                      stride=stride, padding=padding, groups=in_channels),
            nn.Conv1d(in_channels, out_channels, 1),
            nn.InstanceNorm1d(out_channels),
            nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.content(x) * self.gate(x)


def _make_branch(kernel_size: int, stride: int, padding: int):
    """Build a 4-GSCM branch with channel progression 1 → 32 → 64 → 128 → 128."""
    return nn.Sequential(
        GatedSeparableConvModule(1, 32, kernel_size, stride, padding),
        GatedSeparableConvModule(32, 64, kernel_size, stride, padding),
        GatedSeparableConvModule(64, 128, kernel_size, stride, padding),
        GatedSeparableConvModule(128, 128, kernel_size, stride, padding),
    )


class CageExtractor(PreTrainedModel):
    config_class = CageExtractorConfig

    def __init__(self, config: CageExtractorConfig):
        super().__init__(config)

        # Three branches processing raw audio directly (paper B.2)
        # Fine-grained: k=3, s=2, p=1
        self.fine_branch = _make_branch(
            kernel_size=config.fine_kernel, stride=2, padding=1)
        # Mid-grained: k=5, s=2, p=2
        self.mid_branch = _make_branch(
            kernel_size=config.mid_kernel, stride=2, padding=2)
        # Coarse-grained: k=7, s=2, p=3
        self.coarse_branch = _make_branch(
            kernel_size=config.coarse_kernel, stride=2, padding=3)

        # Fusion: concat 3 branches (3*128=384) → reduce to 256 → refine
        # GSCM (k=1, s=1, p=0) to reduce channel dimension to 256
        self.fusion_reduce = GatedSeparableConvModule(
            384, 256, kernel_size=1, stride=1, padding=0)
        # GSCM (k=3, s=1, p=1) that maintains the dimensions
        self.fusion_refine = GatedSeparableConvModule(
            256, 256, kernel_size=3, stride=1, padding=1)

        # Decoder: two FC layers (paper B.2)
        self.decoder = nn.Sequential(
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, config.watermark_bits),
        )

    def forward(self, audio: torch.Tensor) -> torch.Tensor:
        """
        Args:
            audio: [B, 1, T] waveform (or [B, T] which gets unsqueezed)
        Returns:
            logits: [B, watermark_bits] raw logits
        """
        if audio.dim() == 2:
            audio = audio.unsqueeze(1)

        # Multi-scale branches (each processes raw audio)
        x_fine = self.fine_branch(audio)      # [B, 128, T']
        x_mid = self.mid_branch(audio)        # [B, 128, T'']
        x_coarse = self.coarse_branch(audio)  # [B, 128, T''']

        # 1D adaptive average pooling to standardize temporal dimension
        pool_t = min(x_fine.shape[-1], x_mid.shape[-1], x_coarse.shape[-1])
        x_fine = nn.functional.adaptive_avg_pool1d(x_fine, pool_t)
        x_mid = nn.functional.adaptive_avg_pool1d(x_mid, pool_t)
        x_coarse = nn.functional.adaptive_avg_pool1d(x_coarse, pool_t)

        # Fusion
        x_fuse = torch.cat([x_fine, x_mid, x_coarse], dim=1)  # [B, 384, pool_t]
        x_fuse = self.fusion_reduce(x_fuse)   # [B, 256, pool_t]
        x_fuse = self.fusion_refine(x_fuse)   # [B, 256, pool_t]

        # Dual-path Pooling: AAP + AMP → arithmetic mean
        x_aap = nn.functional.adaptive_avg_pool1d(x_fuse, 1).squeeze(-1)  # [B, 256]
        x_amp = nn.functional.adaptive_max_pool1d(x_fuse, 1).squeeze(-1)  # [B, 256]
        x_pool = (x_aap + x_amp) / 2

        # Decode
        logits = self.decoder(x_pool)  # [B, watermark_bits]

        return logits